[misc]: address review — stale STG docstring, typos, rename np-less freq-grid helper, add compile-conditions contract test

- ltx2.py: the keep-mask closure docstring still claimed only
  stg_block_idx is perturbed; that gate moved to the eager caller, so
  document the caller contract instead. Note that
  LTX2_PIPELINE_DEBUG_LOG re-specializes compiled graphs per block.
- rename generate_ltx_freq_grid_np -> generate_ltx_freq_grid_float64
  (it no longer uses numpy); update the gemma.py consumer.
- flash_attn_cute.py: fix 'nstead' typo; flash_attn.py: wrap comment
  to 120 cols; launch script: grammar fix.
- new fastvideo/tests/contract/test_ltx2_compile_conditions.py pins
  the class-level _compile_conditions wiring and the is_ltx2_blocks
  matcher so the whole-model fullgraph fallback can't silently return
  (CPU-only, no model instantiation).

apps/dreamverse/scripts/launch/launch_backend_dreamverse.sh

@@ -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.
+# 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}"
+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

fastvideo/attention/backends/flash_attn.py

@@ -180,8 +180,11 @@ class FlashAttentionImpl(AttentionImpl):
         # 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 "
+            # Keep logging (and its Python lru_cache) out of compiled traces
-                                f"bfloat16 for the kernel and restoring on output.")
+            # 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)

fastvideo/attention/utils/flash_attn_cute.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import functools
 from collections.abc import Callable
 
 import torch
@@ -43,15 +42,17 @@ except ImportError:
     _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 instead of tracing through a memoized platform query.
+    return _IS_SM90_OR_NEWER
 
 
 def _use_fa2(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> bool:

fastvideo/models/dits/ltx2.py

@@ -5,14 +5,11 @@ LTX-2 transformer implementation
 
 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
@@ -770,28 +767,30 @@ def _apply_ltx_split_rotary_emb(
     return output
 
 
-@functools.lru_cache(maxsize=5)
+def generate_ltx_freq_grid_float64(
-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 safe
+    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(
+        torch.linspace(
-            np.log(start) / np.log(theta),
+            math.log(start) / math.log(theta),
-            np.log(end) / np.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)
 
 
-@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:
@@ -1012,7 +1011,7 @@ class TransformerArgsPreprocessor:
         x_dtype: torch.dtype,
     ) -> torch.Tensor:
         if self.double_precision_rope:
-            freq_grid_generator = generate_ltx_freq_grid_np
+            freq_grid_generator = generate_ltx_freq_grid_float64
         else:
             freq_grid_generator = generate_ltx_freq_grid_pytorch
         return precompute_ltx_freqs_cis(
@@ -2058,27 +2057,35 @@ class BasicAVTransformerBlock(torch.nn.Module):
             """STG keep-mask for self-attention.
 
             Returns a per-sample multiplier (1 keep, 0 perturb) broadcastable
-            over ``values``. Only the configured ``stg_block_idx`` is perturbed.
+            over ``values``. Applied unconditionally: the eager caller
+            (``_process_transformer_blocks``) enforces that only the
+            configured ``stg_block_idx`` block receives a truthy flag.
             Supports both the LTX-2.0 bool flag and an LTX-2.3 per-sample
             tensor ``[B]`` (1/True == perturb). When skip_flag is False/0 the
             multiplier is all ones, reproducing the un-skipped LTX-2.0 path.
             """
             bsz = values.shape[0]
             keep = torch.ones((bsz, ), device=values.device, dtype=values.dtype)
-            if self.idx == self.stg_block_idx:
+            # The ``idx == stg_block_idx`` gate lives in the eager
-                if torch.is_tensor(skip_flag):
+            # ``_process_transformer_blocks`` caller, not here: reading the
-                    if skip_flag.ndim == 0:
+            # per-instance ``self.idx`` inside this compiled forward forces
-                        perturb = skip_flag.reshape(1).expand(bsz)
+            # torch.compile to specialize a distinct graph per block (48x),
-                    else:
+            # blowing the Dynamo recompile limit and defeating regional
-                        if skip_flag.shape[0] != bsz:
+            # compilation. The caller only passes a truthy ``skip_flag`` to the
-                            raise ValueError(
+            # configured STG block, so applying it unconditionally here is
-                                "Per-sample STG mask batch size mismatch: "
+            # equivalent while keeping the graph identical across all 48 blocks.
-                                f"got {skip_flag.shape[0]}, expected {bsz}")
+            if torch.is_tensor(skip_flag):
-                        perturb = skip_flag.reshape(bsz)
+                if skip_flag.ndim == 0:
-                    keep = 1.0 - perturb.to(device=values.device,
+                    perturb = skip_flag.reshape(1).expand(bsz)
-                                            dtype=values.dtype)
+                else:
-                elif bool(skip_flag):
+                    if skip_flag.shape[0] != bsz:
-                    keep.zero_()
+                        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:
@@ -2277,6 +2284,9 @@ class BasicAVTransformerBlock(torch.nn.Module):
                 1 + ascale_mlp) + ashift_mlp
             ax = ax + self.audio_ff(ax_scaled) * agate_mlp
 
+        # Debug-only: reading ``self.idx`` (and .item() syncs) here means
+        # enabling this env var re-specializes the compiled graph per block,
+        # defeating the shared-graph regional compilation.
         if os.getenv("LTX2_PIPELINE_DEBUG_LOG", "0") == "1":
             video_sum = vx.float().sum().item() if vx is not None else 0.0
             audio_sum = ax.float().sum().item() if ax is not None else 0.0
@@ -2656,8 +2666,10 @@ class LTXModel(torch.nn.Module):
         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
+            # Preserve the original STG block condition in eager code so it does not introduce a block index guard in the compiled forward
-            skip_a_sa = idx in skip_audio_self_attn_block_set
+            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,
@@ -2768,6 +2780,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)

fastvideo/models/encoders/gemma.py

@@ -15,7 +15,7 @@ from fastvideo.models.dits.ltx2 import (
     FeedForward,
     LTXRopeType,
     apply_ltx_rotary_emb,
-    generate_ltx_freq_grid_np,
+    generate_ltx_freq_grid_float64,
     generate_ltx_freq_grid_pytorch,
     precompute_ltx_freqs_cis,
 )
@@ -330,7 +330,7 @@ class Embeddings1DConnector(nn.Module):
         )
         indices_grid = indices_grid[None, None, :]
         freq_grid_generator = (
-            generate_ltx_freq_grid_np
+            generate_ltx_freq_grid_float64
             if self.double_precision_rope
             else generate_ltx_freq_grid_pytorch
         )

fastvideo/tests/contract/test_ltx2_compile_conditions.py

@@ -0,0 +1,35 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Contract test: LTX2's per-block torch.compile wiring stays intact.
+
+``ComposedPipelineBase._compile_with_conditions`` compiles the submodules
+matched by ``model._compile_conditions`` and falls back to compiling the
+whole model as one fullgraph when the list is empty. LTX2 shipped with the
+conditions declared on its arch config but never exposed on the model
+class, so it silently took the whole-model fallback (enormous unportable
+graphs, 10+ minute cold compiles). These assertions pin the class-level
+wiring and the block matcher so a refactor can't reintroduce the fallback
+unnoticed. CPU-only; no model instantiation.
+"""
+
+from fastvideo.configs.models.dits.ltx2 import LTX2VideoConfig, is_ltx2_blocks
+
+
+def test_ltx2_model_exposes_compile_conditions():
+    from fastvideo.models.dits.ltx2 import LTX2Transformer3DModel
+
+    conditions = LTX2Transformer3DModel._compile_conditions
+    assert conditions, (
+        "LTX2Transformer3DModel._compile_conditions is empty: "
+        "torch.compile falls back to a whole-model fullgraph")
+    assert is_ltx2_blocks in conditions
+    assert LTX2VideoConfig()._compile_conditions == conditions
+
+
+def test_is_ltx2_blocks_matches_only_top_level_blocks():
+    assert is_ltx2_blocks("transformer_blocks.0", None)
+    assert is_ltx2_blocks("model.transformer_blocks.47", None)
+    # Submodules and lookalike names must not match — compiling them would
+    # nest compiled regions inside the per-block graphs.
+    assert not is_ltx2_blocks("transformer_blocks", None)
+    assert not is_ltx2_blocks("transformer_blocks.0.attn1", None)
+    assert not is_ltx2_blocks("audio_transformer_blocks.0", None)