Composable parallel strategies for vLLM-Omni.
A small, runtime-agnostic vocabulary for declaring a stage's parallel layout (:class:StrategySpec + mesh/routing/aggregation patterns), plus the logic to translate a per-role stack of specs into concrete engine sizing and overlay it onto merged stage configs.
The spec/pattern types are deliberately pure data with no engine imports; the translation and apply logic lives alongside them and consumes vLLM-Omni's own config objects.
Modules:
| Name | Description |
|---|---|
aggregation | Closed aggregation-pattern hierarchy for declarative parallel strategies. |
apply | Apply declarative parallel strategies onto merged stage configs. |
routing | Closed routing-pattern hierarchy for declarative parallel strategies. |
spec | Core declarative types for parallel-strategy specification. |
strategy_loader | Load per-role parallel strategies from a strategy file. |
translator | Translate a |
module-attribute ¶
AGGREGATION_PATTERN_VARIANTS: tuple[
type[AggregationPattern], ...
] = (
TakeRank,
Union,
GatherDim,
AllGather,
StitchSpatial,
StitchPipeline,
Combine,
FanInByStage,
)
module-attribute ¶
MESH_AXIS_KINDS: tuple[str, ...] = tuple(
(k.value) for k in MeshAxisKind
)
module-attribute ¶
ROUTING_PATTERN_VARIANTS: tuple[
type[RoutingPattern], ...
] = (
Broadcast,
PartitionByHash,
ShardSequence,
ShardSpatial,
RouteByExpert,
PipelineMicrobatch,
DuplicateWithCondTag,
RouteByStage,
)
Bases: ValueError
Raised when two aggregation inputs conflict for the same request id.
dataclass ¶
dataclass ¶
Bases: AggregationPattern
All ranks materialize the same aggregate result across the axis.
Bases: ValueError
Error for an invalid or unsupported strategy spec.
A single error type (rather than a tree of subclasses) keeps the public surface small and consistent with the rest of the codebase; the specific cause is in the message and is logged before the raise so it is visible even when the type is unavailable to a caller (e.g. across a server boundary).
Strategies that are valid but not built yet — key-stable / affinity routing today — raise NotImplementedError instead, to distinguish "we haven't implemented this" from "your config is wrong".
dataclass ¶
Bases: RoutingPattern
Every worker in the group sees the whole batch (TP, dense EP).
dataclass ¶
Bases: AggregationPattern
Custom reducer combination (CFG). Adapter resolves reducer_id.
dataclass ¶
Bases: RoutingPattern
Duplicate batch with a conditional tag (CFG).
dataclass ¶
Bases: AggregationPattern
Fan-in results routed by stage (stage replicas).
dataclass ¶
Bases: AggregationPattern
Gather along a dimension and materialize on one designated rank.
Default designated rank is axis-rank 0 unless adapter overrides.
dataclass ¶
L3 kernel slot referenced by StrategySpec; resolved by the kernel registry.
dataclass ¶
L2 hook slot referenced by StrategySpec; resolved by the model walker.
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.
dataclass ¶
Declares one axis of a process mesh (kind + size).
dataclass ¶
Result of translating a spec stack into omni parallel sizing.
class-attribute instance-attribute ¶
l1_owners: Mapping[MeshAxisKind, L1Owner] = field(
default_factory=dict
)
Return per-stage kwargs keyed by real OmniEngineArgs/EngineArgs field names.
Two derived values are intentionally not emitted here because they are not per-stage engine args:
stage_replica_size— a per-stage deploynum_replicasknob (StageDeployConfig); the deploy layer consumes it separately.omni_lb_policy— a pipeline-wide load-balancer policy the engine reads once at construction (seeStrategyApplyResult.omni_lb_policy); it is applied at the orchestrator level, not folded into per-stage args.
dataclass ¶
Bases: RoutingPattern
Stable hash partition by request key (DP, HSDP).
dataclass ¶
Bases: RoutingPattern
Pipeline parallelism microbatch routing (PP).
dataclass ¶
Bases: RoutingPattern
Sparse MoE expert routing.
dataclass ¶
Bases: RoutingPattern
Route requests by pipeline stage policy (stage replicas).
The random/round_robin/least_queue policies describe stateless load balancing that a runtime (e.g. an engine's built-in DP load balancer, or the omni StagePool balancer) can realize. hash describes deterministic, key-stable routing that a load balancer cannot honor on its own (see PartitionByHash).
dataclass ¶
dataclass ¶
Bases: RoutingPattern
Shard along a sequence dimension (SP-Ulysses, SP-Ring, CP).
dataclass ¶
Bases: ValueError
Raised when merged hook/kernel specs have conflicting declarations.
dataclass ¶
dataclass ¶
Bases: ValueError
Error for applying a strategy onto stage configs.
A single error type (rather than per-cause subclasses) keeps the public surface small and consistent with the rest of the codebase; the specific cause — an unmatched role, a conflict with an explicit deploy value, or a device-count mismatch — is in the message and is logged before the raise so it stays visible even when the type is unavailable to a caller (e.g. across a server boundary).
dataclass ¶
Outcome of applying strategies to a stage list.
stages is the same list passed in (mutated in place and returned for convenience). omni_lb_policy is the pipeline-wide load-balancer policy derived from any stage_replica axes (None if none declared one); the caller wires it into the orchestrator, since omni does not read it from stage configs.
per_role_config is keyed by the role key the caller supplied — a model_stage name (str). per_stage_config is the same information keyed by the resolved integer stage_id — handy for re-validating a stage after later layers (e.g. CLI overrides) have run.
dataclass ¶
Declarative contract for one parallelism scheme (data + hook/kernel slots).
class-attribute instance-attribute ¶
layer_hook_specs: tuple[LayerHookSpec, ...] = ()
dataclass ¶
Bases: AggregationPattern
Select one rank's result (TP — all ranks agree).
dataclass ¶
Bases: AggregationPattern
Union of disjoint replica results (DP, HSDP).
apply_strategy_specs(
stages: list[Any],
strategy_specs: Mapping[
RoleKey, Sequence[StrategySpec]
],
) -> StrategyApplyResult
Overlay per-role strategy specs onto a merged stage list.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
stages | list[Any] | the | required |
strategy_specs | Mapping[RoleKey, Sequence[StrategySpec]] | maps a role (a | required |
Returns:
| Name | Type | Description |
|---|---|---|
A | StrategyApplyResult | class: |
StrategyApplyResult | pipeline-wide |
Raises:
| Type | Description |
|---|---|
StrategyApplyError | a role matched zero or multiple stages, a derived value conflicts with an explicit deploy value (including two roles deriving different |
AxisTranslationError | the spec stack is invalid/unsupported. |
NotImplementedError | the spec stack requests routing not built yet. |
check_device_layout(
devices: Any,
*,
tensor_parallel_size: int,
data_parallel_size: int,
pipeline_parallel_size: int,
num_replicas: int,
role: RoleKey,
) -> None
Validate a stage's device count against its world size (× replicas in pool mode).
Raises :class:StrategyApplyError when the declared devices count is neither the per-replica world size nor the full num_replicas pool. Exposed publicly so the resolved (post-CLI) layout can be re-checked after later override layers, not just the strategy-derived snapshot.
translate_strategy_stack(
specs: Sequence[StrategySpec],
) -> OmniParallelConfig
Translate a spec stack into an OmniParallelConfig.
Supported kinds: dp (engine data parallel), tp, pp, (dense) ep, and stage_replica (omni replicas). Raises NotImplementedError for deferred (affinity / key-stable) routing, and :class:AxisTranslationError for any other invalid spec (a kind not yet translatable, a repeated kind, an owner incompatible with the axis kind, or unsupported routing). The EP degree must equal tensor_parallel_size * data_parallel_size.