From e3cdf6b1a3a17ea4ab2dc6e420fbe15c402013fa Mon Sep 17 00:00:00 2001 From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com> Date: Fri, 27 Feb 2026 08:11:47 +0800 Subject: [PATCH] Support CP packed format in unsharder in dump comparator (#19461) --- .../comparator/aligner/unsharder/executor.py | 59 ++++++++- .../comparator/aligner/unsharder/planner.py | 45 ++++++- .../comparator/aligner/unsharder/types.py | 8 +- .../aligner/unsharder/test_executor.py | 116 ++++++++++++++++++ 4 files changed, 221 insertions(+), 7 deletions(-) diff --git a/python/sglang/srt/debug_utils/comparator/aligner/unsharder/executor.py b/python/sglang/srt/debug_utils/comparator/aligner/unsharder/executor.py index c9877cf76..a033730bc 100644 --- a/python/sglang/srt/debug_utils/comparator/aligner/unsharder/executor.py +++ b/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 diff --git a/python/sglang/srt/debug_utils/comparator/aligner/unsharder/planner.py b/python/sglang/srt/debug_utils/comparator/aligner/unsharder/planner.py index 2f713e212..5a525d334 100644 --- a/python/sglang/srt/debug_utils/comparator/aligner/unsharder/planner.py +++ b/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) diff --git a/python/sglang/srt/debug_utils/comparator/aligner/unsharder/types.py b/python/sglang/srt/debug_utils/comparator/aligner/unsharder/types.py index b3cc24eea..a15300388 100644 --- a/python/sglang/srt/debug_utils/comparator/aligner/unsharder/types.py +++ b/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"), ] diff --git a/test/registered/debug_utils/comparator/aligner/unsharder/test_executor.py b/test/registered/debug_utils/comparator/aligner/unsharder/test_executor.py index d8357082a..1a7864682 100644 --- a/test/registered/debug_utils/comparator/aligner/unsharder/test_executor.py +++ b/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__]))