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

> Accelerate LTX-2 transformer inference using torch.compile. This guide shows how to compile transformer blocks for faster diffusion models and reduce Python overhead.

- Repository: [Lightricks/LTX-2](https://github.com/Lightricks/LTX-2)
- Tags: performance
- Published: 2026-08-15

---

**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](https://github.com/Lightricks/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`](https://github.com/Lightricks/LTX-2/blob/main/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`](https://github.com/Lightricks/LTX-2/blob/main/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`](https://github.com/Lightricks/LTX-2/blob/main/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:

```bash

# 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`](https://github.com/Lightricks/LTX-2/blob/main/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`](https://github.com/Lightricks/LTX-2/blob/main/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:

```python
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`](https://github.com/Lightricks/LTX-2/blob/main/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`](https://github.com/Lightricks/LTX-2/blob/main/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:

```python
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`](https://github.com/Lightricks/LTX-2/blob/main/packages/ltx-core/src/ltx_core/model/transformer/compiling.py) | Core compilation machinery |
| [`packages/ltx-pipelines/src/ltx_pipelines/utils/args.py`](https://github.com/Lightricks/LTX-2/blob/main/packages/ltx-pipelines/src/ltx_pipelines/utils/args.py) | CLI `--compile` flag implementation |
| [`packages/ltx-trainer/src/ltx_trainer/trainer.py`](https://github.com/Lightricks/LTX-2/blob/main/packages/ltx-trainer/src/ltx_trainer/trainer.py) | Training-time compilation logging |
| [`packages/ltx-core/src/ltx_core/model/transformer/attention.py`](https://github.com/Lightricks/LTX-2/blob/main/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`](https://github.com/Lightricks/LTX-2/blob/main/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.