Accelerating LTX-2 Transformer Inference with `torch.compile`: A Complete Guide

LTX-2 accelerates transformer-based diffusion models by compiling individual transformer blocks with torch.compile, with optional CUDA-graph capture to eliminate Python overhead.

The LTX-2 video generation framework from Lightricks achieves dramatic inference speedups through deep torch.compile integration. According to the LTX-2 source code, the implementation spans CLI configuration, dynamic shape handling, and per-block compilation with optional full-model CUDA-graph capture. This guide walks through the architecture, configuration options, and practical usage patterns.

How torch.compile Is Integrated in LTX-2

The acceleration system operates through three coordinated layers:

Layer	Purpose	Key Source
CLI flag – --compile	User-facing switch that creates a CompilationConfig	ltx_pipelines/utils/args.py (lines 889-1004)
Configuration object – CompilationConfig	Holds Dynamo/Inductor options and LTX-specific knobs	ltx_core/model/transformer/compiling.py (lines 22-48)
Compilation runner – compile_transformer_captured / _compile_blocks	Wraps blocks with torch.compile and manages CUDA-graph capture	ltx_core/model/transformer/compiling.py (lines 61-90)

CLI Configuration: The --compile Flag

The entry point for accelerating LTX-2 transformer inference with torch.compile is the --compile command-line argument. When invoked without parameters, it uses sensible defaults (backend="inductor", mode=None).

Users can override any CompilationConfig field via KEY=VALUE syntax:

# Basic usage with defaults

python -m ltx_pipelines.ti2vid_two_stages --compile

# Reduce-overhead mode with fullgraph enforcement

python -m ltx_pipelines.ti2vid_two_stages \
  --compile mode=reduce-overhead fullgraph=true

# Maximum performance: CUDA-graph capture with GPU-resident weights

python -m ltx_pipelines.ti2vid_two_stages \
  --compile capture=true keeps_gpu_resident_weights=true

The argument parsing logic in ltx_pipelines/utils/args.py implements a custom CompileAction class that converts these key-value pairs into a structured CompilationConfig object.

Configuration Options in CompilationConfig

The CompilationConfig dataclass in ltx_core/model/transformer/compiling.py exposes fine-grained control over the compilation process:

Parameter	Default	Description
mode	None	Dynamo optimization mode: "default", "reduce-overhead", "max-autotune"
backend	"inductor"	Compiler backend: "inductor", "eager", "aot_eager"
fullgraph	False	If True, requires entire computation graph to be compilable (no graph breaks)
dynamic	False	Enable dynamic shape support globally
seq_dim_dynamic	False	Mark sequence dimension as dynamic specifically (avoids recompilation across frame counts)
capture	False	Capture entire block-loop as CUDA graph for steady-state inference
recompile_perturbed_block	False	Whether to recompile when classifier-free guidance perturbations change

Setting seq_dim_dynamic=True is particularly important for video generation workloads. It calls torch._dynamo.mark_dynamic on the sequence dimension, allowing a single compiled artifact to handle variable-length sequences without recompilation.

Per-Block Compilation with _compile_blocks

The core compilation happens in compile_transformer_captured and its helper _compile_blocks. This approach compiles each transformer block individually rather than the entire model at once:

from ltx_core.model.transformer.compiling import (
    CompilationConfig, compile_transformer_captured
)
from ltx_core.model.transformer.model import LTXModel

# Configure for production inference

cfg = CompilationConfig(
    mode="reduce-overhead",
    backend="inductor",
    fullgraph=True,
    seq_dim_dynamic=True,  # Handle variable frame counts

    capture=False,         # Set True for CUDA-graph mode

)

model = LTXModel.from_checkpoint("path/to/checkpoint")
model = compile_transformer_captured(model, cfg)

Behind the scenes, _compile_blocks iterates over model.transformer_blocks and replaces each with torch.compile(block, **config). The Python loop calling these blocks remains in eager mode, but the heavy computation inside each block becomes optimized Inductor kernels.

CUDA-Graph Capture for Maximum Throughput

When capture=True, LTX-2 enters its highest-performance mode. After block-wise compilation, the entire model._process_transformer_blocks loop is recorded once using CudaGraphRunner. Subsequent inference calls replay this captured graph entirely on the GPU.

Requirements for capture mode:

• keeps_gpu_resident_weights=True — weights must remain in GPU memory
• Input shapes must remain stable (the graph is recorded for specific tensor sizes)
• seq_dim_dynamic should typically be False for capture, or used with fixed sequence lengths

This "single-graph" runtime eliminates Python interpreter overhead and kernel launch latency, delivering the lowest possible per-frame latency for production deployments.

Handling Perturbations: CompiledBlockPerturbationsProcessor

LTX-2 supports classifier-free guidance and other per-block perturbations that modify inputs dynamically. The CompiledBlockPerturbationsProcessor class bridges this flexibility with compilation:

• Prepares perturbed inputs before they enter compiled regions
• Preserves the ability to apply skip masks and guidance modifications
• Maintains compiled performance for the core computation

This design allows recompile_perturbed_block=False (the default), avoiding recompilation when guidance scale changes between inference steps.

Safety Guards and Known Limitations

The LTX-2 source code includes several protective measures:

FSDP + torch.compile warning — The trainer in ltx_trainer/trainer.py logs a warning that this combination is experimental and may hang on the first iteration.

Selective kernel exclusion — Certain low-level kernels are wrapped with @torch.compiler.disable to prevent problematic graph breaks. One example is found in ltx_core/model/transformer/attention.py, where stage-5 CuTe launches are excluded from compilation.

Verifying Compilation Status

Debug compiled blocks with this inspection pattern:

for i, blk in enumerate(model.transformer_blocks):
    is_compiled = hasattr(blk, '__torch_compile_success__')
    print(f"Block {i}: compiled? {is_compiled}")

The __torch_compile_success__ attribute is set by PyTorch's compiler when a module successfully compiles.

Key Source Files

File	Role
packages/ltx-core/src/ltx_core/model/transformer/compiling.py	Core compilation machinery
packages/ltx-pipelines/src/ltx_pipelines/utils/args.py	CLI --compile flag implementation
packages/ltx-trainer/src/ltx_trainer/trainer.py	Training-time compilation logging
packages/ltx-core/src/ltx_core/model/transformer/attention.py	Example of @torch.compiler.disable usage

Summary

• Accelerating LTX-2 transformer inference with torch.compile is enabled via the --compile CLI flag, which creates a CompilationConfig with flexible key-value overrides
• The system compiles individual transformer blocks rather than the full model, preserving compatibility with LTX-2's perturbation-based guidance API
• seq_dim_dynamic=True marks the sequence dimension as dynamic, avoiding recompilation across different frame counts
• capture=True enables CUDA-graph mode for maximum steady-state throughput, eliminating Python overhead entirely
• Safety guards in the trainer prevent problematic combinations like FSDP with compilation, and selective @torch.compiler.disable annotations exclude problematic kernels

Frequently Asked Questions

What is the fastest configuration for LTX-2 inference?

Use mode=max-autotune capture=true with keeps_gpu_resident_weights=true. This compiles with autotuned kernels and captures the entire transformer loop as a CUDA graph. Note that this requires fixed input shapes and sufficient GPU memory to keep all weights resident.

Why does LTX-2 compile blocks individually instead of the whole model?

Per-block compilation preserves the flexibility of LTX-2's perturbation system for classifier-free guidance. The CompiledBlockPerturbationsProcessor can modify inputs between blocks while the heavy computation inside each block runs as optimized kernels. Full-model compilation would struggle with these dynamic perturbations.

How do I avoid recompilation when changing video length?

Set seq_dim_dynamic=true in your --compile arguments or CompilationConfig. This calls torch._dynamo.mark_dynamic on the sequence dimension, allowing PyTorch to generate a single compiled artifact that handles variable sequence lengths without recompilation.

Can I use torch.compile during training with FSDP?

The LTX-2 trainer logs a warning against this combination in ltx_trainer/trainer.py. FSDP with torch.compile is marked as experimental and may hang on the first iteration. For training, consider compilation only on the individual blocks without FSDP sharding, or use eager mode until the integration stabilizes.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}