chore: import upstream snapshot with attribution

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