diff --git a/python/sglang/srt/debug_utils/comparator/aligner/axis_aligner.py b/python/sglang/srt/debug_utils/comparator/aligner/axis_aligner.py index 3a95ebb67..34451d293 100644 --- a/python/sglang/srt/debug_utils/comparator/aligner/axis_aligner.py +++ b/python/sglang/srt/debug_utils/comparator/aligner/axis_aligner.py @@ -5,7 +5,7 @@ from typing import Optional import torch from einops import rearrange -from sglang.srt.debug_utils.comparator.dims import ( +from sglang.srt.debug_utils.comparator.dims_spec import ( _FUSED_NAME_SEP, DimSpec, _SingletonDimUtil, diff --git a/python/sglang/srt/debug_utils/comparator/aligner/entrypoint/planner.py b/python/sglang/srt/debug_utils/comparator/aligner/entrypoint/planner.py index 7cf0fbd7a..673701f90 100644 --- a/python/sglang/srt/debug_utils/comparator/aligner/entrypoint/planner.py +++ b/python/sglang/srt/debug_utils/comparator/aligner/entrypoint/planner.py @@ -23,8 +23,10 @@ from sglang.srt.debug_utils.comparator.aligner.unsharder.parallel_info import ( from sglang.srt.debug_utils.comparator.aligner.unsharder.planner import ( compute_unsharder_plan, ) -from sglang.srt.debug_utils.comparator.dims import ( +from sglang.srt.debug_utils.comparator.dims_spec import ( DimSpec, + DimsSpec, + ParallelAxis, _SingletonDimUtil, parse_dims, ) @@ -106,12 +108,15 @@ def compute_per_step_sub_plans( if dims_str is None: return [] - dim_specs: list[DimSpec] = _SingletonDimUtil.filter_out(parse_dims(dims_str).dims) + dims_spec: DimsSpec = parse_dims(dims_str) + dim_specs: list[DimSpec] = _SingletonDimUtil.filter_out(dims_spec.dims) + replicated_axes: frozenset[ParallelAxis] = dims_spec.replicated_axes parallel_infos = [normalize_parallel_info(meta) for meta in metas] unsharder_plans = compute_unsharder_plan( dim_specs=dim_specs, parallel_infos=parallel_infos, + explicit_replicated_axes=replicated_axes, thd_global_seq_lens=thd_global_seq_lens, ) reorderer_plans = compute_reorderer_plans( diff --git a/python/sglang/srt/debug_utils/comparator/aligner/reorderer/executor.py b/python/sglang/srt/debug_utils/comparator/aligner/reorderer/executor.py index a44f5edce..20b2338fe 100644 --- a/python/sglang/srt/debug_utils/comparator/aligner/reorderer/executor.py +++ b/python/sglang/srt/debug_utils/comparator/aligner/reorderer/executor.py @@ -7,7 +7,7 @@ from sglang.srt.debug_utils.comparator.aligner.reorderer.types import ( ZigzagToNaturalParams, ZigzagToNaturalThdParams, ) -from sglang.srt.debug_utils.comparator.dims import ( +from sglang.srt.debug_utils.comparator.dims_spec import ( resolve_dim_by_name, strip_dim_names, ) diff --git a/python/sglang/srt/debug_utils/comparator/aligner/reorderer/planner.py b/python/sglang/srt/debug_utils/comparator/aligner/reorderer/planner.py index c8f03362c..bbede55f5 100644 --- a/python/sglang/srt/debug_utils/comparator/aligner/reorderer/planner.py +++ b/python/sglang/srt/debug_utils/comparator/aligner/reorderer/planner.py @@ -6,7 +6,7 @@ from sglang.srt.debug_utils.comparator.aligner.reorderer.types import ( ZigzagToNaturalThdParams, ) from sglang.srt.debug_utils.comparator.aligner.unsharder.types import AxisInfo -from sglang.srt.debug_utils.comparator.dims import ( +from sglang.srt.debug_utils.comparator.dims_spec import ( SEQ_DIM_NAME, TOKEN_DIM_NAME, DimSpec, diff --git a/python/sglang/srt/debug_utils/comparator/aligner/token_aligner/concat_steps/executor.py b/python/sglang/srt/debug_utils/comparator/aligner/token_aligner/concat_steps/executor.py index 400274826..201367d5e 100644 --- a/python/sglang/srt/debug_utils/comparator/aligner/token_aligner/concat_steps/executor.py +++ b/python/sglang/srt/debug_utils/comparator/aligner/token_aligner/concat_steps/executor.py @@ -4,7 +4,7 @@ from typing import Optional import torch -from sglang.srt.debug_utils.comparator.dims import ( +from sglang.srt.debug_utils.comparator.dims_spec import ( SEQ_DIM_NAME, TOKEN_DIM_NAME, ) diff --git a/python/sglang/srt/debug_utils/comparator/aligner/token_aligner/smart/aux_loader.py b/python/sglang/srt/debug_utils/comparator/aligner/token_aligner/smart/aux_loader.py index 8650b0d71..a63c26978 100644 --- a/python/sglang/srt/debug_utils/comparator/aligner/token_aligner/smart/aux_loader.py +++ b/python/sglang/srt/debug_utils/comparator/aligner/token_aligner/smart/aux_loader.py @@ -24,7 +24,7 @@ from sglang.srt.debug_utils.comparator.aligner.token_aligner.smart.types import from sglang.srt.debug_utils.comparator.aligner.unsharder.parallel_info import ( normalize_parallel_info, ) -from sglang.srt.debug_utils.comparator.dims import ( +from sglang.srt.debug_utils.comparator.dims_spec import ( ParallelAxis, TokenLayout, apply_dim_names, diff --git a/python/sglang/srt/debug_utils/comparator/aligner/token_aligner/smart/aux_plugins.py b/python/sglang/srt/debug_utils/comparator/aligner/token_aligner/smart/aux_plugins.py index 5d0ad9ae8..8e497e9ec 100644 --- a/python/sglang/srt/debug_utils/comparator/aligner/token_aligner/smart/aux_plugins.py +++ b/python/sglang/srt/debug_utils/comparator/aligner/token_aligner/smart/aux_plugins.py @@ -11,7 +11,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.srt.debug_utils.comparator.log_sink import log_sink from sglang.srt.debug_utils.comparator.output_types import InfoLog diff --git a/python/sglang/srt/debug_utils/comparator/aligner/token_aligner/smart/executor.py b/python/sglang/srt/debug_utils/comparator/aligner/token_aligner/smart/executor.py index c8755d81d..98a4cca7d 100644 --- a/python/sglang/srt/debug_utils/comparator/aligner/token_aligner/smart/executor.py +++ b/python/sglang/srt/debug_utils/comparator/aligner/token_aligner/smart/executor.py @@ -7,7 +7,7 @@ from sglang.srt.debug_utils.comparator.aligner.token_aligner.smart.types import TokenAlignerPlan, TokenLocator, ) -from sglang.srt.debug_utils.comparator.dims import ( +from sglang.srt.debug_utils.comparator.dims_spec import ( BATCH_DIM_NAME, SEQ_DIM_NAME, TOKEN_DIM_NAME, diff --git a/python/sglang/srt/debug_utils/comparator/aligner/token_aligner/smart/types.py b/python/sglang/srt/debug_utils/comparator/aligner/token_aligner/smart/types.py index 20fde589c..116c411d5 100644 --- a/python/sglang/srt/debug_utils/comparator/aligner/token_aligner/smart/types.py +++ b/python/sglang/srt/debug_utils/comparator/aligner/token_aligner/smart/types.py @@ -5,7 +5,7 @@ from typing import NamedTuple, Optional, Union from pydantic import model_validator -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, _check_equal_lengths, 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 b5700a36d..788f36579 100644 --- a/python/sglang/srt/debug_utils/comparator/aligner/unsharder/executor.py +++ b/python/sglang/srt/debug_utils/comparator/aligner/unsharder/executor.py @@ -11,7 +11,7 @@ from sglang.srt.debug_utils.comparator.aligner.unsharder.types import ( UnsharderParams, UnsharderPlan, ) -from sglang.srt.debug_utils.comparator.dims import ( +from sglang.srt.debug_utils.comparator.dims_spec import ( ParallelAxis, resolve_dim_by_name, ) diff --git a/python/sglang/srt/debug_utils/comparator/aligner/unsharder/parallel_info.py b/python/sglang/srt/debug_utils/comparator/aligner/unsharder/parallel_info.py index 321554c82..2ff4a005d 100644 --- a/python/sglang/srt/debug_utils/comparator/aligner/unsharder/parallel_info.py +++ b/python/sglang/srt/debug_utils/comparator/aligner/unsharder/parallel_info.py @@ -1,7 +1,7 @@ from typing import Optional 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 _PARALLEL_INFO_KEYS = ("sglang_parallel_info", "megatron_parallel_info") 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 5e7fb14ce..cde2a9a59 100644 --- a/python/sglang/srt/debug_utils/comparator/aligner/unsharder/planner.py +++ b/python/sglang/srt/debug_utils/comparator/aligner/unsharder/planner.py @@ -10,7 +10,7 @@ from sglang.srt.debug_utils.comparator.aligner.unsharder.types import ( UnsharderParams, UnsharderPlan, ) -from sglang.srt.debug_utils.comparator.dims import ( +from sglang.srt.debug_utils.comparator.dims_spec import ( TOKEN_DIM_NAME, DimSpec, ParallelAxis, @@ -32,6 +32,7 @@ def compute_unsharder_plan( dim_specs: list[DimSpec], parallel_infos: list[dict[ParallelAxis, AxisInfo]], *, + explicit_replicated_axes: frozenset[ParallelAxis] = frozenset(), thd_global_seq_lens: Optional[list[int]] = None, ) -> list[UnsharderPlan]: if not parallel_infos: @@ -54,7 +55,21 @@ def compute_unsharder_plan( reversed_sharded_modifiers = [ (name, m) for name, m in reversed_sharded_modifiers if m.axis in sharded_axes ] - replicated_axes: set[ParallelAxis] = all_axes - sharded_axes + + # RECOMPUTE_PSEUDO is always implicitly replicated (system-injected, not user-facing) + auto_replicated: frozenset[ParallelAxis] = frozenset( + {ParallelAxis.RECOMPUTE_PSEUDO} & all_axes + ) + effective_replicated: frozenset[ParallelAxis] = ( + explicit_replicated_axes | auto_replicated + ) + + _validate_explicit_replicated( + explicit_replicated_axes=effective_replicated, + sharded_axes=sharded_axes, + all_axes=all_axes, + ) + replicated_axes: frozenset[ParallelAxis] = effective_replicated if not sharded_axes and not replicated_axes: return [] @@ -96,6 +111,37 @@ def compute_unsharder_plan( return plans +def _validate_explicit_replicated( + *, + explicit_replicated_axes: frozenset[ParallelAxis], + sharded_axes: set[ParallelAxis], + all_axes: set[ParallelAxis], +) -> None: + """Validate explicit replicated declarations against sharded axes and parallel_infos.""" + invalid: frozenset[ParallelAxis] = explicit_replicated_axes - all_axes + if invalid: + invalid_names: str = ", ".join(sorted(a.value for a in invalid)) + raise ValueError( + f"Declared replicated axes {{{invalid_names}}} not found in parallel_infos " + f"(active axes: {{{', '.join(sorted(a.value for a in all_axes))}}})" + ) + + conflict: set[ParallelAxis] = explicit_replicated_axes & sharded_axes + if conflict: + conflict_names: str = ", ".join(sorted(a.value for a in conflict)) + raise ValueError( + f"Axes {{{conflict_names}}} declared as both sharded and replicated" + ) + + undeclared: set[ParallelAxis] = all_axes - sharded_axes - explicit_replicated_axes + if undeclared: + undeclared_names: str = ", ".join(sorted(a.value for a in undeclared)) + raise ValueError( + f"Axes {{{undeclared_names}}} are active (axis_size > 1) but not declared " + f"in dims. Annotate as sharded in dim spec or as '# axis:replicated'." + ) + + def _validate( *, axes_to_validate: set[ParallelAxis], 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 19c752d3a..270090419 100644 --- a/python/sglang/srt/debug_utils/comparator/aligner/unsharder/types.py +++ b/python/sglang/srt/debug_utils/comparator/aligner/unsharder/types.py @@ -4,7 +4,7 @@ from typing import Annotated, Literal, Union from pydantic import Field, model_validator -from sglang.srt.debug_utils.comparator.dims import ParallelAxis +from sglang.srt.debug_utils.comparator.dims_spec import ParallelAxis from sglang.srt.debug_utils.comparator.utils import _FrozenBase diff --git a/python/sglang/srt/debug_utils/comparator/bundle_comparator.py b/python/sglang/srt/debug_utils/comparator/bundle_comparator.py index 3f69241c1..f4ce41ed5 100644 --- a/python/sglang/srt/debug_utils/comparator/bundle_comparator.py +++ b/python/sglang/srt/debug_utils/comparator/bundle_comparator.py @@ -18,7 +18,7 @@ from sglang.srt.debug_utils.comparator.aligner.entrypoint.types import AlignerPl from sglang.srt.debug_utils.comparator.aligner.token_aligner.smart.types import ( TokenAlignerPlan, ) -from sglang.srt.debug_utils.comparator.dims import ( +from sglang.srt.debug_utils.comparator.dims_spec import ( SEQ_DIM_NAME, TOKEN_DIM_NAME, apply_dim_names, diff --git a/python/sglang/srt/debug_utils/comparator/dims.py b/python/sglang/srt/debug_utils/comparator/dims.py deleted file mode 100644 index ec5638bf7..000000000 --- a/python/sglang/srt/debug_utils/comparator/dims.py +++ /dev/null @@ -1,338 +0,0 @@ -from __future__ import annotations - -import re -from enum import Enum -from typing import Optional - -import torch - -from sglang.srt.debug_utils.comparator.utils import _FrozenBase - -TOKEN_DIM_NAME: str = "t" -BATCH_DIM_NAME: str = "b" -SEQ_DIM_NAME: str = "s" -SQUEEZE_DIM_NAME: str = "1" - - -class TokenLayout(Enum): - T = "t" # single flat token dim - BS = "bs" # separate batch + seq dims, need collapse - - -class ParallelAxis(Enum): - TP = "tp" - CP = "cp" - EP = "ep" - SP = "sp" - RECOMPUTE_PSEUDO = "recompute_pseudo" - - -class Ordering(Enum): - ZIGZAG = "zigzag" - NATURAL = "natural" - - -class Reduction(Enum): - PARTIAL = "partial" - - -class ParallelModifier(_FrozenBase): - axis: ParallelAxis - ordering: Optional[Ordering] = None - reduction: Optional[Reduction] = None - - -_FUSED_NAME_SEP: str = "___" - - -class DimSpec(_FrozenBase): - name: str - parallel_modifiers: list[ParallelModifier] = [] - - @property - def sub_dims(self) -> list[str]: - """Sub-dim names. Fused: ``["num_heads", "head_dim"]``; plain: ``["h"]``.""" - return self.name.split("*") - - @property - def is_fused(self) -> bool: - return len(self.sub_dims) > 1 - - @property - def sanitized_name(self) -> str: - """Name safe for PyTorch named tensors (``*`` → ``___``).""" - if self.is_fused: - return _FUSED_NAME_SEP.join(self.sub_dims) - return self.name - - -class DimsSpec(_FrozenBase): - """Parsed result of a full dims string like ``"b s h[tp] # dp:=moe_dp"``.""" - - dims: list[DimSpec] - dp_group_alias: Optional[str] = None - - -class DimsSpec(_FrozenBase): - """Parsed result of a full dims string like ``"b s h(tp) # dp:=moe_dp"``.""" - - dims: list[DimSpec] - dp_group_alias: Optional[str] = None - - -class _SingletonDimUtil: - """Utilities for squeeze dims (name="1") and their singleton tensor-name mapping.""" - - PREFIX: str = "singleton" - - @staticmethod - def is_squeeze(spec: DimSpec) -> bool: - return spec.name == SQUEEZE_DIM_NAME - - @staticmethod - def filter_out(dim_specs: list[DimSpec]) -> list[DimSpec]: - return [s for s in dim_specs if not _SingletonDimUtil.is_squeeze(s)] - - @staticmethod - def make_name(index: int) -> str: - return f"{_SingletonDimUtil.PREFIX}{index}" - - @staticmethod - def is_singleton_name(name: str) -> bool: - return ( - name.startswith(_SingletonDimUtil.PREFIX) - and name[len(_SingletonDimUtil.PREFIX) :].isdigit() - ) - - @staticmethod - def sanitize_names(names: list[str]) -> list[str]: - """Replace '1' with 'singleton0', 'singleton1', ... for named tensor compatibility.""" - result: list[str] = [] - sq_idx: int = 0 - - for name in names: - if name == SQUEEZE_DIM_NAME: - result.append(_SingletonDimUtil.make_name(sq_idx)) - sq_idx += 1 - else: - result.append(name) - - return result - - -_DIM_PATTERN = re.compile(r"^(?P[a-zA-Z_]\w*)(?:\[(?P[^\]]+)\])?$") - -_FUSED_DIM_PATTERN = re.compile(r"^\((?P[^)]+)\)(?:\[(?P[^\]]+)\])?$") - -_SUB_DIM_NAME_PATTERN = re.compile(r"^[a-zA-Z_]\w*$") - -_AXIS_LOOKUP: dict[str, ParallelAxis] = {m.value: m for m in ParallelAxis} -_QUALIFIER_LOOKUP: dict[str, Ordering | Reduction] = { - **{m.value: m for m in Ordering}, - **{m.value: m for m in Reduction}, -} - - -def _parse_modifier_token(modifier_token: str, dim_token: str) -> ParallelModifier: - """Parse 'sp', 'cp:zigzag', 'tp:partial', or 'cp:zigzag+partial' → ParallelModifier. - - Format: ``axis`` or ``axis:qual`` or ``axis:qual+qual``. - Colon separates axis from qualifiers; ``+`` separates multiple qualifiers. - """ - axis_str: str - qualifiers_str: str - if ":" in modifier_token: - axis_str, qualifiers_str = modifier_token.split(":", maxsplit=1) - else: - axis_str, qualifiers_str = modifier_token, "" - - axis_str = axis_str.strip() - axis: Optional[ParallelAxis] = _AXIS_LOOKUP.get(axis_str) - if axis is None: - raise ValueError( - f"Unknown axis {axis_str!r} in modifier {modifier_token!r} " - f"of dim spec: {dim_token!r}" - ) - - ordering: Optional[Ordering] = None - reduction: Optional[Reduction] = None - - for q_str in (q.strip() for q in qualifiers_str.split("+") if q.strip()): - qualifier: Optional[Ordering | Reduction] = _QUALIFIER_LOOKUP.get(q_str) - if qualifier is None: - raise ValueError( - f"Unknown qualifier {q_str!r} in modifier " - f"{modifier_token!r} of dim spec: {dim_token!r}" - ) - if isinstance(qualifier, Ordering): - if ordering is not None: - raise ValueError( - f"Multiple ordering values in modifier " - f"{modifier_token!r} of dim spec: {dim_token!r}" - ) - ordering = qualifier - else: - if reduction is not None: - raise ValueError( - f"Multiple reduction values in modifier " - f"{modifier_token!r} of dim spec: {dim_token!r}" - ) - reduction = qualifier - - return ParallelModifier(axis=axis, ordering=ordering, reduction=reduction) - - -def parse_dim(token: str) -> DimSpec: - if token == SQUEEZE_DIM_NAME: - return DimSpec(name=SQUEEZE_DIM_NAME) - - fused_match = _FUSED_DIM_PATTERN.match(token) - if fused_match is not None: - return _parse_fused_dim(token=token, fused_match=fused_match) - - return _parse_single_dim(token) - - -def _parse_single_dim(token: str) -> DimSpec: - match = _DIM_PATTERN.match(token) - if match is None: - raise ValueError(f"Invalid dim token: {token!r}") - - name: str = match.group("name") - modifiers: list[ParallelModifier] = _parse_modifiers( - modifiers_str=match.group("modifiers"), dim_token=token - ) - return DimSpec(name=name, parallel_modifiers=modifiers) - - -def _parse_fused_dim(*, token: str, fused_match: re.Match[str]) -> DimSpec: - inner: str = fused_match.group("inner") - modifiers_str: Optional[str] = fused_match.group("modifiers") - - sub_names: list[str] = [s.strip() for s in inner.split("*")] - for sub_name in sub_names: - if not _SUB_DIM_NAME_PATTERN.match(sub_name): - raise ValueError( - f"Invalid sub-dim {sub_name!r} in fused dim token: {token!r}" - ) - - if len(sub_names) != len(set(sub_names)): - raise ValueError(f"Duplicate sub-dim names in fused dim token: {token!r}") - - if len(sub_names) < 2: - raise ValueError( - f"Fused dim must have at least 2 sub-dims, got {len(sub_names)} in: {token!r}" - ) - - fused_name: str = "*".join(sub_names) - modifiers: list[ParallelModifier] = _parse_modifiers( - modifiers_str=modifiers_str, dim_token=token - ) - return DimSpec(name=fused_name, parallel_modifiers=modifiers) - - -def _parse_modifiers( - *, modifiers_str: Optional[str], dim_token: str -) -> list[ParallelModifier]: - if modifiers_str is None: - return [] - - modifiers: list[ParallelModifier] = [] - seen_axes: set[ParallelAxis] = set() - - for modifier_token in (p.strip() for p in modifiers_str.split(",")): - modifier: ParallelModifier = _parse_modifier_token(modifier_token, dim_token) - if modifier.axis in seen_axes: - raise ValueError( - f"Duplicate axis {modifier.axis.value!r} in dim spec: {dim_token!r}" - ) - seen_axes.add(modifier.axis) - modifiers.append(modifier) - - return modifiers - - -def parse_dims(dims_str: str) -> DimsSpec: - """Parse ``"b s[cp:zigzag] h[tp] d # dp:=moe_dp"`` → :class:`DimsSpec`. - - The shape part (before ``#``) produces :pyattr:`DimsSpec.dims`. - The declaration part (after ``#``) is scanned for ``dp:=`` - which populates :pyattr:`DimsSpec.dp_group_alias`. - """ - parts: list[str] = dims_str.split("#", maxsplit=1) - raw: str = parts[0] - - if not raw.strip(): - raise ValueError("dims string must not be empty") - - dims: list[DimSpec] = [parse_dim(token) for token in raw.strip().split()] - - # Collect all semantic names (expanding fused sub-dims) for duplicate detection - semantic_names: list[str] = [] - for spec in dims: - if _SingletonDimUtil.is_squeeze(spec): - continue - semantic_names.extend(spec.sub_dims) - - if len(semantic_names) != len(set(semantic_names)): - duplicates = sorted({n for n in semantic_names if semantic_names.count(n) > 1}) - raise ValueError(f"Duplicate dim names: {duplicates}") - - dp_group_alias: Optional[str] = ( - _extract_dp_group_alias(parts[1]) if len(parts) > 1 else None - ) - - return DimsSpec(dims=dims, dp_group_alias=dp_group_alias) - - -def resolve_dim_names(dims_str: str) -> list[str]: - """Parse dims string and return tensor-compatible names ('1' → 'singleton0', ...).""" - specs: list[DimSpec] = parse_dims(dims_str).dims - names: list[str] = [spec.sanitized_name for spec in specs] - return _SingletonDimUtil.sanitize_names(names) - - -def find_dim_index(dim_specs: list[DimSpec], name: str) -> Optional[int]: - """Find index by name. Accepts both ``*``-form and ``___``-form for fused dims.""" - for i, spec in enumerate(dim_specs): - if spec.name == name or spec.sanitized_name == name: - return i - return None - - -def resolve_dim_by_name(tensor: torch.Tensor, name: str) -> int: - if tensor.names[0] is None: - raise ValueError(f"Tensor has no names, cannot resolve {name!r}") - - names: tuple[Optional[str], ...] = tensor.names - try: - return list(names).index(name) - except ValueError: - raise ValueError(f"Dim name {name!r} not in tensor names {names}") - - -def apply_dim_names(tensor: torch.Tensor, dim_names: list[str]) -> torch.Tensor: - if tensor.ndim != len(dim_names): - raise ValueError( - f"dims metadata mismatch: tensor has {tensor.ndim} dims (shape {list(tensor.shape)}) " - f"but dims string specifies {len(dim_names)} names {dim_names}. " - f"Please fix the dims string in the dumper.dump() call to match the actual tensor shape." - ) - return tensor.refine_names(*dim_names) - - -def strip_dim_names(tensor: torch.Tensor) -> torch.Tensor: - return tensor.rename(None) - - -_DP_ALIAS_PATTERN = re.compile(r"^dp:=(\w+)$") - - -def _extract_dp_group_alias(declaration_part: str) -> Optional[str]: - """Scan the ``#`` declaration section for a ``dp:=`` token.""" - for token in declaration_part.strip().split(): - match = _DP_ALIAS_PATTERN.match(token) - if match is not None: - return match.group(1) - - return None diff --git a/python/sglang/srt/debug_utils/comparator/dims_spec/__init__.py b/python/sglang/srt/debug_utils/comparator/dims_spec/__init__.py new file mode 100644 index 000000000..6d6480209 --- /dev/null +++ b/python/sglang/srt/debug_utils/comparator/dims_spec/__init__.py @@ -0,0 +1,49 @@ +from sglang.srt.debug_utils.comparator.dims_spec.dim_parser import parse_dim +from sglang.srt.debug_utils.comparator.dims_spec.dims_parser import ( + _SingletonDimUtil, + parse_dims, + resolve_dim_names, +) +from sglang.srt.debug_utils.comparator.dims_spec.tensor_naming import ( + apply_dim_names, + find_dim_index, + resolve_dim_by_name, + strip_dim_names, +) +from sglang.srt.debug_utils.comparator.dims_spec.types import ( + _FUSED_NAME_SEP, + BATCH_DIM_NAME, + SEQ_DIM_NAME, + SQUEEZE_DIM_NAME, + TOKEN_DIM_NAME, + DimSpec, + DimsSpec, + Ordering, + ParallelAxis, + ParallelModifier, + Reduction, + TokenLayout, +) + +__all__ = [ + "BATCH_DIM_NAME", + "SEQ_DIM_NAME", + "SQUEEZE_DIM_NAME", + "TOKEN_DIM_NAME", + "DimsSpec", + "DimSpec", + "Ordering", + "ParallelAxis", + "ParallelModifier", + "Reduction", + "TokenLayout", + "_FUSED_NAME_SEP", + "_SingletonDimUtil", + "apply_dim_names", + "find_dim_index", + "parse_dim", + "parse_dims", + "resolve_dim_by_name", + "resolve_dim_names", + "strip_dim_names", +] diff --git a/python/sglang/srt/debug_utils/comparator/dims_spec/comment_parser.py b/python/sglang/srt/debug_utils/comparator/dims_spec/comment_parser.py new file mode 100644 index 000000000..222c5ecb1 --- /dev/null +++ b/python/sglang/srt/debug_utils/comparator/dims_spec/comment_parser.py @@ -0,0 +1,59 @@ +from __future__ import annotations + +import re +from typing import NamedTuple, Optional + +from sglang.srt.debug_utils.comparator.dims_spec.types import ( + _AXIS_LOOKUP, + ParallelAxis, +) + +_DP_ALIAS_PATTERN = re.compile(r"^dp:=(\w+)$") +_REPLICATED_PATTERN = re.compile(r"^(\w+):replicated$") + + +class _CommentSuffix(NamedTuple): + dp_group_alias: Optional[str] = None + replicated_axes: frozenset[ParallelAxis] = frozenset() + + +def _parse_comment_suffix(declaration_part: str) -> _CommentSuffix: + """Parse the ``#`` comment section for dp alias and replicated declarations.""" + dp_group_alias: Optional[str] = None + replicated_axes: set[ParallelAxis] = set() + + for token in declaration_part.strip().split(): + dp_match = _DP_ALIAS_PATTERN.match(token) + if dp_match is not None: + if dp_group_alias is not None: + raise ValueError( + f"Duplicate dp alias declaration: already have {dp_group_alias!r}, " + f"got {dp_match.group(1)!r}" + ) + dp_group_alias = dp_match.group(1) + continue + + repl_match = _REPLICATED_PATTERN.match(token) + if repl_match is not None: + axis_str: str = repl_match.group(1) + axis: Optional[ParallelAxis] = _AXIS_LOOKUP.get(axis_str) + if axis is None: + raise ValueError( + f"Unknown axis {axis_str!r} in replicated declaration: {token!r}" + ) + if axis in replicated_axes: + raise ValueError( + f"Duplicate replicated declaration for axis {axis_str!r}" + ) + replicated_axes.add(axis) + continue + + raise ValueError( + f"Unrecognized token {token!r} in # comment section. " + f"Expected 'dp:=' or ':replicated'." + ) + + return _CommentSuffix( + dp_group_alias=dp_group_alias, + replicated_axes=frozenset(replicated_axes), + ) diff --git a/python/sglang/srt/debug_utils/comparator/dims_spec/dim_parser.py b/python/sglang/srt/debug_utils/comparator/dims_spec/dim_parser.py new file mode 100644 index 000000000..5aac65be2 --- /dev/null +++ b/python/sglang/srt/debug_utils/comparator/dims_spec/dim_parser.py @@ -0,0 +1,68 @@ +from __future__ import annotations + +import re +from typing import Optional + +from sglang.srt.debug_utils.comparator.dims_spec.modifier_parser import ( + _parse_modifiers, +) +from sglang.srt.debug_utils.comparator.dims_spec.types import ( + SQUEEZE_DIM_NAME, + DimSpec, + ParallelModifier, +) + +_DIM_PATTERN = re.compile(r"^(?P[a-zA-Z_]\w*)(?:\[(?P[^\]]+)\])?$") + +_FUSED_DIM_PATTERN = re.compile(r"^\((?P[^)]+)\)(?:\[(?P[^\]]+)\])?$") + +_SUB_DIM_NAME_PATTERN = re.compile(r"^[a-zA-Z_]\w*$") + + +def parse_dim(token: str) -> DimSpec: + if token == SQUEEZE_DIM_NAME: + return DimSpec(name=SQUEEZE_DIM_NAME) + + fused_match = _FUSED_DIM_PATTERN.match(token) + if fused_match is not None: + return _parse_fused_dim(token=token, fused_match=fused_match) + + return _parse_single_dim(token) + + +def _parse_single_dim(token: str) -> DimSpec: + match = _DIM_PATTERN.match(token) + if match is None: + raise ValueError(f"Invalid dim token: {token!r}") + + name: str = match.group("name") + modifiers: list[ParallelModifier] = _parse_modifiers( + modifiers_str=match.group("modifiers"), dim_token=token + ) + return DimSpec(name=name, parallel_modifiers=modifiers) + + +def _parse_fused_dim(*, token: str, fused_match: re.Match[str]) -> DimSpec: + inner: str = fused_match.group("inner") + modifiers_str: Optional[str] = fused_match.group("modifiers") + + sub_names: list[str] = [s.strip() for s in inner.split("*")] + for sub_name in sub_names: + if not _SUB_DIM_NAME_PATTERN.match(sub_name): + raise ValueError( + f"Invalid sub-dim {sub_name!r} in fused dim token: {token!r}" + ) + + if len(sub_names) != len(set(sub_names)): + raise ValueError(f"Duplicate sub-dim names in fused dim token: {token!r}") + + if len(sub_names) < 2: + raise ValueError( + f"Fused dim must have at least 2 sub-dims, got {len(sub_names)} in: {token!r}" + ) + + fused_name: str = "*".join(sub_names) + modifiers: list[ParallelModifier] = _parse_modifiers( + modifiers_str=modifiers_str, dim_token=token + ) + return DimSpec(name=fused_name, parallel_modifiers=modifiers) diff --git a/python/sglang/srt/debug_utils/comparator/dims_spec/dims_parser.py b/python/sglang/srt/debug_utils/comparator/dims_spec/dims_parser.py new file mode 100644 index 000000000..4e0a90895 --- /dev/null +++ b/python/sglang/srt/debug_utils/comparator/dims_spec/dims_parser.py @@ -0,0 +1,113 @@ +from __future__ import annotations + +from typing import Optional + +from sglang.srt.debug_utils.comparator.dims_spec.comment_parser import ( + _CommentSuffix, + _parse_comment_suffix, +) +from sglang.srt.debug_utils.comparator.dims_spec.dim_parser import parse_dim +from sglang.srt.debug_utils.comparator.dims_spec.types import ( + SQUEEZE_DIM_NAME, + DimSpec, + DimsSpec, + ParallelAxis, +) + + +class _SingletonDimUtil: + """Utilities for squeeze dims (name="1") and their singleton tensor-name mapping.""" + + PREFIX: str = "singleton" + + @staticmethod + def is_squeeze(spec: DimSpec) -> bool: + return spec.name == SQUEEZE_DIM_NAME + + @staticmethod + def filter_out(dim_specs: list[DimSpec]) -> list[DimSpec]: + return [s for s in dim_specs if not _SingletonDimUtil.is_squeeze(s)] + + @staticmethod + def make_name(index: int) -> str: + return f"{_SingletonDimUtil.PREFIX}{index}" + + @staticmethod + def is_singleton_name(name: str) -> bool: + return ( + name.startswith(_SingletonDimUtil.PREFIX) + and name[len(_SingletonDimUtil.PREFIX) :].isdigit() + ) + + @staticmethod + def sanitize_names(names: list[str]) -> list[str]: + """Replace '1' with 'singleton0', 'singleton1', ... for named tensor compatibility.""" + result: list[str] = [] + sq_idx: int = 0 + + for name in names: + if name == SQUEEZE_DIM_NAME: + result.append(_SingletonDimUtil.make_name(sq_idx)) + sq_idx += 1 + else: + result.append(name) + + return result + + +def parse_dims(dims_str: str) -> DimsSpec: + """Parse ``"b s[cp:zigzag] h[tp] d # dp:=moe_dp ep:replicated"`` → :class:`DimsSpec`. + + The shape part (before ``#``) produces :pyattr:`DimsSpec.dims`. + The declaration part (after ``#``) is scanned for: + - ``dp:=`` → :pyattr:`DimsSpec.dp_group_alias` + - ``axis:replicated`` → :pyattr:`DimsSpec.replicated_axes` + """ + parts: list[str] = dims_str.split("#", maxsplit=1) + raw: str = parts[0] + + if not raw.strip(): + raise ValueError("dims string must not be empty") + + dims: list[DimSpec] = [parse_dim(token) for token in raw.strip().split()] + + # Collect all semantic names (expanding fused sub-dims) for duplicate detection + semantic_names: list[str] = [] + for spec in dims: + if _SingletonDimUtil.is_squeeze(spec): + continue + semantic_names.extend(spec.sub_dims) + + if len(semantic_names) != len(set(semantic_names)): + duplicates = sorted({n for n in semantic_names if semantic_names.count(n) > 1}) + raise ValueError(f"Duplicate dim names: {duplicates}") + + comment_suffix: _CommentSuffix = ( + _parse_comment_suffix(parts[1]) if len(parts) > 1 else _CommentSuffix() + ) + dp_group_alias: Optional[str] = comment_suffix.dp_group_alias + replicated_axes: frozenset[ParallelAxis] = comment_suffix.replicated_axes + + sharded_axes: set[ParallelAxis] = { + m.axis for spec in dims for m in spec.parallel_modifiers + } + conflict: frozenset[ParallelAxis] = replicated_axes & sharded_axes + if conflict: + conflict_names: str = ", ".join(sorted(a.value for a in conflict)) + raise ValueError( + f"Axes declared as both sharded (in dim spec) and replicated " + f"(in # declaration): {conflict_names}" + ) + + return DimsSpec( + dims=dims, + dp_group_alias=dp_group_alias, + replicated_axes=replicated_axes, + ) + + +def resolve_dim_names(dims_str: str) -> list[str]: + """Parse dims string and return tensor-compatible names ('1' → 'singleton0', ...).""" + specs: list[DimSpec] = parse_dims(dims_str).dims + names: list[str] = [spec.sanitized_name for spec in specs] + return _SingletonDimUtil.sanitize_names(names) diff --git a/python/sglang/srt/debug_utils/comparator/dims_spec/modifier_parser.py b/python/sglang/srt/debug_utils/comparator/dims_spec/modifier_parser.py new file mode 100644 index 000000000..c1ecfa887 --- /dev/null +++ b/python/sglang/srt/debug_utils/comparator/dims_spec/modifier_parser.py @@ -0,0 +1,84 @@ +from __future__ import annotations + +from typing import Optional + +from sglang.srt.debug_utils.comparator.dims_spec.types import ( + _AXIS_LOOKUP, + _QUALIFIER_LOOKUP, + Ordering, + ParallelAxis, + ParallelModifier, + Reduction, +) + + +def _parse_modifier_token(modifier_token: str, dim_token: str) -> ParallelModifier: + """Parse 'sp', 'cp:zigzag', 'tp:partial', or 'cp:zigzag+partial' → ParallelModifier. + + Format: ``axis`` or ``axis:qual`` or ``axis:qual+qual``. + Colon separates axis from qualifiers; ``+`` separates multiple qualifiers. + """ + axis_str: str + qualifiers_str: str + if ":" in modifier_token: + axis_str, qualifiers_str = modifier_token.split(":", maxsplit=1) + else: + axis_str, qualifiers_str = modifier_token, "" + + axis_str = axis_str.strip() + axis: Optional[ParallelAxis] = _AXIS_LOOKUP.get(axis_str) + if axis is None: + raise ValueError( + f"Unknown axis {axis_str!r} in modifier {modifier_token!r} " + f"of dim spec: {dim_token!r}" + ) + + ordering: Optional[Ordering] = None + reduction: Optional[Reduction] = None + + for q_str in (q.strip() for q in qualifiers_str.split("+") if q.strip()): + if q_str == "sharded": + continue + qualifier: Optional[Ordering | Reduction] = _QUALIFIER_LOOKUP.get(q_str) + if qualifier is None: + raise ValueError( + f"Unknown qualifier {q_str!r} in modifier " + f"{modifier_token!r} of dim spec: {dim_token!r}" + ) + if isinstance(qualifier, Ordering): + if ordering is not None: + raise ValueError( + f"Multiple ordering values in modifier " + f"{modifier_token!r} of dim spec: {dim_token!r}" + ) + ordering = qualifier + else: + if reduction is not None: + raise ValueError( + f"Multiple reduction values in modifier " + f"{modifier_token!r} of dim spec: {dim_token!r}" + ) + reduction = qualifier + + return ParallelModifier(axis=axis, ordering=ordering, reduction=reduction) + + +def _parse_modifiers( + *, modifiers_str: Optional[str], dim_token: str +) -> list[ParallelModifier]: + if modifiers_str is None: + return [] + + modifiers: list[ParallelModifier] = [] + seen_axes: set[ParallelAxis] = set() + + for modifier_token in (p.strip() for p in modifiers_str.split(",")): + modifier: ParallelModifier = _parse_modifier_token(modifier_token, dim_token) + if modifier.axis in seen_axes: + raise ValueError( + f"Duplicate axis {modifier.axis.value!r} in dim spec: {dim_token!r}" + ) + seen_axes.add(modifier.axis) + modifiers.append(modifier) + + return modifiers diff --git a/python/sglang/srt/debug_utils/comparator/dims_spec/tensor_naming.py b/python/sglang/srt/debug_utils/comparator/dims_spec/tensor_naming.py new file mode 100644 index 000000000..0f06ebadc --- /dev/null +++ b/python/sglang/srt/debug_utils/comparator/dims_spec/tensor_naming.py @@ -0,0 +1,40 @@ +from __future__ import annotations + +from typing import Optional + +import torch + +from sglang.srt.debug_utils.comparator.dims_spec.types import DimSpec + + +def find_dim_index(dim_specs: list[DimSpec], name: str) -> Optional[int]: + """Find index by name. Accepts both ``*``-form and ``___``-form for fused dims.""" + for i, spec in enumerate(dim_specs): + if spec.name == name or spec.sanitized_name == name: + return i + return None + + +def resolve_dim_by_name(tensor: torch.Tensor, name: str) -> int: + if tensor.names[0] is None: + raise ValueError(f"Tensor has no names, cannot resolve {name!r}") + + names: tuple[Optional[str], ...] = tensor.names + try: + return list(names).index(name) + except ValueError: + raise ValueError(f"Dim name {name!r} not in tensor names {names}") + + +def apply_dim_names(tensor: torch.Tensor, dim_names: list[str]) -> torch.Tensor: + if tensor.ndim != len(dim_names): + raise ValueError( + f"dims metadata mismatch: tensor has {tensor.ndim} dims (shape {list(tensor.shape)}) " + f"but dims string specifies {len(dim_names)} names {dim_names}. " + f"Please fix the dims string in the dumper.dump() call to match the actual tensor shape." + ) + return tensor.refine_names(*dim_names) + + +def strip_dim_names(tensor: torch.Tensor) -> torch.Tensor: + return tensor.rename(None) diff --git a/python/sglang/srt/debug_utils/comparator/dims_spec/types.py b/python/sglang/srt/debug_utils/comparator/dims_spec/types.py new file mode 100644 index 000000000..afb792076 --- /dev/null +++ b/python/sglang/srt/debug_utils/comparator/dims_spec/types.py @@ -0,0 +1,77 @@ +from __future__ import annotations + +from enum import Enum +from typing import Optional + +from sglang.srt.debug_utils.comparator.utils import _FrozenBase + +TOKEN_DIM_NAME: str = "t" +BATCH_DIM_NAME: str = "b" +SEQ_DIM_NAME: str = "s" +SQUEEZE_DIM_NAME: str = "1" + + +class TokenLayout(Enum): + T = "t" # single flat token dim + BS = "bs" # separate batch + seq dims, need collapse + + +class ParallelAxis(Enum): + TP = "tp" + CP = "cp" + EP = "ep" + SP = "sp" + RECOMPUTE_PSEUDO = "recompute_pseudo" + + +class Ordering(Enum): + ZIGZAG = "zigzag" + NATURAL = "natural" + + +class Reduction(Enum): + PARTIAL = "partial" + + +class ParallelModifier(_FrozenBase): + axis: ParallelAxis + ordering: Optional[Ordering] = None + reduction: Optional[Reduction] = None + + +_AXIS_LOOKUP: dict[str, ParallelAxis] = {m.value: m for m in ParallelAxis} +_QUALIFIER_LOOKUP: dict[str, Ordering | Reduction] = { + **{m.value: m for m in Ordering}, + **{m.value: m for m in Reduction}, +} + +_FUSED_NAME_SEP: str = "___" + + +class DimSpec(_FrozenBase): + name: str + parallel_modifiers: list[ParallelModifier] = [] + + @property + def sub_dims(self) -> list[str]: + """Sub-dim names. Fused: ``["num_heads", "head_dim"]``; plain: ``["h"]``.""" + return self.name.split("*") + + @property + def is_fused(self) -> bool: + return len(self.sub_dims) > 1 + + @property + def sanitized_name(self) -> str: + """Name safe for PyTorch named tensors (``*`` → ``___``).""" + if self.is_fused: + return _FUSED_NAME_SEP.join(self.sub_dims) + return self.name + + +class DimsSpec(_FrozenBase): + """Parsed result of a full dims string like ``"b s h[tp] # dp:=moe_dp"``.""" + + dims: list[DimSpec] + dp_group_alias: Optional[str] = None + replicated_axes: frozenset[ParallelAxis] = frozenset() diff --git a/test/registered/debug_utils/comparator/aligner/entrypoint/test_executor.py b/test/registered/debug_utils/comparator/aligner/entrypoint/test_executor.py index 1f685ab83..ca0c5bed6 100644 --- a/test/registered/debug_utils/comparator/aligner/entrypoint/test_executor.py +++ b/test/registered/debug_utils/comparator/aligner/entrypoint/test_executor.py @@ -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 diff --git a/test/registered/debug_utils/comparator/aligner/entrypoint/test_planner.py b/test/registered/debug_utils/comparator/aligner/entrypoint/test_planner.py index df93b325b..139adeaf9 100644 --- a/test/registered/debug_utils/comparator/aligner/entrypoint/test_planner.py +++ b/test/registered/debug_utils/comparator/aligner/entrypoint/test_planner.py @@ -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 diff --git a/test/registered/debug_utils/comparator/aligner/reorderer/test_executor.py b/test/registered/debug_utils/comparator/aligner/reorderer/test_executor.py index bdfb4a7e4..811cc3d68 100644 --- a/test/registered/debug_utils/comparator/aligner/reorderer/test_executor.py +++ b/test/registered/debug_utils/comparator/aligner/reorderer/test_executor.py @@ -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) diff --git a/test/registered/debug_utils/comparator/aligner/reorderer/test_planner.py b/test/registered/debug_utils/comparator/aligner/reorderer/test_planner.py index 17cc2e5bf..1ecb0dcd1 100644 --- a/test/registered/debug_utils/comparator/aligner/reorderer/test_planner.py +++ b/test/registered/debug_utils/comparator/aligner/reorderer/test_planner.py @@ -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) diff --git a/test/registered/debug_utils/comparator/aligner/token_aligner/test_aux_plugins.py b/test/registered/debug_utils/comparator/aligner/token_aligner/test_aux_plugins.py index 95f04059d..e53473738 100644 --- a/test/registered/debug_utils/comparator/aligner/token_aligner/test_aux_plugins.py +++ b/test/registered/debug_utils/comparator/aligner/token_aligner/test_aux_plugins.py @@ -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) diff --git a/test/registered/debug_utils/comparator/aligner/token_aligner/test_executor.py b/test/registered/debug_utils/comparator/aligner/token_aligner/test_executor.py index c3565c3be..ec4325642 100644 --- a/test/registered/debug_utils/comparator/aligner/token_aligner/test_executor.py +++ b/test/registered/debug_utils/comparator/aligner/token_aligner/test_executor.py @@ -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 diff --git a/test/registered/debug_utils/comparator/aligner/token_aligner/test_planner.py b/test/registered/debug_utils/comparator/aligner/token_aligner/test_planner.py index 3fb3037ea..180780c4e 100644 --- a/test/registered/debug_utils/comparator/aligner/token_aligner/test_planner.py +++ b/test/registered/debug_utils/comparator/aligner/token_aligner/test_planner.py @@ -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 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 e435b7b50..8bfbd1e74 100644 --- a/test/registered/debug_utils/comparator/aligner/unsharder/test_executor.py +++ b/test/registered/debug_utils/comparator/aligner/unsharder/test_executor.py @@ -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) diff --git a/test/registered/debug_utils/comparator/aligner/unsharder/test_parallel_info.py b/test/registered/debug_utils/comparator/aligner/unsharder/test_parallel_info.py index 0ef84e35c..2388e9552 100644 --- a/test/registered/debug_utils/comparator/aligner/unsharder/test_parallel_info.py +++ b/test/registered/debug_utils/comparator/aligner/unsharder/test_parallel_info.py @@ -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) diff --git a/test/registered/debug_utils/comparator/aligner/unsharder/test_planner.py b/test/registered/debug_utils/comparator/aligner/unsharder/test_planner.py index 1e02abb46..855f59600 100644 --- a/test/registered/debug_utils/comparator/aligner/unsharder/test_planner.py +++ b/test/registered/debug_utils/comparator/aligner/unsharder/test_planner.py @@ -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: diff --git a/test/registered/debug_utils/comparator/dims_spec/__init__.py b/test/registered/debug_utils/comparator/dims_spec/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/test/registered/debug_utils/comparator/dims_spec/test_dim_parser.py b/test/registered/debug_utils/comparator/dims_spec/test_dim_parser.py new file mode 100644 index 000000000..d04a6fb1c --- /dev/null +++ b/test/registered/debug_utils/comparator/dims_spec/test_dim_parser.py @@ -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__])) diff --git a/test/registered/debug_utils/comparator/dims_spec/test_dims_parser.py b/test/registered/debug_utils/comparator/dims_spec/test_dims_parser.py new file mode 100644 index 000000000..ab1e17d76 --- /dev/null +++ b/test/registered/debug_utils/comparator/dims_spec/test_dims_parser.py @@ -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__])) diff --git a/test/registered/debug_utils/comparator/dims_spec/test_tensor_naming.py b/test/registered/debug_utils/comparator/dims_spec/test_tensor_naming.py new file mode 100644 index 000000000..2719760a8 --- /dev/null +++ b/test/registered/debug_utils/comparator/dims_spec/test_tensor_naming.py @@ -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__])) diff --git a/test/registered/debug_utils/comparator/dims_spec/test_types.py b/test/registered/debug_utils/comparator/dims_spec/test_types.py new file mode 100644 index 000000000..bc0e627b7 --- /dev/null +++ b/test/registered/debug_utils/comparator/dims_spec/test_types.py @@ -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__])) diff --git a/test/registered/debug_utils/comparator/test_dims.py b/test/registered/debug_utils/comparator/test_dims.py deleted file mode 100644 index 578c43004..000000000 --- a/test/registered/debug_utils/comparator/test_dims.py +++ /dev/null @@ -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__])) diff --git a/test/registered/debug_utils/comparator/test_entrypoint.py b/test/registered/debug_utils/comparator/test_entrypoint.py index 8bcb0f1eb..ff3839ced 100644 --- a/test/registered/debug_utils/comparator/test_entrypoint.py +++ b/test/registered/debug_utils/comparator/test_entrypoint.py @@ -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, diff --git a/test/registered/debug_utils/comparator/test_model_validation.py b/test/registered/debug_utils/comparator/test_model_validation.py index 3308231c1..82306af28 100644 --- a/test/registered/debug_utils/comparator/test_model_validation.py +++ b/test/registered/debug_utils/comparator/test_model_validation.py @@ -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,