Skip to content

vllm_omni.diffusion.models.cosmos3.mixed_precision

Scheduled A16 execution for quantized Cosmos3 checkpoints.

Modules:

Name Description
checkpoint

Checkpoint discovery for the Cosmos3 diffusion-step precision policy.
config

Configuration parsing and denoising-step policy for mixed precision.
runtime

Format-agnostic schedule state, linear discovery, and dispatch.
strategy

Format-specific live-weight validation and reference A16 materialization.

Cosmos3MixedPrecisionConfig dataclass

Validated precision policy shared by every quantization strategy.

first_steps class-attribute instance-attribute 

first_steps: int = 3

last_steps class-attribute instance-attribute 

last_steps: int = 3

reasoner class-attribute instance-attribute 

reasoner: ReasonerPolicy = 'a16'

from_additional_config classmethod 

from_additional_config(
    additional_config: Mapping[str, object] | None,
) -> Cosmos3MixedPrecisionConfig | None

Parse Cosmos3 fields from vLLM-Omni's additional configuration.

resolve_additional_config classmethod 

resolve_additional_config(
    additional_config: Mapping[str, object] | None,
) -> tuple[bool, Cosmos3MixedPrecisionConfig | None]

Return whether an override was supplied and its effective policy.

The separate presence bit distinguishes an absent runtime override from an explicit disable. That distinction matters when a checkpoint owns a default policy.

use_high_precision 

use_high_precision(step_index: int, num_steps: int) -> bool

Select the high path for an actual scheduler-step index.

Cosmos3MixedPrecisionRuntime

Own one transformer's schedule state and wrapped linear inventory.

config instance-attribute 

config = config

install 

install(transformer: Module) -> None

reset 

reset() -> None

set_step 

set_step(step_index: int, num_steps: int) -> None

use_high_precision 

use_high_precision(path: PrecisionPath) -> bool

resolve_mixed_precision_config 

resolve_mixed_precision_config(
    od_config: object,
) -> tuple[
    Cosmos3MixedPrecisionConfig | None, PolicySource
]

Resolve an explicit runtime override before the checkpoint default.