Skip to content

vllm_omni.config.composable_parallel.aggregation

Closed aggregation-pattern hierarchy for declarative parallel strategies.

An aggregation pattern describes how the per-worker results of one mesh axis are recombined into a single logical result (e.g. take one tensor-parallel rank's output, union disjoint data-parallel shards). The hierarchy is closed: subclasses may only be defined in this module.

AGGREGATION_PATTERN_VARIANTS module-attribute 

AGGREGATION_PATTERN_VARIANTS: tuple[
    type[AggregationPattern], ...
] = (
    TakeRank,
    Union,
    GatherDim,
    AllGather,
    StitchSpatial,
    StitchPipeline,
    Combine,
    FanInByStage,
)

ReducerFn module-attribute 

ReducerFn = Callable[..., object]

AggregationConflictError

Bases: ValueError

Raised when two aggregation inputs conflict for the same request id.

AggregationPattern dataclass

Base class for closed aggregation patterns. Subclass only within this module.

pattern_kind 

pattern_kind() -> str

AllGather dataclass

Bases: AggregationPattern

All ranks materialize the same aggregate result across the axis.

dim class-attribute instance-attribute 

dim: int | None = None

pattern_kind 

pattern_kind() -> str

Combine dataclass

Bases: AggregationPattern

Custom reducer combination (CFG). Adapter resolves reducer_id.

reducer_id class-attribute instance-attribute 

reducer_id: str = 'default'

pattern_kind 

pattern_kind() -> str

FanInByStage dataclass

Bases: AggregationPattern

Fan-in results routed by stage (stage replicas).

pattern_kind 

pattern_kind() -> str

GatherDim dataclass

Bases: AggregationPattern

Gather along a dimension and materialize on one designated rank.

Default designated rank is axis-rank 0 unless adapter overrides.

dim class-attribute instance-attribute 

dim: int = 1

pattern_kind 

pattern_kind() -> str

StitchPipeline dataclass

Bases: AggregationPattern

Stitch pipeline stage outputs (PP).

pattern_kind 

pattern_kind() -> str

StitchSpatial dataclass

Bases: AggregationPattern

Stitch spatial patches (VAE-PP).

grid class-attribute instance-attribute 

grid: tuple[int, int] = (1, 1)

pattern_kind 

pattern_kind() -> str

TakeRank dataclass

Bases: AggregationPattern

Select one rank's result (TP — all ranks agree).

rank class-attribute instance-attribute 

rank: int = 0

pattern_kind 

pattern_kind() -> str

Union dataclass

Bases: AggregationPattern

Union of disjoint replica results (DP, HSDP).

pattern_kind 

pattern_kind() -> str