Support context parallel zigzag reordering in dump comparator (#19281)

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

@@ -0,0 +1,82 @@
+from typing import Literal
+
+import torch
+
+from sglang.srt.debug_utils.comparator.aligner.unshard.types import AxisInfo
+from sglang.srt.debug_utils.comparator.dims import DimSpec, Ordering, ParallelAxis
+from sglang.srt.debug_utils.comparator.utils import _FrozenBase
+
+
+class ZigzagToNaturalParams(_FrozenBase):
+    op: Literal["zigzag_to_natural"] = "zigzag_to_natural"
+    dim: int
+    cp_size: int
+
+
+ReorderParams = ZigzagToNaturalParams
+
+
+class ReorderPlan(_FrozenBase):
+    params: ReorderParams
+
+
+_ALLOWED_ZIGZAG_DIM_NAMES: set[str] = {"s"}
+
+
+def compute_reorder_plans(
+    dim_specs: list[DimSpec],
+    parallel_infos: list[dict[ParallelAxis, AxisInfo]],
+) -> list[ReorderPlan]:
+    plans: list[ReorderPlan] = []
+
+    for dim_index, spec in enumerate(dim_specs):
+        if (
+            spec.ordering is not None
+            and spec.ordering != Ordering.NATURAL
+            and spec.parallel is not None
+        ):
+            if spec.name not in _ALLOWED_ZIGZAG_DIM_NAMES:
+                raise ValueError(
+                    f"Zigzag ordering is only supported on sequence dims "
+                    f"(bshd/sbhd format, dim name must be one of "
+                    f"{sorted(_ALLOWED_ZIGZAG_DIM_NAMES)}), "
+                    f"but got dim name {spec.name!r} in {spec}"
+                )
+
+            assert spec.ordering == Ordering.ZIGZAG
+            axis_size: int = parallel_infos[0][spec.parallel].axis_size
+            plans.append(
+                ReorderPlan(
+                    params=ZigzagToNaturalParams(dim=dim_index, cp_size=axis_size),
+                )
+            )
+
+    return plans
+
+
+def execute_reorder_plan(
+    plan: ReorderPlan,
+    tensors: list[torch.Tensor],
+) -> list[torch.Tensor]:
+    return [
+        _reorder_zigzag_to_natural(
+            tensor, dim=plan.params.dim, cp_size=plan.params.cp_size
+        )
+        for tensor in tensors
+    ]
+
+
+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).
+    """
+    num_chunks: int = cp_size * 2
+    chunks: tuple[torch.Tensor, ...] = tensor.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)
+    ]
+    return torch.cat([chunks[i] for i in order], dim=dim)

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

@@ -1,6 +1,6 @@
 import torch
 
-from sglang.srt.debug_utils.comparator.unshard.types import (
+from sglang.srt.debug_utils.comparator.aligner.unshard.types import (
     ConcatParams,
     UnshardParams,
     UnshardPlan,

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

@@ -1,7 +1,7 @@
 from typing import Optional
 
+from sglang.srt.debug_utils.comparator.aligner.unshard.types import AxisInfo
 from sglang.srt.debug_utils.comparator.dims import ParallelAxis
-from sglang.srt.debug_utils.comparator.unshard.types import AxisInfo
 
 _PARALLEL_INFO_KEYS = ("sglang_parallel_info", "megatron_parallel_info")
 

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

@@ -1,13 +1,13 @@
 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 (
+from sglang.srt.debug_utils.comparator.aligner.unshard.types import (
     AxisInfo,
     ConcatParams,
     UnshardParams,
     UnshardPlan,
 )
+from sglang.srt.debug_utils.comparator.dims import DimSpec, ParallelAxis
 
 # _CoordsList[tensor_index][axis] =
 #     the axis_rank (shard position) of the tensor_index-th tensor along `axis`
@@ -127,8 +127,4 @@ def _resolve_unshard_params(*, spec: DimSpec, dim_index: int) -> UnshardParams:
         raise NotImplementedError(
             f"Unshard for reduction={spec.reduction} not yet implemented (Phase 2)"
         )
-    if spec.ordering is not None and spec.ordering != Ordering.NATURAL:
-        raise NotImplementedError(
-            f"Unshard for ordering={spec.ordering} not yet implemented (Phase 2)"
-        )
     return ConcatParams(dim=dim_index)

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

@@ -16,11 +16,6 @@ class ConcatParams(_FrozenBase):
     dim: int
 
 
-# Phase 2: add ReduceSumParams, CpZigzagParams here, then change UnshardParams to:
-#   UnshardParams = Annotated[
-#       Union[ConcatParams, ReduceSumParams, CpZigzagParams],
-#       Field(discriminator="op"),
-#   ]
 UnshardParams = ConcatParams
 
 
@@ -33,8 +28,3 @@ class UnshardPlan(_FrozenBase):
     #   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,
-# reordering) will add their own plan types here.
-Plan = UnshardPlan

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

@@ -1,22 +1,33 @@
 from pathlib import Path
-from typing import Any, Optional
+from typing import Any, Optional, Union
 
 import torch
 
+from sglang.srt.debug_utils.comparator.aligner.reorder import (
+    ReorderPlan,
+    compute_reorder_plans,
+    execute_reorder_plan,
+)
+from sglang.srt.debug_utils.comparator.aligner.unshard.executor import (
+    execute_unshard_plan,
+)
+from sglang.srt.debug_utils.comparator.aligner.unshard.parallel_info import (
+    normalize_parallel_info,
+)
+from sglang.srt.debug_utils.comparator.aligner.unshard.planner import (
+    compute_unshard_plan,
+)
+from sglang.srt.debug_utils.comparator.aligner.unshard.types import UnshardPlan
 from sglang.srt.debug_utils.comparator.dims import parse_dims
 from sglang.srt.debug_utils.comparator.output_types import (
     ComparisonRecord,
     SkipRecord,
 )
 from sglang.srt.debug_utils.comparator.tensor_comparison.compare import compare_tensors
-from sglang.srt.debug_utils.comparator.unshard.executor import execute_unshard_plan
-from sglang.srt.debug_utils.comparator.unshard.parallel_info import (
-    normalize_parallel_info,
-)
-from sglang.srt.debug_utils.comparator.unshard.planner import compute_unshard_plan
-from sglang.srt.debug_utils.comparator.unshard.types import Plan, UnshardPlan
 from sglang.srt.debug_utils.dump_loader import ValueWithMeta
 
+Plan = Union[UnshardPlan, ReorderPlan]
+
 
 def process_tensor_group(
     *,
@@ -83,7 +94,13 @@ 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]
 
-    return compute_unshard_plan(dim_specs=dim_specs, parallel_infos=parallel_infos)
+    unshard_plans = compute_unshard_plan(
+        dim_specs=dim_specs, parallel_infos=parallel_infos
+    )
+    reorder_plans = compute_reorder_plans(
+        dim_specs=dim_specs, parallel_infos=parallel_infos
+    )
+    return [*unshard_plans, *reorder_plans]
 
 
 def _extract_tensors(
@@ -106,10 +123,16 @@ def _execute_plans(
 
     current = tensors
     for plan in plans:
-        if isinstance(plan, UnshardPlan):
-            current = execute_unshard_plan(plan, current)
-        else:
-            raise NotImplementedError(f"Unknown {plan=}")
+        current = _execute_plan(current, plan)
 
     assert len(current) == 1
     return current[0]
+
+
+def _execute_plan(tensors, plan):
+    if isinstance(plan, UnshardPlan):
+        return execute_unshard_plan(plan, tensors)
+    elif isinstance(plan, ReorderPlan):
+        return execute_reorder_plan(plan, tensors)
+    else:
+        raise NotImplementedError(f"Unknown {plan=}")