vllm_omni.diffusion.models.cosmos3.mixed_precision.runtime ¶

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

PrecisionPath `module-attribute` ¶

PrecisionPath = Literal['reasoner', 'generation']

Bases: LinearMethodBase

Dispatch between a checkpoint-native method and dense A16.

base_method = base_method

module_name = module_name

path = path

runtime = runtime

strategy = strategy

apply(
    layer: Module, x: Tensor, bias: Tensor | None = None
) -> Tensor

create_weights(*args, **kwargs) -> None

process_weights_after_loading(layer: Module) -> None

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

config = config

install(transformer: Module) -> None

reset() -> None

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

use_high_precision(path: PrecisionPath) -> bool