[perf]: enable per-block torch.compile for LTX2 with persistent cache

[Someone45]

Cache LTX2 graphs and allow all 48 blocks to share one compiled graph. Also added compile to the dreamverse backend launch.

Purpose

Dreamverse LTX2 startup compilation cost is very high (10+ minutes) and we are able to cache it, reducing subsequent start up below 2 minutes.

LTX2 was the only major DiT still compiled as a single whole-model fullgraph, because LTX2Transformer3DModel never wired the _compile_conditions declared on its arch config. That produced enormous, unportable graphs (~8k guards) and long cold compiles. Even with per-block (regional) compilation turned on, the cold warmup died deterministically and the Inductor/AOTAutograd artifacts could not be reused across processes, so every launch paid the full autotune/codegen cost. This PR fixes both: makes the 48 transformer blocks share a single compiled graph, and makes those graphs disk-cache-serializable so a second launch on the same box reloads them.

Changes

• LTX2Transformer3DModel now picks up _compile_conditions from its arch config, so the 48 transformer blocks get compiled individually instead of the whole model as one giant fullgraph. LTX2 was the last major DiT still doing the latter.
• We made the block forward free of per-instance state so all 48 blocks share one compiled graph, as it was otherwise hitting the Dynamo limit at block 17. Moved that check up into the eager _process_transformer_blocks loop
• Made the graphs safe to cache. AOTAutograd's serializer was being tripped and forcing a fresh compile from the lru_cache on the rotary freq-grid helpers, the functools.cache on _sm90_or_newer, and a dtype mismatch in the RMSNorm fallback
• Turned on the caching in the launch script with the paths to cache to

Test Plan

Validated on a B200 (sm100), torch 2.12.0+cu130. A script runs the dreamverse worker's initialize() + warmup(), the streams continuation chunks. Two configs: 576×320 and 1920×1088 at 121 frame chunks (NVFP4, FlashAttention 4), both with max-autotune-no-cudagraphs and the persistent cache.

export TORCHINDUCTOR_CACHE_DIR=$PWD/.cache/inductor TRITON_CACHE_DIR=$PWD/.cache/triton
export TORCHINDUCTOR_FX_GRAPH_CACHE=1 TORCHINDUCTOR_AUTOGRAD_CACHE=1

Cold vs. reload warmup: ran the dreamverse worker's initialize() + warmup() to build the cache from cold, then again in a fresh process pointed at the same cache dirs to verify it loads it from disk.
Warm state: after a single warmup, we streamed chunks and timed each 121 frame 1080p chunk.

Test Results

Cold vs. reload warmup (persistent Inductor/Triton + FX-graph + AOTAutograd cache):
576x320 SDPA autotune: cold 557.5s -> reload 58.3s (9.6x)
1080p FA4 autotune: cold 716.0s -> reload 76.3s (9.4x)
(0 recompile-limit hits, 0 autograd_cache_bypass, 0 failures)

Warm steady-state (1080p, FA4, fully warm continuation chunks):
121-frame 1080p chunk generation ~5.3s median (range 5.24-5.45s), e2e 5.26-5.47s
-> within the 5-7s/chunk target; 0 recompiles on continuation chunks

Residual reload ~76s = Dynamo tracing (uncached on torch 2.12) + model load + first inference

[github-actions]

Welcome to FastVideo! Thanks for your first pull request.

How our CI works:

PRs run a two-tier CI system:

1. Pre-commit — formatting (yapf), linting (ruff), type checking (mypy). Runs immediately on every PR.
2. Fastcheck — core GPU tests (encoders, VAEs, transformers, kernels, unit tests). Runs automatically via Buildkite on relevant file changes (~10-15 min).
3. Full Suite — integration tests, training pipelines, SSIM regression. Runs only when a reviewer adds the ready label.

Before your PR is reviewed:

• pre-commit run --all-files passes locally
• You've added or updated tests for your changes
• The PR description explains what and why

If pre-commit fails, a bot comment will explain how to fix it. Fastcheck and Full Suite results appear in the Checks section below.

Useful links:

• Contributing Guide
• Development Roadmap
• Slack

[mergify]

Merge Protections

🔴 1 of 1 protections blocking · waiting on 👀 reviews and 🤖 CI

	Protection	Waiting on
🔴	PR merge requirements	👀 reviews and 🤖 CI

🔴 PR merge requirements

Waiting for

• #approved-reviews-by>=1
• check-success=full-suite-passed

This rule is failing.

• #approved-reviews-by>=1
• check-success=full-suite-passed
• check-success=fastcheck-passed
• check-success~=pre-commit
• title~=(?i)^\[(feat|feature|bugfix|fix|refactor|perf|ci|doc|docs|misc|chore|kernel|new.?model|skill|skills|infra)\]

[gemini-code-assist]

Code Review

This pull request introduces optimizations to make the codebase compatible with torch.compile, including setting up persistent cache directories in the launch script, avoiding logging and platform queries inside compiled traces, and refactoring model blocks to prevent graph recompilations. The review feedback highlights three important improvements: lazily initializing the device capability check to avoid premature CUDA initialization on GPU 0, implementing a compilation-safe caching mechanism for rotary frequency generation to prevent eager-mode performance regressions, and gracefully handling potential directory creation failures in the launch script.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

Querying device capability at import time (via current_platform.has_device_capability(90)) is a known anti-pattern in PyTorch distributed applications. If this module is imported before torch.cuda.set_device(local_rank) is called (which is very common in multi-GPU training/inference setups), it will initialize CUDA on the default device (device 0). This can cause unexpected memory allocation on GPU 0 (the "GPU 0 memory leak") or even runtime errors. We should initialize _IS_SM90_OR_NEWER lazily on the first call to _sm90_or_newer().

Suggested change

	_IS_SM90_OR_NEWER = current_platform.has_device_capability(90)
	def _check_dropout(dropout_p: float) -> None:
	if dropout_p != 0.0:
	raise NotImplementedError(f"flash_attn.cute does not support dropout (got dropout_p={dropout_p})")
	@functools.cache
	def _sm90_or_newer() -> bool:
	return current_platform.has_device_capability(90)
	# compiled callers only read a boolean nstead of tracing through a memoized platform query.
	return _IS_SM90_OR_NEWER
	_IS_SM90_OR_NEWER = None
	def _check_dropout(dropout_p: float) -> None:
	if dropout_p != 0.0:
	raise NotImplementedError(f"flash_attn.cute does not support dropout (got dropout_p={dropout_p})")
	def _sm90_or_newer() -> bool:
	# compiled callers only read a boolean nstead of tracing through a memoized platform query.
	global _IS_SM90_OR_NEWER
	if _IS_SM90_OR_NEWER is None:
	_IS_SM90_OR_NEWER = current_platform.has_device_capability(90)
	return _IS_SM90_OR_NEWER

[gemini-code-assist]

Removing lru_cache entirely causes a performance regression in eager mode because the frequency grid is re-allocated and re-computed on every single forward pass (denoising step). We can suggest a compilation-safe caching mechanism using a simple dictionary cache that is bypassed when compiling (i.e., when torch.compiler.is_compiling() is active).

_FREQ_GRID_NP_CACHE = {}


def generate_ltx_freq_grid_np(
    positional_embedding_theta: float, positional_embedding_max_pos_count: int, inner_dim: int
) -> torch.Tensor:
    """Generate LTX-2 rotary frequencies with float64 precision.

    Implemented in pure torch and free of Python memoization so it is suitable for callers that may be captured and reused by torch.compile.
    """
    key = (positional_embedding_theta, positional_embedding_max_pos_count, inner_dim)
    if not torch.compiler.is_compiling() and key in _FREQ_GRID_NP_CACHE:
        return _FREQ_GRID_NP_CACHE[key]

    theta = positional_embedding_theta
    start = 1
    end = theta
    n_elem = 2 * positional_embedding_max_pos_count
    pow_indices = torch.pow(
        theta,
        torch.linspace(
            math.log(start) / math.log(theta),
            math.log(end) / math.log(theta),
            inner_dim // n_elem,
            dtype=torch.float64,
        ),
    )
    res = (pow_indices * math.pi / 2).to(dtype=torch.float32)
    if not torch.compiler.is_compiling():
        _FREQ_GRID_NP_CACHE[key] = res
    return res

[gemini-code-assist]

In bash scripts running with set -e, any command that returns a non-zero exit status will cause the script to immediately exit. If mkdir -p fails (for example, due to permission restrictions or a read-only filesystem in containerized or HPC environments), the entire launch script will crash and prevent the backend from starting. Since the persistent cache is an optional performance optimization, we should handle any failure of mkdir -p gracefully so that the backend can still launch and run without caching.

Suggested change

	mkdir -p "${TORCHINDUCTOR_CACHE_DIR}" "${TRITON_CACHE_DIR}"
	mkdir -p "${TORCHINDUCTOR_CACHE_DIR}" "${TRITON_CACHE_DIR}" || echo "[launch-demo] Warning: Could not create cache directories. Proceeding without persistent cache."

[mergify]

Pre-commit checks failed

Hi @Someone45, the pre-commit checks have failed. To fix them locally:

# Install pre-commit if you haven't already
uv pip install pre-commit
pre-commit install

# Run all checks and auto-fix what's possible
pre-commit run --all-files

Common fixes:

• yapf: yapf -i <file> (formatting)
• ruff: ruff check --fix <file> (linting)
• codespell: codespell --write-changes <file> (spelling)

After fixing, commit and push the changes. The checks will re-run automatically.

For future commits, pre-commit will run automatically on changed files before each commit.

[mergify]

This PR has merge conflicts with the base branch. Please rebase:

git fetch origin main
git rebase origin/main
# Resolve any conflicts, then:
git push --force-with-lease