Make reorderer support packed format with CP in dump comparator (#19462)

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

@@ -1,6 +1,12 @@
+from typing import Optional
+
 import torch
 
-from sglang.srt.debug_utils.comparator.aligner.reorderer.types import ReordererPlan
+from sglang.srt.debug_utils.comparator.aligner.reorderer.types import (
+    ReordererPlan,
+    ZigzagToNaturalParams,
+    ZigzagToNaturalThdParams,
+)
 from sglang.srt.debug_utils.comparator.dims import (
     resolve_dim_by_name,
     strip_dim_names,
@@ -11,12 +17,66 @@ def execute_reorderer_plan(
     plan: ReordererPlan,
     tensors: list[torch.Tensor],
 ) -> list[torch.Tensor]:
-    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
+    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.
+    """
+    stripped: torch.Tensor = strip_dim_names(tensor)
+    names: tuple[Optional[str], ...] = tensor.names
+
+    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 = result.refine_names(*names)
+    return result
+
 
 def _reorder_zigzag_to_natural(
     tensor: torch.Tensor, *, dim: int, cp_size: int
@@ -27,7 +87,7 @@ def _reorder_zigzag_to_natural(
     (megatron/core/ssm/mamba_context_parallel.py:360-373).
     """
     stripped: torch.Tensor = strip_dim_names(tensor)
-    names: tuple = tensor.names
+    names: tuple[Optional[str], ...] = tensor.names
 
     num_chunks: int = cp_size * 2
     chunks: tuple[torch.Tensor, ...] = stripped.chunk(num_chunks, dim=dim)

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

@@ -1,21 +1,27 @@
+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 import (
     SEQ_DIM_NAME,
+    TOKEN_DIM_NAME,
     DimSpec,
     Ordering,
     ParallelAxis,
 )
 
-_ALLOWED_ZIGZAG_DIM_NAMES: set[str] = {SEQ_DIM_NAME}
+_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] = []
 
@@ -28,17 +34,35 @@ def compute_reorderer_plans(
             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"(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(
-                ReordererPlan(
-                    params=ZigzagToNaturalParams(dim_name=spec.name, cp_size=axis_size),
+            if spec.ordering != Ordering.ZIGZAG:
+                raise ValueError(
+                    f"Unsupported ordering {spec.ordering!r} for dim {spec.name!r}"
                 )
-            )
+            axis_size: int = parallel_infos[0][spec.parallel].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

@@ -1,4 +1,6 @@
-from typing import Literal
+from typing import Annotated, Literal, Union
+
+from pydantic import Field
 
 from sglang.srt.debug_utils.comparator.utils import _FrozenBase
 
@@ -9,7 +11,17 @@ class ZigzagToNaturalParams(_FrozenBase):
     cp_size: int
 
 
-ReordererParams = ZigzagToNaturalParams
+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):