Add aligner entrypoint and bundle handler in dump comparator (#19375)

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

@@ -0,0 +1,105 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Optional
+
+import torch
+
+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.unsharder.executor import (
+    execute_unsharder_plan,
+)
+from sglang.srt.debug_utils.comparator.aligner.unsharder.types import UnsharderPlan
+from sglang.srt.debug_utils.comparator.utils import Pair
+
+
+@dataclass(frozen=True)
+class AlignerResult:
+    tensors: Optional[Pair[torch.Tensor]]
+    failed_side_xy: Optional[str]  # "x" or "y"; None if success
+
+
+def execute_aligner_plan(
+    *,
+    tensors_pair: Pair[list[torch.Tensor]],
+    plan: AlignerPlan,
+) -> AlignerResult:
+    """Execute unified unshard/reorder per side, then combine."""
+
+    step_tensors_x: dict[int, torch.Tensor] = _execute_step_plans(
+        tensors=tensors_pair.x, step_plans=plan.per_step_plans.x
+    )
+    step_tensors_y: dict[int, torch.Tensor] = _execute_step_plans(
+        tensors=tensors_pair.y, step_plans=plan.per_step_plans.y
+    )
+
+    if not step_tensors_x or not step_tensors_y:
+        failed_side_xy: str = "x" if not step_tensors_x else "y"
+        return AlignerResult(tensors=None, failed_side_xy=failed_side_xy)
+
+    assert len(step_tensors_x) == 1 and len(step_tensors_y) == 1
+    combined = Pair(
+        x=list(step_tensors_x.values())[0],
+        y=list(step_tensors_y.values())[0],
+    )
+
+    return AlignerResult(tensors=combined, failed_side_xy=None)
+
+
+def _execute_step_plans(
+    tensors: list[torch.Tensor],
+    step_plans: list[AlignerPerStepPlan],
+) -> dict[int, torch.Tensor]:
+    result: dict[int, torch.Tensor] = {}
+
+    for step_plan in step_plans:
+        step_tensors: list[torch.Tensor] = [
+            tensors[i] for i in step_plan.input_object_indices
+        ]
+        tensor: Optional[torch.Tensor] = execute_sub_plans(
+            tensors=step_tensors, plans=step_plan.sub_plans
+        )
+        if tensor is not None:
+            result[step_plan.step] = tensor
+
+    return result
+
+
+def execute_sub_plans(
+    tensors: list[torch.Tensor],
+    plans: list[AlignerPerStepSubPlan],
+) -> Optional[torch.Tensor]:
+    if not tensors:
+        return None
+
+    if not plans:
+        if len(tensors) != 1:
+            return None
+        return tensors[0]
+
+    current = tensors
+    for plan in plans:
+        current = execute_sub_plan(tensors=current, plan=plan)
+
+    assert len(current) == 1
+    return current[0]
+
+
+def execute_sub_plan(
+    tensors: list[torch.Tensor],
+    plan: AlignerPerStepSubPlan,
+) -> list[torch.Tensor]:
+    if isinstance(plan, UnsharderPlan):
+        return execute_unsharder_plan(plan, tensors)
+    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,75 @@
+from __future__ import annotations
+
+from typing import Any
+
+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.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 import parse_dims
+from sglang.srt.debug_utils.comparator.utils import Pair
+
+
+def compute_aligner_plan(
+    *,
+    metas_pair: Pair[list[dict[str, Any]]],
+) -> AlignerPlan:
+    return AlignerPlan(
+        per_step_plans=metas_pair.map(
+            lambda metas: _compute_per_step_plans(metas=metas)
+        ),
+    )
+
+
+def _compute_per_step_plans(metas: list[dict[str, Any]]) -> 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]
+        plans: list[AlignerPerStepSubPlan] = compute_per_step_sub_plans(
+            metas=step_metas
+        )
+        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]],
+) -> list[AlignerPerStepSubPlan]:
+    if not metas or len(metas) == 1:
+        return []
+
+    dims_str = metas[0].get("dims")
+    if dims_str is None:
+        return []
+
+    dim_specs = parse_dims(dims_str)
+    parallel_infos = [normalize_parallel_info(meta) for meta in metas]
+
+    unsharder_plans = compute_unsharder_plan(
+        dim_specs=dim_specs, parallel_infos=parallel_infos
+    )
+    reorderer_plans = compute_reorderer_plans(
+        dim_specs=dim_specs, parallel_infos=parallel_infos
+    )
+    return [*unsharder_plans, *reorderer_plans]

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

@@ -0,0 +1,22 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Union
+
+from sglang.srt.debug_utils.comparator.aligner.reorderer.types import ReordererPlan
+from sglang.srt.debug_utils.comparator.aligner.unsharder.types import UnsharderPlan
+from sglang.srt.debug_utils.comparator.utils import Pair
+
+AlignerPerStepSubPlan = Union[UnsharderPlan, ReordererPlan]
+
+
+@dataclass(frozen=True)
+class AlignerPerStepPlan:
+    step: int
+    input_object_indices: list[int]
+    sub_plans: list[AlignerPerStepSubPlan]
+
+
+@dataclass(frozen=True)
+class AlignerPlan:
+    per_step_plans: Pair[list[AlignerPerStepPlan]]

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

@@ -0,0 +1,105 @@
+"""Compare two tensor bundles."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any, Union
+
+import torch
+
+from sglang.srt.debug_utils.comparator.aligner.entrypoint.executor import (
+    AlignerResult,
+    execute_aligner_plan,
+)
+from sglang.srt.debug_utils.comparator.aligner.entrypoint.planner import (
+    compute_aligner_plan,
+)
+from sglang.srt.debug_utils.comparator.aligner.entrypoint.types import AlignerPlan
+from sglang.srt.debug_utils.comparator.output_types import (
+    ComparisonRecord,
+    SkipRecord,
+)
+from sglang.srt.debug_utils.comparator.tensor_comparator.comparator import (
+    compare_tensor_pair,
+)
+from sglang.srt.debug_utils.comparator.utils import Pair
+from sglang.srt.debug_utils.comparator.warning_sink import warning_sink
+from sglang.srt.debug_utils.dump_loader import ValueWithMeta
+
+_FAILED_SIDE_MAP: dict[str, str] = {"x": "baseline", "y": "target"}
+
+
+def compare_bundle_pair(
+    *,
+    name: str,
+    filenames_pair: Pair[list[str]],
+    baseline_path: Path,
+    target_path: Path,
+    diff_threshold: float,
+) -> Union[ComparisonRecord, SkipRecord]:
+    with warning_sink.context() as collected_warnings:
+        result = _compare_bundle_pair_raw(
+            name=name,
+            filenames_pair=filenames_pair,
+            baseline_path=baseline_path,
+            target_path=target_path,
+            diff_threshold=diff_threshold,
+        )
+
+    return result.model_copy(update={"warnings": collected_warnings})
+
+
+def _compare_bundle_pair_raw(
+    *,
+    name: str,
+    filenames_pair: Pair[list[str]],
+    baseline_path: Path,
+    target_path: Path,
+    diff_threshold: float,
+) -> Union[ComparisonRecord, SkipRecord]:
+    # 1. Load (tensor + meta, ungrouped)
+    valid_pair: Pair[list[ValueWithMeta]] = Pair(
+        x=_load_valid_tensors(filenames=filenames_pair.x, base_path=baseline_path),
+        y=_load_valid_tensors(filenames=filenames_pair.y, base_path=target_path),
+    )
+
+    if not valid_pair.x or not valid_pair.y:
+        reason = "baseline_load_failed" if not valid_pair.x else "target_load_failed"
+        return SkipRecord(name=name, reason=reason)
+
+    # 2. Plan (meta only, no tensor)
+    metas_pair: Pair[list[dict[str, Any]]] = valid_pair.map(
+        lambda items: [it.meta for it in items]
+    )
+    plan: AlignerPlan = compute_aligner_plan(metas_pair=metas_pair)
+
+    # 3. Execute (tensor + plan only, no meta)
+    tensors_pair: Pair[list[torch.Tensor]] = valid_pair.map(
+        lambda items: [it.value for it in items]
+    )
+    aligner_result: AlignerResult = execute_aligner_plan(
+        tensors_pair=tensors_pair, plan=plan
+    )
+
+    if aligner_result.tensors is None:
+        assert aligner_result.failed_side_xy is not None
+        side_name: str = _FAILED_SIDE_MAP[aligner_result.failed_side_xy]
+        reason = f"{side_name}_load_failed"
+        return SkipRecord(name=name, reason=reason)
+
+    # 4. Compare
+    info = compare_tensor_pair(
+        x_baseline=aligner_result.tensors.x,
+        x_target=aligner_result.tensors.y,
+        name=name,
+        diff_threshold=diff_threshold,
+    )
+    return ComparisonRecord(**info.model_dump())
+
+
+def _load_valid_tensors(filenames: list[str], base_path: Path) -> list[ValueWithMeta]:
+    return [
+        x
+        for f in filenames
+        if isinstance((x := ValueWithMeta.load(base_path / f)).value, torch.Tensor)
+    ]

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

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

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

@@ -1,17 +1,26 @@
+from __future__ import annotations
+
 import argparse
 from pathlib import Path
+from typing import Iterator, Union
 
 import polars as pl
 
+from sglang.srt.debug_utils.comparator.bundle_comparator import compare_bundle_pair
+from sglang.srt.debug_utils.comparator.bundle_matcher import (
+    TensorBundleInfo,
+    match_bundles,
+)
 from sglang.srt.debug_utils.comparator.output_types import (
+    ComparisonRecord,
     ConfigRecord,
+    SkipRecord,
     SummaryRecord,
     print_record,
 )
-from sglang.srt.debug_utils.comparator.pipeline import process_tensor_group
-from sglang.srt.debug_utils.dump_loader import filter_rows, read_meta
-
-_NON_KEY_COLS = {"dump_index", "filename"}
+from sglang.srt.debug_utils.comparator.utils import Pair
+from sglang.srt.debug_utils.comparator.warning_sink import warning_sink
+from sglang.srt.debug_utils.dump_loader import read_meta
 
 
 def main() -> None:
@@ -20,6 +29,32 @@ def main() -> None:
 
 
 def run(args: argparse.Namespace) -> None:
+    print_record(
+        ConfigRecord.from_args(args),
+        output_format=args.output_format,
+    )
+
+    warning_sink.set_output_format(args.output_format)
+
+    dfs: Pair[pl.DataFrame] = _read_df(args)
+
+    bundle_info_pairs: list[Pair[TensorBundleInfo]] = match_bundles(
+        dfs=dfs,
+        skip_keys=_compute_skip_keys(args),
+    )
+
+    comparison_records = _compare_bundle_pairs(
+        bundle_info_pairs=bundle_info_pairs,
+        baseline_path=Path(args.baseline_path),
+        target_path=Path(args.target_path),
+        diff_threshold=args.diff_threshold,
+    )
+    _consume_comparison_records(
+        comparison_records=comparison_records, output_format=args.output_format
+    )
+
+
+def _read_df(args: argparse.Namespace) -> Pair[pl.DataFrame]:
     df_baseline = read_meta(args.baseline_path)
 
     df_target = read_meta(args.target_path)
@@ -30,37 +65,54 @@ def run(args: argparse.Namespace) -> None:
         df_target = df_target.filter(pl.col("filename").str.contains(args.filter))
     assert all(c in df_target.columns for c in ["rank", "step", "dump_index", "name"])
 
-    print_record(
-        ConfigRecord.from_args(args),
-        output_format=args.output_format,
-    )
+    return Pair(x=df_baseline, y=df_target)
 
-    counts: dict[str, int] = {"passed": 0, "failed": 0, "skipped": 0}
-    grouping: str = args.grouping
 
-    non_key_cols = _NON_KEY_COLS | ({"rank"} if grouping == "logical" else set())
-    key_cols = [c for c in df_target.columns if c not in non_key_cols]
-    tensor_group_keys = df_target.unique(subset=key_cols)
+def _compute_skip_keys(args: argparse.Namespace) -> set[str]:
+    skip_keys: set[str] = {"dump_index", "filename"}
+    if args.grouping == "logical":
+        skip_keys |= {"rank"}
+    return skip_keys
 
-    for tensor_group_key in tensor_group_keys.iter_rows(named=True):
-        conditions = {k: tensor_group_key[k] for k in key_cols}
-        baseline_rows = filter_rows(df_baseline, conditions=conditions)
-        target_rows = filter_rows(df_target, conditions=conditions)
 
-        record = process_tensor_group(
-            name=tensor_group_key["name"],
-            baseline_filenames=[r["filename"] for r in baseline_rows],
-            target_filenames=[r["filename"] for r in target_rows],
-            baseline_path=Path(args.baseline_path),
-            target_path=Path(args.target_path),
-            diff_threshold=args.diff_threshold,
+def _compare_bundle_pairs(
+    *,
+    bundle_info_pairs: list[Pair[TensorBundleInfo]],
+    baseline_path: Path,
+    target_path: Path,
+    diff_threshold: float,
+) -> Iterator[Union[ComparisonRecord, SkipRecord]]:
+    for bundle_info_pair in bundle_info_pairs:
+        if not bundle_info_pair.y:
+            continue
+
+        name: str = bundle_info_pair.y[0].name
+        filenames_pair: Pair[list[str]] = bundle_info_pair.map(
+            lambda infos: [info.filename for info in infos]
         )
+        yield compare_bundle_pair(
+            name=name,
+            filenames_pair=filenames_pair,
+            baseline_path=baseline_path,
+            target_path=target_path,
+            diff_threshold=diff_threshold,
+        )
+
+
+def _consume_comparison_records(
+    *,
+    comparison_records: Iterator[Union[ComparisonRecord, SkipRecord]],
+    output_format: str,
+) -> None:
+    counts: dict[str, int] = {"passed": 0, "failed": 0, "skipped": 0}
+
+    for record in comparison_records:
         counts[record.category] += 1
-        print_record(record, output_format=args.output_format)
+        print_record(record, output_format=output_format)
 
     print_record(
         SummaryRecord(total=sum(counts.values()), **counts),
-        output_format=args.output_format,
+        output_format=output_format,
     )
 
 

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

@@ -1,169 +0,0 @@
-from pathlib import Path
-from typing import Any, Optional, Union
-
-import torch
-
-from sglang.srt.debug_utils.comparator.aligner.reorderer.executor import (
-    execute_reorderer_plan,
-)
-from sglang.srt.debug_utils.comparator.aligner.reorderer.planner import (
-    compute_reorderer_plans,
-)
-from sglang.srt.debug_utils.comparator.aligner.reorderer.types import ReordererPlan
-from sglang.srt.debug_utils.comparator.aligner.unsharder.executor import (
-    execute_unsharder_plan,
-)
-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.aligner.unsharder.types import UnsharderPlan
-from sglang.srt.debug_utils.comparator.dims import parse_dims
-from sglang.srt.debug_utils.comparator.output_types import (
-    AnyWarning,
-    ComparisonRecord,
-    SkipRecord,
-)
-from sglang.srt.debug_utils.comparator.tensor_comparator.comparator import (
-    compare_tensor_pair,
-)
-from sglang.srt.debug_utils.comparator.warning_sink import warning_sink
-from sglang.srt.debug_utils.dump_loader import ValueWithMeta
-
-Plan = Union[UnsharderPlan, ReordererPlan]
-
-
-def process_tensor_group(
-    *,
-    name: str,
-    baseline_filenames: list[str],
-    target_filenames: list[str],
-    baseline_path: Path,
-    target_path: Path,
-    diff_threshold: float,
-) -> ComparisonRecord | SkipRecord:
-    with warning_sink.context() as collected_warnings:
-        return _process_tensor_group_raw(
-            name=name,
-            baseline_filenames=baseline_filenames,
-            target_filenames=target_filenames,
-            baseline_path=baseline_path,
-            target_path=target_path,
-            diff_threshold=diff_threshold,
-            collected_warnings=collected_warnings,
-        )
-
-
-def _process_tensor_group_raw(
-    *,
-    name: str,
-    baseline_filenames: list[str],
-    target_filenames: list[str],
-    baseline_path: Path,
-    target_path: Path,
-    diff_threshold: float,
-    collected_warnings: list[AnyWarning],
-) -> ComparisonRecord | SkipRecord:
-    b_tensors = _load_tensors(baseline_filenames, baseline_path)
-    t_tensors = _load_tensors(target_filenames, target_path)
-
-    b_plans, t_plans = _compute_plans(
-        baseline_metas=[item.meta for item in b_tensors],
-        target_metas=[item.meta for item in t_tensors],
-    )
-
-    b_extracted = _extract_tensors(b_tensors)
-    t_extracted = _extract_tensors(t_tensors)
-    del b_tensors, t_tensors
-
-    b_tensor = _execute_plans(b_extracted, b_plans)
-    t_tensor = _execute_plans(t_extracted, t_plans)
-
-    if b_tensor is None or t_tensor is None:
-        reason = "baseline_load_failed" if b_tensor is None else "target_load_failed"
-        return SkipRecord(name=name, reason=reason, warnings=collected_warnings)
-
-    info = compare_tensor_pair(
-        x_baseline=b_tensor,
-        x_target=t_tensor,
-        name=name,
-        diff_threshold=diff_threshold,
-    )
-
-    return ComparisonRecord(**info.model_dump(), warnings=collected_warnings)
-
-
-def _load_tensors(filenames: list[str], base_path: Path) -> list[ValueWithMeta]:
-    return [ValueWithMeta.load(base_path / f) for f in filenames]
-
-
-def _compute_plans(
-    *,
-    baseline_metas: list[dict[str, Any]],
-    target_metas: list[dict[str, Any]],
-) -> tuple[list[Plan], list[Plan]]:
-    """This function deliberately takes metadata, since plan computation must never inspect actual tensor data."""
-    return (
-        _compute_plans_for_group(baseline_metas),
-        _compute_plans_for_group(target_metas),
-    )
-
-
-def _compute_plans_for_group(metas: list[dict[str, Any]]) -> list[Plan]:
-    if not metas or len(metas) == 1:
-        return []
-
-    dims_str = metas[0].get("dims")
-    if dims_str is None:
-        return []
-
-    dim_specs = parse_dims(dims_str)
-    parallel_infos = [normalize_parallel_info(meta) for meta in metas]
-
-    unsharder_plans = compute_unsharder_plan(
-        dim_specs=dim_specs, parallel_infos=parallel_infos
-    )
-    reorderer_plans = compute_reorderer_plans(
-        dim_specs=dim_specs, parallel_infos=parallel_infos
-    )
-    return [*unsharder_plans, *reorderer_plans]
-
-
-def _extract_tensors(
-    loaded: list[ValueWithMeta],
-) -> Optional[list[torch.Tensor]]:
-    return [value for item in loaded if isinstance(value := item.value, torch.Tensor)]
-
-
-def _execute_plans(
-    tensors: list[torch.Tensor],
-    plans: list[Plan],
-) -> Optional[torch.Tensor]:
-    if not tensors:
-        return None
-
-    if not plans:
-        if len(tensors) != 1:
-            return None
-        return tensors[0]
-
-    current = tensors
-    for plan in plans:
-        current = _execute_plan(current, plan)
-
-    assert len(current) == 1
-    return current[0]
-
-
-def _execute_plan(
-    tensors: list[torch.Tensor],
-    plan: Plan,
-) -> list[torch.Tensor]:
-    if isinstance(plan, UnsharderPlan):
-        return execute_unsharder_plan(plan, tensors)
-    elif isinstance(plan, ReordererPlan):
-        return execute_reorderer_plan(plan, tensors)
-    else:
-        raise NotImplementedError(f"Unknown {plan=}")

test/registered/debug_utils/comparator/aligner/entrypoint/test_executor.py

@@ -0,0 +1,199 @@
+import sys
+from typing import Optional
+
+import pytest
+import torch
+
+from sglang.srt.debug_utils.comparator.aligner.entrypoint.executor import (
+    AlignerResult,
+    _execute_step_plans,
+    execute_aligner_plan,
+    execute_sub_plan,
+    execute_sub_plans,
+)
+from sglang.srt.debug_utils.comparator.aligner.entrypoint.types import (
+    AlignerPerStepPlan,
+    AlignerPlan,
+)
+from sglang.srt.debug_utils.comparator.aligner.unsharder.types import (
+    ConcatParams,
+    UnsharderPlan,
+)
+from sglang.srt.debug_utils.comparator.dims import ParallelAxis
+from sglang.srt.debug_utils.comparator.utils import Pair
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=15, suite="default", nightly=True)
+
+
+class TestExecuteSubPlans:
+    def test_empty_tensors_returns_none(self) -> None:
+        result: Optional[torch.Tensor] = execute_sub_plans(tensors=[], plans=[])
+        assert result is None
+
+    def test_no_plans_single_tensor_passthrough(self) -> None:
+        tensor: torch.Tensor = torch.tensor([1.0, 2.0, 3.0])
+        result: Optional[torch.Tensor] = execute_sub_plans(tensors=[tensor], plans=[])
+        assert result is not None
+        assert torch.equal(result, tensor)
+
+    def test_no_plans_multiple_tensors_returns_none(self) -> None:
+        tensors: list[torch.Tensor] = [
+            torch.tensor([1.0]),
+            torch.tensor([2.0]),
+        ]
+        result: Optional[torch.Tensor] = execute_sub_plans(tensors=tensors, plans=[])
+        assert result is None
+
+    def test_with_unsharder_plan(self) -> None:
+        t0: torch.Tensor = torch.tensor([[1.0, 2.0]])
+        t1: torch.Tensor = torch.tensor([[3.0, 4.0]])
+
+        plan = UnsharderPlan(
+            axis=ParallelAxis.TP,
+            params=ConcatParams(dim=1),
+            groups=[[0, 1]],
+        )
+
+        result: Optional[torch.Tensor] = execute_sub_plans(
+            tensors=[t0, t1], plans=[plan]
+        )
+
+        assert result is not None
+        expected: torch.Tensor = torch.tensor([[1.0, 2.0, 3.0, 4.0]])
+        assert torch.equal(result, expected)
+
+
+class TestExecuteSubPlan:
+    def test_unknown_plan_type_raises(self) -> None:
+        class _FakePlan:
+            pass
+
+        with pytest.raises(NotImplementedError, match="Unknown"):
+            execute_sub_plan(tensors=[torch.tensor([1.0])], plan=_FakePlan())  # type: ignore[arg-type]
+
+
+class TestExecuteStepPlans:
+    def test_step_with_none_result_omitted(self) -> None:
+        tensors: list[torch.Tensor] = [
+            torch.tensor([1.0]),
+            torch.tensor([2.0]),
+        ]
+
+        step_plan = AlignerPerStepPlan(
+            step=0,
+            input_object_indices=[0, 1],
+            sub_plans=[],
+        )
+
+        result: dict[int, torch.Tensor] = _execute_step_plans(
+            tensors=tensors, step_plans=[step_plan]
+        )
+
+        assert result == {}
+
+    def test_single_step_passthrough(self) -> None:
+        tensor: torch.Tensor = torch.tensor([1.0, 2.0])
+
+        step_plan = AlignerPerStepPlan(
+            step=5,
+            input_object_indices=[0],
+            sub_plans=[],
+        )
+
+        result: dict[int, torch.Tensor] = _execute_step_plans(
+            tensors=[tensor], step_plans=[step_plan]
+        )
+
+        assert 5 in result
+        assert torch.equal(result[5], tensor)
+
+
+class TestExecuteAlignerPlan:
+    def _make_step_plan(self, *, step: int, indices: list[int]) -> AlignerPerStepPlan:
+        return AlignerPerStepPlan(step=step, input_object_indices=indices, sub_plans=[])
+
+    def test_x_side_empty_returns_failed_x(self) -> None:
+        plan = AlignerPlan(
+            per_step_plans=Pair(
+                x=[self._make_step_plan(step=0, indices=[0, 1])],
+                y=[self._make_step_plan(step=0, indices=[0])],
+            ),
+        )
+
+        tensors_pair: Pair[list[torch.Tensor]] = Pair(
+            x=[torch.tensor([1.0]), torch.tensor([2.0])],
+            y=[torch.tensor([3.0])],
+        )
+
+        result: AlignerResult = execute_aligner_plan(
+            tensors_pair=tensors_pair, plan=plan
+        )
+
+        assert result.tensors is None
+        assert result.failed_side_xy == "x"
+
+    def test_y_side_empty_returns_failed_y(self) -> None:
+        plan = AlignerPlan(
+            per_step_plans=Pair(
+                x=[self._make_step_plan(step=0, indices=[0])],
+                y=[self._make_step_plan(step=0, indices=[0, 1])],
+            ),
+        )
+
+        tensors_pair: Pair[list[torch.Tensor]] = Pair(
+            x=[torch.tensor([1.0])],
+            y=[torch.tensor([2.0]), torch.tensor([3.0])],
+        )
+
+        result: AlignerResult = execute_aligner_plan(
+            tensors_pair=tensors_pair, plan=plan
+        )
+
+        assert result.tensors is None
+        assert result.failed_side_xy == "y"
+
+    def test_single_step(self) -> None:
+        plan = AlignerPlan(
+            per_step_plans=Pair(
+                x=[self._make_step_plan(step=0, indices=[0])],
+                y=[self._make_step_plan(step=0, indices=[0])],
+            ),
+        )
+
+        t_x: torch.Tensor = torch.tensor([1.0, 2.0])
+        t_y: torch.Tensor = torch.tensor([3.0, 4.0])
+        tensors_pair: Pair[list[torch.Tensor]] = Pair(x=[t_x], y=[t_y])
+
+        result: AlignerResult = execute_aligner_plan(
+            tensors_pair=tensors_pair, plan=plan
+        )
+
+        assert result.tensors is not None
+        assert result.failed_side_xy is None
+        assert torch.equal(result.tensors.x, t_x)
+        assert torch.equal(result.tensors.y, t_y)
+
+    def test_success_returns_none_failed_side(self) -> None:
+        plan = AlignerPlan(
+            per_step_plans=Pair(
+                x=[self._make_step_plan(step=0, indices=[0])],
+                y=[self._make_step_plan(step=0, indices=[0])],
+            ),
+        )
+
+        tensors_pair: Pair[list[torch.Tensor]] = Pair(
+            x=[torch.tensor([10.0])],
+            y=[torch.tensor([20.0])],
+        )
+
+        result: AlignerResult = execute_aligner_plan(
+            tensors_pair=tensors_pair, plan=plan
+        )
+
+        assert result.failed_side_xy is None
+        assert result.tensors is not None
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__]))

test/registered/debug_utils/comparator/aligner/entrypoint/test_planner.py

@@ -0,0 +1,146 @@
+import sys
+from typing import Any, Optional
+
+import pytest
+
+from sglang.srt.debug_utils.comparator.aligner.entrypoint.planner import (
+    _compute_per_step_plans,
+    compute_aligner_plan,
+    compute_per_step_sub_plans,
+)
+from sglang.srt.debug_utils.comparator.aligner.entrypoint.types import (
+    AlignerPerStepPlan,
+    AlignerPerStepSubPlan,
+    AlignerPlan,
+)
+from sglang.srt.debug_utils.comparator.aligner.reorderer.types import ReordererPlan
+from sglang.srt.debug_utils.comparator.aligner.unsharder.types import UnsharderPlan
+from sglang.srt.debug_utils.comparator.utils import Pair
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=15, suite="default", nightly=True)
+
+
+def _make_meta(
+    *,
+    step: int = 0,
+    dims: Optional[str] = None,
+    tp_rank: int = 0,
+    tp_size: int = 1,
+    cp_rank: int = 0,
+    cp_size: int = 1,
+) -> dict[str, Any]:
+    meta: dict[str, Any] = {"step": step}
+    if dims is not None:
+        meta["dims"] = dims
+    meta["sglang_parallel_info"] = {
+        "tp_rank": tp_rank,
+        "tp_size": tp_size,
+        "cp_rank": cp_rank,
+        "cp_size": cp_size,
+    }
+    return meta
+
+
+class TestComputePerStepSubPlans:
+    def test_empty_metas(self) -> None:
+        result: list[AlignerPerStepSubPlan] = compute_per_step_sub_plans(metas=[])
+        assert result == []
+
+    def test_single_meta(self) -> None:
+        result: list[AlignerPerStepSubPlan] = compute_per_step_sub_plans(
+            metas=[_make_meta(dims="b h(tp)", tp_size=2)]
+        )
+        assert result == []
+
+    def test_dims_none(self) -> None:
+        result: list[AlignerPerStepSubPlan] = compute_per_step_sub_plans(
+            metas=[
+                _make_meta(tp_rank=0, tp_size=2),
+                _make_meta(tp_rank=1, tp_size=2),
+            ]
+        )
+        assert result == []
+
+    def test_tp_sharded_returns_unsharder_plan(self) -> None:
+        result: list[AlignerPerStepSubPlan] = compute_per_step_sub_plans(
+            metas=[
+                _make_meta(dims="b h(tp)", tp_rank=0, tp_size=2),
+                _make_meta(dims="b h(tp)", tp_rank=1, tp_size=2),
+            ]
+        )
+        assert len(result) >= 1
+        unsharder_plans: list[UnsharderPlan] = [
+            p for p in result if isinstance(p, UnsharderPlan)
+        ]
+        assert len(unsharder_plans) >= 1
+
+    def test_zigzag_returns_both_plans(self) -> None:
+        result: list[AlignerPerStepSubPlan] = compute_per_step_sub_plans(
+            metas=[
+                _make_meta(dims="b s(cp,zigzag) h", cp_rank=0, cp_size=2),
+                _make_meta(dims="b s(cp,zigzag) h", cp_rank=1, cp_size=2),
+            ]
+        )
+        unsharder_plans: list[UnsharderPlan] = [
+            p for p in result if isinstance(p, UnsharderPlan)
+        ]
+        reorderer_plans: list[ReordererPlan] = [
+            p for p in result if isinstance(p, ReordererPlan)
+        ]
+        assert len(unsharder_plans) >= 1
+        assert len(reorderer_plans) >= 1
+
+
+class TestComputePerStepPlans:
+    def test_groups_by_step(self) -> None:
+        metas: list[dict[str, Any]] = [
+            _make_meta(step=0, tp_rank=0, tp_size=2),
+            _make_meta(step=0, tp_rank=1, tp_size=2),
+            _make_meta(step=1, tp_rank=0, tp_size=1),
+        ]
+        result: list[AlignerPerStepPlan] = _compute_per_step_plans(metas=metas)
+
+        assert len(result) == 2
+        assert result[0].step == 0
+        assert result[0].input_object_indices == [0, 1]
+        assert result[1].step == 1
+        assert result[1].input_object_indices == [2]
+
+    def test_sorted_by_step(self) -> None:
+        metas: list[dict[str, Any]] = [
+            _make_meta(step=2),
+            _make_meta(step=0),
+            _make_meta(step=1),
+        ]
+        result: list[AlignerPerStepPlan] = _compute_per_step_plans(metas=metas)
+
+        steps: list[int] = [p.step for p in result]
+        assert steps == [0, 1, 2]
+
+    def test_single_meta_per_step_empty_sub_plans(self) -> None:
+        metas: list[dict[str, Any]] = [
+            _make_meta(step=0),
+            _make_meta(step=1),
+        ]
+        result: list[AlignerPerStepPlan] = _compute_per_step_plans(metas=metas)
+
+        assert len(result) == 2
+        assert all(plan.sub_plans == [] for plan in result)
+
+
+class TestComputeAlignerPlan:
+    def test_wraps_both_sides(self) -> None:
+        metas_x: list[dict[str, Any]] = [_make_meta(step=0)]
+        metas_y: list[dict[str, Any]] = [_make_meta(step=0)]
+
+        plan: AlignerPlan = compute_aligner_plan(
+            metas_pair=Pair(x=metas_x, y=metas_y),
+        )
+
+        assert len(plan.per_step_plans.x) == 1
+        assert len(plan.per_step_plans.y) == 1
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__]))

test/registered/debug_utils/comparator/test_bundle_matcher.py

@@ -0,0 +1,150 @@
+import sys
+from typing import Any
+
+import polars as pl
+import pytest
+
+from sglang.srt.debug_utils.comparator.bundle_matcher import (
+    TensorBundleInfo,
+    TensorFileInfo,
+    _rows_to_tensor_infos,
+    match_bundles,
+)
+from sglang.srt.debug_utils.comparator.utils import Pair
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=15, suite="default", nightly=True)
+
+
+def _make_row(
+    *, name: str, step: int = 0, rank: int = 0, filename: str | None = None
+) -> dict[str, Any]:
+    if filename is None:
+        filename = f"name={name}___step={step}___rank={rank}.pt"
+    return {"name": name, "step": step, "rank": rank, "filename": filename}
+
+
+def _make_df(rows: list[dict[str, Any]]) -> pl.DataFrame:
+    return pl.DataFrame(rows)
+
+
+class TestMatchBundles:
+    def test_single_tensor_single_step(self) -> None:
+        target_df: pl.DataFrame = _make_df([_make_row(name="t_a")])
+        baseline_df: pl.DataFrame = _make_df([_make_row(name="t_a")])
+
+        results: list[Pair[TensorBundleInfo]] = match_bundles(
+            dfs=Pair(x=baseline_df, y=target_df),
+            skip_keys={"filename"},
+        )
+
+        assert len(results) == 1
+        assert len(results[0].x) == 1
+        assert len(results[0].y) == 1
+        assert results[0].y[0].name == "t_a"
+
+    def test_multiple_names_separate_bundles(self) -> None:
+        target_df: pl.DataFrame = _make_df([
+            _make_row(name="t_a"),
+            _make_row(name="t_b"),
+        ])
+        baseline_df: pl.DataFrame = _make_df([
+            _make_row(name="t_a"),
+            _make_row(name="t_b"),
+        ])
+
+        results: list[Pair[TensorBundleInfo]] = match_bundles(
+            dfs=Pair(x=baseline_df, y=target_df),
+            skip_keys={"filename"},
+        )
+
+        assert len(results) == 2
+        result_names: list[str] = [r.y[0].name for r in results]
+        assert "t_a" in result_names
+        assert "t_b" in result_names
+
+    def test_skip_rank_groups_across_ranks(self) -> None:
+        target_df: pl.DataFrame = _make_df([
+            _make_row(name="t_a", rank=0),
+            _make_row(name="t_a", rank=1),
+        ])
+        baseline_df: pl.DataFrame = _make_df([
+            _make_row(name="t_a", rank=0),
+            _make_row(name="t_a", rank=1),
+        ])
+
+        results: list[Pair[TensorBundleInfo]] = match_bundles(
+            dfs=Pair(x=baseline_df, y=target_df),
+            skip_keys={"filename", "rank"},
+        )
+
+        assert len(results) == 1
+        assert len(results[0].y) == 2
+
+    def test_baseline_missing_tensor(self) -> None:
+        target_df: pl.DataFrame = _make_df([
+            _make_row(name="t_a"),
+            _make_row(name="t_extra"),
+        ])
+        baseline_df: pl.DataFrame = _make_df([
+            _make_row(name="t_a"),
+        ])
+
+        results: list[Pair[TensorBundleInfo]] = match_bundles(
+            dfs=Pair(x=baseline_df, y=target_df),
+            skip_keys={"filename"},
+        )
+
+        assert len(results) == 2
+        extra_pair: Pair[TensorBundleInfo] = [
+            r for r in results if r.y[0].name == "t_extra"
+        ][0]
+        assert extra_pair.x == []
+
+    def test_empty_target_returns_empty(self) -> None:
+        target_df: pl.DataFrame = _make_df([])
+        baseline_df: pl.DataFrame = _make_df([_make_row(name="t_a")])
+
+        results: list[Pair[TensorBundleInfo]] = match_bundles(
+            dfs=Pair(x=baseline_df, y=target_df),
+            skip_keys={"filename"},
+        )
+
+        assert results == []
+
+    def test_skip_step_groups_across_steps(self) -> None:
+        target_df: pl.DataFrame = _make_df([
+            _make_row(name="t_a", step=0),
+            _make_row(name="t_a", step=1),
+        ])
+        baseline_df: pl.DataFrame = _make_df([
+            _make_row(name="t_a", step=0),
+            _make_row(name="t_a", step=1),
+        ])
+
+        results: list[Pair[TensorBundleInfo]] = match_bundles(
+            dfs=Pair(x=baseline_df, y=target_df),
+            skip_keys={"filename", "step"},
+        )
+
+        assert len(results) == 1
+        assert len(results[0].y) == 2
+
+
+class TestRowsToTensorInfos:
+    def test_filters_extra_columns(self) -> None:
+        rows: list[dict[str, Any]] = [
+            {"filename": "a.pt", "name": "t_a", "step": 0, "rank": 7}
+        ]
+        infos: list[TensorFileInfo] = _rows_to_tensor_infos(rows)
+
+        assert len(infos) == 1
+        assert infos[0] == TensorFileInfo(filename="a.pt", name="t_a", step=0)
+
+    def test_empty_rows(self) -> None:
+        infos: list[TensorFileInfo] = _rows_to_tensor_infos([])
+        assert infos == []
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__]))