ff6f0da271
Signed-off-by: Kushan A. <kahmadian@nvidia.com> Signed-off-by: kushanam <kushanam@users.noreply.github.com> Signed-off-by: Kushan A <kahmadian@nvidia.com> Signed-off-by: kahmadian <kahmadian@users.noreply.github.com> Co-authored-by: kushanam <kushanam@users.noreply.github.com> Co-authored-by: kahmadian <kahmadian@users.noreply.github.com>
137 lines
4.2 KiB
Python
137 lines
4.2 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Core declarative types for parallel-strategy specification.
|
|
|
|
A :class:`StrategySpec` declares one parallelism scheme: the mesh axis it sizes
|
|
(kind + degree), how a batch is routed across that axis, and how per-worker
|
|
results are aggregated back. A *stack* of specs (one per axis) fully describes a
|
|
stage's parallel layout in a runtime-agnostic, data-only form that can be
|
|
translated into concrete engine sizing.
|
|
|
|
The mesh-axis kinds below enumerate every parallelism dimension the contract can
|
|
describe. Only a subset is wired end-to-end today (see ``translator.py``):
|
|
|
|
* **Wired** (translatable now): ``tp``, ``dp``, ``pp``, ``ep``,
|
|
``stage_replica``.
|
|
* **Reserved** (declared in the type system but not yet translatable — the
|
|
translator raises ``AxisTranslationError`` for them): ``sp_ulysses``,
|
|
``sp_ring``, ``cfg``, ``vae_pp``, ``hsdp``, ``stage_pp``, ``cp``.
|
|
|
|
Reserved kinds fail fast at translation time rather than silently doing nothing,
|
|
so declaring one is an explicit "not yet" rather than a no-op.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import Mapping
|
|
from dataclasses import dataclass, field
|
|
from enum import Enum
|
|
from typing import Any, Literal
|
|
|
|
from vllm_omni.config.composable_parallel.aggregation import AggregationPattern
|
|
from vllm_omni.config.composable_parallel.routing import RoutingPattern
|
|
|
|
|
|
class MeshAxisKind(str, Enum):
|
|
"""Enumerates every parallelism dimension the strategy contract can describe.
|
|
|
|
Subclassing ``str`` keeps existing string comparisons and dict keys working
|
|
(``MeshAxisKind.TP == "tp"`` and both hash the same), while making the type
|
|
refactor-safe: a typo'd kind is an ``AttributeError`` on the enum rather than
|
|
a silently-wrong bare string.
|
|
"""
|
|
|
|
TP = "tp"
|
|
DP = "dp"
|
|
PP = "pp"
|
|
EP = "ep"
|
|
SP_ULYSSES = "sp_ulysses"
|
|
SP_RING = "sp_ring"
|
|
CFG = "cfg"
|
|
VAE_PP = "vae_pp"
|
|
HSDP = "hsdp"
|
|
STAGE_PP = "stage_pp"
|
|
STAGE_REPLICA = "stage_replica"
|
|
CP = "cp"
|
|
|
|
|
|
# Bare string values, kept as a plain tuple so membership checks and error
|
|
# messages read as the kind names (``"tp"``) rather than enum reprs. ``str``-enum
|
|
# equality means both ``"tp"`` and ``MeshAxisKind.TP`` satisfy ``in`` checks.
|
|
MESH_AXIS_KINDS: tuple[str, ...] = tuple(k.value for k in MeshAxisKind)
|
|
|
|
HookCategory = Literal[
|
|
"linear",
|
|
"attention_pre",
|
|
"attention_post",
|
|
"ffn_pre",
|
|
"ffn_post",
|
|
"moe_dispatch",
|
|
"other",
|
|
]
|
|
|
|
HOOK_CATEGORY_ORDER: dict[HookCategory, int] = {
|
|
"linear": 0,
|
|
"attention_pre": 1,
|
|
"attention_post": 2,
|
|
"ffn_pre": 3,
|
|
"ffn_post": 4,
|
|
"moe_dispatch": 5,
|
|
"other": 6,
|
|
}
|
|
|
|
|
|
class SpecMergeConflictError(ValueError):
|
|
"""Raised when merged hook/kernel specs have conflicting declarations."""
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class MeshAxisSpec:
|
|
"""Declares one axis of a process mesh (kind + size)."""
|
|
|
|
kind: MeshAxisKind
|
|
size: int
|
|
|
|
def __post_init__(self) -> None:
|
|
if self.kind not in MESH_AXIS_KINDS:
|
|
raise ValueError(f"MeshAxisSpec.kind must be one of {MESH_AXIS_KINDS}, got {self.kind!r}")
|
|
if self.size <= 0:
|
|
raise ValueError(f"MeshAxisSpec.size must be > 0, got {self.size}")
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class LayerHookSpec:
|
|
"""L2 hook slot referenced by StrategySpec; resolved by the model walker."""
|
|
|
|
hook_id: str
|
|
target: str | None = None
|
|
category: HookCategory = "other"
|
|
priority: int = 0
|
|
axis_index: int = 0
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class KernelSpec:
|
|
"""L3 kernel slot referenced by StrategySpec; resolved by the kernel registry."""
|
|
|
|
kernel_id: str
|
|
target: str | None = None
|
|
category: HookCategory = "other"
|
|
priority: int = 0
|
|
axis_index: int = 0
|
|
group_axis_kind: MeshAxisKind | None = None
|
|
requires_collective: bool = False
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class StrategySpec:
|
|
"""Declarative contract for one parallelism scheme (data + hook/kernel slots)."""
|
|
|
|
name: str
|
|
mesh_axis: MeshAxisSpec
|
|
routing: RoutingPattern
|
|
aggregation: AggregationPattern
|
|
layer_hook_specs: tuple[LayerHookSpec, ...] = ()
|
|
kernel_specs: tuple[KernelSpec, ...] = ()
|
|
shard_extension: Mapping[str, Any] = field(default_factory=dict)
|