Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Someone45 wants to merge 1 commit into
base: main
Choose a base branch
from
+53 −33
Open

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 10 additions & 0 deletions apps/dreamverse/scripts/launch/launch_backend_dreamverse.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,16 @@ export FASTVIDEO_GENERATION_SEGMENT_CAP="${FASTVIDEO_GENERATION_SEGMENT_CAP:-6}"
export FASTVIDEO_PROMPT_AUTO_SLEEP_MS="${FASTVIDEO_PROMPT_AUTO_SLEEP_MS:-120}"
export FASTVIDEO_PROMPT_AUTO_TIMEOUT_MS="${FASTVIDEO_PROMPT_AUTO_TIMEOUT_MS:-1800}"

# Persistent torch.compile cache so warmup pays the full autotune/codegen cost only on the first launch.
# This allows later launches on the same box reload the Inductor FX-graph + AOTAutograd + Triton artifacts from disk
export DREAMVERSE_TORCH_COMPILE_CACHE_ROOT="${DREAMVERSE_TORCH_COMPILE_CACHE_ROOT:-${HOME}/.cache/dreamverse/torch_compile}"
export TORCHINDUCTOR_CACHE_DIR="${TORCHINDUCTOR_CACHE_DIR:-${DREAMVERSE_TORCH_COMPILE_CACHE_ROOT}/inductor}"
export TRITON_CACHE_DIR="${TRITON_CACHE_DIR:-${DREAMVERSE_TORCH_COMPILE_CACHE_ROOT}/triton}"
export TORCHINDUCTOR_FX_GRAPH_CACHE="${TORCHINDUCTOR_FX_GRAPH_CACHE:-1}"
export TORCHINDUCTOR_AUTOGRAD_CACHE="${TORCHINDUCTOR_AUTOGRAD_CACHE:-1}"
mkdir -p "${TORCHINDUCTOR_CACHE_DIR}" "${TRITON_CACHE_DIR}"

@gemini-code-assist gemini-code-assist Bot Jul 14, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

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."

echo "[launch-demo] torch.compile cache: ${DREAMVERSE_TORCH_COMPILE_CACHE_ROOT} (first launch builds it, later launches reload it)"

cd "${DREAMVERSE_ROOT}"

if ! command -v dreamverse-server >/dev/null 2>&1; then
Expand Down
6 changes: 4 additions & 2 deletions fastvideo/attention/backends/flash_attn.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -180,8 +180,10 @@ def forward(
# SP). Cast through bf16 and restore, matching TORCH_SDPA's tolerance.
orig_dtype = query.dtype
if orig_dtype not in (torch.float16, torch.bfloat16):
logger.warning_once(f"FLASH_ATTN received {orig_dtype} inputs; casting to "
f"bfloat16 for the kernel and restoring on output.")
# we keep logging and its Python lru_cache out of compiled traces so that the AOTAutograd cache can serialize.
if not torch.compiler.is_compiling():
logger.warning_once(f"FLASH_ATTN received {orig_dtype} inputs; casting to "
f"bfloat16 for the kernel and restoring on output.")
query = query.to(torch.bfloat16)
key = key.to(torch.bfloat16)
value = value.to(torch.bfloat16)
Expand Down
7 changes: 4 additions & 3 deletions fastvideo/attention/utils/flash_attn_cute.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,5 @@
from __future__ import annotations

import functools
from collections.abc import Callable

import torch
Expand Down Expand Up @@ -43,15 +42,17 @@
_flash_attn_2_func = None
_flash_attn_2_varlen_func = None

_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
Comment on lines +45 to +55

@gemini-code-assist gemini-code-assist Bot Jul 14, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

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



def _use_fa2(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> bool:
Expand Down
63 changes: 35 additions & 28 deletions fastvideo/models/dits/ltx2.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,14 +5,11 @@

from dataclasses import dataclass, replace
from enum import Enum
import functools
import math
import os
from pathlib import Path
from typing import Any, Optional, Tuple, Callable

import numpy as np

import torch
import torch.nn as nn
from einops import rearrange, repeat
Expand Down Expand Up @@ -770,28 +767,29 @@ def _apply_ltx_split_rotary_emb(
return output


@functools.lru_cache(maxsize=5)
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 high-precision numpy."""
"""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.
"""
theta = positional_embedding_theta
start = 1
end = theta
n_elem = 2 * positional_embedding_max_pos_count
pow_indices = np.power(
pow_indices = torch.pow(
theta,
np.linspace(
np.log(start) / np.log(theta),
np.log(end) / np.log(theta),
torch.linspace(
math.log(start) / math.log(theta),
math.log(end) / math.log(theta),
inner_dim // n_elem,
dtype=np.float64,
dtype=torch.float64,
),
)
return torch.tensor(pow_indices * math.pi / 2, dtype=torch.float32)
return (pow_indices * math.pi / 2).to(dtype=torch.float32)
Comment on lines 770 to +790

@gemini-code-assist gemini-code-assist Bot Jul 14, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

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



@functools.lru_cache(maxsize=5)
def generate_ltx_freq_grid_pytorch(
positional_embedding_theta: float, positional_embedding_max_pos_count: int, inner_dim: int
) -> torch.Tensor:
Expand Down Expand Up @@ -2065,20 +2063,26 @@ def _build_attn_keep_mask(
"""
bsz = values.shape[0]
keep = torch.ones((bsz, ), device=values.device, dtype=values.dtype)
if self.idx == self.stg_block_idx:
if torch.is_tensor(skip_flag):
if skip_flag.ndim == 0:
perturb = skip_flag.reshape(1).expand(bsz)
else:
if skip_flag.shape[0] != bsz:
raise ValueError(
"Per-sample STG mask batch size mismatch: "
f"got {skip_flag.shape[0]}, expected {bsz}")
perturb = skip_flag.reshape(bsz)
keep = 1.0 - perturb.to(device=values.device,
dtype=values.dtype)
elif bool(skip_flag):
keep.zero_()
# The ``idx == stg_block_idx`` gate lives in the eager
# ``_process_transformer_blocks`` caller, not here: reading the
# per-instance ``self.idx`` inside this compiled forward forces
# torch.compile to specialize a distinct graph per block (48x),
# blowing the Dynamo recompile limit and defeating regional
# compilation. The caller only passes a truthy ``skip_flag`` to the
# configured STG block, so applying it unconditionally here is
# equivalent while keeping the graph identical across all 48 blocks.
if torch.is_tensor(skip_flag):
if skip_flag.ndim == 0:
perturb = skip_flag.reshape(1).expand(bsz)
else:
if skip_flag.shape[0] != bsz:
raise ValueError(
"Per-sample STG mask batch size mismatch: "
f"got {skip_flag.shape[0]}, expected {bsz}")
perturb = skip_flag.reshape(bsz)
keep = 1.0 - perturb.to(device=values.device, dtype=values.dtype)
elif bool(skip_flag):
keep.zero_()
return keep.view(bsz, *([1] * (values.ndim - 1)))

if run_vx:
Expand Down Expand Up @@ -2656,8 +2660,10 @@ def _process_transformer_blocks(
skip_audio_self_attn_block_set = set(skip_audio_self_attn_blocks or [])

for idx, block in enumerate(self.transformer_blocks):
skip_v_sa = idx in skip_video_self_attn_block_set
skip_a_sa = idx in skip_audio_self_attn_block_set
# Preserve the original STG block condition in eager code so it does not introduce a block index guard in the compiled forward
is_stg_block = idx == block.stg_block_idx
skip_v_sa = is_stg_block and idx in skip_video_self_attn_block_set
skip_a_sa = is_stg_block and idx in skip_audio_self_attn_block_set
video, audio = block(
video=video,
audio=audio,
Expand Down Expand Up @@ -2768,6 +2774,7 @@ class LTX2Transformer3DModel(BaseDiT):
reverse_param_names_mapping = LTX2VideoConfig().reverse_param_names_mapping
lora_param_names_mapping = LTX2VideoConfig().lora_param_names_mapping
_fsdp_shard_conditions = LTX2VideoConfig()._fsdp_shard_conditions
_compile_conditions = LTX2VideoConfig()._compile_conditions

def __init__(self, config: LTX2VideoConfig, hf_config: dict[str, Any]):
super().__init__(config=config, hf_config=hf_config)
Expand Down
Loading