chore: import upstream snapshot with attribution

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

@@ -0,0 +1,9 @@
+from sglang.srt.debug_utils.comparator.aligner.entrypoint.traced_types import (  # noqa: F401
+    TracedAlignerPlan,
+)
+from sglang.srt.debug_utils.comparator.aligner.entrypoint.types import (  # noqa: F401
+    AlignerPlan,
+)
+from sglang.srt.debug_utils.comparator.output_types import ComparisonTensorRecord
+
+ComparisonTensorRecord.model_rebuild()

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

@@ -0,0 +1,4 @@
+from sglang.srt.debug_utils.comparator.entrypoint import main
+
+if __name__ == "__main__":
+    main()

python/sglang/srt/debug_utils/comparator/aligner/axis_aligner.py

@@ -0,0 +1,219 @@
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+from einops import rearrange
+
+from sglang.srt.debug_utils.comparator.dims_spec import (
+    _FUSED_NAME_SEP,
+    SEQ_DIM_NAME,
+    TOKEN_DIM_NAME,
+    DimSpec,
+    _SingletonDimUtil,
+    parse_dims,
+    without_dim_names,
+)
+from sglang.srt.debug_utils.comparator.log_sink import log_sink
+from sglang.srt.debug_utils.comparator.utils import Pair, _FrozenBase
+
+# --- types ---
+
+
+class AxisAlignerPlan(_FrozenBase):
+    pattern: Pair[Optional[str]]  # einops pattern per side, None = no-op
+
+
+# --- planner ---
+
+
+def compute_axis_aligner_plan(
+    dims_str_pair: Pair[Optional[str]],
+) -> Optional[AxisAlignerPlan]:
+    if dims_str_pair.x is None or dims_str_pair.y is None:
+        return None
+
+    dims_pair: Pair[str] = Pair(x=dims_str_pair.x, y=dims_str_pair.y)
+    specs_pair: Pair[list[DimSpec]] = dims_pair.map(lambda s: parse_dims(s).dims)
+
+    if not _semantic_names_match(specs_pair):
+        return None
+
+    # Canonical dim order follows y; fused groups stay fused (flatten, not unflatten).
+    canonical_order: Optional[list[str]] = _build_canonical_order(specs_pair)
+    if canonical_order is None:
+        return None
+
+    pattern: Pair[Optional[str]] = specs_pair.map(
+        lambda specs: _build_side_pattern(specs=specs, canonical_order=canonical_order)
+    )
+
+    if pattern.x is None and pattern.y is None:
+        return None
+
+    return AxisAlignerPlan(pattern=pattern)
+
+
+_SEQ_DIM_EQUIVALENCES: frozenset[frozenset[str]] = frozenset(
+    {
+        frozenset({SEQ_DIM_NAME, TOKEN_DIM_NAME}),  # s ≡ t
+    }
+)
+
+
+def _normalize_dim_name(name: str) -> str:
+    for equiv_set in _SEQ_DIM_EQUIVALENCES:
+        if name in equiv_set:
+            return min(equiv_set)
+    return name
+
+
+def _semantic_names_match(specs_pair: Pair[list[DimSpec]]) -> bool:
+    """Check that both sides share the same semantic name set (ignoring squeeze dims)."""
+    names_pair: Pair[list[str]] = specs_pair.map(_expand_and_skip_squeeze)
+
+    if set(map(_normalize_dim_name, names_pair.x)) == set(
+        map(_normalize_dim_name, names_pair.y)
+    ):
+        return True
+
+    # Local import to avoid circular dependency:
+    # output_types -> aligner/entrypoint/types -> axis_aligner -> output_types
+    from sglang.srt.debug_utils.comparator.output_types import ErrorLog
+
+    log_sink.add(
+        ErrorLog(
+            category="axis_aligner_dim_mismatch",
+            message=(
+                f"AxisAligner: dim name sets differ (x={names_pair.x}, y={names_pair.y}), "
+                f"skipping axis swap"
+            ),
+        )
+    )
+    return False
+
+
+def _expand_and_skip_squeeze(specs: list[DimSpec]) -> list[str]:
+    """Expand DimSpecs to flat semantic names, skipping squeeze dims."""
+    return [
+        name
+        for spec in specs
+        if not _SingletonDimUtil.is_squeeze(spec)
+        for name in spec.sub_dims
+    ]
+
+
+def _build_canonical_order(specs_pair: Pair[list[DimSpec]]) -> Optional[list[str]]:
+    """Build canonical dim order following y, preferring fused representation.
+
+    Each element is either a plain name (``"c"``) or a fused placeholder (``"a___b"``).
+    Fused groups from *either* side are merged — the separate side must flatten.
+    Squeeze dims are excluded.
+
+    Returns ``None`` if the two sides have overlapping but incompatible fused groups
+    (e.g. x fuses ``(a*b)`` while y fuses ``(b*c)``).
+    """
+    # Map each sub-dim name → (placeholder, siblings) from both sides
+    fused_lookup: dict[str, tuple[str, frozenset[str]]] = {}
+    for spec in (*specs_pair.x, *specs_pair.y):
+        if spec.is_fused:
+            placeholder: str = spec.sanitized_name
+            siblings: frozenset[str] = frozenset(spec.sub_dims)
+            for sub_name in spec.sub_dims:
+                existing: Optional[tuple[str, frozenset[str]]] = fused_lookup.get(
+                    sub_name
+                )
+                if existing is not None and existing[1] != siblings:
+                    from sglang.srt.debug_utils.comparator.output_types import ErrorLog
+
+                    log_sink.add(
+                        ErrorLog(
+                            category="axis_aligner_fused_conflict",
+                            message=(
+                                f"AxisAligner: overlapping fused groups for sub-dim {sub_name!r} "
+                                f"({existing[0]} vs {placeholder}), skipping axis alignment"
+                            ),
+                        )
+                    )
+                    return None
+                fused_lookup.setdefault(sub_name, (placeholder, siblings))
+
+    result: list[str] = []
+    consumed: set[str] = set()
+
+    for spec in specs_pair.y:
+        if _SingletonDimUtil.is_squeeze(spec):
+            continue
+
+        names: list[str] = spec.sub_dims
+        if any(n in consumed for n in names):
+            continue
+
+        entry: Optional[tuple[str, frozenset[str]]] = fused_lookup.get(names[0])
+        if entry is not None:
+            fused_placeholder, sibs = entry
+            result.append(fused_placeholder)
+            consumed.update(sibs)
+        else:
+            result.append(_normalize_dim_name(spec.name))
+            consumed.update(names)
+
+    return result
+
+
+def _build_side_pattern(
+    *, specs: list[DimSpec], canonical_order: list[str]
+) -> Optional[str]:
+    """Build an einops pattern for one side to reach ``canonical_order``.
+
+    Fused specs become their placeholder; separate specs that belong to a fused group
+    stay as individual names on the LHS and become ``(a b)`` on the RHS (einops flatten).
+    Squeeze dims (``1``) appear on the LHS but are dropped from the RHS.
+    """
+    source_tokens: list[str] = [spec.sanitized_name for spec in specs]
+
+    # Map normalized dim names back to this side's original names so that
+    # einops patterns use consistent identifiers on LHS and RHS.
+    norm_to_original: dict[str, str] = {
+        _normalize_dim_name(spec.name): spec.name for spec in specs
+    }
+
+    def _to_side_name(token: str) -> str:
+        return norm_to_original.get(token, token)
+
+    # Build per-side target: replace fused placeholders with ``(a b)`` only if this side
+    # has the sub-dims as separate (non-fused) names in the source
+    fused_placeholders: set[str] = {
+        spec.sanitized_name for spec in specs if spec.is_fused
+    }
+    translated_order: list[str] = [_to_side_name(t) for t in canonical_order]
+    target_tokens: list[str] = [
+        (
+            f"({t.replace(_FUSED_NAME_SEP, ' ')})"
+            if _FUSED_NAME_SEP in t and t not in fused_placeholders
+            else t
+        )
+        for t in translated_order
+    ]
+
+    if source_tokens == target_tokens:
+        return None
+
+    return f"{' '.join(source_tokens)} -> {' '.join(target_tokens)}"
+
+
+# --- executor ---
+
+
+def execute_axis_aligner_plan(
+    tensor: torch.Tensor, plan: AxisAlignerPlan, *, side: str
+) -> torch.Tensor:
+    if side not in ("x", "y"):
+        raise ValueError(f"side must be 'x' or 'y', got {side!r}")
+
+    pattern: Optional[str] = plan.pattern.x if side == "x" else plan.pattern.y
+
+    if pattern is not None:
+        tensor = rearrange(without_dim_names(tensor), pattern)
+
+    return tensor

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

@@ -0,0 +1,212 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import NamedTuple, Optional
+
+import torch
+
+from sglang.srt.debug_utils.comparator.aligner.axis_aligner import (
+    execute_axis_aligner_plan,
+)
+from sglang.srt.debug_utils.comparator.aligner.entrypoint.traced_types import (
+    TracedAlignerPlan,
+    TracedSidePlan,
+    TracedStepPlan,
+    TracedSubPlan,
+)
+from sglang.srt.debug_utils.comparator.aligner.entrypoint.types import (
+    AlignerPerStepPlan,
+    AlignerPerStepSubPlan,
+    AlignerPlan,
+)
+from sglang.srt.debug_utils.comparator.aligner.reorderer.executor import (
+    execute_reorderer_plan,
+)
+from sglang.srt.debug_utils.comparator.aligner.reorderer.types import ReordererPlan
+from sglang.srt.debug_utils.comparator.aligner.token_aligner.concat_steps import (
+    execute_token_aligner_concat_steps,
+)
+from sglang.srt.debug_utils.comparator.aligner.token_aligner.smart.executor import (
+    execute_token_aligner,
+)
+from sglang.srt.debug_utils.comparator.aligner.unsharder.executor import (
+    UnsharderResult,
+    execute_unsharder_plan,
+)
+from sglang.srt.debug_utils.comparator.aligner.unsharder.types import UnsharderPlan
+from sglang.srt.debug_utils.comparator.output_types import (
+    ReplicatedCheckResult,
+    ShapeSnapshot,
+)
+from sglang.srt.debug_utils.comparator.utils import Pair
+
+
+class StepPlansResult(NamedTuple):
+    tensors: dict[int, torch.Tensor]
+    checks: list[ReplicatedCheckResult]
+    traced_side: TracedSidePlan
+
+
+class SubPlansResult(NamedTuple):
+    tensor: Optional[torch.Tensor]
+    checks: list[ReplicatedCheckResult]
+    snapshots: list[ShapeSnapshot]
+
+
+@dataclass(frozen=True)
+class AlignerResult:
+    tensors: Optional[Pair[torch.Tensor]]
+    failed_side_xy: Optional[str]  # "x" or "y"; None if success
+    replicated_checks: list[ReplicatedCheckResult] = field(default_factory=list)
+    traced_plan: Optional[TracedAlignerPlan] = None
+
+
+def execute_aligner_plan(
+    *,
+    tensors_pair: Pair[list[torch.Tensor]],
+    plan: AlignerPlan,
+) -> AlignerResult:
+    """Execute unified unshard/reorder + token-align."""
+    all_checks: list[ReplicatedCheckResult] = []
+
+    # Per-side: unshard + reorder -> dict[step, tensor]
+    result_x: StepPlansResult = _execute_step_plans(
+        tensors=tensors_pair.x, step_plans=plan.per_step_plans.x
+    )
+    all_checks.extend(result_x.checks)
+
+    result_y: StepPlansResult = _execute_step_plans(
+        tensors=tensors_pair.y, step_plans=plan.per_step_plans.y
+    )
+    all_checks.extend(result_y.checks)
+
+    traced_plan: TracedAlignerPlan = TracedAlignerPlan(
+        plan=plan,
+        per_side=Pair(x=result_x.traced_side, y=result_y.traced_side),
+    )
+
+    if not result_x.tensors or not result_y.tensors:
+        failed_side_xy: str = "x" if not result_x.tensors else "y"
+        return AlignerResult(
+            tensors=None,
+            failed_side_xy=failed_side_xy,
+            replicated_checks=all_checks,
+            traced_plan=traced_plan,
+        )
+
+    # Cross-side: token alignment (or direct extraction for single-step)
+    step_pair: Pair[dict[int, torch.Tensor]] = Pair(
+        x=result_x.tensors, y=result_y.tensors
+    )
+    combined: Pair[torch.Tensor]
+    if plan.token_aligner_mode == "concat_steps":
+        combined = execute_token_aligner_concat_steps(tensor_of_step_pair=step_pair)
+    elif plan.token_aligner_mode == "smart":
+        assert plan.token_aligner_plan is not None
+        combined = execute_token_aligner(
+            plan=plan.token_aligner_plan,
+            tensor_of_step_pair=step_pair,
+        )
+    else:
+        assert len(result_x.tensors) == 1 and len(result_y.tensors) == 1
+        combined = Pair(
+            x=list(result_x.tensors.values())[0],
+            y=list(result_y.tensors.values())[0],
+        )
+
+    # Cross-side: axis alignment (squeeze singletons + rearrange dim order)
+    if (aligner_plan := plan.axis_aligner_plan) is not None:
+        combined = Pair(
+            x=execute_axis_aligner_plan(tensor=combined.x, plan=aligner_plan, side="x"),
+            y=execute_axis_aligner_plan(tensor=combined.y, plan=aligner_plan, side="y"),
+        )
+
+    return AlignerResult(
+        tensors=combined,
+        failed_side_xy=None,
+        replicated_checks=all_checks,
+        traced_plan=traced_plan,
+    )
+
+
+def _execute_step_plans(
+    tensors: list[torch.Tensor],
+    step_plans: list[AlignerPerStepPlan],
+) -> StepPlansResult:
+    result: dict[int, torch.Tensor] = {}
+    all_checks: list[ReplicatedCheckResult] = []
+    traced_steps: list[TracedStepPlan] = []
+
+    for step_plan in step_plans:
+        step_tensors: list[torch.Tensor] = [
+            tensors[i] for i in step_plan.input_object_indices
+        ]
+        sub_result: SubPlansResult = execute_sub_plans(
+            tensors=step_tensors, plans=step_plan.sub_plans
+        )
+        all_checks.extend(sub_result.checks)
+
+        traced_subs: list[TracedSubPlan] = [
+            TracedSubPlan(plan=sub_plan, snapshot=snapshot)
+            for sub_plan, snapshot in zip(step_plan.sub_plans, sub_result.snapshots)
+        ]
+        traced_steps.append(
+            TracedStepPlan(
+                step=step_plan.step,
+                input_object_indices=step_plan.input_object_indices,
+                sub_plans=traced_subs,
+            )
+        )
+
+        if sub_result.tensor is not None:
+            result[step_plan.step] = sub_result.tensor
+
+    return StepPlansResult(
+        tensors=result,
+        checks=all_checks,
+        traced_side=TracedSidePlan(step_plans=traced_steps),
+    )
+
+
+def execute_sub_plans(
+    tensors: list[torch.Tensor],
+    plans: list[AlignerPerStepSubPlan],
+) -> SubPlansResult:
+    if not tensors:
+        return SubPlansResult(tensor=None, checks=[], snapshots=[])
+
+    if not plans:
+        if len(tensors) != 1:
+            return SubPlansResult(tensor=None, checks=[], snapshots=[])
+        return SubPlansResult(tensor=tensors[0], checks=[], snapshots=[])
+
+    current: list[torch.Tensor] = tensors
+    all_checks: list[ReplicatedCheckResult] = []
+    all_snapshots: list[ShapeSnapshot] = []
+    for plan in plans:
+        input_shapes: list[list[int]] = [list(t.shape) for t in current]
+        current, checks = execute_sub_plan(tensors=current, plan=plan)
+        output_shapes: list[list[int]] = [list(t.shape) for t in current]
+        all_checks.extend(checks)
+        all_snapshots.append(
+            ShapeSnapshot(
+                input_shapes=input_shapes,
+                output_shapes=output_shapes,
+            )
+        )
+
+    assert len(current) == 1
+    return SubPlansResult(tensor=current[0], checks=all_checks, snapshots=all_snapshots)
+
+
+def execute_sub_plan(
+    tensors: list[torch.Tensor],
+    plan: AlignerPerStepSubPlan,
+) -> tuple[list[torch.Tensor], list[ReplicatedCheckResult]]:
+    if isinstance(plan, UnsharderPlan):
+        unsharder_result: UnsharderResult = execute_unsharder_plan(plan, tensors)
+        return unsharder_result.tensors, unsharder_result.replicated_checks
+    elif isinstance(plan, ReordererPlan):
+        return execute_reorderer_plan(plan, tensors), []
+    else:
+        raise NotImplementedError(f"Unknown {plan=}")

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

@@ -0,0 +1,134 @@
+from __future__ import annotations
+
+from typing import Any, Optional
+
+from sglang.srt.debug_utils.comparator.aligner.axis_aligner import (
+    AxisAlignerPlan,
+    compute_axis_aligner_plan,
+)
+from sglang.srt.debug_utils.comparator.aligner.entrypoint.types import (
+    AlignerPerStepPlan,
+    AlignerPerStepSubPlan,
+    AlignerPlan,
+)
+from sglang.srt.debug_utils.comparator.aligner.reorderer.planner import (
+    compute_reorderer_plans,
+)
+from sglang.srt.debug_utils.comparator.aligner.token_aligner.smart.types import (
+    TokenAlignerPlan,
+)
+from sglang.srt.debug_utils.comparator.aligner.unsharder.parallel_info import (
+    normalize_parallel_info,
+)
+from sglang.srt.debug_utils.comparator.aligner.unsharder.planner import (
+    compute_unsharder_plan,
+)
+from sglang.srt.debug_utils.comparator.dims_spec import (
+    DimSpec,
+    DimsSpec,
+    ParallelAxis,
+    _SingletonDimUtil,
+    parse_dims,
+)
+from sglang.srt.debug_utils.comparator.utils import Pair
+
+
+def compute_aligner_plan(
+    *,
+    metas_pair: Pair[list[dict[str, Any]]],
+    token_aligner_mode: Optional[str],
+    token_aligner_plan: Optional[TokenAlignerPlan],
+    thd_seq_lens_by_step_pair: Pair[Optional[dict[int, list[int]]]] = Pair(
+        x=None, y=None
+    ),
+) -> AlignerPlan:
+    dims_str_pair: Pair[Optional[str]] = metas_pair.map(
+        lambda metas: metas[0].get("dims") if metas else None
+    )
+    axis_aligner_plan: Optional[AxisAlignerPlan] = compute_axis_aligner_plan(
+        dims_str_pair=dims_str_pair
+    )
+
+    return AlignerPlan(
+        per_step_plans=Pair(
+            x=_compute_per_step_plans(
+                metas=metas_pair.x,
+                thd_seq_lens_by_step=thd_seq_lens_by_step_pair.x,
+            ),
+            y=_compute_per_step_plans(
+                metas=metas_pair.y,
+                thd_seq_lens_by_step=thd_seq_lens_by_step_pair.y,
+            ),
+        ),
+        token_aligner_mode=token_aligner_mode,
+        token_aligner_plan=token_aligner_plan,
+        axis_aligner_plan=axis_aligner_plan,
+    )
+
+
+def _compute_per_step_plans(
+    metas: list[dict[str, Any]],
+    *,
+    thd_seq_lens_by_step: Optional[dict[int, list[int]]] = None,
+) -> list[AlignerPerStepPlan]:
+    step_to_input_indices: dict[int, list[int]] = {}
+    for i, meta in enumerate(metas):
+        step: int = int(meta["step"])
+        step_to_input_indices.setdefault(step, []).append(i)
+
+    result: list[AlignerPerStepPlan] = []
+    for step in sorted(step_to_input_indices):
+        input_indices: list[int] = step_to_input_indices[step]
+        step_metas: list[dict[str, Any]] = [metas[idx] for idx in input_indices]
+        step_seq_lens: Optional[list[int]] = (
+            thd_seq_lens_by_step.get(step) if thd_seq_lens_by_step is not None else None
+        )
+        plans: list[AlignerPerStepSubPlan] = compute_per_step_sub_plans(
+            metas=step_metas,
+            thd_global_seq_lens=step_seq_lens,
+        )
+        result.append(
+            AlignerPerStepPlan(
+                step=step, input_object_indices=input_indices, sub_plans=plans
+            )
+        )
+
+    return result
+
+
+def compute_per_step_sub_plans(
+    metas: list[dict[str, Any]],
+    *,
+    thd_global_seq_lens: Optional[list[int]] = None,
+) -> list[AlignerPerStepSubPlan]:
+    if not metas or len(metas) == 1:
+        return []
+
+    dims_str = metas[0].get("dims")
+    if dims_str is None:
+        return []
+
+    dims_spec: DimsSpec = parse_dims(dims_str)
+    dim_specs: list[DimSpec] = _SingletonDimUtil.filter_out(dims_spec.dims)
+    replicated_axes: frozenset[ParallelAxis] = dims_spec.replicated_axes
+    parallel_infos = [normalize_parallel_info(meta) for meta in metas]
+
+    dp_axis: ParallelAxis = (
+        ParallelAxis(dims_spec.dp_group_alias)
+        if dims_spec.dp_group_alias
+        else ParallelAxis.DP
+    )
+
+    unsharder_plans = compute_unsharder_plan(
+        dim_specs=dim_specs,
+        parallel_infos=parallel_infos,
+        explicit_replicated_axes=replicated_axes,
+        thd_global_seq_lens=thd_global_seq_lens,
+        dp_filtered_axis=dims_spec.dp_axis,
+    )
+    reorderer_plans = compute_reorderer_plans(
+        dim_specs=dim_specs,
+        parallel_infos=parallel_infos,
+        thd_global_seq_lens=thd_global_seq_lens,
+    )
+    return [*unsharder_plans, *reorderer_plans]

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

@@ -0,0 +1,37 @@
+"""Traced wrapper types that embed execution traces (ShapeSnapshots) into plan nodes.
+
+These types are created *after* execution, pairing each sub-plan with its
+observed shape snapshot so that downstream formatters never need to manually
+zip plan + trace by index.
+"""
+
+from __future__ import annotations
+
+from typing import Optional
+
+from sglang.srt.debug_utils.comparator.aligner.entrypoint.types import (
+    AlignerPerStepSubPlan,
+    AlignerPlan,
+)
+from sglang.srt.debug_utils.comparator.output_types import ShapeSnapshot
+from sglang.srt.debug_utils.comparator.utils import Pair, _StrictBase
+
+
+class TracedSubPlan(_StrictBase):
+    plan: AlignerPerStepSubPlan
+    snapshot: Optional[ShapeSnapshot] = None
+
+
+class TracedStepPlan(_StrictBase):
+    step: int
+    input_object_indices: list[int]
+    sub_plans: list[TracedSubPlan]
+
+
+class TracedSidePlan(_StrictBase):
+    step_plans: list[TracedStepPlan]
+
+
+class TracedAlignerPlan(_StrictBase):
+    plan: AlignerPlan
+    per_side: Pair[TracedSidePlan]

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

@@ -0,0 +1,31 @@
+from __future__ import annotations
+
+from typing import Annotated, Optional, Union
+
+from pydantic import Discriminator
+
+from sglang.srt.debug_utils.comparator.aligner.axis_aligner import AxisAlignerPlan
+from sglang.srt.debug_utils.comparator.aligner.reorderer.types import ReordererPlan
+from sglang.srt.debug_utils.comparator.aligner.token_aligner.smart.types import (
+    TokenAlignerPlan,
+)
+from sglang.srt.debug_utils.comparator.aligner.unsharder.types import UnsharderPlan
+from sglang.srt.debug_utils.comparator.utils import Pair, _FrozenBase
+
+AlignerPerStepSubPlan = Annotated[
+    Union[UnsharderPlan, ReordererPlan],
+    Discriminator("type"),
+]
+
+
+class AlignerPerStepPlan(_FrozenBase):
+    step: int
+    input_object_indices: list[int]
+    sub_plans: list[AlignerPerStepSubPlan]
+
+
+class AlignerPlan(_FrozenBase):
+    per_step_plans: Pair[list[AlignerPerStepPlan]]
+    token_aligner_mode: Optional[str] = None  # "concat_steps" | "smart" | None
+    token_aligner_plan: Optional[TokenAlignerPlan] = None
+    axis_aligner_plan: Optional[AxisAlignerPlan] = None

python/sglang/srt/debug_utils/comparator/aligner/reorderer/executor.py

@@ -0,0 +1,103 @@
+from typing import Optional
+
+import torch
+
+from sglang.srt.debug_utils.comparator.aligner.reorderer.types import (
+    ReordererPlan,
+    ZigzagToNaturalParams,
+    ZigzagToNaturalThdParams,
+)
+from sglang.srt.debug_utils.comparator.dims_spec import (
+    apply_dim_names,
+    get_dim_names,
+    resolve_dim_by_name,
+    without_dim_names,
+)
+
+
+def execute_reorderer_plan(
+    plan: ReordererPlan,
+    tensors: list[torch.Tensor],
+) -> list[torch.Tensor]:
+    if isinstance(plan.params, ZigzagToNaturalThdParams):
+        thd_dim: int = resolve_dim_by_name(tensors[0], plan.params.dim_name)
+        return [
+            _reorder_zigzag_to_natural_thd(
+                tensor,
+                dim=thd_dim,
+                cp_size=plan.params.cp_size,
+                seq_lens=plan.params.seq_lens,
+            )
+            for tensor in tensors
+        ]
+
+    if isinstance(plan.params, ZigzagToNaturalParams):
+        dim: int = resolve_dim_by_name(tensors[0], plan.params.dim_name)
+        return [
+            _reorder_zigzag_to_natural(tensor, dim=dim, cp_size=plan.params.cp_size)
+            for tensor in tensors
+        ]
+
+    raise ValueError(f"Unsupported reorderer params type: {type(plan.params).__name__}")
+
+
+def _reorder_zigzag_to_natural_thd(
+    tensor: torch.Tensor, *, dim: int, cp_size: int, seq_lens: list[int]
+) -> torch.Tensor:
+    """Undo CP zigzag interleaving for THD (packed-seq) format.
+
+    Each seq in seq_lens is independently reordered from zigzag to natural order
+    along the given dim.
+    """
+    names: tuple[Optional[str], ...] = get_dim_names(tensor)
+    stripped: torch.Tensor = without_dim_names(tensor)
+
+    split_sizes: list[int] = list(seq_lens)
+    remainder: int = stripped.shape[dim] - sum(split_sizes)
+    if remainder < 0:
+        raise ValueError(
+            f"sum(seq_lens)={sum(split_sizes)} exceeds tensor dim size "
+            f"{stripped.shape[dim]} along dim={dim}"
+        )
+    if remainder > 0:
+        split_sizes.append(remainder)
+
+    segments: list[torch.Tensor] = list(stripped.split(split_sizes, dim=dim))
+
+    reordered_segments: list[torch.Tensor] = [
+        _reorder_zigzag_to_natural(seg, dim=dim, cp_size=cp_size)
+        for seg in segments[: len(seq_lens)]
+    ]
+
+    # Tail padding — pass through unchanged
+    if remainder > 0:
+        reordered_segments.append(segments[-1])
+
+    result: torch.Tensor = torch.cat(reordered_segments, dim=dim)
+
+    if names[0] is not None:
+        result = apply_dim_names(result, list(names))
+    return result
+
+
+def _reorder_zigzag_to_natural(
+    tensor: torch.Tensor, *, dim: int, cp_size: int
+) -> torch.Tensor:
+    """Undo CP zigzag interleaving, restoring natural chunk order.
+
+    Generalized from Megatron-LM _undo_attention_load_balancing
+    (megatron/core/ssm/mamba_context_parallel.py:360-373).
+    """
+    names: tuple[Optional[str], ...] = get_dim_names(tensor)
+    stripped: torch.Tensor = without_dim_names(tensor)
+
+    num_chunks: int = cp_size * 2
+    chunks: tuple[torch.Tensor, ...] = stripped.chunk(num_chunks, dim=dim)
+    order: list[int] = [2 * i for i in range(cp_size)] + [
+        num_chunks - 2 * i - 1 for i in range(cp_size)
+    ]
+    result: torch.Tensor = torch.cat([chunks[i] for i in order], dim=dim)
+
+    if names[0] is not None:
+        result = apply_dim_names(result, list(names))
+    return result

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

@@ -0,0 +1,67 @@
+from typing import Optional
+
+from sglang.srt.debug_utils.comparator.aligner.reorderer.types import (
+    ReordererPlan,
+    ZigzagToNaturalParams,
+    ZigzagToNaturalThdParams,
+)
+from sglang.srt.debug_utils.comparator.aligner.unsharder.types import AxisInfo
+from sglang.srt.debug_utils.comparator.dims_spec import (
+    SEQ_DIM_NAME,
+    TOKEN_DIM_NAME,
+    DimSpec,
+    Ordering,
+    ParallelAxis,
+)
+
+_ALLOWED_ZIGZAG_DIM_NAMES: set[str] = {SEQ_DIM_NAME, TOKEN_DIM_NAME}
+
+
+def compute_reorderer_plans(
+    dim_specs: list[DimSpec],
+    parallel_infos: list[dict[ParallelAxis, AxisInfo]],
+    *,
+    thd_global_seq_lens: Optional[list[int]] = None,
+) -> list[ReordererPlan]:
+    plans: list[ReordererPlan] = []
+
+    for spec in dim_specs:
+        for modifier in spec.parallel_modifiers:
+            if modifier.ordering is None or modifier.ordering == Ordering.NATURAL:
+                continue
+
+            if spec.name not in _ALLOWED_ZIGZAG_DIM_NAMES:
+                raise ValueError(
+                    f"Zigzag ordering is only supported on sequence dims "
+                    f"(dim name must be one of "
+                    f"{sorted(_ALLOWED_ZIGZAG_DIM_NAMES)}), "
+                    f"but got dim name {spec.name!r} in {spec}"
+                )
+
+            if modifier.ordering != Ordering.ZIGZAG:
+                raise ValueError(
+                    f"Unsupported ordering {modifier.ordering!r} for dim {spec.name!r}"
+                )
+            axis_size: int = parallel_infos[0][modifier.axis].axis_size
+
+            if spec.name == TOKEN_DIM_NAME:
+                if thd_global_seq_lens is None:
+                    raise ValueError(
+                        "thd_global_seq_lens is required for zigzag reorder on 't' dimension"
+                    )
+                params = ZigzagToNaturalThdParams(
+                    dim_name=spec.name,
+                    cp_size=axis_size,
+                    seq_lens=thd_global_seq_lens,
+                )
+            elif spec.name == SEQ_DIM_NAME:
+                params = ZigzagToNaturalParams(dim_name=spec.name, cp_size=axis_size)
+            else:
+                raise ValueError(
+                    f"Unsupported zigzag dim name {spec.name!r}, "
+                    f"expected one of {sorted(_ALLOWED_ZIGZAG_DIM_NAMES)}"
+                )
+
+            plans.append(ReordererPlan(params=params))
+
+    return plans

python/sglang/srt/debug_utils/comparator/aligner/reorderer/types.py

@@ -0,0 +1,29 @@
+from typing import Annotated, Literal, Union
+
+from pydantic import Field
+
+from sglang.srt.debug_utils.comparator.utils import _FrozenBase
+
+
+class ZigzagToNaturalParams(_FrozenBase):
+    op: Literal["zigzag_to_natural"] = "zigzag_to_natural"
+    dim_name: str
+    cp_size: int
+
+
+class ZigzagToNaturalThdParams(_FrozenBase):
+    op: Literal["zigzag_to_natural_thd"] = "zigzag_to_natural_thd"
+    dim_name: str
+    cp_size: int
+    seq_lens: list[int]  # unshard-ed per-seq token counts, e.g. [100, 64, 92]
+
+
+ReordererParams = Annotated[
+    Union[ZigzagToNaturalParams, ZigzagToNaturalThdParams],
+    Field(discriminator="op"),
+]
+
+
+class ReordererPlan(_FrozenBase):
+    type: Literal["reorderer"] = "reorderer"
+    params: ReordererParams

python/sglang/srt/debug_utils/comparator/aligner/token_aligner/concat_steps/__init__.py

@@ -0,0 +1,7 @@
+from sglang.srt.debug_utils.comparator.aligner.token_aligner.concat_steps.executor import (
+    execute_token_aligner_concat_steps,
+)
+
+__all__ = [
+    "execute_token_aligner_concat_steps",
+]

python/sglang/srt/debug_utils/comparator/aligner/token_aligner/concat_steps/executor.py

@@ -0,0 +1,45 @@
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+
+from sglang.srt.debug_utils.comparator.dims_spec import (
+    SEQ_DIM_NAME,
+    TOKEN_DIM_NAME,
+    get_dim_names,
+)
+from sglang.srt.debug_utils.comparator.utils import Pair
+
+_UNNAMED_TOKEN_DIM_FALLBACK: int = 0
+
+
+def execute_token_aligner_concat_steps(
+    tensor_of_step_pair: Pair[dict[int, torch.Tensor]],
+) -> Pair[torch.Tensor]:
+    """Concat all steps in order, then truncate to min(total_x, total_y) tokens."""
+    some_tensor: torch.Tensor = next(iter(tensor_of_step_pair.x.values()))
+    token_dim: int = _resolve_token_dim(some_tensor)
+
+    concatenated: Pair[torch.Tensor] = tensor_of_step_pair.map(
+        lambda d: _concat_steps(d, dim=token_dim)
+    )
+    common: int = min(concatenated.x.shape[token_dim], concatenated.y.shape[token_dim])
+    return concatenated.map(lambda t: t.narrow(dim=token_dim, start=0, length=common))
+
+
+def _resolve_token_dim(tensor: torch.Tensor) -> int:
+    """Find the token/seq dim index. Falls back to dim 0 for unnamed tensors or
+    tensors without a recognised token/seq dim."""
+    names: tuple[Optional[str], ...] = get_dim_names(tensor)
+    if names[0] is None:
+        return _UNNAMED_TOKEN_DIM_FALLBACK
+    for candidate in (TOKEN_DIM_NAME, SEQ_DIM_NAME):
+        if candidate in names:
+            return list(names).index(candidate)
+
+    return _UNNAMED_TOKEN_DIM_FALLBACK
+
+
+def _concat_steps(tensor_of_step: dict[int, torch.Tensor], *, dim: int) -> torch.Tensor:
+    return torch.cat([tensor_of_step[s] for s in sorted(tensor_of_step)], dim=dim)

python/sglang/srt/debug_utils/comparator/aligner/token_aligner/concat_steps/thd_seq_lens_loader.py

@@ -0,0 +1,43 @@
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Optional
+
+import polars as pl
+
+from sglang.srt.debug_utils.comparator.aligner.token_aligner.smart.aux_loader import (
+    _detect_plugin,
+    _load_and_align_aux_tensor,
+)
+from sglang.srt.debug_utils.comparator.aligner.token_aligner.smart.aux_plugins import (
+    _AuxFrameworkPlugin,
+)
+
+
+def load_thd_seq_lens_only(
+    dump_path: Path, df: pl.DataFrame
+) -> Optional[dict[int, list[int]]]:
+    plugin: Optional[_AuxFrameworkPlugin] = _detect_plugin(df, dump_path=dump_path)
+    if plugin is None or not plugin.cp_sharded_names:
+        return None
+
+    non_cp_tensor_names: set[str] = (
+        set(df["name"].unique().to_list()) & plugin.tensor_names
+    ) - plugin.cp_sharded_names
+    steps: list[int] = sorted(df["step"].unique().to_list())
+
+    result: dict[int, list[int]] = {}
+    for step in steps:
+        step_data: dict[str, object] = {}
+        for name in non_cp_tensor_names:
+            tensor = _load_and_align_aux_tensor(
+                name=name, step=step, df=df, dump_path=dump_path, plugin=plugin
+            )
+            if tensor is not None:
+                step_data[name] = tensor
+
+        seq_lens: Optional[list[int]] = plugin.extract_global_seq_lens(step_data)
+        if seq_lens is not None:
+            result[step] = seq_lens
+
+    return result or None

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

@@ -0,0 +1,132 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Literal, Optional
+
+import polars as pl
+
+from sglang.srt.debug_utils.comparator.aligner.token_aligner.concat_steps.thd_seq_lens_loader import (
+    load_thd_seq_lens_only,
+)
+from sglang.srt.debug_utils.comparator.aligner.token_aligner.smart.aux_loader import (
+    has_aux_tensors,
+    load_and_normalize_aux,
+)
+from sglang.srt.debug_utils.comparator.aligner.token_aligner.smart.planner import (
+    compute_token_aligner_plan,
+)
+from sglang.srt.debug_utils.comparator.aligner.token_aligner.smart.seq_info_builder import (
+    build_seqs_info,
+)
+from sglang.srt.debug_utils.comparator.aligner.token_aligner.smart.types import (
+    TokenAlignerGlobalAux,
+    TokenAlignerPlan,
+    TokenAlignerSeqsInfo,
+)
+from sglang.srt.debug_utils.comparator.log_sink import log_sink
+from sglang.srt.debug_utils.comparator.output_types import InfoLog
+from sglang.srt.debug_utils.comparator.utils import Pair
+
+_NONE_THD: Pair[Optional[dict[int, list[int]]]] = Pair(x=None, y=None)
+
+
+TokenAlignerMode = Literal["concat_steps", "smart"]
+
+
+@dataclass(frozen=True)
+class TokenAlignerResult:
+    """Result of token aligner computation, bundling mode + plan with THD metadata."""
+
+    mode: Optional[TokenAlignerMode]
+    plan: Optional[TokenAlignerPlan]
+    thd_seq_lens_by_step_pair: Pair[Optional[dict[int, list[int]]]]
+
+
+def compute_maybe_token_aligner_result(
+    *,
+    dir_pair: Pair[Path],
+    dfs: Pair[pl.DataFrame],
+    token_aligner_mode: Optional[TokenAlignerMode],
+) -> TokenAlignerResult:
+    if token_aligner_mode is None:
+        return TokenAlignerResult(
+            mode=None, plan=None, thd_seq_lens_by_step_pair=_NONE_THD
+        )
+
+    if token_aligner_mode == "concat_steps":
+        thd_pair: Pair[Optional[dict[int, list[int]]]] = _load_thd_seq_lens_pair(
+            dir_pair=dir_pair, dfs=dfs
+        )
+        return TokenAlignerResult(
+            mode="concat_steps", plan=None, thd_seq_lens_by_step_pair=thd_pair
+        )
+    elif token_aligner_mode == "smart":
+        if not (has_aux_tensors(dfs.x) and has_aux_tensors(dfs.y)):
+            log_sink.add(
+                InfoLog(
+                    category="aux_tensors_missing",
+                    message="Aux tensors missing, skipping token alignment",
+                )
+            )
+            return TokenAlignerResult(
+                mode=None, plan=None, thd_seq_lens_by_step_pair=_NONE_THD
+            )
+
+        return _build_smart_result(dir_pair=dir_pair, dfs=dfs)
+    else:
+        raise NotImplementedError(f"Unknown {token_aligner_mode=}")
+
+
+def _build_smart_result(
+    *,
+    dir_pair: Pair[Path],
+    dfs: Pair[pl.DataFrame],
+) -> TokenAlignerResult:
+    """Load aux tensors, build token indices, and compute the alignment plan."""
+    aux_pair: Pair[Optional[TokenAlignerGlobalAux]] = Pair(
+        x=load_and_normalize_aux(dump_path=dir_pair.x, df=dfs.x),
+        y=load_and_normalize_aux(dump_path=dir_pair.y, df=dfs.y),
+    )
+
+    thd_seq_lens_by_step_pair: Pair[Optional[dict[int, list[int]]]] = aux_pair.map(
+        lambda aux: aux.thd_seq_lens_by_step if aux is not None else None
+    )
+
+    if aux_pair.x is None or aux_pair.y is None:
+        log_sink.add(
+            InfoLog(
+                category="framework_detection_failed",
+                message="Framework detection failed, skipping token alignment",
+            )
+        )
+        return TokenAlignerResult(
+            mode=None,
+            plan=None,
+            thd_seq_lens_by_step_pair=thd_seq_lens_by_step_pair,
+        )
+
+    global_aux: Pair[TokenAlignerGlobalAux] = Pair(x=aux_pair.x, y=aux_pair.y)
+
+    seqs_info: Pair[TokenAlignerSeqsInfo] = global_aux.map(build_seqs_info)
+
+    plan: Optional[TokenAlignerPlan] = compute_token_aligner_plan(
+        seqs_info_pair=seqs_info
+    )
+    return TokenAlignerResult(
+        mode="smart",
+        plan=plan,
+        thd_seq_lens_by_step_pair=thd_seq_lens_by_step_pair,
+    )
+
+
+def _load_thd_seq_lens_pair(
+    *,
+    dir_pair: Pair[Path],
+    dfs: Pair[pl.DataFrame],
+) -> Pair[Optional[dict[int, list[int]]]]:
+    """Load only thd_seq_lens for each side (lightweight, no full aux loading)."""
+    return Pair(
+        x=load_thd_seq_lens_only(dump_path=dir_pair.x, df=dfs.x),
+        y=load_thd_seq_lens_only(dump_path=dir_pair.y, df=dfs.y),
+    )

python/sglang/srt/debug_utils/comparator/aligner/token_aligner/smart/aux_loader.py

@@ -0,0 +1,286 @@
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any, Optional
+
+import polars as pl
+import torch
+
+from sglang.srt.debug_utils.comparator.aligner.entrypoint.executor import (
+    execute_sub_plans,
+)
+from sglang.srt.debug_utils.comparator.aligner.entrypoint.planner import (
+    compute_per_step_sub_plans,
+)
+from sglang.srt.debug_utils.comparator.aligner.token_aligner.smart.aux_plugins import (
+    AUX_NAMES,
+    _AuxFrameworkPlugin,
+    _plugins,
+)
+from sglang.srt.debug_utils.comparator.aligner.token_aligner.smart.types import (
+    TokenAlignerGlobalAux,
+    TokenAlignerStepAux,
+)
+from sglang.srt.debug_utils.comparator.aligner.unsharder.parallel_info import (
+    normalize_parallel_info,
+)
+from sglang.srt.debug_utils.comparator.dims_spec import (
+    ParallelAxis,
+    TokenLayout,
+    apply_dim_names,
+    resolve_dim_names,
+    without_dim_names,
+)
+from sglang.srt.debug_utils.comparator.dp_utils import filter_to_non_empty_dp_rank
+from sglang.srt.debug_utils.comparator.log_sink import log_sink
+from sglang.srt.debug_utils.comparator.output_types import ErrorLog, InfoLog
+from sglang.srt.debug_utils.dump_loader import ValueWithMeta, filter_rows
+
+# re-export for existing callers
+__all__ = [
+    "AUX_NAMES",
+    "has_aux_tensors",
+    "load_and_normalize_aux",
+]
+
+
+def load_and_normalize_aux(
+    dump_path: Path, df: pl.DataFrame
+) -> Optional[TokenAlignerGlobalAux]:
+    """Bootstrap: load, unshard, and normalize auxiliary tensors for one side."""
+    plugin: Optional[_AuxFrameworkPlugin] = _detect_plugin(df, dump_path=dump_path)
+    if plugin is None:
+        return None
+
+    available_names: set[str] = set(df["name"].unique().to_list()) & plugin.all_names
+    steps: list[int] = sorted(df["step"].unique().to_list())
+    tensor_names: set[str] = available_names & plugin.tensor_names
+    non_tensor_names: set[str] = available_names & plugin.non_tensor_names
+
+    steps_data: dict[int, dict[str, object]] = {}
+    thd_seq_lens_by_step: dict[int, list[int]] = {}
+    for step in steps:
+        step_data, thd_seq_lens = _load_step_data(
+            step=step,
+            tensor_names=tensor_names,
+            non_tensor_names=non_tensor_names,
+            df=df,
+            dump_path=dump_path,
+            plugin=plugin,
+        )
+        if step_data:
+            steps_data[step] = step_data
+        if thd_seq_lens is not None:
+            thd_seq_lens_by_step[step] = thd_seq_lens
+
+    layout: TokenLayout = plugin.detect_layout(steps_data)
+
+    step_auxs: dict[int, TokenAlignerStepAux] = {
+        step: plugin.compute_step_aux(step_data, layout=layout, step=step)
+        for step, step_data in steps_data.items()
+    }
+
+    return TokenAlignerGlobalAux(
+        step_auxs=step_auxs,
+        framework=plugin.name,
+        layout=layout,
+        thd_seq_lens_by_step=thd_seq_lens_by_step or None,
+    )
+
+
+def has_aux_tensors(df: pl.DataFrame) -> bool:
+    """Check if the DataFrame contains the minimum auxiliary tensors for alignment."""
+    names: set[str] = set(df["name"].unique().to_list())
+    return any(plugin.has_required_names(names) for plugin in _plugins)
+
+
+def _detect_plugin(df: pl.DataFrame, dump_path: Path) -> Optional[_AuxFrameworkPlugin]:
+    names: set[str] = set(df["name"].unique().to_list())
+
+    for plugin in _plugins:
+        if names & plugin.discriminating_names:
+            return plugin
+
+    first_row: dict = df.row(0, named=True)
+    value: ValueWithMeta = ValueWithMeta.load(dump_path / first_row["filename"])
+
+    for plugin in _plugins:
+        if f"{plugin.name}_parallel_info" in value.meta:
+            return plugin
+
+    return None
+
+
+def _load_step_data(
+    *,
+    step: int,
+    tensor_names: set[str],
+    non_tensor_names: set[str],
+    df: pl.DataFrame,
+    dump_path: Path,
+    plugin: _AuxFrameworkPlugin,
+) -> tuple[dict[str, object], Optional[list[int]]]:
+    """Load all tensor and non-tensor aux values for a single step.
+
+    Two-pass loading: non-CP-sharded tensors first (to obtain cu_seqlens_q
+    for seq_lens), then CP-sharded tensors with seq_lens for THD unshard/reorder.
+
+    Returns (step_data, thd_global_seq_lens).
+    """
+    result: dict[str, object] = {}
+
+    # Pass 0: non-tensor values
+    for name in non_tensor_names:
+        value = _load_non_tensor_aux(name=name, step=step, df=df, dump_path=dump_path)
+        if value is not None:
+            result[name] = value
+
+    # Pass 1: non-CP-sharded tensors (e.g. cu_seqlens_q, seq_lens)
+    non_cp_tensor_names: set[str] = tensor_names - plugin.cp_sharded_names
+    cp_tensor_names: set[str] = tensor_names & plugin.cp_sharded_names
+
+    for name in non_cp_tensor_names:
+        tensor = _load_and_align_aux_tensor(
+            name=name, step=step, df=df, dump_path=dump_path, plugin=plugin
+        )
+        if tensor is not None:
+            result[name] = tensor
+
+    # Derive global seq_lens for THD unshard (framework-specific extraction)
+    thd_global_seq_lens: Optional[list[int]] = plugin.extract_global_seq_lens(result)
+
+    # Pass 2: CP-sharded tensors (input_ids, position_ids, etc.)
+    for name in cp_tensor_names:
+        tensor = _load_and_align_aux_tensor(
+            name=name,
+            step=step,
+            df=df,
+            dump_path=dump_path,
+            plugin=plugin,
+            thd_global_seq_lens=thd_global_seq_lens,
+        )
+        if tensor is not None:
+            result[name] = tensor
+
+    return result, thd_global_seq_lens
+
+
+def _load_non_tensor_aux(
+    *, name: str, step: int, df: pl.DataFrame, dump_path: Path
+) -> Optional[object]:
+    """Load a non-tensor auxiliary value for a step, validating consistency across ranks."""
+    rows = filter_rows(df, conditions={"name": name, "step": step})
+    if not rows:
+        return None
+
+    loaded: list[ValueWithMeta] = [
+        ValueWithMeta.load(dump_path / r["filename"]) for r in rows
+    ]
+    loaded = filter_to_non_empty_dp_rank(loaded, dp_axis=ParallelAxis.DP)
+
+    if len(loaded) > 1:
+        first_value = loaded[0].value
+        for i, item in enumerate(loaded[1:], start=1):
+            if item.value != first_value:
+                log_sink.add(
+                    ErrorLog(
+                        category=f"{name}_mismatch",
+                        message=(
+                            f"{name} mismatch across ranks: rank 0 has {first_value}, "
+                            f"rank {i} has {item.value}"
+                        ),
+                    )
+                )
+                break
+
+    return loaded[0].value
+
+
+def _load_and_align_aux_tensor(
+    *,
+    name: str,
+    step: int,
+    df: pl.DataFrame,
+    dump_path: Path,
+    plugin: _AuxFrameworkPlugin,
+    thd_global_seq_lens: Optional[list[int]] = None,
+) -> Optional[torch.Tensor]:
+    """Load an auxiliary tensor for (name, step), align if needed."""
+    rows = filter_rows(df, conditions={"name": name, "step": step})
+    if not rows:
+        return None
+
+    loaded: list[ValueWithMeta] = [
+        ValueWithMeta.load(dump_path / r["filename"]) for r in rows
+    ]
+    loaded = filter_to_non_empty_dp_rank(loaded, dp_axis=ParallelAxis.DP)
+
+    tensors: list[torch.Tensor] = [
+        item.value for item in loaded if isinstance(item.value, torch.Tensor)
+    ]
+    if not tensors:
+        return None
+
+    if len(tensors) == 1:
+        return tensors[0]
+
+    metas: list[dict[str, Any]] = [item.meta for item in loaded]
+    metas = _ensure_dims_in_metas(
+        name=name, plugin=plugin, metas=metas, ndim=tensors[0].ndim
+    )
+
+    sub_plans = compute_per_step_sub_plans(
+        metas=metas,
+        thd_global_seq_lens=(
+            thd_global_seq_lens if name in plugin.cp_sharded_names else None
+        ),
+    )
+    if sub_plans:
+        dims_str: Optional[str] = metas[0].get("dims")
+        if dims_str is not None:
+            dim_names: list[str] = resolve_dim_names(dims_str)
+            tensors = [apply_dim_names(t, dim_names) for t in tensors]
+
+        sub_result = execute_sub_plans(tensors=tensors, plans=sub_plans)
+        assert sub_result.tensor is not None
+        return without_dim_names(
+            sub_result.tensor
+        )  # strip named dims before returning to plugin
+
+    log_sink.add(
+        InfoLog(
+            category="aux_no_dims",
+            message=(
+                f"aux tensor '{name}' has {len(tensors)} ranks "
+                f"but no dims metadata, using rank 0 only"
+            ),
+        )
+    )
+    return tensors[0]
+
+
+def _ensure_dims_in_metas(
+    *,
+    name: str,
+    plugin: _AuxFrameworkPlugin,
+    metas: list[dict[str, Any]],
+    ndim: int,
+) -> list[dict[str, Any]]:
+    """Inject inferred dims into metas if not already present.
+
+    Returns metas unchanged if dims is already set, or a new list with dims
+    injected if inference succeeds for CP-sharded tensors.
+    """
+    if metas[0].get("dims") is not None:
+        return metas
+
+    parallel_infos = [normalize_parallel_info(m) for m in metas]
+    has_cp: bool = any(ParallelAxis.CP in info for info in parallel_infos)
+    if not has_cp:
+        return metas
+
+    if name in plugin.cp_sharded_names:
+        inferred_dims: str = plugin.infer_cp_sharded_dims(name=name, ndim=ndim)
+        return [{**m, "dims": inferred_dims} for m in metas]
+
+    return metas

python/sglang/srt/debug_utils/comparator/aligner/token_aligner/smart/aux_plugins.py

@@ -0,0 +1,292 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import Optional
+
+import torch
+
+from sglang.srt.debug_utils.comparator.aligner.token_aligner.smart.types import (
+    PositionalSeqId,
+    SeqId,
+    SGLangSeqId,
+    TokenAlignerStepAux,
+)
+from sglang.srt.debug_utils.comparator.dims_spec import TokenLayout
+from sglang.srt.debug_utils.comparator.log_sink import log_sink
+from sglang.srt.debug_utils.comparator.output_types import InfoLog
+
+# ── plugin ABC ─────────────────────────────────────────────────────
+
+
+class _AuxFrameworkPlugin(ABC):
+    @property
+    @abstractmethod
+    def name(self) -> str: ...
+
+    @property
+    @abstractmethod
+    def tensor_names(self) -> frozenset[str]: ...
+
+    @property
+    @abstractmethod
+    def non_tensor_names(self) -> frozenset[str]: ...
+
+    @property
+    def cp_sharded_names(self) -> frozenset[str]:
+        return frozenset()
+
+    @property
+    def discriminating_names(self) -> frozenset[str]:
+        """Field names unique to this framework (excluding shared names like input_ids)."""
+        return frozenset()
+
+    @abstractmethod
+    def detect_layout(self, raw: dict[int, dict[str, object]]) -> TokenLayout: ...
+
+    @abstractmethod
+    def compute_step_aux(
+        self, step_data: dict[str, object], *, layout: TokenLayout, step: int
+    ) -> TokenAlignerStepAux: ...
+
+    @abstractmethod
+    def has_required_names(self, names: set[str]) -> bool:
+        """Whether the minimum set of aux names needed for alignment is present."""
+        ...
+
+    @property
+    def all_names(self) -> frozenset[str]:
+        return self.tensor_names | self.non_tensor_names
+
+    def extract_global_seq_lens(
+        self, step_data: dict[str, object]
+    ) -> Optional[list[int]]:
+        """Extract per-seq token counts from loaded step data.
+
+        Returns None if this framework doesn't support THD / no relevant data available.
+        """
+        return None
+
+    def infer_cp_sharded_dims(self, name: str, ndim: int) -> str:
+        """Infer dims string for a CP-sharded aux tensor based on its ndim."""
+        raise NotImplementedError(
+            f"infer_cp_sharded_dims not implemented for {type(self).__name__}"
+        )
+
+
+# ── sglang plugin ─────────────────────────────────────────────────
+
+
+class _SGLangPlugin(_AuxFrameworkPlugin):
+    @property
+    def name(self) -> str:
+        return "sglang"
+
+    @property
+    def tensor_names(self) -> frozenset[str]:
+        return frozenset({"input_ids", "positions", "seq_lens", "req_pool_indices"})
+
+    @property
+    def non_tensor_names(self) -> frozenset[str]:
+        return frozenset({"rids"})
+
+    @property
+    def cp_sharded_names(self) -> frozenset[str]:
+        return frozenset({"input_ids", "positions"})
+
+    @property
+    def discriminating_names(self) -> frozenset[str]:
+        return frozenset({"seq_lens", "positions", "req_pool_indices", "rids"})
+
+    def has_required_names(self, names: set[str]) -> bool:
+        return "input_ids" in names and "seq_lens" in names
+
+    def detect_layout(self, raw: dict[int, dict[str, object]]) -> TokenLayout:
+        return TokenLayout.T
+
+    def extract_global_seq_lens(
+        self, step_data: dict[str, object]
+    ) -> Optional[list[int]]:
+        if not self.cp_sharded_names:
+            return None
+
+        seq_lens = step_data.get("seq_lens")
+        if not isinstance(seq_lens, torch.Tensor):
+            return None
+
+        return seq_lens.tolist()
+
+    def infer_cp_sharded_dims(self, name: str, ndim: int) -> str:
+        """Infer dims for CP-sharded aux tensors.
+
+        NOTE: assumes zigzag ordering — natural-order CP without explicit dims
+        will be mishandled. Callers should set dims explicitly for non-zigzag CP.
+        """
+        if ndim == 1:
+            return "t[cp:zigzag]"
+        raise ValueError(
+            f"SGLang: cannot infer dims for CP-sharded '{name}' with ndim={ndim}"
+        )
+
+    def compute_step_aux(
+        self, step_data: dict[str, object], *, layout: TokenLayout, step: int
+    ) -> TokenAlignerStepAux:
+        input_ids = step_data["input_ids"]
+        positions = step_data["positions"]
+        seq_lens = step_data["seq_lens"]
+        rids_raw = step_data.get("rids")
+
+        assert isinstance(
+            input_ids, torch.Tensor
+        ), f"input_ids: expected Tensor, got {type(input_ids)}"
+        assert isinstance(
+            positions, torch.Tensor
+        ), f"positions: expected Tensor, got {type(positions)}"
+        assert isinstance(
+            seq_lens, torch.Tensor
+        ), f"seq_lens: expected Tensor, got {type(seq_lens)}"
+
+        seq_lens_list: list[int] = seq_lens.tolist()
+        num_seqs: int = len(seq_lens_list)
+
+        seq_ids: list[SeqId]
+        if rids_raw is not None and isinstance(rids_raw, (list, tuple)):
+            seq_ids = [SGLangSeqId(rid=str(r)) for r in rids_raw]
+        else:
+            seq_ids = [PositionalSeqId(step=step, seq_index=i) for i in range(num_seqs)]
+
+        return TokenAlignerStepAux(
+            input_ids=input_ids.tolist(),
+            positions=positions.tolist(),
+            seq_lens=seq_lens_list,
+            seq_ids=seq_ids,
+        )
+
+
+# ── megatron plugin ───────────────────────────────────────────────
+
+
+class _MegatronPlugin(_AuxFrameworkPlugin):
+    @property
+    def name(self) -> str:
+        return "megatron"
+
+    @property
+    def tensor_names(self) -> frozenset[str]:
+        return frozenset({"input_ids", "position_ids", "cu_seqlens_q", "cu_seqlens_kv"})
+
+    @property
+    def non_tensor_names(self) -> frozenset[str]:
+        return frozenset({"qkv_format"})
+
+    @property
+    def cp_sharded_names(self) -> frozenset[str]:
+        return frozenset({"input_ids", "position_ids"})
+
+    @property
+    def discriminating_names(self) -> frozenset[str]:
+        return frozenset({"cu_seqlens_q", "cu_seqlens_kv", "qkv_format"})
+
+    def has_required_names(self, names: set[str]) -> bool:
+        return "input_ids" in names
+
+    def extract_global_seq_lens(
+        self, step_data: dict[str, object]
+    ) -> Optional[list[int]]:
+        if not self.cp_sharded_names:
+            return None
+
+        cu_seqlens_q = step_data.get("cu_seqlens_q")
+        if not isinstance(cu_seqlens_q, torch.Tensor):
+            return None
+
+        return (cu_seqlens_q[1:] - cu_seqlens_q[:-1]).tolist()
+
+    def infer_cp_sharded_dims(self, name: str, ndim: int) -> str:
+        """Infer dims for CP-sharded aux tensors.
+
+        NOTE: assumes zigzag ordering — natural-order CP without explicit dims
+        will be mishandled. Callers should set dims explicitly for non-zigzag CP.
+        """
+        if ndim == 1:
+            return "t[cp:zigzag]"
+        if ndim == 2:
+            return "b s[cp:zigzag]"
+        raise ValueError(
+            f"Megatron: cannot infer dims for CP-sharded '{name}' with ndim={ndim}"
+        )
+
+    def detect_layout(self, raw: dict[int, dict[str, object]]) -> TokenLayout:
+        for step_data in raw.values():
+            if (qkv_format := step_data.get("qkv_format")) is not None:
+                fmt = qkv_format if isinstance(qkv_format, str) else str(qkv_format)
+                if "bshd" in fmt.lower():
+                    return TokenLayout.BS
+                return TokenLayout.T
+
+            input_ids = step_data.get("input_ids")
+            if isinstance(input_ids, torch.Tensor) and input_ids.ndim == 2:
+                return TokenLayout.BS
+
+        log_sink.add(
+            InfoLog(
+                category="layout_detection_fallback",
+                message=(
+                    "Megatron layout detection: no qkv_format or 2D input_ids found, "
+                    "falling back to T"
+                ),
+            )
+        )
+        return TokenLayout.T
+
+    def compute_step_aux(
+        self, step_data: dict[str, object], *, layout: TokenLayout, step: int
+    ) -> TokenAlignerStepAux:
+        input_ids: torch.Tensor = step_data["input_ids"]
+        is_bshd: bool = layout == TokenLayout.BS
+
+        # BSHD [B, S] → flat [B*S]; THD [T] stays as-is
+        flat_ids: list[int] = input_ids.reshape(-1).tolist()
+
+        if (cu_seqlens_q := step_data.get("cu_seqlens_q")) is not None:
+            seq_lens_list: list[int] = (cu_seqlens_q[1:] - cu_seqlens_q[:-1]).tolist()
+        elif is_bshd:
+            seq_lens_list = [input_ids.shape[1]] * input_ids.shape[0]
+        else:
+            seq_lens_list = [input_ids.shape[0]]
+
+        if (position_ids := step_data.get("position_ids")) is not None:
+            flat_positions: list[int] = position_ids.reshape(-1).tolist()
+        elif is_bshd:
+            flat_positions = list(range(input_ids.shape[1])) * input_ids.shape[0]
+        else:
+            flat_positions = _infer_positions(
+                seq_lens=torch.tensor(seq_lens_list)
+            ).tolist()
+
+        num_seqs: int = len(seq_lens_list)
+        seq_ids: list[SeqId] = [
+            PositionalSeqId(step=step, seq_index=seq_index)
+            for seq_index in range(num_seqs)
+        ]
+
+        return TokenAlignerStepAux(
+            input_ids=flat_ids,
+            positions=flat_positions,
+            seq_lens=seq_lens_list,
+            seq_ids=seq_ids,
+        )
+
+
+# ── plugin registry ───────────────────────────────────────────────
+
+_plugins: list[_AuxFrameworkPlugin] = [_SGLangPlugin(), _MegatronPlugin()]
+
+AUX_NAMES: frozenset[str] = frozenset().union(*(p.all_names for p in _plugins))
+
+
+# ── helpers ────────────────────────────────────────────────────────
+
+
+def _infer_positions(*, seq_lens: torch.Tensor) -> torch.Tensor:
+    """Infer positions when position_ids is missing (THD only)."""
+    return torch.cat([torch.arange(int(slen.item())) for slen in seq_lens])

python/sglang/srt/debug_utils/comparator/aligner/token_aligner/smart/executor.py

@@ -0,0 +1,150 @@
+from __future__ import annotations
+
+import torch
+from einops import rearrange
+
+from sglang.srt.debug_utils.comparator.aligner.token_aligner.smart.types import (
+    TokenAlignerPlan,
+    TokenLocator,
+)
+from sglang.srt.debug_utils.comparator.dims_spec import (
+    BATCH_DIM_NAME,
+    SEQ_DIM_NAME,
+    TOKEN_DIM_NAME,
+    TokenLayout,
+    apply_dim_names,
+    get_dim_names,
+    resolve_dim_by_name,
+)
+from sglang.srt.debug_utils.comparator.utils import Pair
+
+_UNNAMED_TOKEN_DIM_FALLBACK: int = 0
+
+
+def execute_token_aligner(
+    plan: TokenAlignerPlan,
+    tensor_of_step_pair: Pair[dict[int, torch.Tensor]],
+) -> Pair[torch.Tensor]:
+    flat_pair: Pair[dict[int, torch.Tensor]] = Pair(
+        x=_collapse_bs_to_t(
+            tensor_of_step=tensor_of_step_pair.x, layout=plan.layouts.x
+        ),
+        y=_collapse_bs_to_t(
+            tensor_of_step=tensor_of_step_pair.y, layout=plan.layouts.y
+        ),
+    )
+
+    if not plan.locators.x.steps:
+        return Pair(
+            x=_make_empty(tensor_of_step=flat_pair.x),
+            y=_make_empty(tensor_of_step=flat_pair.y),
+        )
+
+    return Pair(
+        x=_extract_and_stack_tokens(
+            tensor_of_step=flat_pair.x, locator=plan.locators.x
+        ),
+        y=_extract_and_stack_tokens(
+            tensor_of_step=flat_pair.y, locator=plan.locators.y
+        ),
+    )
+
+
+# ── BS → T preprocessing ─────────────────────────────────────────
+
+
+def _collapse_bs_to_t(
+    *,
+    tensor_of_step: dict[int, torch.Tensor],
+    layout: TokenLayout,
+) -> dict[int, torch.Tensor]:
+    """Collapse B and S dims into a single flat token dim (always batch-major).
+
+    Handles both ``b s`` and ``s b`` orderings correctly via einops rearrange.
+    Returns the original tensors unchanged if layout is T.
+    """
+    if layout != TokenLayout.BS:
+        return tensor_of_step
+
+    some_tensor: torch.Tensor = next(iter(tensor_of_step.values()))
+    batch_dim: int = _resolve_dim_or_fallback(some_tensor, BATCH_DIM_NAME)
+    seq_dim: int = _resolve_dim_or_fallback(some_tensor, SEQ_DIM_NAME)
+
+    if abs(batch_dim - seq_dim) != 1:
+        raise ValueError(
+            f"BS dims must be adjacent: "
+            f"{BATCH_DIM_NAME}={batch_dim}, "
+            f"{SEQ_DIM_NAME}={seq_dim}"
+        )
+
+    lhs_pattern, rhs_pattern, new_names = _build_bs_collapse_pattern(
+        names=list(get_dim_names(some_tensor)),
+        batch_dim=batch_dim,
+        seq_dim=seq_dim,
+    )
+
+    result: dict[int, torch.Tensor] = {}
+    for step, tensor in tensor_of_step.items():
+        collapsed: torch.Tensor = rearrange(tensor, f"{lhs_pattern} -> {rhs_pattern}")
+        collapsed = apply_dim_names(collapsed, [n for n in new_names if n is not None])
+        result[step] = collapsed
+
+    return result
+
+
+def _build_bs_collapse_pattern(
+    *,
+    names: list[str | None],
+    batch_dim: int,
+    seq_dim: int,
+) -> tuple[str, str, list[str | None]]:
+    """Build einops lhs/rhs patterns and output dim names for BS→T collapse.
+
+    Always produces batch-major order ``(b s)`` regardless of input ordering.
+    Uses the tensor's own dim names as einops axis names.
+    """
+    lo: int = min(batch_dim, seq_dim)
+    hi: int = max(batch_dim, seq_dim)
+
+    lhs: str = " ".join(names)  # type: ignore[arg-type]
+
+    rhs_names: list[str] = list(names[:lo]) + [f"({BATCH_DIM_NAME} {SEQ_DIM_NAME})"] + list(names[hi + 1 :])  # type: ignore[misc]
+    rhs: str = " ".join(rhs_names)
+
+    new_names: list[str | None] = (
+        list(names[:lo]) + [TOKEN_DIM_NAME] + list(names[hi + 1 :])
+    )
+
+    return lhs, rhs, new_names
+
+
+# ── core logic (T layout only) ───────────────────────────────────
+
+
+def _resolve_dim_or_fallback(tensor: torch.Tensor, name: str) -> int:
+    if get_dim_names(tensor)[0] is None:
+        return _UNNAMED_TOKEN_DIM_FALLBACK
+    return resolve_dim_by_name(tensor, name)
+
+
+def _make_empty(*, tensor_of_step: dict[int, torch.Tensor]) -> torch.Tensor:
+    dummy: torch.Tensor = next(iter(tensor_of_step.values()))
+    token_dim: int = _resolve_dim_or_fallback(dummy, TOKEN_DIM_NAME)
+    shape: list[int] = list(dummy.shape)
+    shape[token_dim] = 0
+    return torch.empty(shape, dtype=dummy.dtype)
+
+
+def _extract_and_stack_tokens(
+    *,
+    tensor_of_step: dict[int, torch.Tensor],
+    locator: TokenLocator,
+) -> torch.Tensor:
+    some_tensor: torch.Tensor = next(iter(tensor_of_step.values()))
+    token_dim: int = _resolve_dim_or_fallback(some_tensor, TOKEN_DIM_NAME)
+
+    tokens: list[torch.Tensor] = [
+        tensor_of_step[s].select(dim=token_dim, index=i)
+        for s, i in zip(locator.steps, locator.token_index_in_step)
+    ]
+    return torch.stack(tokens, dim=token_dim)

python/sglang/srt/debug_utils/comparator/aligner/token_aligner/smart/planner.py

@@ -0,0 +1,135 @@
+from __future__ import annotations
+
+from collections import defaultdict
+from typing import NamedTuple, Optional
+
+from sglang.srt.debug_utils.comparator.aligner.token_aligner.smart.types import (
+    SeqId,
+    TokenAlignerPlan,
+    TokenAlignerSeqInfo,
+    TokenAlignerSeqsInfo,
+    TokenLocator,
+)
+from sglang.srt.debug_utils.comparator.utils import Pair
+
+
+def compute_token_aligner_plan(
+    seqs_info_pair: Pair[TokenAlignerSeqsInfo],
+) -> TokenAlignerPlan:
+    """Compute a token alignment plan from two side token seqs_info_pair."""
+    matched_pairs: list[tuple[SeqId, SeqId]] = _match_sequences(
+        seqs=Pair(x=seqs_info_pair.x.sequences, y=seqs_info_pair.y.sequences)
+    )
+
+    _empty = TokenLocator(steps=[], token_index_in_step=[])
+    locator_x: TokenLocator = _empty
+    locator_y: TokenLocator = _empty
+
+    for seq_id_x, seq_id_y in matched_pairs:
+        rec: Pair[TokenAlignerSeqInfo] = Pair(
+            x=seqs_info_pair.x.sequences[seq_id_x],
+            y=seqs_info_pair.y.sequences[seq_id_y],
+        )
+
+        # positions is validated to be [0, 1, ..., N-1], so position == index
+        # and the common range is simply [0, min(len_x, len_y)).
+        common_len: int = min(len(rec.x.positions), len(rec.y.positions))
+
+        x_ids = rec.x.input_ids[:common_len]
+        y_ids = rec.y.input_ids[:common_len]
+        assert x_ids == y_ids, f"{seq_id_x=} {seq_id_y=} {x_ids=} {y_ids=}"
+
+        locator_x = locator_x + TokenLocator(
+            steps=rec.x.locator.steps[:common_len],
+            token_index_in_step=rec.x.locator.token_index_in_step[:common_len],
+        )
+        locator_y = locator_y + TokenLocator(
+            steps=rec.y.locator.steps[:common_len],
+            token_index_in_step=rec.y.locator.token_index_in_step[:common_len],
+        )
+
+    return TokenAlignerPlan(
+        locators=Pair(x=locator_x, y=locator_y),
+        layouts=seqs_info_pair.map(lambda s: s.layout),
+    )
+
+
+# -------------------- Sequence matcher --------------------
+
+
+def _match_sequences(
+    seqs: Pair[dict[SeqId, TokenAlignerSeqInfo]],
+) -> list[tuple[SeqId, SeqId]]:
+    """For each y (target) sequence, find a matching x (baseline) sequence.
+
+    Two-pass: exact match first, then prefix match for remaining.
+    """
+    x_lookup: dict[tuple[int, ...], list[SeqId]] = defaultdict(list)
+    for seq_id, rec in seqs.x.items():
+        x_lookup[tuple(rec.input_ids)].append(seq_id)
+
+    claimed_x_ids: set[SeqId] = set()
+    matched_seq_id_pairs: list[tuple[SeqId, SeqId]] = []
+
+    for seq_id_y in sorted(seqs.y.keys()):
+        seq_y: TokenAlignerSeqInfo = seqs.y[seq_id_y]
+
+        matched_x: Optional[SeqId] = _find_matching_x_exact(
+            seq_y=seq_y, x_lookup=x_lookup, claimed_x_ids=claimed_x_ids
+        )
+        if matched_x is None:
+            matched_x = _find_matching_x_prefix(
+                seq_y=seq_y, x_seqs=seqs.x, claimed_x_ids=claimed_x_ids
+            )
+
+        if matched_x is not None:
+            matched_seq_id_pairs.append((matched_x, seq_id_y))
+            claimed_x_ids.add(matched_x)
+
+    return matched_seq_id_pairs
+
+
+def _find_matching_x_exact(
+    *,
+    seq_y: TokenAlignerSeqInfo,
+    x_lookup: dict[tuple[int, ...], list[SeqId]],
+    claimed_x_ids: set[SeqId],
+) -> Optional[SeqId]:
+    """Find an x sequence with identical input_ids."""
+    ids_y_key: tuple[int, ...] = tuple(seq_y.input_ids)
+    candidates: list[SeqId] = x_lookup.get(ids_y_key, [])
+    for candidate in candidates:
+        if candidate not in claimed_x_ids:
+            return candidate
+    return None
+
+
+class _PrefixCandidate(NamedTuple):
+    seq_id_x: SeqId
+    overlap_len: int
+
+
+def _find_matching_x_prefix(
+    *,
+    seq_y: TokenAlignerSeqInfo,
+    x_seqs: dict[SeqId, TokenAlignerSeqInfo],
+    claimed_x_ids: set[SeqId],
+) -> Optional[SeqId]:
+    """Find the x sequence with the longest prefix relationship to y."""
+    ids_y: list[int] = seq_y.input_ids
+    candidates: list[_PrefixCandidate] = [
+        _PrefixCandidate(
+            seq_id_x=seq_id_x, overlap_len=min(len(seq_x.input_ids), len(ids_y))
+        )
+        for seq_id_x, seq_x in x_seqs.items()
+        if seq_id_x not in claimed_x_ids and _is_prefix_pair(seq_x.input_ids, ids_y)
+    ]
+    if not candidates:
+        return None
+    return max(candidates, key=lambda c: c.overlap_len).seq_id_x
+
+
+def _is_prefix_pair(a: list[int], b: list[int]) -> bool:
+    """True if a is a prefix of b, or b is a prefix of a."""
+    shorter_len: int = min(len(a), len(b))
+    return a[:shorter_len] == b[:shorter_len]

python/sglang/srt/debug_utils/comparator/aligner/token_aligner/smart/seq_info_builder.py

@@ -0,0 +1,81 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+from sglang.srt.debug_utils.comparator.aligner.token_aligner.smart.types import (
+    SeqId,
+    TokenAlignerGlobalAux,
+    TokenAlignerSeqInfo,
+    TokenAlignerSeqsInfo,
+    TokenAlignerStepAux,
+    TokenLocator,
+)
+
+
+@dataclass
+class _SeqInfoAccumulator:
+    """Mutable accumulator for building TokenAlignerSeqInfo without per-step validation."""
+
+    input_ids: list[int] = field(default_factory=list)
+    positions: list[int] = field(default_factory=list)
+    steps: list[int] = field(default_factory=list)
+    token_index_in_step: list[int] = field(default_factory=list)
+
+    def extend(
+        self,
+        *,
+        input_ids: list[int],
+        positions: list[int],
+        steps: list[int],
+        token_index_in_step: list[int],
+    ) -> None:
+        self.input_ids.extend(input_ids)
+        self.positions.extend(positions)
+        self.steps.extend(steps)
+        self.token_index_in_step.extend(token_index_in_step)
+
+    def build(self) -> TokenAlignerSeqInfo:
+        return TokenAlignerSeqInfo(
+            input_ids=self.input_ids,
+            positions=self.positions,
+            locator=TokenLocator(
+                steps=self.steps,
+                token_index_in_step=self.token_index_in_step,
+            ),
+        )
+
+
+def build_seqs_info(global_aux: TokenAlignerGlobalAux) -> TokenAlignerSeqsInfo:
+    """Build sequence info for one side from its auxiliary tensors."""
+    return TokenAlignerSeqsInfo(
+        sequences=_build_token_aligner_seq_infos(global_aux),
+        layout=global_aux.layout,
+    )
+
+
+def _build_token_aligner_seq_infos(
+    global_aux: TokenAlignerGlobalAux,
+) -> dict[SeqId, TokenAlignerSeqInfo]:
+    """Build token index for any framework/layout using seq_ids for identity tracking."""
+    accum: dict[SeqId, _SeqInfoAccumulator] = {}
+
+    for step in sorted(global_aux.step_auxs.keys()):
+        aux: TokenAlignerStepAux = global_aux.step_auxs[step]
+
+        offset: int = 0
+        for seq_index, seq_len in enumerate(aux.seq_lens):
+            seq_id: SeqId = aux.seq_ids[seq_index]
+
+            if seq_id not in accum:
+                accum[seq_id] = _SeqInfoAccumulator()
+
+            accum[seq_id].extend(
+                input_ids=aux.input_ids[offset : offset + seq_len],
+                positions=aux.positions[offset : offset + seq_len],
+                steps=[step] * seq_len,
+                token_index_in_step=list(range(offset, offset + seq_len)),
+            )
+
+            offset += seq_len
+
+    return {seq_id: acc.build() for seq_id, acc in accum.items()}

python/sglang/srt/debug_utils/comparator/aligner/token_aligner/smart/types.py

@@ -0,0 +1,128 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import NamedTuple, Optional, Union
+
+from pydantic import model_validator
+
+from sglang.srt.debug_utils.comparator.dims_spec import TokenLayout
+from sglang.srt.debug_utils.comparator.utils import (
+    Pair,
+    _check_equal_lengths,
+    _FrozenBase,
+)
+
+
+class SGLangSeqId(NamedTuple):
+    rid: str
+
+
+class PositionalSeqId(NamedTuple):
+    step: int
+    seq_index: int
+
+
+SeqId = Union[SGLangSeqId, PositionalSeqId]
+
+
+@dataclass(frozen=True)
+class TokenAlignerStepAux:
+    """Normalized auxiliary tensors for a single step (framework-agnostic)."""
+
+    input_ids: list[int]  # [num_tokens]
+    positions: list[int]  # [num_tokens]
+    seq_lens: list[int]  # [num_seqs]
+    seq_ids: list[SeqId]  # [num_seqs] — sequence identity
+
+    def __post_init__(self) -> None:
+        _check_equal_lengths(input_ids=self.input_ids, positions=self.positions)
+        _check_equal_lengths(seq_lens=self.seq_lens, seq_ids=self.seq_ids)
+
+        token_count: int = sum(self.seq_lens)
+        if token_count != len(self.input_ids):
+            raise ValueError(
+                f"sum(seq_lens)={token_count} != len(input_ids)={len(self.input_ids)}"
+            )
+
+
+@dataclass(frozen=True)
+class TokenAlignerGlobalAux:
+    """Auxiliary tensors for one side across all steps + side-level metadata."""
+
+    step_auxs: dict[int, TokenAlignerStepAux]
+    framework: str  # "sglang" | "megatron"
+    layout: TokenLayout
+    thd_seq_lens_by_step: Optional[dict[int, list[int]]] = field(default=None)
+
+
+class TokenLocator(_FrozenBase):
+    """Locates tokens within a multi-step tensor store.
+
+    token i is at tensor_of_step[steps[i]][token_index_in_step[i]].
+    """
+
+    steps: list[int]
+    token_index_in_step: list[int]
+
+    def __add__(self, other: TokenLocator) -> TokenLocator:
+        return TokenLocator(
+            steps=self.steps + other.steps,
+            token_index_in_step=self.token_index_in_step + other.token_index_in_step,
+        )
+
+
+class TokenAlignerSeqInfo(_FrozenBase):
+    """Information for a sequence, containing information to locate all the tokens inside the sequence."""
+
+    # All these fields are of shape (num_tokens_in_seq,)
+    input_ids: list[int]
+    positions: list[int]
+    locator: TokenLocator
+
+    @model_validator(mode="after")
+    def _validate_fields(self) -> TokenAlignerSeqInfo:
+        n: int = len(self.input_ids)
+        _check_equal_lengths(
+            input_ids=self.input_ids,
+            positions=self.positions,
+            locator_steps=self.locator.steps,
+            locator_token_index_in_step=self.locator.token_index_in_step,
+        )
+
+        if self.positions != list(range(n)):
+            raise ValueError(
+                f"positions must be [0, 1, ..., {n - 1}], got {self.positions}"
+            )
+
+        return self
+
+    def __add__(self, other: TokenAlignerSeqInfo) -> TokenAlignerSeqInfo:
+        return TokenAlignerSeqInfo(
+            input_ids=self.input_ids + other.input_ids,
+            positions=self.positions + other.positions,
+            locator=self.locator + other.locator,
+        )
+
+
+class TokenAlignerSeqsInfo(_FrozenBase):
+    """All sequences for one side across all steps."""
+
+    sequences: dict[SeqId, TokenAlignerSeqInfo]
+    layout: TokenLayout
+
+
+class TokenAlignerPlan(_FrozenBase):
+    """Token alignment plan. locators.x[i] and locators.y[i] correspond to the same logical token."""
+
+    locators: Pair[TokenLocator]
+    layouts: Pair[TokenLayout]
+
+    @model_validator(mode="after")
+    def _validate_fields(self) -> TokenAlignerPlan:
+        _check_equal_lengths(
+            locators_x_steps=self.locators.x.steps,
+            locators_x_token_index_in_step=self.locators.x.token_index_in_step,
+            locators_y_steps=self.locators.y.steps,
+            locators_y_token_index_in_step=self.locators.y.token_index_in_step,
+        )
+        return self

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

@@ -0,0 +1,190 @@
+from dataclasses import dataclass, field
+from typing import Optional
+
+import torch
+
+from sglang.srt.debug_utils.comparator.aligner.unsharder.types import (
+    ConcatParams,
+    CpThdConcatParams,
+    PickParams,
+    ReduceSumParams,
+    UnsharderParams,
+    UnsharderPlan,
+)
+from sglang.srt.debug_utils.comparator.dims_spec import (
+    ParallelAxis,
+    apply_dim_names,
+    get_dim_names,
+    resolve_dim_by_name,
+    without_dim_names,
+)
+from sglang.srt.debug_utils.comparator.output_types import ReplicatedCheckResult
+from sglang.srt.debug_utils.comparator.tensor_comparator.comparator import compute_diff
+
+_REPLICATED_ATOL: float = 1e-6
+
+
+@dataclass(frozen=True)
+class UnsharderResult:
+    tensors: list[torch.Tensor]
+    replicated_checks: list[ReplicatedCheckResult] = field(default_factory=list)
+
+
+def execute_unsharder_plan(
+    plan: UnsharderPlan,
+    tensors: list[torch.Tensor],
+) -> UnsharderResult:
+    result_tensors: list[torch.Tensor] = []
+    all_checks: list[ReplicatedCheckResult] = []
+
+    for group_idx, group in enumerate(plan.groups):
+        group_tensors = [tensors[i] for i in group]
+        tensor, checks = _apply_unshard(
+            plan.params,
+            group_tensors,
+            axis=plan.axis,
+            group_index=group_idx,
+        )
+        result_tensors.append(tensor)
+        all_checks.extend(checks)
+
+    return UnsharderResult(tensors=result_tensors, replicated_checks=all_checks)
+
+
+def _apply_unshard(
+    params: UnsharderParams,
+    ordered_tensors: list[torch.Tensor],
+    *,
+    axis: ParallelAxis,
+    group_index: int,
+) -> tuple[torch.Tensor, list[ReplicatedCheckResult]]:
+    if isinstance(params, PickParams):
+        checks: list[ReplicatedCheckResult] = _verify_replicated_group(
+            ordered_tensors,
+            axis=axis,
+            group_index=group_index,
+        )
+        return ordered_tensors[0], checks
+
+    if isinstance(params, ConcatParams):
+        dim: int = resolve_dim_by_name(ordered_tensors[0], params.dim_name)
+        names: tuple[Optional[str], ...] = get_dim_names(ordered_tensors[0])
+        result = torch.cat(ordered_tensors, dim=dim)
+        if names[0] is not None:
+            result = apply_dim_names(result, list(names))
+        return result, []
+
+    if isinstance(params, CpThdConcatParams):
+        thd_dim: int = resolve_dim_by_name(ordered_tensors[0], params.dim_name)
+        return (
+            _thd_concat(
+                ordered_tensors,
+                dim=thd_dim,
+                seq_lens_per_rank=params.seq_lens_per_rank,
+            ),
+            [],
+        )
+
+    if isinstance(params, ReduceSumParams):
+        names: tuple[Optional[str], ...] = get_dim_names(ordered_tensors[0])
+        stripped: list[torch.Tensor] = [without_dim_names(t) for t in ordered_tensors]
+        result: torch.Tensor = torch.stack(stripped).sum(dim=0)
+        if names[0] is not None:
+            result = apply_dim_names(result, list(names))
+        return result, []
+
+    raise ValueError(f"Unsupported unshard operation: {type(params).__name__}")
+
+
+def _verify_replicated_group(
+    ordered_tensors: list[torch.Tensor],
+    *,
+    axis: ParallelAxis,
+    group_index: int,
+) -> list[ReplicatedCheckResult]:
+    baseline: torch.Tensor = ordered_tensors[0].float()
+
+    return [
+        _check_replicated_pair(
+            baseline=baseline,
+            other=ordered_tensors[i],
+            axis=axis,
+            group_index=group_index,
+            compared_index=i,
+        )
+        for i in range(1, len(ordered_tensors))
+    ]
+
+
+def _check_replicated_pair(
+    *,
+    baseline: torch.Tensor,
+    other: torch.Tensor,
+    axis: ParallelAxis,
+    group_index: int,
+    compared_index: int,
+) -> ReplicatedCheckResult:
+    other_float: torch.Tensor = other.float()
+
+    if baseline.shape != other_float.shape:
+        passed = False
+        diff_info = None
+    else:
+        diff_info = compute_diff(
+            x_baseline=baseline,
+            x_target=other_float,
+            predicate=f"max_abs <= {_REPLICATED_ATOL}",
+        )
+        passed = diff_info.passed
+
+    return ReplicatedCheckResult(
+        axis=axis.value,
+        group_index=group_index,
+        compared_index=compared_index,
+        baseline_index=0,
+        passed=passed,
+        atol=_REPLICATED_ATOL,
+        diff=diff_info,
+    )
+
+
+def _thd_concat(
+    ordered_tensors: list[torch.Tensor],
+    *,
+    dim: int,
+    seq_lens_per_rank: list[int],
+) -> torch.Tensor:
+    """Per-seq concat across ranks for THD format.
+
+    Each rank holds segments of each seq packed contiguously:
+      rank_data = [seq0_tokens | seq1_tokens | ... | pad_tokens]
+
+    This function splits each rank by seq_lens, then interleaves across ranks
+    per-seq: [seqA_r0 + seqA_r1 + ... | seqB_r0 + seqB_r1 + ... | tail_pad].
+    """
+    names: tuple[Optional[str], ...] = get_dim_names(ordered_tensors[0])
+    stripped: list[torch.Tensor] = [without_dim_names(t) for t in ordered_tensors]
+
+    # Split each rank into [seq0, seq1, ..., tail_remainder]
+    split_sizes: list[int] = list(seq_lens_per_rank)
+    remainder: int = stripped[0].shape[dim] - sum(split_sizes)
+    if remainder < 0:
+        raise ValueError(
+            f"sum(seq_lens_per_rank)={sum(split_sizes)} exceeds tensor dim size "
+            f"{stripped[0].shape[dim]} along dim={dim}"
+        )
+    if remainder > 0:
+        split_sizes.append(remainder)
+    per_rank_splits: list[tuple[torch.Tensor, ...]] = [
+        t.split(split_sizes, dim=dim) for t in stripped
+    ]
+
+    # Per-seq concat across ranks, then concatenate all seqs
+    result: torch.Tensor = torch.cat(
+        [torch.cat(rank_parts, dim=dim) for rank_parts in zip(*per_rank_splits)],
+        dim=dim,
+    )
+
+    if names[0] is not None:
+        result = apply_dim_names(result, list(names))
+    return result

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

@@ -0,0 +1,45 @@
+from typing import Optional
+
+from sglang.srt.debug_utils.comparator.aligner.unsharder.types import AxisInfo
+from sglang.srt.debug_utils.comparator.dims_spec 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 and not _is_error_sentinel(value):
+            if info is not None:
+                raise ValueError(
+                    f"Meta contains multiple parallel_info keys among {_PARALLEL_INFO_KEYS}"
+                )
+            info = value
+
+    if info is None:
+        info = {}
+
+    result: dict[ParallelAxis, AxisInfo] = {}
+    for axis in ParallelAxis:
+        axis_rank = info.get(f"{axis.value}_rank")
+        axis_size = info.get(f"{axis.value}_size")
+
+        # Recompute pseudo-axis lives at top-level meta, not inside parallel_info
+        if axis_rank is None:
+            axis_rank = meta.get(f"{axis.value}_rank")
+            axis_size = meta.get(f"{axis.value}_size")
+
+        if axis_rank is not None and axis_size is not None and axis_size > 1:
+            result[axis] = AxisInfo(
+                axis_rank=axis_rank,
+                axis_size=axis_size,
+            )
+
+    return result

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

@@ -0,0 +1,373 @@
+from collections import defaultdict
+from typing import NamedTuple, Optional
+
+from sglang.srt.debug_utils.comparator.aligner.unsharder.types import (
+    AxisInfo,
+    ConcatParams,
+    CpThdConcatParams,
+    PickParams,
+    ReduceSumParams,
+    UnsharderParams,
+    UnsharderPlan,
+)
+from sglang.srt.debug_utils.comparator.dims_spec import (
+    TOKEN_DIM_NAME,
+    DimSpec,
+    ParallelAxis,
+    ParallelModifier,
+)
+
+# _CoordsList[tensor_index][axis] =
+#     the axis_rank (shard position) of the tensor_index-th tensor along `axis`
+#     (e.g. coords[2] = {TP: 3} means tensor 2 is the 3rd shard in TP axis)
+_CoordsList = list[dict[ParallelAxis, int]]
+
+
+class _GroupResult(NamedTuple):
+    groups: list[list[int]]
+    projected_coords: _CoordsList
+
+
+def compute_unsharder_plan(
+    dim_specs: list[DimSpec],
+    parallel_infos: list[dict[ParallelAxis, AxisInfo]],
+    *,
+    explicit_replicated_axes: frozenset[ParallelAxis] = frozenset(),
+    thd_global_seq_lens: Optional[list[int]] = None,
+    dp_filtered_axis: Optional[ParallelAxis] = None,
+) -> list[UnsharderPlan]:
+    if not parallel_infos:
+        raise ValueError("parallel_infos must not be empty")
+
+    # Within each dim spec, reverse modifier order: innermost shard (rightmost) unshards first.
+    reversed_sharded_modifiers: list[tuple[str, ParallelModifier]] = [
+        (spec.sanitized_name, m)
+        for spec in dim_specs
+        for m in reversed(spec.parallel_modifiers)
+    ]
+
+    sharded_axes_raw: set[ParallelAxis] = {
+        m.axis for _, m in reversed_sharded_modifiers
+    }
+    all_axes: set[ParallelAxis] = {axis for info in parallel_infos for axis in info}
+
+    # axis annotated in dims but absent from all parallel_infos -> axis_size=1, skip
+    sharded_axes: set[ParallelAxis] = sharded_axes_raw & all_axes
+    reversed_sharded_modifiers = [
+        (name, m) for name, m in reversed_sharded_modifiers if m.axis in sharded_axes
+    ]
+
+    # RECOMPUTE_PSEUDO is always implicitly replicated (system-injected, not user-facing)
+    auto_replicated: frozenset[ParallelAxis] = frozenset(
+        {ParallelAxis.RECOMPUTE_PSEUDO} & all_axes
+    )
+    effective_replicated: frozenset[ParallelAxis] = (
+        explicit_replicated_axes | auto_replicated
+    )
+
+    _validate_explicit_replicated(
+        explicit_replicated_axes=effective_replicated,
+        sharded_axes=sharded_axes,
+        all_axes=all_axes,
+        parallel_infos=parallel_infos,
+        dp_filtered_axis=dp_filtered_axis,
+    )
+    replicated_axes: frozenset[ParallelAxis] = effective_replicated
+
+    if not sharded_axes and not replicated_axes:
+        return []
+
+    _validate(
+        axes_to_validate=sharded_axes | replicated_axes,
+        parallel_infos=parallel_infos,
+    )
+
+    current_coords: _CoordsList = [
+        {axis: info[axis].axis_rank for axis in sharded_axes | replicated_axes}
+        for info in parallel_infos
+    ]
+
+    axis_and_params: list[tuple[ParallelAxis, UnsharderParams]] = [
+        (axis, PickParams()) for axis in sorted(replicated_axes, key=lambda a: a.value)
+    ] + [
+        (
+            modifier.axis,
+            _resolve_unshard_params(
+                modifier=modifier,
+                dim_name=dim_name,
+                parallel_infos=parallel_infos,
+                thd_global_seq_lens=thd_global_seq_lens,
+            ),
+        )
+        for dim_name, modifier in reversed_sharded_modifiers
+    ]
+
+    plans: list[UnsharderPlan] = []
+    for axis, params in axis_and_params:
+        result = _group_and_project(
+            current_coords=current_coords,
+            target_axis=axis,
+        )
+        plans.append(UnsharderPlan(axis=axis, params=params, groups=result.groups))
+        current_coords = result.projected_coords
+
+    return plans
+
+
+def _validate_explicit_replicated(
+    *,
+    explicit_replicated_axes: frozenset[ParallelAxis],
+    sharded_axes: set[ParallelAxis],
+    all_axes: set[ParallelAxis],
+    parallel_infos: list[dict[ParallelAxis, AxisInfo]],
+    dp_filtered_axis: Optional[ParallelAxis] = None,
+) -> None:
+    """Validate explicit replicated declarations against sharded axes and parallel_infos."""
+    invalid: frozenset[ParallelAxis] = explicit_replicated_axes - all_axes
+    if invalid:
+        invalid_names: str = ", ".join(sorted(a.value for a in invalid))
+        raise ValueError(
+            f"Declared replicated axes {{{invalid_names}}} not found in parallel_infos "
+            f"(active axes: {{{', '.join(sorted(a.value for a in all_axes))}}})"
+        )
+
+    conflict: set[ParallelAxis] = explicit_replicated_axes & sharded_axes
+    if conflict:
+        conflict_names: str = ", ".join(sorted(a.value for a in conflict))
+        raise ValueError(
+            f"Axes {{{conflict_names}}} declared as both sharded and replicated"
+        )
+
+    _validate_replicated_axes_orthogonal(
+        explicit_replicated_axes=explicit_replicated_axes,
+        parallel_infos=parallel_infos,
+    )
+
+    candidate_axes: set[ParallelAxis] = (
+        all_axes - sharded_axes - explicit_replicated_axes
+    )
+    implicitly_replicated: frozenset[ParallelAxis] = _compute_dependent_axes(
+        parent_axes=explicit_replicated_axes,
+        candidate_axes=candidate_axes,
+        parallel_infos=parallel_infos,
+    )
+    implicitly_sharded: frozenset[ParallelAxis] = _compute_dependent_axes(
+        parent_axes=sharded_axes,
+        candidate_axes=candidate_axes - implicitly_replicated,
+        parallel_infos=parallel_infos,
+    )
+
+    declared_axes: frozenset[ParallelAxis] = frozenset(
+        sharded_axes
+        | explicit_replicated_axes
+        | implicitly_replicated
+        | implicitly_sharded
+        | ({dp_filtered_axis} if dp_filtered_axis is not None else set())
+    )
+    undeclared: set[ParallelAxis] = all_axes - declared_axes
+
+    jointly_determined: frozenset[ParallelAxis] = frozenset(
+        child
+        for child in undeclared
+        if _is_jointly_determined(
+            parallel_infos, parent_axes=declared_axes, child=child
+        )
+    )
+    undeclared -= jointly_determined
+
+    if undeclared:
+        undeclared_names: str = ", ".join(sorted(a.value for a in undeclared))
+        raise ValueError(
+            f"Axes {{{undeclared_names}}} are active (axis_size > 1) but not declared "
+            f"in dims. Annotate as sharded in dim spec or as '# axis:replicated'."
+        )
+
+
+def _validate_replicated_axes_orthogonal(
+    *,
+    explicit_replicated_axes: frozenset[ParallelAxis],
+    parallel_infos: list[dict[ParallelAxis, AxisInfo]],
+) -> None:
+    """Every pair of explicitly replicated axes must be fully orthogonal (no dependency)."""
+    axes: list[ParallelAxis] = sorted(explicit_replicated_axes, key=lambda a: a.value)
+    if len(axes) < 2:
+        return
+
+    violations: list[str] = []
+    for i, axis_a in enumerate(axes):
+        for axis_b in axes[i + 1 :]:
+            for parent, child in [(axis_a, axis_b), (axis_b, axis_a)]:
+                if _is_dependent_axis(parallel_infos, parent=parent, child=child):
+                    violations.append(
+                        f"'{parent.value}' determines '{child.value}' — "
+                        f"remove '{child.value}:replicated'"
+                    )
+
+    if violations:
+        details = "; ".join(violations)
+        raise ValueError(
+            f"Explicitly-replicated axes overlap (not orthogonal): {details}"
+        )
+
+
+def _validate(
+    *,
+    axes_to_validate: set[ParallelAxis],
+    parallel_infos: list[dict[ParallelAxis, AxisInfo]],
+) -> None:
+    """Check that every rank has all axes, sizes are consistent, and ranks are complete."""
+    axis_sizes: dict[ParallelAxis, int] = {}
+
+    for world_rank, parallel_info in enumerate(parallel_infos):
+        for axis in axes_to_validate:
+            if axis not in parallel_info:
+                raise ValueError(
+                    f"world_rank={world_rank} missing parallel_info for "
+                    f"axis {axis.value!r}"
+                )
+
+            axis_info = parallel_info[axis]
+            if axis not in axis_sizes:
+                axis_sizes[axis] = axis_info.axis_size
+            elif axis_info.axis_size != axis_sizes[axis]:
+                raise ValueError(
+                    f"Inconsistent axis_size for {axis.value}: "
+                    f"expected {axis_sizes[axis]}, got {axis_info.axis_size} "
+                    f"at world_rank={world_rank}"
+                )
+
+    for axis, expected_size in axis_sizes.items():
+        seen_ranks = {info[axis].axis_rank for info in parallel_infos}
+        if seen_ranks != set(range(expected_size)):
+            raise ValueError(
+                f"axis_rank coverage for {axis.value} is incomplete: "
+                f"got {sorted(seen_ranks)}, expected 0..{expected_size - 1}"
+            )
+
+
+def _compute_dependent_axes(
+    parent_axes: set[ParallelAxis] | frozenset[ParallelAxis],
+    candidate_axes: set[ParallelAxis],
+    parallel_infos: list[dict[ParallelAxis, AxisInfo]],
+) -> frozenset[ParallelAxis]:
+    """Return candidate axes whose rank is uniquely determined by some parent axis."""
+    return frozenset(
+        child
+        for child in candidate_axes
+        if any(
+            _is_dependent_axis(parallel_infos, parent=parent, child=child)
+            for parent in parent_axes
+        )
+    )
+
+
+def _is_jointly_determined(
+    parallel_infos: list[dict[ParallelAxis, AxisInfo]],
+    *,
+    parent_axes: frozenset[ParallelAxis],
+    child: ParallelAxis,
+) -> bool:
+    """True if child's rank is uniquely determined by the joint tuple of parent ranks.
+
+    Unlike ``_is_dependent_axis`` which checks single-parent dependency, this
+    checks whether the *combination* of all parent axes jointly determines the
+    child.  For example, ``edp_rank`` may not be a function of ``tp_rank`` alone
+    or ``cp_rank`` alone, but it *is* a function of ``(tp_rank, cp_rank)``.
+
+    Parent axes that are absent from *every* info are ignored (they carry no
+    information — e.g. DP with size 1 filtered by ``normalize_parallel_info``).
+    However, a parent axis present in *some* infos but missing from an info
+    that contains the child makes the determination incomplete → ``False``.
+    """
+    if not parent_axes:
+        return False
+
+    active_parents: frozenset[ParallelAxis] = frozenset(
+        ax for ax in parent_axes if any(ax in info for info in parallel_infos)
+    )
+    if not active_parents:
+        return False
+
+    mapping: dict[frozenset, int] = {}
+    for info in parallel_infos:
+        if child not in info:
+            continue
+        if not active_parents.issubset(info):
+            return False
+        parent_key = frozenset((ax, info[ax].axis_rank) for ax in active_parents)
+        child_rank: int = info[child].axis_rank
+        if mapping.setdefault(parent_key, child_rank) != child_rank:
+            return False
+
+    return bool(mapping)
+
+
+def _is_dependent_axis(
+    parallel_infos: list[dict[ParallelAxis, AxisInfo]],
+    *,
+    parent: ParallelAxis,
+    child: ParallelAxis,
+) -> bool:
+    """True if child's rank is uniquely determined by parent's rank."""
+    parent_rank_to_child_rank: dict[int, int] = {}
+    for info in parallel_infos:
+        if parent not in info or child not in info:
+            continue
+        parent_rank = info[parent].axis_rank
+        child_rank = info[child].axis_rank
+        if parent_rank_to_child_rank.setdefault(parent_rank, child_rank) != child_rank:
+            return False
+    return True
+
+
+def _group_and_project(
+    *,
+    current_coords: _CoordsList,
+    target_axis: ParallelAxis,
+) -> _GroupResult:
+    """Group tensors by other-axes coords, sort within group by target_axis rank."""
+    # buckets[coords_excluding_target] = [(axis_rank, tensor_index), ...]
+    # e.g. when target_axis=CP: buckets[{(TP,0)}] = [(0, 1), (1, 3)]
+    #   means tensor 1 (CP rank 0) and tensor 3 (CP rank 1) share TP rank 0
+    buckets: dict[frozenset, list[tuple[int, int]]] = defaultdict(list)
+
+    for idx, coords in enumerate(current_coords):
+        key = frozenset((k, v) for k, v in coords.items() if k != target_axis)
+        buckets[key].append((coords[target_axis], idx))
+
+    groups: list[list[int]] = []
+    projected: _CoordsList = []
+    for key in sorted(buckets, key=lambda k: sorted((a.value, v) for a, v in k)):
+        entries = sorted(buckets[key])
+        groups.append([idx for _, idx in entries])
+        projected.append(dict(key))
+
+    return _GroupResult(groups=groups, projected_coords=projected)
+
+
+def _resolve_unshard_params(
+    *,
+    modifier: ParallelModifier,
+    dim_name: str,
+    parallel_infos: list[dict[ParallelAxis, AxisInfo]],
+    thd_global_seq_lens: Optional[list[int]] = None,
+) -> UnsharderParams:
+    if modifier.reduction is not None:
+        return ReduceSumParams()
+
+    if (
+        dim_name == TOKEN_DIM_NAME
+        and modifier.axis == ParallelAxis.CP
+        and thd_global_seq_lens is not None
+    ):
+        axis_size: int = parallel_infos[0][modifier.axis].axis_size
+        for s in thd_global_seq_lens:
+            if s % axis_size != 0:
+                raise ValueError(
+                    f"THD seq_len {s} is not divisible by cp_size {axis_size}. "
+                    f"Sequences must be padded to a multiple of cp_size for CP zigzag."
+                )
+        seq_lens_per_rank: list[int] = [s // axis_size for s in thd_global_seq_lens]
+        return CpThdConcatParams(dim_name=dim_name, seq_lens_per_rank=seq_lens_per_rank)
+
+    return ConcatParams(dim_name=dim_name)

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

@@ -0,0 +1,60 @@
+from __future__ import annotations
+
+from typing import Annotated, Literal, Union
+
+from pydantic import Field, model_validator
+
+from sglang.srt.debug_utils.comparator.dims_spec import ParallelAxis
+from sglang.srt.debug_utils.comparator.utils import _FrozenBase
+
+
+class AxisInfo(_FrozenBase):
+    axis_rank: int
+    axis_size: int
+
+    @model_validator(mode="after")
+    def _validate_bounds(self) -> AxisInfo:
+        if self.axis_size <= 0:
+            raise ValueError(f"axis_size must be > 0, got {self.axis_size}")
+        if not (0 <= self.axis_rank < self.axis_size):
+            raise ValueError(
+                f"axis_rank must be in [0, {self.axis_size}), got {self.axis_rank}"
+            )
+        return self
+
+
+class ConcatParams(_FrozenBase):
+    op: Literal["concat"] = "concat"
+    dim_name: str
+
+
+class CpThdConcatParams(_FrozenBase):
+    op: Literal["cp_thd_concat"] = "cp_thd_concat"
+    dim_name: str
+    seq_lens_per_rank: list[int]  # per-seq token count on each rank, e.g. [50, 32, 46]
+
+
+class PickParams(_FrozenBase):
+    op: Literal["pick"] = "pick"
+
+
+class ReduceSumParams(_FrozenBase):
+    op: Literal["reduce_sum"] = "reduce_sum"
+
+
+UnsharderParams = Annotated[
+    Union[ConcatParams, CpThdConcatParams, PickParams, ReduceSumParams],
+    Field(discriminator="op"),
+]
+
+
+class UnsharderPlan(_FrozenBase):
+    type: Literal["unsharder"] = "unsharder"
+    axis: ParallelAxis
+    params: UnsharderParams
+    # groups[i] = indices in the input tensor list, which will be operated (e.g. concat) into i-th output tensor.
+    #
+    # Multistep example (CP=2, TP=2, 4 input tensors):
+    #   plan[0] (CP): groups=[[0,2],[1,3]]  — 4 tensors → 2 tensors
+    #   plan[1] (TP): groups=[[0,1]]        — 2 tensors → 1 tensor
+    groups: list[list[int]]

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

@@ -0,0 +1,436 @@
+"""Compare two tensor bundles."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any, Optional, Union
+
+import torch
+
+from sglang.srt.debug_utils.comparator.aligner.entrypoint.executor import (
+    AlignerResult,
+    execute_aligner_plan,
+)
+from sglang.srt.debug_utils.comparator.aligner.entrypoint.planner import (
+    compute_aligner_plan,
+)
+from sglang.srt.debug_utils.comparator.aligner.entrypoint.types import AlignerPlan
+from sglang.srt.debug_utils.comparator.aligner.token_aligner.smart.types import (
+    TokenAlignerPlan,
+)
+from sglang.srt.debug_utils.comparator.dims_spec import (
+    SEQ_DIM_NAME,
+    TOKEN_DIM_NAME,
+    ParallelAxis,
+    apply_dim_names,
+    get_dim_names,
+    parse_dims,
+    resolve_dim_names,
+    without_dim_names,
+)
+from sglang.srt.debug_utils.comparator.dp_utils import filter_to_non_empty_dp_rank
+from sglang.srt.debug_utils.comparator.log_sink import log_sink
+from sglang.srt.debug_utils.comparator.meta_overrider import MetaOverrider
+from sglang.srt.debug_utils.comparator.output_types import (
+    BundleFileInfo,
+    BundleSideInfo,
+    ComparisonNonTensorRecord,
+    ComparisonSkipRecord,
+    ComparisonTensorRecord,
+    ErrorLog,
+    _split_logs,
+)
+from sglang.srt.debug_utils.comparator.tensor_comparator.comparator import (
+    FailureDisplayBudget,
+    compare_tensor_pair,
+    compute_tensor_info,
+)
+from sglang.srt.debug_utils.comparator.threshold_dsl import DiffThresholdRule
+from sglang.srt.debug_utils.comparator.utils import Pair
+from sglang.srt.debug_utils.dump_loader import LOAD_FAILED, ValueWithMeta
+
+
+def _build_skip_from_one_empty_side(
+    *, name: str, pair: Pair[list[ValueWithMeta]]
+) -> ComparisonSkipRecord:
+    """Build a skip record when one side of *pair* is empty.
+
+    The non-empty side's tensor info is attached to the record.
+    """
+    assert not pair.x or not pair.y
+    if not pair.x:
+        reason, available_side, available_items = (
+            "baseline_load_failed",
+            "target",
+            pair.y,
+        )
+    else:
+        reason, available_side, available_items = (
+            "target_load_failed",
+            "baseline",
+            pair.x,
+        )
+
+    tensor_items: list[ValueWithMeta] = [
+        it for it in available_items if isinstance(it.value, torch.Tensor)
+    ]
+    if not tensor_items:
+        return ComparisonSkipRecord(name=name, reason=reason)
+
+    first_tensor: torch.Tensor = tensor_items[0].value
+    tensor_info = compute_tensor_info(first_tensor, include_sample=True)
+    metas: list[dict[str, Any]] = [it.meta for it in tensor_items]
+    bundle_info: BundleSideInfo = _collect_bundle_side_info(
+        items=tensor_items, metas=metas
+    )
+
+    return ComparisonSkipRecord(
+        name=name,
+        reason=reason,
+        available_side=available_side,  # type: ignore[arg-type]
+        available_tensor_info=tensor_info,
+        available_bundle_info=bundle_info,
+    )
+
+
+def _collect_bundle_side_info(
+    items: list[ValueWithMeta],
+    metas: list[dict[str, Any]],
+) -> BundleSideInfo:
+    from sglang.srt.debug_utils.comparator.display import (
+        PARALLEL_INFO_KEYS,
+        _extract_parallel_info,
+    )
+
+    files: list[BundleFileInfo] = []
+    for item, meta in zip(items, metas):
+        assert isinstance(item.value, torch.Tensor)
+        tensor: torch.Tensor = item.value
+
+        parallel_info: dict[str, str] = {}
+        for key in PARALLEL_INFO_KEYS:
+            _extract_parallel_info(row_data=parallel_info, info=meta.get(key, {}))
+
+        files.append(
+            BundleFileInfo(
+                shape=list(tensor.shape),
+                dtype=str(tensor.dtype),
+                rank=meta.get("rank"),
+                parallel_info=parallel_info if parallel_info else None,
+                filename=meta.get("filename"),
+            )
+        )
+
+    dims: Optional[str] = metas[0].get("dims") if metas else None
+    return BundleSideInfo(num_files=len(files), files=files, dims=dims)
+
+
+def compare_bundle_pair(
+    *,
+    name: str,
+    filenames_pair: Pair[list[str]],
+    dir_pair: Pair[Path],
+    token_aligner_mode: Optional[str],
+    token_aligner_plan: Optional[TokenAlignerPlan],
+    diff_threshold_rules: Optional[list[DiffThresholdRule]] = None,
+    failure_display_budget: Optional[FailureDisplayBudget] = None,
+    thd_seq_lens_by_step_pair: Pair[Optional[dict[int, list[int]]]] = Pair(
+        x=None, y=None
+    ),
+    viz_output_dir: Optional[Path] = None,
+    compute_per_token: bool = False,
+    meta_overrider: Optional[MetaOverrider] = None,
+) -> Union[ComparisonTensorRecord, ComparisonSkipRecord, ComparisonNonTensorRecord]:
+    with log_sink.context() as collected_logs:
+        result = _compare_bundle_pair_inner(
+            name=name,
+            filenames_pair=filenames_pair,
+            dir_pair=dir_pair,
+            token_aligner_mode=token_aligner_mode,
+            token_aligner_plan=token_aligner_plan,
+            diff_threshold_rules=diff_threshold_rules,
+            failure_display_budget=failure_display_budget,
+            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,
+        )
+
+    errors, infos = _split_logs(collected_logs)
+    return result.model_copy(update={"errors": errors, "infos": infos})
+
+
+def _compare_bundle_pair_inner(
+    *,
+    name: str,
+    filenames_pair: Pair[list[str]],
+    dir_pair: Pair[Path],
+    token_aligner_mode: Optional[str],
+    token_aligner_plan: Optional[TokenAlignerPlan],
+    diff_threshold_rules: Optional[list[DiffThresholdRule]] = None,
+    failure_display_budget: Optional[FailureDisplayBudget] = None,
+    thd_seq_lens_by_step_pair: Pair[Optional[dict[int, list[int]]]] = Pair(
+        x=None, y=None
+    ),
+    viz_output_dir: Optional[Path] = None,
+    compute_per_token: bool = False,
+    meta_overrider: Optional[MetaOverrider] = None,
+) -> Union[ComparisonTensorRecord, ComparisonSkipRecord, ComparisonNonTensorRecord]:
+    # 1. Load all successfully loaded values
+    all_pair: Pair[list[ValueWithMeta]] = Pair(
+        x=_load_all_values(filenames=filenames_pair.x, base_path=dir_pair.x),
+        y=_load_all_values(filenames=filenames_pair.y, base_path=dir_pair.y),
+    )
+
+    if not all_pair.x or not all_pair.y:
+        return _build_skip_from_one_empty_side(name=name, pair=all_pair)
+
+    # 1b. Dims override: patch meta["dims"] before DP filter reads it
+    # (--override-dims may add ``# dp:=moe_dp``, so it must run first)
+    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
+            ],
+        )
+
+    # 1c. DP filter: keep only the non-empty dp_rank
+    all_pair = all_pair.map(
+        lambda items: filter_to_non_empty_dp_rank(
+            items, dp_axis=_extract_dp_axis_from_items(items)
+        )
+    )
+
+    # 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]
+    )
+    if has_non_tensor:
+        return _compare_bundle_pair_non_tensor_type(name=name, value_pair=all_pair)
+
+    # 3. All values are tensors → tensor comparison path
+    return _compare_bundle_pair_tensor_type(
+        name=name,
+        valid_pair=all_pair,
+        token_aligner_mode=token_aligner_mode,
+        token_aligner_plan=token_aligner_plan,
+        diff_threshold_rules=diff_threshold_rules,
+        failure_display_budget=failure_display_budget,
+        thd_seq_lens_by_step_pair=thd_seq_lens_by_step_pair,
+        viz_output_dir=viz_output_dir,
+        compute_per_token=compute_per_token,
+    )
+
+
+def _extract_dp_axis_from_items(items: list[ValueWithMeta]) -> ParallelAxis:
+    """Extract dp axis from the first item's ``meta["dims"]``."""
+    if not items:
+        return ParallelAxis.DP
+    dims_str: Optional[str] = items[0].meta.get("dims")
+    if dims_str is None:
+        return ParallelAxis.DP
+    return parse_dims(dims_str).dp_axis
+
+
+def _compare_bundle_pair_tensor_type(
+    *,
+    name: str,
+    valid_pair: Pair[list[ValueWithMeta]],
+    token_aligner_mode: Optional[str],
+    token_aligner_plan: Optional[TokenAlignerPlan],
+    diff_threshold_rules: Optional[list[DiffThresholdRule]] = None,
+    failure_display_budget: Optional[FailureDisplayBudget] = None,
+    thd_seq_lens_by_step_pair: Pair[Optional[dict[int, list[int]]]] = Pair(
+        x=None, y=None
+    ),
+    viz_output_dir: Optional[Path] = None,
+    compute_per_token: bool = False,
+) -> Union[ComparisonTensorRecord, ComparisonSkipRecord]:
+    if not valid_pair.x or not valid_pair.y:
+        return _build_skip_from_one_empty_side(name=name, pair=valid_pair)
+
+    # Plan (meta only, no tensor)
+    metas_pair: Pair[list[dict[str, Any]]] = valid_pair.map(
+        lambda items: [it.meta for it in items]
+    )
+    plan: AlignerPlan = compute_aligner_plan(
+        metas_pair=metas_pair,
+        token_aligner_mode=token_aligner_mode,
+        token_aligner_plan=token_aligner_plan,
+        thd_seq_lens_by_step_pair=thd_seq_lens_by_step_pair,
+    )
+
+    # Collect raw bundle info before alignment
+    raw_bundle_info: Pair[BundleSideInfo] = Pair(
+        x=_collect_bundle_side_info(items=valid_pair.x, metas=metas_pair.x),
+        y=_collect_bundle_side_info(items=valid_pair.y, metas=metas_pair.y),
+    )
+
+    # Apply dim names to tensors, then execute
+    tensors_pair: Pair[list[torch.Tensor]] = Pair(
+        x=_apply_dim_names_from_meta(
+            tensors=[it.value for it in valid_pair.x],
+            metas=metas_pair.x,
+        ),
+        y=_apply_dim_names_from_meta(
+            tensors=[it.value for it in valid_pair.y],
+            metas=metas_pair.y,
+        ),
+    )
+    aligner_result: AlignerResult = execute_aligner_plan(
+        tensors_pair=tensors_pair, plan=plan
+    )
+    replicated_checks = aligner_result.replicated_checks
+
+    if aligner_result.tensors is None:
+        assert aligner_result.failed_side_xy is not None
+        failed_xy: str = aligner_result.failed_side_xy
+        pair_with_failed_emptied: Pair[list[ValueWithMeta]] = Pair(
+            x=[] if failed_xy == "x" else valid_pair.x,
+            y=[] if failed_xy == "y" else valid_pair.y,
+        )
+        return _build_skip_from_one_empty_side(name=name, pair=pair_with_failed_emptied)
+
+    # Resolve seq_dim for per-token computation
+    seq_dim: Optional[int] = (
+        _resolve_seq_dim(aligner_result.tensors.y) if compute_per_token else None
+    )
+
+    # Compare
+    aligned_baseline: torch.Tensor = without_dim_names(aligner_result.tensors.x)
+    aligned_target: torch.Tensor = without_dim_names(aligner_result.tensors.y)
+
+    info = compare_tensor_pair(
+        x_baseline=aligned_baseline,
+        x_target=aligned_target,
+        name=name,
+        diff_threshold_rules=diff_threshold_rules,
+        failure_display_budget=failure_display_budget,
+        seq_dim=seq_dim,
+    )
+    record = ComparisonTensorRecord(
+        **info.model_dump(),
+        traced_plan=aligner_result.traced_plan,
+        replicated_checks=replicated_checks,
+        raw_bundle_info=raw_bundle_info,
+    )
+
+    if viz_output_dir is not None:
+        _try_generate_viz(
+            baseline=aligned_baseline,
+            target=aligned_target,
+            name=name,
+            viz_output_dir=viz_output_dir,
+        )
+
+    return record
+
+
+def _try_generate_viz(
+    *,
+    baseline: torch.Tensor,
+    target: torch.Tensor,
+    name: str,
+    viz_output_dir: Path,
+) -> None:
+    from sglang.srt.debug_utils.comparator.visualizer import (
+        generate_comparison_figure,
+    )
+    from sglang.srt.debug_utils.comparator.visualizer.preprocessing import (
+        _sanitize_filename,
+    )
+
+    filename: str = _sanitize_filename(name) + ".png"
+    output_path: Path = viz_output_dir / filename
+
+    try:
+        generate_comparison_figure(
+            baseline=baseline,
+            target=target,
+            name=name,
+            output_path=output_path,
+        )
+    except Exception as exc:
+        log_sink.add(
+            ErrorLog(
+                category="visualizer",
+                message=f"Visualization failed for {name}: {exc}",
+            )
+        )
+
+
+def _resolve_seq_dim(tensor: torch.Tensor) -> Optional[int]:
+    """Find the token/seq dimension index from the tensor's named dims."""
+    names: tuple[Optional[str], ...] = get_dim_names(tensor)
+    if names[0] is None:
+        return None
+    for target_name in (TOKEN_DIM_NAME, SEQ_DIM_NAME):
+        if target_name in names:
+            return list(names).index(target_name)
+
+    return None
+
+
+def _compare_bundle_pair_non_tensor_type(
+    *,
+    name: str,
+    value_pair: Pair[list[ValueWithMeta]],
+) -> ComparisonNonTensorRecord:
+    baseline_value: Any = value_pair.x[0].value
+    target_value: Any = value_pair.y[0].value
+
+    try:
+        values_equal: bool = bool(baseline_value == target_value)
+    except Exception:
+        values_equal = False
+
+    return ComparisonNonTensorRecord(
+        name=name,
+        baseline_value=repr(baseline_value),
+        target_value=repr(target_value),
+        baseline_type=type(baseline_value).__name__,
+        target_type=type(target_value).__name__,
+        values_equal=values_equal,
+    )
+
+
+def _apply_dim_names_from_meta(
+    *,
+    tensors: list[torch.Tensor],
+    metas: list[dict[str, Any]],
+) -> list[torch.Tensor]:
+    if not metas:
+        return tensors
+
+    dims_str: Optional[str] = metas[0].get("dims")
+    if dims_str is None:
+        return tensors
+
+    dim_names: list[str] = resolve_dim_names(dims_str)
+    return [apply_dim_names(t, dim_names) for t in tensors]
+
+
+def _load_all_values(filenames: list[str], base_path: Path) -> list[ValueWithMeta]:
+    result: list[ValueWithMeta] = []
+    for f in filenames:
+        item: ValueWithMeta = ValueWithMeta.load(base_path / f)
+        if item.value is LOAD_FAILED:
+            log_sink.add(
+                ErrorLog(
+                    category="load_failed",
+                    message=f"Failed to load tensor file: {f}",
+                )
+            )
+            continue
+        result.append(item)
+    return result

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

@@ -0,0 +1,46 @@
+from __future__ import annotations
+
+import dataclasses
+from dataclasses import dataclass
+from typing import Any
+
+import polars as pl
+
+from sglang.srt.debug_utils.comparator.utils import Pair
+from sglang.srt.debug_utils.dump_loader import filter_rows
+
+
+@dataclass(frozen=True)
+class TensorFileInfo:
+    filename: str
+    name: str
+    step: int
+
+
+TensorBundleInfo = list[TensorFileInfo]
+
+
+def match_bundles(
+    *,
+    dfs: Pair[pl.DataFrame],
+    skip_keys: set[str],
+) -> list[Pair[TensorBundleInfo]]:
+    match_key_cols: list[str] = [c for c in dfs.y.columns if c not in skip_keys]
+    unique_keys: pl.DataFrame = dfs.y.select(match_key_cols).unique(maintain_order=True)
+
+    results: list[Pair[TensorBundleInfo]] = []
+    for key_values in unique_keys.iter_rows(named=True):
+        result = dfs.map(
+            lambda df: _rows_to_tensor_infos(filter_rows(df, conditions=key_values))
+        )
+        results.append(result)
+
+    return results
+
+
+def _rows_to_tensor_infos(rows: list[dict[str, Any]]) -> list[TensorFileInfo]:
+    tensor_info_fields: set[str] = {f.name for f in dataclasses.fields(TensorFileInfo)}
+    return [
+        TensorFileInfo(**{k: v for k, v in row.items() if k in tensor_info_fields})
+        for row in rows
+    ]

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

@@ -0,0 +1,51 @@
+from sglang.srt.debug_utils.comparator.dims_spec.dim_parser import parse_dim
+from sglang.srt.debug_utils.comparator.dims_spec.dims_parser import (
+    _SingletonDimUtil,
+    parse_dims,
+    resolve_dim_names,
+)
+from sglang.srt.debug_utils.comparator.dims_spec.tensor_naming import (
+    apply_dim_names,
+    find_dim_index,
+    get_dim_names,
+    resolve_dim_by_name,
+    without_dim_names,
+)
+from sglang.srt.debug_utils.comparator.dims_spec.types import (
+    _FUSED_NAME_SEP,
+    BATCH_DIM_NAME,
+    SEQ_DIM_NAME,
+    SQUEEZE_DIM_NAME,
+    TOKEN_DIM_NAME,
+    DimSpec,
+    DimsSpec,
+    Ordering,
+    ParallelAxis,
+    ParallelModifier,
+    Reduction,
+    TokenLayout,
+)
+
+__all__ = [
+    "BATCH_DIM_NAME",
+    "SEQ_DIM_NAME",
+    "SQUEEZE_DIM_NAME",
+    "TOKEN_DIM_NAME",
+    "DimsSpec",
+    "DimSpec",
+    "Ordering",
+    "ParallelAxis",
+    "ParallelModifier",
+    "Reduction",
+    "TokenLayout",
+    "_FUSED_NAME_SEP",
+    "_SingletonDimUtil",
+    "apply_dim_names",
+    "find_dim_index",
+    "get_dim_names",
+    "parse_dim",
+    "parse_dims",
+    "resolve_dim_by_name",
+    "resolve_dim_names",
+    "without_dim_names",
+]

python/sglang/srt/debug_utils/comparator/dims_spec/comment_parser.py

@@ -0,0 +1,59 @@
+from __future__ import annotations
+
+import re
+from typing import NamedTuple, Optional
+
+from sglang.srt.debug_utils.comparator.dims_spec.types import (
+    _AXIS_LOOKUP,
+    ParallelAxis,
+)
+
+_DP_ALIAS_PATTERN = re.compile(r"^dp:=(\w+)$")
+_REPLICATED_PATTERN = re.compile(r"^(\w+):replicated$")
+
+
+class _CommentSuffix(NamedTuple):
+    dp_group_alias: Optional[str] = None
+    replicated_axes: frozenset[ParallelAxis] = frozenset()
+
+
+def _parse_comment_suffix(declaration_part: str) -> _CommentSuffix:
+    """Parse the ``#`` comment section for dp alias and replicated declarations."""
+    dp_group_alias: Optional[str] = None
+    replicated_axes: set[ParallelAxis] = set()
+
+    for token in declaration_part.strip().split():
+        dp_match = _DP_ALIAS_PATTERN.match(token)
+        if dp_match is not None:
+            if dp_group_alias is not None:
+                raise ValueError(
+                    f"Duplicate dp alias declaration: already have {dp_group_alias!r}, "
+                    f"got {dp_match.group(1)!r}"
+                )
+            dp_group_alias = dp_match.group(1)
+            continue
+
+        repl_match = _REPLICATED_PATTERN.match(token)
+        if repl_match is not None:
+            axis_str: str = repl_match.group(1)
+            axis: Optional[ParallelAxis] = _AXIS_LOOKUP.get(axis_str)
+            if axis is None:
+                raise ValueError(
+                    f"Unknown axis {axis_str!r} in replicated declaration: {token!r}"
+                )
+            if axis in replicated_axes:
+                raise ValueError(
+                    f"Duplicate replicated declaration for axis {axis_str!r}"
+                )
+            replicated_axes.add(axis)
+            continue
+
+        raise ValueError(
+            f"Unrecognized token {token!r} in # comment section. "
+            f"Expected 'dp:=<group>' or '<axis>:replicated'."
+        )
+
+    return _CommentSuffix(
+        dp_group_alias=dp_group_alias,
+        replicated_axes=frozenset(replicated_axes),
+    )

python/sglang/srt/debug_utils/comparator/dims_spec/dim_parser.py

@@ -0,0 +1,68 @@
+from __future__ import annotations
+
+import re
+from typing import Optional
+
+from sglang.srt.debug_utils.comparator.dims_spec.modifier_parser import (
+    _parse_modifiers,
+)
+from sglang.srt.debug_utils.comparator.dims_spec.types import (
+    SQUEEZE_DIM_NAME,
+    DimSpec,
+    ParallelModifier,
+)
+
+_DIM_PATTERN = re.compile(r"^(?P<name>[a-zA-Z_]\w*)(?:\[(?P<modifiers>[^\]]+)\])?$")
+
+_FUSED_DIM_PATTERN = re.compile(r"^\((?P<inner>[^)]+)\)(?:\[(?P<modifiers>[^\]]+)\])?$")
+
+_SUB_DIM_NAME_PATTERN = re.compile(r"^[a-zA-Z_]\w*$")
+
+
+def parse_dim(token: str) -> DimSpec:
+    if token == SQUEEZE_DIM_NAME:
+        return DimSpec(name=SQUEEZE_DIM_NAME)
+
+    fused_match = _FUSED_DIM_PATTERN.match(token)
+    if fused_match is not None:
+        return _parse_fused_dim(token=token, fused_match=fused_match)
+
+    return _parse_single_dim(token)
+
+
+def _parse_single_dim(token: str) -> DimSpec:
+    match = _DIM_PATTERN.match(token)
+    if match is None:
+        raise ValueError(f"Invalid dim token: {token!r}")
+
+    name: str = match.group("name")
+    modifiers: list[ParallelModifier] = _parse_modifiers(
+        modifiers_str=match.group("modifiers"), dim_token=token
+    )
+    return DimSpec(name=name, parallel_modifiers=modifiers)
+
+
+def _parse_fused_dim(*, token: str, fused_match: re.Match[str]) -> DimSpec:
+    inner: str = fused_match.group("inner")
+    modifiers_str: Optional[str] = fused_match.group("modifiers")
+
+    sub_names: list[str] = [s.strip() for s in inner.split("*")]
+    for sub_name in sub_names:
+        if not _SUB_DIM_NAME_PATTERN.match(sub_name):
+            raise ValueError(
+                f"Invalid sub-dim {sub_name!r} in fused dim token: {token!r}"
+            )
+
+    if len(sub_names) != len(set(sub_names)):
+        raise ValueError(f"Duplicate sub-dim names in fused dim token: {token!r}")
+
+    if len(sub_names) < 2:
+        raise ValueError(
+            f"Fused dim must have at least 2 sub-dims, got {len(sub_names)} in: {token!r}"
+        )
+
+    fused_name: str = "*".join(sub_names)
+    modifiers: list[ParallelModifier] = _parse_modifiers(
+        modifiers_str=modifiers_str, dim_token=token
+    )
+    return DimSpec(name=fused_name, parallel_modifiers=modifiers)

python/sglang/srt/debug_utils/comparator/dims_spec/dims_parser.py

@@ -0,0 +1,113 @@
+from __future__ import annotations
+
+from typing import Optional
+
+from sglang.srt.debug_utils.comparator.dims_spec.comment_parser import (
+    _CommentSuffix,
+    _parse_comment_suffix,
+)
+from sglang.srt.debug_utils.comparator.dims_spec.dim_parser import parse_dim
+from sglang.srt.debug_utils.comparator.dims_spec.types import (
+    SQUEEZE_DIM_NAME,
+    DimSpec,
+    DimsSpec,
+    ParallelAxis,
+)
+
+
+class _SingletonDimUtil:
+    """Utilities for squeeze dims (name="1") and their singleton tensor-name mapping."""
+
+    PREFIX: str = "singleton"
+
+    @staticmethod
+    def is_squeeze(spec: DimSpec) -> bool:
+        return spec.name == SQUEEZE_DIM_NAME
+
+    @staticmethod
+    def filter_out(dim_specs: list[DimSpec]) -> list[DimSpec]:
+        return [s for s in dim_specs if not _SingletonDimUtil.is_squeeze(s)]
+
+    @staticmethod
+    def make_name(index: int) -> str:
+        return f"{_SingletonDimUtil.PREFIX}{index}"
+
+    @staticmethod
+    def is_singleton_name(name: str) -> bool:
+        return (
+            name.startswith(_SingletonDimUtil.PREFIX)
+            and name[len(_SingletonDimUtil.PREFIX) :].isdigit()
+        )
+
+    @staticmethod
+    def sanitize_names(names: list[str]) -> list[str]:
+        """Replace '1' with 'singleton0', 'singleton1', ... for named tensor compatibility."""
+        result: list[str] = []
+        sq_idx: int = 0
+
+        for name in names:
+            if name == SQUEEZE_DIM_NAME:
+                result.append(_SingletonDimUtil.make_name(sq_idx))
+                sq_idx += 1
+            else:
+                result.append(name)
+
+        return result
+
+
+def parse_dims(dims_str: str) -> DimsSpec:
+    """Parse ``"b s[cp:zigzag] h[tp] d # dp:=moe_dp ep:replicated"`` → :class:`DimsSpec`.
+
+    The shape part (before ``#``) produces :pyattr:`DimsSpec.dims`.
+    The declaration part (after ``#``) is scanned for:
+    - ``dp:=<group>`` → :pyattr:`DimsSpec.dp_group_alias`
+    - ``axis:replicated`` → :pyattr:`DimsSpec.replicated_axes`
+    """
+    parts: list[str] = dims_str.split("#", maxsplit=1)
+    raw: str = parts[0]
+
+    if not raw.strip():
+        raise ValueError("dims string must not be empty")
+
+    dims: list[DimSpec] = [parse_dim(token) for token in raw.strip().split()]
+
+    # Collect all semantic names (expanding fused sub-dims) for duplicate detection
+    semantic_names: list[str] = []
+    for spec in dims:
+        if _SingletonDimUtil.is_squeeze(spec):
+            continue
+        semantic_names.extend(spec.sub_dims)
+
+    if len(semantic_names) != len(set(semantic_names)):
+        duplicates = sorted({n for n in semantic_names if semantic_names.count(n) > 1})
+        raise ValueError(f"Duplicate dim names: {duplicates}")
+
+    comment_suffix: _CommentSuffix = (
+        _parse_comment_suffix(parts[1]) if len(parts) > 1 else _CommentSuffix()
+    )
+    dp_group_alias: Optional[str] = comment_suffix.dp_group_alias
+    replicated_axes: frozenset[ParallelAxis] = comment_suffix.replicated_axes
+
+    sharded_axes: set[ParallelAxis] = {
+        m.axis for spec in dims for m in spec.parallel_modifiers
+    }
+    conflict: frozenset[ParallelAxis] = replicated_axes & sharded_axes
+    if conflict:
+        conflict_names: str = ", ".join(sorted(a.value for a in conflict))
+        raise ValueError(
+            f"Axes declared as both sharded (in dim spec) and replicated "
+            f"(in # declaration): {conflict_names}"
+        )
+
+    return DimsSpec(
+        dims=dims,
+        dp_group_alias=dp_group_alias,
+        replicated_axes=replicated_axes,
+    )
+
+
+def resolve_dim_names(dims_str: str) -> list[str]:
+    """Parse dims string and return tensor-compatible names ('1' → 'singleton0', ...)."""
+    specs: list[DimSpec] = parse_dims(dims_str).dims
+    names: list[str] = [spec.sanitized_name for spec in specs]
+    return _SingletonDimUtil.sanitize_names(names)

python/sglang/srt/debug_utils/comparator/dims_spec/modifier_parser.py

@@ -0,0 +1,84 @@
+from __future__ import annotations
+
+from typing import Optional
+
+from sglang.srt.debug_utils.comparator.dims_spec.types import (
+    _AXIS_LOOKUP,
+    _QUALIFIER_LOOKUP,
+    Ordering,
+    ParallelAxis,
+    ParallelModifier,
+    Reduction,
+)
+
+
+def _parse_modifier_token(modifier_token: str, dim_token: str) -> ParallelModifier:
+    """Parse 'sp', 'cp:zigzag', 'tp:partial', or 'cp:zigzag+partial' → ParallelModifier.
+
+    Format: ``axis`` or ``axis:qual`` or ``axis:qual+qual``.
+    Colon separates axis from qualifiers; ``+`` separates multiple qualifiers.
+    """
+    axis_str: str
+    qualifiers_str: str
+    if ":" in modifier_token:
+        axis_str, qualifiers_str = modifier_token.split(":", maxsplit=1)
+    else:
+        axis_str, qualifiers_str = modifier_token, ""
+
+    axis_str = axis_str.strip()
+    axis: Optional[ParallelAxis] = _AXIS_LOOKUP.get(axis_str)
+    if axis is None:
+        raise ValueError(
+            f"Unknown axis {axis_str!r} in modifier {modifier_token!r} "
+            f"of dim spec: {dim_token!r}"
+        )
+
+    ordering: Optional[Ordering] = None
+    reduction: Optional[Reduction] = None
+
+    for q_str in (q.strip() for q in qualifiers_str.split("+") if q.strip()):
+        if q_str == "sharded":
+            continue
+        qualifier: Optional[Ordering | Reduction] = _QUALIFIER_LOOKUP.get(q_str)
+        if qualifier is None:
+            raise ValueError(
+                f"Unknown qualifier {q_str!r} in modifier "
+                f"{modifier_token!r} of dim spec: {dim_token!r}"
+            )
+        if isinstance(qualifier, Ordering):
+            if ordering is not None:
+                raise ValueError(
+                    f"Multiple ordering values in modifier "
+                    f"{modifier_token!r} of dim spec: {dim_token!r}"
+                )
+            ordering = qualifier
+        else:
+            if reduction is not None:
+                raise ValueError(
+                    f"Multiple reduction values in modifier "
+                    f"{modifier_token!r} of dim spec: {dim_token!r}"
+                )
+            reduction = qualifier
+
+    return ParallelModifier(axis=axis, ordering=ordering, reduction=reduction)
+
+
+def _parse_modifiers(
+    *, modifiers_str: Optional[str], dim_token: str
+) -> list[ParallelModifier]:
+    if modifiers_str is None:
+        return []
+
+    modifiers: list[ParallelModifier] = []
+    seen_axes: set[ParallelAxis] = set()
+
+    for modifier_token in (p.strip() for p in modifiers_str.split(",")):
+        modifier: ParallelModifier = _parse_modifier_token(modifier_token, dim_token)
+        if modifier.axis in seen_axes:
+            raise ValueError(
+                f"Duplicate axis {modifier.axis.value!r} in dim spec: {dim_token!r}"
+            )
+        seen_axes.add(modifier.axis)
+        modifiers.append(modifier)
+
+    return modifiers

python/sglang/srt/debug_utils/comparator/dims_spec/tensor_naming.py

@@ -0,0 +1,56 @@
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+
+from sglang.srt.debug_utils.comparator.dims_spec.types import DimSpec
+
+_DIM_NAMES_ATTR = "_dim_names"
+
+
+def find_dim_index(dim_specs: list[DimSpec], name: str) -> Optional[int]:
+    """Find index by name. Accepts both ``*``-form and ``___``-form for fused dims."""
+    for i, spec in enumerate(dim_specs):
+        if spec.name == name or spec.sanitized_name == name:
+            return i
+    return None
+
+
+def get_dim_names(tensor: torch.Tensor) -> tuple[Optional[str], ...]:
+    """Get dimension names attached to a tensor.
+
+    Returns a tuple of ``None`` values if no names are attached.
+    """
+    names = getattr(tensor, _DIM_NAMES_ATTR, None)
+    if names is not None:
+        return names
+    return (None,) * tensor.ndim
+
+
+def resolve_dim_by_name(tensor: torch.Tensor, name: str) -> int:
+    names = get_dim_names(tensor)
+    if names[0] is None:
+        raise ValueError(f"Tensor has no names, cannot resolve {name!r}")
+
+    try:
+        return list(names).index(name)
+    except ValueError:
+        raise ValueError(f"Dim name {name!r} not in tensor names {names}")
+
+
+def apply_dim_names(tensor: torch.Tensor, dim_names: list[str]) -> torch.Tensor:
+    if tensor.ndim != len(dim_names):
+        raise ValueError(
+            f"dims metadata mismatch: tensor has {tensor.ndim} dims (shape {list(tensor.shape)}) "
+            f"but dims string specifies {len(dim_names)} names {dim_names}. "
+            f"Please fix the dims string in the dumper.dump() call to match the actual tensor shape."
+        )
+    view = torch.ops.aten.alias(tensor)
+    view._dim_names = tuple(dim_names)
+    return view
+
+
+def without_dim_names(tensor: torch.Tensor) -> torch.Tensor:
+    # Returns a new view without _dim_names; the original tensor is not modified.
+    return torch.ops.aten.alias(tensor)

python/sglang/srt/debug_utils/comparator/dims_spec/types.py

@@ -0,0 +1,94 @@
+from __future__ import annotations
+
+from enum import Enum
+from typing import Optional
+
+from sglang.srt.debug_utils.comparator.utils import _FrozenBase
+
+TOKEN_DIM_NAME: str = "t"
+BATCH_DIM_NAME: str = "b"
+SEQ_DIM_NAME: str = "s"
+SQUEEZE_DIM_NAME: str = "1"
+
+
+class TokenLayout(Enum):
+    T = "t"  # single flat token dim
+    BS = "bs"  # separate batch + seq dims, need collapse
+
+
+# TODO: allow arbitrary string
+class ParallelAxis(Enum):
+    TP = "tp"
+    CP = "cp"
+    EP = "ep"
+    SP = "sp"
+    DP = "dp"
+    ETP = "etp"
+    EDP = "edp"
+    ATTN_TP = "attn_tp"
+    ATTN_DP = "attn_dp"
+    MOE_EP = "moe_ep"
+    MOE_TP = "moe_tp"
+    MOE_DP = "moe_dp"
+    RECOMPUTE_PSEUDO = "recompute_pseudo"
+
+
+class Ordering(Enum):
+    ZIGZAG = "zigzag"
+    NATURAL = "natural"
+
+
+class Reduction(Enum):
+    PARTIAL = "partial"
+
+
+class ParallelModifier(_FrozenBase):
+    axis: ParallelAxis
+    ordering: Optional[Ordering] = None
+    reduction: Optional[Reduction] = None
+
+
+_AXIS_LOOKUP: dict[str, ParallelAxis] = {m.value: m for m in ParallelAxis}
+_QUALIFIER_LOOKUP: dict[str, Ordering | Reduction] = {
+    **{m.value: m for m in Ordering},
+    **{m.value: m for m in Reduction},
+}
+
+_FUSED_NAME_SEP: str = "___"
+
+
+class DimSpec(_FrozenBase):
+    name: str
+    parallel_modifiers: list[ParallelModifier] = []
+
+    @property
+    def sub_dims(self) -> list[str]:
+        """Sub-dim names. Fused: ``["num_heads", "head_dim"]``; plain: ``["h"]``."""
+        return self.name.split("*")
+
+    @property
+    def is_fused(self) -> bool:
+        return len(self.sub_dims) > 1
+
+    @property
+    def sanitized_name(self) -> str:
+        """Name safe for PyTorch named tensors (``*`` → ``___``)."""
+        if self.is_fused:
+            return _FUSED_NAME_SEP.join(self.sub_dims)
+        return self.name
+
+
+class DimsSpec(_FrozenBase):
+    """Parsed result of a full dims string like ``"b s h[tp] # dp:=moe_dp"``."""
+
+    dims: list[DimSpec]
+    dp_group_alias: Optional[str] = None
+    replicated_axes: frozenset[ParallelAxis] = frozenset()
+
+    @property
+    def dp_axis(self) -> ParallelAxis:
+        return (
+            ParallelAxis(self.dp_group_alias)
+            if self.dp_group_alias
+            else ParallelAxis.DP
+        )

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

@@ -0,0 +1,144 @@
+from __future__ import annotations
+
+from collections import defaultdict
+from io import StringIO
+from pathlib import Path
+from typing import Any, Optional
+
+import polars as pl
+
+from sglang.srt.debug_utils.comparator.output_types import (
+    InputIdsRecord,
+    RankInfoRecord,
+)
+from sglang.srt.debug_utils.comparator.report_sink import report_sink
+from sglang.srt.debug_utils.dump_loader import LOAD_FAILED, ValueWithMeta
+
+PARALLEL_INFO_KEYS: list[str] = ["sglang_parallel_info", "megatron_parallel_info"]
+
+
+def emit_display_records(
+    *,
+    df: pl.DataFrame,
+    dump_dir: Path,
+    label: str,
+    tokenizer: Any,
+) -> None:
+    rank_rows: Optional[list[dict[str, Any]]] = _collect_rank_info(
+        df, dump_dir=dump_dir
+    )
+    if rank_rows is not None:
+        report_sink.add(RankInfoRecord(label=label, rows=rank_rows))
+
+    input_ids_rows: Optional[list[dict[str, Any]]] = _collect_input_ids_and_positions(
+        df, dump_dir=dump_dir, tokenizer=tokenizer
+    )
+    if input_ids_rows is not None:
+        report_sink.add(InputIdsRecord(label=label, rows=input_ids_rows))
+
+
+def _render_polars_as_text(df: pl.DataFrame, *, title: Optional[str] = None) -> str:
+    from rich.console import Console
+    from rich.table import Table
+
+    table = Table(title=title)
+    for col in df.columns:
+        table.add_column(col)
+    for row in df.iter_rows():
+        table.add_row(*[str(v) for v in row])
+
+    buf = StringIO()
+    Console(file=buf, force_terminal=False, width=200).print(table)
+    return buf.getvalue().rstrip("\n")
+
+
+def _render_polars_as_rich_table(
+    df: pl.DataFrame, *, title: Optional[str] = None
+) -> Any:
+    from rich.table import Table
+
+    table = Table(title=title)
+    for col in df.columns:
+        table.add_column(col)
+    for row in df.iter_rows():
+        table.add_row(*[str(v) for v in row])
+    return table
+
+
+def _collect_rank_info(
+    df: pl.DataFrame, dump_dir: Path
+) -> Optional[list[dict[str, Any]]]:
+    unique_rows: pl.DataFrame = (
+        df.filter(pl.col("name") == "input_ids")
+        .sort("rank")
+        .unique(subset=["rank"], keep="first")
+    )
+    if unique_rows.is_empty():
+        return None
+
+    table_rows: list[dict[str, Any]] = []
+    for row in unique_rows.to_dicts():
+        meta: dict[str, Any] = ValueWithMeta.load(dump_dir / row["filename"]).meta
+
+        row_data: dict[str, Any] = {"rank": row["rank"]}
+        for key in PARALLEL_INFO_KEYS:
+            _extract_parallel_info(row_data=row_data, info=meta.get(key, {}))
+        table_rows.append(row_data)
+
+    return table_rows or None
+
+
+def _collect_input_ids_and_positions(
+    df: pl.DataFrame,
+    dump_dir: Path,
+    *,
+    tokenizer: Any = None,
+) -> Optional[list[dict[str, Any]]]:
+    filtered: pl.DataFrame = df.filter(pl.col("name").is_in(["input_ids", "positions"]))
+    if filtered.is_empty():
+        return None
+
+    data_by_step_rank: dict[tuple[int, int], dict[str, Any]] = defaultdict(dict)
+    for row in filtered.to_dicts():
+        key: tuple[int, int] = (row["step"], row["rank"])
+        item: ValueWithMeta = ValueWithMeta.load(dump_dir / row["filename"])
+        if item.value is not LOAD_FAILED:
+            data_by_step_rank[key][row["name"]] = item.value
+
+    table_rows: list[dict[str, Any]] = []
+    for (step, rank), data in sorted(data_by_step_rank.items()):
+        ids = data.get("input_ids")
+        pos = data.get("positions")
+
+        ids_list: Optional[list[int]] = (
+            ids.flatten().tolist() if ids is not None else None
+        )
+
+        row_data: dict[str, Any] = {
+            "step": step,
+            "rank": rank,
+            "num_tokens": len(ids_list) if ids_list is not None else None,
+            "input_ids": str(ids_list) if ids_list is not None else "N/A",
+            "positions": str(pos.flatten().tolist()) if pos is not None else "N/A",
+        }
+
+        if tokenizer is not None and ids_list is not None:
+            row_data["decoded_text"] = repr(
+                tokenizer.decode(ids_list, skip_special_tokens=False)
+            )
+
+        table_rows.append(row_data)
+
+    return table_rows or None
+
+
+def _extract_parallel_info(row_data: dict[str, Any], info: dict[str, Any]) -> None:
+    if not info or info.get("error"):
+        return
+
+    for key in sorted(info.keys()):
+        if key.endswith("_rank"):
+            base: str = key[:-5]
+            size_key: str = f"{base}_size"
+            if size_key in info:
+                row_data[base] = f"{info[key]}/{info[size_key]}"

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

@@ -0,0 +1,100 @@
+"""DP filtering: keep only the non-empty dp_rank items."""
+
+from __future__ import annotations
+
+from collections import defaultdict
+from typing import Optional
+
+import torch
+
+from sglang.srt.debug_utils.comparator.dims_spec import ParallelAxis
+from sglang.srt.debug_utils.dump_loader import ValueWithMeta
+
+_PARALLEL_INFO_KEYS = ("sglang_parallel_info", "megatron_parallel_info")
+
+
+def filter_to_non_empty_dp_rank(
+    items: list[ValueWithMeta],
+    *,
+    dp_axis: ParallelAxis,
+) -> list[ValueWithMeta]:
+    """Filter items to the single non-empty dp_rank.
+
+    - dp_size <= 1: return items unchanged.
+    - dp_size > 1: group by dp_rank, assert exactly one group has non-empty
+      tensors, return that group.
+
+    *dp_axis* determines which rank/size fields to look up (e.g.
+    ``ParallelAxis.MOE_DP`` → ``moe_dp_rank`` / ``moe_dp_size``).
+    If the fields are absent the filter is a noop (items returned unchanged).
+    """
+    if not items:
+        return items
+
+    dp_info: Optional[tuple[int, int]] = _extract_dp_info(
+        items[0].meta, dp_axis=dp_axis
+    )
+    if dp_info is None:
+        return items
+
+    _dp_rank, dp_size = dp_info
+    if dp_size <= 1:
+        return items
+
+    has_any_tensor: bool = any(isinstance(item.value, torch.Tensor) for item in items)
+    if not has_any_tensor:
+        return items
+
+    groups: dict[int, list[ValueWithMeta]] = defaultdict(list)
+    for item in items:
+        item_dp: Optional[tuple[int, int]] = _extract_dp_info(
+            item.meta, dp_axis=dp_axis
+        )
+        rank: int = item_dp[0] if item_dp is not None else 0
+        groups[rank].append(item)
+
+    non_empty_ranks: list[int] = [
+        rank for rank, group in groups.items() if _group_has_data(group)
+    ]
+
+    assert len(non_empty_ranks) == 1, (
+        f"Expected exactly 1 non-empty dp_rank, got {len(non_empty_ranks)}: "
+        f"ranks={non_empty_ranks}"
+    )
+
+    return groups[non_empty_ranks[0]]
+
+
+def _extract_dp_info(
+    meta: dict,
+    *,
+    dp_axis: ParallelAxis,
+) -> Optional[tuple[int, int]]:
+    """Extract (dp_rank, dp_size) from meta's parallel_info block.
+
+    *dp_axis* determines which fields to look up: e.g.
+    ``ParallelAxis.DP`` → ``dp_rank``/``dp_size``,
+    ``ParallelAxis.MOE_DP`` → ``moe_dp_rank``/``moe_dp_size``.
+    """
+    rank_field: str = f"{dp_axis.value}_rank"
+    size_field: str = f"{dp_axis.value}_size"
+
+    for key in _PARALLEL_INFO_KEYS:
+        info = meta.get(key)
+        if not isinstance(info, dict) or not info:
+            continue
+
+        dp_rank = info.get(rank_field)
+        dp_size = info.get(size_field)
+        if dp_rank is not None and dp_size is not None:
+            return (int(dp_rank), int(dp_size))
+
+    return None
+
+
+def _group_has_data(group: list[ValueWithMeta]) -> bool:
+    """Check if any tensor in the group is non-empty (numel > 0)."""
+    return any(
+        isinstance(item.value, torch.Tensor) and item.value.numel() > 0
+        for item in group
+    )

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

@@ -0,0 +1,475 @@
+from __future__ import annotations
+
+import argparse
+import sys
+import traceback as _traceback_module
+from pathlib import Path
+from typing import Any, Iterator, Optional, Union
+
+import polars as pl
+
+from sglang.srt.debug_utils.comparator.aligner.token_aligner.entrypoint import (
+    TokenAlignerResult,
+    compute_maybe_token_aligner_result,
+)
+from sglang.srt.debug_utils.comparator.aligner.token_aligner.smart.aux_loader import (
+    AUX_NAMES,
+)
+from sglang.srt.debug_utils.comparator.aligner.token_aligner.smart.types import (
+    TokenAlignerPlan,
+)
+from sglang.srt.debug_utils.comparator.bundle_comparator import compare_bundle_pair
+from sglang.srt.debug_utils.comparator.bundle_matcher import (
+    TensorBundleInfo,
+    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 (
+    ComparisonErrorRecord,
+    ComparisonNonTensorRecord,
+    ComparisonSkipRecord,
+    ComparisonTensorRecord,
+    ConfigRecord,
+    RecordLocation,
+    SummaryRecord,
+)
+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.report_sink import report_sink
+from sglang.srt.debug_utils.comparator.tensor_comparator.comparator import (
+    DEFAULT_PREDICATE,
+    FailureDisplayBudget,
+)
+from sglang.srt.debug_utils.comparator.threshold_dsl import (
+    DiffThresholdRule,
+    parse_diff_threshold_rules,
+)
+from sglang.srt.debug_utils.comparator.utils import (
+    Pair,
+    auto_descend_dir,
+    compute_exit_code,
+)
+from sglang.srt.debug_utils.dump_loader import read_meta, read_tokenizer_path
+
+_DEFAULT_SKIP_KEYS: set[str] = {"dump_index", "filename"}
+
+_DIMS_DEBUG_HINT: str = (
+    "\nHint: If this is a dims annotation issue, do NOT re-run expensive dumps.\n"
+    "Use --override-dims at comparison time, e.g.:\n"
+    '  python -m sglang.srt.debug_utils.comparator --override-dims "tensor_name:b s h[tp] d"\n'
+    "(Use --override-baseline-dims / --override-target-dims for per-side overrides.\n"
+    " Use --override-config for bulk overrides via YAML file.)"
+)
+
+
+def main() -> None:
+    args = parse_args(sys.argv[1:])
+    sys.exit(run(args))
+
+
+def run(args: argparse.Namespace) -> int:
+    report_sink.configure(
+        output_format=args.output_format,
+        report_path=None,
+        verbosity=args.verbosity,
+    )
+
+    dir_pair: Pair[Path] = Pair(
+        x=auto_descend_dir(Path(args.baseline_path), label="baseline_path"),
+        y=auto_descend_dir(Path(args.target_path), label="target_path"),
+    )
+    viz_output_dir: Optional[Path] = (
+        Path(args.viz_output_dir) if args.viz_bundle_details else None
+    )
+    visualize_per_token: Optional[Path] = (
+        Path(args.visualize_per_token) if args.visualize_per_token else None
+    )
+    override_config: Optional[Path] = (
+        Path(args.override_config) if args.override_config else None
+    )
+
+    report_path: Optional[Path] = _resolve_report_path(
+        target_path=dir_pair.y,
+        report_path_arg=args.report_path,
+    )
+    report_sink.configure(
+        output_format=args.output_format,
+        report_path=report_path,
+        verbosity=args.verbosity,
+    )
+
+    try:
+        report_sink.add(ConfigRecord(config=vars(args)))
+
+        dfs: Pair[pl.DataFrame] = _read_df(
+            dir_pair=dir_pair,
+            start_step=args.start_step,
+            end_step=args.end_step,
+            filter_pattern=args.filter,
+        )
+
+        tokenizer: Any = _maybe_load_tokenizer(
+            tokenizer_arg=args.tokenizer, dir_pair=dir_pair
+        )
+        for label, df, dump_dir in [
+            ("baseline", dfs.x, dir_pair.x),
+            ("target", dfs.y, dir_pair.y),
+        ]:
+            emit_display_records(
+                df=df, dump_dir=dump_dir, label=label, tokenizer=tokenizer
+            )
+
+        ta_result: TokenAlignerResult = compute_maybe_token_aligner_result(
+            dir_pair=dir_pair,
+            dfs=dfs,
+            token_aligner_mode=args.token_aligner,
+        )
+
+        if ta_result.mode == "smart":
+            dfs = dfs.map(lambda df: df.filter(~pl.col("name").is_in(AUX_NAMES)))
+
+        skip_keys: set[str] = _DEFAULT_SKIP_KEYS | set(args.grouping_skip_keys or [])
+        bundle_info_pairs: list[Pair[TensorBundleInfo]] = match_bundles(
+            dfs=dfs, skip_keys=skip_keys
+        )
+
+        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=override_config,
+        )
+
+        comparison_records = _compare_bundle_pairs(
+            bundle_info_pairs=bundle_info_pairs,
+            dir_pair=dir_pair,
+            token_aligner_mode=ta_result.mode,
+            token_aligner_plan=ta_result.plan,
+            diff_threshold_rules=parse_diff_threshold_rules(
+                args.diff_threshold, default_predicate=DEFAULT_PREDICATE
+            ),
+            failure_display_budget=FailureDisplayBudget(),
+            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,
+        )
+        summary, skipped_names, failed_names, errored_names = (
+            _consume_comparison_records(
+                comparison_records=comparison_records,
+                visualize_per_token=visualize_per_token,
+            )
+        )
+        return compute_exit_code(
+            summary,
+            allow_skipped_pattern=args.allow_skipped_pattern,
+            skipped_names=skipped_names,
+            allow_failed_pattern=args.allow_failed_pattern,
+            failed_names=failed_names,
+            errored_names=errored_names,
+        )
+    finally:
+        report_sink.close()
+        if report_path is not None:
+            print(f"Report: {report_path}", file=sys.stderr)
+
+
+def _resolve_report_path(
+    *, target_path: Path, report_path_arg: Optional[str]
+) -> Optional[Path]:
+    if report_path_arg is not None:
+        return Path(report_path_arg) if report_path_arg else None
+    return target_path / "comparator_report.jsonl"
+
+
+def _maybe_load_tokenizer(*, tokenizer_arg: Optional[str], dir_pair: Pair[Path]) -> Any:
+    tokenizer_path: Optional[str] = tokenizer_arg
+
+    if tokenizer_path is None:
+        for directory in [dir_pair.x, dir_pair.y]:
+            tokenizer_path = read_tokenizer_path(directory)
+            if tokenizer_path is not None:
+                break
+
+    if tokenizer_path is None:
+        return None
+
+    try:
+        from transformers import AutoTokenizer
+
+        return AutoTokenizer.from_pretrained(tokenizer_path)
+    except Exception:
+        return None
+
+
+def _read_df(
+    *,
+    dir_pair: Pair[Path],
+    start_step: int,
+    end_step: int,
+    filter_pattern: Optional[str],
+) -> Pair[pl.DataFrame]:
+    df_baseline = read_meta(dir_pair.x)
+
+    df_target = read_meta(dir_pair.y)
+    df_target = df_target.filter(
+        (pl.col("step") >= start_step) & (pl.col("step") <= end_step)
+    )
+    if filter_pattern:
+        df_target = df_target.filter(pl.col("filename").str.contains(filter_pattern))
+    assert all(c in df_target.columns for c in ["rank", "step", "dump_index", "name"])
+
+    return Pair(x=df_baseline, y=df_target)
+
+
+def _compare_bundle_pairs(
+    *,
+    bundle_info_pairs: list[Pair[TensorBundleInfo]],
+    dir_pair: Pair[Path],
+    token_aligner_mode: Optional[str],
+    token_aligner_plan: Optional[TokenAlignerPlan],
+    diff_threshold_rules: Optional[list[DiffThresholdRule]] = None,
+    failure_display_budget: Optional[FailureDisplayBudget] = None,
+    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[
+        ComparisonTensorRecord,
+        ComparisonSkipRecord,
+        ComparisonNonTensorRecord,
+        ComparisonErrorRecord,
+    ]
+]:
+    for bundle_info_pair in bundle_info_pairs:
+        if not bundle_info_pair.y:
+            continue
+
+        name: str = bundle_info_pair.y[0].name
+        filenames_pair: Pair[list[str]] = bundle_info_pair.map(
+            lambda infos: [info.filename for info in infos]
+        )
+
+        record: Union[
+            ComparisonTensorRecord,
+            ComparisonSkipRecord,
+            ComparisonNonTensorRecord,
+            ComparisonErrorRecord,
+        ]
+        try:
+            record = compare_bundle_pair(
+                name=name,
+                filenames_pair=filenames_pair,
+                dir_pair=dir_pair,
+                token_aligner_mode=token_aligner_mode,
+                token_aligner_plan=token_aligner_plan,
+                diff_threshold_rules=diff_threshold_rules,
+                failure_display_budget=failure_display_budget,
+                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,
+            )
+        except Exception as exc:
+            tb = _traceback_module.format_exc()
+            record = ComparisonErrorRecord(
+                name=name,
+                exception_type=type(exc).__name__,
+                exception_message=str(exc),
+                traceback_str=f"{_DIMS_DEBUG_HINT}\n\n{tb}",
+            )
+
+        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[
+            ComparisonTensorRecord,
+            ComparisonSkipRecord,
+            ComparisonNonTensorRecord,
+            ComparisonErrorRecord,
+        ]
+    ],
+    visualize_per_token: Optional[Path] = None,
+) -> tuple[SummaryRecord, list[str], list[str], list[str]]:
+    counts: dict[str, int] = {"passed": 0, "failed": 0, "skipped": 0, "errored": 0}
+    collected_comparisons: list[ComparisonTensorRecord] = []
+    skipped_names: list[str] = []
+    failed_names: list[str] = []
+    errored_names: list[str] = []
+
+    for record in comparison_records:
+        counts[record.category] += 1
+        report_sink.add(record)
+        if isinstance(record, ComparisonSkipRecord) and record.category == "skipped":
+            skipped_names.append(record.name)
+        if record.category == "failed":
+            failed_names.append(record.name)
+        if isinstance(record, ComparisonErrorRecord):
+            errored_names.append(record.name)
+        if visualize_per_token is not None and isinstance(
+            record, ComparisonTensorRecord
+        ):
+            collected_comparisons.append(record)
+
+    summary: SummaryRecord = SummaryRecord(total=sum(counts.values()), **counts)
+    report_sink.add(summary)
+
+    if visualize_per_token is not None and collected_comparisons:
+        generate_per_token_heatmap(
+            records=collected_comparisons,
+            output_path=visualize_per_token,
+        )
+
+    return summary, skipped_names, failed_names, errored_names
+
+
+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)
+    parser.add_argument("--start-step", type=int, default=0)
+    parser.add_argument("--end-step", type=int, default=1000000)
+    parser.add_argument(
+        "--diff-threshold",
+        nargs="*",
+        default=None,
+        metavar="REGEX PREDICATE",
+        help="Per-tensor pass criterion. Either a single float shorthand "
+        "(0.0085 == '.*' 'rel <= 0.0085'), or (regex predicate) pairs, e.g. "
+        "--diff-threshold '.*expert.*' 'rel <= 0.0085 or max_abs <= 1e-3' '.*' 'rel <= 0.0085'. "
+        "A tensor uses the first fullmatching regex's predicate -- a boolean expression "
+        "over rel/max_abs/mean_abs with < <= > >= and and/or. A tensor matching no "
+        "pattern is an error. Default: 'rel <= 1e-3' for every tensor.",
+    )
+    parser.add_argument(
+        "--filter", type=str, default=None, help="Regex to filter filenames (include)"
+    )
+    parser.add_argument(
+        "--output-format",
+        type=str,
+        choices=["text", "json"],
+        default="text",
+        help="Output format: text (default) or json (JSONL, one JSON object per line)",
+    )
+    parser.add_argument(
+        "--verbosity",
+        type=str,
+        choices=["minimal", "normal", "verbose"],
+        default="normal",
+        help="Output verbosity: minimal (1 line per tensor), normal (compact lifecycle), "
+        "verbose (full detail). Default: normal",
+    )
+    parser.add_argument(
+        "--preset",
+        type=str,
+        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=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",
+        type=str,
+        default=None,
+        help="Tokenizer path for decoding input_ids (auto-discovered from dump metadata if not set)",
+    )
+    parser.add_argument(
+        "--viz-bundle-details",
+        action="store_true",
+        default=False,
+        help="Generate comparison heatmap/histogram PNG for each compared tensor",
+    )
+    parser.add_argument(
+        "--viz-output-dir",
+        type=str,
+        default="/tmp/comparator_viz/",
+        help="Output directory for visualization PNGs (default: /tmp/comparator_viz/)",
+    )
+    parser.add_argument(
+        "--visualize-per-token",
+        type=str,
+        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.)",
+    )
+    parser.add_argument(
+        "--allow-skipped-pattern",
+        type=str,
+        default=".*",
+        help="Regex pattern for tensor names allowed to be skipped. "
+        "Default '.*' allows all skips. Use '^$' to forbid all skips.",
+    )
+    parser.add_argument(
+        "--allow-failed-pattern",
+        type=str,
+        default=None,
+        help="Regex pattern for tensor names allowed to fail without affecting exit code. "
+        "Default None (all failures affect exit code).",
+    )
+
+    # Report output
+    parser.add_argument(
+        "--report-path",
+        type=str,
+        default=None,
+        help="Path for JSONL report (default: <target-path>/comparator_report.jsonl). "
+        "Pass empty string '' to disable.",
+    )
+
+    return parser.parse_args(argv)

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

@@ -0,0 +1,37 @@
+from __future__ import annotations
+
+from contextlib import contextmanager
+from typing import Generator
+
+from sglang.srt.debug_utils.comparator.output_types import BaseLog
+
+
+class LogSink:
+    def __init__(self) -> None:
+        self._stack: list[list[BaseLog]] = []
+
+    @contextmanager
+    def context(self) -> Generator[list[BaseLog], None, None]:
+        bucket: list[BaseLog] = []
+        self._stack.append(bucket)
+        try:
+            yield bucket
+        finally:
+            popped = self._stack.pop()
+            assert popped is bucket
+
+    def add(self, log: BaseLog) -> None:
+        if self._stack:
+            self._stack[-1].append(log)
+        else:
+            from sglang.srt.debug_utils.comparator.output_types import (
+                LogRecord,
+                _split_logs,
+            )
+            from sglang.srt.debug_utils.comparator.report_sink import report_sink
+
+            errors, infos = _split_logs([log])
+            report_sink.add(LogRecord(errors=errors, infos=infos))
+
+
+log_sink = LogSink()

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

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

@@ -0,0 +1,398 @@
+"""Formatting functions for comparator output records.
+
+Extracted from output_types.py to separate data-structure definitions
+from rendering / formatting logic.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Literal
+
+from rich.console import Group
+from rich.markup import escape
+from rich.panel import Panel
+
+from sglang.srt.debug_utils.comparator.tensor_comparator.formatter import (
+    format_comparison,
+    format_replicated_checks,
+)
+
+if TYPE_CHECKING:
+    from rich.console import RenderableType
+
+    from sglang.srt.debug_utils.comparator.aligner.entrypoint.traced_types import (
+        TracedAlignerPlan,
+        TracedSubPlan,
+    )
+    from sglang.srt.debug_utils.comparator.aligner.entrypoint.types import AlignerPlan
+    from sglang.srt.debug_utils.comparator.output_types import (
+        ComparisonErrorRecord,
+        ComparisonNonTensorRecord,
+        ComparisonSkipRecord,
+        ComparisonTensorRecord,
+        ConfigRecord,
+        ErrorLog,
+        InfoLog,
+        LogRecord,
+        SummaryRecord,
+        _OutputRecord,
+        _TableRecord,
+    )
+
+Verbosity = Literal["minimal", "normal", "verbose"]
+
+
+# ── Record-level rendering (body + logs) ─────────────────────────────
+
+
+def _render_record_rich(
+    record: _OutputRecord, *, verbosity: Verbosity = "normal"
+) -> RenderableType:
+    body: RenderableType = record._format_rich_body(verbosity=verbosity)
+
+    log_lines: list[str] = _format_log_lines_rich(
+        errors=record.errors, infos=record.infos
+    )
+
+    if not log_lines:
+        return body
+
+    log_block: str = "\n".join(log_lines)
+    if isinstance(body, str):
+        return body + "\n" + log_block
+    return Group(body, log_block)
+
+
+def _render_record_text(record: _OutputRecord) -> str:
+    body: str = record._format_body()
+
+    log_suffix: str = _format_log_lines_text(errors=record.errors, infos=record.infos)
+
+    if log_suffix:
+        body += "\n" + log_suffix
+
+    return body
+
+
+def _format_log_lines_rich(
+    *, errors: list[ErrorLog], infos: list[InfoLog]
+) -> list[str]:
+    lines: list[str] = []
+
+    if errors:
+        lines.extend(f"  [red]✗ {e.to_text()}[/]" for e in errors)
+    if infos:
+        lines.extend(f"  [dim]ℹ {i.to_text()}[/]" for i in infos)
+
+    return lines
+
+
+def _format_log_lines_text(*, errors: list[ErrorLog], infos: list[InfoLog]) -> str:
+    lines: list[str] = []
+
+    if errors:
+        lines.extend(f"  ✗ {e.to_text()}" for e in errors)
+    if infos:
+        lines.extend(f"  ℹ {i.to_text()}" for i in infos)
+
+    return "\n".join(lines)
+
+
+# ── ConfigRecord ──────────────────────────────────────────────────────
+
+
+def _format_config_body(record: ConfigRecord) -> str:
+    return f"Config: {record.config}"
+
+
+def _format_config_rich_body(
+    record: ConfigRecord, verbosity: Verbosity = "normal"
+) -> RenderableType:
+    lines: list[str] = [f"  [bold]{k}[/] : {v}" for k, v in record.config.items()]
+    return Panel("\n".join(lines), title="Comparator Config", border_style="cyan")
+
+
+# ── ComparisonSkipRecord ─────────────────────────────────────────────
+
+
+def _format_skip_body(record: ComparisonSkipRecord) -> str:
+    text: str = (
+        f"Skip: {record.name}{record._format_location_suffix()} ({record.reason})"
+    )
+    if record.available_side is not None and record.available_tensor_info is not None:
+        info = record.available_tensor_info
+        text += f"\n  {record.available_side}: shape={info.shape} dtype={info.dtype}"
+        text += (
+            f" mean={info.stats.mean:.4f} std={info.stats.std:.4f}"
+            f" range=[{info.stats.min:.4f}, {info.stats.max:.4f}]"
+        )
+        if info.sample is not None:
+            text += f"\n  sample: {info.sample}"
+    return text
+
+
+def _format_skip_rich_body(
+    record: ComparisonSkipRecord, verbosity: Verbosity = "normal"
+) -> RenderableType:
+    suffix: str = record._format_location_suffix()
+    header: str = (
+        f"[dim]⊘ {escape(record.name)}{suffix} ── skipped ({escape(record.reason)})[/]"
+    )
+
+    if (
+        verbosity == "minimal"
+        or record.available_side is None
+        or record.available_tensor_info is None
+    ):
+        return header
+
+    info = record.available_tensor_info
+    side: str = record.available_side
+    dtype_str: str = info.dtype.replace("torch.", "")
+
+    lines: list[str] = [header]
+
+    # Bundle info line
+    if record.available_bundle_info is not None:
+        bi = record.available_bundle_info
+        shapes: list[list[int]] = [f.shape for f in bi.files]
+        unique_shapes: set[str] = {str(s) for s in shapes}
+        shape_desc: str = (
+            escape(str(shapes[0])) if len(unique_shapes) == 1 else "mixed shapes"
+        )
+        dims_part: str = f"  [dim]dims: {bi.dims}[/]" if bi.dims else ""
+        lines.append(
+            f"   {side:8s}  [cyan]{bi.num_files} files[/]"
+            f" × {shape_desc} {dtype_str}{dims_part}"
+        )
+    else:
+        lines.append(f"   {side:8s}  {escape(str(info.shape))} {dtype_str}")
+
+    # Stats line (compact single-side)
+    stats = info.stats
+    range_str: str = escape(f"[{stats.min:.4f}, {stats.max:.4f}]")
+    lines.append(
+        f"   [dim]stats[/]     mean={stats.mean:.4f}  std={stats.std:.4f}"
+        f"  range={range_str}"
+    )
+
+    # Sample
+    if info.sample is not None:
+        lines.append(f"   [dim]sample[/]    {escape(info.sample)}")
+
+    return "\n".join(lines)
+
+
+# ── ComparisonErrorRecord ────────────────────────────────────────────
+
+
+def _format_error_body(record: ComparisonErrorRecord) -> str:
+    prefix: str = record._format_location_prefix()
+    return (
+        f"{prefix}Error: {record.name} ({record.exception_type})\n"
+        f"{record.exception_message}\n"
+        f"{record.traceback_str}"
+    )
+
+
+def _format_error_rich_body(
+    record: ComparisonErrorRecord, verbosity: Verbosity = "normal"
+) -> RenderableType:
+    prefix: str = record._format_location_prefix_rich()
+    name: str = escape(record.name)
+    header: str = (
+        f"{prefix}[bold red]{name} ── errored ({escape(record.exception_type)}): "
+        f"{escape(record.exception_message)}[/]"
+    )
+    if verbosity == "minimal":
+        return header
+    return header + f"\n[dim]{escape(record.traceback_str)}[/]"
+
+
+# ── _TableRecord ─────────────────────────────────────────────────────
+
+
+def _format_table_body(record: _TableRecord) -> str:
+    import polars as pl
+
+    from sglang.srt.debug_utils.comparator.display import _render_polars_as_text
+
+    return _render_polars_as_text(
+        pl.DataFrame(record.rows), title=record._table_title()
+    )
+
+
+def _format_table_rich_body(
+    record: _TableRecord, verbosity: Verbosity = "normal"
+) -> RenderableType:
+    import polars as pl
+
+    from sglang.srt.debug_utils.comparator.display import (
+        _render_polars_as_rich_table,
+    )
+
+    return _render_polars_as_rich_table(
+        pl.DataFrame(record.rows), title=record._table_title()
+    )
+
+
+# ── ComparisonTensorRecord ───────────────────────────────────────────
+
+
+def _format_tensor_comparison_body(record: ComparisonTensorRecord) -> str:
+    body: str = record._format_location_prefix() + format_comparison(record)
+    if record.replicated_checks:
+        body += "\n" + format_replicated_checks(record.replicated_checks)
+    if record.traced_plan is not None:
+        body += "\n" + _format_aligner_plan(record.traced_plan)
+    return body
+
+
+def _format_tensor_comparison_rich_body(
+    record: ComparisonTensorRecord, verbosity: Verbosity = "normal"
+) -> RenderableType:
+    from sglang.srt.debug_utils.comparator.tensor_comparator.formatter import (
+        format_comparison_rich,
+    )
+
+    return record._format_location_prefix_rich() + format_comparison_rich(
+        record=record, verbosity=verbosity
+    )
+
+
+# ── ComparisonNonTensorRecord ────────────────────────────────────────
+
+
+def _format_non_tensor_body(record: ComparisonNonTensorRecord) -> str:
+    suffix: str = record._format_location_suffix()
+    if record.values_equal:
+        return f"NonTensor: {record.name}{suffix} = {record.baseline_value} ({record.baseline_type}) [equal]"
+    return (
+        f"NonTensor: {record.name}{suffix}\n"
+        f"  baseline = {record.baseline_value} ({record.baseline_type})\n"
+        f"  target   = {record.target_value} ({record.target_type})"
+    )
+
+
+def _format_non_tensor_rich_body(
+    record: ComparisonNonTensorRecord, verbosity: Verbosity = "normal"
+) -> RenderableType:
+    suffix: str = record._format_location_suffix()
+    name: str = escape(record.name)
+    baseline_val: str = escape(record.baseline_value)
+    target_val: str = escape(record.target_value)
+
+    if record.values_equal:
+        return (
+            f"═ {name}{suffix} = {baseline_val} "
+            f"({record.baseline_type}) [green]✓[/]"
+        )
+    return (
+        f"═ [bold red]{name}{suffix}[/]\n"
+        f"  baseline = {baseline_val} ({record.baseline_type})\n"
+        f"  target   = {target_val} ({record.target_type})"
+    )
+
+
+# ── SummaryRecord ────────────────────────────────────────────────────
+
+
+def _format_summary_body(record: SummaryRecord) -> str:
+    text: str = (
+        f"Summary: {record.passed} passed, {record.failed} failed, "
+        f"{record.skipped} skipped (total {record.total})"
+    )
+    if record.errored > 0:
+        text += f", {record.errored} errored"
+    return text
+
+
+def _format_summary_rich_body(
+    record: SummaryRecord, verbosity: Verbosity = "normal"
+) -> RenderableType:
+    text: str = (
+        f"[bold green]{record.passed} passed[/] │ "
+        f"[bold red]{record.failed} failed[/] │ "
+        f"[yellow]{record.skipped} skipped[/] │ "
+        f"{record.total} total"
+    )
+    if record.errored > 0:
+        text += f" │ [bold red]{record.errored} errored[/]"
+    return Panel(text, title="SUMMARY", border_style="bold")
+
+
+# ── LogRecord ────────────────────────────────────────────────────────
+
+
+def _format_log_body(record: LogRecord) -> str:
+    return ""
+
+
+# ── Standalone helpers ───────────────────────────────────────────────
+
+
+def _format_aligner_plan(traced_plan: TracedAlignerPlan) -> str:
+    lines: list[str] = ["Aligner Plan:"]
+
+    for side_label, traced_side in [
+        ("baseline", traced_plan.per_side.x),
+        ("target", traced_plan.per_side.y),
+    ]:
+        if not traced_side.step_plans:
+            lines.append(f"  {side_label}: (no steps)")
+            continue
+
+        step_summaries: list[str] = []
+        for traced_step in traced_side.step_plans:
+            sub_strs: list[str] = [
+                _format_sub_plan_text(traced_sub)
+                for traced_sub in traced_step.sub_plans
+            ]
+            summary: str = ", ".join(sub_strs) if sub_strs else "passthrough"
+            step_summaries.append(f"step={traced_step.step}: {summary}")
+        lines.append(f"  {side_label}: [{'; '.join(step_summaries)}]")
+
+    lines.extend(_format_cross_side_plan_text(traced_plan.plan))
+    return "\n".join(lines)
+
+
+def _format_sub_plan_text(traced_sub: TracedSubPlan) -> str:
+    from sglang.srt.debug_utils.comparator.aligner.reorderer.types import ReordererPlan
+    from sglang.srt.debug_utils.comparator.aligner.unsharder.types import UnsharderPlan
+
+    sub = traced_sub.plan
+    qualifier: str = ""
+    if isinstance(sub, UnsharderPlan):
+        qualifier = f"({sub.axis.value})"
+    elif isinstance(sub, ReordererPlan):
+        qualifier = f"({sub.params.op})"
+
+    sub_desc: str = f"{sub.type}{qualifier}"
+
+    if traced_sub.snapshot is not None:
+        snap = traced_sub.snapshot
+        in_count: int = len(snap.input_shapes)
+        out_count: int = len(snap.output_shapes)
+        in_shape: str = str(snap.input_shapes[0]) if snap.input_shapes else "?"
+        out_shape: str = str(snap.output_shapes[0]) if snap.output_shapes else "?"
+        sub_desc += f" {in_count}x{in_shape} -> {out_count}x{out_shape}"
+
+    return sub_desc
+
+
+def _format_cross_side_plan_text(plan: AlignerPlan) -> list[str]:
+    lines: list[str] = []
+
+    if plan.token_aligner_plan is not None:
+        num_tokens: int = len(plan.token_aligner_plan.locators.x.steps)
+        lines.append(f"  token_aligner: {num_tokens} tokens aligned")
+
+    if plan.axis_aligner_plan is not None:
+        parts: list[str] = []
+        if plan.axis_aligner_plan.pattern.x:
+            parts.append(f"x: {plan.axis_aligner_plan.pattern.x}")
+        if plan.axis_aligner_plan.pattern.y:
+            parts.append(f"y: {plan.axis_aligner_plan.pattern.y}")
+        lines.append(f"  axis_aligner: {', '.join(parts)}")
+
+    return lines

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

@@ -0,0 +1,324 @@
+from __future__ import annotations
+
+from abc import abstractmethod
+from typing import TYPE_CHECKING, Annotated, Any, Literal, Optional, Union
+
+from pydantic import ConfigDict, Discriminator, Field, TypeAdapter, model_validator
+from rich.console import Group, RenderableType
+from rich.markup import escape
+
+from sglang.srt.debug_utils.comparator.output_formatter import (  # noqa: F401 — re-export
+    _format_aligner_plan as _format_aligner_plan,
+)
+from sglang.srt.debug_utils.comparator.output_formatter import (
+    _format_config_body,
+    _format_config_rich_body,
+    _format_error_body,
+    _format_error_rich_body,
+    _format_log_body,
+    _format_non_tensor_body,
+    _format_non_tensor_rich_body,
+    _format_skip_body,
+    _format_skip_rich_body,
+    _format_summary_body,
+    _format_summary_rich_body,
+    _format_table_body,
+    _format_table_rich_body,
+    _format_tensor_comparison_body,
+    _format_tensor_comparison_rich_body,
+    _render_record_rich,
+    _render_record_text,
+)
+from sglang.srt.debug_utils.comparator.tensor_comparator.types import (
+    DiffInfo,
+    TensorComparisonInfo,
+    TensorInfo,
+)
+from sglang.srt.debug_utils.comparator.utils import Pair, _StrictBase
+
+if TYPE_CHECKING:
+    from sglang.srt.debug_utils.comparator.aligner.entrypoint.traced_types import (
+        TracedAlignerPlan,
+    )
+    from sglang.srt.debug_utils.comparator.report_sink import Verbosity
+
+
+class BaseLog(_StrictBase):
+    category: str
+    message: str
+
+    def to_text(self) -> str:
+        return self.message
+
+
+class ErrorLog(BaseLog):
+    kind: Literal["error"] = "error"
+
+
+class InfoLog(BaseLog):
+    kind: Literal["info"] = "info"
+
+
+AnyLog = Annotated[Union[ErrorLog, InfoLog], Discriminator("kind")]
+
+
+def _split_logs(logs: list[BaseLog]) -> tuple[list[ErrorLog], list[InfoLog]]:
+    errors: list[ErrorLog] = [log for log in logs if isinstance(log, ErrorLog)]
+    infos: list[InfoLog] = [log for log in logs if isinstance(log, InfoLog)]
+    return errors, infos
+
+
+class ReplicatedCheckResult(_StrictBase):
+    axis: str
+    group_index: int
+    compared_index: int
+    baseline_index: int
+    passed: bool
+    atol: float
+    diff: Optional[DiffInfo] = None
+
+
+class BundleFileInfo(_StrictBase):
+    """Per-file info within a bundle (one rank's raw tensor)."""
+
+    shape: list[int]
+    dtype: str
+    rank: Optional[int] = None
+    parallel_info: Optional[dict[str, str]] = None  # e.g. {"tp": "0/4", "ep": "1/2"}
+    filename: Optional[str] = None
+
+
+class BundleSideInfo(_StrictBase):
+    num_files: int
+    files: list[BundleFileInfo]
+    dims: Optional[str] = None  # e.g. "b s h(tp) d"
+
+
+class ShapeSnapshot(_StrictBase):
+    input_shapes: list[list[int]]
+    output_shapes: list[list[int]]
+
+
+class _OutputRecord(_StrictBase):
+    errors: list[ErrorLog] = Field(default_factory=list)
+    infos: list[InfoLog] = Field(default_factory=list)
+
+    @abstractmethod
+    def _format_body(self) -> str: ...
+
+    def _format_rich_body(self, verbosity: Verbosity = "normal") -> RenderableType:
+        return self._format_body()
+
+    def to_rich(self, verbosity: Verbosity = "normal") -> RenderableType:
+        return _render_record_rich(self, verbosity=verbosity)
+
+    def to_text(self) -> str:
+        return _render_record_text(self)
+
+
+class RecordLocation(_StrictBase):
+    step: Optional[int] = None
+
+
+class _BaseComparisonRecord(_OutputRecord):
+    location: RecordLocation = Field(default_factory=RecordLocation)
+
+    def to_rich(self, verbosity: Verbosity = "normal") -> RenderableType:
+        result = _render_record_rich(self, verbosity=verbosity)
+        if isinstance(result, str):
+            return result + "\n"
+        return Group(result, "")
+
+    def _format_location_prefix(self) -> str:
+        if self.location.step is not None:
+            return f"[step={self.location.step}] "
+        return ""
+
+    def _format_location_prefix_rich(self) -> str:
+        if self.location.step is not None:
+            return escape(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]
+
+    def _format_body(self) -> str:
+        return _format_config_body(self)
+
+    def _format_rich_body(self, verbosity: Verbosity = "normal") -> RenderableType:
+        return _format_config_rich_body(self, verbosity=verbosity)
+
+
+class ComparisonSkipRecord(_BaseComparisonRecord):
+    type: Literal["comparison_skip"] = "comparison_skip"
+    name: str
+    reason: str
+    available_side: Optional[Literal["baseline", "target"]] = None
+    available_tensor_info: Optional[TensorInfo] = None
+    available_bundle_info: Optional[BundleSideInfo] = None
+
+    @property
+    def category(self) -> str:
+        if self.errors:
+            return "failed"
+        return "skipped"
+
+    def _format_body(self) -> str:
+        return _format_skip_body(self)
+
+    def _format_rich_body(self, verbosity: Verbosity = "normal") -> RenderableType:
+        return _format_skip_rich_body(self, verbosity=verbosity)
+
+
+class ComparisonErrorRecord(_BaseComparisonRecord):
+    type: Literal["comparison_error"] = "comparison_error"
+    name: str
+    exception_type: str
+    exception_message: str
+    traceback_str: str
+
+    @property
+    def category(self) -> str:
+        return "errored"
+
+    def _format_body(self) -> str:
+        return _format_error_body(self)
+
+    def _format_rich_body(self, verbosity: Verbosity = "normal") -> RenderableType:
+        return _format_error_rich_body(self, verbosity=verbosity)
+
+
+class _TableRecord(_OutputRecord):
+    label: str
+    rows: list[dict[str, Any]]
+
+    @abstractmethod
+    def _table_title(self) -> str: ...
+
+    def _format_body(self) -> str:
+        return _format_table_body(self)
+
+    def _format_rich_body(self, verbosity: Verbosity = "normal") -> RenderableType:
+        return _format_table_rich_body(self, verbosity=verbosity)
+
+
+class RankInfoRecord(_TableRecord):
+    type: Literal["rank_info"] = "rank_info"
+
+    def _table_title(self) -> str:
+        return f"{self.label} ranks"
+
+
+class InputIdsRecord(_TableRecord):
+    type: Literal["input_ids"] = "input_ids"
+
+    def _table_title(self) -> str:
+        return f"{self.label} input_ids & positions"
+
+
+class ComparisonTensorRecord(TensorComparisonInfo, _BaseComparisonRecord):
+    model_config = ConfigDict(extra="forbid", defer_build=True)
+
+    type: Literal["comparison_tensor"] = "comparison_tensor"
+    traced_plan: Optional[TracedAlignerPlan] = None
+    replicated_checks: list[ReplicatedCheckResult] = Field(default_factory=list)
+    raw_bundle_info: Optional[Pair[BundleSideInfo]] = None
+
+    @property
+    def category(self) -> str:
+        if self.errors:
+            return "failed"
+        if any(not check.passed for check in self.replicated_checks):
+            return "failed"
+        return "passed" if self.diff is not None and self.diff.passed else "failed"
+
+    def _format_body(self) -> str:
+        return _format_tensor_comparison_body(self)
+
+    def _format_rich_body(self, verbosity: Verbosity = "normal") -> RenderableType:
+        return _format_tensor_comparison_rich_body(self, verbosity=verbosity)
+
+
+class ComparisonNonTensorRecord(_BaseComparisonRecord):
+    type: Literal["comparison_non_tensor"] = "comparison_non_tensor"
+    name: str
+    baseline_value: str
+    target_value: str
+    baseline_type: str
+    target_type: str
+    values_equal: bool
+
+    @property
+    def category(self) -> str:
+        if self.errors:
+            return "failed"
+        return "passed" if self.values_equal else "failed"
+
+    def _format_body(self) -> str:
+        return _format_non_tensor_body(self)
+
+    def _format_rich_body(self, verbosity: Verbosity = "normal") -> RenderableType:
+        return _format_non_tensor_rich_body(self, verbosity=verbosity)
+
+
+class SummaryRecord(_OutputRecord):
+    type: Literal["summary"] = "summary"
+    total: int
+    passed: int
+    failed: int
+    skipped: int
+    errored: int = 0
+
+    @model_validator(mode="after")
+    def _validate_totals(self) -> SummaryRecord:
+        expected: int = self.passed + self.failed + self.skipped + self.errored
+        if self.total != expected:
+            raise ValueError(
+                f"total={self.total} != passed({self.passed}) + failed({self.failed}) "
+                f"+ skipped({self.skipped}) + errored({self.errored}) = {expected}"
+            )
+        return self
+
+    def _format_body(self) -> str:
+        return _format_summary_body(self)
+
+    def _format_rich_body(self, verbosity: Verbosity = "normal") -> RenderableType:
+        return _format_summary_rich_body(self, verbosity=verbosity)
+
+
+class LogRecord(_OutputRecord):
+    type: Literal["log"] = "log"
+
+    def _format_body(self) -> str:
+        return _format_log_body(self)
+
+
+AnyRecord = Annotated[
+    Union[
+        ConfigRecord,
+        RankInfoRecord,
+        InputIdsRecord,
+        ComparisonSkipRecord,
+        ComparisonErrorRecord,
+        ComparisonTensorRecord,
+        ComparisonNonTensorRecord,
+        SummaryRecord,
+        LogRecord,
+    ],
+    Discriminator("type"),
+]
+
+
+def _get_any_record_adapter() -> TypeAdapter:
+    return TypeAdapter(AnyRecord)
+
+
+def parse_record_json(json_str: str | bytes) -> AnyRecord:
+    return _get_any_record_adapter().validate_json(json_str)

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

@@ -0,0 +1,83 @@
+"""Per-token relative difference heatmap generator.
+
+Produces a single PNG with rows = tensor names, columns = token positions,
+color = log10(rel_diff).
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Optional
+
+from sglang.srt.debug_utils.comparator.output_types import ComparisonTensorRecord
+
+
+def generate_per_token_heatmap(
+    *,
+    records: list[ComparisonTensorRecord],
+    output_path: Path,
+) -> Optional[Path]:
+    """Generate a per-token relative difference heatmap PNG.
+
+    Returns the output path if a file was written, or None if no data was available.
+    """
+    rows_data: list[tuple[str, list[float]]] = _collect_per_token_data(records=records)
+    if not rows_data:
+        return None
+
+    _render_heatmap(rows_data=rows_data, output_path=output_path)
+    return output_path
+
+
+def _collect_per_token_data(
+    *,
+    records: list[ComparisonTensorRecord],
+) -> list[tuple[str, list[float]]]:
+    rows: list[tuple[str, list[float]]] = []
+    for record in records:
+        if record.diff is None or record.diff.per_token_rel_diff is None:
+            continue
+        rows.append((record.name, record.diff.per_token_rel_diff))
+    return rows
+
+
+def _render_heatmap(
+    *,
+    rows_data: list[tuple[str, list[float]]],
+    output_path: Path,
+) -> None:
+    import matplotlib
+    import numpy as np
+
+    matplotlib.use("Agg")
+    import matplotlib.pyplot as plt
+
+    max_len: int = max(len(vals) for _, vals in rows_data)
+    labels: list[str] = [label for label, _ in rows_data]
+
+    matrix: np.ndarray = np.full((len(rows_data), max_len), np.nan, dtype=np.float64)
+    for i, (_, vals) in enumerate(rows_data):
+        matrix[i, : len(vals)] = vals
+
+    fig_width: float = max(12.0, max_len * 0.15)
+    fig_height: float = max(6.0, len(rows_data) * 0.3)
+    fig, ax = plt.subplots(figsize=(fig_width, fig_height))
+
+    im = ax.imshow(
+        np.log10(matrix + 1e-10), aspect="auto", cmap="hot", interpolation="nearest"
+    )
+
+    ax.set_xlabel("Token Position")
+    ax.set_ylabel("Tensor")
+    ax.set_yticks(range(len(labels)))
+    ax.set_yticklabels(labels, fontsize=8)
+
+    colorbar = fig.colorbar(im, ax=ax)
+    colorbar.set_label("log10(rel_diff)")
+
+    ax.set_title("Per-Token Relative Difference Heatmap")
+    fig.tight_layout()
+
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    fig.savefig(str(output_path), dpi=150)
+    plt.close(fig)

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 :]

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

@@ -0,0 +1,91 @@
+from __future__ import annotations
+
+import os
+import sys
+from pathlib import Path
+from typing import IO, Literal, Optional
+
+from rich.console import Console
+
+from sglang.srt.debug_utils.comparator.output_types import _OutputRecord
+
+Verbosity = Literal["minimal", "normal", "verbose"]
+
+
+class ReportSink:
+    """Unified entry point for all record output."""
+
+    def __init__(self) -> None:
+        self._output_format: str = "text"
+        self._verbosity: Verbosity = "normal"
+        self._report_file: Optional[IO[str]] = None
+        self._report_path: Optional[Path] = None
+        self._console: Optional[Console] = None
+
+    @property
+    def verbosity(self) -> Verbosity:
+        return self._verbosity
+
+    def configure(
+        self,
+        *,
+        output_format: str = "text",
+        report_path: Optional[Path] = None,
+        verbosity: Verbosity = "normal",
+    ) -> None:
+        self._output_format = output_format
+        self._verbosity = verbosity
+
+        if report_path is not None:
+            try:
+                report_path.parent.mkdir(parents=True, exist_ok=True)
+                self._report_file = open(report_path, "w", encoding="utf-8")
+                self._report_path = report_path
+            except OSError as exc:
+                print(
+                    f"Warning: cannot open report file {report_path}: {exc}",
+                    file=sys.stderr,
+                )
+
+    def add(self, record: _OutputRecord) -> None:
+        self._print_to_stdout(record)
+
+        if self._report_file is not None:
+            self._report_file.write(record.model_dump_json())
+            self._report_file.write("\n")
+            self._report_file.flush()
+
+    def close(self) -> None:
+        if self._report_file is not None:
+            self._report_file.close()
+            self._report_file = None
+
+    @property
+    def report_path(self) -> Optional[Path]:
+        return self._report_path
+
+    def _reset(self) -> None:
+        self.close()
+        self._output_format = "text"
+        self._verbosity = "normal"
+        self._report_path = None
+        self._console = None
+
+    def _get_console(self) -> Console:
+        if self._console is None:
+            try:
+                width = os.get_terminal_size().columns
+            except OSError:
+                width = 200
+            self._console = Console(force_terminal=True, width=width)
+        return self._console
+
+    def _print_to_stdout(self, record: _OutputRecord) -> None:
+        if self._output_format == "json":
+            print(record.model_dump_json())
+        else:
+            console: Console = self._get_console()
+            console.print(record.to_rich(verbosity=self._verbosity))
+
+
+report_sink = ReportSink()

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

@@ -0,0 +1,3 @@
+from sglang.srt.debug_utils.comparator.tensor_comparator.comparator import (
+    compare_tensor_pair,
+)

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

@@ -0,0 +1,250 @@
+from dataclasses import dataclass
+from typing import Optional
+
+import torch
+
+from sglang.srt.debug_utils.comparator.tensor_comparator.types import (
+    DEFAULT_PERCENTILES,
+    DiffInfo,
+    TensorComparisonInfo,
+    TensorInfo,
+    TensorStats,
+)
+from sglang.srt.debug_utils.comparator.threshold_dsl import (
+    DiffThresholdRule,
+    evaluate_predicate,
+    parse_predicate,
+    resolve_predicate,
+)
+from sglang.srt.debug_utils.comparator.utils import (
+    Pair,
+    argmax_coord,
+    calc_per_token_rel_diff,
+    calc_rel_diff,
+    compute_smaller_dtype,
+    try_unify_shape,
+)
+from sglang.srt.debug_utils.dumper import get_truncated_value
+
+QUANTILE_NUMEL_THRESHOLD = 10_000_000
+SAMPLE_DIFF_THRESHOLD = 1e-3
+DEFAULT_PREDICATE: str = "rel <= 0.001"
+
+
+@dataclass
+class FailureDisplayBudget:
+    max_detail: int = 50
+    num_emitted: int = 0
+
+    def take(self) -> bool:
+        if self.max_detail < 0:
+            return True
+        if self.num_emitted >= self.max_detail:
+            return False
+        self.num_emitted += 1
+        return True
+
+
+def compute_tensor_info(
+    tensor: torch.Tensor,
+    *,
+    include_sample: bool = False,
+    include_percentiles: bool = True,
+) -> TensorInfo:
+    """Compute TensorInfo (shape, dtype, stats, optional sample) for a single tensor."""
+    stats: TensorStats = _compute_tensor_stats(
+        tensor.float(), include_percentiles=include_percentiles
+    )
+    sample: Optional[str] = (
+        str(get_truncated_value(tensor.float())) if include_sample else None
+    )
+    return TensorInfo(
+        shape=list(tensor.shape),
+        dtype=str(tensor.dtype),
+        stats=stats,
+        sample=sample,
+    )
+
+
+def compare_tensor_pair(
+    x_baseline: torch.Tensor,
+    x_target: torch.Tensor,
+    name: str = "",
+    diff_threshold_rules: Optional[list[DiffThresholdRule]] = None,
+    seq_dim: Optional[int] = None,
+    failure_display_budget: Optional[FailureDisplayBudget] = None,
+) -> TensorComparisonInfo:
+    predicate = resolve_predicate(
+        name, diff_threshold_rules, default_predicate=DEFAULT_PREDICATE
+    )
+
+    x_baseline_original = x_baseline
+    x_baseline = try_unify_shape(x_baseline, target_shape=x_target.shape)
+    unified_shape = list(x_baseline.shape)
+
+    baseline_original_dtype = x_baseline.dtype
+    target_original_dtype = x_target.dtype
+
+    x_baseline_f = x_baseline.float()
+    x_target_f = x_target.float()
+
+    shape_mismatch = x_baseline_f.shape != x_target_f.shape
+
+    diff: Optional[DiffInfo] = None
+    diff_downcast: Optional[DiffInfo] = None
+    downcast_dtype: Optional[torch.dtype] = None
+
+    if not shape_mismatch:
+        diff = compute_diff(
+            x_baseline=x_baseline_f,
+            x_target=x_target_f,
+            predicate=predicate,
+            seq_dim=seq_dim,
+            include_percentiles=False,
+        )
+
+    is_failure = shape_mismatch or (diff is not None and not diff.passed)
+    needs_detail = is_failure and (
+        failure_display_budget is None or failure_display_budget.take()
+    )
+
+    baseline_info: TensorInfo = compute_tensor_info(
+        x_baseline_original, include_percentiles=needs_detail
+    )
+    target_info: TensorInfo = compute_tensor_info(
+        x_target, include_percentiles=needs_detail
+    )
+
+    if not shape_mismatch and needs_detail:
+        diff = compute_diff(
+            x_baseline=x_baseline_f,
+            x_target=x_target_f,
+            predicate=predicate,
+            seq_dim=seq_dim,
+            include_percentiles=True,
+        )
+
+    if diff is not None:
+        needs_sample = diff.max_abs_diff > SAMPLE_DIFF_THRESHOLD
+        if needs_sample:
+            baseline_info.sample = str(get_truncated_value(x_baseline_f))
+            target_info.sample = str(get_truncated_value(x_target_f))
+
+        if baseline_original_dtype != target_original_dtype:
+            downcast_dtype = compute_smaller_dtype(
+                Pair(x=baseline_original_dtype, y=target_original_dtype)
+            )
+            if downcast_dtype is not None:
+                diff_downcast = compute_diff(
+                    x_baseline=x_baseline_f.to(downcast_dtype),
+                    x_target=x_target_f.to(downcast_dtype),
+                    predicate=predicate,
+                    include_percentiles=needs_detail,
+                )
+
+    return TensorComparisonInfo(
+        name=name,
+        baseline=baseline_info,
+        target=target_info,
+        unified_shape=unified_shape,
+        shape_mismatch=shape_mismatch,
+        diff=diff,
+        diff_downcast=diff_downcast,
+        downcast_dtype=str(downcast_dtype) if downcast_dtype is not None else None,
+    )
+
+
+def _compute_tensor_stats(
+    x: torch.Tensor, *, include_percentiles: bool = True
+) -> 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 = (
+        include_percentiles and x.numel() < QUANTILE_NUMEL_THRESHOLD
+    )
+    return TensorStats(
+        mean=torch.mean(x).item(),
+        abs_mean=torch.mean(x.abs()).item(),
+        std=torch.std(x).item(),
+        min=torch.min(x).item(),
+        max=torch.max(x).item(),
+        percentiles=_compute_percentiles(x, include=include_quantiles),
+    )
+
+
+def _compute_percentiles(x: torch.Tensor, *, include: bool) -> dict[int, float]:
+    if not include:
+        return {}
+    import numpy as np
+
+    arr = x.detach().float().numpy().ravel()
+    values = np.percentile(arr, list(DEFAULT_PERCENTILES))
+    return {p: float(v) for p, v in zip(DEFAULT_PERCENTILES, values)}
+
+
+def compute_diff(
+    x_baseline: torch.Tensor,
+    x_target: torch.Tensor,
+    predicate: str = DEFAULT_PREDICATE,
+    seq_dim: Optional[int] = None,
+    include_percentiles: bool = True,
+) -> 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,
+            predicate=predicate,
+            passed=True,
+        )
+
+    raw_abs_diff = (x_target - x_baseline).abs()
+    max_diff_coord = argmax_coord(raw_abs_diff)
+
+    max_abs_diff = raw_abs_diff.max().item()
+    rel_diff = (
+        0.0 if max_abs_diff == 0.0 else calc_rel_diff(x_target, x_baseline).item()
+    )
+    mean_abs_diff = raw_abs_diff.mean().item()
+
+    include_quantiles: bool = (
+        include_percentiles and raw_abs_diff.numel() < QUANTILE_NUMEL_THRESHOLD
+    )
+
+    per_token_rel_diff: Optional[list[float]] = None
+    if seq_dim is not None and x_baseline.dim() > seq_dim:
+        per_token_rel_diff = calc_per_token_rel_diff(
+            x_baseline, x_target, seq_dim=seq_dim
+        ).tolist()
+
+    return DiffInfo(
+        rel_diff=rel_diff,
+        max_abs_diff=max_abs_diff,
+        mean_abs_diff=mean_abs_diff,
+        abs_diff_percentiles=_compute_percentiles(
+            raw_abs_diff, include=include_quantiles
+        ),
+        max_diff_coord=list(max_diff_coord),
+        baseline_at_max=x_baseline[max_diff_coord].item(),
+        target_at_max=x_target[max_diff_coord].item(),
+        predicate=predicate,
+        passed=evaluate_predicate(
+            parse_predicate(predicate),
+            rel=rel_diff,
+            max_abs=max_abs_diff,
+            mean_abs=mean_abs_diff,
+        ),
+        per_token_rel_diff=per_token_rel_diff,
+    )

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

@@ -0,0 +1,522 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Literal, Optional
+
+from rich.markup import escape
+
+from sglang.srt.debug_utils.comparator.aligner.reorderer.types import ReordererPlan
+from sglang.srt.debug_utils.comparator.aligner.unsharder.types import UnsharderPlan
+from sglang.srt.debug_utils.comparator.tensor_comparator.types import (
+    DiffInfo,
+    TensorComparisonInfo,
+    TensorInfo,
+    TensorStats,
+)
+
+if TYPE_CHECKING:
+    from sglang.srt.debug_utils.comparator.aligner.entrypoint.traced_types import (
+        TracedAlignerPlan,
+        TracedSubPlan,
+    )
+    from sglang.srt.debug_utils.comparator.aligner.entrypoint.types import AlignerPlan
+    from sglang.srt.debug_utils.comparator.output_types import (
+        BundleSideInfo,
+        ComparisonTensorRecord,
+        ReplicatedCheckResult,
+        ShapeSnapshot,
+    )
+    from sglang.srt.debug_utils.comparator.utils import Pair
+
+Verbosity = Literal["minimal", "normal", "verbose"]
+
+
+def _esc_shape(shape: Optional[list[int]]) -> str:
+    return escape(str(shape))
+
+
+def _strip_torch_prefix(dtype: str) -> str:
+    return dtype.replace("torch.", "")
+
+
+# ---------------------------------------------------------------------------
+# Number formatting
+# ---------------------------------------------------------------------------
+
+
+def _fmt_val(value: float) -> str:
+    return f"{value:.2e}"
+
+
+def _fmt_diff_colored(diff: float, *, threshold: float = 1e-2) -> str:
+    formatted: str = f"{diff:+.2e}"
+    if abs(diff) >= threshold:
+        return f"[yellow]{formatted}[/]"
+    return f"[dim]{formatted}[/]"
+
+
+# ---------------------------------------------------------------------------
+# Passed / color / marker helper
+# ---------------------------------------------------------------------------
+
+
+def _category_marker(category: str) -> tuple[bool, str, str]:
+    passed: bool = category == "passed"
+    color: str = "green" if passed else "red"
+    marker: str = f"[{color}]✅[/]" if passed else f"[{color}]❌[/]"
+    return passed, color, marker
+
+
+# ---------------------------------------------------------------------------
+# Stats formatting helpers (shared between compact / verbose)
+# ---------------------------------------------------------------------------
+
+
+_STAT_HEADER = (
+    f"      [dim]{'':10s} {'baseline':>10s}   {'target':>10s}       {'Δ':s}[/]"
+)
+
+
+def _format_stat_line(stat_name: str, val_b: float, val_t: float, diff: float) -> str:
+    return f"      [blue]{stat_name:10s}[/] {val_b:>10.4f}   {val_t:>10.4f}   {_fmt_diff_colored(diff)}"
+
+
+# ---------------------------------------------------------------------------
+# Old text-only formatters (kept for to_text() backward compatibility)
+# ---------------------------------------------------------------------------
+
+
+def format_comparison(info: TensorComparisonInfo) -> str:
+    lines: list[str] = []
+    baseline = info.baseline
+    target = info.target
+
+    dtype_marker = "" if baseline.dtype == target.dtype else "🟠"
+    lines.append(
+        f"Raw "
+        f"[shape] {baseline.shape} vs {target.shape}\t"
+        f"[{dtype_marker}dtype] {baseline.dtype} vs {target.dtype}"
+    )
+
+    if info.unified_shape != baseline.shape:
+        lines.append(
+            f"Unify shape: {baseline.shape} -> {info.unified_shape} "
+            f"(to match {target.shape})"
+        )
+
+    lines.append(
+        f"After unify "
+        f"[shape] {info.unified_shape} vs {target.shape}\t"
+        f"[dtype] {baseline.dtype} vs {target.dtype}"
+    )
+
+    lines.extend(_format_stats_comparison(baseline=baseline.stats, target=target.stats))
+
+    if info.shape_mismatch:
+        lines.append("⚠️ Shape mismatch")
+        return "\n".join(lines)
+
+    if info.diff is not None:
+        lines.extend(_format_diff(diff=info.diff))
+
+    if info.diff_downcast is not None and info.downcast_dtype is not None:
+        lines.extend(
+            _format_diff(
+                diff=info.diff_downcast,
+                prefix_text=f"When downcast to {info.downcast_dtype}: ",
+            )
+        )
+
+    if baseline.sample is not None:
+        lines.append(f"x_baseline(sample)={baseline.sample}")
+    if target.sample is not None:
+        lines.append(f"x_target(sample)={target.sample}")
+
+    return "\n".join(lines)
+
+
+def format_replicated_checks(checks: list[ReplicatedCheckResult]) -> str:
+    lines: list[str] = ["Replicated checks:"]
+
+    for check in checks:
+        marker: str = "✅" if check.passed else "❌"
+
+        if check.diff is not None:
+            detail: str = (
+                f"rel_diff={check.diff.rel_diff:.6e} "
+                f"max_abs_diff={check.diff.max_abs_diff:.6e} "
+                f"mean_abs_diff={check.diff.mean_abs_diff:.6e}"
+            )
+        else:
+            detail = "n/a diff"
+
+        lines.append(
+            f"  {marker} axis={check.axis} group={check.group_index} "
+            f"idx={check.compared_index} vs {check.baseline_index}: "
+            f"{detail}"
+        )
+
+    return "\n".join(lines)
+
+
+def _format_stats_comparison(baseline: TensorStats, target: TensorStats) -> list[str]:
+    lines: list[str] = []
+
+    for stat_name in TensorStats.model_fields:
+        if stat_name == "percentiles":
+            continue
+        value_baseline: float = getattr(baseline, stat_name)
+        value_target: float = getattr(target, stat_name)
+        lines.append(
+            f"[{stat_name}] {value_baseline:.4f} vs {value_target:.4f} "
+            f"(diff: {value_target - value_baseline:.4f})"
+        )
+
+    for p in sorted(set(baseline.percentiles) & set(target.percentiles)):
+        value_baseline = baseline.percentiles[p]
+        value_target = target.percentiles[p]
+        lines.append(
+            f"[p{p}] {value_baseline:.4f} vs {value_target:.4f} "
+            f"(diff: {value_target - value_baseline:.4f})"
+        )
+
+    return lines
+
+
+def _format_diff(diff: DiffInfo, prefix_text: str = "") -> list[str]:
+    marker: str = "✅" if diff.passed else "❌"
+    lines: list[str] = [
+        prefix_text
+        + f"{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}",
+    ]
+
+    if diff.abs_diff_percentiles:
+        quantile_parts: list[str] = [
+            f"p{p}={value:.4f}"
+            for p, value in sorted(diff.abs_diff_percentiles.items())
+        ]
+        lines.append("[abs_diff] " + " ".join(quantile_parts))
+
+    return lines
+
+
+# ---------------------------------------------------------------------------
+# New Rich markup formatters
+# ---------------------------------------------------------------------------
+
+
+def format_comparison_rich(
+    record: ComparisonTensorRecord,
+    verbosity: Verbosity = "normal",
+) -> str:
+    if verbosity == "minimal":
+        return _format_comparison_minimal(record)
+
+    return _format_comparison_normal_or_verbose(
+        record=record,
+        verbose=(verbosity == "verbose"),
+    )
+
+
+def _format_comparison_minimal(record: ComparisonTensorRecord) -> str:
+    passed, color, marker = _category_marker(record.category)
+
+    name_part: str = f"[bold {color}]{escape(record.name):30s}[/]"
+    if record.diff is not None:
+        return f"{marker} {name_part} rel_diff={_fmt_val(record.diff.rel_diff)}"
+    elif record.shape_mismatch:
+        return f"{marker} {name_part} [yellow]shape mismatch[/]"
+    else:
+        return f"{marker} {name_part}"
+
+
+def _format_comparison_normal_or_verbose(
+    *,
+    record: ComparisonTensorRecord,
+    verbose: bool,
+) -> str:
+    passed, color, marker = _category_marker(record.category)
+
+    baseline: TensorInfo = record.baseline
+    target: TensorInfo = record.target
+    aligned_shape: str = _esc_shape(record.unified_shape)
+    dtype_str: str = _strip_torch_prefix(baseline.dtype)
+
+    lines: list[str] = []
+
+    # L0: Header
+    lines.append(
+        f"{marker} [bold {color}]{escape(record.name)}[/] "
+        f"[dim cyan]── {dtype_str}  {aligned_shape}[/]"
+    )
+
+    # L1: Key metrics
+    if record.diff is not None:
+        diff: DiffInfo = record.diff
+        rel_style: str = f"bold {color}" if not passed else color
+        lines.append(
+            f"   [{rel_style}]rel_diff={_fmt_val(diff.rel_diff)}[/]"
+            f"  max_abs={_fmt_val(diff.max_abs_diff)}"
+            f"  mean_abs={_fmt_val(diff.mean_abs_diff)}"
+        )
+
+        if not passed:
+            lines.append(
+                f"   max_abs @ {_esc_shape(diff.max_diff_coord)}: "
+                f"baseline={diff.baseline_at_max}  target={diff.target_at_max}"
+            )
+    elif record.shape_mismatch:
+        lines.append("   [yellow]⚠ Shape mismatch[/]")
+
+    # Downcast info
+    if record.diff_downcast is not None and record.downcast_dtype is not None:
+        dc: DiffInfo = record.diff_downcast
+        dc_marker: str = "[green]✅[/]" if dc.passed else "[red]❌[/]"
+        lines.append(
+            f"   {dc_marker} downcast to {record.downcast_dtype}: "
+            f"rel_diff={_fmt_val(dc.rel_diff)}"
+        )
+
+    # Bundle section
+    if record.raw_bundle_info is not None:
+        lines.append("   [dim]Bundle[/]")
+        lines.extend(
+            _format_bundle_section(bundle_info=record.raw_bundle_info, verbose=verbose)
+        )
+
+    # Plan section
+    if record.traced_plan is not None:
+        lines.append("   [dim]Plan[/]")
+        lines.extend(
+            _format_plan_section_rich(
+                traced_plan=record.traced_plan,
+                verbose=verbose,
+            )
+        )
+
+    # Aligned section
+    lines.append("   [dim]Aligned[/]")
+    lines.append(
+        f"      {_esc_shape(record.unified_shape)} vs {_esc_shape(target.shape)}"
+        f"   {baseline.dtype} vs {target.dtype}"
+    )
+
+    # Stats section
+    lines.append("   [dim]Stats[/]")
+    lines.extend(
+        _format_stats_rich(
+            baseline=baseline.stats, target=target.stats, verbose=verbose
+        )
+    )
+
+    show_detail: bool = verbose or not passed
+
+    # Abs diff percentiles
+    if show_detail and record.diff is not None and record.diff.abs_diff_percentiles:
+        lines.append("   [dim]Abs Diff Percentiles[/]")
+        lines.append("      " + _format_abs_diff_percentiles_rich(record.diff))
+
+    # Samples
+    if show_detail and baseline.sample is not None:
+        lines.append("   [dim]Samples[/]")
+        lines.append(f"      baseline  {escape(baseline.sample)}")
+        if target.sample is not None:
+            lines.append(f"      target    {escape(target.sample)}")
+
+    # Replicated checks
+    if show_detail and record.replicated_checks:
+        lines.append("   [dim]Replicated Checks[/]")
+        for check in record.replicated_checks:
+            chk_marker: str = "[green]✅[/]" if check.passed else "[red]❌[/]"
+            if check.diff is not None:
+                lines.append(
+                    f"      {chk_marker} axis={check.axis}  group={check.group_index}"
+                    f"  idx={check.compared_index} vs {check.baseline_index}"
+                    f"  rel_diff={_fmt_val(check.diff.rel_diff)}"
+                    f"  max_abs={_fmt_val(check.diff.max_abs_diff)}"
+                )
+            else:
+                lines.append(
+                    f"      {chk_marker} axis={check.axis}  group={check.group_index}"
+                    f"  idx={check.compared_index} vs {check.baseline_index}: n/a"
+                )
+
+    return "\n".join(lines)
+
+
+def _format_bundle_section(
+    bundle_info: Pair[BundleSideInfo], *, verbose: bool = False
+) -> list[str]:
+    lines: list[str] = []
+
+    for label, side in [("baseline", bundle_info.x), ("target", bundle_info.y)]:
+        if not side.files:
+            lines.append(f"      {label:8s}  [dim](no files)[/]")
+            continue
+
+        dtype_desc: str = _strip_torch_prefix(side.files[0].dtype)
+
+        if verbose:
+            dims_part: str = f"  dims: {side.dims}" if side.dims else ""
+            lines.append(
+                f"      {label:8s}  [cyan]{side.num_files} files[/]"
+                f" {dtype_desc}{dims_part}"
+            )
+
+            for idx, f in enumerate(side.files):
+                rank_part: str = f"rank={f.rank}" if f.rank is not None else ""
+                par_part: str = ""
+                if f.parallel_info:
+                    par_part = " " + " ".join(
+                        f"{k}={v}" for k, v in f.parallel_info.items()
+                    )
+                file_part: str = f"  [dim]{escape(f.filename)}[/]" if f.filename else ""
+                lines.append(
+                    f"         [{idx}] {_esc_shape(f.shape)}  {rank_part}{par_part}{file_part}"
+                )
+        else:
+            shapes: list[list[int]] = [f.shape for f in side.files]
+            unique_shapes: set[str] = {str(s) for s in shapes}
+            shape_desc: str
+            if len(unique_shapes) == 1:
+                shape_desc = _esc_shape(shapes[0])
+            else:
+                shape_desc = "mixed shapes"
+
+            dims_part = f"  [dim]dims: {side.dims}[/]" if side.dims else ""
+            lines.append(
+                f"      {label:8s}  [cyan]{side.num_files} files[/]"
+                f" × {shape_desc} {dtype_desc}{dims_part}"
+            )
+
+    return lines
+
+
+def _format_plan_section_rich(
+    *,
+    traced_plan: TracedAlignerPlan,
+    verbose: bool = False,
+) -> list[str]:
+    lines: list[str] = []
+
+    for side_label, traced_side in [
+        ("baseline", traced_plan.per_side.x),
+        ("target", traced_plan.per_side.y),
+    ]:
+        if not traced_side.step_plans:
+            lines.append(f"      {side_label:8s}  [dim](passthrough)[/]")
+            continue
+
+        parts: list[str] = [
+            _format_sub_plan_rich(traced_sub)
+            for traced_step in traced_side.step_plans
+            for traced_sub in traced_step.sub_plans
+        ]
+        lines.append(f"      {side_label:8s}  " + " → ".join(parts))
+
+    lines.extend(_format_cross_side_plan_rich(traced_plan.plan))
+    return lines
+
+
+def _format_sub_plan_rich(traced_sub: TracedSubPlan) -> str:
+    sub = traced_sub.plan
+    snapshot: Optional[ShapeSnapshot] = traced_sub.snapshot
+
+    op_name: str = sub.type
+    qualifier: str = ""
+    if isinstance(sub, UnsharderPlan):
+        qualifier = f"({sub.axis.value})"
+    elif isinstance(sub, ReordererPlan):
+        qualifier = f"({sub.params.op})"
+
+    shape_change: str = ""
+    if snapshot:
+        in_count: int = len(snapshot.input_shapes)
+        out_count: int = len(snapshot.output_shapes)
+        in_shape: str = (
+            _esc_shape(snapshot.input_shapes[0]) if snapshot.input_shapes else "?"
+        )
+        out_shape: str = (
+            _esc_shape(snapshot.output_shapes[0]) if snapshot.output_shapes else "?"
+        )
+        shape_change = f" ({in_count}×{in_shape} → {out_count}×{out_shape})"
+
+    return f"[magenta]{op_name}{qualifier}[/]{shape_change}"
+
+
+def _format_cross_side_plan_rich(plan: AlignerPlan) -> list[str]:
+    lines: list[str] = []
+
+    if plan.token_aligner_plan is not None:
+        num_tokens: int = len(plan.token_aligner_plan.locators.x.steps)
+        lines.append(f"      token_aligner  [dim]{num_tokens} tokens[/]")
+
+    if plan.axis_aligner_plan is not None:
+        parts: list[str] = []
+        if plan.axis_aligner_plan.pattern.x:
+            parts.append(f"x={plan.axis_aligner_plan.pattern.x}")
+        if plan.axis_aligner_plan.pattern.y:
+            parts.append(f"y={plan.axis_aligner_plan.pattern.y}")
+        if parts:
+            lines.append(f"      axis_aligner  [dim]{', '.join(parts)}[/]")
+        else:
+            lines.append("      axis_aligner  [dim](no-op)[/]")
+
+    return lines
+
+
+def _format_stats_rich(
+    *,
+    baseline: TensorStats,
+    target: TensorStats,
+    verbose: bool = False,
+) -> list[str]:
+    lines: list[str] = [_STAT_HEADER]
+
+    if verbose:
+        # All stat fields
+        for stat_name in TensorStats.model_fields:
+            if stat_name == "percentiles":
+                continue
+            val_b: float = getattr(baseline, stat_name)
+            val_t: float = getattr(target, stat_name)
+            lines.append(_format_stat_line(stat_name, val_b, val_t, val_t - val_b))
+
+        # Percentiles
+        for p in sorted(set(baseline.percentiles) & set(target.percentiles)):
+            val_b = baseline.percentiles[p]
+            val_t = target.percentiles[p]
+            lines.append(_format_stat_line(f"p{p}", val_b, val_t, val_t - val_b))
+    else:
+        # Compact: mean, std, range, then percentiles
+        for stat_name in ("mean", "std"):
+            val_b = getattr(baseline, stat_name)
+            val_t = getattr(target, stat_name)
+            lines.append(_format_stat_line(stat_name, val_b, val_t, val_t - val_b))
+
+        # Range line: combine min/max (escape brackets to avoid Rich markup)
+        range_baseline: str = escape(f"[{baseline.min:.4f}, {baseline.max:.4f}]")
+        range_target: str = escape(f"[{target.min:.4f}, {target.max:.4f}]")
+        lines.append(f"      [blue]{'range':10s}[/] {range_baseline}   {range_target}")
+
+        # Percentiles (compact: same as verbose)
+        for p in sorted(set(baseline.percentiles) & set(target.percentiles)):
+            val_b = baseline.percentiles[p]
+            val_t = target.percentiles[p]
+            lines.append(_format_stat_line(f"p{p}", val_b, val_t, val_t - val_b))
+
+    return lines
+
+
+def _format_abs_diff_percentiles_rich(diff: DiffInfo) -> str:
+    parts: list[str] = []
+    for p, value in sorted(diff.abs_diff_percentiles.items()):
+        formatted: str = f"p{p}={_fmt_val(value)}"
+        if p >= 99 and value > 0.1:
+            formatted = f"[yellow]{formatted}[/]"
+        parts.append(formatted)
+    return "  ".join(parts)

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

@@ -0,0 +1,45 @@
+from typing import Optional
+
+from sglang.srt.debug_utils.comparator.utils import _StrictBase
+
+DEFAULT_PERCENTILES: tuple[int, ...] = (1, 5, 50, 95, 99)
+
+
+class TensorStats(_StrictBase):
+    mean: float
+    abs_mean: float
+    std: float
+    min: float
+    max: float
+    percentiles: dict[int, float] = {}
+
+
+class TensorInfo(_StrictBase):
+    shape: list[int]
+    dtype: str
+    stats: TensorStats
+    sample: Optional[str] = None
+
+
+class DiffInfo(_StrictBase):
+    rel_diff: float
+    max_abs_diff: float
+    mean_abs_diff: float
+    abs_diff_percentiles: dict[int, float] = {}
+    max_diff_coord: list[int]
+    baseline_at_max: float
+    target_at_max: float
+    predicate: str = ""
+    passed: bool
+    per_token_rel_diff: Optional[list[float]] = None
+
+
+class TensorComparisonInfo(_StrictBase):
+    name: str
+    baseline: TensorInfo
+    target: TensorInfo
+    unified_shape: Optional[list[int]]
+    shape_mismatch: bool
+    diff: Optional[DiffInfo] = None
+    diff_downcast: Optional[DiffInfo] = None
+    downcast_dtype: Optional[str] = None

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

@@ -0,0 +1,84 @@
+import re
+from dataclasses import dataclass
+from functools import lru_cache
+from types import CodeType
+from typing import Optional
+
+ALLOWED_NAMES: tuple[str, ...] = ("rel", "max_abs", "mean_abs")
+
+_EVAL_GLOBALS: dict = {"__builtins__": {}}
+_DUMMY_ENV: dict[str, float] = {name: 1.0 for name in ALLOWED_NAMES}
+
+
+@dataclass(frozen=True)
+class DiffThresholdRule:
+    pattern: str
+    predicate: str
+
+
+def parse_diff_threshold_rules(
+    raw: Optional[list[str]], *, default_predicate: str
+) -> list[DiffThresholdRule]:
+    if not raw:
+        return [DiffThresholdRule(".*", default_predicate)]
+    if len(raw) == 1:
+        try:
+            value = float(raw[0])
+        except ValueError as e:
+            raise ValueError(
+                f"--diff-threshold with a single argument must be a float shorthand "
+                f"(e.g. 0.0085); got {raw[0]!r}. For per-regex predicates pass "
+                f"(regex predicate) pairs."
+            ) from e
+        return [DiffThresholdRule(".*", f"rel <= {value}")]
+    if len(raw) % 2 != 0:
+        raise ValueError(
+            f"--diff-threshold expects a single float shorthand or (regex predicate) "
+            f"pairs; got an odd number of arguments: {raw}"
+        )
+    rules = [DiffThresholdRule(raw[i], raw[i + 1]) for i in range(0, len(raw), 2)]
+    for rule in rules:
+        parse_predicate(rule.predicate)
+    return rules
+
+
+def resolve_predicate(
+    name: str,
+    diff_threshold_rules: Optional[list[DiffThresholdRule]],
+    *,
+    default_predicate: str,
+) -> str:
+    if not diff_threshold_rules:
+        return default_predicate
+    for rule in diff_threshold_rules:
+        if re.fullmatch(rule.pattern, name):
+            return rule.predicate
+    raise ValueError(
+        f"tensor {name!r} matched no --diff-threshold pattern "
+        f"({[rule.pattern for rule in diff_threshold_rules]}); add a catch-all '.*' rule or a matching pattern."
+    )
+
+
+@lru_cache(maxsize=None)
+def parse_predicate(expr: str) -> CodeType:
+    try:
+        code = compile(expr, "<predicate>", "eval")
+    except SyntaxError as e:
+        raise ValueError(f"invalid predicate {expr!r}: {e}") from e
+    try:
+        eval(code, _EVAL_GLOBALS, dict(_DUMMY_ENV))
+    except Exception as e:
+        raise ValueError(
+            f"invalid predicate {expr!r}: {e}; allowed names are {ALLOWED_NAMES}."
+        ) from e
+    return code
+
+
+def evaluate_predicate(
+    code: CodeType, *, rel: float, max_abs: float, mean_abs: float
+) -> bool:
+    return bool(
+        eval(
+            code, _EVAL_GLOBALS, {"rel": rel, "max_abs": max_abs, "mean_abs": mean_abs}
+        )
+    )

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

@@ -0,0 +1,165 @@
+from __future__ import annotations
+
+import functools
+import re
+from pathlib import Path
+from typing import TYPE_CHECKING, Callable, Generic, Optional, Tuple, TypeVar
+
+import torch
+from pydantic import BaseModel, ConfigDict
+
+_T = TypeVar("_T")
+_U = TypeVar("_U")
+
+
+def _check_equal_lengths(**named_lists: list) -> None:
+    lengths: dict[str, int] = {name: len(lst) for name, lst in named_lists.items()}
+    unique: set[int] = set(lengths.values())
+    if len(unique) > 1:
+        details: str = ", ".join(f"{name}={length}" for name, length in lengths.items())
+        raise ValueError(f"Length mismatch: {details}")
+
+
+def auto_descend_dir(directory: Path, label: str) -> Path:
+    """If directory has no .pt files but exactly one subdirectory does, descend into it.
+
+    Raises ValueError when the layout is ambiguous (>=2 subdirs with .pt)
+    or when no .pt data is found at all.
+    """
+    if any(directory.glob("*.pt")):
+        return directory
+
+    candidates: list[Path] = [
+        sub for sub in directory.iterdir() if sub.is_dir() and any(sub.glob("*.pt"))
+    ]
+
+    if len(candidates) >= 2:
+        names: str = ", ".join(sorted(c.name for c in candidates))
+        raise ValueError(
+            f"{label}: directory {directory} has no .pt files at top level "
+            f"and multiple subdirectories contain data ({names}). "
+            f"Please specify the exact subdirectory."
+        )
+
+    if len(candidates) == 0:
+        raise ValueError(
+            f"{label}: no .pt files found in {directory} or any of its subdirectories."
+        )
+
+    resolved: Path = candidates[0]
+
+    from sglang.srt.debug_utils.comparator.log_sink import log_sink
+    from sglang.srt.debug_utils.comparator.output_types import InfoLog
+
+    log_sink.add(
+        InfoLog(
+            category="auto_descend",
+            message=f"auto-descend {label}: {directory} -> {resolved}",
+        )
+    )
+    return resolved
+
+
+class _StrictBase(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+
+class _FrozenBase(BaseModel):
+    model_config = ConfigDict(frozen=True, extra="forbid")
+
+
+class Pair(_FrozenBase, Generic[_T]):
+    x: _T
+    y: _T
+
+    def map(self, fn: Callable[[_T], _U]) -> Pair[_U]:
+        return Pair(x=fn(self.x), y=fn(self.y))
+
+
+def argmax_coord(x: torch.Tensor) -> Tuple[int, ...]:
+    flat_idx = x.argmax()
+    return tuple(idx.item() for idx in torch.unravel_index(flat_idx, x.shape))
+
+
+def compute_smaller_dtype(
+    dtypes: Pair[torch.dtype],
+) -> Optional[torch.dtype]:
+    info_dict = {
+        (torch.float32, torch.bfloat16): torch.bfloat16,
+        # ... add more ...
+    }
+    return info_dict.get((dtypes.x, dtypes.y)) or info_dict.get((dtypes.y, dtypes.x))
+
+
+def try_unify_shape(x: torch.Tensor, target_shape: torch.Size) -> torch.Tensor:
+    x_shape = x.shape
+    num_dim_to_remove = len(x_shape) - len(target_shape)
+    if (x_shape[num_dim_to_remove:] == target_shape) and all(
+        val == 1 for val in x_shape[:num_dim_to_remove]
+    ):
+        return functools.reduce(lambda a, _: a.squeeze(0), range(num_dim_to_remove), x)
+
+    return x
+
+
+# Copied from DeepGEMM
+def calc_rel_diff(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+    x, y = x.double(), y.double()
+    denominator = (x * x + y * y).sum()
+    sim = 2 * (x * y).sum() / denominator
+    return 1 - sim
+
+
+def calc_per_token_rel_diff(
+    x: torch.Tensor, y: torch.Tensor, *, seq_dim: int
+) -> torch.Tensor:
+    """Cosine-distance-like metric per token position.
+
+    Sums over all dims except seq_dim.
+    """
+    x, y = x.double(), y.double()
+    other_dims: list[int] = [d for d in range(x.dim()) if d != seq_dim]
+
+    if other_dims:
+        denominator: torch.Tensor = (x * x + y * y).sum(dim=other_dims)
+        sim: torch.Tensor = 2 * (x * y).sum(dim=other_dims) / (denominator + 1e-10)
+    else:
+        denominator = x * x + y * y
+        sim = 2 * (x * y) / (denominator + 1e-10)
+
+    return (1 - sim).float()
+
+
+if TYPE_CHECKING:
+    from sglang.srt.debug_utils.comparator.output_types import SummaryRecord
+
+
+def compute_exit_code(
+    summary: SummaryRecord,
+    *,
+    allow_skipped_pattern: str,
+    skipped_names: list[str],
+    allow_failed_pattern: Optional[str],
+    failed_names: list[str],
+    errored_names: Optional[list[str]] = None,
+) -> int:
+    if summary.passed == 0:
+        return 1
+
+    if errored_names:
+        return 1
+
+    if not _is_all_match_pattern(pattern=allow_failed_pattern, strings=failed_names):
+        return 1
+
+    if not _is_all_match_pattern(pattern=allow_skipped_pattern, strings=skipped_names):
+        return 1
+
+    return 0
+
+
+def _is_all_match_pattern(*, pattern: Optional[str], strings: list[str]) -> bool:
+    if pattern is None:
+        return len(strings) == 0
+    compiled: re.Pattern[str] = re.compile(pattern)
+    return all(compiled.fullmatch(s) for s in strings)

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

@@ -0,0 +1,3 @@
+from sglang.srt.debug_utils.comparator.visualizer.figure import (  # noqa: F401
+    generate_comparison_figure,
+)

python/sglang/srt/debug_utils/comparator/visualizer/figure.py

@@ -0,0 +1,116 @@
+"""Main orchestration logic for comparison figure generation."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Callable, Optional
+
+import numpy as np
+import torch
+
+from sglang.srt.debug_utils.comparator.visualizer.preprocessing import (
+    _preprocess_tensor,
+)
+
+
+@dataclass(frozen=True)
+class _PanelContext:
+    baseline_2d: torch.Tensor
+    target_2d: torch.Tensor
+    diff: Optional[torch.Tensor]  # None when shapes differ
+    name: str
+
+
+@dataclass(frozen=True)
+class _Panel:
+    label: str
+    requires_diff: bool
+    draw: Callable[[np.ndarray, int, _PanelContext], Optional[str]]
+
+
+def _build_panels() -> list[_Panel]:
+    from sglang.srt.debug_utils.comparator.visualizer.panels import (
+        _draw_baseline_heatmap,
+        _draw_diff_heatmap,
+        _draw_diff_histogram,
+        _draw_hist2d,
+        _draw_sampled,
+        _draw_target_heatmap,
+    )
+
+    return [
+        _Panel(
+            label="Baseline Heatmap", requires_diff=False, draw=_draw_baseline_heatmap
+        ),
+        _Panel(label="Target Heatmap", requires_diff=False, draw=_draw_target_heatmap),
+        _Panel(label="Abs Diff Heatmap", requires_diff=True, draw=_draw_diff_heatmap),
+        _Panel(label="Abs Diff Hist", requires_diff=True, draw=_draw_diff_histogram),
+        _Panel(label="Hist2D", requires_diff=True, draw=_draw_hist2d),
+        _Panel(label="Sampled", requires_diff=True, draw=_draw_sampled),
+    ]
+
+
+def generate_comparison_figure(
+    *,
+    baseline: torch.Tensor,
+    target: torch.Tensor,
+    name: str,
+    output_path: Path,
+) -> None:
+    """Generate a multi-panel comparison PNG for a baseline/target tensor pair.
+
+    Panels (6 rows x 2 cols, left=normal, right=log10):
+      Row 0: Baseline heatmap
+      Row 1: Target heatmap
+      Row 2: Abs Diff heatmap
+      Row 3: Abs Diff histogram
+      Row 4: Hist2D scatter (baseline vs target density)
+      Row 5: Sampled scatter (10k sampled mini-heatmap)
+    """
+    import matplotlib.pyplot as plt
+
+    baseline_f: torch.Tensor = baseline.detach().cpu().float()
+    target_f: torch.Tensor = target.detach().cpu().float()
+
+    can_diff: bool = baseline_f.shape == target_f.shape
+
+    baseline_2d: torch.Tensor = _preprocess_tensor(baseline_f)
+    target_2d: torch.Tensor = _preprocess_tensor(target_f)
+
+    diff: Optional[torch.Tensor] = (baseline_2d - target_2d).abs() if can_diff else None
+
+    ctx = _PanelContext(
+        baseline_2d=baseline_2d,
+        target_2d=target_2d,
+        diff=diff,
+        name=name,
+    )
+
+    panels: list[_Panel] = _build_panels()
+    active: list[_Panel] = [p for p in panels if not p.requires_diff or can_diff]
+
+    nrows: int = len(active)
+    ncols: int = 2
+    fig, axes = plt.subplots(nrows, ncols, figsize=(5 * ncols, 3.5 * nrows))
+    if nrows == 1:
+        axes = axes.reshape(1, -1)
+
+    stats_lines: list[str] = []
+    for i, panel in enumerate(active):
+        stats_line: Optional[str] = panel.draw(axes, i, ctx)
+        if stats_line is not None:
+            stats_lines.append(stats_line)
+
+    num_stats: int = len(stats_lines)
+    title_height: float = 0.015 * num_stats + 0.015
+    fig.suptitle(
+        "\n".join(stats_lines),
+        fontsize=9,
+        family="monospace",
+        y=1 - title_height / 2,
+    )
+    plt.tight_layout(rect=[0, 0, 1, 1 - title_height])
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    plt.savefig(str(output_path), dpi=150, bbox_inches="tight")
+    plt.close(fig)

python/sglang/srt/debug_utils/comparator/visualizer/panels.py

@@ -0,0 +1,226 @@
+"""Panel draw functions for tensor comparison visualization."""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import numpy as np
+import torch
+
+from sglang.srt.debug_utils.comparator.visualizer.figure import _PanelContext
+from sglang.srt.debug_utils.comparator.visualizer.preprocessing import (
+    _SCATTER_SAMPLE_SIZE,
+    _format_log_ticks,
+    _format_stats,
+    _maybe_downsample_numpy,
+    _safe_hist,
+    _to_log10,
+)
+
+
+def _draw_baseline_heatmap(
+    axes: np.ndarray, row_idx: int, ctx: _PanelContext
+) -> Optional[str]:
+    _draw_heatmap_pair(
+        axes, row_idx=row_idx, t=ctx.baseline_2d, title=f"{ctx.name} Baseline"
+    )
+    return _format_stats("Baseline", ctx.baseline_2d)
+
+
+def _draw_target_heatmap(
+    axes: np.ndarray, row_idx: int, ctx: _PanelContext
+) -> Optional[str]:
+    _draw_heatmap_pair(
+        axes, row_idx=row_idx, t=ctx.target_2d, title=f"{ctx.name} Target"
+    )
+    return _format_stats("Target", ctx.target_2d)
+
+
+def _draw_diff_heatmap(
+    axes: np.ndarray, row_idx: int, ctx: _PanelContext
+) -> Optional[str]:
+    assert ctx.diff is not None
+    _draw_heatmap_pair(axes, row_idx=row_idx, t=ctx.diff, title=f"{ctx.name} Abs Diff")
+    return _format_stats("Abs Diff", ctx.diff)
+
+
+def _draw_diff_histogram(
+    axes: np.ndarray, row_idx: int, ctx: _PanelContext
+) -> Optional[str]:
+    assert ctx.diff is not None
+    _draw_histogram_pair(
+        axes, row_idx=row_idx, diff=ctx.diff, label=f"{ctx.name} Abs Diff"
+    )
+    return None
+
+
+def _draw_hist2d(axes: np.ndarray, row_idx: int, ctx: _PanelContext) -> Optional[str]:
+    _draw_scatter_hist2d(
+        axes,
+        row_idx=row_idx,
+        baseline=ctx.baseline_2d,
+        target=ctx.target_2d,
+        label=ctx.name,
+    )
+    return None
+
+
+def _draw_sampled(axes: np.ndarray, row_idx: int, ctx: _PanelContext) -> Optional[str]:
+    _draw_scatter_sampled(
+        axes,
+        row_idx=row_idx,
+        baseline=ctx.baseline_2d,
+        target=ctx.target_2d,
+        label=ctx.name,
+    )
+    return None
+
+
+# ────────────────────── internal drawing helpers ──────────────────────
+
+
+def _draw_heatmap_pair(
+    axes: np.ndarray,
+    *,
+    row_idx: int,
+    t: torch.Tensor,
+    title: str,
+) -> None:
+    import matplotlib.pyplot as plt
+
+    ax_normal = axes[row_idx, 0]
+    ax_log = axes[row_idx, 1]
+
+    im = ax_normal.imshow(t.numpy(), aspect="auto", cmap="viridis")
+    ax_normal.set_title(title)
+    plt.colorbar(im, ax=ax_normal)
+
+    im_log = ax_log.imshow(_to_log10(t).numpy(), aspect="auto", cmap="viridis")
+    ax_log.set_title(f"{title} (Log10)")
+    cbar = plt.colorbar(im_log, ax=ax_log)
+    _format_log_ticks(cbar.ax, axis="y")
+
+
+def _draw_histogram_pair(
+    axes: np.ndarray,
+    *,
+    row_idx: int,
+    diff: torch.Tensor,
+    label: str,
+) -> None:
+
+    ax_normal = axes[row_idx, 0]
+    ax_log = axes[row_idx, 1]
+
+    diff_flat: np.ndarray = _maybe_downsample_numpy(diff.flatten())
+
+    _safe_hist(ax_normal, diff_flat, bins=100, edgecolor="none")
+    ax_normal.set_title(f"{label} Histogram")
+    ax_normal.set_xlabel("Abs Diff")
+    ax_normal.set_ylabel("Count")
+
+    log_flat: np.ndarray = np.log10(np.abs(diff_flat) + 1e-10)
+    _safe_hist(ax_log, log_flat, bins=100, edgecolor="none")
+    ax_log.set_title(f"{label} Histogram (Log10)")
+    ax_log.set_xlabel("Abs Diff")
+    ax_log.set_ylabel("Count")
+    _format_log_ticks(ax_log, axis="x")
+
+
+def _draw_scatter_hist2d(
+    axes: np.ndarray,
+    *,
+    row_idx: int,
+    baseline: torch.Tensor,
+    target: torch.Tensor,
+    label: str,
+) -> None:
+    import matplotlib.pyplot as plt
+
+    ax_normal = axes[row_idx, 0]
+    ax_log = axes[row_idx, 1]
+
+    b_flat: np.ndarray = _maybe_downsample_numpy(baseline.flatten())
+    t_flat: np.ndarray = _maybe_downsample_numpy(target.flatten())
+    min_len: int = min(len(b_flat), len(t_flat))
+    b_flat = b_flat[:min_len]
+    t_flat = t_flat[:min_len]
+
+    # Normal scale
+    lim: float = float(max(np.abs(b_flat).max(), np.abs(t_flat).max())) * 1.05
+    if lim == 0:
+        lim = 1.0
+    _h, _xe, _ye, im = ax_normal.hist2d(
+        b_flat,
+        t_flat,
+        bins=200,
+        range=[[-lim, lim], [-lim, lim]],
+        cmap="viridis",
+        norm="log",
+    )
+    ax_normal.plot([-lim, lim], [-lim, lim], "r--", linewidth=0.5)
+    ax_normal.set_title(f"{label} Hist2D")
+    ax_normal.set_xlabel("Baseline")
+    ax_normal.set_ylabel("Target")
+    ax_normal.set_aspect("equal")
+    plt.colorbar(im, ax=ax_normal)
+
+    # Log scale
+    b_log: np.ndarray = np.log10(np.abs(b_flat) + 1e-10)
+    t_log: np.ndarray = np.log10(np.abs(t_flat) + 1e-10)
+    vmin: float = float(min(b_log.min(), t_log.min())) - 0.5
+    vmax: float = float(max(b_log.max(), t_log.max())) + 0.5
+    _h2, _xe2, _ye2, im2 = ax_log.hist2d(
+        b_log,
+        t_log,
+        bins=200,
+        range=[[vmin, vmax], [vmin, vmax]],
+        cmap="viridis",
+        norm="log",
+    )
+    ax_log.plot([vmin, vmax], [vmin, vmax], "r--", linewidth=0.5)
+    ax_log.set_title(f"{label} Hist2D (Log10 Abs)")
+    ax_log.set_xlabel("Baseline")
+    ax_log.set_ylabel("Target")
+    ax_log.set_aspect("equal")
+    plt.colorbar(im2, ax=ax_log)
+    _format_log_ticks(ax_log, axis="both")
+
+
+def _draw_scatter_sampled(
+    axes: np.ndarray,
+    *,
+    row_idx: int,
+    baseline: torch.Tensor,
+    target: torch.Tensor,
+    label: str,
+) -> None:
+    import matplotlib.pyplot as plt
+
+    ax_baseline = axes[row_idx, 0]
+    ax_target = axes[row_idx, 1]
+
+    b_flat: np.ndarray = baseline.flatten().numpy()
+    t_flat: np.ndarray = target.flatten().numpy()
+
+    n_samples: int = min(_SCATTER_SAMPLE_SIZE, len(b_flat))
+    rng: np.random.Generator = np.random.default_rng(seed=42)
+    indices: np.ndarray = np.sort(rng.choice(len(b_flat), n_samples, replace=False))
+    b_sampled: np.ndarray = b_flat[indices]
+    t_sampled: np.ndarray = t_flat[indices]
+
+    side: int = int(np.sqrt(n_samples))
+    n_use: int = side * side
+    b_2d: np.ndarray = b_sampled[:n_use].reshape(side, side)
+    t_2d: np.ndarray = t_sampled[:n_use].reshape(side, side)
+
+    vmin: float = float(min(b_2d.min(), t_2d.min()))
+    vmax: float = float(max(b_2d.max(), t_2d.max()))
+
+    im_b = ax_baseline.imshow(b_2d, aspect="auto", cmap="viridis", vmin=vmin, vmax=vmax)
+    ax_baseline.set_title(f"{label} Baseline (10k sampled)")
+    plt.colorbar(im_b, ax=ax_baseline)
+
+    im_t = ax_target.imshow(t_2d, aspect="auto", cmap="viridis", vmin=vmin, vmax=vmax)
+    ax_target.set_title(f"{label} Target (10k sampled)")
+    plt.colorbar(im_t, ax=ax_target)

python/sglang/srt/debug_utils/comparator/visualizer/preprocessing.py

@@ -0,0 +1,101 @@
+"""Tensor preprocessing and utility functions for visualization."""
+
+from __future__ import annotations
+
+import math
+import re
+
+import numpy as np
+import torch
+
+_DOWNSAMPLE_THRESHOLD: int = 10_000_000
+_SCATTER_SAMPLE_SIZE: int = 10_000
+
+
+def _preprocess_tensor(tensor: torch.Tensor) -> torch.Tensor:
+    t: torch.Tensor = tensor.squeeze()
+
+    while t.ndim < 2:
+        t = t.unsqueeze(0)
+    if t.ndim > 2:
+        t = t.reshape(-1, t.shape[-1])
+
+    t = _reshape_to_balanced_aspect(t)
+    return t
+
+
+def _reshape_to_balanced_aspect(
+    t: torch.Tensor, max_ratio: float = 5.0
+) -> torch.Tensor:
+    assert t.ndim == 2
+
+    h, w = t.shape
+    ratio: float = h / w if w > 0 else float("inf")
+
+    if 1 / max_ratio <= ratio <= max_ratio:
+        return t
+
+    total: int = h * w
+    target_side: int = int(math.sqrt(total))
+
+    for new_h in range(target_side, 0, -1):
+        if total % new_h == 0:
+            new_w: int = total // new_h
+            new_ratio: float = new_h / new_w
+            if 1 / max_ratio <= new_ratio <= max_ratio:
+                return t.reshape(new_h, new_w)
+
+    return t.reshape(1, -1)
+
+
+# ────────────────────── utility ──────────────────────
+
+
+def _to_log10(t: torch.Tensor) -> torch.Tensor:
+    return t.abs().clamp(min=1e-10).log10()
+
+
+def _format_log_ticks(ax: object, axis: str = "both") -> None:
+    from matplotlib.ticker import FuncFormatter
+
+    formatter = FuncFormatter(
+        lambda x, _: f"1e{int(x)}" if x == int(x) else f"1e{x:.1f}"
+    )
+    if axis in ("x", "both"):
+        ax.xaxis.set_major_formatter(formatter)
+    if axis in ("y", "both"):
+        ax.yaxis.set_major_formatter(formatter)
+
+
+def _format_stats(name: str, t: torch.Tensor) -> str:
+    return (
+        f"{name}: shape={tuple(t.shape)}, "
+        f"min={t.min().item():.4g}, max={t.max().item():.4g}, "
+        f"mean={t.mean().item():.4g}, std={t.std().item():.4g}"
+    )
+
+
+def _safe_hist(
+    ax: object, data: np.ndarray, *, bins: int = 100, **kwargs: object
+) -> None:
+    data_f64: np.ndarray = data.astype(np.float64)
+    try:
+        ax.hist(data_f64, bins=bins, **kwargs)
+    except ValueError:
+        ax.hist(data_f64, bins=max(1, len(np.unique(data_f64[:1000]))), **kwargs)
+
+
+def _maybe_downsample_numpy(
+    t: torch.Tensor,
+    max_elements: int = _DOWNSAMPLE_THRESHOLD,
+) -> np.ndarray:
+    if t.numel() <= max_elements:
+        return t.numpy()
+
+    rng: np.random.Generator = np.random.default_rng(seed=0)
+    indices: np.ndarray = rng.choice(t.numel(), max_elements, replace=False)
+    return t.numpy()[indices]
+
+
+def _sanitize_filename(name: str) -> str:
+    return re.sub(r"[/\.\s]+", "_", name).strip("_")