Skip to content

vllm_omni.diffusion.models.ltx2.ops.residual_adaln ¶

Exact LTX DiffVAE residual, RMSNorm, and AdaLN Triton fusions.

logger `module-attribute` ¶

logger = logging.getLogger(__name__)

try_residual_add3_exact ¶

try_residual_add3_exact(
    x: Tensor,
    residual_a: Tensor,
    residual_b: Tensor,
    residual_c: Tensor,
) -> Tensor | None

Fuse three ordered BF16 residual additions into one pointwise pass.

try_residual_rms_norm_modulate_exact ¶

try_residual_rms_norm_modulate_exact(
    x: Tensor,
    residual_a: Tensor,
    residual_b: Tensor | None,
    norm_weight: Tensor,
    scale: Tensor,
    shift: Tensor,
    eps: float,
) -> Tensor | None

Fuse one or two ordered BF16 residuals with RMSNorm and AdaLN.