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"),
 ]
 

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

@@ -13,7 +13,9 @@ from sglang.srt.debug_utils.comparator.aligner.unsharder.planner import (
 )
 from sglang.srt.debug_utils.comparator.aligner.unsharder.types import (
     AxisInfo,
+    CpThdConcatParams,
     PickParams,
+    UnsharderPlan,
 )
 from sglang.srt.debug_utils.comparator.dims import (
     DimSpec,
@@ -505,5 +507,119 @@ class TestVerifyReplicatedGroup:
         assert warnings[0].differing_index == 1
 
 
+class TestThdCpConcat:
+    def test_single_seq(self) -> None:
+        """Single seq THD unshard: 2 ranks → per-seq concat."""
+        rank0 = torch.tensor([1, 2, 3]).refine_names("t")
+        rank1 = torch.tensor([4, 5, 6]).refine_names("t")
+
+        plan = UnsharderPlan(
+            axis=ParallelAxis.CP,
+            params=CpThdConcatParams(dim_name="t", seq_lens_per_rank=[3]),
+            groups=[[0, 1]],
+        )
+        with warning_sink.context():
+            result = execute_unsharder_plan(plan, [rank0, rank1])
+
+        assert len(result) == 1
+        expected = torch.tensor([1, 2, 3, 4, 5, 6])
+        assert torch.equal(result[0].rename(None), expected)
+
+    def test_multi_seq(self) -> None:
+        """Multi-seq THD unshard: 2 ranks, seq_lens=[50, 32, 46]."""
+        # rank0: [seqA_r0(50) | seqB_r0(32) | pad_r0(46)]
+        # rank1: [seqA_r1(50) | seqB_r1(32) | pad_r1(46)]
+        seq_a_r0 = torch.arange(0, 50)
+        seq_b_r0 = torch.arange(100, 132)
+        pad_r0 = torch.full((46,), -1)
+        rank0 = torch.cat([seq_a_r0, seq_b_r0, pad_r0]).refine_names("t")
+
+        seq_a_r1 = torch.arange(50, 100)
+        seq_b_r1 = torch.arange(132, 164)
+        pad_r1 = torch.full((46,), -2)
+        rank1 = torch.cat([seq_a_r1, seq_b_r1, pad_r1]).refine_names("t")
+
+        plan = UnsharderPlan(
+            axis=ParallelAxis.CP,
+            params=CpThdConcatParams(dim_name="t", seq_lens_per_rank=[50, 32, 46]),
+            groups=[[0, 1]],
+        )
+        with warning_sink.context():
+            result = execute_unsharder_plan(plan, [rank0, rank1])
+
+        assert len(result) == 1
+        unsharded: torch.Tensor = result[0].rename(None)
+
+        # seqA: r0(50) + r1(50) = 100 tokens, values 0..99
+        assert torch.equal(unsharded[:100], torch.cat([seq_a_r0, seq_a_r1]))
+        # seqB: r0(32) + r1(32) = 64 tokens
+        assert torch.equal(unsharded[100:164], torch.cat([seq_b_r0, seq_b_r1]))
+        # pad: r0(46) + r1(46) = 92 tokens
+        assert torch.equal(unsharded[164:256], torch.cat([pad_r0, pad_r1]))
+
+    def test_with_hidden_dim(self) -> None:
+        """THD unshard with trailing hidden dim: shape [T, H]."""
+        torch.manual_seed(42)
+        hidden: int = 4
+        # rank0: [seqA_r0(3, 4) | seqB_r0(2, 4)]
+        # rank1: [seqA_r1(3, 4) | seqB_r1(2, 4)]
+        seq_a_r0 = torch.randn(3, hidden)
+        seq_b_r0 = torch.randn(2, hidden)
+        rank0 = torch.cat([seq_a_r0, seq_b_r0]).refine_names("t", "h")
+
+        seq_a_r1 = torch.randn(3, hidden)
+        seq_b_r1 = torch.randn(2, hidden)
+        rank1 = torch.cat([seq_a_r1, seq_b_r1]).refine_names("t", "h")
+
+        plan = UnsharderPlan(
+            axis=ParallelAxis.CP,
+            params=CpThdConcatParams(dim_name="t", seq_lens_per_rank=[3, 2]),
+            groups=[[0, 1]],
+        )
+        with warning_sink.context():
+            result = execute_unsharder_plan(plan, [rank0, rank1])
+
+        assert len(result) == 1
+        unsharded: torch.Tensor = result[0].rename(None)
+
+        assert unsharded.shape == (10, hidden)
+        assert torch.equal(unsharded[:6], torch.cat([seq_a_r0, seq_a_r1]))
+        assert torch.equal(unsharded[6:10], torch.cat([seq_b_r0, seq_b_r1]))
+
+    def test_with_leading_batch_dim(self) -> None:
+        """THD unshard with leading batch dim: shape [B, T, H], t is dim=1."""
+        torch.manual_seed(42)
+        batch: int = 2
+        hidden: int = 4
+        # rank0: [seqA_r0(3) | seqB_r0(2)] per batch item
+        # rank1: [seqA_r1(3) | seqB_r1(2)] per batch item
+        seq_a_r0 = torch.randn(batch, 3, hidden)
+        seq_b_r0 = torch.randn(batch, 2, hidden)
+        rank0 = torch.cat([seq_a_r0, seq_b_r0], dim=1).refine_names("b", "t", "h")
+
+        seq_a_r1 = torch.randn(batch, 3, hidden)
+        seq_b_r1 = torch.randn(batch, 2, hidden)
+        rank1 = torch.cat([seq_a_r1, seq_b_r1], dim=1).refine_names("b", "t", "h")
+
+        plan = UnsharderPlan(
+            axis=ParallelAxis.CP,
+            params=CpThdConcatParams(dim_name="t", seq_lens_per_rank=[3, 2]),
+            groups=[[0, 1]],
+        )
+        with warning_sink.context():
+            result = execute_unsharder_plan(plan, [rank0, rank1])
+
+        assert len(result) == 1
+        unsharded: torch.Tensor = result[0].rename(None)
+
+        assert unsharded.shape == (batch, 10, hidden)
+        # seqA: r0(3) + r1(3) = 6 tokens per batch
+        assert torch.equal(unsharded[:, :6, :], torch.cat([seq_a_r0, seq_a_r1], dim=1))
+        # seqB: r0(2) + r1(2) = 4 tokens per batch
+        assert torch.equal(
+            unsharded[:, 6:10, :], torch.cat([seq_b_r0, seq_b_r1], dim=1)
+        )
+
+
 if __name__ == "__main__":
     sys.exit(pytest.main([__file__]))