Support CP packed format in unsharder in dump comparator (#19461)

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

@@ -1,12 +1,18 @@
+from typing import Optional
+
 import torch
 
 from sglang.srt.debug_utils.comparator.aligner.unsharder.types import (
     ConcatParams,
+    CpThdConcatParams,
     PickParams,
     UnsharderParams,
     UnsharderPlan,
 )
-from sglang.srt.debug_utils.comparator.dims import ParallelAxis, resolve_dim_by_name
+from sglang.srt.debug_utils.comparator.dims import (
+    ParallelAxis,
+    resolve_dim_by_name,
+)
 from sglang.srt.debug_utils.comparator.output_types import ReplicatedMismatchWarning
 from sglang.srt.debug_utils.comparator.warning_sink import warning_sink
 
@@ -49,7 +55,14 @@ def _apply_unshard(
         dim: int = resolve_dim_by_name(ordered_tensors[0], params.dim_name)
         return torch.cat(ordered_tensors, dim=dim)
 
-    # Phase 2: ReduceSumParams, CpZigzagParams
+    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,
+        )
+
     raise ValueError(f"Unsupported unshard operation: {type(params).__name__}")
 
 
@@ -73,3 +86,45 @@ def _verify_replicated_group(
                     max_abs_diff=(baseline - other).abs().max().item(),
                 )
             )
+
+
+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], ...] = ordered_tensors[0].names
+    stripped: list[torch.Tensor] = [t.rename(None) 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 = result.refine_names(*names)
+    return result

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

@@ -1,14 +1,19 @@
 from collections import defaultdict
-from typing import NamedTuple
+from typing import NamedTuple, Optional
 
 from sglang.srt.debug_utils.comparator.aligner.unsharder.types import (
     AxisInfo,
     ConcatParams,
+    CpThdConcatParams,
     PickParams,
     UnsharderParams,
     UnsharderPlan,
 )
-from sglang.srt.debug_utils.comparator.dims import DimSpec, ParallelAxis
+from sglang.srt.debug_utils.comparator.dims import (
+    TOKEN_DIM_NAME,
+    DimSpec,
+    ParallelAxis,
+)
 
 # _CoordsList[tensor_index][axis] =
 #     the axis_rank (shard position) of the tensor_index-th tensor along `axis`
@@ -24,6 +29,8 @@ class _GroupResult(NamedTuple):
 def compute_unsharder_plan(
     dim_specs: list[DimSpec],
     parallel_infos: list[dict[ParallelAxis, AxisInfo]],
+    *,
+    thd_global_seq_lens: Optional[list[int]] = None,
 ) -> list[UnsharderPlan]:
     if not parallel_infos:
         raise ValueError("parallel_infos must not be empty")
@@ -58,7 +65,14 @@ def compute_unsharder_plan(
     axis_and_params: list[tuple[ParallelAxis, UnsharderParams]] = [
         (axis, PickParams()) for axis in sorted(replicated_axes, key=lambda a: a.value)
     ] + [
-        (axis, _resolve_unshard_params(spec=spec))
+        (
+            axis,
+            _resolve_unshard_params(
+                spec=spec,
+                parallel_infos=parallel_infos,
+                thd_global_seq_lens=thd_global_seq_lens,
+            ),
+        )
         for axis, spec in sharded_axis_infos.items()
     ]
 
@@ -134,9 +148,32 @@ def _group_and_project(
     return _GroupResult(groups=groups, projected_coords=projected)
 
 
-def _resolve_unshard_params(*, spec: DimSpec) -> UnsharderParams:
+def _resolve_unshard_params(
+    *,
+    spec: DimSpec,
+    parallel_infos: list[dict[ParallelAxis, AxisInfo]],
+    thd_global_seq_lens: Optional[list[int]] = None,
+) -> UnsharderParams:
     if spec.reduction is not None:
         raise NotImplementedError(
             f"Unshard for reduction={spec.reduction} not yet implemented (Phase 2)"
         )
+
+    if spec.name == TOKEN_DIM_NAME and thd_global_seq_lens is not None:
+        if spec.parallel is None:
+            raise ValueError(
+                f"THD unshard requires a parallel axis on dim '{spec.name}', but got None"
+            )
+        axis_size: int = parallel_infos[0][spec.parallel].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=spec.name, seq_lens_per_rank=seq_lens_per_rank
+        )
+
     return ConcatParams(dim_name=spec.name)

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

@@ -28,12 +28,18 @@ class ConcatParams(_FrozenBase):
     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"
 
 
 UnsharderParams = Annotated[
-    Union[ConcatParams, PickParams],
+    Union[ConcatParams, CpThdConcatParams, PickParams],
     Field(discriminator="op"),
 ]