chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Support flattened dims in dump comparator (#19678)

Browse Source
This commit is contained in:
fzyzcjy
2026-03-02 18:43:01 +08:00
committed by GitHub
parent 15e83eea61
commit a70dd11011
14 changed files with 872 additions and 260 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/debug_utils/comparator/aligner/unsharder/test_executor.py
+40 -16
View File
@@ -33,7 +33,7 @@ register_cpu_ci(est_time=10, suite="default", nightly=True)
def _name_tensors(
tensors: list[torch.Tensor], dim_specs: list[DimSpec]
) -> list[torch.Tensor]:
names: list[str] = [s.name for s in dim_specs]
names: list[str] = [s.sanitized_name for s in dim_specs]
return [t.refine_names(*names) for t in tensors]
@@ -42,7 +42,7 @@ class TestExecuteUnsharderPlan:
full_tensor = torch.randn(2, 8, 16)
shards = list(full_tensor.chunk(4, dim=1))
dim_specs = parse_dims("b h(tp) d").dims
dim_specs = parse_dims("b h[tp] d").dims
parallel_infos = [
{ParallelAxis.TP: AxisInfo(axis_rank=i, axis_size=4)} for i in range(4)
]
@@ -67,7 +67,7 @@ class TestExecuteUnsharderPlan:
{ParallelAxis.TP: AxisInfo(axis_rank=3, axis_size=4)},
{ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=4)},
]
dim_specs = parse_dims("h(tp) d").dims
dim_specs = parse_dims("h[tp] d").dims
plans = compute_unsharder_plan(dim_specs, parallel_infos)
assert len(plans) == 1
@@ -95,7 +95,7 @@ class TestExecuteUnsharderPlan:
shards_a = list(full_a.chunk(4, dim=0))
shards_b = list(full_b.chunk(4, dim=0))
dim_specs = parse_dims("s(cp) h(tp)").dims
dim_specs = parse_dims("s[cp] h[tp]").dims
parallel_infos = []
for cp_rank in range(2):
for tp_rank in range(4):
@@ -145,7 +145,7 @@ class TestExecuteUnsharderPlan:
}
)
dim_specs = parse_dims("b s(cp) h(tp)").dims
dim_specs = parse_dims("b s[cp] h[tp]").dims
plans = compute_unsharder_plan(dim_specs, parallel_infos)
assert len(plans) == 2
@@ -187,7 +187,7 @@ class TestExecuteUnsharderPlan:
}
)
dim_specs = parse_dims("b s(cp) h(tp)").dims
dim_specs = parse_dims("b s[cp] h[tp]").dims
plans = compute_unsharder_plan(dim_specs, parallel_infos)
assert len(plans) == 2
@@ -241,7 +241,7 @@ class TestExecuteUnsharderPlan:
}
)
dim_specs = parse_dims("b e(ep) s(cp) h(tp)").dims
dim_specs = parse_dims("b e[ep] s[cp] h[tp]").dims
plans = compute_unsharder_plan(dim_specs, parallel_infos)
assert len(plans) == 3
@@ -290,7 +290,7 @@ class TestExecuteUnsharderPlan:
}
)
dim_specs = parse_dims("b e(ep) s(cp) h(tp)").dims
dim_specs = parse_dims("b e[ep] s[cp] h[tp]").dims
plans = compute_unsharder_plan(dim_specs, parallel_infos)
assert len(plans) == 3
@@ -326,7 +326,7 @@ class TestPickOperation:
def test_pick_multiple_groups(self) -> None:
"""PickParams with multiple groups picks one from each."""
dim_specs = parse_dims("h(tp)").dims
dim_specs = parse_dims("h[tp]").dims
parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [
{
ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
@@ -361,7 +361,7 @@ class TestPickOperation:
assert all(c.passed for c in unsharder_result.replicated_checks)
def test_replicated_tp_sharded_cp_e2e(self) -> None:
"""CP2 TP2, dims='b s(cp) d': replicated TP pick + sharded CP concat round-trip."""
"""CP2 TP2, dims='b s[cp] d': replicated TP pick + sharded CP concat round-trip."""
torch.manual_seed(42)
full_tensor = torch.randn(4, 8, 16)
cp_chunks = list(full_tensor.chunk(2, dim=1))
@@ -378,7 +378,7 @@ class TestPickOperation:
}
)
dim_specs = parse_dims("b s(cp) d").dims
dim_specs = parse_dims("b s[cp] d").dims
plans = compute_unsharder_plan(dim_specs, parallel_infos)
assert len(plans) == 2
@@ -654,7 +654,7 @@ class TestReduceSum:
part_a = full_tensor * 0.6
part_b = full_tensor * 0.4
dim_specs = parse_dims("h(tp:partial) d").dims
dim_specs = parse_dims("h[tp:partial] d").dims
parallel_infos = [
{ParallelAxis.TP: AxisInfo(axis_rank=i, axis_size=2)} for i in range(2)
]
@@ -676,7 +676,7 @@ class TestReduceSum:
full_tensor = torch.randn(4, 8)
parts: list[torch.Tensor] = [full_tensor * 0.25 for _ in range(4)]
dim_specs = parse_dims("h(tp:partial) d").dims
dim_specs = parse_dims("h[tp:partial] d").dims
parallel_infos = [
{ParallelAxis.TP: AxisInfo(axis_rank=i, axis_size=4)} for i in range(4)
]
@@ -710,7 +710,7 @@ class TestReduceSum:
}
)
dim_specs = parse_dims("b s(cp) h(tp:partial)").dims
dim_specs = parse_dims("b s[cp] h[tp:partial]").dims
plans = compute_unsharder_plan(dim_specs, parallel_infos)
assert len(plans) == 2
@@ -739,7 +739,7 @@ class TestReduceSum:
{ParallelAxis.TP: AxisInfo(axis_rank=3, axis_size=4)},
{ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=4)},
]
dim_specs = parse_dims("h(tp:partial) d").dims
dim_specs = parse_dims("h[tp:partial] d").dims
plans = compute_unsharder_plan(dim_specs, parallel_infos)
named_parts: list[torch.Tensor] = _name_tensors(parts, dim_specs)
@@ -752,7 +752,7 @@ class TestReduceSum:
def test_reduce_preserves_named_dims(self) -> None:
"""Named tensor dimensions are preserved through reduce_sum."""
dim_specs = parse_dims("h(tp:partial) d").dims
dim_specs = parse_dims("h[tp:partial] d").dims
part_a = torch.randn(4, 8).refine_names("h", "d")
part_b = torch.randn(4, 8).refine_names("h", "d")
@@ -1028,5 +1028,29 @@ class TestReduceSum:
)
class TestFusedDimExecutor:
def test_fused_tp2_concat(self) -> None:
"""Fused dim "t (num_heads*head_dim)[tp]": TP=2 concat on fused axis."""
torch.manual_seed(42)
full_tensor = torch.randn(4, 128) # t=4, nh*hd=128
shards = list(full_tensor.chunk(2, dim=1))
dim_specs = parse_dims("t (num_heads*head_dim)[tp]").dims
parallel_infos = [
{ParallelAxis.TP: AxisInfo(axis_rank=i, axis_size=2)} for i in range(2)
]
plans = compute_unsharder_plan(dim_specs, parallel_infos)
assert len(plans) == 1
named_shards: list[torch.Tensor] = _name_tensors(shards, dim_specs)
unsharder_result: UnsharderResult = execute_unsharder_plan(
plans[0], named_shards
)
assert len(unsharder_result.tensors) == 1
assert torch.allclose(unsharder_result.tensors[0].rename(None), full_tensor)
if __name__ == "__main__":
sys.exit(pytest.main([__file__]))
test/registered/debug_utils/comparator/aligner/unsharder/test_planner.py
+85 -32
View File
@@ -19,7 +19,7 @@ register_cpu_ci(est_time=10, suite="default", nightly=True)
class TestComputeUnsharderPlan:
def test_tp4_plan(self) -> None:
dim_specs = parse_dims("b s h(tp) d").dims
dim_specs = parse_dims("b s h[tp] d").dims
parallel_infos = [
{ParallelAxis.TP: AxisInfo(axis_rank=i, axis_size=4)} for i in range(4)
]
@@ -31,7 +31,7 @@ class TestComputeUnsharderPlan:
assert plans[0].groups == [[0, 1, 2, 3]]
def test_inconsistent_axis_size_raises(self) -> None:
dim_specs = parse_dims("h(tp)").dims
dim_specs = parse_dims("h[tp]").dims
parallel_infos = [
{ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=4)},
{ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=2)},
@@ -41,7 +41,7 @@ class TestComputeUnsharderPlan:
def test_missing_axis_in_all_parallel_infos_skipped(self) -> None:
"""Axis in dims but absent from all parallel_infos -> axis_size=1, auto-skip."""
dim_specs = parse_dims("h(tp)").dims
dim_specs = parse_dims("h[tp]").dims
parallel_infos = [{ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2)}]
# TP not in any parallel_info → skipped; CP is replicated but only 1 rank
# with size=2 → incomplete coverage
@@ -49,13 +49,13 @@ class TestComputeUnsharderPlan:
compute_unsharder_plan(dim_specs, parallel_infos)
def test_empty_parallel_infos_raises(self) -> None:
dim_specs = parse_dims("h(tp)").dims
dim_specs = parse_dims("h[tp]").dims
with pytest.raises(ValueError, match="must not be empty"):
compute_unsharder_plan(dim_specs, [])
def test_scrambled_world_ranks(self) -> None:
"""world_rank order != axis_rank order."""
dim_specs = parse_dims("h(tp)").dims
dim_specs = parse_dims("h[tp]").dims
parallel_infos = [
{ParallelAxis.TP: AxisInfo(axis_rank=2, axis_size=4)},
{ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=4)},
@@ -74,7 +74,7 @@ class TestComputeUnsharderPlan:
def test_multi_axis_plan(self) -> None:
"""Multi-axis (TP + CP) produces a 2-step plan."""
dim_specs = parse_dims("s(cp) h(tp)").dims
dim_specs = parse_dims("s[cp] h[tp]").dims
parallel_infos = [
{
ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
@@ -101,7 +101,7 @@ class TestComputeUnsharderPlan:
def test_cp_tp_plan(self) -> None:
"""CP=2 + TP=4 produces correct 2-step plan with correct groups."""
dim_specs = parse_dims("s(cp) h(tp)").dims
dim_specs = parse_dims("s[cp] h[tp]").dims
parallel_infos = []
for cp_rank in range(2):
for tp_rank in range(4):
@@ -129,7 +129,7 @@ class TestComputeUnsharderPlan:
def test_cp_tp_scrambled_ranks(self) -> None:
"""Scrambled rank assignment still produces correct plan."""
dim_specs = parse_dims("s(cp) h(tp)").dims
dim_specs = parse_dims("s[cp] h[tp]").dims
parallel_infos = [
{
ParallelAxis.CP: AxisInfo(axis_rank=1, axis_size=2),
@@ -165,7 +165,7 @@ class TestComputeUnsharderPlan:
def test_axis_rank_coverage_incomplete_raises(self) -> None:
"""TP size=4 but only ranks 0,1,3 provided (missing rank 2)."""
dim_specs = parse_dims("h(tp)").dims
dim_specs = parse_dims("h[tp]").dims
parallel_infos = [
{ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=4)},
{ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=4)},
@@ -175,7 +175,7 @@ class TestComputeUnsharderPlan:
compute_unsharder_plan(dim_specs, parallel_infos)
def test_reduction_partial_returns_reduce_sum(self) -> None:
dim_specs = parse_dims("h(tp:partial)").dims
dim_specs = parse_dims("h[tp:partial]").dims
parallel_infos = [
{ParallelAxis.TP: AxisInfo(axis_rank=i, axis_size=2)} for i in range(2)
]
@@ -188,7 +188,7 @@ class TestComputeUnsharderPlan:
def test_reduction_partial_tp4(self) -> None:
"""TP=4 with partial reduction produces a single ReduceSumParams step."""
dim_specs = parse_dims("h(tp:partial)").dims
dim_specs = parse_dims("h[tp:partial]").dims
parallel_infos = [
{ParallelAxis.TP: AxisInfo(axis_rank=i, axis_size=4)} for i in range(4)
]
@@ -200,7 +200,7 @@ class TestComputeUnsharderPlan:
def test_multi_axis_with_reduction_on_one(self) -> None:
"""CP concat + TP reduce produces a 2-step plan."""
dim_specs = parse_dims("s(cp) h(tp:partial)").dims
dim_specs = parse_dims("s[cp] h[tp:partial]").dims
parallel_infos: list[dict[ParallelAxis, AxisInfo]] = []
for cp_rank in range(2):
for tp_rank in range(2):
@@ -221,7 +221,7 @@ class TestComputeUnsharderPlan:
def test_reduction_scrambled_ranks(self) -> None:
"""Scrambled world_rank order with partial reduction."""
dim_specs = parse_dims("h(tp:partial)").dims
dim_specs = parse_dims("h[tp:partial]").dims
parallel_infos = [
{ParallelAxis.TP: AxisInfo(axis_rank=2, axis_size=4)},
{ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=4)},
@@ -235,7 +235,7 @@ class TestComputeUnsharderPlan:
assert plans[0].groups == [[1, 3, 0, 2]]
def test_ordering_zigzag_accepted(self) -> None:
dim_specs = parse_dims("s(cp:zigzag)").dims
dim_specs = parse_dims("s[cp:zigzag]").dims
parallel_infos = [
{ParallelAxis.CP: AxisInfo(axis_rank=i, axis_size=2)} for i in range(2)
]
@@ -244,7 +244,7 @@ class TestComputeUnsharderPlan:
assert plans[0].axis == ParallelAxis.CP
def test_ordering_natural_accepted(self) -> None:
dim_specs = parse_dims("s(cp:natural)").dims
dim_specs = parse_dims("s[cp:natural]").dims
parallel_infos = [
{ParallelAxis.CP: AxisInfo(axis_rank=i, axis_size=2)} for i in range(2)
]
@@ -254,7 +254,7 @@ class TestComputeUnsharderPlan:
def test_three_axis_plan(self) -> None:
"""EP=2 + CP=2 + TP=2 produces a 3-step plan."""
dim_specs = parse_dims("b e(ep) s(cp) h(tp)").dims
dim_specs = parse_dims("b e[ep] s[cp] h[tp]").dims
parallel_infos: list[dict[ParallelAxis, AxisInfo]] = []
for ep_rank in range(2):
for cp_rank in range(2):
@@ -289,8 +289,8 @@ class TestComputeUnsharderPlan:
assert len(plans[2].groups[0]) == 2
def test_same_dim_cp_sp_plan(self) -> None:
"""t(cp:zigzag,sp) with CP=2 SP=2: SP unshards first (inner), then CP."""
dim_specs = parse_dims("t(cp:zigzag,sp) 1 h").dims
"""t[cp:zigzag,sp] with CP=2 SP=2: SP unshards first (inner), then CP."""
dim_specs = parse_dims("t[cp:zigzag,sp] 1 h").dims
parallel_infos: list[dict[ParallelAxis, AxisInfo]] = []
for cp_rank in range(2):
for sp_rank in range(2):
@@ -323,12 +323,12 @@ class TestComputeUnsharderPlan:
assert len(cp_plan.groups[0]) == 2
def test_same_dim_cp_sp_with_thd(self) -> None:
"""t(cp:zigzag,sp) with THD: SP → ConcatParams, CP → CpThdConcatParams."""
"""t[cp:zigzag,sp] with THD: SP → ConcatParams, CP → CpThdConcatParams."""
from sglang.srt.debug_utils.comparator.aligner.unsharder.types import (
CpThdConcatParams,
)
dim_specs = parse_dims("t(cp:zigzag,sp) h").dims
dim_specs = parse_dims("t[cp:zigzag,sp] h").dims
parallel_infos: list[dict[ParallelAxis, AxisInfo]] = []
for cp_rank in range(2):
for sp_rank in range(2):
@@ -360,8 +360,8 @@ class TestComputeUnsharderPlan:
assert cp_plan.params.seq_lens_per_rank == [50, 32]
def test_sp_in_dims_but_not_in_parallel_info(self) -> None:
"""s(sp) in dims but SP absent from parallel_info (SP disabled), should auto-skip."""
dim_specs = parse_dims("s(sp) b h(tp)").dims
"""s[sp] in dims but SP absent from parallel_info (SP disabled), should auto-skip."""
dim_specs = parse_dims("s[sp] b h[tp]").dims
parallel_infos = [
{ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2)},
{ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=2)},
@@ -371,15 +371,15 @@ class TestComputeUnsharderPlan:
assert plans[0].axis == ParallelAxis.TP
def test_all_dims_sharded_but_single_gpu(self) -> None:
"""Single GPU (TP=1, CP=1), dims has s(cp) h(tp) but parallel_info is empty."""
dim_specs = parse_dims("b s(cp) h(tp) d").dims
"""Single GPU (TP=1, CP=1), dims has s[cp] h[tp] but parallel_info is empty."""
dim_specs = parse_dims("b s[cp] h[tp] d").dims
parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [{}]
plans = compute_unsharder_plan(dim_specs, parallel_infos)
assert plans == []
def test_sharded_axis_missing_from_rank_raises(self) -> None:
"""A world_rank missing a sharded axis raises ValueError."""
dim_specs = parse_dims("s(cp) h(tp)").dims
dim_specs = parse_dims("s[cp] h[tp]").dims
parallel_infos = [
{
ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
@@ -396,8 +396,8 @@ class TestComputeUnsharderPlan:
class TestReplicatedAxes:
def test_replicated_tp_with_sharded_cp(self) -> None:
"""CP2 TP2, dims='b s(cp) d' → PickPlan(TP) + ConcatPlan(CP)."""
dim_specs = parse_dims("b s(cp) d").dims
"""CP2 TP2, dims='b s[cp] d' → PickPlan(TP) + ConcatPlan(CP)."""
dim_specs = parse_dims("b s[cp] d").dims
parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [
{
ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
@@ -458,8 +458,8 @@ class TestReplicatedAxes:
assert axes == {ParallelAxis.CP, ParallelAxis.TP}
def test_multiple_replicated_one_sharded(self) -> None:
"""CP2 TP2 EP2, dims='h(tp)' → PickPlan(CP) + PickPlan(EP) + ConcatPlan(TP)."""
dim_specs = parse_dims("h(tp)").dims
"""CP2 TP2 EP2, dims='h[tp]' → PickPlan(CP) + PickPlan(EP) + ConcatPlan(TP)."""
dim_specs = parse_dims("h[tp]").dims
parallel_infos: list[dict[ParallelAxis, AxisInfo]] = []
for cp_rank in range(2):
for ep_rank in range(2):
@@ -486,7 +486,7 @@ class TestReplicatedAxes:
def test_replicated_scrambled_ranks(self) -> None:
"""Scrambled world_rank order with replicated axis."""
dim_specs = parse_dims("h(tp)").dims
dim_specs = parse_dims("h[tp]").dims
parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [
{
ParallelAxis.CP: AxisInfo(axis_rank=1, axis_size=2),
@@ -515,7 +515,7 @@ class TestReplicatedAxes:
def test_replicated_axis_inconsistent_size_raises(self) -> None:
"""Replicated axis with inconsistent sizes raises ValueError."""
dim_specs = parse_dims("h(tp)").dims
dim_specs = parse_dims("h[tp]").dims
parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [
{
ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
@@ -531,7 +531,7 @@ class TestReplicatedAxes:
def test_replicated_axis_missing_from_rank_raises(self) -> None:
"""A rank missing a replicated axis that other ranks have raises ValueError."""
dim_specs = parse_dims("h(tp)").dims
dim_specs = parse_dims("h[tp]").dims
parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [
{
ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
@@ -560,5 +560,58 @@ class TestReplicatedAxes:
assert plans[0].groups == [[0, 1]]
class TestComputeUnsharderPlanFusedDims:
def test_fused_dim_tp2(self) -> None:
"""Fused dim "(num_heads*head_dim)[tp]" should unshard on the fused tensor name."""
dim_specs = parse_dims("t (num_heads*head_dim)[tp]").dims
parallel_infos = [
{ParallelAxis.TP: AxisInfo(axis_rank=i, axis_size=2)} for i in range(2)
]
plans = compute_unsharder_plan(dim_specs, parallel_infos)
assert len(plans) == 1
assert plans[0].axis == ParallelAxis.TP
assert isinstance(plans[0].params, ConcatParams)
assert plans[0].params.dim_name == "num_heads___head_dim"
assert plans[0].groups == [[0, 1]]
def test_fused_dim_modifier_on_second_sub(self) -> None:
"""Modifier on fused dim: "(a*b)[tp]" should produce concat plan."""
dim_specs = parse_dims("t (a*b)[tp]").dims
parallel_infos = [
{ParallelAxis.TP: AxisInfo(axis_rank=i, axis_size=2)} for i in range(2)
]
plans = compute_unsharder_plan(dim_specs, parallel_infos)
assert len(plans) == 1
assert plans[0].axis == ParallelAxis.TP
assert isinstance(plans[0].params, ConcatParams)
assert plans[0].params.dim_name == "a___b"
def test_fused_dim_no_modifier(self) -> None:
"""Fused dim without any modifier should have no unshard plans (beyond replicated)."""
dim_specs = parse_dims("t (a*b)").dims
parallel_infos = [
{ParallelAxis.TP: AxisInfo(axis_rank=i, axis_size=2)} for i in range(2)
]
plans = compute_unsharder_plan(dim_specs, parallel_infos)
# TP not annotated in dims → replicated → pick
assert len(plans) == 1
assert isinstance(plans[0].params, PickParams)
def test_fused_dim_with_reduction(self) -> None:
"""Fused dim with partial reduction: "(a*b)[tp:partial]"."""
dim_specs = parse_dims("t (a*b)[tp:partial]").dims
parallel_infos = [
{ParallelAxis.TP: AxisInfo(axis_rank=i, axis_size=2)} for i in range(2)
]
plans = compute_unsharder_plan(dim_specs, parallel_infos)
assert len(plans) == 1
assert plans[0].axis == ParallelAxis.TP
assert isinstance(plans[0].params, ReduceSumParams)
if __name__ == "__main__":
sys.exit(pytest.main([__file__]))