Skip to content

vllm_omni.config.composable_parallel.spec

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.

HOOK_CATEGORY_ORDER module-attribute 

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,
}

HookCategory module-attribute 

HookCategory = Literal[
    "linear",
    "attention_pre",
    "attention_post",
    "ffn_pre",
    "ffn_post",
    "moe_dispatch",
    "other",
]

MESH_AXIS_KINDS module-attribute 

MESH_AXIS_KINDS: tuple[str, ...] = tuple(
    (k.value) for k in MeshAxisKind
)

KernelSpec dataclass

L3 kernel slot referenced by StrategySpec; resolved by the kernel registry.

axis_index class-attribute instance-attribute 

axis_index: int = 0

category class-attribute instance-attribute 

category: HookCategory = 'other'

group_axis_kind class-attribute instance-attribute 

group_axis_kind: MeshAxisKind | None = None

kernel_id instance-attribute 

kernel_id: str

priority class-attribute instance-attribute 

priority: int = 0

requires_collective class-attribute instance-attribute 

requires_collective: bool = False

target class-attribute instance-attribute 

target: str | None = None

LayerHookSpec dataclass

L2 hook slot referenced by StrategySpec; resolved by the model walker.

axis_index class-attribute instance-attribute 

axis_index: int = 0

category class-attribute instance-attribute 

category: HookCategory = 'other'

hook_id instance-attribute 

hook_id: str

priority class-attribute instance-attribute 

priority: int = 0

target class-attribute instance-attribute 

target: str | None = None

MeshAxisKind

Bases: 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.

CFG class-attribute instance-attribute 

CFG = 'cfg'

CP class-attribute instance-attribute 

CP = 'cp'

DP class-attribute instance-attribute 

DP = 'dp'

EP class-attribute instance-attribute 

EP = 'ep'

HSDP class-attribute instance-attribute 

HSDP = 'hsdp'

PP class-attribute instance-attribute 

PP = 'pp'

SP_RING class-attribute instance-attribute 

SP_RING = 'sp_ring'

SP_ULYSSES class-attribute instance-attribute 

SP_ULYSSES = 'sp_ulysses'

STAGE_PP class-attribute instance-attribute 

STAGE_PP = 'stage_pp'

STAGE_REPLICA class-attribute instance-attribute 

STAGE_REPLICA = 'stage_replica'

TP class-attribute instance-attribute 

TP = 'tp'

VAE_PP class-attribute instance-attribute 

VAE_PP = 'vae_pp'

MeshAxisSpec dataclass

Declares one axis of a process mesh (kind + size).

kind instance-attribute 

kind: MeshAxisKind

size instance-attribute 

size: int

SpecMergeConflictError

Bases: ValueError

Raised when merged hook/kernel specs have conflicting declarations.

StrategySpec dataclass

Declarative contract for one parallelism scheme (data + hook/kernel slots).

aggregation instance-attribute 

aggregation: AggregationPattern

kernel_specs class-attribute instance-attribute 

kernel_specs: tuple[KernelSpec, ...] = ()

layer_hook_specs class-attribute instance-attribute 

layer_hook_specs: tuple[LayerHookSpec, ...] = ()

mesh_axis instance-attribute 

mesh_axis: MeshAxisSpec

name instance-attribute 

name: str

routing instance-attribute 

routing: RoutingPattern

shard_extension class-attribute instance-attribute 

shard_extension: Mapping[str, Any] = field(
    default_factory=dict
)