vllm_omni.diffusion.models.ltx2.ops ¶

Shared operator selection for LTX-2 eager kernels and their tests.

Modules:

Name	Description
`fna`	TileLang FNA for the LTX-2.5 DiffVAE stage-5 schedule.
`qk_rms_norm`	Bit-exact multi-row Q/K RMSNorm for the LTX diffusion VAE decoder.
`residual_adaln`	Exact LTX DiffVAE residual, RMSNorm, and AdaLN Triton fusions.
`rope_tables`	Cached 3D RoPE tables for LTX DiffVAE kernels.
`swiglu`	Bit-exact workspace-reusing SwiGLU for the LTX diffusion VAE decoder.

LTX2VAEOperatorSet `dataclass` ¶

fna: Callable[..., Tensor] | None

qk_norm_rope: Callable[..., tuple[Tensor, Tensor] | None]

residual_add: Callable[..., Tensor | None]

residual_norm: Callable[..., Tensor | None]

swiglu: Callable[..., Tensor | None]

is_ltx2_fna_eligible(tensor: Tensor) -> bool

Keep the strict, Hopper-tuned TileLang FNA schedule on SM90.

is_ltx2_fusion_eligible(tensor: Tensor) -> bool

Allow self-verifying pointwise fusions on qualified CUDA devices.

resolve_ltx2_vae_operators(
    device: device,
) -> LTX2VAEOperatorSet | None

Select the production operators, or return None for unsupported devices.