Support non orthogonal parallel axes and explicit replication annotation in dump comparator (#19679)
This commit is contained in:
fzyzcjy
@@ -22,7 +22,7 @@ from sglang.srt.debug_utils.comparator.aligner.unsharder.types import (
|
||||
ConcatParams,
|
||||
UnsharderPlan,
|
||||
)
|
||||
from sglang.srt.debug_utils.comparator.dims import ParallelAxis, TokenLayout
|
||||
from sglang.srt.debug_utils.comparator.dims_spec import ParallelAxis, TokenLayout
|
||||
from sglang.srt.debug_utils.comparator.utils import Pair
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
|
||||
@@ -21,7 +21,7 @@ from sglang.srt.debug_utils.comparator.aligner.unsharder.types import (
|
||||
CpThdConcatParams,
|
||||
UnsharderPlan,
|
||||
)
|
||||
from sglang.srt.debug_utils.comparator.dims import TokenLayout
|
||||
from sglang.srt.debug_utils.comparator.dims_spec import TokenLayout
|
||||
from sglang.srt.debug_utils.comparator.utils import Pair
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
|
||||
@@ -19,7 +19,7 @@ from sglang.srt.debug_utils.comparator.aligner.unsharder.types import (
|
||||
CpThdConcatParams,
|
||||
UnsharderPlan,
|
||||
)
|
||||
from sglang.srt.debug_utils.comparator.dims import ParallelAxis
|
||||
from sglang.srt.debug_utils.comparator.dims_spec import ParallelAxis
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
register_cpu_ci(est_time=10, suite="default", nightly=True)
|
||||
|
||||
@@ -17,7 +17,11 @@ from sglang.srt.debug_utils.comparator.aligner.unsharder.planner import (
|
||||
compute_unsharder_plan,
|
||||
)
|
||||
from sglang.srt.debug_utils.comparator.aligner.unsharder.types import AxisInfo
|
||||
from sglang.srt.debug_utils.comparator.dims import DimSpec, ParallelAxis, parse_dims
|
||||
from sglang.srt.debug_utils.comparator.dims_spec import (
|
||||
DimSpec,
|
||||
ParallelAxis,
|
||||
parse_dims,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
register_cpu_ci(est_time=10, suite="default", nightly=True)
|
||||
|
||||
@@ -13,7 +13,7 @@ from sglang.srt.debug_utils.comparator.aligner.token_aligner.smart.types import
|
||||
SGLangSeqId,
|
||||
TokenAlignerStepAux,
|
||||
)
|
||||
from sglang.srt.debug_utils.comparator.dims import TokenLayout
|
||||
from sglang.srt.debug_utils.comparator.dims_spec import TokenLayout
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
register_cpu_ci(est_time=15, suite="default", nightly=True)
|
||||
|
||||
@@ -21,7 +21,7 @@ from sglang.srt.debug_utils.comparator.aligner.token_aligner.smart.types import
|
||||
TokenAlignerStepAux,
|
||||
TokenLocator,
|
||||
)
|
||||
from sglang.srt.debug_utils.comparator.dims import TokenLayout
|
||||
from sglang.srt.debug_utils.comparator.dims_spec import TokenLayout
|
||||
from sglang.srt.debug_utils.comparator.utils import Pair
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
|
||||
@@ -19,7 +19,7 @@ from sglang.srt.debug_utils.comparator.aligner.token_aligner.smart.types import
|
||||
TokenAlignerStepAux,
|
||||
TokenLocator,
|
||||
)
|
||||
from sglang.srt.debug_utils.comparator.dims import TokenLayout
|
||||
from sglang.srt.debug_utils.comparator.dims_spec import TokenLayout
|
||||
from sglang.srt.debug_utils.comparator.utils import Pair
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
|
||||
@@ -19,7 +19,7 @@ from sglang.srt.debug_utils.comparator.aligner.unsharder.types import (
|
||||
ReduceSumParams,
|
||||
UnsharderPlan,
|
||||
)
|
||||
from sglang.srt.debug_utils.comparator.dims import (
|
||||
from sglang.srt.debug_utils.comparator.dims_spec import (
|
||||
DimSpec,
|
||||
ParallelAxis,
|
||||
parse_dims,
|
||||
@@ -307,13 +307,16 @@ class TestPickOperation:
|
||||
def test_pick_single_group(self) -> None:
|
||||
"""PickParams picks the first tensor from a single group."""
|
||||
tensor = torch.randn(4, 8)
|
||||
dim_specs = parse_dims("h d").dims
|
||||
dim_specs = parse_dims("h d # tp:replicated").dims
|
||||
replicated = frozenset({ParallelAxis.TP})
|
||||
parallel_infos = [
|
||||
{ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2)},
|
||||
{ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=2)},
|
||||
]
|
||||
|
||||
plans = compute_unsharder_plan(dim_specs, parallel_infos)
|
||||
plans = compute_unsharder_plan(
|
||||
dim_specs, parallel_infos, explicit_replicated_axes=replicated
|
||||
)
|
||||
assert len(plans) == 1
|
||||
assert isinstance(plans[0].params, PickParams)
|
||||
|
||||
@@ -326,7 +329,8 @@ 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] # cp:replicated").dims
|
||||
replicated = frozenset({ParallelAxis.CP})
|
||||
parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [
|
||||
{
|
||||
ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
|
||||
@@ -346,7 +350,9 @@ class TestPickOperation:
|
||||
},
|
||||
]
|
||||
|
||||
plans = compute_unsharder_plan(dim_specs, parallel_infos)
|
||||
plans = compute_unsharder_plan(
|
||||
dim_specs, parallel_infos, explicit_replicated_axes=replicated
|
||||
)
|
||||
pick_plans = [p for p in plans if isinstance(p.params, PickParams)]
|
||||
assert len(pick_plans) == 1
|
||||
assert pick_plans[0].axis == ParallelAxis.CP
|
||||
@@ -361,7 +367,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 # tp:replicated': 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,8 +384,11 @@ class TestPickOperation:
|
||||
}
|
||||
)
|
||||
|
||||
dim_specs = parse_dims("b s[cp] d").dims
|
||||
plans = compute_unsharder_plan(dim_specs, parallel_infos)
|
||||
dim_specs = parse_dims("b s[cp] d # tp:replicated").dims
|
||||
replicated = frozenset({ParallelAxis.TP})
|
||||
plans = compute_unsharder_plan(
|
||||
dim_specs, parallel_infos, explicit_replicated_axes=replicated
|
||||
)
|
||||
assert len(plans) == 2
|
||||
|
||||
current: list[torch.Tensor] = _name_tensors(tensors, dim_specs)
|
||||
@@ -391,7 +400,7 @@ class TestPickOperation:
|
||||
assert torch.allclose(current[0].rename(None), full_tensor)
|
||||
|
||||
def test_fully_replicated_e2e(self) -> None:
|
||||
"""CP2 TP2, dims='b h d': fully replicated -> 2 pick steps -> 1 tensor."""
|
||||
"""CP2 TP2, dims='b h d # cp:replicated tp:replicated': fully replicated -> 2 pick steps -> 1 tensor."""
|
||||
torch.manual_seed(42)
|
||||
full_tensor = torch.randn(4, 8, 16)
|
||||
|
||||
@@ -407,8 +416,11 @@ class TestPickOperation:
|
||||
}
|
||||
)
|
||||
|
||||
dim_specs = parse_dims("b h d").dims
|
||||
plans = compute_unsharder_plan(dim_specs, parallel_infos)
|
||||
dim_specs = parse_dims("b h d # cp:replicated tp:replicated").dims
|
||||
replicated = frozenset({ParallelAxis.CP, ParallelAxis.TP})
|
||||
plans = compute_unsharder_plan(
|
||||
dim_specs, parallel_infos, explicit_replicated_axes=replicated
|
||||
)
|
||||
assert len(plans) == 2
|
||||
assert all(isinstance(p.params, PickParams) for p in plans)
|
||||
|
||||
@@ -472,12 +484,15 @@ class TestVerifyReplicatedGroup:
|
||||
|
||||
def test_execute_returns_replicated_checks(self) -> None:
|
||||
"""execute_unsharder_plan returns replicated checks for mismatch."""
|
||||
dim_specs = parse_dims("h d").dims
|
||||
dim_specs = parse_dims("h d # tp:replicated").dims
|
||||
replicated = frozenset({ParallelAxis.TP})
|
||||
parallel_infos = [
|
||||
{ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2)},
|
||||
{ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=2)},
|
||||
]
|
||||
plans = compute_unsharder_plan(dim_specs, parallel_infos)
|
||||
plans = compute_unsharder_plan(
|
||||
dim_specs, parallel_infos, explicit_replicated_axes=replicated
|
||||
)
|
||||
|
||||
tensor_a = torch.zeros(4)
|
||||
tensor_b = torch.ones(4)
|
||||
|
||||
@@ -6,7 +6,7 @@ from sglang.srt.debug_utils.comparator.aligner.unsharder.parallel_info import (
|
||||
normalize_parallel_info,
|
||||
)
|
||||
from sglang.srt.debug_utils.comparator.aligner.unsharder.types import AxisInfo
|
||||
from sglang.srt.debug_utils.comparator.dims import ParallelAxis
|
||||
from sglang.srt.debug_utils.comparator.dims_spec import ParallelAxis
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
register_cpu_ci(est_time=10, suite="default", nightly=True)
|
||||
|
||||
@@ -11,7 +11,7 @@ from sglang.srt.debug_utils.comparator.aligner.unsharder.types import (
|
||||
PickParams,
|
||||
ReduceSumParams,
|
||||
)
|
||||
from sglang.srt.debug_utils.comparator.dims import ParallelAxis, parse_dims
|
||||
from sglang.srt.debug_utils.comparator.dims_spec import ParallelAxis, parse_dims
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
register_cpu_ci(est_time=10, suite="default", nightly=True)
|
||||
@@ -40,12 +40,11 @@ class TestComputeUnsharderPlan:
|
||||
compute_unsharder_plan(dim_specs, parallel_infos)
|
||||
|
||||
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."""
|
||||
"""Axis in dims but absent from all parallel_infos -> axis_size=1, auto-skip.
|
||||
But CP is active and undeclared → raises undeclared error."""
|
||||
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
|
||||
with pytest.raises(ValueError, match="axis_rank coverage"):
|
||||
with pytest.raises(ValueError, match="not declared"):
|
||||
compute_unsharder_plan(dim_specs, parallel_infos)
|
||||
|
||||
def test_empty_parallel_infos_raises(self) -> None:
|
||||
@@ -394,9 +393,208 @@ class TestComputeUnsharderPlan:
|
||||
compute_unsharder_plan(dim_specs, parallel_infos)
|
||||
|
||||
|
||||
class TestReplicatedAxes:
|
||||
class TestExplicitReplicatedAxes:
|
||||
def test_replicated_tp_with_sharded_cp(self) -> None:
|
||||
"""CP2 TP2, dims='b s[cp] d' → PickPlan(TP) + ConcatPlan(CP)."""
|
||||
"""CP2 TP2, dims='b s[cp] d # tp:replicated' → PickPlan(TP) + ConcatPlan(CP)."""
|
||||
dim_specs = parse_dims("b s[cp] d # tp:replicated").dims
|
||||
replicated = frozenset({ParallelAxis.TP})
|
||||
parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [
|
||||
{
|
||||
ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
|
||||
ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2),
|
||||
},
|
||||
{
|
||||
ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
|
||||
ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=2),
|
||||
},
|
||||
{
|
||||
ParallelAxis.CP: AxisInfo(axis_rank=1, axis_size=2),
|
||||
ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2),
|
||||
},
|
||||
{
|
||||
ParallelAxis.CP: AxisInfo(axis_rank=1, axis_size=2),
|
||||
ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=2),
|
||||
},
|
||||
]
|
||||
plans = compute_unsharder_plan(
|
||||
dim_specs, parallel_infos, explicit_replicated_axes=replicated
|
||||
)
|
||||
|
||||
assert len(plans) == 2
|
||||
assert plans[0].axis == ParallelAxis.TP
|
||||
assert isinstance(plans[0].params, PickParams)
|
||||
assert len(plans[0].groups) == 2
|
||||
for group in plans[0].groups:
|
||||
assert len(group) == 2
|
||||
|
||||
assert plans[1].axis == ParallelAxis.CP
|
||||
assert isinstance(plans[1].params, ConcatParams)
|
||||
assert plans[1].params.dim_name == "s"
|
||||
|
||||
def test_fully_replicated(self) -> None:
|
||||
"""CP2 TP2, dims='b h d # cp:replicated tp:replicated' → PickPlan(CP) + PickPlan(TP)."""
|
||||
dim_specs = parse_dims("b h d # cp:replicated tp:replicated").dims
|
||||
replicated = frozenset({ParallelAxis.CP, ParallelAxis.TP})
|
||||
parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [
|
||||
{
|
||||
ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
|
||||
ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2),
|
||||
},
|
||||
{
|
||||
ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
|
||||
ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=2),
|
||||
},
|
||||
{
|
||||
ParallelAxis.CP: AxisInfo(axis_rank=1, axis_size=2),
|
||||
ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2),
|
||||
},
|
||||
{
|
||||
ParallelAxis.CP: AxisInfo(axis_rank=1, axis_size=2),
|
||||
ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=2),
|
||||
},
|
||||
]
|
||||
plans = compute_unsharder_plan(
|
||||
dim_specs, parallel_infos, explicit_replicated_axes=replicated
|
||||
)
|
||||
|
||||
assert len(plans) == 2
|
||||
assert all(isinstance(p.params, PickParams) for p in plans)
|
||||
axes = {p.axis for p in plans}
|
||||
assert axes == {ParallelAxis.CP, ParallelAxis.TP}
|
||||
|
||||
def test_multiple_replicated_one_sharded(self) -> None:
|
||||
"""CP2 TP2 EP2, dims='h[tp] # cp:replicated ep:replicated'."""
|
||||
dim_specs = parse_dims("h[tp] # cp:replicated ep:replicated").dims
|
||||
replicated = frozenset({ParallelAxis.CP, ParallelAxis.EP})
|
||||
parallel_infos: list[dict[ParallelAxis, AxisInfo]] = []
|
||||
for cp_rank in range(2):
|
||||
for ep_rank in range(2):
|
||||
for tp_rank in range(2):
|
||||
parallel_infos.append(
|
||||
{
|
||||
ParallelAxis.CP: AxisInfo(axis_rank=cp_rank, axis_size=2),
|
||||
ParallelAxis.EP: AxisInfo(axis_rank=ep_rank, axis_size=2),
|
||||
ParallelAxis.TP: AxisInfo(axis_rank=tp_rank, axis_size=2),
|
||||
}
|
||||
)
|
||||
|
||||
plans = compute_unsharder_plan(
|
||||
dim_specs, parallel_infos, explicit_replicated_axes=replicated
|
||||
)
|
||||
|
||||
assert len(plans) == 3
|
||||
pick_plans = [p for p in plans if isinstance(p.params, PickParams)]
|
||||
concat_plans = [p for p in plans if isinstance(p.params, ConcatParams)]
|
||||
assert len(pick_plans) == 2
|
||||
assert len(concat_plans) == 1
|
||||
assert concat_plans[0].axis == ParallelAxis.TP
|
||||
|
||||
replicated_axes_in_plan = {p.axis for p in pick_plans}
|
||||
assert replicated_axes_in_plan == {ParallelAxis.CP, ParallelAxis.EP}
|
||||
|
||||
def test_replicated_scrambled_ranks(self) -> None:
|
||||
"""Scrambled world_rank order with explicit replicated axis."""
|
||||
dim_specs = parse_dims("h[tp] # cp:replicated").dims
|
||||
replicated = frozenset({ParallelAxis.CP})
|
||||
parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [
|
||||
{
|
||||
ParallelAxis.CP: AxisInfo(axis_rank=1, axis_size=2),
|
||||
ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2),
|
||||
},
|
||||
{
|
||||
ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
|
||||
ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=2),
|
||||
},
|
||||
{
|
||||
ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
|
||||
ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2),
|
||||
},
|
||||
{
|
||||
ParallelAxis.CP: AxisInfo(axis_rank=1, axis_size=2),
|
||||
ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=2),
|
||||
},
|
||||
]
|
||||
plans = compute_unsharder_plan(
|
||||
dim_specs, parallel_infos, explicit_replicated_axes=replicated
|
||||
)
|
||||
|
||||
assert len(plans) == 2
|
||||
assert plans[0].axis == ParallelAxis.CP
|
||||
assert isinstance(plans[0].params, PickParams)
|
||||
assert plans[1].axis == ParallelAxis.TP
|
||||
assert isinstance(plans[1].params, ConcatParams)
|
||||
|
||||
def test_replicated_axis_inconsistent_size_raises(self) -> None:
|
||||
"""Replicated axis with inconsistent sizes raises ValueError."""
|
||||
dim_specs = parse_dims("h[tp] # cp:replicated").dims
|
||||
replicated = frozenset({ParallelAxis.CP})
|
||||
parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [
|
||||
{
|
||||
ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
|
||||
ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2),
|
||||
},
|
||||
{
|
||||
ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=4),
|
||||
ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=2),
|
||||
},
|
||||
]
|
||||
with pytest.raises(ValueError, match="Inconsistent axis_size"):
|
||||
compute_unsharder_plan(
|
||||
dim_specs, parallel_infos, explicit_replicated_axes=replicated
|
||||
)
|
||||
|
||||
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] # cp:replicated").dims
|
||||
replicated = frozenset({ParallelAxis.CP})
|
||||
parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [
|
||||
{
|
||||
ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
|
||||
ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2),
|
||||
},
|
||||
{
|
||||
# missing CP — replicated axis absent from this rank
|
||||
ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=2),
|
||||
},
|
||||
]
|
||||
with pytest.raises(ValueError, match="missing parallel_info"):
|
||||
compute_unsharder_plan(
|
||||
dim_specs, parallel_infos, explicit_replicated_axes=replicated
|
||||
)
|
||||
|
||||
def test_recompute_pseudo_auto_replicated(self) -> None:
|
||||
"""RECOMPUTE_PSEUDO is auto-replicated without explicit declaration."""
|
||||
dim_specs = parse_dims("h d").dims
|
||||
parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [
|
||||
{ParallelAxis.RECOMPUTE_PSEUDO: AxisInfo(axis_rank=0, axis_size=2)},
|
||||
{ParallelAxis.RECOMPUTE_PSEUDO: AxisInfo(axis_rank=1, axis_size=2)},
|
||||
]
|
||||
plans = compute_unsharder_plan(dim_specs, parallel_infos)
|
||||
|
||||
assert len(plans) == 1
|
||||
assert plans[0].axis == ParallelAxis.RECOMPUTE_PSEUDO
|
||||
assert isinstance(plans[0].params, PickParams)
|
||||
assert plans[0].groups == [[0, 1]]
|
||||
|
||||
def test_recompute_pseudo_explicit_replicated_also_works(self) -> None:
|
||||
"""RECOMPUTE_PSEUDO with explicit # recompute_pseudo:replicated also works."""
|
||||
dim_specs = parse_dims("h d # recompute_pseudo:replicated").dims
|
||||
replicated = frozenset({ParallelAxis.RECOMPUTE_PSEUDO})
|
||||
parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [
|
||||
{ParallelAxis.RECOMPUTE_PSEUDO: AxisInfo(axis_rank=0, axis_size=2)},
|
||||
{ParallelAxis.RECOMPUTE_PSEUDO: AxisInfo(axis_rank=1, axis_size=2)},
|
||||
]
|
||||
plans = compute_unsharder_plan(
|
||||
dim_specs, parallel_infos, explicit_replicated_axes=replicated
|
||||
)
|
||||
|
||||
assert len(plans) == 1
|
||||
assert plans[0].axis == ParallelAxis.RECOMPUTE_PSEUDO
|
||||
assert isinstance(plans[0].params, PickParams)
|
||||
assert plans[0].groups == [[0, 1]]
|
||||
|
||||
def test_undeclared_active_axis_raises(self) -> None:
|
||||
"""Active axis not declared as sharded or replicated raises ValueError."""
|
||||
dim_specs = parse_dims("b s[cp] d").dims
|
||||
parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [
|
||||
{
|
||||
@@ -416,149 +614,88 @@ class TestReplicatedAxes:
|
||||
ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=2),
|
||||
},
|
||||
]
|
||||
plans = compute_unsharder_plan(dim_specs, parallel_infos)
|
||||
|
||||
assert len(plans) == 2
|
||||
assert plans[0].axis == ParallelAxis.TP
|
||||
assert isinstance(plans[0].params, PickParams)
|
||||
assert len(plans[0].groups) == 2
|
||||
for group in plans[0].groups:
|
||||
assert len(group) == 2
|
||||
|
||||
assert plans[1].axis == ParallelAxis.CP
|
||||
assert isinstance(plans[1].params, ConcatParams)
|
||||
assert plans[1].params.dim_name == "s"
|
||||
|
||||
def test_fully_replicated(self) -> None:
|
||||
"""CP2 TP2, dims='b h d' → PickPlan(CP) + PickPlan(TP)."""
|
||||
dim_specs = parse_dims("b h d").dims
|
||||
parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [
|
||||
{
|
||||
ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
|
||||
ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2),
|
||||
},
|
||||
{
|
||||
ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
|
||||
ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=2),
|
||||
},
|
||||
{
|
||||
ParallelAxis.CP: AxisInfo(axis_rank=1, axis_size=2),
|
||||
ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2),
|
||||
},
|
||||
{
|
||||
ParallelAxis.CP: AxisInfo(axis_rank=1, axis_size=2),
|
||||
ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=2),
|
||||
},
|
||||
]
|
||||
plans = compute_unsharder_plan(dim_specs, parallel_infos)
|
||||
|
||||
assert len(plans) == 2
|
||||
assert all(isinstance(p.params, PickParams) for p in plans)
|
||||
axes = {p.axis for p in plans}
|
||||
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
|
||||
parallel_infos: list[dict[ParallelAxis, AxisInfo]] = []
|
||||
for cp_rank in range(2):
|
||||
for ep_rank in range(2):
|
||||
for tp_rank in range(2):
|
||||
parallel_infos.append(
|
||||
{
|
||||
ParallelAxis.CP: AxisInfo(axis_rank=cp_rank, axis_size=2),
|
||||
ParallelAxis.EP: AxisInfo(axis_rank=ep_rank, axis_size=2),
|
||||
ParallelAxis.TP: AxisInfo(axis_rank=tp_rank, axis_size=2),
|
||||
}
|
||||
)
|
||||
|
||||
plans = compute_unsharder_plan(dim_specs, parallel_infos)
|
||||
|
||||
assert len(plans) == 3
|
||||
pick_plans = [p for p in plans if isinstance(p.params, PickParams)]
|
||||
concat_plans = [p for p in plans if isinstance(p.params, ConcatParams)]
|
||||
assert len(pick_plans) == 2
|
||||
assert len(concat_plans) == 1
|
||||
assert concat_plans[0].axis == ParallelAxis.TP
|
||||
|
||||
replicated_axes = {p.axis for p in pick_plans}
|
||||
assert replicated_axes == {ParallelAxis.CP, ParallelAxis.EP}
|
||||
|
||||
def test_replicated_scrambled_ranks(self) -> None:
|
||||
"""Scrambled world_rank order with replicated axis."""
|
||||
dim_specs = parse_dims("h[tp]").dims
|
||||
parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [
|
||||
{
|
||||
ParallelAxis.CP: AxisInfo(axis_rank=1, axis_size=2),
|
||||
ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2),
|
||||
},
|
||||
{
|
||||
ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
|
||||
ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=2),
|
||||
},
|
||||
{
|
||||
ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
|
||||
ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2),
|
||||
},
|
||||
{
|
||||
ParallelAxis.CP: AxisInfo(axis_rank=1, axis_size=2),
|
||||
ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=2),
|
||||
},
|
||||
]
|
||||
plans = compute_unsharder_plan(dim_specs, parallel_infos)
|
||||
|
||||
assert len(plans) == 2
|
||||
assert plans[0].axis == ParallelAxis.CP
|
||||
assert isinstance(plans[0].params, PickParams)
|
||||
assert plans[1].axis == ParallelAxis.TP
|
||||
assert isinstance(plans[1].params, ConcatParams)
|
||||
|
||||
def test_replicated_axis_inconsistent_size_raises(self) -> None:
|
||||
"""Replicated axis with inconsistent sizes raises ValueError."""
|
||||
dim_specs = parse_dims("h[tp]").dims
|
||||
parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [
|
||||
{
|
||||
ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
|
||||
ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2),
|
||||
},
|
||||
{
|
||||
ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=4),
|
||||
ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=2),
|
||||
},
|
||||
]
|
||||
with pytest.raises(ValueError, match="Inconsistent axis_size"):
|
||||
with pytest.raises(ValueError, match="tp.*not declared"):
|
||||
compute_unsharder_plan(dim_specs, parallel_infos)
|
||||
|
||||
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
|
||||
parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [
|
||||
{
|
||||
ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
|
||||
ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2),
|
||||
},
|
||||
{
|
||||
# missing CP — replicated axis absent from this rank
|
||||
ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=2),
|
||||
},
|
||||
def test_replicated_not_in_parallel_infos_raises(self) -> None:
|
||||
"""Declaring replicated axis not in parallel_infos raises ValueError."""
|
||||
dim_specs = parse_dims("h[tp] # ep:replicated").dims
|
||||
replicated = frozenset({ParallelAxis.EP})
|
||||
parallel_infos = [
|
||||
{ParallelAxis.TP: AxisInfo(axis_rank=i, axis_size=2)} for i in range(2)
|
||||
]
|
||||
with pytest.raises(ValueError, match="missing parallel_info"):
|
||||
compute_unsharder_plan(dim_specs, parallel_infos)
|
||||
with pytest.raises(ValueError, match="not found in parallel_infos"):
|
||||
compute_unsharder_plan(
|
||||
dim_specs, parallel_infos, explicit_replicated_axes=replicated
|
||||
)
|
||||
|
||||
def test_recompute_pseudo_replicated(self) -> None:
|
||||
"""RECOMPUTE_PSEUDO with no dim annotation → replicated → PickParams."""
|
||||
dim_specs = parse_dims("h d").dims
|
||||
parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [
|
||||
{ParallelAxis.RECOMPUTE_PSEUDO: AxisInfo(axis_rank=0, axis_size=2)},
|
||||
{ParallelAxis.RECOMPUTE_PSEUDO: AxisInfo(axis_rank=1, axis_size=2)},
|
||||
def test_explicit_replicated_conflicts_with_sharded_raises(self) -> None:
|
||||
"""Planner-level defense: replicated overlaps sharded → ValueError."""
|
||||
dim_specs = parse_dims("h[tp]").dims
|
||||
replicated = frozenset({ParallelAxis.TP})
|
||||
parallel_infos = [
|
||||
{ParallelAxis.TP: AxisInfo(axis_rank=i, axis_size=2)} for i in range(2)
|
||||
]
|
||||
with pytest.raises(ValueError, match="both sharded and replicated"):
|
||||
compute_unsharder_plan(
|
||||
dim_specs, parallel_infos, explicit_replicated_axes=replicated
|
||||
)
|
||||
|
||||
|
||||
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.RECOMPUTE_PSEUDO
|
||||
assert isinstance(plans[0].params, PickParams)
|
||||
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 modifier + explicit replicated TP → PickParams."""
|
||||
dim_specs = parse_dims("t (a*b) # tp:replicated").dims
|
||||
replicated = frozenset({ParallelAxis.TP})
|
||||
parallel_infos = [
|
||||
{ParallelAxis.TP: AxisInfo(axis_rank=i, axis_size=2)} for i in range(2)
|
||||
]
|
||||
plans = compute_unsharder_plan(
|
||||
dim_specs, parallel_infos, explicit_replicated_axes=replicated
|
||||
)
|
||||
|
||||
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)
|
||||
|
||||
|
||||
class TestComputeUnsharderPlanFusedDims:
|
||||
def test_fused_dim_tp2(self) -> None:
|
||||
|
||||
@@ -0,0 +1,151 @@
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
from sglang.srt.debug_utils.comparator.dims_spec import (
|
||||
DimSpec,
|
||||
Ordering,
|
||||
ParallelAxis,
|
||||
ParallelModifier,
|
||||
Reduction,
|
||||
parse_dim,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
register_cpu_ci(est_time=5, suite="default", nightly=True)
|
||||
|
||||
|
||||
class TestParseDim:
|
||||
def test_plain_name(self) -> None:
|
||||
assert parse_dim("b") == DimSpec(name="b")
|
||||
|
||||
def test_parallel_axis(self) -> None:
|
||||
assert parse_dim("h[tp]") == DimSpec(
|
||||
name="h",
|
||||
parallel_modifiers=[ParallelModifier(axis=ParallelAxis.TP)],
|
||||
)
|
||||
|
||||
def test_all_parallel_axes(self) -> None:
|
||||
assert parse_dim("a[tp]").parallel_modifiers[0].axis == ParallelAxis.TP
|
||||
assert parse_dim("a[cp]").parallel_modifiers[0].axis == ParallelAxis.CP
|
||||
assert parse_dim("a[ep]").parallel_modifiers[0].axis == ParallelAxis.EP
|
||||
assert parse_dim("a[sp]").parallel_modifiers[0].axis == ParallelAxis.SP
|
||||
|
||||
def test_ordering(self) -> None:
|
||||
assert (
|
||||
parse_dim("s[cp:zigzag]").parallel_modifiers[0].ordering == Ordering.ZIGZAG
|
||||
)
|
||||
assert (
|
||||
parse_dim("s[cp:natural]").parallel_modifiers[0].ordering
|
||||
== Ordering.NATURAL
|
||||
)
|
||||
|
||||
def test_reduction(self) -> None:
|
||||
assert (
|
||||
parse_dim("h[tp:partial]").parallel_modifiers[0].reduction
|
||||
== Reduction.PARTIAL
|
||||
)
|
||||
|
||||
def test_all_qualifiers(self) -> None:
|
||||
assert parse_dim("s[cp:zigzag+partial]") == DimSpec(
|
||||
name="s",
|
||||
parallel_modifiers=[
|
||||
ParallelModifier(
|
||||
axis=ParallelAxis.CP,
|
||||
ordering=Ordering.ZIGZAG,
|
||||
reduction=Reduction.PARTIAL,
|
||||
),
|
||||
],
|
||||
)
|
||||
|
||||
def test_multi_axis(self) -> None:
|
||||
result: DimSpec = parse_dim("t[cp:zigzag,sp]")
|
||||
assert result.name == "t"
|
||||
assert len(result.parallel_modifiers) == 2
|
||||
assert result.parallel_modifiers[0] == ParallelModifier(
|
||||
axis=ParallelAxis.CP, ordering=Ordering.ZIGZAG
|
||||
)
|
||||
assert result.parallel_modifiers[1] == ParallelModifier(axis=ParallelAxis.SP)
|
||||
|
||||
def test_invalid_token_raises(self) -> None:
|
||||
with pytest.raises(ValueError, match="Invalid dim token"):
|
||||
parse_dim("h[]")
|
||||
with pytest.raises(ValueError, match="Invalid dim token"):
|
||||
parse_dim("h[tp[x]]")
|
||||
|
||||
def test_unknown_axis_raises(self) -> None:
|
||||
with pytest.raises(ValueError, match="Unknown axis"):
|
||||
parse_dim("h[xyz]")
|
||||
|
||||
def test_unknown_qualifier_raises(self) -> None:
|
||||
with pytest.raises(ValueError, match="Unknown qualifier"):
|
||||
parse_dim("h[tp:foobar]")
|
||||
|
||||
def test_multiple_ordering_raises(self) -> None:
|
||||
with pytest.raises(ValueError, match="Multiple ordering"):
|
||||
parse_dim("s[cp:zigzag+natural]")
|
||||
|
||||
def test_multiple_reduction_raises(self) -> None:
|
||||
with pytest.raises(ValueError, match="Multiple reduction"):
|
||||
parse_dim("h[tp:partial+partial]")
|
||||
|
||||
def test_duplicate_axis_raises(self) -> None:
|
||||
with pytest.raises(ValueError, match="Duplicate axis"):
|
||||
parse_dim("h[tp,tp]")
|
||||
|
||||
def test_squeeze_dim(self) -> None:
|
||||
assert parse_dim("1") == DimSpec(name="1")
|
||||
|
||||
def test_squeeze_dim_rejects_modifiers(self) -> None:
|
||||
with pytest.raises(ValueError, match="Invalid dim token"):
|
||||
parse_dim("1[tp]")
|
||||
|
||||
|
||||
class TestParseFusedDim:
|
||||
def test_basic_fused(self) -> None:
|
||||
result: DimSpec = parse_dim("(num_heads*head_dim)")
|
||||
assert result.name == "num_heads*head_dim"
|
||||
assert result.parallel_modifiers == []
|
||||
assert result.is_fused
|
||||
assert result.sub_dims == ["num_heads", "head_dim"]
|
||||
|
||||
def test_fused_with_modifier(self) -> None:
|
||||
result: DimSpec = parse_dim("(num_heads*head_dim)[tp]")
|
||||
assert result.name == "num_heads*head_dim"
|
||||
assert result.parallel_modifiers == [ParallelModifier(axis=ParallelAxis.TP)]
|
||||
assert result.sub_dims == ["num_heads", "head_dim"]
|
||||
|
||||
def test_three_way_fused(self) -> None:
|
||||
result: DimSpec = parse_dim("(a*b*c)")
|
||||
assert result.name == "a*b*c"
|
||||
assert len(result.sub_dims) == 3
|
||||
assert result.sub_dims == ["a", "b", "c"]
|
||||
|
||||
def test_three_way_fused_with_modifier(self) -> None:
|
||||
result: DimSpec = parse_dim("(a*b*c)[tp]")
|
||||
assert result.parallel_modifiers == [ParallelModifier(axis=ParallelAxis.TP)]
|
||||
assert len(result.sub_dims) == 3
|
||||
|
||||
def test_fused_with_complex_modifier(self) -> None:
|
||||
result: DimSpec = parse_dim("(a*b)[cp:zigzag]")
|
||||
assert result.parallel_modifiers == [
|
||||
ParallelModifier(axis=ParallelAxis.CP, ordering=Ordering.ZIGZAG)
|
||||
]
|
||||
assert result.sub_dims == ["a", "b"]
|
||||
|
||||
def test_regular_dim_not_fused(self) -> None:
|
||||
result: DimSpec = parse_dim("h[tp]")
|
||||
assert not result.is_fused
|
||||
assert result.sub_dims == ["h"]
|
||||
|
||||
def test_fused_duplicate_sub_names_raises(self) -> None:
|
||||
with pytest.raises(ValueError, match="Duplicate sub-dim"):
|
||||
parse_dim("(a*a)")
|
||||
|
||||
def test_fused_invalid_sub_dim_raises(self) -> None:
|
||||
with pytest.raises(ValueError, match="Invalid sub-dim"):
|
||||
parse_dim("(a*1)")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__]))
|
||||
@@ -0,0 +1,294 @@
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
from sglang.srt.debug_utils.comparator.dims_spec import (
|
||||
SQUEEZE_DIM_NAME,
|
||||
DimSpec,
|
||||
DimsSpec,
|
||||
Ordering,
|
||||
ParallelAxis,
|
||||
ParallelModifier,
|
||||
_SingletonDimUtil,
|
||||
parse_dims,
|
||||
resolve_dim_names,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
register_cpu_ci(est_time=5, suite="default", nightly=True)
|
||||
|
||||
|
||||
class TestSingletonDimUtilFilterOut:
|
||||
def test_no_squeeze(self) -> None:
|
||||
specs: list[DimSpec] = parse_dims("t h d").dims
|
||||
assert _SingletonDimUtil.filter_out(specs) == specs
|
||||
|
||||
def test_with_squeeze(self) -> None:
|
||||
specs: list[DimSpec] = parse_dims("t 1 h").dims
|
||||
filtered: list[DimSpec] = _SingletonDimUtil.filter_out(specs)
|
||||
assert len(filtered) == 2
|
||||
assert filtered[0].name == "t"
|
||||
assert filtered[1].name == "h"
|
||||
|
||||
def test_all_squeeze(self) -> None:
|
||||
specs: list[DimSpec] = parse_dims("1 1").dims
|
||||
assert _SingletonDimUtil.filter_out(specs) == []
|
||||
|
||||
|
||||
class TestSingletonDimUtilIsSqueeze:
|
||||
def test_squeeze(self) -> None:
|
||||
assert _SingletonDimUtil.is_squeeze(DimSpec(name=SQUEEZE_DIM_NAME)) is True
|
||||
|
||||
def test_non_squeeze(self) -> None:
|
||||
assert _SingletonDimUtil.is_squeeze(DimSpec(name="t")) is False
|
||||
|
||||
|
||||
class TestSingletonDimUtilMakeName:
|
||||
def test_indices(self) -> None:
|
||||
assert _SingletonDimUtil.make_name(0) == "singleton0"
|
||||
assert _SingletonDimUtil.make_name(1) == "singleton1"
|
||||
assert _SingletonDimUtil.make_name(99) == "singleton99"
|
||||
|
||||
|
||||
class TestSingletonDimUtilSanitizeNames:
|
||||
def test_no_squeeze(self) -> None:
|
||||
assert _SingletonDimUtil.sanitize_names(["t", "h", "d"]) == ["t", "h", "d"]
|
||||
|
||||
def test_single_squeeze(self) -> None:
|
||||
assert _SingletonDimUtil.sanitize_names(["t", "1", "h"]) == [
|
||||
"t",
|
||||
"singleton0",
|
||||
"h",
|
||||
]
|
||||
|
||||
def test_multiple_squeeze(self) -> None:
|
||||
assert _SingletonDimUtil.sanitize_names(["1", "t", "1", "h"]) == [
|
||||
"singleton0",
|
||||
"t",
|
||||
"singleton1",
|
||||
"h",
|
||||
]
|
||||
|
||||
def test_empty(self) -> None:
|
||||
assert _SingletonDimUtil.sanitize_names([]) == []
|
||||
|
||||
|
||||
class TestParseDims:
|
||||
def test_multi_dims(self) -> None:
|
||||
assert parse_dims("b s h d").dims == [
|
||||
DimSpec(name="b"),
|
||||
DimSpec(name="s"),
|
||||
DimSpec(name="h"),
|
||||
DimSpec(name="d"),
|
||||
]
|
||||
|
||||
def test_single_dim(self) -> None:
|
||||
assert parse_dims("t").dims == [DimSpec(name="t")]
|
||||
|
||||
def test_mixed_annotated(self) -> None:
|
||||
assert parse_dims("b s[cp:zigzag] h[tp] d").dims == [
|
||||
DimSpec(name="b"),
|
||||
DimSpec(
|
||||
name="s",
|
||||
parallel_modifiers=[
|
||||
ParallelModifier(axis=ParallelAxis.CP, ordering=Ordering.ZIGZAG),
|
||||
],
|
||||
),
|
||||
DimSpec(
|
||||
name="h",
|
||||
parallel_modifiers=[ParallelModifier(axis=ParallelAxis.TP)],
|
||||
),
|
||||
DimSpec(name="d"),
|
||||
]
|
||||
|
||||
def test_empty_string_raises(self) -> None:
|
||||
with pytest.raises(ValueError, match="empty"):
|
||||
parse_dims("")
|
||||
|
||||
def test_whitespace_only_raises(self) -> None:
|
||||
with pytest.raises(ValueError, match="empty"):
|
||||
parse_dims(" ")
|
||||
|
||||
def test_duplicate_name_raises(self) -> None:
|
||||
with pytest.raises(ValueError, match="Duplicate"):
|
||||
parse_dims("h h")
|
||||
|
||||
def test_with_squeeze_dims(self) -> None:
|
||||
dims: list[DimSpec] = parse_dims("t 1 h").dims
|
||||
assert len(dims) == 3
|
||||
assert dims[0] == DimSpec(name="t")
|
||||
assert dims[1] == DimSpec(name="1")
|
||||
assert dims[2] == DimSpec(name="h")
|
||||
|
||||
def test_multiple_squeeze_dims_no_duplicate_error(self) -> None:
|
||||
dims: list[DimSpec] = parse_dims("t 1 h 1 d").dims
|
||||
assert len(dims) == 5
|
||||
assert dims[1] == DimSpec(name="1")
|
||||
assert dims[3] == DimSpec(name="1")
|
||||
|
||||
|
||||
class TestParseDimsWithFused:
|
||||
def test_fused_in_dims(self) -> None:
|
||||
result: DimsSpec = parse_dims("t (num_heads*head_dim)[tp]")
|
||||
assert len(result.dims) == 2
|
||||
assert result.dims[0] == DimSpec(name="t")
|
||||
assert result.dims[1].is_fused
|
||||
assert result.dims[1].name == "num_heads*head_dim"
|
||||
|
||||
def test_fused_and_regular_mixed(self) -> None:
|
||||
result: DimsSpec = parse_dims("t (num_heads*head_dim)[tp] d")
|
||||
assert len(result.dims) == 3
|
||||
assert not result.dims[0].is_fused
|
||||
assert result.dims[1].is_fused
|
||||
assert not result.dims[2].is_fused
|
||||
|
||||
def test_fused_sub_name_conflicts_with_regular_raises(self) -> None:
|
||||
with pytest.raises(ValueError, match="Duplicate"):
|
||||
parse_dims("t num_heads (num_heads*head_dim)")
|
||||
|
||||
def test_multiple_fused_dims(self) -> None:
|
||||
result: DimsSpec = parse_dims("(a*b) (c*d)")
|
||||
assert len(result.dims) == 2
|
||||
assert result.dims[0].is_fused
|
||||
assert result.dims[1].is_fused
|
||||
|
||||
def test_cross_fused_duplicate_sub_name_raises(self) -> None:
|
||||
with pytest.raises(ValueError, match="Duplicate"):
|
||||
parse_dims("(a*b) (c*a)")
|
||||
|
||||
|
||||
class TestParseDimsWithHash:
|
||||
"""parse_dims strips the ``#`` declaration section from dims."""
|
||||
|
||||
def test_shape_dims_unchanged(self) -> None:
|
||||
assert parse_dims("b s h[tp] # dp:=moe_dp").dims == parse_dims("b s h[tp]").dims
|
||||
|
||||
def test_dp_group_alias_extracted(self) -> None:
|
||||
assert parse_dims("b s h[tp] # dp:=moe_dp").dp_group_alias == "moe_dp"
|
||||
|
||||
def test_no_hash_no_alias(self) -> None:
|
||||
assert parse_dims("b s h[tp]").dp_group_alias is None
|
||||
|
||||
def test_whitespace_around_hash(self) -> None:
|
||||
assert parse_dims("t h # dp:=foo ").dims == parse_dims("t h").dims
|
||||
assert parse_dims("t h # dp:=foo ").dp_group_alias == "foo"
|
||||
|
||||
def test_multiple_declarations_picks_dp(self) -> None:
|
||||
result: DimsSpec = parse_dims("t h[tp] # dp:=moe_dp ep:replicated")
|
||||
assert result.dims == parse_dims("t h[tp]").dims
|
||||
assert result.dp_group_alias == "moe_dp"
|
||||
assert result.replicated_axes == frozenset({ParallelAxis.EP})
|
||||
|
||||
def test_no_dp_alias_token(self) -> None:
|
||||
result: DimsSpec = parse_dims("t h[tp] # ep:replicated")
|
||||
assert result.dp_group_alias is None
|
||||
assert result.replicated_axes == frozenset({ParallelAxis.EP})
|
||||
|
||||
|
||||
class TestDpGroupAlias:
|
||||
def test_basic(self) -> None:
|
||||
assert parse_dims("b s h[tp] # dp:=moe_dp").dp_group_alias == "moe_dp"
|
||||
|
||||
def test_no_hash_returns_none(self) -> None:
|
||||
assert parse_dims("t h").dp_group_alias is None
|
||||
|
||||
def test_no_dp_alias_token(self) -> None:
|
||||
assert parse_dims("t h[tp] # ep:replicated").dp_group_alias is None
|
||||
|
||||
def test_multiple_tokens_picks_dp(self) -> None:
|
||||
assert (
|
||||
parse_dims("b s # ep:replicated dp:=custom_dp").dp_group_alias
|
||||
== "custom_dp"
|
||||
)
|
||||
|
||||
|
||||
class TestExplicitReplicatedAxes:
|
||||
def test_single_replicated(self) -> None:
|
||||
result: DimsSpec = parse_dims("b s h[tp] d # ep:replicated")
|
||||
assert result.replicated_axes == frozenset({ParallelAxis.EP})
|
||||
|
||||
def test_explicit_sharded_equivalent(self) -> None:
|
||||
assert parse_dims("b s h[tp:sharded] d").dims == parse_dims("b s h[tp] d").dims
|
||||
|
||||
def test_multiple_replicated(self) -> None:
|
||||
result: DimsSpec = parse_dims("b s h[tp] d # ep:replicated cp:replicated")
|
||||
assert result.replicated_axes == frozenset({ParallelAxis.EP, ParallelAxis.CP})
|
||||
|
||||
def test_dp_alias_and_replicated_coexist(self) -> None:
|
||||
result: DimsSpec = parse_dims("b s h[tp] d # dp:=moe_dp ep:replicated")
|
||||
assert result.dp_group_alias == "moe_dp"
|
||||
assert result.replicated_axes == frozenset({ParallelAxis.EP})
|
||||
|
||||
def test_no_hash_replicated_empty(self) -> None:
|
||||
result: DimsSpec = parse_dims("b s h[tp] d")
|
||||
assert result.replicated_axes == frozenset()
|
||||
|
||||
def test_hash_without_replicated(self) -> None:
|
||||
result: DimsSpec = parse_dims("b s h[tp] d # dp:=moe_dp")
|
||||
assert result.replicated_axes == frozenset()
|
||||
|
||||
def test_replicated_conflicts_with_sharded_raises(self) -> None:
|
||||
with pytest.raises(ValueError, match="both sharded.*and replicated"):
|
||||
parse_dims("b s h[tp] d # tp:replicated")
|
||||
|
||||
def test_unknown_axis_in_replicated_raises(self) -> None:
|
||||
with pytest.raises(ValueError, match="Unknown axis"):
|
||||
parse_dims("b s h[tp] d # xyz:replicated")
|
||||
|
||||
def test_duplicate_replicated_declaration_raises(self) -> None:
|
||||
with pytest.raises(ValueError, match="Duplicate replicated"):
|
||||
parse_dims("b s h d # ep:replicated ep:replicated")
|
||||
|
||||
def test_unrecognized_token_in_comment_raises(self) -> None:
|
||||
with pytest.raises(ValueError, match="Unrecognized token"):
|
||||
parse_dims("b s h[tp] d # ep:replicatd")
|
||||
|
||||
def test_duplicate_dp_alias_raises(self) -> None:
|
||||
with pytest.raises(ValueError, match="Duplicate dp alias"):
|
||||
parse_dims("b s h d # dp:=foo dp:=bar")
|
||||
|
||||
|
||||
class TestResolveDimNames:
|
||||
def test_no_squeeze(self) -> None:
|
||||
assert resolve_dim_names("t h d") == ["t", "h", "d"]
|
||||
|
||||
def test_single_squeeze(self) -> None:
|
||||
assert resolve_dim_names("t 1 h") == ["t", "singleton0", "h"]
|
||||
|
||||
def test_multiple_squeeze(self) -> None:
|
||||
assert resolve_dim_names("1 t 1 h") == [
|
||||
"singleton0",
|
||||
"t",
|
||||
"singleton1",
|
||||
"h",
|
||||
]
|
||||
|
||||
|
||||
class TestResolveDimNamesWithFused:
|
||||
def test_fused_dim_uses_triple_underscore(self) -> None:
|
||||
assert resolve_dim_names("t (num_heads*head_dim)") == [
|
||||
"t",
|
||||
"num_heads___head_dim",
|
||||
]
|
||||
|
||||
def test_fused_with_regular_dims(self) -> None:
|
||||
assert resolve_dim_names("t (num_heads*head_dim)[tp] d") == [
|
||||
"t",
|
||||
"num_heads___head_dim",
|
||||
"d",
|
||||
]
|
||||
|
||||
def test_three_way_fused(self) -> None:
|
||||
assert resolve_dim_names("(a*b*c)") == ["a___b___c"]
|
||||
|
||||
def test_fused_with_squeeze(self) -> None:
|
||||
assert resolve_dim_names("t 1 (a*b)") == ["t", "singleton0", "a___b"]
|
||||
|
||||
|
||||
class TestResolveDimNamesWithHash:
|
||||
def test_hash_stripped(self) -> None:
|
||||
assert resolve_dim_names("t h # dp:=moe_dp") == ["t", "h"]
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__]))
|
||||
@@ -0,0 +1,96 @@
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.srt.debug_utils.comparator.dims_spec import (
|
||||
DimSpec,
|
||||
apply_dim_names,
|
||||
find_dim_index,
|
||||
parse_dims,
|
||||
resolve_dim_by_name,
|
||||
strip_dim_names,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
register_cpu_ci(est_time=5, suite="default", nightly=True)
|
||||
|
||||
|
||||
class TestFindDimIndex:
|
||||
def test_found(self) -> None:
|
||||
specs: list[DimSpec] = parse_dims("b s h d").dims
|
||||
assert find_dim_index(specs, "s") == 1
|
||||
|
||||
def test_not_found(self) -> None:
|
||||
specs: list[DimSpec] = parse_dims("b s h d").dims
|
||||
assert find_dim_index(specs, "t") is None
|
||||
|
||||
def test_first_dim(self) -> None:
|
||||
specs: list[DimSpec] = parse_dims("t h d").dims
|
||||
assert find_dim_index(specs, "t") == 0
|
||||
|
||||
def test_last_dim(self) -> None:
|
||||
specs: list[DimSpec] = parse_dims("b s h d").dims
|
||||
assert find_dim_index(specs, "d") == 3
|
||||
|
||||
def test_with_modifiers(self) -> None:
|
||||
specs: list[DimSpec] = parse_dims("b s[cp:zigzag] h[tp] d").dims
|
||||
assert find_dim_index(specs, "h") == 2
|
||||
|
||||
def test_empty_list(self) -> None:
|
||||
assert find_dim_index([], "t") is None
|
||||
|
||||
|
||||
class TestResolveDimByName:
|
||||
def test_resolve_found(self) -> None:
|
||||
tensor: torch.Tensor = torch.randn(2, 3, 4).refine_names("b", "s", "h")
|
||||
assert resolve_dim_by_name(tensor, "b") == 0
|
||||
assert resolve_dim_by_name(tensor, "s") == 1
|
||||
assert resolve_dim_by_name(tensor, "h") == 2
|
||||
|
||||
def test_resolve_not_found_raises(self) -> None:
|
||||
tensor: torch.Tensor = torch.randn(2, 3).refine_names("b", "s")
|
||||
with pytest.raises(ValueError, match="not in tensor names"):
|
||||
resolve_dim_by_name(tensor, "h")
|
||||
|
||||
def test_resolve_unnamed_raises(self) -> None:
|
||||
tensor: torch.Tensor = torch.randn(2, 3)
|
||||
with pytest.raises(ValueError, match="no names"):
|
||||
resolve_dim_by_name(tensor, "b")
|
||||
|
||||
|
||||
class TestApplyDimNames:
|
||||
def test_apply(self) -> None:
|
||||
tensor: torch.Tensor = torch.randn(2, 3, 4)
|
||||
named: torch.Tensor = apply_dim_names(tensor, ["b", "s", "h"])
|
||||
assert named.names == ("b", "s", "h")
|
||||
assert named.shape == (2, 3, 4)
|
||||
|
||||
def test_apply_preserves_data(self) -> None:
|
||||
tensor: torch.Tensor = torch.randn(2, 3)
|
||||
named: torch.Tensor = apply_dim_names(tensor, ["x", "y"])
|
||||
assert torch.equal(strip_dim_names(named), tensor)
|
||||
|
||||
def test_ndim_mismatch_gives_clear_error(self) -> None:
|
||||
tensor: torch.Tensor = torch.randn(10, 1, 128)
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match=r"dims metadata mismatch.*3 dims.*shape \[10, 1, 128\].*2 names \['t', 'num_experts'\].*fix the dims string",
|
||||
):
|
||||
apply_dim_names(tensor, ["t", "num_experts"])
|
||||
|
||||
|
||||
class TestStripDimNames:
|
||||
def test_strip(self) -> None:
|
||||
tensor: torch.Tensor = torch.randn(2, 3).refine_names("a", "b")
|
||||
stripped: torch.Tensor = strip_dim_names(tensor)
|
||||
assert stripped.names == (None, None)
|
||||
|
||||
def test_strip_already_unnamed(self) -> None:
|
||||
tensor: torch.Tensor = torch.randn(2, 3)
|
||||
stripped: torch.Tensor = strip_dim_names(tensor)
|
||||
assert stripped.names == (None, None)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__]))
|
||||
@@ -0,0 +1,27 @@
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
from sglang.srt.debug_utils.comparator.dims_spec import (
|
||||
BATCH_DIM_NAME,
|
||||
SEQ_DIM_NAME,
|
||||
TOKEN_DIM_NAME,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
register_cpu_ci(est_time=5, suite="default", nightly=True)
|
||||
|
||||
|
||||
class TestDimConstants:
|
||||
def test_token_dim_name(self) -> None:
|
||||
assert TOKEN_DIM_NAME == "t"
|
||||
|
||||
def test_batch_dim_name(self) -> None:
|
||||
assert BATCH_DIM_NAME == "b"
|
||||
|
||||
def test_seq_dim_name(self) -> None:
|
||||
assert SEQ_DIM_NAME == "s"
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__]))
|
||||
@@ -1,488 +0,0 @@
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.srt.debug_utils.comparator.dims import (
|
||||
BATCH_DIM_NAME,
|
||||
SEQ_DIM_NAME,
|
||||
SQUEEZE_DIM_NAME,
|
||||
TOKEN_DIM_NAME,
|
||||
DimSpec,
|
||||
DimsSpec,
|
||||
Ordering,
|
||||
ParallelAxis,
|
||||
ParallelModifier,
|
||||
Reduction,
|
||||
_SingletonDimUtil,
|
||||
apply_dim_names,
|
||||
find_dim_index,
|
||||
parse_dim,
|
||||
parse_dims,
|
||||
resolve_dim_by_name,
|
||||
resolve_dim_names,
|
||||
strip_dim_names,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
register_cpu_ci(est_time=10, suite="default", nightly=True)
|
||||
|
||||
|
||||
class TestParseDim:
|
||||
def test_plain_name(self) -> None:
|
||||
assert parse_dim("b") == DimSpec(name="b")
|
||||
|
||||
def test_parallel_axis(self) -> None:
|
||||
assert parse_dim("h[tp]") == DimSpec(
|
||||
name="h",
|
||||
parallel_modifiers=[ParallelModifier(axis=ParallelAxis.TP)],
|
||||
)
|
||||
|
||||
def test_all_parallel_axes(self) -> None:
|
||||
assert parse_dim("a[tp]").parallel_modifiers[0].axis == ParallelAxis.TP
|
||||
assert parse_dim("a[cp]").parallel_modifiers[0].axis == ParallelAxis.CP
|
||||
assert parse_dim("a[ep]").parallel_modifiers[0].axis == ParallelAxis.EP
|
||||
assert parse_dim("a[sp]").parallel_modifiers[0].axis == ParallelAxis.SP
|
||||
|
||||
def test_ordering(self) -> None:
|
||||
assert (
|
||||
parse_dim("s[cp:zigzag]").parallel_modifiers[0].ordering == Ordering.ZIGZAG
|
||||
)
|
||||
assert (
|
||||
parse_dim("s[cp:natural]").parallel_modifiers[0].ordering
|
||||
== Ordering.NATURAL
|
||||
)
|
||||
|
||||
def test_reduction(self) -> None:
|
||||
assert (
|
||||
parse_dim("h[tp:partial]").parallel_modifiers[0].reduction
|
||||
== Reduction.PARTIAL
|
||||
)
|
||||
|
||||
def test_all_qualifiers(self) -> None:
|
||||
assert parse_dim("s[cp:zigzag+partial]") == DimSpec(
|
||||
name="s",
|
||||
parallel_modifiers=[
|
||||
ParallelModifier(
|
||||
axis=ParallelAxis.CP,
|
||||
ordering=Ordering.ZIGZAG,
|
||||
reduction=Reduction.PARTIAL,
|
||||
),
|
||||
],
|
||||
)
|
||||
|
||||
def test_multi_axis(self) -> None:
|
||||
result: DimSpec = parse_dim("t[cp:zigzag,sp]")
|
||||
assert result.name == "t"
|
||||
assert len(result.parallel_modifiers) == 2
|
||||
assert result.parallel_modifiers[0] == ParallelModifier(
|
||||
axis=ParallelAxis.CP, ordering=Ordering.ZIGZAG
|
||||
)
|
||||
assert result.parallel_modifiers[1] == ParallelModifier(axis=ParallelAxis.SP)
|
||||
|
||||
def test_invalid_token_raises(self) -> None:
|
||||
with pytest.raises(ValueError, match="Invalid dim token"):
|
||||
parse_dim("h[]")
|
||||
with pytest.raises(ValueError, match="Invalid dim token"):
|
||||
parse_dim("h[tp[x]]")
|
||||
|
||||
def test_unknown_axis_raises(self) -> None:
|
||||
with pytest.raises(ValueError, match="Unknown axis"):
|
||||
parse_dim("h[xyz]")
|
||||
|
||||
def test_unknown_qualifier_raises(self) -> None:
|
||||
with pytest.raises(ValueError, match="Unknown qualifier"):
|
||||
parse_dim("h[tp:foobar]")
|
||||
|
||||
def test_multiple_ordering_raises(self) -> None:
|
||||
with pytest.raises(ValueError, match="Multiple ordering"):
|
||||
parse_dim("s[cp:zigzag+natural]")
|
||||
|
||||
def test_multiple_reduction_raises(self) -> None:
|
||||
with pytest.raises(ValueError, match="Multiple reduction"):
|
||||
parse_dim("h[tp:partial+partial]")
|
||||
|
||||
def test_duplicate_axis_raises(self) -> None:
|
||||
with pytest.raises(ValueError, match="Duplicate axis"):
|
||||
parse_dim("h[tp,tp]")
|
||||
|
||||
def test_squeeze_dim(self) -> None:
|
||||
assert parse_dim("1") == DimSpec(name="1")
|
||||
|
||||
def test_squeeze_dim_rejects_modifiers(self) -> None:
|
||||
with pytest.raises(ValueError, match="Invalid dim token"):
|
||||
parse_dim("1[tp]")
|
||||
|
||||
def test_squeeze_dim(self) -> None:
|
||||
assert parse_dim("1") == DimSpec(name="1")
|
||||
|
||||
def test_squeeze_dim_rejects_modifiers(self) -> None:
|
||||
with pytest.raises(ValueError, match="Invalid dim token"):
|
||||
parse_dim("1(tp)")
|
||||
|
||||
def test_squeeze_dim(self) -> None:
|
||||
assert parse_dim("1") == DimSpec(name="1")
|
||||
|
||||
def test_squeeze_dim_rejects_modifiers(self) -> None:
|
||||
with pytest.raises(ValueError, match="Invalid dim token"):
|
||||
parse_dim("1(tp)")
|
||||
|
||||
|
||||
class TestParseDims:
|
||||
def test_multi_dims(self) -> None:
|
||||
assert parse_dims("b s h d").dims == [
|
||||
DimSpec(name="b"),
|
||||
DimSpec(name="s"),
|
||||
DimSpec(name="h"),
|
||||
DimSpec(name="d"),
|
||||
]
|
||||
|
||||
def test_single_dim(self) -> None:
|
||||
assert parse_dims("t").dims == [DimSpec(name="t")]
|
||||
|
||||
def test_mixed_annotated(self) -> None:
|
||||
assert parse_dims("b s[cp:zigzag] h[tp] d").dims == [
|
||||
DimSpec(name="b"),
|
||||
DimSpec(
|
||||
name="s",
|
||||
parallel_modifiers=[
|
||||
ParallelModifier(axis=ParallelAxis.CP, ordering=Ordering.ZIGZAG),
|
||||
],
|
||||
),
|
||||
DimSpec(
|
||||
name="h",
|
||||
parallel_modifiers=[ParallelModifier(axis=ParallelAxis.TP)],
|
||||
),
|
||||
DimSpec(name="d"),
|
||||
]
|
||||
|
||||
def test_empty_string_raises(self) -> None:
|
||||
with pytest.raises(ValueError, match="empty"):
|
||||
parse_dims("")
|
||||
|
||||
def test_whitespace_only_raises(self) -> None:
|
||||
with pytest.raises(ValueError, match="empty"):
|
||||
parse_dims(" ")
|
||||
|
||||
def test_duplicate_name_raises(self) -> None:
|
||||
with pytest.raises(ValueError, match="Duplicate"):
|
||||
parse_dims("h h")
|
||||
|
||||
def test_with_squeeze_dims(self) -> None:
|
||||
dims: list[DimSpec] = parse_dims("t 1 h").dims
|
||||
assert len(dims) == 3
|
||||
assert dims[0] == DimSpec(name="t")
|
||||
assert dims[1] == DimSpec(name="1")
|
||||
assert dims[2] == DimSpec(name="h")
|
||||
|
||||
def test_multiple_squeeze_dims_no_duplicate_error(self) -> None:
|
||||
dims: list[DimSpec] = parse_dims("t 1 h 1 d").dims
|
||||
assert len(dims) == 5
|
||||
assert dims[1] == DimSpec(name="1")
|
||||
assert dims[3] == DimSpec(name="1")
|
||||
|
||||
|
||||
class TestDimConstants:
|
||||
def test_token_dim_name(self) -> None:
|
||||
assert TOKEN_DIM_NAME == "t"
|
||||
|
||||
def test_batch_dim_name(self) -> None:
|
||||
assert BATCH_DIM_NAME == "b"
|
||||
|
||||
def test_seq_dim_name(self) -> None:
|
||||
assert SEQ_DIM_NAME == "s"
|
||||
|
||||
|
||||
class TestFindDimIndex:
|
||||
def test_found(self) -> None:
|
||||
specs: list[DimSpec] = parse_dims("b s h d").dims
|
||||
assert find_dim_index(specs, "s") == 1
|
||||
|
||||
def test_not_found(self) -> None:
|
||||
specs: list[DimSpec] = parse_dims("b s h d").dims
|
||||
assert find_dim_index(specs, "t") is None
|
||||
|
||||
def test_first_dim(self) -> None:
|
||||
specs: list[DimSpec] = parse_dims("t h d").dims
|
||||
assert find_dim_index(specs, "t") == 0
|
||||
|
||||
def test_last_dim(self) -> None:
|
||||
specs: list[DimSpec] = parse_dims("b s h d").dims
|
||||
assert find_dim_index(specs, "d") == 3
|
||||
|
||||
def test_with_modifiers(self) -> None:
|
||||
specs: list[DimSpec] = parse_dims("b s[cp:zigzag] h[tp] d").dims
|
||||
assert find_dim_index(specs, "h") == 2
|
||||
|
||||
def test_empty_list(self) -> None:
|
||||
assert find_dim_index([], "t") is None
|
||||
|
||||
|
||||
class TestResolveDimByName:
|
||||
def test_resolve_found(self) -> None:
|
||||
tensor: torch.Tensor = torch.randn(2, 3, 4).refine_names("b", "s", "h")
|
||||
assert resolve_dim_by_name(tensor, "b") == 0
|
||||
assert resolve_dim_by_name(tensor, "s") == 1
|
||||
assert resolve_dim_by_name(tensor, "h") == 2
|
||||
|
||||
def test_resolve_not_found_raises(self) -> None:
|
||||
tensor: torch.Tensor = torch.randn(2, 3).refine_names("b", "s")
|
||||
with pytest.raises(ValueError, match="not in tensor names"):
|
||||
resolve_dim_by_name(tensor, "h")
|
||||
|
||||
def test_resolve_unnamed_raises(self) -> None:
|
||||
tensor: torch.Tensor = torch.randn(2, 3)
|
||||
with pytest.raises(ValueError, match="no names"):
|
||||
resolve_dim_by_name(tensor, "b")
|
||||
|
||||
|
||||
class TestApplyDimNames:
|
||||
def test_apply(self) -> None:
|
||||
tensor: torch.Tensor = torch.randn(2, 3, 4)
|
||||
named: torch.Tensor = apply_dim_names(tensor, ["b", "s", "h"])
|
||||
assert named.names == ("b", "s", "h")
|
||||
assert named.shape == (2, 3, 4)
|
||||
|
||||
def test_apply_preserves_data(self) -> None:
|
||||
tensor: torch.Tensor = torch.randn(2, 3)
|
||||
named: torch.Tensor = apply_dim_names(tensor, ["x", "y"])
|
||||
assert torch.equal(strip_dim_names(named), tensor)
|
||||
|
||||
def test_ndim_mismatch_gives_clear_error(self) -> None:
|
||||
tensor: torch.Tensor = torch.randn(10, 1, 128)
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match=r"dims metadata mismatch.*3 dims.*shape \[10, 1, 128\].*2 names \['t', 'num_experts'\].*fix the dims string",
|
||||
):
|
||||
apply_dim_names(tensor, ["t", "num_experts"])
|
||||
|
||||
|
||||
class TestStripDimNames:
|
||||
def test_strip(self) -> None:
|
||||
tensor: torch.Tensor = torch.randn(2, 3).refine_names("a", "b")
|
||||
stripped: torch.Tensor = strip_dim_names(tensor)
|
||||
assert stripped.names == (None, None)
|
||||
|
||||
def test_strip_already_unnamed(self) -> None:
|
||||
tensor: torch.Tensor = torch.randn(2, 3)
|
||||
stripped: torch.Tensor = strip_dim_names(tensor)
|
||||
assert stripped.names == (None, None)
|
||||
|
||||
|
||||
class TestResolveDimNames:
|
||||
def test_no_squeeze(self) -> None:
|
||||
assert resolve_dim_names("t h d") == ["t", "h", "d"]
|
||||
|
||||
def test_single_squeeze(self) -> None:
|
||||
assert resolve_dim_names("t 1 h") == ["t", "singleton0", "h"]
|
||||
|
||||
def test_multiple_squeeze(self) -> None:
|
||||
assert resolve_dim_names("1 t 1 h") == [
|
||||
"singleton0",
|
||||
"t",
|
||||
"singleton1",
|
||||
"h",
|
||||
]
|
||||
|
||||
|
||||
class TestSingletonDimUtilFilterOut:
|
||||
def test_no_squeeze(self) -> None:
|
||||
specs: list[DimSpec] = parse_dims("t h d").dims
|
||||
assert _SingletonDimUtil.filter_out(specs) == specs
|
||||
|
||||
def test_with_squeeze(self) -> None:
|
||||
specs: list[DimSpec] = parse_dims("t 1 h").dims
|
||||
filtered: list[DimSpec] = _SingletonDimUtil.filter_out(specs)
|
||||
assert len(filtered) == 2
|
||||
assert filtered[0].name == "t"
|
||||
assert filtered[1].name == "h"
|
||||
|
||||
def test_all_squeeze(self) -> None:
|
||||
specs: list[DimSpec] = parse_dims("1 1").dims
|
||||
assert _SingletonDimUtil.filter_out(specs) == []
|
||||
|
||||
|
||||
class TestSingletonDimUtilIsSqueeze:
|
||||
def test_squeeze(self) -> None:
|
||||
assert _SingletonDimUtil.is_squeeze(DimSpec(name=SQUEEZE_DIM_NAME)) is True
|
||||
|
||||
def test_non_squeeze(self) -> None:
|
||||
assert _SingletonDimUtil.is_squeeze(DimSpec(name="t")) is False
|
||||
|
||||
|
||||
class TestSingletonDimUtilMakeName:
|
||||
def test_indices(self) -> None:
|
||||
assert _SingletonDimUtil.make_name(0) == "singleton0"
|
||||
assert _SingletonDimUtil.make_name(1) == "singleton1"
|
||||
assert _SingletonDimUtil.make_name(99) == "singleton99"
|
||||
|
||||
|
||||
class TestSingletonDimUtilSanitizeNames:
|
||||
def test_no_squeeze(self) -> None:
|
||||
assert _SingletonDimUtil.sanitize_names(["t", "h", "d"]) == ["t", "h", "d"]
|
||||
|
||||
def test_single_squeeze(self) -> None:
|
||||
assert _SingletonDimUtil.sanitize_names(["t", "1", "h"]) == [
|
||||
"t",
|
||||
"singleton0",
|
||||
"h",
|
||||
]
|
||||
|
||||
def test_multiple_squeeze(self) -> None:
|
||||
assert _SingletonDimUtil.sanitize_names(["1", "t", "1", "h"]) == [
|
||||
"singleton0",
|
||||
"t",
|
||||
"singleton1",
|
||||
"h",
|
||||
]
|
||||
|
||||
def test_empty(self) -> None:
|
||||
assert _SingletonDimUtil.sanitize_names([]) == []
|
||||
|
||||
|
||||
class TestParseDimsWithHash:
|
||||
"""parse_dims strips the ``#`` declaration section from dims."""
|
||||
|
||||
def test_shape_dims_unchanged(self) -> None:
|
||||
assert parse_dims("b s h[tp] # dp:=moe_dp").dims == parse_dims("b s h[tp]").dims
|
||||
|
||||
def test_dp_group_alias_extracted(self) -> None:
|
||||
assert parse_dims("b s h[tp] # dp:=moe_dp").dp_group_alias == "moe_dp"
|
||||
|
||||
def test_no_hash_no_alias(self) -> None:
|
||||
assert parse_dims("b s h[tp]").dp_group_alias is None
|
||||
|
||||
def test_whitespace_around_hash(self) -> None:
|
||||
assert parse_dims("t h # dp:=foo ").dims == parse_dims("t h").dims
|
||||
assert parse_dims("t h # dp:=foo ").dp_group_alias == "foo"
|
||||
|
||||
def test_multiple_declarations_picks_dp(self) -> None:
|
||||
result: DimsSpec = parse_dims("t h[tp] # dp:=moe_dp ep:replicated")
|
||||
assert result.dims == parse_dims("t h[tp]").dims
|
||||
assert result.dp_group_alias == "moe_dp"
|
||||
|
||||
def test_no_dp_alias_token(self) -> None:
|
||||
assert parse_dims("t h[tp] # ep:replicated").dp_group_alias is None
|
||||
|
||||
|
||||
class TestDpGroupAlias:
|
||||
def test_basic(self) -> None:
|
||||
assert parse_dims("b s h[tp] # dp:=moe_dp").dp_group_alias == "moe_dp"
|
||||
|
||||
def test_no_hash_returns_none(self) -> None:
|
||||
assert parse_dims("t h").dp_group_alias is None
|
||||
|
||||
def test_no_dp_alias_token(self) -> None:
|
||||
assert parse_dims("t h[tp] # ep:replicated").dp_group_alias is None
|
||||
|
||||
def test_multiple_tokens_picks_dp(self) -> None:
|
||||
assert (
|
||||
parse_dims("b s # ep:replicated dp:=custom_dp").dp_group_alias
|
||||
== "custom_dp"
|
||||
)
|
||||
|
||||
|
||||
class TestResolveDimNamesWithFused:
|
||||
def test_fused_dim_uses_triple_underscore(self) -> None:
|
||||
assert resolve_dim_names("t (num_heads*head_dim)") == [
|
||||
"t",
|
||||
"num_heads___head_dim",
|
||||
]
|
||||
|
||||
def test_fused_with_regular_dims(self) -> None:
|
||||
assert resolve_dim_names("t (num_heads*head_dim)[tp] d") == [
|
||||
"t",
|
||||
"num_heads___head_dim",
|
||||
"d",
|
||||
]
|
||||
|
||||
def test_three_way_fused(self) -> None:
|
||||
assert resolve_dim_names("(a*b*c)") == ["a___b___c"]
|
||||
|
||||
def test_fused_with_squeeze(self) -> None:
|
||||
assert resolve_dim_names("t 1 (a*b)") == ["t", "singleton0", "a___b"]
|
||||
|
||||
|
||||
class TestResolveDimNamesWithHash:
|
||||
def test_hash_stripped(self) -> None:
|
||||
assert resolve_dim_names("t h # dp:=moe_dp") == ["t", "h"]
|
||||
|
||||
|
||||
class TestParseFusedDim:
|
||||
def test_basic_fused(self) -> None:
|
||||
result: DimSpec = parse_dim("(num_heads*head_dim)")
|
||||
assert result.name == "num_heads*head_dim"
|
||||
assert result.parallel_modifiers == []
|
||||
assert result.is_fused
|
||||
assert result.sub_dims == ["num_heads", "head_dim"]
|
||||
|
||||
def test_fused_with_modifier(self) -> None:
|
||||
result: DimSpec = parse_dim("(num_heads*head_dim)[tp]")
|
||||
assert result.name == "num_heads*head_dim"
|
||||
assert result.parallel_modifiers == [ParallelModifier(axis=ParallelAxis.TP)]
|
||||
assert result.sub_dims == ["num_heads", "head_dim"]
|
||||
|
||||
def test_three_way_fused(self) -> None:
|
||||
result: DimSpec = parse_dim("(a*b*c)")
|
||||
assert result.name == "a*b*c"
|
||||
assert len(result.sub_dims) == 3
|
||||
assert result.sub_dims == ["a", "b", "c"]
|
||||
|
||||
def test_three_way_fused_with_modifier(self) -> None:
|
||||
result: DimSpec = parse_dim("(a*b*c)[tp]")
|
||||
assert result.parallel_modifiers == [ParallelModifier(axis=ParallelAxis.TP)]
|
||||
assert len(result.sub_dims) == 3
|
||||
|
||||
def test_fused_with_complex_modifier(self) -> None:
|
||||
result: DimSpec = parse_dim("(a*b)[cp:zigzag]")
|
||||
assert result.parallel_modifiers == [
|
||||
ParallelModifier(axis=ParallelAxis.CP, ordering=Ordering.ZIGZAG)
|
||||
]
|
||||
assert result.sub_dims == ["a", "b"]
|
||||
|
||||
def test_regular_dim_not_fused(self) -> None:
|
||||
result: DimSpec = parse_dim("h[tp]")
|
||||
assert not result.is_fused
|
||||
assert result.sub_dims == ["h"]
|
||||
|
||||
def test_fused_duplicate_sub_names_raises(self) -> None:
|
||||
with pytest.raises(ValueError, match="Duplicate sub-dim"):
|
||||
parse_dim("(a*a)")
|
||||
|
||||
def test_fused_invalid_sub_dim_raises(self) -> None:
|
||||
with pytest.raises(ValueError, match="Invalid sub-dim"):
|
||||
parse_dim("(a*1)")
|
||||
|
||||
|
||||
class TestParseDimsWithFused:
|
||||
def test_fused_in_dims(self) -> None:
|
||||
result: DimsSpec = parse_dims("t (num_heads*head_dim)[tp]")
|
||||
assert len(result.dims) == 2
|
||||
assert result.dims[0] == DimSpec(name="t")
|
||||
assert result.dims[1].is_fused
|
||||
assert result.dims[1].name == "num_heads*head_dim"
|
||||
|
||||
def test_fused_and_regular_mixed(self) -> None:
|
||||
result: DimsSpec = parse_dims("t (num_heads*head_dim)[tp] d")
|
||||
assert len(result.dims) == 3
|
||||
assert not result.dims[0].is_fused
|
||||
assert result.dims[1].is_fused
|
||||
assert not result.dims[2].is_fused
|
||||
|
||||
def test_fused_sub_name_conflicts_with_regular_raises(self) -> None:
|
||||
with pytest.raises(ValueError, match="Duplicate"):
|
||||
parse_dims("t num_heads (num_heads*head_dim)")
|
||||
|
||||
def test_multiple_fused_dims(self) -> None:
|
||||
result: DimsSpec = parse_dims("(a*b) (c*d)")
|
||||
assert len(result.dims) == 2
|
||||
assert result.dims[0].is_fused
|
||||
assert result.dims[1].is_fused
|
||||
|
||||
def test_cross_fused_duplicate_sub_name_raises(self) -> None:
|
||||
with pytest.raises(ValueError, match="Duplicate"):
|
||||
parse_dims("(a*b) (c*a)")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__]))
|
||||
@@ -1738,7 +1738,7 @@ class TestEntrypointReplicatedAxis:
|
||||
cp_size=2,
|
||||
tp_size=2,
|
||||
seq_dim=1,
|
||||
dims_str="b s[cp] d",
|
||||
dims_str="b s[cp] d # tp:replicated",
|
||||
)
|
||||
|
||||
argv = _make_argv(
|
||||
@@ -1777,7 +1777,7 @@ class TestEntrypointReplicatedAxis:
|
||||
cp_size=2,
|
||||
tp_size=2,
|
||||
seq_dim=1,
|
||||
dims_str="b s[cp] d",
|
||||
dims_str="b s[cp] d # tp:replicated",
|
||||
tp_noise=0.5,
|
||||
)
|
||||
|
||||
@@ -1813,7 +1813,7 @@ class TestEntrypointReplicatedAxis:
|
||||
cp_size=2,
|
||||
tp_size=2,
|
||||
seq_dim=1,
|
||||
dims_str="b s[cp] d",
|
||||
dims_str="b s[cp] d # tp:replicated",
|
||||
tp_noise=0.5,
|
||||
)
|
||||
_create_replicated_tp_sharded_cp_dumps(
|
||||
@@ -1823,7 +1823,7 @@ class TestEntrypointReplicatedAxis:
|
||||
cp_size=2,
|
||||
tp_size=2,
|
||||
seq_dim=1,
|
||||
dims_str="b s[cp] d",
|
||||
dims_str="b s[cp] d # tp:replicated",
|
||||
tp_noise=0.5,
|
||||
)
|
||||
|
||||
@@ -1862,7 +1862,7 @@ class TestEntrypointReplicatedAxis:
|
||||
rank=0,
|
||||
name="attn_out",
|
||||
tensor=torch.randn(4, 4, 6),
|
||||
dims="b s[cp] d",
|
||||
dims="b s[cp] d # tp:replicated",
|
||||
parallel_info={
|
||||
"cp_rank": 0,
|
||||
"cp_size": 2,
|
||||
@@ -1876,7 +1876,7 @@ class TestEntrypointReplicatedAxis:
|
||||
rank=1,
|
||||
name="attn_out",
|
||||
tensor=torch.randn(4, 4, 3),
|
||||
dims="b s[cp] d",
|
||||
dims="b s[cp] d # tp:replicated",
|
||||
parallel_info={
|
||||
"cp_rank": 0,
|
||||
"cp_size": 2,
|
||||
@@ -1890,7 +1890,7 @@ class TestEntrypointReplicatedAxis:
|
||||
rank=2,
|
||||
name="attn_out",
|
||||
tensor=torch.randn(4, 4, 6),
|
||||
dims="b s[cp] d",
|
||||
dims="b s[cp] d # tp:replicated",
|
||||
parallel_info={
|
||||
"cp_rank": 1,
|
||||
"cp_size": 2,
|
||||
@@ -1904,7 +1904,7 @@ class TestEntrypointReplicatedAxis:
|
||||
rank=3,
|
||||
name="attn_out",
|
||||
tensor=torch.randn(4, 4, 3),
|
||||
dims="b s[cp] d",
|
||||
dims="b s[cp] d # tp:replicated",
|
||||
parallel_info={
|
||||
"cp_rank": 1,
|
||||
"cp_size": 2,
|
||||
|
||||
@@ -20,7 +20,7 @@ from sglang.srt.debug_utils.comparator.aligner.unsharder.types import (
|
||||
ConcatParams,
|
||||
UnsharderPlan,
|
||||
)
|
||||
from sglang.srt.debug_utils.comparator.dims import ParallelAxis, TokenLayout
|
||||
from sglang.srt.debug_utils.comparator.dims_spec import ParallelAxis, TokenLayout
|
||||
from sglang.srt.debug_utils.comparator.output_types import (
|
||||
ErrorLog,
|
||||
NonTensorComparisonRecord,
|
||||
|
||||
Reference in New Issue
Block a user