Support multi axis unsharding in dump comparator (#19280)

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

@@ -82,9 +82,8 @@ def _compute_plans_for_group(metas: list[dict[str, Any]]) -> list[Plan]:
 
     dim_specs = parse_dims(dims_str)
     parallel_infos = [normalize_parallel_info(meta) for meta in metas]
-    plan = compute_unshard_plan(dim_specs=dim_specs, parallel_infos=parallel_infos)
 
-    return [plan] if plan is not None else []
+    return compute_unshard_plan(dim_specs=dim_specs, parallel_infos=parallel_infos)
 
 
 def _extract_tensors(
@@ -105,15 +104,12 @@ def _execute_plans(
             return None
         return tensors[0]
 
-    assert len(plans) <= 1, "multi-plan not supported yet"
-
+    current = tensors
     for plan in plans:
         if isinstance(plan, UnshardPlan):
-            # TODO: incorrect `tensors_by_world_rank` if multi UnshardPlan
-            tensors = execute_unshard_plan(
-                plan, tensors_by_world_rank=dict(enumerate(tensors))
-            )
+            current = execute_unshard_plan(plan, current)
         else:
             raise NotImplementedError(f"Unknown {plan=}")
 
-    return tensors
+    assert len(current) == 1
+    return current[0]

python/sglang/srt/debug_utils/comparator/unshard/executor.py

@@ -2,26 +2,25 @@ import torch
 
 from sglang.srt.debug_utils.comparator.unshard.types import (
     ConcatParams,
+    UnshardParams,
     UnshardPlan,
 )
 
 
 def execute_unshard_plan(
     plan: UnshardPlan,
-    tensors_by_world_rank: dict[int, torch.Tensor],
-) -> torch.Tensor:
-    ordered_tensors = [
-        tensors_by_world_rank[world_rank]
-        for world_rank in plan.world_ranks_by_axis_rank
-    ]
-
-    return _apply_unshard(plan, ordered_tensors)
+    tensors: list[torch.Tensor],
+) -> list[torch.Tensor]:
+    result: list[torch.Tensor] = []
+    for group in plan.groups:
+        group_tensors = [tensors[i] for i in group]
+        result.append(_apply_unshard(plan.params, group_tensors))
+    return result
 
 
 def _apply_unshard(
-    plan: UnshardPlan, ordered_tensors: list[torch.Tensor]
+    params: UnshardParams, ordered_tensors: list[torch.Tensor]
 ) -> torch.Tensor:
-    params = plan.params
     if isinstance(params, ConcatParams):
         return _unshard_concat(ordered_tensors, dim=params.dim)
     # Phase 2: ReduceSumParams, CpZigzagParams

python/sglang/srt/debug_utils/comparator/unshard/planner.py

@@ -1,4 +1,5 @@
-from typing import Optional
+from collections import defaultdict
+from typing import NamedTuple
 
 from sglang.srt.debug_utils.comparator.dims import DimSpec, Ordering, ParallelAxis
 from sglang.srt.debug_utils.comparator.unshard.types import (
@@ -8,64 +9,117 @@ from sglang.srt.debug_utils.comparator.unshard.types import (
     UnshardPlan,
 )
 
+# _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_unshard_plan(
     dim_specs: list[DimSpec],
     parallel_infos: list[dict[ParallelAxis, AxisInfo]],
-) -> Optional[UnshardPlan]:
+) -> list[UnshardPlan]:
     if not parallel_infos:
         raise ValueError("parallel_infos must not be empty")
 
-    sharded_axes: dict[ParallelAxis, tuple[int, DimSpec]] = {}
-    for dim_idx, spec in enumerate(dim_specs):
-        if spec.parallel is not None:
-            sharded_axes[spec.parallel] = (dim_idx, spec)
+    sharded_axis_infos: dict[ParallelAxis, tuple[int, DimSpec]] = {
+        spec.parallel: (dim_idx, spec)
+        for dim_idx, spec in enumerate(dim_specs)
+        if spec.parallel is not None
+    }
+    if not sharded_axis_infos:
+        return []
 
-    if len(sharded_axes) > 1:
-        raise NotImplementedError(
-            f"Multi-axis unshard is not supported. "
-            f"Got {len(sharded_axes)} sharded axes: {sorted(a.value for a in sharded_axes)}"
+    _validate(sharded_axes=set(sharded_axis_infos), parallel_infos=parallel_infos)
+
+    current_coords: _CoordsList = [
+        {axis: info[axis].axis_rank for axis in sharded_axis_infos}
+        for info in parallel_infos
+    ]
+
+    plans: list[UnshardPlan] = []
+    for axis, (dim_index, spec) in sharded_axis_infos.items():
+        result = _group_and_project(
+            current_coords=current_coords,
+            target_axis=axis,
         )
 
-    if not sharded_axes:
-        return None
+        plans.append(
+            UnshardPlan(
+                axis=axis,
+                params=_resolve_unshard_params(spec=spec, dim_index=dim_index),
+                groups=result.groups,
+            )
+        )
 
-    axis_name, (dim_idx, spec) = next(iter(sharded_axes.items()))
+        current_coords = result.projected_coords
 
-    expected_size: Optional[int] = None
-    rank_to_world: dict[int, int] = {}
+    return plans
 
-    for world_rank, pinfo in enumerate(parallel_infos):
-        if axis_name not in pinfo:
-            continue
 
-        ainfo = pinfo[axis_name]
+def _validate(
+    *,
+    sharded_axes: set[ParallelAxis],
+    parallel_infos: list[dict[ParallelAxis, AxisInfo]],
+) -> None:
+    """Check that every rank has all sharded axes, sizes are consistent, and ranks are complete."""
+    axis_sizes: dict[ParallelAxis, int] = {}
 
-        if expected_size is None:
-            expected_size = ainfo.axis_size
-        elif ainfo.axis_size != expected_size:
+    for world_rank, parallel_info in enumerate(parallel_infos):
+        for axis in sharded_axes:
+            if axis not in parallel_info:
+                raise ValueError(
+                    f"world_rank={world_rank} missing parallel_info for "
+                    f"sharded 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"Inconsistent axis_size for {axis_name.value}: "
-                f"expected {expected_size}, got {ainfo.axis_size} "
-                f"at world_rank={world_rank}"
+                f"axis_rank coverage for {axis.value} is incomplete: "
+                f"got {sorted(seen_ranks)}, expected 0..{expected_size - 1}"
             )
 
-        rank_to_world.setdefault(ainfo.axis_rank, world_rank)
 
-    if expected_size is None:
-        raise ValueError(f"No parallel_info found for sharded axis {axis_name.value!r}")
+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)
 
-    if set(rank_to_world.keys()) != set(range(expected_size)):
-        raise ValueError(
-            f"axis_rank coverage for {axis_name.value} is incomplete: "
-            f"got {sorted(rank_to_world.keys())}, expected 0..{expected_size - 1}"
-        )
+    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))
 
-    return UnshardPlan(
-        axis=spec.parallel,
-        params=_resolve_unshard_params(spec=spec, dim_index=dim_idx),
-        world_ranks_by_axis_rank=[rank_to_world[i] for i in range(expected_size)],
-    )
+    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(*, spec: DimSpec, dim_index: int) -> UnshardParams:

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

@@ -27,7 +27,12 @@ UnshardParams = ConcatParams
 class UnshardPlan(_FrozenBase):
     axis: ParallelAxis
     params: UnshardParams
-    world_ranks_by_axis_rank: list[int]
+    # 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]]
 
 
 # Union of all plan types. Future pipeline components (e.g. reduction,