chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 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/reorderer/test_planner.py
+15 -15
View File
@@ -25,8 +25,8 @@ register_cpu_ci(est_time=10, suite="default", nightly=True)
class TestComputeReordererPlans:
def test_compute_reorderer_plans_zigzag(self) -> None:
"""s(cp:zigzag) produces a ReordererPlan."""
dim_specs = parse_dims("b s(cp:zigzag) h(tp)").dims
"""s[cp:zigzag] produces a ReordererPlan."""
dim_specs = parse_dims("b s[cp:zigzag] h[tp]").dims
parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [
{
ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
@@ -43,8 +43,8 @@ class TestComputeReordererPlans:
assert plans[0].params.cp_size == 2
def test_compute_reorderer_plans_thd_zigzag(self) -> None:
"""t(cp:zigzag) produces a ZigzagToNaturalThdParams plan."""
dim_specs = parse_dims("t(cp:zigzag) h(tp)").dims
"""t[cp:zigzag] produces a ZigzagToNaturalThdParams plan."""
dim_specs = parse_dims("t[cp:zigzag] h[tp]").dims
parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [
{
ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
@@ -64,8 +64,8 @@ class TestComputeReordererPlans:
assert plans[0].params.seq_lens == [100, 64, 92]
def test_non_seq_dim_still_raises(self) -> None:
"""Zigzag on non-sequence/non-token dim (e.g. h(cp:zigzag)) raises ValueError."""
dim_specs = parse_dims("h(cp:zigzag) d").dims
"""Zigzag on non-sequence/non-token dim (e.g. h[cp:zigzag]) raises ValueError."""
dim_specs = parse_dims("h[cp:zigzag] d").dims
parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [
{ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2)},
]
@@ -73,8 +73,8 @@ class TestComputeReordererPlans:
compute_reorderer_plans(dim_specs=dim_specs, parallel_infos=parallel_infos)
def test_thd_zigzag_without_seq_lens_raises(self) -> None:
"""t(cp:zigzag) without thd_global_seq_lens raises ValueError."""
dim_specs = parse_dims("t(cp:zigzag) h(tp)").dims
"""t[cp:zigzag] without thd_global_seq_lens raises ValueError."""
dim_specs = parse_dims("t[cp:zigzag] h[tp]").dims
parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [
{
ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
@@ -85,8 +85,8 @@ class TestComputeReordererPlans:
compute_reorderer_plans(dim_specs=dim_specs, parallel_infos=parallel_infos)
def test_thd_natural_no_reorder(self) -> None:
"""t(cp:natural) and t(cp) produce no reorder plans."""
for dims_str in ["t(cp:natural) h(tp)", "t(cp) h(tp)"]:
"""t[cp:natural] and t[cp] produce no reorder plans."""
for dims_str in ["t[cp:natural] h[tp]", "t[cp] h[tp]"]:
dim_specs = parse_dims(dims_str).dims
parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [
{
@@ -100,8 +100,8 @@ class TestComputeReordererPlans:
assert plans == []
def test_compute_reorderer_plans_natural(self) -> None:
"""s(cp) and s(cp:natural) produce no reorder plans."""
for dims_str in ["b s(cp) h(tp)", "b s(cp:natural) h(tp)"]:
"""s[cp] and s[cp:natural] produce no reorder plans."""
for dims_str in ["b s[cp] h[tp]", "b s[cp:natural] h[tp]"]:
dim_specs = parse_dims(dims_str).dims
parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [
{
@@ -141,7 +141,7 @@ class TestCpZigzagTpE2E:
}
)
dim_specs: list[DimSpec] = parse_dims("b s(cp:zigzag) h(tp)").dims
dim_specs: list[DimSpec] = parse_dims("b s[cp:zigzag] h[tp]").dims
dim_names: list[str] = [s.name for s in dim_specs]
unsharder_plans = compute_unsharder_plan(
@@ -167,7 +167,7 @@ class TestCpZigzagTpE2E:
class TestCpZigzagSpSameDimE2E:
"""E2E test for t(cp:zigzag,sp) — two axes sharding the same token dim."""
"""E2E test for t[cp:zigzag,sp] — two axes sharding the same token dim."""
def test_cp2_sp2_zigzag_e2e(self) -> None:
"""CP=2 zigzag + SP=2 on same token dim: full unshard + reorder round-trip.
@@ -215,7 +215,7 @@ class TestCpZigzagSpSameDimE2E:
}
)
dim_specs: list[DimSpec] = parse_dims("t(cp:zigzag,sp) h").dims
dim_specs: list[DimSpec] = parse_dims("t[cp:zigzag,sp] h").dims
dim_names: list[str] = [s.name for s in dim_specs]
unsharder_plans = compute_unsharder_plan(