Automated PR - 2026-05-28

README.md

@@ -77,7 +77,7 @@ Download the following models from the [LTX-2.3 HuggingFace repository](https://
 
 * **Use DistilledPipeline** - Fastest inference with only 8 predefined sigmas (8 steps stage 1, 4 steps stage 2)
 * **Enable FP8 quantization** - Enables lower memory footprint: `--quantization fp8-cast` (CLI) or `quantization=QuantizationPolicy.fp8_cast()` (Python). Fp8-cast should be used with bf16 checkpoints, it shall downcast them on the fly. For Hopper GPUs with TensorRT-LLM, use `--quantization fp8-scaled-mm` for FP8 scaled matrix multiplication. Fp8-scaled-mm should be used with fp8 checkpoints.
-* **Install attention optimizations** - Use xFormers (`uv sync --extra xformers`) or [Flash Attention 3](https://github.com/Dao-AILab/flash-attention) for Hopper GPUs
+* **Install attention optimizations** - On datacenter Blackwell GPUs (B200), install FlashAttention 4 manually: `uv pip install 'flash-attn-4==4.0.0b9'` (this specific revision is the one we have verified against torch 2.9.1+cu128; newer betas have known issues on consumer Blackwell). On other CUDA GPUs (including Hopper), use xFormers (`uv sync --extra xformers`).
 * **Use gradient estimation** - Reduce inference steps from 40 to 20-30 while maintaining quality (see [pipeline documentation](packages/ltx-pipelines/README.md#denoising-loop-optimization))
 * **Skip memory cleanup** - If you have sufficient VRAM, disable automatic memory cleanup between stages for faster processing
 * **Choose single-stage pipeline** - Use `TI2VidOneStagePipeline` for faster generation when high resolution isn't required

packages/ltx-core/README.md

@@ -125,39 +125,26 @@ Uses NVIDIA TensorRT-LLM's `cublas_scaled_mm` for efficient FP8 matrix multiplic
 **Usage with QuantizationPolicy:**
 
 ```python
-from ltx_core.quantization import QuantizationPolicy
+from ltx_core.quantization.fp8_scaled_mm import build_policy as build_fp8_scaled_mm_policy
 
-# Dynamic input quantization (no calibration needed)
+# Discovers the layer set from the checkpoint's .weight_scale tensors
-policy = QuantizationPolicy.fp8_scaled_mm()
+policy = build_fp8_scaled_mm_policy("/path/to/checkpoint.safetensors")
-
-# Static input quantization with calibration file
-policy = QuantizationPolicy.fp8_scaled_mm(calibration_amax_path="/path/to/amax.json")
 ```
 
-The policy provides `sd_ops` and `module_ops` that can be passed to the model builder:
+The policy carries `sd_ops`, `module_ops`, and `fuse_rule` that are passed to the model builder:
 
 ```python
+import torch
 from ltx_core.loader import SingleGPUModelBuilder
 
 builder = SingleGPUModelBuilder(
-    model=model,
+    model_class_configurator=MyModelConfigurator,
-    device=device,
+    model_path="/path/to/checkpoint.safetensors",
-    sd_ops=policy.sd_ops,
+    model_sd_ops=policy.sd_ops,
     module_ops=policy.module_ops,
+    fuse_rule=policy.fuse_rule,
 )
-builder.load(checkpoint_path)
+model = builder.build(device=torch.device("cuda"))
-```
-
-**Calibration File Format** (for static input quantization):
-
-```json
-{
-  "amax_values": {
-    "transformer_blocks.0.attn.to_q.input_quantizer": 12.5,
-    "transformer_blocks.0.attn.to_k.input_quantizer": 8.3,
-    ...
-  }
-}
 ```
 
 #### FP8 Cast
@@ -165,7 +152,9 @@ builder.load(checkpoint_path)
 A simpler approach that casts weights to FP8 for storage and upcasts during inference:
 
 ```python
-policy = QuantizationPolicy.fp8_cast()
+from ltx_core.quantization.fp8_cast import build_policy as build_fp8_cast_policy
+
+policy = build_fp8_cast_policy("/path/to/checkpoint.safetensors")
 ```
 
 For complete, production-ready pipeline implementations that combine these building blocks, see the [`ltx-pipelines`](../ltx-pipelines/) package.

packages/ltx-core/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "ltx-core"
-version = "1.1.3"
+version = "1.1.4"
 description = "Core implementation of Lightricks' LTX-2 model"
 readme = "README.md"
 requires-python = ">=3.10"

packages/ltx-core/src/ltx_core/block_streaming/builder.py

@@ -3,7 +3,6 @@
 from __future__ import annotations
 
 import logging
-from collections.abc import Callable
 from dataclasses import dataclass, field, replace
 from typing import Generic
 
@@ -17,13 +16,14 @@ from ltx_core.block_streaming.provider import WeightsProvider
 from ltx_core.block_streaming.source import DiskWeightSource, PinnedWeightSource, WeightSource
 from ltx_core.block_streaming.utils import allocate_layout_views, derive_layout, make_block_key, resolve_attr
 from ltx_core.block_streaming.wrapper import BlockStreamingWrapper
-from ltx_core.loader.fuse_loras import aggregate_lora_products, fuse_lora_weights
+from ltx_core.loader.fuse_loras import FuseRule, bf16_fuse_rule, fuse_lora_weights
 from ltx_core.loader.helpers import create_meta_model, load_state_dict, read_model_config
 from ltx_core.loader.module_ops import ModuleOps
 from ltx_core.loader.primitives import (
     LoraPathStrengthAndSDOps,
     LoraStateDictWithStrength,
     ModelBuilderProtocol,
+    StateDict,
     StateDictLoader,
 )
 from ltx_core.loader.registry import DummyRegistry, Registry
@@ -50,14 +50,13 @@ class StreamingModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType])
         loras: LoRA adapters fused into weights at load time.
         model_loader: Strategy for reading checkpoint metadata.
         registry: Shared cache for loaded state dicts.
+        fuse_rule: Per-policy LoRA merge rule. Defaults to ``bf16_fuse_rule``;
+            use ``fp8_cast_fuse_rule`` for fp8_cast streaming so the pinned
+            buffers receive correctly-quantized weights.
         blocks_attr: Dotted path to the ``nn.ModuleList`` (e.g.
-            ``"velocity_model.transformer_blocks"``).
+            ``"transformer_blocks"``).
         blocks_prefix: State-dict key prefix for block weights
             (e.g. ``"transformer_blocks"``).
-        state_dict_prefix: Wrapper offset prepended to keys when loading into
-            the meta model (e.g. ``"velocity_model."`` when wrapped by ``X0Model``).
-        model_wrapper: Optional callable wrapping the model
-            (e.g. ``X0Model``).
     """
 
     model_class_configurator: type[ModelConfigurator[ModelType]]
@@ -67,12 +66,11 @@ class StreamingModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType])
     loras: tuple[LoraPathStrengthAndSDOps, ...] = field(default_factory=tuple)
     model_loader: StateDictLoader = field(default_factory=SafetensorsModelStateDictLoader)
     registry: Registry = field(default_factory=DummyRegistry)
+    fuse_rule: FuseRule = bf16_fuse_rule
 
     # Streaming-specific
     blocks_attr: str = ""
     blocks_prefix: str = ""
-    state_dict_prefix: str = ""
-    model_wrapper: Callable[[ModelType], nn.Module] | None = None
 
     def with_sd_ops(self, sd_ops: SDOps | None) -> StreamingModelBuilder:
         return replace(self, model_sd_ops=sd_ops)
@@ -83,6 +81,9 @@ class StreamingModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType])
     def with_loras(self, loras: tuple[LoraPathStrengthAndSDOps, ...]) -> StreamingModelBuilder:
         return replace(self, loras=loras)
 
+    def with_fuse_rule(self, fuse_rule: FuseRule) -> StreamingModelBuilder:
+        return replace(self, fuse_rule=fuse_rule)
+
     def model_config(self) -> dict:
         """Read model configuration from the checkpoint metadata."""
         return read_model_config(self.model_path, self.model_loader)
@@ -113,8 +114,6 @@ class StreamingModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType])
 
         config = read_model_config(self.model_path, self.model_loader)
         meta_model: nn.Module = create_meta_model(self.model_class_configurator, config, self.module_ops)
-        if self.model_wrapper is not None:
-            meta_model = self.model_wrapper(meta_model)
         meta_model.eval()
 
         blocks = resolve_attr(meta_model, self.blocks_attr)
@@ -142,7 +141,15 @@ class StreamingModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType])
             target_device,
             reuse_barrier=lambda event: copy_stream.wait_event(event),
         )
-        provider = WeightsProvider(gpu_pool, copy_stream, target_device, source, lora_sources, self.blocks_prefix)
+        provider = WeightsProvider(
+            gpu_pool,
+            copy_stream,
+            target_device,
+            source,
+            lora_sources,
+            self.blocks_prefix,
+            fuse_rule=self.fuse_rule,
+        )
         return BlockStreamingWrapper(
             model=meta_model,
             blocks=blocks,
@@ -190,7 +197,9 @@ class StreamingModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType])
         pinned_blocks = allocate_layout_views(blocks_layout, pin_memory=True)
 
         should_sync = False
-        for key, fused in fuse_lora_weights(model_sd, lora_sd_and_strengths, dtype=None, preserve_input_device=False):
+        for key, fused in fuse_lora_weights(
+            model_sd, lora_sd_and_strengths, fuse_rule=self.fuse_rule, preserve_input_device=False
+        ):
             if key in pinned_blocks:
                 pinned_blocks[key].copy_(fused, non_blocking=True)
                 model_sd.sd[key] = None
@@ -216,7 +225,7 @@ class StreamingModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType])
         }
 
         non_block_sd: dict[str, torch.Tensor] = {
-            self.state_dict_prefix + model_key: model_sd.sd[model_key].to(device=target_device, dtype=dtype)
+            model_key: model_sd.sd[model_key].to(device=target_device, dtype=dtype)
             for _sft_key, model_key in non_block_keys
         }
 
@@ -235,7 +244,7 @@ class StreamingModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType])
         non_block_keys: list[tuple[str, str]],
     ) -> tuple[WeightSource, list[LoraSource]]:
         """Create a DiskWeightSource backed by a DiskBlockReader for lazy loading.
-        Derives the shared pool layout from the meta model's block 0 — this
+        Derives the shared pool layout from the meta model's block 0 - this
         relies on module_ops (e.g. fp8_cast) leaving the meta param dtype in
         sync with the post-sd_ops checkpoint dtype.
         """
@@ -248,8 +257,8 @@ class StreamingModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType])
             target_device,
             dtype,
             sd_ops=self.model_sd_ops,
-            key_prefix=self.state_dict_prefix,
             lora_sources=lora_sources,
+            fuse_rule=self.fuse_rule,
         )
 
         blocks = resolve_attr(meta_model, self.blocks_attr)
@@ -271,28 +280,6 @@ class StreamingModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType])
         source = DiskWeightSource(cpu_pool, block_reader)
         return source, lora_sources
 
-    # ------------------------------------------------------------------
-    # Helpers
-    # ------------------------------------------------------------------
-
-    @staticmethod
-    def _fuse_lora_delta(
-        model_key: str,
-        tensor: torch.Tensor,
-        lora_sources: list[LoraSource],
-    ) -> torch.Tensor:
-        """Add all matching LoRA deltas to *tensor* in-place via ``addmm_``."""
-        if not lora_sources or not model_key.endswith(".weight"):
-            return tensor
-        prefix = model_key[: -len(".weight")]
-        products = (
-            ab
-            for ab in (s.get_ab(prefix, device=tensor.device, dtype=tensor.dtype) for s in lora_sources)
-            if ab is not None
-        )
-        aggregate_lora_products(products, out=tensor)
-        return tensor
-
     @staticmethod
     @torch.inference_mode()
     def _load_non_block_weights(
@@ -302,21 +289,36 @@ class StreamingModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType])
         device: torch.device,
         dtype: torch.dtype,
         sd_ops: SDOps | None = None,
-        key_prefix: str = "",
         lora_sources: list[LoraSource] | None = None,
+        fuse_rule: FuseRule = bf16_fuse_rule,
     ) -> None:
-        """Load non-block weights into *model* on *device*."""
+        """Load non-block weights into *model* on *device*, applying ``sd_ops`` and fusing LoRAs.
-        state_dict: dict[str, torch.Tensor] = {}
+        Fusion goes through :func:`fuse_lora_weights` under *fuse_rule* so the
-        sources = lora_sources or []
+        bf16 rounding pattern matches the non-streaming and block-streaming
+        paths, and any quantization-specific fuse rule the builder configures
+        is honored here as well.
+        """
+        non_block_sd: dict[str, torch.Tensor] = {}
         for sft_key, model_key in non_block_keys:
             tensor = reader.get_tensor(sft_key).to(device=device, dtype=dtype)
-            tensor = StreamingModelBuilder._fuse_lora_delta(model_key, tensor, sources)
             if sd_ops is not None:
                 for kv in sd_ops.apply_to_key_value(model_key, tensor):
-                    state_dict[key_prefix + kv.new_key] = kv.new_value
+                    non_block_sd[kv.new_key] = kv.new_value
-                continue
+            else:
-            state_dict[key_prefix + model_key] = tensor
+                non_block_sd[model_key] = tensor
-        model.load_state_dict(state_dict, strict=False, assign=True)
+
+        if lora_sources:
+            lora_sd_and_strengths = [src.as_state_dict_with_strength() for src in lora_sources]
+            non_block_state = StateDict(sd=non_block_sd, device=device, size=0, dtype={dtype})
+            for key, fused in fuse_lora_weights(
+                non_block_state,
+                lora_sd_and_strengths,
+                fuse_rule=fuse_rule,
+                preserve_input_device=True,
+            ):
+                non_block_sd[key] = fused
+
+        model.load_state_dict(non_block_sd, strict=False, assign=True)
 
 
 def _scan_checkpoint_keys(

packages/ltx-core/src/ltx_core/block_streaming/disk.py

@@ -9,6 +9,7 @@ import torch
 
 from ltx_core.block_streaming.utils import allocate_layout_views, make_block_key
 from ltx_core.loader.fuse_loras import LoraProduct
+from ltx_core.loader.primitives import LoraStateDictWithStrength, StateDict
 from ltx_core.loader.sd_ops import SDOps
 
 _SAFETENSORS_DTYPE_TO_TORCH: dict[str, torch.dtype] = {
@@ -125,6 +126,22 @@ class LoraSource:
                 b_view.copy_(handle.get_tensor(b_keys[prefix]))
                 self._pinned_ab[prefix] = (a_view, b_view)
 
+    def as_state_dict_with_strength(self) -> LoraStateDictWithStrength:
+        """Return a :class:`LoraStateDictWithStrength` view of the pinned A/B factors.
+        Lets :func:`fuse_lora_weights` consume disk-streaming LoRAs without
+        re-reading the safetensors file or re-applying ``sd_ops``.
+        """
+        sd: dict[str, torch.Tensor] = {}
+        for prefix, (a, b) in self._pinned_ab.items():
+            sd[f"{prefix}.lora_A.weight"] = a
+            sd[f"{prefix}.lora_B.weight"] = b
+        size = sum(t.numel() * t.element_size() for t in sd.values())
+        dtypes = {t.dtype for t in sd.values()}
+        return LoraStateDictWithStrength(
+            StateDict(sd=sd, device=torch.device("cpu"), size=size, dtype=dtypes),
+            self.strength,
+        )
+
     def get_ab(
         self,
         param_prefix: str,

packages/ltx-core/src/ltx_core/block_streaming/provider.py

@@ -9,8 +9,10 @@ import torch
 from ltx_core.block_streaming.disk import LoraSource
 from ltx_core.block_streaming.pool import WeightPool
 from ltx_core.block_streaming.source import WeightSource
-from ltx_core.block_streaming.utils import FP8_DTYPES
+from ltx_core.loader.fuse_loras import FuseRule, aggregate_lora_products, bf16_fuse_rule
-from ltx_core.loader.fuse_loras import aggregate_lora_products, fuse_cast_fp8_weight
+from ltx_core.loader.primitives import StateDict
+
+_EMPTY_STATE_DICT = StateDict(sd={}, device=torch.device("cpu"), size=0, dtype=set())
 
 
 def _contiguous_byte_view(weights: dict[str, torch.Tensor]) -> torch.Tensor | None:
@@ -43,6 +45,8 @@ class WeightsProvider:
         source: Pinned CPU weight source.
         lora_sources: LoRA adapters fused on H2D copy.
         blocks_prefix: State-dict prefix for LoRA key matching.
+        fuse_rule: Per-policy LoRA merge rule (must be streaming-compatible:
+            no companion-key emission). Defaults to ``bf16_fuse_rule``.
     """
 
     def __init__(
@@ -53,6 +57,7 @@ class WeightsProvider:
         source: WeightSource,
         lora_sources: list[LoraSource] | None = None,
         blocks_prefix: str = "",
+        fuse_rule: FuseRule = bf16_fuse_rule,
     ) -> None:
         self._copy_stream = copy_stream
         self._pool = pool
@@ -62,6 +67,7 @@ class WeightsProvider:
         self._source = source
         self._lora_sources = lora_sources or []
         self._blocks_prefix = blocks_prefix
+        self._fuse_rule = fuse_rule
 
     def get(self, idx: int) -> dict[str, torch.Tensor]:
         """Return GPU weights for block *idx*. Does H2D copy on miss."""
@@ -126,22 +132,19 @@ class WeightsProvider:
         return len(self._cache)
 
     def _fuse_block_loras(self, idx: int, weights: dict[str, torch.Tensor]) -> None:
-        """Fuse LoRA deltas directly into GPU block weights."""
+        """Fuse LoRA deltas directly into GPU block weights via ``fuse_rule``."""
+        agg_dtype = self._fuse_rule.aggregation_dtype
         for name, tensor in weights.items():
             if not name.endswith(".weight"):
                 continue
             prefix = f"{self._blocks_prefix}.{idx}.{name}".removesuffix(".weight")
-            is_fp8 = tensor.dtype in FP8_DTYPES
-            agg_dtype = torch.bfloat16 if is_fp8 else tensor.dtype
             products = (
                 ab
                 for ab in (s.get_ab(prefix, device=self._target_device, dtype=agg_dtype) for s in self._lora_sources)
                 if ab is not None
             )
-            aggregated = aggregate_lora_products(products, agg_dtype)
+            deltas = aggregate_lora_products(products, agg_dtype)
-            if aggregated is None:
+            if deltas is None:
                 continue
-            if is_fp8:
+            fused = self._fuse_rule(name, tensor, deltas, _EMPTY_STATE_DICT)
-                tensor.copy_(fuse_cast_fp8_weight(aggregated, tensor, tensor.dtype))
+            tensor.copy_(fused[name])
-            else:
-                tensor.add_(aggregated)

packages/ltx-core/src/ltx_core/loader/attention_ops.py

@@ -0,0 +1,42 @@
+"""Builder ops for swapping attention backends on a meta model before load."""
+
+import torch
+
+from ltx_core.loader.module_ops import ModuleOps
+from ltx_core.model.transformer.attention import (
+    Attention,
+    AttentionCallable,
+    AttentionFunction,
+    MaskedAttentionCallable,
+    MaskedAttentionFunction,
+)
+
+
+def set_attention_module_op(
+    attention: AttentionFunction | AttentionCallable | None = None,
+    masked_attention: MaskedAttentionFunction | MaskedAttentionCallable | None = None,
+) -> ModuleOps:
+    """Build a ``ModuleOps`` that overrides the attention callables on every
+    ``Attention`` submodule of a model. Applied via ``create_meta_model`` so
+    the meta model is mutated before weight loading. Matcher returns False
+    for models with no ``Attention`` submodules, so the op is a no-op there.
+    Either or both slots may be supplied; *None* leaves that slot untouched.
+    """
+    fn = attention.to_callable() if isinstance(attention, AttentionFunction) else attention
+    masked_fn = (
+        masked_attention.to_callable() if isinstance(masked_attention, MaskedAttentionFunction) else masked_attention
+    )
+
+    def matcher(model: torch.nn.Module) -> bool:
+        return any(isinstance(m, Attention) for m in model.modules())
+
+    def mutator(model: torch.nn.Module) -> torch.nn.Module:
+        for module in model.modules():
+            if isinstance(module, Attention):
+                if fn is not None:
+                    module.attention_function = fn
+                if masked_fn is not None:
+                    module.masked_attention_function = masked_fn
+        return model
+
+    return ModuleOps(name="set_attention_backend", matcher=matcher, mutator=mutator)

packages/ltx-core/src/ltx_core/loader/fuse_loras.py

@@ -1,12 +1,10 @@
-from collections.abc import Iterable, Iterator
+from collections.abc import Callable, Iterable, Iterator
+from dataclasses import dataclass
 from typing import NamedTuple
 
 import torch
 
-from ltx_core.loader.kernels import TRITON_AVAILABLE
 from ltx_core.loader.primitives import LoraStateDictWithStrength, StateDict
-from ltx_core.quantization.fp8_cast import fused_add_round_launch
-from ltx_core.quantization.fp8_scaled_mm import quantize_weight_to_fp8_per_tensor
 
 
 class LoraProduct(NamedTuple):
@@ -17,6 +15,62 @@ class LoraProduct(NamedTuple):
     strength: float
 
 
+#: Signature for a fuse callable used by :class:`FuseRule`.
+#:
+#: Args:
+#:     key: The state-dict key being fused (e.g. ``"...layers.0.attn.q.weight"``).
+#:     weight: The current value at ``key`` from ``model_sd``, on the fusion device.
+#:     deltas: The pre-aggregated LoRA delta for ``key``, in ``aggregation_dtype``.
+#:     model_sd: The full state dict, for rules that need companion keys
+#:         (e.g. an existing ``.weight_scale``).
+#:
+#: Returns a dict of state-dict keys to overwrite -- at minimum ``{key: new_weight}``,
+#: plus any companion keys (e.g. an updated ``.weight_scale``) the policy needs to
+#: keep in sync.
+FuseFn = Callable[[str, torch.Tensor, torch.Tensor, StateDict], dict[str, torch.Tensor]]
+
+
+@dataclass(frozen=True)
+class FuseRule:
+    """Fuse an aggregated LoRA delta into one weight key.
+    Each policy supplies its own rule (see ``QuantizationPolicy.fuse_rule``);
+    ``fuse_lora_weights`` is policy-agnostic boilerplate around it.
+    Attributes:
+        aggregation_dtype: Dtype callers must pre-aggregate LoRA deltas in
+            before invoking the rule.
+        fuse_fn: Callable that applies the pre-aggregated deltas to the weight
+            (and any companion keys) and returns a dict of keys to overwrite —
+            at minimum ``{key: new_weight}``, plus any companion keys (e.g. an
+            updated ``.weight_scale`` for scaled-FP8 layouts) the policy needs
+            to keep in sync.
+    """
+
+    aggregation_dtype: torch.dtype
+    fuse_fn: FuseFn
+
+    def __call__(
+        self,
+        key: str,
+        weight: torch.Tensor,
+        deltas: torch.Tensor,
+        model_sd: StateDict,
+    ) -> dict[str, torch.Tensor]:
+        return self.fuse_fn(key, weight, deltas, model_sd)
+
+
+def _bf16_fuse(
+    key: str,
+    weight: torch.Tensor,
+    deltas: torch.Tensor,
+    model_sd: StateDict,  # noqa: ARG001
+) -> dict[str, torch.Tensor]:
+    deltas.add_(weight)
+    return {key: deltas.to(dtype=weight.dtype)}
+
+
+bf16_fuse_rule = FuseRule(aggregation_dtype=torch.bfloat16, fuse_fn=_bf16_fuse)
+
+
 def _get_device() -> torch.device:
     if torch.cuda.is_available():
         return torch.device("cuda", torch.cuda.current_device())
@@ -25,19 +79,16 @@ def _get_device() -> torch.device:
 
 def aggregate_lora_products(
     products: Iterable[LoraProduct],
-    dtype: torch.dtype | None = None,
+    dtype: torch.dtype,
-    *,
-    out: torch.Tensor | None = None,
 ) -> torch.Tensor | None:
     """Accumulate ``sum((B * strength) @ A)`` across :class:`LoraProduct` items.
-    If ``out`` is provided, ``addmm_`` accumulates directly into it — caller
+    The first product materializes a freshly-allocated aggregator via
-    ensures A/B dtypes and devices match ``out``. Otherwise the first product
+    ``torch.matmul(B * strength, A).to(dtype)`` -- preserving the
-    materializes the ``(out, in)``-shape aggregator at ``dtype``; subsequent
+    ``(B * strength) @ A`` rounding pattern. Subsequent products use
-    products use ``addmm_`` to avoid allocating the full intermediate delta.
+    ``addmm_`` to avoid allocating the full intermediate delta.
-    Returns ``out`` (or the new aggregator), or ``None`` if ``products`` was empty
+    Returns the aggregator, or ``None`` if ``products`` was empty.
-    and ``out`` was not given.
     """
-    aggregated = out
+    aggregated: torch.Tensor | None = None
     for product in products:
         if aggregated is None:
             aggregated = torch.matmul(product.b * product.strength, product.a).to(dtype=dtype)
@@ -46,62 +97,34 @@ def aggregate_lora_products(
     return aggregated
 
 
-def fuse_cast_fp8_weight(
-    delta_bf16: torch.Tensor,
-    weight_fp8: torch.Tensor,
-    target_dtype: torch.dtype,
-) -> torch.Tensor:
-    """Return ``(delta_bf16 + dequantize(weight_fp8)).to(target_dtype)``.
-    CUDA with Triton uses stochastic rounding; otherwise uses a deterministic bf16 add.
-    ``delta_bf16`` is the bf16 accumulator and is mutated in place.
-    """
-    if delta_bf16.dtype != torch.bfloat16:
-        raise ValueError(f"delta_bf16 must be bfloat16, got {delta_bf16.dtype}")
-    if str(weight_fp8.device).startswith("cuda") and TRITON_AVAILABLE:
-        fused_add_round_launch(delta_bf16, weight_fp8, seed=0)
-    else:
-        delta_bf16.add_(weight_fp8.to(dtype=torch.bfloat16))
-    return delta_bf16.to(dtype=target_dtype)
-
-
 def fuse_lora_weights(
     model_sd: StateDict,
     lora_sd_and_strengths: list[LoraStateDictWithStrength],
-    dtype: torch.dtype | None = None,
+    fuse_rule: FuseRule = bf16_fuse_rule,
     preserve_input_device: bool = True,
 ) -> Iterator[tuple[str, torch.Tensor]]:
     """Yield ``(key, fused_tensor)`` for each weight modified by at least one LoRA.
-    For scaled-FP8 weights, this includes both the updated ``.weight`` tensor
+    The fusion math is delegated to ``fuse_rule``.
-    and its corresponding ``.weight_scale`` tensor.
+    Output dtypes are the rule's responsibility.
     When ``preserve_input_device`` is False, fused tensors are yielded on the device
     used for fusion; caller is responsible for moving them to their final
     destination.
     """
-    for key, original_weight in model_sd.sd.items():
+    fusion_device = _get_device()
-        if original_weight is None or key.endswith(".weight_scale"):
+    for key in _affected_weight_keys(lora_sd_and_strengths):
+        original_weight = model_sd.sd.get(key)
+        if original_weight is None:
             continue
-        original_device = original_weight.device
-        weight = original_weight.to(device=_get_device())
-        target_dtype = dtype if dtype is not None else weight.dtype
-        deltas_dtype = target_dtype if target_dtype not in [torch.float8_e4m3fn, torch.float8_e5m2] else torch.bfloat16
 
-        deltas = _aggregate_deltas(lora_sd_and_strengths, key, deltas_dtype, weight.device)
+        products = _products_for_sd_key(lora_sd_and_strengths, key, fuse_rule.aggregation_dtype, fusion_device)
+        deltas = aggregate_lora_products(products, fuse_rule.aggregation_dtype)
         if deltas is None:
             continue
 
-        scale_key = key.replace(".weight", ".weight_scale") if key.endswith(".weight") else None
+        original_device = original_weight.device
-        is_scaled_fp8 = scale_key is not None and scale_key in model_sd.sd
+        weight = original_weight.to(device=fusion_device)
 
-        if weight.dtype == torch.float8_e4m3fn:
+        fused = fuse_rule(key, weight, deltas, model_sd)
-            if is_scaled_fp8:
-                fused = _fuse_delta_with_scaled_fp8(deltas, weight, key, scale_key, model_sd)
-            else:
-                fused = {key: fuse_cast_fp8_weight(deltas, weight, target_dtype)}
-        elif weight.dtype == torch.bfloat16:
-            deltas.add_(weight)
-            fused = {key: deltas.to(dtype=target_dtype)}
-        else:
-            raise ValueError(f"Unsupported dtype: {weight.dtype}")
 
         for k, v in fused.items():
             yield k, v.to(device=original_device) if preserve_input_device else v
@@ -110,53 +133,46 @@ def fuse_lora_weights(
 def apply_loras(
     model_sd: StateDict,
     lora_sd_and_strengths: list[LoraStateDictWithStrength],
-    dtype: torch.dtype | None = None,
+    fuse_rule: FuseRule = bf16_fuse_rule,
     destination_sd: StateDict | None = None,
 ) -> StateDict:
     """Fuse LoRAs into ``model_sd`` and place the results in ``destination_sd``.
     When ``destination_sd`` is provided, the fused tensors are placed directly into it.
     """
+    fused_iter = fuse_lora_weights(
+        model_sd,
+        lora_sd_and_strengths,
+        fuse_rule=fuse_rule,
+    )
     if destination_sd is not None:
-        for key, fused in fuse_lora_weights(model_sd, lora_sd_and_strengths, dtype):
+        for key, fused in fused_iter:
             destination_sd.sd[key] = fused
         return destination_sd
 
-    fused = dict(fuse_lora_weights(model_sd, lora_sd_and_strengths, dtype))
+    fused = dict(fused_iter)
     sd = {k: (fused[k] if k in fused else v.clone()) for k, v in model_sd.sd.items()}
     return StateDict(sd, model_sd.device, model_sd.size, model_sd.dtype)
 
 
-def _aggregate_deltas(
+def _affected_weight_keys(lora_sd_and_strengths: list[LoraStateDictWithStrength]) -> set[str]:
-    lora_sd_and_strengths: list[LoraStateDictWithStrength], key: str, dtype: torch.dtype, device: torch.device
+    """Return the set of ``.weight`` keys touched by at least one LoRA in the list."""
-) -> torch.Tensor | None:
+    suffix = ".lora_A.weight"
+    return {k[: -len(suffix)] + ".weight" for lsd, _ in lora_sd_and_strengths for k in lsd.sd if k.endswith(suffix)}
+
+
+def _products_for_sd_key(
+    lora_sd_and_strengths: list[LoraStateDictWithStrength],
+    key: str,
+    dtype: torch.dtype,
+    device: torch.device,
+) -> Iterator[LoraProduct]:
+    """Yield :class:`LoraProduct` items matching *key* across state-dict-backed LoRAs."""
     prefix = key[: -len(".weight")]
     key_a = f"{prefix}.lora_A.weight"
     key_b = f"{prefix}.lora_B.weight"
-
-    def _ab_products() -> Iterator[LoraProduct]:
     for lsd, coef in lora_sd_and_strengths:
         if key_a not in lsd.sd or key_b not in lsd.sd:
             continue
         a = lsd.sd[key_a].to(device=device, dtype=dtype, non_blocking=True)
         b = lsd.sd[key_b].to(device=device, dtype=dtype, non_blocking=True)
         yield LoraProduct(a, b, coef)
-
-    return aggregate_lora_products(_ab_products(), dtype)
-
-
-def _fuse_delta_with_scaled_fp8(
-    deltas: torch.Tensor,
-    weight: torch.Tensor,
-    key: str,
-    scale_key: str,
-    model_sd: StateDict,
-) -> dict[str, torch.Tensor]:
-    """Dequantize scaled FP8 weight, add LoRA delta, and re-quantize."""
-    weight_scale = model_sd.sd[scale_key]
-
-    original_weight = weight.to(torch.float32) * weight_scale
-
-    new_weight = original_weight + deltas.to(torch.float32)
-
-    new_fp8_weight, new_weight_scale = quantize_weight_to_fp8_per_tensor(new_weight)
-    return {key: new_fp8_weight, scale_key: new_weight_scale}

packages/ltx-core/src/ltx_core/loader/primitives.py

@@ -10,6 +10,7 @@ from ltx_core.loader.sd_ops import SDOps
 from ltx_core.model.model_protocol import ModelType
 
 if TYPE_CHECKING:
+    from ltx_core.loader.fuse_loras import FuseRule
     from ltx_core.loader.registry import Registry
 
 
@@ -110,6 +111,10 @@ class ModelBuilderProtocol(BuilderProtocol[ModelType], Protocol[ModelType]):
         """Return a copy of this builder that loads LoRA weights onto the given device."""
         ...
 
+    def with_fuse_rule(self, fuse_rule: "FuseRule") -> "ModelBuilderProtocol[ModelType]":
+        """Return a copy of this builder with the given LoRA fuse rule (e.g. from a quantization policy)."""
+        ...
+
     def build(
         self, device: torch.device | None = None, dtype: torch.dtype | None = None, **kwargs: object
     ) -> ModelType:

packages/ltx-core/src/ltx_core/loader/single_gpu_model_builder.py

@@ -5,7 +5,7 @@ from typing import Generic
 import torch
 from torch import nn
 
-from ltx_core.loader.fuse_loras import apply_loras
+from ltx_core.loader.fuse_loras import FuseRule, apply_loras, bf16_fuse_rule
 from ltx_core.loader.helpers import create_meta_model, load_state_dict, read_model_config
 from ltx_core.loader.module_ops import ModuleOps
 from ltx_core.loader.primitives import (
@@ -46,6 +46,7 @@ def _load_model_weights(
     dtype: torch.dtype | None,
     model_sd_ops: SDOps | None = None,
     lora_load_device: torch.device | None = None,
+    fuse_rule: FuseRule = bf16_fuse_rule,
 ) -> None:
     """Load base weights and fuse LoRAs into *meta_model* in-place."""
     if lora_load_device is None:
@@ -57,7 +58,7 @@ def _load_model_weights(
     if not lora_strengths or (min(lora_strengths) == 0 and max(lora_strengths) == 0):
         sd = model_sd.sd
         if dtype is not None:
-            sd = {key: value.to(dtype=dtype) for key, value in model_sd.sd.items()}
+            sd = {key: value.to(dtype=dtype) for key, value in sd.items()}
         meta_model.load_state_dict(sd, strict=False, assign=True)
         return
 
@@ -68,10 +69,13 @@ def _load_model_weights(
     final_sd = apply_loras(
         model_sd=model_sd,
         lora_sd_and_strengths=lora_sd_and_strengths,
-        dtype=dtype,
+        fuse_rule=fuse_rule,
         destination_sd=model_sd if isinstance(registry, DummyRegistry) else None,
     )
-    meta_model.load_state_dict(final_sd.sd, strict=False, assign=True)
+    fused_sd = final_sd.sd
+    if dtype is not None:
+        fused_sd = {key: value.to(dtype=dtype) for key, value in fused_sd.items()}
+    meta_model.load_state_dict(fused_sd, strict=False, assign=True)
 
 
 @dataclass(frozen=True)
@@ -91,6 +95,7 @@ class SingleGPUModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType],
             ``torch.device("cpu")``, which keeps LoRA weights in CPU memory and transfers them to
             the target GPU sequentially during fusion, reducing peak GPU memory usage compared to
             loading all LoRA weights directly onto the GPU at once.
+        fuse_rule: Per-policy LoRA merge rule. Defaults to ``bf16_fuse_rule``;
     """
 
     model_class_configurator: type[ModelConfigurator[ModelType]]
@@ -101,6 +106,7 @@ class SingleGPUModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType],
     model_loader: StateDictLoader = field(default_factory=SafetensorsModelStateDictLoader)
     registry: Registry = field(default_factory=DummyRegistry)
     lora_load_device: torch.device = field(default_factory=lambda: torch.device("cpu"))
+    fuse_rule: FuseRule = bf16_fuse_rule
 
     def lora(self, lora_path: str, strength: float, sd_ops: SDOps) -> "SingleGPUModelBuilder":
         return replace(self, loras=(*self.loras, LoraPathStrengthAndSDOps(lora_path, strength, sd_ops)))
@@ -120,6 +126,9 @@ class SingleGPUModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType],
     def with_lora_load_device(self, device: torch.device) -> "SingleGPUModelBuilder":
         return replace(self, lora_load_device=device)
 
+    def with_fuse_rule(self, fuse_rule: FuseRule) -> "SingleGPUModelBuilder":
+        return replace(self, fuse_rule=fuse_rule)
+
     def model_config(self) -> dict:
         return read_model_config(self.model_path, self.model_loader)
 
@@ -158,5 +167,6 @@ class SingleGPUModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType],
             dtype=dtype,
             model_sd_ops=self.model_sd_ops,
             lora_load_device=self.lora_load_device,
+            fuse_rule=self.fuse_rule,
         )
         return self._return_model(meta_model, device)

packages/ltx-core/src/ltx_core/modality_tiling.py

@@ -21,11 +21,14 @@ from ltx_core.types import VideoLatentShape
 @dataclass(frozen=True)
 class TilingContext:
     """Opaque context produced by :meth:`VideoModalityTilingHelper.tile_modality`.
-    Carries the token-level keep mask and per-conditioning-token blend
+    Carries the token-level keep indices and per-conditioning-token blend
     weights needed by :meth:`~VideoModalityTilingHelper.blend`.
     """
 
-    keep_mask: torch.Tensor
+    keep_indices: torch.Tensor
+    """``(num_kept,)`` int64 — sorted indices of tokens the tile processes."""
+    num_total_tokens: int
+    """Total number of tokens in the full (untiled) sequence."""
     cond_blend_weights: torch.Tensor | None
     """``(num_kept_cond,)`` — weight for each kept conditioning token,
     equal to ``1 / num_tiles_that_keep_this_token``.  ``None`` when
@@ -81,14 +84,32 @@ class VideoModalityTilingHelper:
             A ``(tiled_modality, context)`` tuple.  Pass *context* to
             :meth:`blend` together with the model output.
         """
-        keep_mask = self._keep_mask(modality, tile)
+        device = modality.positions.device
+        gen_indices = self._generated_token_indices(tile, device=device)
+        num_total = modality.latent.shape[1]
+
+        cond_blend_weights: torch.Tensor | None = None
+        if num_total > self._num_generated_tokens:
+            keep_per_tile_cond = self._all_tiles_cond_keep(modality)  # (num_tiles, num_cond) bool
+            tile_idx = next((i for i, t in enumerate(self._tiles) if t.in_coords == tile.in_coords), None)
+            if tile_idx is None:
+                raise ValueError(
+                    f"Tile with in_coords={tile.in_coords} is not in this helper's tile set; "
+                    f"pass a tile obtained from `helper.tiles`."
+                )
+            my_cond_keep = keep_per_tile_cond[tile_idx]
+            cond_indices = self._num_generated_tokens + my_cond_keep.nonzero(as_tuple=False).squeeze(1)
+            keep_indices = torch.cat([gen_indices, cond_indices])
+            total_keepers = keep_per_tile_cond.sum(dim=0).float()  # (num_cond,)
+            cond_blend_weights = 1.0 / total_keepers[my_cond_keep]
+        else:
+            keep_indices = gen_indices
 
         tile_attention_mask = None
         if modality.attention_mask is not None:
-            keep_indices = keep_mask.nonzero(as_tuple=False).squeeze(1)
             tile_attention_mask = modality.attention_mask[:, keep_indices, :][:, :, keep_indices]
 
-        positions = modality.positions[:, :, keep_mask, :]
+        positions = modality.positions[:, :, keep_indices, :]
         if normalize_positions:
             num_tile_gen = self._tile_generated_token_count(tile)
             gen_pos = positions[:, :, :num_tile_gen, :]  # (B, 3, num_tile_gen, 2)
@@ -97,26 +118,15 @@ class VideoModalityTilingHelper:
 
         tiled = replace(
             modality,
-            latent=modality.latent[:, keep_mask, :],
+            latent=modality.latent[:, keep_indices, :],
-            timesteps=modality.timesteps[:, keep_mask],
+            timesteps=modality.timesteps[:, keep_indices],
             positions=positions,
             attention_mask=tile_attention_mask,
         )
 
-        cond_blend_weights = None
+        return tiled, TilingContext(
-        num_total = modality.latent.shape[1]
+            keep_indices=keep_indices, num_total_tokens=num_total, cond_blend_weights=cond_blend_weights
-        if num_total > self._num_generated_tokens:
+        )
-            cond_keep = keep_mask[self._num_generated_tokens :]
-            # Count how many tiles keep each conditioning token.
-            cond_counts = torch.zeros(cond_keep.sum(), dtype=torch.float32)
-            for t in self._tiles:
-                other_mask = self._keep_mask(modality, t)
-                other_cond = other_mask[self._num_generated_tokens :]
-                # Map other tile's kept cond tokens into this tile's kept subset.
-                cond_counts += other_cond[cond_keep].float()
-            cond_blend_weights = 1.0 / cond_counts
-
-        return tiled, TilingContext(keep_mask=keep_mask, cond_blend_weights=cond_blend_weights)
 
     # -- blend -------------------------------------------------------------
 
@@ -147,9 +157,9 @@ class VideoModalityTilingHelper:
         """
         batch, _, dim = tile_to_blend.shape
         num_tile_gen = self._tile_generated_token_count(tile)
-        gen_indices = self._generated_token_indices(tile)
+        gen_indices = self._generated_token_indices(tile, device=tile_to_blend.device)
 
-        num_total_tokens = context.keep_mask.shape[0]
+        num_total_tokens = context.num_total_tokens
         expected_shape = (batch, num_total_tokens, dim)
 
         if output is not None:
@@ -168,8 +178,7 @@ class VideoModalityTilingHelper:
         # Scatter kept conditioning tokens, weighted by 1/N where N is
         # the number of tiles that keep each token (so they sum to 1).
         if num_total_tokens > self._num_generated_tokens and context.cond_blend_weights is not None:
-            cond_keep = context.keep_mask[self._num_generated_tokens :]
+            cond_indices = context.keep_indices[context.keep_indices >= self._num_generated_tokens]
-            cond_indices = self._num_generated_tokens + cond_keep.nonzero(as_tuple=False).squeeze(1)
             weights = context.cond_blend_weights.to(device=tile_to_blend.device, dtype=tile_to_blend.dtype)
             result[:, cond_indices, :] += tile_to_blend[:, num_tile_gen:, :] * weights[None, :, None]
 
@@ -189,46 +198,42 @@ class VideoModalityTilingHelper:
         )
         return self._patchifier.get_token_count(tile_shape)
 
-    def _generated_token_indices(self, tile: Tile) -> torch.Tensor:
+    def _generated_token_indices(self, tile: Tile, device: torch.device | None = None) -> torch.Tensor:
         """Flat token indices of *tile*'s generated tokens in the full sequence."""
         frame_slice, height_slice, width_slice = tile.in_coords
-        f = torch.arange(frame_slice.start, frame_slice.stop)
+        f = torch.arange(frame_slice.start, frame_slice.stop, device=device)
-        h = torch.arange(height_slice.start, height_slice.stop)
+        h = torch.arange(height_slice.start, height_slice.stop, device=device)
-        w = torch.arange(width_slice.start, width_slice.stop)
+        w = torch.arange(width_slice.start, width_slice.stop, device=device)
         return (
             f[:, None, None] * self._latent_shape.height * self._latent_shape.width
             + h[None, :, None] * self._latent_shape.width
             + w[None, None, :]
         ).reshape(-1)
 
-    def _keep_mask(self, modality: Modality, tile: Tile) -> torch.Tensor:
+    def _all_tiles_cond_keep(self, modality: Modality) -> torch.Tensor:
-        """Boolean mask ``(num_total_tokens,)`` — True for tokens the tile processes.
+        """Vectorized (num_tiles, num_cond) bool: which tiles keep each conditioning token.
-        Generated tokens are selected by grid position.  Conditioning
+        A conditioning token is kept by a tile when its ``[start, end)`` interval
-        tokens are kept when their ``[start, end)`` intervals overlap
+        overlaps the tile in all three dimensions, or when it has a negative time
-        the tile in all three dimensions, or when they have a negative
+        coordinate (reference token).
-        time coordinate (reference tokens).
         """
-        num_total = modality.latent.shape[1]
-        mask = torch.zeros(num_total, dtype=torch.bool)
-
-        gen_indices = self._generated_token_indices(tile)
-        mask[gen_indices] = True
-
-        if num_total > self._num_generated_tokens:
-            gen_positions = modality.positions[:, :, gen_indices, :]  # (B, 3, num_tile_gen, 2)
-            tile_start = gen_positions[..., 0].amin(dim=2)  # (B, 3)
-            tile_end = gen_positions[..., 1].amax(dim=2)  # (B, 3)
-
         cond_positions = modality.positions[:, :, self._num_generated_tokens :, :]  # (B, 3, num_cond, 2)
+        device = cond_positions.device
 
-            overlaps = (cond_positions[..., 0] < tile_end.unsqueeze(2)) & (
+        # Per-tile (start, end) bounds along each axis; small Python loop (num_tiles <= ~16).
-                cond_positions[..., 1] > tile_start.unsqueeze(2)
+        starts_list: list[torch.Tensor] = []
-            )  # (B, 3, num_cond)
+        ends_list: list[torch.Tensor] = []
-            overlaps_all_dims = overlaps.all(dim=1)  # (B, num_cond)
+        for t in self._tiles:
+            gen_idx = self._generated_token_indices(t, device=device)
+            gen_positions = modality.positions[:, :, gen_idx, :]  # (B, 3, num_tile_gen, 2)
+            starts_list.append(gen_positions[..., 0].amin(dim=2))  # (B, 3)
+            ends_list.append(gen_positions[..., 1].amax(dim=2))  # (B, 3)
+        tile_starts = torch.stack(starts_list, dim=0)  # (num_tiles, B, 3)
+        tile_ends = torch.stack(ends_list, dim=0)  # (num_tiles, B, 3)
 
-            has_negative_time = cond_positions[:, 0, :, 0] < 0  # (B, num_cond)
+        cond_starts = cond_positions[..., 0]  # (B, 3, num_cond)
-
+        cond_ends = cond_positions[..., 1]  # (B, 3, num_cond)
-            keep_cond = (overlaps_all_dims | has_negative_time).any(dim=0)  # (num_cond,)
+        # Broadcast: (1, B, 3, num_cond) vs (num_tiles, B, 3, 1) -> (num_tiles, B, 3, num_cond).
-            mask[self._num_generated_tokens :] = keep_cond
+        overlaps = (cond_starts[None] < tile_ends[..., None]) & (cond_ends[None] > tile_starts[..., None])
-
+        overlaps_all_dims = overlaps.all(dim=2)  # (num_tiles, B, num_cond)
-        return mask
+        has_negative_time = (cond_positions[:, 0, :, 0] < 0)[None]  # (1, B, num_cond)
+        return (overlaps_all_dims | has_negative_time).any(dim=1)  # (num_tiles, num_cond)

packages/ltx-core/src/ltx_core/model/transformer/attention.py

@@ -1,12 +1,38 @@
+import functools
+import logging
+from dataclasses import dataclass, field
 from enum import Enum
 from typing import Protocol
 
 import torch
+from torch.nn.attention import SDPBackend, sdpa_kernel
+
+from ltx_core.model.transformer.ops import (
+    GatedAttentionCallable,
+    PreAttentionCallable,
+    PytorchGatedAttention,
+    PytorchPreAttention,
+)
+from ltx_core.model.transformer.rope import LTXRopeType
+
+logger = logging.getLogger(__name__)
+
+
+def _torch_default_sdpa_priority() -> list[SDPBackend]:
+    """Fetch torch's current default SDPA priority order at runtime.
+    Used as the default for ``PytorchAttention`` so the wrapper-always
+    code path matches torch's native dispatch order without hard-coding it
+    (which would drift if torch updates the default).
+    ``torch._C._get_sdp_priority_order`` is a private API; we accept that
+    risk because the project pins ``torch`` in the lockfile, so any
+    rename/removal surfaces on a controlled torch bump rather than silently.
+    """
+    return [SDPBackend(p) for p in torch._C._get_sdp_priority_order()]
 
-from ltx_core.model.transformer.rope import LTXRopeType, apply_rotary_emb
 
 memory_efficient_attention = None
 flash_attn_interface = None
+flash_attn_4_func = None
 try:
     from xformers.ops import memory_efficient_attention
 except ImportError:
@@ -17,15 +43,44 @@ try:
         import flash_attn_interface
 except ImportError:
     flash_attn_interface = None
+try:
+    from flash_attn.cute import flash_attn_func as flash_attn_4_func
+except ImportError:
+    flash_attn_4_func = None
 
 
 class AttentionCallable(Protocol):
+    """Unmasked attention. Backends without a mask kernel (FA3/FA4) implement only
+    this protocol; backends that support masks too (Pytorch/SDPA, xFormers) are
+    structurally usable here and as :class:`MaskedAttentionCallable`."""
+
+    def __call__(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, heads: int) -> torch.Tensor: ...
+
+
+class MaskedAttentionCallable(Protocol):
+    """Masked attention. Mask is required (not optional) -- the caller has already
+    decided this is the masked path and chosen a backend that can serve it. Used
+    by :class:`Attention` when its forward receives a non-None ``mask``."""
+
     def __call__(
-        self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, heads: int, mask: torch.Tensor | None = None
+        self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, heads: int, mask: torch.Tensor
     ) -> torch.Tensor: ...
 
 
 class PytorchAttention(AttentionCallable):
+    def __init__(self, priority: list[SDPBackend] | None = None) -> None:
+        # priority=None -> snapshot torch's default SDPA priority at construction.
+        # Always passed through ``sdpa_kernel(..., set_priority=True)`` so the
+        # call site is uniform regardless of how the priority was chosen.
+        self._priority = priority if priority is not None else _torch_default_sdpa_priority()
+
+    @property
+    def label(self) -> str:
+        """Human-readable identifier (used in the AUTOMATIC selection log).
+        Encodes the SDPA priority list so a single-backend pin reads differently
+        from the full-priority dispatcher walk."""
+        return f"SDPA[{'>'.join(b.name for b in self._priority)}]"
+
     def __call__(
         self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, heads: int, mask: torch.Tensor | None = None
     ) -> torch.Tensor:
@@ -41,12 +96,17 @@ class PytorchAttention(AttentionCallable):
             if mask.ndim == 3:
                 mask = mask.unsqueeze(1)
 
-        out = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=0.0, is_causal=False)
+        with sdpa_kernel(self._priority, set_priority=True):
+            out = torch.nn.functional.scaled_dot_product_attention(
+                q, k, v, attn_mask=mask, dropout_p=0.0, is_causal=False
+            )
         out = out.transpose(1, 2).reshape(b, -1, heads * dim_head)
         return out
 
 
 class XFormersAttention(AttentionCallable):
+    label = "xFormers"
+
     def __call__(
         self,
         q: torch.Tensor,
@@ -92,13 +152,14 @@ class XFormersAttention(AttentionCallable):
 
 
 class FlashAttention3(AttentionCallable):
+    label = "FlashAttention3"
+
     def __call__(
         self,
         q: torch.Tensor,
         k: torch.Tensor,
         v: torch.Tensor,
         heads: int,
-        mask: torch.Tensor | None = None,
     ) -> torch.Tensor:
         if flash_attn_interface is None:
             raise RuntimeError("FlashAttention3 was selected but `FlashAttention3` is not installed.")
@@ -108,32 +169,275 @@ class FlashAttention3(AttentionCallable):
 
         q, k, v = (t.view(b, -1, heads, dim_head) for t in (q, k, v))
 
-        if mask is not None:
-            raise NotImplementedError("Mask is not supported for FlashAttention3")
-
         out = flash_attn_interface.flash_attn_func(q.to(v.dtype), k.to(v.dtype), v)
         out = out.reshape(b, -1, heads * dim_head)
         return out
 
 
+class FlashAttention4(AttentionCallable):
+    label = "FlashAttention4"
+
+    def __call__(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        heads: int,
+    ) -> torch.Tensor:
+        if flash_attn_4_func is None:
+            raise RuntimeError("FlashAttention4 was selected but `flash-attn-4` is not installed.")
+
+        b, _, dim_head = q.shape
+        dim_head //= heads
+
+        q, k, v = (t.view(b, -1, heads, dim_head) for t in (q, k, v))
+
+        out, _ = flash_attn_4_func(q.to(v.dtype), k.to(v.dtype), v)
+        out = out.reshape(b, -1, heads * dim_head)
+        return out
+
+
+# --- Automatic selection -----------------------------------------------------
+# AUTOMATIC inspects installed extras and the GPU arch and returns the fastest
+# usable callable for each path. The selection runs once per process (cached)
+# and logs the resulting label once. The unmasked and masked picks are
+# independent: each calls its own helper and may end up on different backends
+# (e.g. FA3 unmasked + xFormers masked on H100).
+
+
+def _sdpa_can_use(backend: SDPBackend, *, with_mask: bool) -> bool:
+    """Ask torch whether *backend* can run with the given mask shape.
+    ``MATH`` is the universal SDPA fallback (pure PyTorch ops, no kernel
+    requirements) so it returns True everywhere, CPU included. The other
+    backends use ``torch.backends.cuda.can_use_*`` capability checks (no GPU
+    compute, no synchronization) and are False without CUDA. The probe shapes
+    are small but realistic enough to surface constraints (head dim, dtype)
+    that the per-backend rules care about.
+    """
+    if backend is SDPBackend.MATH:
+        return True
+    if not torch.cuda.is_available():
+        return False
+    q = torch.empty(1, 4, 128, 64, device="cuda", dtype=torch.bfloat16)
+    k = torch.empty(1, 4, 128, 64, device="cuda", dtype=torch.bfloat16)
+    v = torch.empty(1, 4, 128, 64, device="cuda", dtype=torch.bfloat16)
+    mask = torch.zeros(1, 4, 128, 128, device="cuda", dtype=torch.bfloat16) if with_mask else None
+    params = torch.backends.cuda.SDPAParams(q, k, v, mask, 0.0, False, False)
+    if backend is SDPBackend.CUDNN_ATTENTION:
+        return torch.backends.cuda.can_use_cudnn_attention(params, debug=False)
+    if backend is SDPBackend.FLASH_ATTENTION:
+        return torch.backends.cuda.can_use_flash_attention(params, debug=False)
+    if backend is SDPBackend.EFFICIENT_ATTENTION:
+        return torch.backends.cuda.can_use_efficient_attention(params, debug=False)
+    return False
+
+
+_SDPA_FULL_PRIORITY: tuple[SDPBackend, ...] = (
+    SDPBackend.CUDNN_ATTENTION,
+    SDPBackend.FLASH_ATTENTION,
+    SDPBackend.EFFICIENT_ATTENTION,
+    SDPBackend.MATH,
+)
+
+
+def _sdpa_full_priority() -> PytorchAttention:
+    """Hand SDPA the full backend priority order; let torch's dispatcher pick at call time.
+    ``sdpa_kernel(_SDPA_FULL_PRIORITY, set_priority=True)`` enables all four
+    backends and orders them; torch then walks the order at call time and picks
+    the first backend whose ``can_use_*`` check passes for the actual
+    shapes/dtype/mask. FLASH is rejected automatically when a mask is present;
+    CUDNN may be rejected under deterministic mode; MATH is the universal
+    fallback. Probing per-backend usability up front from generic probe shapes
+    cannot anticipate the variety of real call sites (e.g. broadcast key-only
+    masks, large head dim), so we defer the choice to the dispatcher.
+    """
+    return PytorchAttention(priority=list(_SDPA_FULL_PRIORITY))
+
+
+def _select_primary_attention() -> AttentionCallable:
+    """Pick the fastest unmasked attention based on installed extras and GPU arch.
+    Priority by arch:
+    - Hopper (sm_90, H100): FA3 / xFormers (mutually exclusive at import) > FA4 > SDPA.
+    - Datacenter Blackwell (sm_100, B200): FA4 > SDPA. FA4 is intentionally *not*
+      picked on consumer Blackwell (sm_120) -- known regressions in newer
+      FA4 betas; users who want it on sm_120 must opt in explicitly.
+    - Everywhere else (Ada, Ampere, CPU): SDPA with the full backend priority
+      list -- torch's runtime dispatcher picks the best fit at call time.
+    """
+    if torch.cuda.is_available():
+        major, _ = torch.cuda.get_device_capability(0)
+        if major == 9:
+            if flash_attn_interface is not None:
+                return FlashAttention3()
+            if memory_efficient_attention is not None:
+                return XFormersAttention()
+            if flash_attn_4_func is not None:
+                return FlashAttention4()
+        if major == 10 and flash_attn_4_func is not None:
+            return FlashAttention4()
+    return _sdpa_full_priority()
+
+
+def _select_masked_attention() -> MaskedAttentionCallable:
+    """Pick a mask-aware attention. Prefers xFormers when installed; else SDPA with
+    the full priority list (the dispatcher rejects FLASH automatically when a
+    mask is present and walks past it)."""
+    if memory_efficient_attention is not None:
+        return XFormersAttention()
+    return _sdpa_full_priority()
+
+
+@functools.cache
+def automatic_attention() -> AttentionCallable:
+    """Cached AUTOMATIC pick for the unmasked path. Logs the chosen label once
+    per process."""
+    fn = _select_primary_attention()
+    logger.info("Automatic attention selected: %s", fn.label)
+    return fn
+
+
+@functools.cache
+def automatic_masked_attention() -> MaskedAttentionCallable:
+    """Cached AUTOMATIC pick for the masked path. Logs the chosen label once
+    per process."""
+    fn = _select_masked_attention()
+    logger.info("Automatic masked attention selected: %s", fn.label)
+    return fn
+
+
+def _resolve_sdpa_variant(backend: SDPBackend, name: str, *, with_mask: bool) -> PytorchAttention:
+    """Build a single-backend ``PytorchAttention`` pin, raising if the backend
+    can't actually serve the call on this machine. Used by both
+    :meth:`AttentionFunction.to_callable` and :meth:`MaskedAttentionFunction.to_callable`;
+    ``with_mask`` differs between the two so the capability check considers
+    the protocol the caller intends to use. Not used for ``MATH`` -- MATH is
+    the universal fallback and would falsely fail the CUDA-only probe on CPU.
+    """
+    if not _sdpa_can_use(backend, with_mask=with_mask):
+        raise RuntimeError(
+            f"{name} selected but the SDPA {backend.name} backend is not usable on this machine "
+            "(either no CUDA, the backend rejected the probe shapes, or "
+            "torch.use_deterministic_algorithms(True) excluded it)."
+        )
+    return PytorchAttention(priority=[backend])
+
+
 class AttentionFunction(Enum):
     PYTORCH = "pytorch"
     XFORMERS = "xformers"
     FLASH_ATTENTION_3 = "flash_attention_3"
-    DEFAULT = "default"
+    FLASH_ATTENTION_4 = "flash_attention_4"
+    SDPA_CUDNN = "sdpa_cudnn"
+    SDPA_FLASH = "sdpa_flash"
+    SDPA_EFFICIENT = "sdpa_efficient"
+    SDPA_MATH = "sdpa_math"
+    # Pick the fastest unmasked backend for the current GPU/extras combo; see
+    # :func:`automatic_attention`. Default for :class:`AttentionOps`.
+    AUTOMATIC = "automatic"
 
-    def to_callable(self) -> AttentionCallable:
+    def to_callable(self) -> AttentionCallable:  # noqa: PLR0911
         """Resolve to a concrete callable. Use this at module init time so that
-        torch.compile can trace through the attention call without graph breaks."""
+        torch.compile can trace through the attention call without graph breaks.
-        if self is AttentionFunction.PYTORCH:
+        Every non-AUTOMATIC variant raises :class:`RuntimeError` when the backend
+        isn't usable on this machine -- missing package or SDPA backend rejected
+        on this hardware (e.g. cuDNN under ``torch.use_deterministic_algorithms``).
+        Opting in means "this kernel or fail loudly". ``AUTOMATIC`` returns the
+        cached :func:`automatic_attention` instance so the once-per-process log
+        fires only on the first resolution.
+        """
+        match self:
+            case AttentionFunction.AUTOMATIC:
+                return automatic_attention()
+            case AttentionFunction.PYTORCH:
                 return PytorchAttention()
-        elif self is AttentionFunction.XFORMERS:
+            case AttentionFunction.XFORMERS:
+                if memory_efficient_attention is None:
+                    raise RuntimeError("AttentionFunction.XFORMERS selected but `xformers` is not installed.")
                 return XFormersAttention()
-        elif self is AttentionFunction.FLASH_ATTENTION_3:
+            case AttentionFunction.FLASH_ATTENTION_3:
+                if flash_attn_interface is None:
+                    raise RuntimeError(
+                        "AttentionFunction.FLASH_ATTENTION_3 selected but `flash-attn-3` is not installed."
+                    )
                 return FlashAttention3()
-        else:
+            case AttentionFunction.FLASH_ATTENTION_4:
-            # Default behavior: XFormers if installed else - PyTorch
+                if flash_attn_4_func is None:
-            return XFormersAttention() if memory_efficient_attention is not None else PytorchAttention()
+                    raise RuntimeError(
+                        "AttentionFunction.FLASH_ATTENTION_4 selected but `flash-attn-4` is not installed."
+                    )
+                return FlashAttention4()
+            case AttentionFunction.SDPA_MATH:
+                return PytorchAttention(priority=[SDPBackend.MATH])
+            case AttentionFunction.SDPA_CUDNN:
+                return _resolve_sdpa_variant(
+                    SDPBackend.CUDNN_ATTENTION, "AttentionFunction.SDPA_CUDNN", with_mask=False
+                )
+            case AttentionFunction.SDPA_FLASH:
+                return _resolve_sdpa_variant(
+                    SDPBackend.FLASH_ATTENTION, "AttentionFunction.SDPA_FLASH", with_mask=False
+                )
+            case AttentionFunction.SDPA_EFFICIENT:
+                return _resolve_sdpa_variant(
+                    SDPBackend.EFFICIENT_ATTENTION, "AttentionFunction.SDPA_EFFICIENT", with_mask=False
+                )
+
+
+class MaskedAttentionFunction(Enum):
+    """Backends usable on the masked path. Mirrors :class:`AttentionFunction` minus
+    the variants the torch SDPA dispatcher (or the wrapped kernel) rejects with a
+    mask: ``SDPA_FLASH`` -- FLASH kernel cannot serve an additive ``attn_mask``;
+    ``FLASH_ATTENTION_3``/``FLASH_ATTENTION_4`` -- neither has a mask kernel at all.
+    Keeping them out makes "this backend cannot mask" a type error, not a runtime one."""
+
+    PYTORCH = "pytorch"
+    XFORMERS = "xformers"
+    SDPA_CUDNN = "sdpa_cudnn"
+    SDPA_EFFICIENT = "sdpa_efficient"
+    SDPA_MATH = "sdpa_math"
+    # Pick the fastest mask-capable backend for the current extras combo; see
+    # :func:`automatic_masked_attention`. Default for the masked slot of
+    # :class:`AttentionOps`.
+    AUTOMATIC = "automatic"
+
+    def to_callable(self) -> MaskedAttentionCallable:
+        """Resolve to a concrete masked callable. Same backend classes as
+        :meth:`AttentionFunction.to_callable`; the protocol returned just exposes
+        the masked call signature.
+        Non-AUTOMATIC variants raise :class:`RuntimeError` when the backend isn't
+        usable for the masked path on this machine. SDPA probes run with
+        ``with_mask=True`` so the capability check considers the protocol the
+        caller will actually use."""
+        match self:
+            case MaskedAttentionFunction.AUTOMATIC:
+                return automatic_masked_attention()
+            case MaskedAttentionFunction.PYTORCH:
+                return PytorchAttention()
+            case MaskedAttentionFunction.XFORMERS:
+                if memory_efficient_attention is None:
+                    raise RuntimeError("MaskedAttentionFunction.XFORMERS selected but `xformers` is not installed.")
+                return XFormersAttention()
+            case MaskedAttentionFunction.SDPA_MATH:
+                return PytorchAttention(priority=[SDPBackend.MATH])
+            case MaskedAttentionFunction.SDPA_CUDNN:
+                return _resolve_sdpa_variant(
+                    SDPBackend.CUDNN_ATTENTION, "MaskedAttentionFunction.SDPA_CUDNN", with_mask=True
+                )
+            case MaskedAttentionFunction.SDPA_EFFICIENT:
+                return _resolve_sdpa_variant(
+                    SDPBackend.EFFICIENT_ATTENTION, "MaskedAttentionFunction.SDPA_EFFICIENT", with_mask=True
+                )
+
+
+@dataclass(frozen=True)
+class AttentionOps:
+    """Pluggable callables consumed by :class:`Attention`."""
+
+    attention_function: AttentionCallable = field(default_factory=lambda: AttentionFunction.AUTOMATIC.to_callable())
+    masked_attention_function: MaskedAttentionCallable = field(
+        default_factory=lambda: MaskedAttentionFunction.AUTOMATIC.to_callable()
+    )
+    preattention_function: PreAttentionCallable = field(default_factory=PytorchPreAttention)
+    gated_attention_function: GatedAttentionCallable = field(default_factory=PytorchGatedAttention)
 
 
 class Attention(torch.nn.Module):
@@ -145,16 +449,17 @@ class Attention(torch.nn.Module):
         dim_head: int = 64,
         norm_eps: float = 1e-6,
         rope_type: LTXRopeType = LTXRopeType.SPLIT,
-        attention_function: AttentionCallable | AttentionFunction = AttentionFunction.DEFAULT,
+        ops: AttentionOps | None = None,
         apply_gated_attention: bool = False,
     ) -> None:
         super().__init__()
+        if ops is None:
+            ops = AttentionOps()
         self.rope_type = rope_type
-        self.attention_function = (
+        self.attention_function = ops.attention_function
-            attention_function.to_callable()
+        self.masked_attention_function = ops.masked_attention_function
-            if isinstance(attention_function, AttentionFunction)
+        self.preattention_function = ops.preattention_function
-            else attention_function
+        self.gated_attention_function = ops.gated_attention_function
-        )
 
         inner_dim = dim_head * heads
         context_dim = query_dim if context_dim is None else context_dim
@@ -196,7 +501,9 @@ class Attention(torch.nn.Module):
             context: Key/value context tensor of shape ``(B, S, context_dim)``.
                 Falls back to ``x`` (self-attention) when *None*.
             mask: Optional attention mask. Interpretation depends on the attention
-                backend (additive bias for xformers/PyTorch SDPA).
+                backend (additive bias for xformers/PyTorch SDPA). A non-None
+                ``mask`` routes to ``masked_attention_function``; ``None`` keeps
+                the unmasked path.
             pe: Rotary positional embeddings applied to both ``q`` and ``k``.
             k_pe: Separate rotary positional embeddings for ``k`` only. When
                 *None*, ``pe`` is reused for keys.
@@ -221,29 +528,17 @@ class Attention(torch.nn.Module):
         else:
             q = self.to_q(x)
             k = self.to_k(context)
-
+            q, k = self.preattention_function(q, k, self, mask, pe, k_pe)
-            q = self.q_norm(q)
+            if mask is None:
-            k = self.k_norm(k)
+                out = self.attention_function(q, k, v, self.heads)  # (B, T, H*D)
-
+            else:
-            if pe is not None:
+                out = self.masked_attention_function(q, k, v, self.heads, mask)
-                q = apply_rotary_emb(q, pe, self.rope_type)
-                k = apply_rotary_emb(k, pe if k_pe is None else k_pe, self.rope_type)
-
-            out = self.attention_function(q, k, v, self.heads, mask)  # (B, T, H*D)
 
             if perturbation_mask is not None:
                 out = out * perturbation_mask + v * (1 - perturbation_mask)
 
         # Apply per-head gating if enabled
         if self.to_gate_logits is not None:
-            gate_logits = self.to_gate_logits(x)  # (B, T, H)
+            out = self.gated_attention_function(x, out, self)
-            b, t, _ = out.shape
-            # Reshape to (B, T, H, D) for per-head gating
-            out = out.view(b, t, self.heads, self.dim_head)
-            # Apply gating: 2 * sigmoid(x) so that zero-init gives identity (2 * 0.5 = 1.0)
-            gates = 2.0 * torch.sigmoid(gate_logits)  # (B, T, H)
-            out = out * gates.unsqueeze(-1)  # (B, T, H, D) * (B, T, H, 1)
-            # Reshape back to (B, T, H*D)
-            out = out.view(b, t, self.heads * self.dim_head)
 
         return self.to_out(out)

packages/ltx-core/src/ltx_core/model/transformer/compiling.py

@@ -1,19 +1,112 @@
+from dataclasses import dataclass, field
+from typing import Any
+
 import torch
 
+from ltx_core.guidance.perturbations import BatchedPerturbationConfig, PerturbationType
 from ltx_core.loader.module_ops import ModuleOps
 from ltx_core.loader.sd_ops import SDOps
 from ltx_core.model.transformer.model import LTXModel
+from ltx_core.model.transformer.transformer_args import BlockPerturbationsProcessor, TransformerArgs
+
+# Defaults applied inside the patched forward. Overriding via CompilationConfig
+# replaces these wholesale; it does not merge.
+_DEFAULT_INDUCTOR_CONFIG: dict[str, Any] = {"unsafe_skip_cache_dynamic_shape_guards": True}
+_DEFAULT_DYNAMO_CONFIG: dict[str, Any] = {"inline_inbuilt_nn_modules": True, "cache_size_limit": 256}
 
 
-def compile_transformer(model: LTXModel) -> LTXModel:
+@dataclass(frozen=True)
-    model.transformer_blocks = torch.nn.ModuleList(torch.compile(m) for m in model.transformer_blocks)
+class CompilationConfig:
+    """``torch.compile`` configuration for transformer blocks. ``None`` keeps eager."""
+
+    mode: str | None = None
+    backend: str = "inductor"
+    fullgraph: bool = False
+    dynamic: bool | None = None
+    inductor_config: dict[str, Any] = field(default_factory=lambda: dict(_DEFAULT_INDUCTOR_CONFIG))
+    dynamo_config: dict[str, Any] = field(default_factory=lambda: dict(_DEFAULT_DYNAMO_CONFIG))
+
+
+class _SeqDynamicMarkingProcessor:
+    """Marks the per-block seq dim dynamic, then delegates to an inner processor.
+    Installed by ``compile_transformer`` so the per-block compile artifact stays
+    shape-polymorphic. Wraps whatever ``block_input_processor`` was already on
+    the model -- callers that customised the processor keep their customisation;
+    only the seq-dim marking is layered on top. Lives outside the compiled
+    region, so ``mark_dynamic`` runs in eager mode on the tensors that are
+    about to cross into the trace.
+    """
+
+    def __init__(self, inner: BlockPerturbationsProcessor) -> None:
+        self.inner = inner
+
+    def __call__(
+        self,
+        args: TransformerArgs,
+        perturbations: BatchedPerturbationConfig,
+        block_idx: int,
+        self_attn_type: PerturbationType,
+        cross_attn_type: PerturbationType,
+    ) -> TransformerArgs:
+        # Positional embeddings are second-from-last regardless of rope type:
+        # split rope is (B, H, T, D//2) -- dim -2 == 2; interleaved rope is (B, T, D)
+        # -- dim -2 == 1. Both work via the negative index.
+        torch._dynamo.mark_dynamic(args.x, 1)
+        cos, sin = args.positional_embeddings
+        torch._dynamo.mark_dynamic(cos, cos.ndim - 2)
+        torch._dynamo.mark_dynamic(sin, sin.ndim - 2)
+        if args.cross_positional_embeddings is not None:
+            cross_cos, cross_sin = args.cross_positional_embeddings
+            torch._dynamo.mark_dynamic(cross_cos, cross_cos.ndim - 2)
+            torch._dynamo.mark_dynamic(cross_sin, cross_sin.ndim - 2)
+        if args.self_attention_mask is not None:
+            # Dense form is (B, 1, T, T); key-padding form (from the SP wrapper)
+            # is (B, 1, 1, T) -- leave the size-1 query dim static so Dynamo
+            # keeps the broadcast.
+            if args.self_attention_mask.shape[2] > 1:
+                torch._dynamo.mark_dynamic(args.self_attention_mask, 2)
+            torch._dynamo.mark_dynamic(args.self_attention_mask, 3)
+        if args.context_mask is not None:
+            torch._dynamo.mark_dynamic(args.context_mask, 2)
+        # `timesteps` / `embedded_timestep` are per-token when conditioning sets a
+        # per-position denoise mask, in which case their dim 1 equals the seq length
+        # and must vary with it. When they're a single timestep broadcast across the
+        # sequence (dim 1 == 1), leaving them static lets Dynamo keep the size-1
+        # broadcast.
+        if args.timesteps.shape[1] > 1:
+            torch._dynamo.mark_dynamic(args.timesteps, 1)
+        if args.embedded_timestep.shape[1] > 1:
+            torch._dynamo.mark_dynamic(args.embedded_timestep, 1)
+        # `cross_scale_shift_timestep` is the cross-attn AdaLN scale/shift input
+        # derived from the own-modality per-token timesteps (denoise_mask * sigma),
+        # so its dim 1 equals the seq length when conditioning is per-token.
+        # `cross_gate_timestep` is the cross-modality sigma scalar -- dim 1 is 1
+        # and broadcasts, leave it static. Same guard pattern as `timesteps`.
+        if args.cross_scale_shift_timestep is not None and args.cross_scale_shift_timestep.shape[1] > 1:
+            torch._dynamo.mark_dynamic(args.cross_scale_shift_timestep, 1)
+        return self.inner(args, perturbations, block_idx, self_attn_type, cross_attn_type)
+
+
+def compile_transformer(model: LTXModel, config: CompilationConfig) -> LTXModel:
+    """Compile each transformer block via ``torch.compile`` with the given settings.
+    The patched forward emits ``torch.compiler.cudagraph_mark_step_begin()`` once
+    per step. Under CUDA-graph-enabling modes (``"reduce-overhead"`` /
+    ``"max-autotune"``) this overrides Dynamo's per-invocation auto-mark
+    heuristic, which would otherwise fire once per compiled block call (48 per
+    forward) and treat each block call as a fresh iteration. Under other modes
+    the mark is a no-op (decrements an unread counter).
+    """
+    model.transformer_blocks = torch.nn.ModuleList(
+        torch.compile(m, mode=config.mode, backend=config.backend, fullgraph=config.fullgraph, dynamic=config.dynamic)
+        for m in model.transformer_blocks
+    )
+    model.block_input_processor = _SeqDynamicMarkingProcessor(inner=model.block_input_processor)
 
     def patched_dynamo_forward(*args, **kwargs) -> tuple[torch.Tensor, torch.Tensor]:
+        torch.compiler.cudagraph_mark_step_begin()
         with (
-            torch._inductor.config.patch(unsafe_skip_cache_dynamic_shape_guards=True),
+            torch._inductor.config.patch(**config.inductor_config),
-            torch._dynamo.config.patch(  # type: ignore[attr-defined]
+            torch._dynamo.config.patch(**config.dynamo_config),  # type: ignore[attr-defined]
-                inline_inbuilt_nn_modules=True, cache_size_limit=256, allow_unspec_int_on_nn_module=True
-            ),
         ):
             return model.forward_without_compilation(*args, **kwargs)
 
@@ -22,11 +115,13 @@ def compile_transformer(model: LTXModel) -> LTXModel:
     return model
 
 
-COMPILE_TRANSFORMER = ModuleOps(
+def build_compile_transformer_op(config: CompilationConfig) -> ModuleOps:
+    """Build a ``ModuleOps`` that compiles transformer blocks with the given settings."""
+    return ModuleOps(
         name="compile_transformer",
         matcher=lambda model: isinstance(model, LTXModel),
-    mutator=lambda model: compile_transformer(model),
+        mutator=lambda model: compile_transformer(model, config),
-)
+    )
 
 
 def modify_sd_ops_for_compilation(original_sd_ops: SDOps, number_of_blocks: int = 48) -> SDOps:

packages/ltx-core/src/ltx_core/model/transformer/modality.py

@@ -17,8 +17,19 @@ class Modality:
             the batch size, *T* is the total number of tokens (noisy +
             conditioning), and *D* is the input dimension.
         timesteps: Per-token timestep embeddings, shape ``(B, T)``.
-        positions: Positional coordinates, shape ``(B, 3, T)`` for video
+        positions: Per-token patch coordinates used to build the RoPE
-            (time, height, width) or ``(B, 1, T)`` for audio.
+            frequencies. With the default ``use_middle_indices_grid=True``,
+            shape is ``(B, n_pos_dims, T, 2)`` where ``n_pos_dims=3`` for
+            video (time, height, width) and ``n_pos_dims=1`` for audio
+            (time); the last dim of size 2 holds the ``[start, end)``
+            index bounds of each patch, and RoPE is evaluated at the
+            *middle* of that range -- hence the flag name. Taking the
+            patch midpoint produces a smoother and more accurate
+            positional signal than indexing by the patch's start when
+            patches span more than one spatial / temporal unit.
+            When ``use_middle_indices_grid=False``, the legacy 3-D form
+            ``(B, n_pos_dims, T)`` of integer positional indices is
+            accepted instead and used as-is (no midpoint derivation).
         context: Text conditioning embeddings from the prompt encoder.
         enabled: Whether this modality is active in the current forward pass.
         context_mask: Optional mask for the text context tokens.
@@ -34,9 +45,10 @@ class Modality:
     )  # Shape: (B, T, D) where B is the batch size, T is the number of tokens, and D is input dimension
     sigma: torch.Tensor  # Shape: (B,). Current sigma value, used for cross-attention timestep calculation.
     timesteps: torch.Tensor  # Shape: (B, T) where T is the number of timesteps
-    positions: (
+    # Shape: (B, n_pos_dims, T, 2) by default (use_middle_indices_grid=True);
-        torch.Tensor
+    # n_pos_dims=3 for video, 1 for audio; last dim holds [start, end) patch bounds.
-    )  # Shape: (B, 3, T) for video, where 3 is the number of dimensions and T is the number of tokens
+    # Legacy form (B, n_pos_dims, T) when use_middle_indices_grid=False.
+    positions: torch.Tensor
     context: torch.Tensor
     enabled: bool = True
     context_mask: torch.Tensor | None = None

packages/ltx-core/src/ltx_core/model/transformer/model.py

@@ -2,13 +2,18 @@ from enum import Enum
 
 import torch
 
-from ltx_core.guidance.perturbations import BatchedPerturbationConfig
+from ltx_core.guidance.perturbations import BatchedPerturbationConfig, PerturbationType
 from ltx_core.model.transformer.adaln import AdaLayerNormSingle, adaln_embedding_coefficient
-from ltx_core.model.transformer.attention import AttentionCallable, AttentionFunction
 from ltx_core.model.transformer.modality import Modality
 from ltx_core.model.transformer.rope import LTXRopeType
-from ltx_core.model.transformer.transformer import BasicAVTransformerBlock, TransformerConfig
+from ltx_core.model.transformer.transformer import (
+    DEFAULT_TRANSFORMER_OPS,
+    BasicAVTransformerBlock,
+    TransformerConfig,
+    TransformerOpsConfig,
+)
 from ltx_core.model.transformer.transformer_args import (
+    BlockPerturbationsProcessor,
     MultiModalTransformerArgsPreprocessor,
     TransformerArgs,
     TransformerArgsPreprocessor,
@@ -45,7 +50,7 @@ class LTXModel(torch.nn.Module):
         num_layers: int = 48,
         cross_attention_dim: int = 4096,
         norm_eps: float = 1e-06,
-        attention_type: AttentionFunction | AttentionCallable = AttentionFunction.DEFAULT,
+        ops: TransformerOpsConfig = DEFAULT_TRANSFORMER_OPS,
         positional_embedding_theta: float = 10000.0,
         positional_embedding_max_pos: list[int] | None = None,
         timestep_scale_multiplier: int = 1000,
@@ -115,9 +120,13 @@ class LTXModel(torch.nn.Module):
             audio_attention_head_dim=audio_attention_head_dim if model_type.is_audio_enabled() else 0,
             audio_cross_attention_dim=audio_cross_attention_dim,
             norm_eps=norm_eps,
-            attention_type=attention_type,
+            ops=ops,
             apply_gated_attention=apply_gated_attention,
         )
+        # Hook for per-block input prep. Compile transforms in `compiling.py`
+        # wrap (not replace) this with a processor that also marks the seq dim
+        # dynamic, so any caller customisation here is preserved as the inner.
+        self.block_input_processor = BlockPerturbationsProcessor()
 
     @property
     def _adaln_embedding_coefficient(self) -> int:
@@ -284,7 +293,7 @@ class LTXModel(torch.nn.Module):
         audio_attention_head_dim: int,
         audio_cross_attention_dim: int,
         norm_eps: float,
-        attention_type: AttentionFunction | AttentionCallable,
+        ops: TransformerOpsConfig,
         apply_gated_attention: bool,
     ) -> None:
         """Initialize transformer blocks for LTX."""
@@ -315,14 +324,13 @@ class LTXModel(torch.nn.Module):
         self.transformer_blocks = torch.nn.ModuleList(
             [
                 BasicAVTransformerBlock(
-                    idx=idx,
                     video=video_config,
                     audio=audio_config,
                     rope_type=self.rope_type,
                     norm_eps=norm_eps,
-                    attention_function=attention_type,
+                    ops=ops,
                 )
-                for idx in range(num_layers)
+                for _ in range(num_layers)
             ]
         )
 
@@ -340,29 +348,44 @@ class LTXModel(torch.nn.Module):
         self,
         video: TransformerArgs | None,
         audio: TransformerArgs | None,
-        perturbations: BatchedPerturbationConfig,
+        perturbations: BatchedPerturbationConfig | None,
-    ) -> tuple[TransformerArgs, TransformerArgs]:
+    ) -> tuple[TransformerArgs | None, TransformerArgs | None]:
-        """Process transformer blocks for LTXAV."""
+        """Process transformer blocks for LTXAV.
+        Per-block perturbation masks are precomputed here and attached to each
+        modality's ``TransformerArgs`` so the block forward has no per-block
+        identity to specialise on — all blocks share a single Dynamo cache slot.
+        """
+        if perturbations is None:
+            batch_size = (video or audio).x.shape[0]
+            perturbations = BatchedPerturbationConfig.empty(batch_size)
+
+        for block_idx, block in enumerate(self.transformer_blocks):
+            if video is not None:
+                video = self.block_input_processor(
+                    video,
+                    perturbations,
+                    block_idx,
+                    self_attn_type=PerturbationType.SKIP_VIDEO_SELF_ATTN,
+                    cross_attn_type=PerturbationType.SKIP_A2V_CROSS_ATTN,
+                )
+            if audio is not None:
+                audio = self.block_input_processor(
+                    audio,
+                    perturbations,
+                    block_idx,
+                    self_attn_type=PerturbationType.SKIP_AUDIO_SELF_ATTN,
+                    cross_attn_type=PerturbationType.SKIP_V2A_CROSS_ATTN,
+                )
 
-        # Process transformer blocks
-        for block in self.transformer_blocks:
             if self._enable_gradient_checkpointing and self.training:
-                # Use gradient checkpointing to save memory during training.
-                # With use_reentrant=False, we can pass dataclasses directly -
-                # PyTorch will track all tensor leaves in the computation graph.
                 video, audio = torch.utils.checkpoint.checkpoint(
                     block,
                     video,
                     audio,
-                    perturbations,
                     use_reentrant=False,
                 )
             else:
-                video, audio = block(
+                video, audio = block(video=video, audio=audio)
-                    video=video,
-                    audio=audio,
-                    perturbations=perturbations,
-                )
 
         return video, audio
 

packages/ltx-core/src/ltx_core/model/transformer/model_configurator.py

@@ -2,10 +2,10 @@ import torch
 
 from ltx_core.loader.sd_ops import SDOps
 from ltx_core.model.model_protocol import ModelConfigurator
-from ltx_core.model.transformer.attention import AttentionFunction
 from ltx_core.model.transformer.model import LTXModel, LTXModelType
 from ltx_core.model.transformer.rope import LTXRopeType
 from ltx_core.model.transformer.text_projection import create_caption_projection
+from ltx_core.model.transformer.transformer import DEFAULT_TRANSFORMER_OPS, TransformerOpsConfig
 from ltx_core.utils import check_config_value
 
 
@@ -16,7 +16,7 @@ class LTXModelConfigurator(ModelConfigurator[LTXModel]):
     """
 
     @classmethod
-    def from_config(cls: type[LTXModel], config: dict) -> LTXModel:
+    def from_config(cls, config: dict, ops: TransformerOpsConfig = DEFAULT_TRANSFORMER_OPS) -> LTXModel:
         # Build caption projections for 19B models (projection handled in transformer).
         caption_projection, audio_caption_projection = _build_caption_projections(config, is_av=True)
 
@@ -40,6 +40,7 @@ class LTXModelConfigurator(ModelConfigurator[LTXModel]):
         check_config_value(config, "share_ff", False)
         check_config_value(config, "av_cross_ada_norm", True)
         check_config_value(config, "use_middle_indices_grid", True)
+        check_config_value(config, "num_attention_heads", config.get("audio_num_attention_heads", float("nan")))
 
         return LTXModel(
             model_type=LTXModelType.AudioVideo,
@@ -50,7 +51,7 @@ class LTXModelConfigurator(ModelConfigurator[LTXModel]):
             num_layers=config.get("num_layers", 48),
             cross_attention_dim=config.get("cross_attention_dim", 4096),
             norm_eps=config.get("norm_eps", 1e-06),
-            attention_type=AttentionFunction(config.get("attention_type", "default")),
+            ops=ops,
             positional_embedding_theta=config.get("positional_embedding_theta", 10000.0),
             positional_embedding_max_pos=config.get("positional_embedding_max_pos", [20, 2048, 2048]),
             timestep_scale_multiplier=config.get("timestep_scale_multiplier", 1000),
@@ -78,7 +79,7 @@ class LTXVideoOnlyModelConfigurator(ModelConfigurator[LTXModel]):
     """
 
     @classmethod
-    def from_config(cls: type[LTXModel], config: dict) -> LTXModel:
+    def from_config(cls, config: dict, ops: TransformerOpsConfig = DEFAULT_TRANSFORMER_OPS) -> LTXModel:
         # Build caption projection for 19B model (projection handled in transformer).
         caption_projection, _ = _build_caption_projections(config, is_av=False)
 
@@ -109,7 +110,7 @@ class LTXVideoOnlyModelConfigurator(ModelConfigurator[LTXModel]):
             num_layers=config.get("num_layers", 48),
             cross_attention_dim=config.get("cross_attention_dim", 4096),
             norm_eps=config.get("norm_eps", 1e-06),
-            attention_type=AttentionFunction(config.get("attention_type", "default")),
+            ops=ops,
             positional_embedding_theta=config.get("positional_embedding_theta", 10000.0),
             positional_embedding_max_pos=config.get("positional_embedding_max_pos", [20, 2048, 2048]),
             timestep_scale_multiplier=config.get("timestep_scale_multiplier", 1000),

packages/ltx-core/src/ltx_core/model/transformer/ops.py

@@ -0,0 +1,106 @@
+from typing import List, Protocol
+
+import torch
+from torch import nn
+
+from ltx_core.model.transformer.rope import apply_rotary_emb
+from ltx_core.utils import rms_norm
+
+
+class PreAttentionCallable(Protocol):
+    def __call__(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        attn_module: nn.Module,
+        mask: torch.Tensor | None,
+        pe: torch.Tensor | None,
+        k_pe: torch.Tensor | None,
+    ) -> tuple[torch.Tensor, torch.Tensor]: ...
+
+
+class PytorchPreAttention(PreAttentionCallable):
+    def __call__(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        attn_module: nn.Module,
+        mask: torch.Tensor | None,  # noqa: ARG002
+        pe: torch.Tensor | None,
+        k_pe: torch.Tensor | None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        q = attn_module.q_norm(q)
+        k = attn_module.k_norm(k)
+        if pe is not None:
+            q = apply_rotary_emb(q, pe, attn_module.rope_type)
+            k = apply_rotary_emb(k, pe if k_pe is None else k_pe, attn_module.rope_type)
+        return q, k
+
+
+class AdaZeroCallable(Protocol):
+    def __call__(
+        self,
+        x: torch.Tensor,
+        eps: float,
+        scale: torch.Tensor,
+        shift: torch.Tensor,
+    ) -> torch.Tensor: ...
+
+
+class PytorchAdaZeroFunction(AdaZeroCallable):
+    def __call__(
+        self,
+        x: torch.Tensor,
+        eps: float,
+        scale: torch.Tensor,
+        shift: torch.Tensor,
+    ) -> torch.Tensor:
+        return rms_norm(x, eps=eps) * (1 + scale) + shift
+
+
+class PostSACallable(Protocol):
+    def __call__(
+        self,
+        x: torch.Tensor,
+        y: torch.Tensor,
+        norm_weights: torch.Tensor | None,
+        eps: float,
+        gate: torch.Tensor,
+    ) -> List[torch.Tensor]: ...
+
+
+class PytorchPostSAFunction(PostSACallable):
+    def __call__(
+        self,
+        x: torch.Tensor,
+        y: torch.Tensor,
+        norm_weights: torch.Tensor | None,
+        eps: float,
+        gate: torch.Tensor,
+    ) -> List[torch.Tensor]:
+        x_fma = x + y * gate
+        return x_fma, rms_norm(x_fma, norm_weights, eps=eps)
+
+
+class GatedAttentionCallable(Protocol):
+    def __call__(
+        self,
+        x: torch.Tensor,
+        attn_out: torch.Tensor,
+        attn_module: nn.Module,
+    ) -> torch.Tensor: ...
+
+
+class PytorchGatedAttention(GatedAttentionCallable):
+    def __call__(
+        self,
+        x: torch.Tensor,
+        attn_out: torch.Tensor,
+        attn_module: nn.Module,
+    ) -> torch.Tensor:
+        gate_logits = attn_module.to_gate_logits(x)  # (B, T, H)
+        b, t, _ = attn_out.shape
+        out = attn_out.view(b, t, attn_module.heads, attn_module.dim_head)
+        gates = 2.0 * torch.sigmoid(gate_logits)  # (B, T, H)
+        out = out * gates.unsqueeze(-1)  # (B, T, H, D) * (B, T, H, 1)
+        return out.view(b, t, attn_module.heads * attn_module.dim_head)

packages/ltx-core/src/ltx_core/model/transformer/rope.py

@@ -43,16 +43,26 @@ def apply_interleaved_rotary_emb(
 def apply_split_rotary_emb(
     input_tensor: torch.Tensor, cos_freqs: torch.Tensor, sin_freqs: torch.Tensor
 ) -> torch.Tensor:
-    needs_reshape = False
+    if sin_freqs.shape != cos_freqs.shape:
-    if input_tensor.ndim != 4 and cos_freqs.ndim == 4:
-        b, h, t, _ = cos_freqs.shape
-        if input_tensor.shape[0] != b:
         raise ValueError(
-                f"apply_split_rotary_emb: input_tensor batch ({input_tensor.shape[0]}) "
+            f"apply_split_rotary_emb: sin_freqs.shape {tuple(sin_freqs.shape)} must equal "
-                f"must equal cos_freqs batch ({b})."
+            f"cos_freqs.shape {tuple(cos_freqs.shape)}."
         )
-        input_tensor = input_tensor.reshape(b, t, h, -1).swapaxes(1, 2)
+    needs_reshape = input_tensor.ndim != 4 and cos_freqs.ndim == 4
-        needs_reshape = True
+    if needs_reshape:
+        b_freq = cos_freqs.shape[0]
+        h = cos_freqs.shape[1]
+        b_in = input_tensor.shape[0]
+        if b_freq not in (1, b_in):
+            raise ValueError(
+                f"apply_split_rotary_emb: cos_freqs batch ({b_freq}) must be 1 "
+                f"(broadcast) or equal input_tensor batch ({b_in})."
+            )
+        # `unflatten` only touches the last dim, keeping the batch and seq dims as
+        # the input tensor's own symbolic ints under torch.compile. `reshape(b_in,
+        # t, h, -1)` would have forced Dynamo to specialise those dims because it
+        # cannot prove `b_in == cos_freqs.shape[0]` and `seq == t` across tensors.
+        input_tensor = input_tensor.unflatten(-1, (h, -1)).transpose(1, 2)
 
     split_input = rearrange(input_tensor, "... (d r) -> ... d r", d=2)
     first_half_input = split_input[..., :1, :]
@@ -67,7 +77,9 @@ def apply_split_rotary_emb(
 
     output = rearrange(output, "... d r -> ... (d r)")
     if needs_reshape:
-        output = output.swapaxes(1, 2).reshape(b, t, -1)
+        # `transpose(1, 2).flatten(-2)` keeps the batch and seq dims symbolic; using
+        # `reshape(b_in, t, -1)` would force Dynamo to specialise both axes.
+        output = output.transpose(1, 2).flatten(-2)
 
     return output
 

packages/ltx-core/src/ltx_core/model/transformer/transformer.py

@@ -1,11 +1,27 @@
-from dataclasses import dataclass, replace
+from dataclasses import dataclass, field, replace
 
 import torch
 
-from ltx_core.guidance.perturbations import BatchedPerturbationConfig, PerturbationType
 from ltx_core.model.transformer.adaln import adaln_embedding_coefficient
-from ltx_core.model.transformer.attention import Attention, AttentionCallable, AttentionFunction
+from ltx_core.model.transformer.attention import (
+    Attention,
+    AttentionCallable,
+    AttentionFunction,
+    AttentionOps,
+    MaskedAttentionCallable,
+    MaskedAttentionFunction,
+)
 from ltx_core.model.transformer.feed_forward import FeedForward
+from ltx_core.model.transformer.ops import (
+    AdaZeroCallable,
+    GatedAttentionCallable,
+    PostSACallable,
+    PreAttentionCallable,
+    PytorchAdaZeroFunction,
+    PytorchGatedAttention,
+    PytorchPostSAFunction,
+    PytorchPreAttention,
+)
 from ltx_core.model.transformer.rope import LTXRopeType
 from ltx_core.model.transformer.transformer_args import TransformerArgs
 from ltx_core.utils import rms_norm
@@ -21,19 +37,68 @@ class TransformerConfig:
     cross_attention_adaln: bool = False
 
 
+@dataclass(frozen=True)
+class TransformerOpsConfig:
+    """Pluggable ops for :class:`BasicAVTransformerBlock`.
+    Use :meth:`from_functions` to construct from enum values or partial overrides
+    without spelling out a full :class:`AttentionOps`.
+    """
+
+    attention_ops: AttentionOps = field(default_factory=AttentionOps)
+    ada_zero_function: AdaZeroCallable = field(default_factory=PytorchAdaZeroFunction)
+    post_sa_function: PostSACallable = field(default_factory=PytorchPostSAFunction)
+
+    @classmethod
+    def from_functions(
+        cls,
+        attention: AttentionFunction | AttentionCallable = AttentionFunction.AUTOMATIC,
+        masked_attention: MaskedAttentionFunction | MaskedAttentionCallable = MaskedAttentionFunction.AUTOMATIC,
+        preattention: PreAttentionCallable | None = None,
+        gated_attention: GatedAttentionCallable | None = None,
+        ada_zero: AdaZeroCallable | None = None,
+        post_sa: PostSACallable | None = None,
+    ) -> "TransformerOpsConfig":
+        """Build a config from individual functions or enums. Each *None* slot
+        falls back to the standard PyTorch implementation."""
+        attention_callable = attention.to_callable() if isinstance(attention, AttentionFunction) else attention
+        masked_callable = (
+            masked_attention.to_callable()
+            if isinstance(masked_attention, MaskedAttentionFunction)
+            else masked_attention
+        )
+        attention_ops = AttentionOps(
+            attention_function=attention_callable,
+            masked_attention_function=masked_callable,
+            preattention_function=preattention if preattention is not None else PytorchPreAttention(),
+            gated_attention_function=(gated_attention if gated_attention is not None else PytorchGatedAttention()),
+        )
+        return cls(
+            attention_ops=attention_ops,
+            ada_zero_function=ada_zero if ada_zero is not None else PytorchAdaZeroFunction(),
+            post_sa_function=post_sa if post_sa is not None else PytorchPostSAFunction(),
+        )
+
+
+# Frozen, so safe to share as a default argument across callers that want the
+# stock PyTorch ops without explicit construction.
+DEFAULT_TRANSFORMER_OPS = TransformerOpsConfig()
+
+
 class BasicAVTransformerBlock(torch.nn.Module):
     def __init__(
         self,
-        idx: int,
         video: TransformerConfig | None = None,
         audio: TransformerConfig | None = None,
         rope_type: LTXRopeType = LTXRopeType.SPLIT,
         norm_eps: float = 1e-6,
-        attention_function: AttentionFunction | AttentionCallable = AttentionFunction.DEFAULT,
+        ops: TransformerOpsConfig | None = None,
     ):
         super().__init__()
 
-        self.idx = idx
+        if ops is None:
+            ops = TransformerOpsConfig()
+        self.ada_zero_function = ops.ada_zero_function
+        self.post_sa_function = ops.post_sa_function
         if video is not None:
             self.attn1 = Attention(
                 query_dim=video.dim,
@@ -42,7 +107,7 @@ class BasicAVTransformerBlock(torch.nn.Module):
                 context_dim=None,
                 rope_type=rope_type,
                 norm_eps=norm_eps,
-                attention_function=attention_function,
+                ops=ops.attention_ops,
                 apply_gated_attention=video.apply_gated_attention,
             )
             self.attn2 = Attention(
@@ -52,7 +117,7 @@ class BasicAVTransformerBlock(torch.nn.Module):
                 dim_head=video.d_head,
                 rope_type=rope_type,
                 norm_eps=norm_eps,
-                attention_function=attention_function,
+                ops=ops.attention_ops,
                 apply_gated_attention=video.apply_gated_attention,
             )
             self.ff = FeedForward(video.dim, dim_out=video.dim)
@@ -67,7 +132,7 @@ class BasicAVTransformerBlock(torch.nn.Module):
                 context_dim=None,
                 rope_type=rope_type,
                 norm_eps=norm_eps,
-                attention_function=attention_function,
+                ops=ops.attention_ops,
                 apply_gated_attention=audio.apply_gated_attention,
             )
             self.audio_attn2 = Attention(
@@ -77,7 +142,7 @@ class BasicAVTransformerBlock(torch.nn.Module):
                 dim_head=audio.d_head,
                 rope_type=rope_type,
                 norm_eps=norm_eps,
-                attention_function=attention_function,
+                ops=ops.attention_ops,
                 apply_gated_attention=audio.apply_gated_attention,
             )
             self.audio_ff = FeedForward(audio.dim, dim_out=audio.dim)
@@ -93,7 +158,7 @@ class BasicAVTransformerBlock(torch.nn.Module):
                 dim_head=audio.d_head,
                 rope_type=rope_type,
                 norm_eps=norm_eps,
-                attention_function=attention_function,
+                ops=ops.attention_ops,
                 apply_gated_attention=video.apply_gated_attention,
             )
 
@@ -105,7 +170,7 @@ class BasicAVTransformerBlock(torch.nn.Module):
                 dim_head=audio.d_head,
                 rope_type=rope_type,
                 norm_eps=norm_eps,
-                attention_function=attention_function,
+                ops=ops.attention_ops,
                 apply_gated_attention=audio.apply_gated_attention,
             )
 
@@ -188,16 +253,10 @@ class BasicAVTransformerBlock(torch.nn.Module):
         self,
         video: TransformerArgs | None,
         audio: TransformerArgs | None,
-        perturbations: BatchedPerturbationConfig | None = None,
     ) -> tuple[TransformerArgs | None, TransformerArgs | None]:
         if video is None and audio is None:
             raise ValueError("At least one of video or audio must be provided")
 
-        batch_size = (video or audio).x.shape[0]
-
-        if perturbations is None:
-            perturbations = BatchedPerturbationConfig.empty(batch_size)
-
         vx = video.x if video is not None else None
         ax = audio.x if audio is not None else None
 
@@ -211,28 +270,18 @@ class BasicAVTransformerBlock(torch.nn.Module):
             vshift_msa, vscale_msa, vgate_msa = self.get_ada_values(
                 self.scale_shift_table, vx.shape[0], video.timesteps, slice(0, 3)
             )
-            norm_vx = rms_norm(vx, eps=self.norm_eps) * (1 + vscale_msa) + vshift_msa
+            norm_vx = self.ada_zero_function(vx, self.norm_eps, vscale_msa, vshift_msa)
             del vshift_msa, vscale_msa
 
-            all_perturbed = perturbations.all_in_batch(PerturbationType.SKIP_VIDEO_SELF_ATTN, self.idx)
+            vx_msa_out = self.attn1(
-            none_perturbed = not perturbations.any_in_batch(PerturbationType.SKIP_VIDEO_SELF_ATTN, self.idx)
-            v_mask = (
-                perturbations.mask_like(PerturbationType.SKIP_VIDEO_SELF_ATTN, self.idx, vx)
-                if not all_perturbed and not none_perturbed
-                else None
-            )
-            vx = (
-                vx
-                + self.attn1(
                 norm_vx,
                 pe=video.positional_embeddings,
                 mask=video.self_attention_mask,
-                    perturbation_mask=v_mask,
+                perturbation_mask=video.self_attn_perturbation_mask,
-                    all_perturbed=all_perturbed,
+                all_perturbed=video.self_attn_all_perturbed,
             )
-                * vgate_msa
+            vx = vx + vx_msa_out * vgate_msa
-            )
+            del vgate_msa, norm_vx, vx_msa_out
-            del vgate_msa, norm_vx, v_mask
             vx = vx + self._apply_text_cross_attention(
                 vx,
                 video.context,
@@ -250,27 +299,17 @@ class BasicAVTransformerBlock(torch.nn.Module):
                 self.audio_scale_shift_table, ax.shape[0], audio.timesteps, slice(0, 3)
             )
 
-            norm_ax = rms_norm(ax, eps=self.norm_eps) * (1 + ascale_msa) + ashift_msa
+            norm_ax = self.ada_zero_function(ax, self.norm_eps, ascale_msa, ashift_msa)
             del ashift_msa, ascale_msa
-            all_perturbed = perturbations.all_in_batch(PerturbationType.SKIP_AUDIO_SELF_ATTN, self.idx)
+            ax_msa_out = self.audio_attn1(
-            none_perturbed = not perturbations.any_in_batch(PerturbationType.SKIP_AUDIO_SELF_ATTN, self.idx)
-            a_mask = (
-                perturbations.mask_like(PerturbationType.SKIP_AUDIO_SELF_ATTN, self.idx, ax)
-                if not all_perturbed and not none_perturbed
-                else None
-            )
-            ax = (
-                ax
-                + self.audio_attn1(
                 norm_ax,
                 pe=audio.positional_embeddings,
                 mask=audio.self_attention_mask,
-                    perturbation_mask=a_mask,
+                perturbation_mask=audio.self_attn_perturbation_mask,
-                    all_perturbed=all_perturbed,
+                all_perturbed=audio.self_attn_all_perturbed,
             )
-                * agate_msa
+            ax = ax + ax_msa_out * agate_msa
-            )
+            del agate_msa, norm_ax, ax_msa_out
-            del agate_msa, norm_ax, a_mask
             ax = ax + self._apply_text_cross_attention(
                 ax,
                 audio.context,
@@ -285,10 +324,11 @@ class BasicAVTransformerBlock(torch.nn.Module):
 
         # Audio - Video cross attention.
         if run_a2v or run_v2a:
-            vx_norm3 = rms_norm(vx, eps=self.norm_eps)
+            # Snapshot vx/ax before A2V mutates vx; V2A's video keys/values must
-            ax_norm3 = rms_norm(ax, eps=self.norm_eps)
+            # use the pre-A2V state so direction order doesn't bias the result.
-
+            vx_pre_av = vx
-            if run_a2v and not perturbations.all_in_batch(PerturbationType.SKIP_A2V_CROSS_ATTN, self.idx):
+            ax_pre_av = ax
+            if run_a2v and not video.cross_attn_skip_all:
                 scale_ca_video_a2v, shift_ca_video_a2v, gate_out_a2v = self.get_av_ca_ada_values(
                     self.scale_shift_table_a2v_ca_video,
                     vx.shape[0],
@@ -296,7 +336,7 @@ class BasicAVTransformerBlock(torch.nn.Module):
                     video.cross_gate_timestep,
                     slice(0, 2),
                 )
-                vx_scaled = vx_norm3 * (1 + scale_ca_video_a2v) + shift_ca_video_a2v
+                a2v_vx_scaled = self.ada_zero_function(vx_pre_av, self.norm_eps, scale_ca_video_a2v, shift_ca_video_a2v)
                 del scale_ca_video_a2v, shift_ca_video_a2v
 
                 scale_ca_audio_a2v, shift_ca_audio_a2v, _ = self.get_av_ca_ada_values(
@@ -306,22 +346,21 @@ class BasicAVTransformerBlock(torch.nn.Module):
                     audio.cross_gate_timestep,
                     slice(0, 2),
                 )
-                ax_scaled = ax_norm3 * (1 + scale_ca_audio_a2v) + shift_ca_audio_a2v
+                a2v_ax_scaled = self.ada_zero_function(ax_pre_av, self.norm_eps, scale_ca_audio_a2v, shift_ca_audio_a2v)
                 del scale_ca_audio_a2v, shift_ca_audio_a2v
-                a2v_mask = perturbations.mask_like(PerturbationType.SKIP_A2V_CROSS_ATTN, self.idx, vx)
                 vx = vx + (
                     self.audio_to_video_attn(
-                        vx_scaled,
+                        a2v_vx_scaled,
-                        context=ax_scaled,
+                        context=a2v_ax_scaled,
                         pe=video.cross_positional_embeddings,
                         k_pe=audio.cross_positional_embeddings,
                     )
                     * gate_out_a2v
-                    * a2v_mask
+                    * video.cross_attn_perturbation_mask
                 )
-                del gate_out_a2v, a2v_mask, vx_scaled, ax_scaled
+                del gate_out_a2v, a2v_vx_scaled, a2v_ax_scaled
 
-            if run_v2a and not perturbations.all_in_batch(PerturbationType.SKIP_V2A_CROSS_ATTN, self.idx):
+            if run_v2a and not audio.cross_attn_skip_all:
                 scale_ca_audio_v2a, shift_ca_audio_v2a, gate_out_v2a = self.get_av_ca_ada_values(
                     self.scale_shift_table_a2v_ca_audio,
                     ax.shape[0],
@@ -329,7 +368,7 @@ class BasicAVTransformerBlock(torch.nn.Module):
                     audio.cross_gate_timestep,
                     slice(2, 4),
                 )
-                ax_scaled = ax_norm3 * (1 + scale_ca_audio_v2a) + shift_ca_audio_v2a
+                v2a_ax_scaled = self.ada_zero_function(ax_pre_av, self.norm_eps, scale_ca_audio_v2a, shift_ca_audio_v2a)
                 del scale_ca_audio_v2a, shift_ca_audio_v2a
                 scale_ca_video_v2a, shift_ca_video_v2a, _ = self.get_av_ca_ada_values(
                     self.scale_shift_table_a2v_ca_video,
@@ -338,28 +377,26 @@ class BasicAVTransformerBlock(torch.nn.Module):
                     video.cross_gate_timestep,
                     slice(2, 4),
                 )
-                vx_scaled = vx_norm3 * (1 + scale_ca_video_v2a) + shift_ca_video_v2a
+                v2a_vx_scaled = self.ada_zero_function(vx_pre_av, self.norm_eps, scale_ca_video_v2a, shift_ca_video_v2a)
                 del scale_ca_video_v2a, shift_ca_video_v2a
-                v2a_mask = perturbations.mask_like(PerturbationType.SKIP_V2A_CROSS_ATTN, self.idx, ax)
                 ax = ax + (
                     self.video_to_audio_attn(
-                        ax_scaled,
+                        v2a_ax_scaled,
-                        context=vx_scaled,
+                        context=v2a_vx_scaled,
                         pe=audio.cross_positional_embeddings,
                         k_pe=video.cross_positional_embeddings,
                     )
                     * gate_out_v2a
-                    * v2a_mask
+                    * audio.cross_attn_perturbation_mask
                 )
-                del gate_out_v2a, v2a_mask, ax_scaled, vx_scaled
+                del gate_out_v2a, v2a_vx_scaled, v2a_ax_scaled
-
+            del vx_pre_av, ax_pre_av
-            del vx_norm3, ax_norm3
 
         if run_vx:
             vshift_mlp, vscale_mlp, vgate_mlp = self.get_ada_values(
                 self.scale_shift_table, vx.shape[0], video.timesteps, slice(3, 6)
             )
-            vx_scaled = rms_norm(vx, eps=self.norm_eps) * (1 + vscale_mlp) + vshift_mlp
+            vx_scaled = self.ada_zero_function(vx, self.norm_eps, vscale_mlp, vshift_mlp)
             vx = vx + self.ff(vx_scaled) * vgate_mlp
 
             del vshift_mlp, vscale_mlp, vgate_mlp, vx_scaled
@@ -368,7 +405,7 @@ class BasicAVTransformerBlock(torch.nn.Module):
             ashift_mlp, ascale_mlp, agate_mlp = self.get_ada_values(
                 self.audio_scale_shift_table, ax.shape[0], audio.timesteps, slice(3, 6)
             )
-            ax_scaled = rms_norm(ax, eps=self.norm_eps) * (1 + ascale_mlp) + ashift_mlp
+            ax_scaled = self.ada_zero_function(ax, self.norm_eps, ascale_mlp, ashift_mlp)
             ax = ax + self.audio_ff(ax_scaled) * agate_mlp
 
             del ashift_mlp, ascale_mlp, agate_mlp, ax_scaled

packages/ltx-core/src/ltx_core/model/transformer/transformer_args.py

@@ -2,6 +2,7 @@ from dataclasses import dataclass, replace
 
 import torch
 
+from ltx_core.guidance.perturbations import BatchedPerturbationConfig, PerturbationType
 from ltx_core.model.transformer.adaln import AdaLayerNormSingle
 from ltx_core.model.transformer.modality import Modality
 from ltx_core.model.transformer.rope import (
@@ -19,8 +20,8 @@ class TransformerArgs:
     context_mask: torch.Tensor
     timesteps: torch.Tensor
     embedded_timestep: torch.Tensor
-    positional_embeddings: torch.Tensor
+    positional_embeddings: tuple[torch.Tensor, torch.Tensor]
-    cross_positional_embeddings: torch.Tensor | None
+    cross_positional_embeddings: tuple[torch.Tensor, torch.Tensor] | None
     cross_scale_shift_timestep: torch.Tensor | None
     cross_gate_timestep: torch.Tensor | None
     enabled: bool
@@ -28,6 +29,59 @@ class TransformerArgs:
     self_attention_mask: torch.Tensor | None = (
         None  # Additive log-space self-attention bias (B, 1, T, T), None = full attention
     )
+    # Per-block perturbation state, precomputed by `LTXModel._process_transformer_blocks`
+    # so the block forward needs no per-block identity. The bool shortcuts
+    # (`*_all_perturbed`, `cross_attn_skip_all`) are Python bools that Dynamo specialises
+    # on — fine because they're stable across denoising steps for a fixed perturbation
+    # config.
+    self_attn_perturbation_mask: torch.Tensor | None = None
+    self_attn_all_perturbed: bool = False
+    cross_attn_perturbation_mask: torch.Tensor | None = None
+    cross_attn_skip_all: bool = False
+
+
+class BlockPerturbationsProcessor:
+    """Per-block preparation of ``TransformerArgs``.
+    The base implementation returns a copy of ``args`` with this block's
+    precomputed perturbation flags and masks attached. Subclasses can layer in
+    operations that must run on each block's inputs but stay outside the
+    compile boundary -- e.g. ``torch._dynamo.mark_dynamic`` for
+    shape-polymorphic block compilation (see ``compiling.py``). Swapping the
+    processor on an ``LTXModel`` instance is how compile transforms opt in to
+    such behaviour without baking it into the model's forward.
+    ``self_attn_perturbation_mask`` is None when all or none of the batch is
+    perturbed (the attention call can take the shortcut path). ``cross_attn_*``
+    is None when every sample skips the cross-attention entirely.
+    """
+
+    def __call__(
+        self,
+        args: "TransformerArgs",
+        perturbations: BatchedPerturbationConfig,
+        block_idx: int,
+        self_attn_type: PerturbationType,
+        cross_attn_type: PerturbationType,
+    ) -> "TransformerArgs":
+        device, dtype = args.x.device, args.x.dtype
+
+        all_self = perturbations.all_in_batch(self_attn_type, block_idx)
+        any_self = perturbations.any_in_batch(self_attn_type, block_idx)
+        self_mask: torch.Tensor | None = None
+        if any_self and not all_self:
+            self_mask = perturbations.mask(self_attn_type, block_idx, device, dtype).view(-1, 1, 1)
+
+        all_cross = perturbations.all_in_batch(cross_attn_type, block_idx)
+        cross_mask: torch.Tensor | None = None
+        if not all_cross:
+            cross_mask = perturbations.mask(cross_attn_type, block_idx, device, dtype).view(-1, 1, 1)
+
+        return replace(
+            args,
+            self_attn_perturbation_mask=self_mask,
+            self_attn_all_perturbed=all_self,
+            cross_attn_perturbation_mask=cross_mask,
+            cross_attn_skip_all=all_cross,
+        )
 
 
 class TransformerArgsPreprocessor:
@@ -98,9 +152,12 @@ class TransformerArgsPreprocessor:
         self, attention_mask: torch.Tensor | None, x_dtype: torch.dtype
     ) -> torch.Tensor | None:
         """Prepare self-attention mask by converting [0,1] values to additive log-space bias.
-        Input shape: (B, T, T) with values in [0, 1].
+        Input shape: 3D ``(B, T_q, T_k)`` with values in [0, 1]. The dense form
-        Output shape: (B, 1, T, T) with 0.0 for full attention and a large negative value
+        is ``(B, T, T)``; broadcastable forms like ``(1, 1, T)`` (key-only
-        for masked positions.
+        padding) or ``(B, 1, T)`` are also valid and yield a correspondingly
+        broadcastable output.
+        Output shape: ``(B, 1, T_q, T_k)`` (heads dim inserted) with 0.0 for
+        full attention and a large negative value for masked positions.
         Positions with attention_mask <= 0 are fully masked (mapped to the dtype's minimum
         representable value). Strictly positive entries are converted via log-space for
         smooth attenuation, with small values clamped for numerical stability.
@@ -120,7 +177,7 @@ class TransformerArgsPreprocessor:
         if positive.any():
             bias[positive] = torch.log(attention_mask[positive].clamp(min=eps)).to(x_dtype)
 
-        return bias.unsqueeze(1)  # (B, 1, T, T) for head broadcast
+        return bias.unsqueeze(1)  # (B, 1, T_q, T_k) for head broadcast
 
     def _prepare_positional_embeddings(
         self,
@@ -244,10 +301,6 @@ class MultiModalTransformerArgsPreprocessor:
             if cross_modality.sigma.ndim != 1:
                 raise ValueError("Cross modality sigma must be a 1D tensor")
 
-        cross_timestep = cross_modality.sigma.view(
-            modality.timesteps.shape[0], 1, *[1] * len(modality.timesteps.shape[2:])
-        )
-
         cross_pe = self.simple_preprocessor._prepare_positional_embeddings(
             positions=modality.positions[:, 0:1, :],
             inner_dim=self.audio_cross_attention_dim,
@@ -258,7 +311,8 @@ class MultiModalTransformerArgsPreprocessor:
         )
 
         cross_scale_shift_timestep, cross_gate_timestep = self._prepare_cross_attention_timestep(
-            timestep=cross_timestep,
+            modality_timesteps=modality.timesteps,
+            cross_modality_sigma=cross_modality.sigma,
             timestep_scale_multiplier=self.simple_preprocessor.timestep_scale_multiplier,
             batch_size=transformer_args.x.shape[0],
             hidden_dtype=modality.latent.dtype,
@@ -273,23 +327,23 @@ class MultiModalTransformerArgsPreprocessor:
 
     def _prepare_cross_attention_timestep(
         self,
-        timestep: torch.Tensor | None,
+        modality_timesteps: torch.Tensor,
+        cross_modality_sigma: torch.Tensor,
         timestep_scale_multiplier: int,
         batch_size: int,
         hidden_dtype: torch.dtype,
     ) -> tuple[torch.Tensor, torch.Tensor]:
-        """Prepare cross attention timestep embeddings."""
+        """Prepare A-V cross-attention AdaLN inputs."""
-        timestep = timestep * timestep_scale_multiplier
-
         av_ca_factor = self.av_ca_timestep_scale_multiplier / timestep_scale_multiplier
 
         scale_shift_timestep, _ = self.cross_scale_shift_adaln(
-            timestep.flatten(),
+            (modality_timesteps * timestep_scale_multiplier).flatten(),
             hidden_dtype=hidden_dtype,
         )
         scale_shift_timestep = scale_shift_timestep.view(batch_size, -1, scale_shift_timestep.shape[-1])
+
         gate_noise_timestep, _ = self.cross_gate_adaln(
-            timestep.flatten() * av_ca_factor,
+            (cross_modality_sigma * timestep_scale_multiplier * av_ca_factor).flatten(),
             hidden_dtype=hidden_dtype,
         )
         gate_noise_timestep = gate_noise_timestep.view(batch_size, -1, gate_noise_timestep.shape[-1])

packages/ltx-core/src/ltx_core/model/video_vae/memory_efficient_decode.py

@@ -56,6 +56,25 @@ if TYPE_CHECKING:
 # ---------------------------------------------------------------------------
 
 
+def _memory_format_of(t: torch.Tensor, prefer_channels_last_3d: bool = False) -> torch.memory_format:
+    """Pick the memory format for a workspace allocation.
+    When ``prefer_channels_last_3d`` is True and ``t`` is 5D, return
+    ``channels_last_3d`` regardless of ``t``'s current strides -- the
+    workspace's ``.copy_(t)`` will transcribe the data into the new layout.
+    This is needed because intermediate tensors inside the decoder (after
+    ``rearrange`` + slice + residual add in ``_upsample_forward_efficient``)
+    are not NHWC-contiguous, so an auto-detect helper would silently fall
+    back to NCHW for every workspace after the first upsample.
+    Otherwise fall back to inspecting ``t``: ``channels_last_3d`` if ``t``
+    already uses it, else contiguous.
+    """
+    if prefer_channels_last_3d and t.dim() == 5:
+        return torch.channels_last_3d
+    if t.dim() == 5 and t.is_contiguous(memory_format=torch.channels_last_3d):
+        return torch.channels_last_3d
+    return torch.contiguous_format
+
+
 def _find_temporal_split_size(num_frames: int) -> int:
     """Find chunk size for in-place temporal convolution.
     The chunk size ensures the last chunk has at least 3 frames
@@ -132,6 +151,7 @@ def inplace_conv3d_temporal_chunked(workspace: torch.Tensor, conv: nn.Conv3d) ->
         workspace.shape[4],
         device=workspace.device,
         dtype=workspace.dtype,
+        memory_format=_memory_format_of(workspace),
     )
     o_buf = torch.empty_like(x_buf)
 
@@ -190,6 +210,7 @@ def _causal_pad(x: torch.Tensor, pad_size: int) -> torch.Tensor:
         x.shape[4],
         device=x.device,
         dtype=x.dtype,
+        memory_format=_memory_format_of(x),
     )
     padded[:, :, pad_size:].copy_(x)
     for i in range(pad_size):
@@ -325,6 +346,7 @@ def _midblock_forward_efficient(
     causal: bool,
     timestep: torch.Tensor | None,
     generator: torch.Generator | None,
+    prefer_channels_last_3d: bool = False,
 ) -> torch.Tensor:
     """Memory-efficient ``UNetMidBlock3D`` forward.
     Allocates a single workspace buffer that is reused across all
@@ -351,6 +373,7 @@ def _midblock_forward_efficient(
         hidden_states.shape[4],
         device=hidden_states.device,
         dtype=hidden_states.dtype,
+        memory_format=_memory_format_of(hidden_states, prefer_channels_last_3d),
     )
 
     for resnet in block.res_blocks:
@@ -366,6 +389,7 @@ def _upsample_forward_efficient(
     block: DepthToSpaceUpsample,
     x: torch.Tensor,
     causal: bool,
+    prefer_channels_last_3d: bool = False,
 ) -> torch.Tensor:
     """Memory-efficient ``DepthToSpaceUpsample`` forward.
     For non-causal mode the input is copied into a workspace and the
@@ -393,6 +417,7 @@ def _upsample_forward_efficient(
     if causal:
         x = _causal_pad_free_and_conv(x, block.conv)
     else:
+        mem_fmt = _memory_format_of(x, prefer_channels_last_3d)
         workspace = torch.empty(
             x.shape[0],
             max(in_channels, out_channels),
@@ -401,11 +426,12 @@ def _upsample_forward_efficient(
             x.shape[4],
             device=x.device,
             dtype=x.dtype,
+            memory_format=mem_fmt,
         )
         workspace[:, :in_channels, 1:-1].copy_(x)
         del x
         inplace_conv3d_temporal_chunked(workspace, conv)
-        x = workspace[:, :out_channels, 1:-1].contiguous()
+        x = workspace[:, :out_channels, 1:-1].contiguous(memory_format=mem_fmt)
         del workspace
 
     x = rearrange(
@@ -418,7 +444,7 @@ def _upsample_forward_efficient(
     if block.stride[0] == 2:
         x = x[:, :, 1:, :, :]
     if block.residual:
-        x = x + x_in
+        x.add_(x_in)
         del x_in
     return x
 
@@ -434,12 +460,14 @@ def _final_norm_and_conv_out(
     causal: bool,
     scaled_timestep: torch.Tensor | None,
     batch_size: int,
+    prefer_channels_last_3d: bool = False,
 ) -> torch.Tensor:
     """Workspace-based final norm + [ada] + SiLU + conv_out + unpatchify."""
     conv_out_mod: CausalConv3d = decoder.conv_out  # type: ignore[assignment]
     conv_out = conv_out_mod.conv
     feature_channels = sample.shape[1]
 
+    mem_fmt = _memory_format_of(sample, prefer_channels_last_3d)
     workspace = torch.empty(
         sample.shape[0],
         max(feature_channels, conv_out.out_channels),
@@ -448,6 +476,7 @@ def _final_norm_and_conv_out(
         sample.shape[4],
         device=sample.device,
         dtype=sample.dtype,
+        memory_format=mem_fmt,
     )
     workspace[:, :feature_channels, 1:-1].copy_(sample)
     del sample
@@ -485,7 +514,7 @@ def _final_norm_and_conv_out(
         del padded
     else:
         inplace_conv3d_temporal_chunked(workspace, conv_out)
-        result = workspace[:, : conv_out.out_channels, 1:-1].contiguous()
+        result = workspace[:, : conv_out.out_channels, 1:-1].contiguous(memory_format=mem_fmt)
         del workspace, interior
 
     return unpatchify(result, patch_size_hw=decoder.patch_size, patch_size_t=1)
@@ -507,6 +536,9 @@ def _memory_efficient_forward(
     ``UNetMidBlock3D`` and ``DepthToSpaceUpsample`` blocks use efficient
     paths; standalone ``ResnetBlock3D`` blocks fall back to the standard
     forward.  The final norm + ada + SiLU + conv_out is also workspace-based.
+    All workspaces are allocated ``channels_last_3d`` so cuDNN's NHWC 3D
+    conv kernels run end-to-end. The caller (:func:`enable_memory_efficient_decode`)
+    is responsible for converting the input sample and decoder weights to NHWC.
     """
     causal = decoder.causal
     batch_size = sample.shape[0]
@@ -537,6 +569,11 @@ def _memory_efficient_forward(
             raise ValueError("'timestep' required when timestep_conditioning=True")
         scaled_timestep = timestep * decoder.timestep_scale_multiplier.to(sample)
 
+    # Workspaces are unconditionally NHWC: rearrange + slice + residual-add
+    # inside _upsample_forward_efficient produces NCHW-default output, so
+    # per-tensor inspection would silently fall back to NCHW for every
+    # workspace after the first upsample.
+
     # --- Up blocks (dispatch to efficient path per block type) ---
     for up_block in decoder.up_blocks:
         if isinstance(up_block, UNetMidBlock3D):
@@ -546,15 +583,16 @@ def _memory_efficient_forward(
                 causal=causal,
                 timestep=scaled_timestep if decoder.timestep_conditioning else None,
                 generator=generator,
+                prefer_channels_last_3d=True,
             )
         elif isinstance(up_block, DepthToSpaceUpsample):
-            sample = _upsample_forward_efficient(up_block, sample, causal=causal)
+            sample = _upsample_forward_efficient(up_block, sample, causal=causal, prefer_channels_last_3d=True)
         elif isinstance(up_block, ResnetBlock3D):
             sample = up_block(sample, causal=causal, generator=generator)
         else:
             sample = up_block(sample, causal=causal)
 
-    return _final_norm_and_conv_out(decoder, sample, causal, scaled_timestep, batch_size)
+    return _final_norm_and_conv_out(decoder, sample, causal, scaled_timestep, batch_size, prefer_channels_last_3d=True)
 
 
 # ---------------------------------------------------------------------------
@@ -564,6 +602,10 @@ def _memory_efficient_forward(
 
 def enable_memory_efficient_decode(decoder: nn.Module) -> nn.Module:
     """Patch a ``VideoDecoder`` to use the memory-efficient forward path.
+    The mem-efficient path runs the decoder in ``channels_last_3d`` memory
+    format: weights and inputs are converted on first call so cuDNN's NHWC
+    3D conv kernels are used (~2x faster, avoids the large vol2col scratch
+    buffer of the NCHW path).
     The original ``forward`` is saved as ``decoder._original_forward`` so
     that it can be restored later with :func:`disable_memory_efficient_decode`.
     """
@@ -577,12 +619,20 @@ def enable_memory_efficient_decode(decoder: nn.Module) -> nn.Module:
         return decoder
 
     original_forward = decoder.forward
+    weights_converted = False
 
     def efficient_forward(
         sample: torch.Tensor,
         timestep: torch.Tensor | None = None,
         generator: torch.Generator | None = None,
     ) -> torch.Tensor:
+        nonlocal weights_converted
+        if sample.dim() == 5:
+            if not weights_converted:
+                # Lazy: weights are real by first-call time (meta -> loader -> here).
+                decoder.to(memory_format=torch.channels_last_3d)
+                weights_converted = True
+            sample = sample.to(memory_format=torch.channels_last_3d)
         return _memory_efficient_forward(decoder, sample, timestep, generator)
 
     decoder._original_forward = original_forward  # type: ignore[attr-defined]

packages/ltx-core/src/ltx_core/quantization/__init__.py

@@ -2,7 +2,9 @@ from ltx_core.quantization.fp8_cast import (
     TRANSFORMER_LINEAR_DOWNCAST_MAP,
     UPCAST_DURING_INFERENCE,
     UpcastWithStochasticRounding,
+    fp8_cast_fuse_rule,
 )
+from ltx_core.quantization.fp8_scaled_mm import fp8_scaled_mm_fuse_rule
 from ltx_core.quantization.policy import QuantizationPolicy
 
 __all__ = [
@@ -10,4 +12,6 @@ __all__ = [
     "UPCAST_DURING_INFERENCE",
     "QuantizationPolicy",
     "UpcastWithStochasticRounding",
+    "fp8_cast_fuse_rule",
+    "fp8_scaled_mm_fuse_rule",
 ]

packages/ltx-core/src/ltx_core/quantization/fp8_cast.py

@@ -1,9 +1,15 @@
+from pathlib import Path
+
+import safetensors
 import torch
 
+from ltx_core.loader.fuse_loras import FuseRule, bf16_fuse_rule
 from ltx_core.loader.kernels import TRITON_AVAILABLE
 from ltx_core.loader.module_ops import ModuleOps
+from ltx_core.loader.primitives import StateDict
 from ltx_core.loader.sd_ops import KeyValueOperationResult, SDOps
 from ltx_core.model.transformer.model import LTXModel
+from ltx_core.quantization.policy import QuantizationPolicy
 
 BLOCK_SIZE = 1024
 
@@ -113,58 +119,65 @@ def _replace_fwd_with_upcast(layer: torch.nn.Linear, with_stochastic_rounding: b
         )
 
 
+# Module-name suffixes for the Linears that participate in fp8 cast. Used by
+# both the upcast matcher and the sd_ops downcast map so the two cannot drift.
+# - ``.to_q`` / ``.to_k`` / ``.to_v`` / ``.to_out.0`` have a leading dot so they
+#   only match the attention Linears at ``...attnN.to_q`` etc.
+# - ``ff.net.0.proj`` / ``ff.net.2`` are intentionally **dotless** so they match
+#   both video FF (``...ff.net.0.proj``) and audio FF (``...audio_ff.net.0.proj``).
+_FP8_CAST_KEY_PREFIX = "transformer_blocks."
+_FP8_CAST_LINEAR_SUFFIXES: tuple[str, ...] = (
+    ".to_q",
+    ".to_k",
+    ".to_v",
+    ".to_out.0",
+    "ff.net.0.proj",
+    "ff.net.2",
+)
+
+
+def _is_fp8_cast_linear(module_name: str) -> bool:
+    """Return True if *module_name* names a Linear that should be fp8-cast."""
+    if _FP8_CAST_KEY_PREFIX not in module_name:
+        return False
+    return any(module_name.endswith(suffix) for suffix in _FP8_CAST_LINEAR_SUFFIXES)
+
+
 def _amend_forward_with_upcast(
     model: torch.nn.Module, with_stochastic_rounding: bool = False, seed: int = 0
 ) -> torch.nn.Module:
     """
-    Replace the forward method of the model's Linear layers to forward
+    Replace the forward method of the fp8-cast Linear layers (per
-    with upcast and optional stochastic rounding.
+    :data:`_FP8_CAST_LINEAR_SUFFIXES`) to forward with upcast and optional
+    stochastic rounding.
+    Only the Linears whose weights are downcast by :data:`TRANSFORMER_LINEAR_DOWNCAST_MAP`
+    are retyped. Linears outside that subset (e.g. ``to_gate_logits``) are left as
+    plain ``nn.Linear`` so the meta-model param dtype matches the loaded checkpoint
+    dtype.
     """
-    for m in model.modules():
+    for name, m in model.named_modules():
-        if isinstance(m, (torch.nn.Linear)):
+        if isinstance(m, torch.nn.Linear) and _is_fp8_cast_linear(name):
             _replace_fwd_with_upcast(m, with_stochastic_rounding, seed)
     return model
 
 
-TRANSFORMER_LINEAR_DOWNCAST_MAP = (
+def _build_transformer_linear_downcast_map() -> SDOps:
-    SDOps("TRANSFORMER_LINEAR_DOWNCAST_MAP")
+    """Build the sd_ops downcast map from the same suffix registry as the matcher."""
-    .with_kv_operation(
+    ops = SDOps("TRANSFORMER_LINEAR_DOWNCAST_MAP")
-        key_prefix="transformer_blocks.", key_suffix=".to_q.weight", operation=_naive_weight_or_bias_downcast
+    for suffix in _FP8_CAST_LINEAR_SUFFIXES:
+        ops = ops.with_kv_operation(
+            key_prefix=_FP8_CAST_KEY_PREFIX,
+            key_suffix=suffix + ".weight",
+            operation=_naive_weight_or_bias_downcast,
+        ).with_kv_operation(
+            key_prefix=_FP8_CAST_KEY_PREFIX,
+            key_suffix=suffix + ".bias",
+            operation=_naive_weight_or_bias_downcast,
         )
-    .with_kv_operation(
+    return ops
-        key_prefix="transformer_blocks.", key_suffix=".to_q.bias", operation=_naive_weight_or_bias_downcast
+
-    )
+
-    .with_kv_operation(
+TRANSFORMER_LINEAR_DOWNCAST_MAP = _build_transformer_linear_downcast_map()
-        key_prefix="transformer_blocks.", key_suffix=".to_k.weight", operation=_naive_weight_or_bias_downcast
-    )
-    .with_kv_operation(
-        key_prefix="transformer_blocks.", key_suffix=".to_k.bias", operation=_naive_weight_or_bias_downcast
-    )
-    .with_kv_operation(
-        key_prefix="transformer_blocks.", key_suffix=".to_v.weight", operation=_naive_weight_or_bias_downcast
-    )
-    .with_kv_operation(
-        key_prefix="transformer_blocks.", key_suffix=".to_v.bias", operation=_naive_weight_or_bias_downcast
-    )
-    .with_kv_operation(
-        key_prefix="transformer_blocks.", key_suffix=".to_out.0.weight", operation=_naive_weight_or_bias_downcast
-    )
-    .with_kv_operation(
-        key_prefix="transformer_blocks.", key_suffix=".to_out.0.bias", operation=_naive_weight_or_bias_downcast
-    )
-    .with_kv_operation(
-        key_prefix="transformer_blocks.", key_suffix="ff.net.0.proj.weight", operation=_naive_weight_or_bias_downcast
-    )
-    .with_kv_operation(
-        key_prefix="transformer_blocks.", key_suffix="ff.net.0.proj.bias", operation=_naive_weight_or_bias_downcast
-    )
-    .with_kv_operation(
-        key_prefix="transformer_blocks.", key_suffix="ff.net.2.weight", operation=_naive_weight_or_bias_downcast
-    )
-    .with_kv_operation(
-        key_prefix="transformer_blocks.", key_suffix="ff.net.2.bias", operation=_naive_weight_or_bias_downcast
-    )
-)
 
 UPCAST_DURING_INFERENCE = ModuleOps(
     name="upcast_fp8_during_linear_forward",
@@ -186,3 +199,139 @@ class UpcastWithStochasticRounding(ModuleOps):
             matcher=lambda model: isinstance(model, LTXModel),
             mutator=lambda model: _amend_forward_with_upcast(model, True, seed),
         )
+
+
+def fuse_cast_fp8_weight(
+    delta_bf16: torch.Tensor,
+    weight_fp8: torch.Tensor,
+) -> torch.Tensor:
+    """Return ``(delta_bf16 + dequantize(weight_fp8)).to(weight_fp8.dtype)``.
+    CUDA with Triton uses stochastic rounding via the fused kernel; otherwise
+    falls back to a deterministic bf16 add. ``delta_bf16`` is the bf16
+    accumulator and is mutated in place.
+    """
+    if delta_bf16.dtype != torch.bfloat16:
+        raise ValueError(f"delta_bf16 must be bfloat16, got {delta_bf16.dtype}")
+    if str(weight_fp8.device).startswith("cuda") and TRITON_AVAILABLE:
+        fused_add_round_launch(delta_bf16, weight_fp8, seed=0)
+    else:
+        delta_bf16.add_(weight_fp8.to(dtype=torch.bfloat16))
+    return delta_bf16.to(dtype=weight_fp8.dtype)
+
+
+def _fp8_cast_fuse(
+    key: str,
+    weight: torch.Tensor,
+    deltas: torch.Tensor,
+    model_sd: StateDict,
+) -> dict[str, torch.Tensor]:
+    """Cast the dequantized FP8 weight + BF16 deltas back to ``weight.dtype``
+    (FP8) via the fused-add-round kernel on CUDA.
+    Only a subset of linears are FP8-downcast (see ``TRANSFORMER_LINEAR_DOWNCAST_MAP``);
+    LoRAs may also target layers left in BF16 (e.g. audio ``add_q/k/v_proj``, cross-modal
+    projections). For those, fall back to a plain BF16 fuse.
+    """
+    if weight.dtype not in (torch.float8_e4m3fn, torch.float8_e5m2):
+        return bf16_fuse_rule(key, weight, deltas, model_sd)
+    return {key: fuse_cast_fp8_weight(deltas, weight)}
+
+
+fp8_cast_fuse_rule = FuseRule(aggregation_dtype=torch.bfloat16, fuse_fn=_fp8_cast_fuse)
+
+
+# Raw safetensors storage prefix shared by every diffusion-transformer
+# parameter (and every prequant `*_scale` sibling). Verified against
+# ltx-2.3-22b-{dev,distilled}-fp8.safetensors: 2924/2924 and 2992/2992 of
+# the scale keys start with this exact prefix.
+_RAW_DIFFUSION_MODEL_PREFIX = "model.diffusion_model."
+
+
+def _read_scales(checkpoint_path: str | Path) -> dict[str, torch.Tensor]:
+    """Return ``{post_rename_param_key: scale_tensor}`` for every prequant
+    ``*_scale`` sibling in *checkpoint_path*.
+    Keys are returned in the post-rename form the loader will pass to the
+    sd-op (e.g. ``transformer_blocks.0.attn1.to_q.weight``) -- the raw
+    ``model.diffusion_model.`` prefix and the ``_scale`` suffix are both
+    stripped. Catches both ``.weight_scale`` and ``.bias_scale``; the
+    latter is absent in the current LTX-2.3 prequant checkpoints but
+    accepted for forward compatibility.
+    """
+    out: dict[str, torch.Tensor] = {}
+    with safetensors.safe_open(str(checkpoint_path), framework="pt", device="cpu") as h:
+        raw_keys = h.keys()
+        for k in raw_keys:
+            if not k.endswith("_scale"):
+                continue
+            if not k.startswith(_RAW_DIFFUSION_MODEL_PREFIX):
+                raise ValueError(
+                    f"Scale key {k!r} does not start with the expected raw prefix {_RAW_DIFFUSION_MODEL_PREFIX!r}"
+                )
+            param_key = k.removeprefix(_RAW_DIFFUSION_MODEL_PREFIX).removesuffix("_scale")
+            out[param_key] = h.get_tensor(k)
+    return out
+
+
+def _build_prequant_fold_sd_ops(scales: dict[str, torch.Tensor]) -> SDOps:
+    """Build sd-ops that fold prequant ``*_scale`` siblings into their parent
+    tensor at load time.
+    *scales* is keyed by the **post-rename** param key (e.g.
+    ``transformer_blocks.0.attn1.to_q.weight``); see :func:`_read_scales`.
+    Four ``with_kv_operation`` entries (symmetric for ``.weight`` and ``.bias``):
+    * ``.weight`` / ``.bias`` -> if a sibling scale exists in *scales*, fold;
+      then delegate to ``TRANSFORMER_LINEAR_DOWNCAST_MAP`` (downcast covered
+      Linears, pass everything else through). Without a scale, delegate
+      directly.
+    * ``.weight_scale`` / ``.bias_scale`` -> drop (the scale is consumed by
+      the fold). Raises if the scale key doesn't correspond to a known
+      entry in *scales* -- that means the file shipped a scale we didn't
+      pre-register, which would silently desync the fold.
+    """
+
+    def _on_param(param_key: str, value: torch.Tensor) -> list[KeyValueOperationResult]:
+        scale = scales.get(param_key)
+        if scale is None:
+            return TRANSFORMER_LINEAR_DOWNCAST_MAP.apply_to_key_value(param_key, value)
+        scale = scale.to(device=value.device)
+        if scale.ndim != 0:
+            raise ValueError(f"Unsupported scale shape {tuple(scale.shape)} for {param_key}")
+        bf16 = (value.to(torch.float32) * scale).to(torch.bfloat16)
+        # Delegate the final fp8-vs-bf16 decision to the downcast map: Linears
+        # outside the fp8 subset (e.g. to_gate_logits) stay bf16 to match the
+        # plain nn.Linear that the upcast matcher leaves untouched.
+        return TRANSFORMER_LINEAR_DOWNCAST_MAP.apply_to_key_value(param_key, bf16)
+
+    def _drop_scale(scale_key: str, _value: torch.Tensor) -> list[KeyValueOperationResult]:
+        param_key = scale_key.removesuffix("_scale")
+        if param_key not in scales:
+            raise ValueError(
+                f"Scale key {scale_key!r} has no matching entry in the prequant scales dict; "
+                f"_read_scales and the loader's rename map have drifted"
+            )
+        return []
+
+    # Register the drop ops first so the dict-membership sanity check is the
+    # earliest sd-op that can fire on a scale key -- we crash on a stray scale
+    # before any silently mismatched fold has a chance to land in the state
+    # dict. Registration order is irrelevant for correctness (no overlap
+    # between matchers) but communicates intent.
+    return (
+        SDOps("FP8_CAST_PREQUANT_AWARE")
+        .with_kv_operation(key_suffix=".weight_scale", operation=_drop_scale)
+        .with_kv_operation(key_suffix=".bias_scale", operation=_drop_scale)
+        .with_kv_operation(key_suffix=".weight", operation=_on_param)
+        .with_kv_operation(key_suffix=".bias", operation=_on_param)
+    )
+
+
+def build_policy(checkpoint_path: str | Path) -> QuantizationPolicy:
+    """FP8 casting with upcasting during inference.
+    *checkpoint_path* is required (mirroring ``fp8_scaled_mm.build_policy``).
+    For prequantized fp8 checkpoints, sibling ``*_scale`` tensors (weight or
+    bias) are folded into the parent at load time.
+    """
+    scales = _read_scales(checkpoint_path)
+    return QuantizationPolicy(
+        sd_ops=_build_prequant_fold_sd_ops(scales),
+        module_ops=(UPCAST_DURING_INFERENCE,),
+        fuse_rule=fp8_cast_fuse_rule,
+    )

packages/ltx-core/src/ltx_core/quantization/fp8_scaled_mm.py

@@ -5,8 +5,11 @@ from typing import Callable
 import torch
 from torch import nn
 
+from ltx_core.loader.fuse_loras import FuseRule, bf16_fuse_rule
 from ltx_core.loader.module_ops import ModuleOps
+from ltx_core.loader.primitives import StateDict
 from ltx_core.model.transformer import LTXModel
+from ltx_core.quantization.policy import QuantizationPolicy
 from ltx_core.quantization.trtllm_scaled_usable import trtllm_scaled_mm_usable
 
 
@@ -173,3 +176,42 @@ def get_fp8_swap_module_ops(checkpoint_path: str) -> tuple[ModuleOps, ...]:
             mutator=lambda model: _swap_linears_to_fp8(model, _should_swap),
         ),
     )
+
+
+def _fp8_scaled_mm_fuse(
+    key: str,
+    weight: torch.Tensor,
+    deltas: torch.Tensor,
+    model_sd: StateDict,
+) -> dict[str, torch.Tensor]:
+    """Dequantize via ``weight.float() * weight_scale``, add the BF16 delta,
+    and re-quantize to FP8 with a fresh per-tensor scale.
+    Layers that were not swapped to scaled FP8 (e.g. small embedder linears
+    excluded from the auto-discovered swap set) stay BF16 and have no
+    ``.weight_scale`` companion -- for those, fall back to a plain bf16 fuse.
+    """
+    scale_key = key.replace(".weight", ".weight_scale")
+    if scale_key not in model_sd.sd:
+        return bf16_fuse_rule(key, weight, deltas, model_sd)
+    weight_scale = model_sd.sd[scale_key]
+    original_weight = weight.to(torch.float32) * weight_scale
+    new_weight = original_weight + deltas.to(torch.float32)
+    new_fp8_weight, new_weight_scale = quantize_weight_to_fp8_per_tensor(new_weight)
+    return {key: new_fp8_weight, scale_key: new_weight_scale}
+
+
+fp8_scaled_mm_fuse_rule = FuseRule(aggregation_dtype=torch.bfloat16, fuse_fn=_fp8_scaled_mm_fuse)
+
+
+def build_policy(checkpoint_path: str) -> QuantizationPolicy:
+    """FP8 scaled matmul for checkpoints pre-quantized with per-tensor scales.
+    The set of layers to swap to ``FP8Linear`` is discovered from the
+    checkpoint's ``.weight_scale`` tensors via suffix-matching against the
+    model's named modules. Requires a pre-quantized checkpoint; for BF16
+    checkpoints, use :func:`ltx_core.quantization.fp8_cast.build_policy`.
+    """
+    return QuantizationPolicy(
+        sd_ops=None,
+        module_ops=get_fp8_swap_module_ops(checkpoint_path),
+        fuse_rule=fp8_scaled_mm_fuse_rule,
+    )

packages/ltx-core/src/ltx_core/quantization/policy.py

@@ -1,48 +1,24 @@
 from dataclasses import dataclass
-from enum import Enum
 
+from ltx_core.loader.fuse_loras import FuseRule, bf16_fuse_rule
 from ltx_core.loader.module_ops import ModuleOps
 from ltx_core.loader.sd_ops import SDOps
-from ltx_core.quantization.fp8_cast import TRANSFORMER_LINEAR_DOWNCAST_MAP, UPCAST_DURING_INFERENCE
+from ltx_core.model.model_protocol import ModelConfigurator
-from ltx_core.quantization.fp8_scaled_mm import get_fp8_swap_module_ops
+from ltx_core.model.transformer.model import LTXModel
 
 
 @dataclass(frozen=True)
 class QuantizationPolicy:
     """Configuration for model quantization during loading.
     Attributes:
-        kind: Discriminator for the policy variant.
         sd_ops: State-dict operations applied to each tensor during load.
         module_ops: Post-load module transformations applied to the meta model.
+        model_configurator: Configurator class to use when constructing the transformer.
+        fuse_rule: How LoRA deltas merge into this policy's weight layout.
+            Default ``bf16_fuse_rule`` is used when no policy is configured.
     """
 
-    class Kind(str, Enum):
-        FP8_CAST = "fp8_cast"
-        FP8_SCALED_MM = "fp8_scaled_mm"
-
-    kind: Kind
     sd_ops: SDOps | None = None
     module_ops: tuple[ModuleOps, ...] = ()
-
+    model_configurator: type[ModelConfigurator[LTXModel]] | None = None
-    @classmethod
+    fuse_rule: FuseRule = bf16_fuse_rule
-    def fp8_cast(cls) -> "QuantizationPolicy":
-        """FP8 casting with upcasting during inference."""
-        return cls(
-            kind=cls.Kind.FP8_CAST,
-            sd_ops=TRANSFORMER_LINEAR_DOWNCAST_MAP,
-            module_ops=(UPCAST_DURING_INFERENCE,),
-        )
-
-    @classmethod
-    def fp8_scaled_mm(cls, checkpoint_path: str) -> "QuantizationPolicy":
-        """FP8 scaled matmul for checkpoints pre-quantized with per-tensor scales.
-        The set of layers to swap to ``FP8Linear`` is discovered from the
-        checkpoint's ``.weight_scale`` tensors via suffix-matching against the
-        model's named modules. Requires a pre-quantized checkpoint; for BF16
-        checkpoints, use :meth:`fp8_cast` instead.
-        """
-        return cls(
-            kind=cls.Kind.FP8_SCALED_MM,
-            sd_ops=None,
-            module_ops=get_fp8_swap_module_ops(checkpoint_path),
-        )

packages/ltx-core/src/ltx_core/text_encoders/gemma/encoders/base_encoder.py

@@ -183,7 +183,7 @@ def module_ops_from_gemma_root(gemma_root: str) -> tuple[ModuleOps, ...]:
         return module
 
     def load_processor(module: GemmaTextEncoder) -> GemmaTextEncoder:
-        image_processor = AutoImageProcessor.from_pretrained(processor_root, local_files_only=True)
+        image_processor = AutoImageProcessor.from_pretrained(processor_root, local_files_only=True, use_fast=False)
         if not module.tokenizer:
             raise ValueError("Tokenizer model operation must be performed before processor model operation")
         module.processor = Gemma3Processor(image_processor=image_processor, tokenizer=module.tokenizer.tokenizer)

packages/ltx-core/src/ltx_core/text_encoders/gemma/encoders/encoder_configurator.py

@@ -1,4 +1,5 @@
 import torch
+import transformers
 from transformers import Gemma3Config
 from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS
 from transformers.models.gemma3 import Gemma3ForConditionalGeneration
@@ -19,6 +20,8 @@ from ltx_core.text_encoders.gemma.feature_extractor import (
     FeatureExtractorV2,
 )
 
+_TRANSFORMERS_V5: bool = int(transformers.__version__.split(".", 1)[0]) >= 5
+
 
 class GemmaTextEncoderConfigurator(ModelConfigurator[GemmaTextEncoder]):
     @classmethod
@@ -100,14 +103,30 @@ def _create_feature_extractor(transformer_config: dict) -> torch.nn.Module:
 
 # --- Split SDOps: Gemma LLM keys vs Embeddings Processor keys ---
 
-GEMMA_LLM_KEY_OPS = (
+
+def _build_gemma_llm_key_ops(*, transformers_v5: bool) -> SDOps:
+    """Build the checkpoint-key remapping for the Gemma multimodal encoder.
+    The vision-tower mapping differs between transformers <5 and >=5 because
+    upstream PR https://github.com/huggingface/transformers/pull/39847 flattened
+    ``Gemma3ForConditionalGeneration.model.vision_tower.vision_model`` into
+    ``model.vision_tower``. Checkpoints continue to ship the legacy
+    ``vision_tower.vision_model.*`` prefix, so we strip the inner ``vision_model.``
+    when targeting v5 and pass it through unchanged for v4.
+    """
+    base = (
         SDOps("GEMMA_LLM_KEY_OPS")
         # 1. Map language model layers (note the double .model prefix)
         .with_matching(prefix="language_model.model.")
         .with_replacement("language_model.model.", "model.model.language_model.")
-    # 2. Map the Vision Tower
+        # 2. Map the Vision Tower (version-dependent — see docstring)
         .with_matching(prefix="vision_tower.")
-    .with_replacement("vision_tower.", "model.model.vision_tower.")
+    )
+    if transformers_v5:
+        base = base.with_replacement("vision_tower.vision_model.", "model.model.vision_tower.")
+    else:
+        base = base.with_replacement("vision_tower.", "model.model.vision_tower.")
+    return (
+        base
         # 3. Map the Multi-Modal Projector
         .with_matching(prefix="multi_modal_projector.")
         .with_replacement("multi_modal_projector.", "model.model.multi_modal_projector.")
@@ -119,7 +138,10 @@ GEMMA_LLM_KEY_OPS = (
             ],
             key_prefix="model.model.language_model.embed_tokens.weight",
         )
-)
+    )
+
+
+GEMMA_LLM_KEY_OPS = _build_gemma_llm_key_ops(transformers_v5=_TRANSFORMERS_V5)
 
 EMBEDDINGS_PROCESSOR_KEY_OPS = (
     SDOps("EMBEDDINGS_PROCESSOR_KEY_OPS")
@@ -152,24 +174,85 @@ VIDEO_ONLY_EMBEDDINGS_PROCESSOR_KEY_OPS = (
 )
 
 
+def _resolve_local_base_freq(config: object) -> float:
+    rope_parameters = getattr(config, "rope_parameters", None)
+    if isinstance(rope_parameters, dict) and "sliding_attention" in rope_parameters:
+        sliding = rope_parameters["sliding_attention"]
+        if isinstance(sliding, dict) and "rope_theta" in sliding:
+            return float(sliding["rope_theta"])
+    if hasattr(config, "rope_local_base_freq"):
+        return float(config.rope_local_base_freq)
+    raise AttributeError(
+        "Gemma text_config exposes neither rope_local_base_freq nor rope_parameters['sliding_attention']['rope_theta']"
+    )
+
+
+def _resolve_full_rope_type(config: object) -> str:
+    rope_parameters = getattr(config, "rope_parameters", None)
+    if isinstance(rope_parameters, dict) and "full_attention" in rope_parameters:
+        full = rope_parameters["full_attention"]
+        if isinstance(full, dict) and "rope_type" in full:
+            return str(full["rope_type"])
+    rope_scaling = getattr(config, "rope_scaling", None)
+    if rope_scaling is not None:
+        if isinstance(rope_scaling, dict):
+            if "rope_type" in rope_scaling:
+                return str(rope_scaling["rope_type"])
+        elif hasattr(rope_scaling, "rope_type"):
+            return str(rope_scaling.rope_type)
+    raise AttributeError(
+        "Gemma text_config exposes neither rope_scaling.rope_type nor rope_parameters['full_attention']['rope_type']"
+    )
+
+
+def _populate_rotary_v4(l_model: torch.nn.Module, config: object) -> None:
+    """transformers <5 layout: separate ``rotary_emb_local`` + ``rotary_emb``."""
+    dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
+    base = _resolve_local_base_freq(config)
+    local_inv = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(dtype=torch.float) / dim))
+    full_inv, _ = ROPE_INIT_FUNCTIONS[_resolve_full_rope_type(config)](config)
+    l_model.rotary_emb_local.register_buffer("inv_freq", local_inv)
+    l_model.rotary_emb.register_buffer("inv_freq", full_inv)
+
+
+def _populate_rotary_v5(l_model: torch.nn.Module, config: object) -> None:
+    """transformers >=5 layout: single ``rotary_emb`` with per-layer-type buffers.
+    Mirrors ``Gemma3PreTrainedModel._init_weights`` for ``Gemma3RotaryEmbedding``
+    so meta-built models reach the same numerical state as a from_pretrained load.
+    """
+    rope_emb = l_model.rotary_emb
+    for layer_type in dict.fromkeys(config.layer_types):
+        rope_params = config.rope_parameters[layer_type]
+        if rope_params is None:
+            continue
+        rope_type = rope_params["rope_type"]
+        if rope_type == "default":
+            inv_freq, attn_scaling = rope_emb.compute_default_rope_parameters(config, layer_type=layer_type)
+        else:
+            inv_freq, attn_scaling = ROPE_INIT_FUNCTIONS[rope_type](config, layer_type=layer_type)
+        rope_emb.register_buffer(f"{layer_type}_inv_freq", inv_freq, persistent=False)
+        rope_emb.register_buffer(f"{layer_type}_original_inv_freq", inv_freq.clone(), persistent=False)
+        setattr(rope_emb, f"{layer_type}_attention_scaling", attn_scaling)
+
+
 def create_and_populate(module: GemmaTextEncoder) -> GemmaTextEncoder:
     model = module.model
-    v_model = model.model.vision_tower.vision_model
+    v_tower = model.model.vision_tower
+    v_model = v_tower.vision_model if hasattr(v_tower, "vision_model") else v_tower
     l_model = model.model.language_model
 
     config = model.config.text_config
-    dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
+
-    base = config.rope_local_base_freq
+    if hasattr(l_model, "rotary_emb_local"):
-    local_rope_freqs = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(dtype=torch.float) / dim))
+        _populate_rotary_v4(l_model, config)
-    inv_freqs, _ = ROPE_INIT_FUNCTIONS[config.rope_scaling["rope_type"]](config)
+    else:
+        _populate_rotary_v5(l_model, config)
 
     positions_length = len(v_model.embeddings.position_ids[0])
     position_ids = torch.arange(positions_length, dtype=torch.long, device="cpu").unsqueeze(0)
     v_model.embeddings.register_buffer("position_ids", position_ids)
-    embed_scale = torch.tensor(model.config.text_config.hidden_size**0.5, device="cpu")
+    embed_scale = torch.tensor(config.hidden_size**0.5, device="cpu")
     l_model.embed_tokens.register_buffer("embed_scale", embed_scale)
-    l_model.rotary_emb_local.register_buffer("inv_freq", local_rope_freqs)
-    l_model.rotary_emb.register_buffer("inv_freq", inv_freqs)
 
     return module
 

packages/ltx-pipelines/README.md

@@ -369,7 +369,9 @@ PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True python -m ltx_pipelines.ti2vid_
 When authoring custom scripts, pass a `QuantizationPolicy` to pipeline classes:
 
 ```python
-from ltx_core.quantization import QuantizationPolicy
+from ltx_core.quantization.fp8_cast import build_policy as build_fp8_cast_policy
+# Alternative:
+# from ltx_core.quantization.fp8_scaled_mm import build_policy as build_fp8_scaled_mm_policy
 
 pipeline = TI2VidTwoStagesPipeline(
     checkpoint_path=ltx_model_path,
@@ -377,7 +379,7 @@ pipeline = TI2VidTwoStagesPipeline(
     spatial_upsampler_path=upsampler_path,
     gemma_root=gemma_root_path,
     loras=[],
-    quantization=QuantizationPolicy.fp8_cast(),  # or QuantizationPolicy.fp8_scaled_mm()
+    quantization=build_fp8_cast_policy(ltx_model_path),
 )
 pipeline(...)
 ```
@@ -414,12 +416,13 @@ def denoising_loop(sigmas, video_state, audio_state, stepper):
         video_state=video_state,
         audio_state=audio_state,
         stepper=stepper,
-        denoise_fn=your_denoise_function,
+        transformer=transformer,
+        denoiser=denoiser,
         ge_gamma=2.0,  # Gradient estimation coefficient
     )
 ```
 
-This allows you to use **20-30 steps instead of 40** while maintaining quality. The gradient estimation function is available in [`pipeline_utils.py`](src/ltx_pipelines/utils/helpers.py).
+This allows you to use **20-30 steps instead of 40** while maintaining quality. The gradient estimation function is defined in [`samplers.py`](src/ltx_pipelines/utils/samplers.py).
 
 ---
 
@@ -435,15 +438,18 @@ This allows you to use **20-30 steps instead of 40** while maintaining quality.
 ## 📖 Example: Image-to-Video
 
 ```python
-from ltx_core.loader import LTXV_LORA_COMFY_RENAMING_MAP, LoraPathStrengthAndSDOps
-from ltx_pipelines.ti2vid_two_stages import TI2VidTwoStagesPipeline
 from ltx_core.components.guiders import MultiModalGuiderParams
+from ltx_core.loader import LTXV_LORA_COMFY_RENAMING_MAP, LoraPathStrengthAndSDOps
+from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
+from ltx_pipelines.ti2vid_two_stages import TI2VidTwoStagesPipeline
+from ltx_pipelines.utils.args import ImageConditioningInput
+from ltx_pipelines.utils.media_io import encode_video
 
 distilled_lora = [
     LoraPathStrengthAndSDOps(
         "/path/to/distilled_lora.safetensors",
         0.6,
-        LTXV_LORA_COMFY_RENAMING_MAP
+        LTXV_LORA_COMFY_RENAMING_MAP,
     ),
 ]
 
@@ -473,19 +479,31 @@ audio_guider_params = MultiModalGuiderParams(
     stg_blocks=[29],
 )
 
-# Generate video from image
+# Generate video from image. The pipeline returns (video_iterator, audio);
-pipeline(
+# the caller is responsible for encoding to file via encode_video().
+num_frames = 121
+frame_rate = 25.0
+tiling_config = TilingConfig.default()
+video, audio = pipeline(
     prompt="A serene landscape with mountains in the background",
-    output_path="output.mp4",
+    negative_prompt="worst quality, low quality, blurry, distorted",
     seed=42,
     height=512,
     width=768,
-    num_frames=121,
+    num_frames=num_frames,
-    frame_rate=25.0,
+    frame_rate=frame_rate,
     num_inference_steps=40,
     video_guider_params=video_guider_params,
     audio_guider_params=audio_guider_params,
-    images=[ImageConditioningInput("input_image.jpg", 0, 1.0, 33)],  # Image at frame 0, strength 1.0, CRF 33
+    images=[ImageConditioningInput("input_image.jpg", 0, 1.0, 33)],  # path, frame_idx=0, strength=1.0, crf=33
+    tiling_config=tiling_config,
+)
+encode_video(
+    video=video,
+    fps=frame_rate,
+    audio=audio,
+    output_path="output.mp4",
+    video_chunks_number=get_video_chunks_number(num_frames, tiling_config),
 )
 ```
 

packages/ltx-pipelines/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "ltx-pipelines"
-version = "1.1.3"
+version = "1.1.4"
 description = "Pipelines implementation for Lightricks' LTX-2 model"
 readme = "README.md"
 requires-python = ">=3.10"

packages/ltx-pipelines/src/ltx_pipelines/a2vid_two_stage.py

@@ -9,6 +9,7 @@ from ltx_core.components.schedulers import LTX2Scheduler
 from ltx_core.loader import LoraPathStrengthAndSDOps
 from ltx_core.loader.registry import Registry
 from ltx_core.model.audio_vae import encode_audio as vae_encode_audio
+from ltx_core.model.transformer.compiling import CompilationConfig
 from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
 from ltx_core.quantization import QuantizationPolicy
 from ltx_core.types import Audio, AudioLatentShape, VideoPixelShape
@@ -52,7 +53,7 @@ class A2VidPipelineTwoStage:
         device: torch.device | None = None,
         quantization: QuantizationPolicy | None = None,
         registry: Registry | None = None,
-        torch_compile: bool = False,
+        compilation_config: CompilationConfig | None = None,
         offload_mode: OffloadMode = OffloadMode.NONE,
     ):
         self.device = device or get_device()
@@ -71,7 +72,7 @@ class A2VidPipelineTwoStage:
             loras=tuple(loras),
             quantization=quantization,
             registry=registry,
-            torch_compile=torch_compile,
+            compilation_config=compilation_config,
             offload_mode=offload_mode,
         )
         stage_2_loras = (*tuple(loras), *tuple(distilled_lora))
@@ -82,7 +83,7 @@ class A2VidPipelineTwoStage:
             loras=stage_2_loras,
             quantization=quantization,
             registry=registry,
-            torch_compile=torch_compile,
+            compilation_config=compilation_config,
             offload_mode=offload_mode,
         )
         self.upsampler = VideoUpsampler(
@@ -238,7 +239,7 @@ class A2VidPipelineTwoStage:
 
 @torch.inference_mode()
 def main() -> None:
-    logging.getLogger().setLevel(logging.INFO)
+    logging.basicConfig(level=logging.INFO)
     parser = default_2_stage_arg_parser()
     parser.add_argument(
         "--audio-path",
@@ -266,7 +267,7 @@ def main() -> None:
         gemma_root=args.gemma_root,
         loras=tuple(args.lora) if args.lora else (),
         quantization=args.quantization,
-        torch_compile=args.compile,
+        compilation_config=args.compile,
         offload_mode=args.offload_mode,
     )
     tiling_config = TilingConfig.default()

packages/ltx-pipelines/src/ltx_pipelines/distilled.py

@@ -6,6 +6,7 @@ import torch
 from ltx_core.components.noisers import GaussianNoiser
 from ltx_core.loader import LoraPathStrengthAndSDOps
 from ltx_core.loader.registry import Registry
+from ltx_core.model.transformer.compiling import CompilationConfig
 from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
 from ltx_core.quantization import QuantizationPolicy
 from ltx_core.types import Audio
@@ -53,7 +54,7 @@ class DistilledPipeline:
         device: torch.device | None = None,
         quantization: QuantizationPolicy | None = None,
         registry: Registry | None = None,
-        torch_compile: bool = False,
+        compilation_config: CompilationConfig | None = None,
         offload_mode: OffloadMode = OffloadMode.NONE,
     ):
         self.device = device or get_device()
@@ -75,7 +76,7 @@ class DistilledPipeline:
             loras=tuple(loras),
             quantization=quantization,
             registry=registry,
-            torch_compile=torch_compile,
+            compilation_config=compilation_config,
             offload_mode=offload_mode,
         )
         self.upsampler = VideoUpsampler(
@@ -180,7 +181,7 @@ class DistilledPipeline:
 
 @torch.inference_mode()
 def main() -> None:
-    logging.getLogger().setLevel(logging.INFO)
+    logging.basicConfig(level=logging.INFO)
     checkpoint_path = detect_checkpoint_path(distilled=True)
     params = detect_params(checkpoint_path)
     parser = default_2_stage_distilled_arg_parser(params=params)
@@ -191,7 +192,7 @@ def main() -> None:
         gemma_root=args.gemma_root,
         loras=tuple(args.lora) if args.lora else (),
         quantization=args.quantization,
-        torch_compile=args.compile,
+        compilation_config=args.compile,
         offload_mode=args.offload_mode,
     )
     tiling_config = TilingConfig.default()

packages/ltx-pipelines/src/ltx_pipelines/hdr_ic_lora.py

@@ -59,6 +59,7 @@ from ltx_pipelines.utils.constants import DISTILLED_SIGMA_VALUES, STAGE_2_DISTIL
 from ltx_pipelines.utils.denoisers import SimpleDenoiser
 from ltx_pipelines.utils.helpers import get_device, modality_from_latent_state
 from ltx_pipelines.utils.media_io import ResizeMode, align_resolution, load_video_conditioning_hdr
+from ltx_pipelines.utils.quantization_factory import QuantizationKind
 from ltx_pipelines.utils.types import ModalitySpec, OffloadMode
 
 logger = logging.getLogger(__name__)
@@ -77,7 +78,7 @@ ALIGNMENT_DIVISOR = 64
 # to the pipeline constructor.
 TILED_VAE_ENCODE_PIXEL_THRESHOLD = 512 * 768
 
-_DEFAULT_QUANTIZATION = QuantizationPolicy.fp8_cast()
+_DEFAULT_QUANTIZATION = QuantizationKind.FP8_CAST
 
 # Default stage-2 configuration: one refinement phase with modest 2-way tiling
 # in every dimension and a short 2-step distilled sigma schedule.
@@ -205,7 +206,7 @@ class HDRICLoraPipeline:
         hdr_lora: str | Path,
         text_embeddings_path: str | Path,
         device: torch.device | None = None,
-        quantization: QuantizationPolicy = _DEFAULT_QUANTIZATION,
+        quantization: QuantizationPolicy | QuantizationKind | None = _DEFAULT_QUANTIZATION,
         registry: Registry | None = None,
         hdr_lora_config: HdrLoraConfig | None = None,
         tiled_vae_encode_pixel_threshold: int = TILED_VAE_ENCODE_PIXEL_THRESHOLD,
@@ -232,6 +233,8 @@ class HDRICLoraPipeline:
         """
         self.device = device or get_device()
         self._tiled_vae_encode_threshold = tiled_vae_encode_pixel_threshold
+        if isinstance(quantization, QuantizationKind):
+            quantization = quantization.to_policy(checkpoint_path=distilled_checkpoint_path)
         if offload_mode != OffloadMode.NONE and quantization is not None:
             logger.info("Offload mode enabled — disabling quantization (not supported with layer streaming).")
             quantization = None

packages/ltx-pipelines/src/ltx_pipelines/ic_lora.py

@@ -7,6 +7,7 @@ from ltx_core.components.noisers import GaussianNoiser
 from ltx_core.conditioning import ConditioningItem
 from ltx_core.loader import LoraPathStrengthAndSDOps
 from ltx_core.loader.registry import Registry
+from ltx_core.model.transformer.compiling import CompilationConfig
 from ltx_core.model.video_vae import TilingConfig, VideoEncoder, get_video_chunks_number
 from ltx_core.quantization import QuantizationPolicy
 from ltx_core.types import Audio, VideoPixelShape
@@ -60,7 +61,7 @@ class ICLoraPipeline:
         device: torch.device | None = None,
         quantization: QuantizationPolicy | None = None,
         registry: Registry | None = None,
-        torch_compile: bool = False,
+        compilation_config: CompilationConfig | None = None,
         offload_mode: OffloadMode = OffloadMode.NONE,
     ):
         self.device = device or get_device()
@@ -82,7 +83,7 @@ class ICLoraPipeline:
             loras=tuple(loras),
             quantization=quantization,
             registry=registry,
-            torch_compile=torch_compile,
+            compilation_config=compilation_config,
             offload_mode=offload_mode,
         )
         self.stage_2 = DiffusionStage(
@@ -92,7 +93,7 @@ class ICLoraPipeline:
             loras=(),
             quantization=quantization,
             registry=registry,
-            torch_compile=torch_compile,
+            compilation_config=compilation_config,
             offload_mode=offload_mode,
         )
         self.upsampler = VideoUpsampler(
@@ -330,7 +331,7 @@ class ICLoraPipeline:
 
 @torch.inference_mode()
 def main() -> None:
-    logging.getLogger().setLevel(logging.INFO)
+    logging.basicConfig(level=logging.INFO)
     checkpoint_path = detect_checkpoint_path(distilled=True)
     params = detect_params(checkpoint_path)
     parser = default_2_stage_distilled_arg_parser(params=params)
@@ -386,7 +387,7 @@ def main() -> None:
         gemma_root=args.gemma_root,
         loras=tuple(args.lora) if args.lora else (),
         quantization=args.quantization,
-        torch_compile=args.compile,
+        compilation_config=args.compile,
         offload_mode=args.offload_mode,
     )
     tiling_config = TilingConfig.default()

packages/ltx-pipelines/src/ltx_pipelines/keyframe_interpolation.py

@@ -12,6 +12,7 @@ from ltx_core.components.noisers import GaussianNoiser
 from ltx_core.components.schedulers import LTX2Scheduler
 from ltx_core.loader import LoraPathStrengthAndSDOps
 from ltx_core.loader.registry import Registry
+from ltx_core.model.transformer.compiling import CompilationConfig
 from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
 from ltx_core.quantization import QuantizationPolicy
 from ltx_core.types import Audio, VideoPixelShape
@@ -62,7 +63,7 @@ class KeyframeInterpolationPipeline:
         device: torch.device | None = None,
         quantization: QuantizationPolicy | None = None,
         registry: Registry | None = None,
-        torch_compile: bool = False,
+        compilation_config: CompilationConfig | None = None,
         offload_mode: OffloadMode = OffloadMode.NONE,
     ):
         self.device = device or get_device()
@@ -80,7 +81,7 @@ class KeyframeInterpolationPipeline:
             loras=tuple(loras),
             quantization=quantization,
             registry=registry,
-            torch_compile=torch_compile,
+            compilation_config=compilation_config,
             offload_mode=offload_mode,
         )
         stage_2_loras = (*tuple(loras), *tuple(distilled_lora))
@@ -91,7 +92,7 @@ class KeyframeInterpolationPipeline:
             loras=stage_2_loras,
             quantization=quantization,
             registry=registry,
-            torch_compile=torch_compile,
+            compilation_config=compilation_config,
             offload_mode=offload_mode,
         )
         self.upsampler = VideoUpsampler(
@@ -233,7 +234,7 @@ class KeyframeInterpolationPipeline:
 
 @torch.inference_mode()
 def main() -> None:
-    logging.getLogger().setLevel(logging.INFO)
+    logging.basicConfig(level=logging.INFO)
     checkpoint_path = detect_checkpoint_path()
     params = detect_params(checkpoint_path)
     parser = default_2_stage_arg_parser(params=params)
@@ -245,7 +246,7 @@ def main() -> None:
         gemma_root=args.gemma_root,
         loras=tuple(args.lora) if args.lora else (),
         quantization=args.quantization,
-        torch_compile=args.compile,
+        compilation_config=args.compile,
         offload_mode=args.offload_mode,
     )
     tiling_config = TilingConfig.default()

packages/ltx-pipelines/src/ltx_pipelines/lipdub.py

@@ -13,6 +13,7 @@ from ltx_core.conditioning import AudioConditionByReferenceLatent
 from ltx_core.loader import LoraPathStrengthAndSDOps
 from ltx_core.loader.registry import Registry
 from ltx_core.model.audio_vae import encode_audio as vae_encode_audio
+from ltx_core.model.transformer.compiling import CompilationConfig
 from ltx_core.model.video_vae import TilingConfig, VideoEncoder, get_video_chunks_number
 from ltx_core.quantization import QuantizationPolicy
 from ltx_core.types import Audio, AudioLatentShape, SpatioTemporalScaleFactors, VideoPixelShape
@@ -59,7 +60,7 @@ class LipDubPipeline:
         device: torch.device | None = None,
         quantization: QuantizationPolicy | None = None,
         registry: Registry | None = None,
-        torch_compile: bool = False,
+        compilation_config: CompilationConfig | None = None,
         offload_mode: OffloadMode = OffloadMode.NONE,
     ) -> None:
         self.device = device or get_device()
@@ -89,7 +90,7 @@ class LipDubPipeline:
             loras=loras,
             quantization=quantization,
             registry=registry,
-            torch_compile=torch_compile,
+            compilation_config=compilation_config,
             offload_mode=offload_mode,
         )
         self.upsampler = VideoUpsampler(
@@ -289,7 +290,7 @@ def patchify_lipdub_audio_reference_latent(
 
 @torch.inference_mode()
 def main() -> None:
-    logging.getLogger().setLevel(logging.INFO)
+    logging.basicConfig(level=logging.INFO)
     checkpoint_path = detect_checkpoint_path(distilled=True)
     params = detect_params(checkpoint_path)
     parser = lipdub_arg_parser(params=params)
@@ -304,7 +305,7 @@ def main() -> None:
         gemma_root=args.gemma_root,
         ic_lora=args.lora[0],
         quantization=args.quantization,
-        torch_compile=args.compile,
+        compilation_config=args.compile,
         offload_mode=args.offload_mode,
     )
     tiling_config = TilingConfig.default()

packages/ltx-pipelines/src/ltx_pipelines/retake.py

@@ -11,6 +11,7 @@ from ltx_core.components.schedulers import LTX2Scheduler
 from ltx_core.conditioning.types.noise_mask_cond import TemporalRegionMask
 from ltx_core.loader import LoraPathStrengthAndSDOps
 from ltx_core.loader.registry import Registry
+from ltx_core.model.transformer.compiling import CompilationConfig
 from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
 from ltx_core.quantization import QuantizationPolicy
 from ltx_core.types import (
@@ -73,7 +74,7 @@ class RetakePipeline:
         quantization: QuantizationPolicy | None = None,
         registry: Registry | None = None,
         distilled: bool = True,
-        torch_compile: bool = False,
+        compilation_config: CompilationConfig | None = None,
         offload_mode: OffloadMode = OffloadMode.NONE,
     ):
         self.device = device or get_device()
@@ -108,7 +109,7 @@ class RetakePipeline:
             loras=tuple(loras),
             quantization=quantization,
             registry=registry,
-            torch_compile=torch_compile,
+            compilation_config=compilation_config,
             offload_mode=offload_mode,
         )
         self.video_decoder = VideoDecoder(
@@ -283,7 +284,7 @@ class RetakePipeline:
 @torch.inference_mode()
 def main() -> None:
     """CLI entry point for retake (regenerate a time region)."""
-    logging.getLogger().setLevel(logging.INFO)
+    logging.basicConfig(level=logging.INFO)
     parser = video_editing_arg_parser(distilled=True)
     parser.description = "Retake: regenerate a time region of a video with LTX-2."
     args = parser.parse_args()
@@ -308,7 +309,7 @@ def main() -> None:
         loras=tuple(args.lora) if args.lora else (),
         quantization=args.quantization,
         distilled=True,
-        torch_compile=args.compile,
+        compilation_config=args.compile,
         offload_mode=args.offload_mode,
     )
     params = detect_params(args.distilled_checkpoint_path)

packages/ltx-pipelines/src/ltx_pipelines/ti2vid_one_stage.py

@@ -12,6 +12,7 @@ from ltx_core.components.noisers import GaussianNoiser
 from ltx_core.components.schedulers import LTX2Scheduler
 from ltx_core.loader import LoraPathStrengthAndSDOps
 from ltx_core.loader.registry import Registry
+from ltx_core.model.transformer.compiling import CompilationConfig
 from ltx_core.model.video_vae.tiling import TilingConfig
 from ltx_core.quantization import QuantizationPolicy
 from ltx_core.types import Audio
@@ -55,7 +56,7 @@ class TI2VidOneStagePipeline:
         device: torch.device | None = None,
         quantization: QuantizationPolicy | None = None,
         registry: Registry | None = None,
-        torch_compile: bool = False,
+        compilation_config: CompilationConfig | None = None,
         offload_mode: OffloadMode = OffloadMode.NONE,
     ):
         self.dtype = torch.bfloat16
@@ -82,7 +83,7 @@ class TI2VidOneStagePipeline:
             loras=tuple(loras),
             quantization=quantization,
             registry=registry,
-            torch_compile=torch_compile,
+            compilation_config=compilation_config,
             offload_mode=offload_mode,
         )
         self.video_decoder = VideoDecoder(
@@ -185,7 +186,7 @@ class TI2VidOneStagePipeline:
 
 @torch.inference_mode()
 def main() -> None:
-    logging.getLogger().setLevel(logging.INFO)
+    logging.basicConfig(level=logging.INFO)
     checkpoint_path = detect_checkpoint_path()
     params = detect_params(checkpoint_path)
     parser = default_1_stage_arg_parser(params=params)
@@ -195,7 +196,7 @@ def main() -> None:
         gemma_root=args.gemma_root,
         loras=tuple(args.lora) if args.lora else (),
         quantization=args.quantization,
-        torch_compile=args.compile,
+        compilation_config=args.compile,
         offload_mode=args.offload_mode,
     )
     video, audio = pipeline(

packages/ltx-pipelines/src/ltx_pipelines/ti2vid_two_stages.py

@@ -12,6 +12,7 @@ from ltx_core.components.noisers import GaussianNoiser
 from ltx_core.components.schedulers import LTX2Scheduler
 from ltx_core.loader import LoraPathStrengthAndSDOps
 from ltx_core.loader.registry import Registry
+from ltx_core.model.transformer.compiling import CompilationConfig
 from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
 from ltx_core.quantization import QuantizationPolicy
 from ltx_core.types import Audio, VideoPixelShape
@@ -61,7 +62,7 @@ class TI2VidTwoStagesPipeline:
         device: torch.device | None = None,
         quantization: QuantizationPolicy | None = None,
         registry: Registry | None = None,
-        torch_compile: bool = False,
+        compilation_config: CompilationConfig | None = None,
         offload_mode: OffloadMode = OffloadMode.NONE,
     ):
         self.device = device or get_device()
@@ -85,7 +86,7 @@ class TI2VidTwoStagesPipeline:
             loras=tuple(loras),
             quantization=quantization,
             registry=registry,
-            torch_compile=torch_compile,
+            compilation_config=compilation_config,
             offload_mode=offload_mode,
         )
         self.stage_2 = DiffusionStage(
@@ -95,7 +96,7 @@ class TI2VidTwoStagesPipeline:
             loras=(*tuple(loras), *distilled_lora),
             quantization=quantization,
             registry=registry,
-            torch_compile=torch_compile,
+            compilation_config=compilation_config,
             offload_mode=offload_mode,
         )
 
@@ -223,7 +224,7 @@ class TI2VidTwoStagesPipeline:
 
 @torch.inference_mode()
 def main() -> None:
-    logging.getLogger().setLevel(logging.INFO)
+    logging.basicConfig(level=logging.INFO)
     checkpoint_path = detect_checkpoint_path()
     params = detect_params(checkpoint_path)
     parser = default_2_stage_arg_parser(params=params)
@@ -235,7 +236,7 @@ def main() -> None:
         gemma_root=args.gemma_root,
         loras=tuple(args.lora) if args.lora else (),
         quantization=args.quantization,
-        torch_compile=args.compile,
+        compilation_config=args.compile,
         offload_mode=args.offload_mode,
     )
     tiling_config = TilingConfig.default()

packages/ltx-pipelines/src/ltx_pipelines/ti2vid_two_stages_hq.py

@@ -9,6 +9,7 @@ from ltx_core.components.noisers import GaussianNoiser
 from ltx_core.components.schedulers import LTX2Scheduler
 from ltx_core.loader import LoraPathStrengthAndSDOps
 from ltx_core.loader.registry import Registry
+from ltx_core.model.transformer.compiling import CompilationConfig
 from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
 from ltx_core.quantization import QuantizationPolicy
 from ltx_core.types import Audio, VideoLatentShape, VideoPixelShape
@@ -60,7 +61,7 @@ class TI2VidTwoStagesHQPipeline:
         device: torch.device | None = None,
         quantization: QuantizationPolicy | None = None,
         registry: Registry | None = None,
-        torch_compile: bool = False,
+        compilation_config: CompilationConfig | None = None,
         offload_mode: OffloadMode = OffloadMode.NONE,
     ):
         self.device = device or get_device()
@@ -95,7 +96,7 @@ class TI2VidTwoStagesHQPipeline:
             loras=(*loras, distilled_lora_stage_1),
             quantization=quantization,
             registry=registry,
-            torch_compile=torch_compile,
+            compilation_config=compilation_config,
             offload_mode=offload_mode,
         )
         self.stage_2 = DiffusionStage(
@@ -105,7 +106,7 @@ class TI2VidTwoStagesHQPipeline:
             loras=(*loras, distilled_lora_stage_2),
             quantization=quantization,
             registry=registry,
-            torch_compile=torch_compile,
+            compilation_config=compilation_config,
             offload_mode=offload_mode,
         )
 
@@ -242,7 +243,7 @@ class TI2VidTwoStagesHQPipeline:
 
 @torch.inference_mode()
 def main() -> None:
-    logging.getLogger().setLevel(logging.INFO)
+    logging.basicConfig(level=logging.INFO)
     parser = hq_2_stage_arg_parser(params=LTX_2_3_HQ_PARAMS)
     args = parser.parse_args()
     pipeline = TI2VidTwoStagesHQPipeline(
@@ -254,7 +255,7 @@ def main() -> None:
         gemma_root=args.gemma_root,
         loras=tuple(args.lora) if args.lora else (),
         quantization=args.quantization,
-        torch_compile=args.compile,
+        compilation_config=args.compile,
         offload_mode=args.offload_mode,
     )
     tiling_config = TilingConfig.default()

packages/ltx-pipelines/src/ltx_pipelines/utils/args.py

@@ -1,9 +1,11 @@
 import argparse
+import json
 from collections.abc import Sequence
 from pathlib import Path
-from typing import NamedTuple
+from typing import Any, NamedTuple
 
 from ltx_core.loader import LTXV_LORA_COMFY_RENAMING_MAP, LoraPathStrengthAndSDOps
+from ltx_core.model.transformer.compiling import CompilationConfig
 from ltx_core.quantization import QuantizationPolicy
 from ltx_pipelines.utils.constants import (
     DEFAULT_IMAGE_CRF,
@@ -13,6 +15,7 @@ from ltx_pipelines.utils.constants import (
     LTX_2_3_PARAMS,
     PipelineParams,
 )
+from ltx_pipelines.utils.quantization_factory import QuantizationKind
 from ltx_pipelines.utils.types import OffloadMode
 
 
@@ -32,7 +35,7 @@ class VideoConditioningAction(argparse.Action):
         option_string: str | None = None,  # noqa: ARG002
     ) -> None:
         path, strength_str = values
-        resolved_path = resolve_path(path)
+        resolved_path = resolve_existing_path(path)
         strength = float(strength_str)
         current = getattr(namespace, self.dest) or []
         current.append((resolved_path, strength))
@@ -58,7 +61,7 @@ class VideoMaskConditioningAction(argparse.Action):
             msg = f"{option_string} requires exactly 2 arguments (MASK_PATH STRENGTH), got {len(values)}"
             raise argparse.ArgumentError(self, msg)
 
-        mask_path = resolve_path(values[0])
+        mask_path = resolve_existing_path(values[0])
         strength = float(values[1])
         setattr(namespace, self.dest, (mask_path, strength))
 
@@ -76,7 +79,7 @@ class ImageAction(argparse.Action):
             raise argparse.ArgumentError(self, msg)
 
         conditioning = ImageConditioningInput(
-            path=resolve_path(values[0]),
+            path=resolve_existing_path(values[0]),
             frame_idx=int(values[1]),
             strength=float(values[2]),
             crf=int(values[3]) if len(values) > 3 else DEFAULT_IMAGE_CRF,
@@ -101,7 +104,7 @@ class LoraAction(argparse.Action):
         path = values[0]
         strength_str = values[1] if len(values) > 1 else str(DEFAULT_LORA_STRENGTH)
 
-        resolved_path = resolve_path(path)
+        resolved_path = resolve_existing_path(path)
         strength = float(strength_str)
 
         current = getattr(namespace, self.dest) or []
@@ -109,11 +112,118 @@ class LoraAction(argparse.Action):
         setattr(namespace, self.dest, current)
 
 
+class CompileAction(argparse.Action):
+    """Parse ``--compile [KEY=VALUE ...]`` into a :class:`CompilationConfig`.
+    The flag is absent           -> ``args.compile`` stays at its default (``None``).
+    The flag is passed alone     -> ``CompilationConfig()`` (vanilla torch defaults).
+    The flag is passed with args -> ``CompilationConfig`` with the given fields overridden.
+    Errors (unknown key, malformed value, duplicate key, empty value) raise
+    :class:`argparse.ArgumentError` so argparse formats them as friendly CLI
+    messages rather than uncaught tracebacks.
+    """
+
+    _ALLOWED_KEYS = frozenset({"mode", "backend", "fullgraph", "dynamic", "inductor_config", "dynamo_config"})
+
+    def __call__(
+        self,
+        parser: argparse.ArgumentParser,  # noqa: ARG002
+        namespace: argparse.Namespace,
+        values: list[str],
+        option_string: str | None = None,  # noqa: ARG002
+    ) -> None:
+        overrides: dict[str, object] = {}
+        for item in values:
+            if "=" not in item:
+                raise argparse.ArgumentError(self, f"expects KEY=VALUE pairs, got: {item!r}")
+            key, _, raw = item.partition("=")
+            key = key.strip()
+            if key not in self._ALLOWED_KEYS:
+                raise argparse.ArgumentError(
+                    self,
+                    f"{key!r} is not a CompilationConfig field; valid keys: {sorted(self._ALLOWED_KEYS)}",
+                )
+            if key in overrides:
+                raise argparse.ArgumentError(self, f"{key} given more than once")
+            if key == "mode":
+                overrides[key] = self._parse_mode(raw)
+            elif key == "backend":
+                overrides[key] = self._parse_non_empty(key, raw)
+            elif key == "fullgraph":
+                overrides[key] = self._parse_bool(key, raw)
+            elif key == "dynamic":
+                overrides[key] = self._parse_dynamic(raw)
+            elif key in ("inductor_config", "dynamo_config"):
+                overrides[key] = self._parse_json_dict(key, raw)
+        setattr(namespace, self.dest, CompilationConfig(**overrides))
+
+    def _parse_mode(self, raw: str) -> str | None:
+        stripped = raw.strip()
+        if not stripped:
+            raise argparse.ArgumentError(self, "mode=... value cannot be empty (use mode=none to clear)")
+        if stripped.lower() == "none":
+            return None
+        return stripped
+
+    def _parse_non_empty(self, key: str, raw: str) -> str:
+        stripped = raw.strip()
+        if not stripped:
+            raise argparse.ArgumentError(self, f"{key}=... value cannot be empty")
+        return stripped
+
+    def _parse_bool(self, key: str, raw: str) -> bool:
+        normalized = raw.strip().lower()
+        if normalized in ("true", "1"):
+            return True
+        if normalized in ("false", "0"):
+            return False
+        raise argparse.ArgumentError(self, f"{key}=... must be true or false; got {raw!r}")
+
+    def _parse_dynamic(self, raw: str) -> bool | None:
+        normalized = raw.strip().lower()
+        if normalized in ("auto", "none"):
+            return None
+        if normalized in ("true", "1"):
+            return True
+        if normalized in ("false", "0"):
+            return False
+        raise argparse.ArgumentError(self, f"dynamic=... must be auto/true/false; got {raw!r}")
+
+    def _parse_json_dict(self, key: str, raw: str) -> dict[str, Any]:
+        # Inline JSON object starts with '{'; otherwise treat the value as a path to a JSON file.
+        stripped = raw.strip()
+        if not stripped:
+            raise argparse.ArgumentError(self, f"{key}=... value cannot be empty")
+        if stripped.startswith("{"):
+            source = stripped
+        else:
+            path = Path(stripped).expanduser()
+            if not path.is_file():
+                raise argparse.ArgumentError(
+                    self, f"{key}=... must be a JSON object or a path to a JSON file; got {raw!r}"
+                )
+            source = path.read_text()
+        try:
+            value = json.loads(source)
+        except json.JSONDecodeError as e:
+            raise argparse.ArgumentError(self, f"{key}=... must be a JSON object; got {raw!r} ({e.msg})") from None
+        if not isinstance(value, dict):
+            raise argparse.ArgumentError(self, f"{key}=... must decode to a JSON object; got {type(value).__name__}")
+        return value
+
+
 def resolve_path(path: str) -> str:
     return str(Path(path).expanduser().resolve().as_posix())
 
 
-QUANTIZATION_POLICIES = ("fp8-cast", "fp8-scaled-mm")
+def resolve_existing_path(path: str) -> str:
+    """Resolve *path* and verify it exists."""
+    resolved = resolve_path(path)
+    if not Path(resolved).exists():
+        raise argparse.ArgumentError(None, f"Path not found: {resolved}")
+    return resolved
+
+
+QUANTIZATION_POLICIES = tuple(k.value for k in QuantizationKind)
 
 
 def _resolve_quantization(namespace: argparse.Namespace) -> None:
@@ -123,16 +233,14 @@ def _resolve_quantization(namespace: argparse.Namespace) -> None:
     name = getattr(namespace, "quantization", None)
     if name is None or isinstance(name, QuantizationPolicy):
         return
-    if name == "fp8-cast":
+    try:
-        namespace.quantization = QuantizationPolicy.fp8_cast()
+        kind = QuantizationKind(name)
+    except ValueError:
         return
-    if name == "fp8-scaled-mm":
     ckpt = getattr(namespace, "checkpoint_path", None) or getattr(namespace, "distilled_checkpoint_path", None)
     if ckpt is None:
-            raise SystemExit(
+        raise SystemExit(f"--quantization {kind.value} requires --checkpoint-path (or --distilled-checkpoint-path).")
-                "--quantization fp8-scaled-mm requires --checkpoint-path (or --distilled-checkpoint-path)."
+    namespace.quantization = kind.to_policy(checkpoint_path=ckpt)
-            )
-        namespace.quantization = QuantizationPolicy.fp8_scaled_mm(ckpt)
 
 
 class _PipelineArgumentParser(argparse.ArgumentParser):
@@ -150,7 +258,7 @@ def detect_checkpoint_path(distilled: bool = False) -> str:
     """Pre-parse argv to extract the checkpoint path before building the full parser."""
     pre = argparse.ArgumentParser(add_help=False)
     flag = "--distilled-checkpoint-path" if distilled else "--checkpoint-path"
-    pre.add_argument(flag, type=resolve_path, required=True)
+    pre.add_argument(flag, type=resolve_existing_path, required=True)
     known, _ = pre.parse_known_args()
     return known.distilled_checkpoint_path if distilled else known.checkpoint_path
 
@@ -163,14 +271,14 @@ def basic_arg_parser(
     if distilled:
         parser.add_argument(
             "--distilled-checkpoint-path",
-            type=resolve_path,
+            type=resolve_existing_path,
             required=True,
             help="Path to LTX-2 distilled model checkpoint (.safetensors file).",
         )
     else:
         parser.add_argument(
             "--checkpoint-path",
-            type=resolve_path,
+            type=resolve_existing_path,
             required=True,
             help="Path to LTX-2 model checkpoint (.safetensors file).",
         )
@@ -185,7 +293,7 @@ def basic_arg_parser(
         )
     parser.add_argument(
         "--gemma-root",
-        type=resolve_path,
+        type=resolve_existing_path,
         required=True,
         help="Path to the root directory containing the Gemma text encoder model files.",
     )
@@ -276,8 +384,20 @@ def basic_arg_parser(
     )
     parser.add_argument(
         "--compile",
-        action="store_true",
+        nargs="*",
-        help="Enable torch.compile for transformer blocks to optimize performance.",
+        action=CompileAction,
+        default=None,
+        metavar="KEY=VALUE",
+        help=(
+            "Enable torch.compile for transformer blocks. Pass alone for defaults, "
+            "or with KEY=VALUE overrides for any CompilationConfig field. "
+            "Keys: mode, backend, fullgraph, dynamic, inductor_config, dynamo_config. "
+            "inductor_config/dynamo_config take JSON objects (inline or a path to a .json file) "
+            "that fully replace the defaults. "
+            "Examples: --compile  or  --compile mode=reduce-overhead  or  "
+            "--compile mode=reduce-overhead fullgraph=true backend=eager  or  "
+            "--compile inductor_config='{\"max_autotune\": true}'"
+        ),
     )
     return parser
 
@@ -340,7 +460,7 @@ def video_editing_arg_parser(
     (no height/width/num-frames; resolution comes from input video). Default is distilled checkpoint only.
     """
     parser = basic_arg_parser(distilled=distilled)
-    parser.add_argument("--video-path", type=resolve_path, required=True, help="Path to the source video.")
+    parser.add_argument("--video-path", type=resolve_existing_path, required=True, help="Path to the source video.")
     parser.add_argument("--start-time", type=float, required=True, help="Start time of the region to regenerate (s).")
     parser.add_argument("--end-time", type=float, required=True, help="End time of the region to regenerate (s).")
     return parser
@@ -462,7 +582,7 @@ def default_1_stage_arg_parser(params: PipelineParams = LTX_2_3_PARAMS) -> argpa
         default=video_guider.skip_step,
         help=(
             "Video skip step N controls periodic skipping during the video diffusion process: "
-            "only steps where step_index % (N + 1) == 0 are processed, all others are skipped "
+            "only steps where step_index %% (N + 1) == 0 are processed, all others are skipped "
             f"(e.g., 0 = no skipping; 1 = skip every other step; 2 = skip 2 of every 3 steps; "
             f"default: {video_guider.skip_step})."
         ),
@@ -522,7 +642,7 @@ def default_1_stage_arg_parser(params: PipelineParams = LTX_2_3_PARAMS) -> argpa
         default=audio_guider.skip_step,
         help=(
             "Audio skip step N controls periodic skipping during the audio diffusion process: "
-            "only steps where step_index % (N + 1) == 0 are processed, all others are skipped "
+            "only steps where step_index %% (N + 1) == 0 are processed, all others are skipped "
             f"(e.g., 0 = no skipping; 1 = skip every other step; 2 = skip 2 of every 3 steps; "
             f"default: {audio_guider.skip_step})."
         ),
@@ -562,7 +682,7 @@ def default_2_stage_arg_parser(params: PipelineParams = LTX_2_3_PARAMS) -> argpa
     )
     parser.add_argument(
         "--spatial-upsampler-path",
-        type=resolve_path,
+        type=resolve_existing_path,
         required=True,
         help=(
             "Path to the spatial upsampler model used to increase the resolution "
@@ -605,7 +725,7 @@ def default_2_stage_distilled_arg_parser(params: PipelineParams = LTX_2_3_PARAMS
             )
     parser.add_argument(
         "--spatial-upsampler-path",
-        type=resolve_path,
+        type=resolve_existing_path,
         required=True,
         help=(
             "Path to the spatial upsampler model used to increase the resolution "

packages/ltx-pipelines/src/ltx_pipelines/utils/blocks.py

@@ -6,6 +6,8 @@ removes the need for :class:`ModelLedger`.
 
 from __future__ import annotations
 
+import copy
+import dataclasses
 import logging
 from collections.abc import Iterator
 from contextlib import AbstractContextManager, contextmanager
@@ -21,6 +23,8 @@ from ltx_core.components.noisers import Noiser
 from ltx_core.components.patchifiers import AudioPatchifier, VideoLatentPatchifier
 from ltx_core.components.protocols import DiffusionStepProtocol
 from ltx_core.loader import SDOps
+from ltx_core.loader.attention_ops import set_attention_module_op
+from ltx_core.loader.fuse_loras import bf16_fuse_rule
 from ltx_core.loader.module_ops import ModuleOps
 from ltx_core.loader.primitives import BuilderProtocol, LoraPathStrengthAndSDOps, ModelBuilderProtocol
 from ltx_core.loader.registry import DummyRegistry, Registry
@@ -42,7 +46,15 @@ from ltx_core.model.transformer import (
     LTXModelConfigurator,
     X0Model,
 )
-from ltx_core.model.transformer.compiling import COMPILE_TRANSFORMER, modify_sd_ops_for_compilation
+from ltx_core.model.transformer.attention import (
+    AttentionCallable,
+    AttentionFunction,
+)
+from ltx_core.model.transformer.compiling import (
+    CompilationConfig,
+    build_compile_transformer_op,
+    modify_sd_ops_for_compilation,
+)
 from ltx_core.model.upsampler import LatentUpsamplerConfigurator, upsample_video
 from ltx_core.model.video_vae import (
     MEMORY_EFFICIENT_DECODE,
@@ -53,7 +65,7 @@ from ltx_core.model.video_vae import (
     VideoEncoder,
     VideoEncoderConfigurator,
 )
-from ltx_core.quantization import QuantizationPolicy
+from ltx_core.quantization import QuantizationPolicy, fp8_cast_fuse_rule
 from ltx_core.text_encoders.gemma import (
     EMBEDDINGS_PROCESSOR_KEY_OPS,
     GEMMA_LLM_KEY_OPS,
@@ -101,6 +113,28 @@ def _chain_quantization(
     return chained_sd_ops, (*module_ops, *quantization.module_ops)
 
 
+def _apply_compile_ops(
+    sd_ops: SDOps,
+    module_ops: tuple[ModuleOps, ...],
+    loras: tuple[LoraPathStrengthAndSDOps, ...],
+    number_of_layers: int,
+    compilation_config: CompilationConfig,
+) -> tuple[SDOps, tuple[ModuleOps, ...], tuple[LoraPathStrengthAndSDOps, ...]]:
+    """Rewrite sd_ops/module_ops/LoRAs for compiled blocks (params land under ``_orig_mod``)."""
+    sd_ops = modify_sd_ops_for_compilation(sd_ops, number_of_layers)
+    compile_op = build_compile_transformer_op(compilation_config)
+    module_ops = (*module_ops, compile_op)
+    loras = tuple(
+        LoraPathStrengthAndSDOps(
+            lora.path,
+            lora.strength,
+            modify_sd_ops_for_compilation(lora.sd_ops, number_of_layers),
+        )
+        for lora in loras
+    )
+    return sd_ops, module_ops, loras
+
+
 @contextmanager
 def _streaming_model(
     builder: StreamingModelBuilder,
@@ -170,79 +204,110 @@ class DiffusionStage:
         loras: tuple[LoraPathStrengthAndSDOps, ...] = (),
         quantization: QuantizationPolicy | None = None,
         registry: Registry | None = None,
-        torch_compile: bool = False,
+        compilation_config: CompilationConfig | None = None,
         offload_mode: OffloadMode = OffloadMode.NONE,
         transformer_builder: ModelBuilderProtocol[LTXModel] | DelegatingBuilder[LTXModel] | None = None,
     ) -> None:
+        self._checkpoint_path = checkpoint_path
         self._dtype = dtype
         self._device = device
         self._quantization = quantization
-        self._torch_compile = torch_compile
+        self._compilation_config = compilation_config
         self._offload_mode = offload_mode
+        configurator = (
+            quantization.model_configurator
+            if quantization is not None and quantization.model_configurator is not None
+            else LTXModelConfigurator
+        )
         if transformer_builder is not None:
             self._transformer_builder = transformer_builder
         else:
             self._transformer_builder = Builder(
                 model_path=checkpoint_path,
-                model_class_configurator=LTXModelConfigurator,
+                model_class_configurator=configurator,
                 model_sd_ops=LTXV_MODEL_COMFY_RENAMING_MAP,
                 loras=tuple(loras),
                 registry=registry or DummyRegistry(),
             )
 
         if offload_mode != OffloadMode.NONE:
-            if torch_compile:
+            if compilation_config is not None:
                 raise ValueError("torch.compile is not supported with layer streaming")
+            # WeightsProvider currently only supports plain bf16 + fp8_cast LoRA fusion
+            # (no companion-key emission). Quantization policies that emit
+            # companion keys (e.g. ``.weight_scale``) cannot be streamed yet.
+            if quantization is not None and quantization.fuse_rule is not fp8_cast_fuse_rule:
+                raise ValueError(
+                    "Block streaming is not supported with this quantization policy "
+                    "(only bf16 and fp8_cast are currently supported)."
+                )
             streaming_sd_ops: SDOps = LTXV_MODEL_COMFY_RENAMING_MAP
             streaming_module_ops: tuple[ModuleOps, ...] = ()
             if quantization is not None:
-                if quantization.kind != QuantizationPolicy.Kind.FP8_CAST:
-                    raise ValueError(
-                        f"Layer streaming supports only QuantizationPolicy.fp8_cast(); "
-                        f"got kind={quantization.kind!r} which produces heterogeneous block layouts."
-                    )
                 streaming_sd_ops, streaming_module_ops = _chain_quantization(
                     streaming_sd_ops, streaming_module_ops, quantization
                 )
             self._streaming_builder = StreamingModelBuilder(
-                model_class_configurator=LTXModelConfigurator,
+                model_class_configurator=configurator,
                 model_path=checkpoint_path,
                 model_sd_ops=streaming_sd_ops,
                 module_ops=streaming_module_ops,
                 loras=tuple(loras),
                 registry=registry or DummyRegistry(),
-                blocks_attr="velocity_model.transformer_blocks",
+                fuse_rule=quantization.fuse_rule if quantization is not None else bf16_fuse_rule,
+                blocks_attr="transformer_blocks",
                 blocks_prefix="transformer_blocks",
-                state_dict_prefix="velocity_model.",
-                model_wrapper=lambda m: X0Model(m).eval(),
             )
 
+    def with_attention(self, attention: AttentionFunction | AttentionCallable | None) -> "DiffusionStage":
+        """Return a new ``DiffusionStage`` that pins the transformer build to ``attention``.
+        Functional: never mutates ``self``. The returned stage shares all other
+        configuration with the original; only the underlying builders' ``module_ops``
+        gain a ``set_attention_module_op(attention)`` entry so subsequent transformer
+        builds use that kernel. ``attention=None`` is a no-op (returns ``self``).
+        """
+        if attention is None:
+            return self
+        op = set_attention_module_op(attention)
+        new = copy.copy(self)
+        new._transformer_builder = self._transformer_builder.with_module_ops(
+            (*self._transformer_builder.module_ops, op),
+        )
+        if self._offload_mode != OffloadMode.NONE:
+            new._streaming_builder = dataclasses.replace(
+                self._streaming_builder,
+                module_ops=(*self._streaming_builder.module_ops, op),
+            )
+        return new
+
     def _build_transformer(self, *, device: torch.device | None = None, **kwargs: object) -> X0Model:
         target = device or self._device
         sd_ops = self._transformer_builder.model_sd_ops
         module_ops = self._transformer_builder.module_ops
         loras = self._transformer_builder.loras
-        if self._torch_compile:
+        if self._compilation_config is not None:
-            module_ops = (*module_ops, COMPILE_TRANSFORMER)
             number_of_layers = self._transformer_builder.model_config()["transformer"]["num_layers"]
-            sd_ops = modify_sd_ops_for_compilation(sd_ops, number_of_layers)
+            sd_ops, module_ops, loras = _apply_compile_ops(
-            loras = tuple(
+                sd_ops, module_ops, loras, number_of_layers, self._compilation_config
-                LoraPathStrengthAndSDOps(
-                    lora.path,
-                    lora.strength,
-                    modify_sd_ops_for_compilation(lora.sd_ops, number_of_layers),
-                )
-                for lora in loras
             )
         if self._quantization is not None:
             sd_ops, module_ops = _chain_quantization(sd_ops, module_ops, self._quantization)
 
         builder = self._transformer_builder.with_module_ops(module_ops).with_sd_ops(sd_ops).with_loras(loras)
+        if self._quantization is not None:
+            builder = builder.with_fuse_rule(self._quantization.fuse_rule)
         return X0Model(builder.build(device=target, **kwargs)).to(target).eval()
 
+    @contextmanager
+    def _streaming_transformer_ctx(self) -> Iterator[X0Model]:
+        with _streaming_model(
+            self._streaming_builder, self._offload_mode, self._device, self._dtype
+        ) as streaming_wrapper:
+            yield X0Model(streaming_wrapper).eval()
+
     def _transformer_ctx(self, **kwargs: object) -> AbstractContextManager:
         if self._offload_mode != OffloadMode.NONE:
-            return _streaming_model(self._streaming_builder, self._offload_mode, self._device, self._dtype)
+            return self._streaming_transformer_ctx()
         return gpu_model(self._build_transformer(**kwargs))
 
     def model_context(self, **kwargs: object) -> AbstractContextManager:
@@ -348,7 +413,17 @@ class DiffusionStage:
             v_shape = VideoLatentShape.from_pixel_shape(pixel_shape)
             video_tools = VideoLatentTools(VideoLatentPatchifier(patch_size=1), v_shape, fps)
 
+        mode = "streaming" if self._offload_mode != OffloadMode.NONE else "standard"
+        logger.info("Building transformer (%s) from %s", mode, self._checkpoint_path)
         with self._transformer_ctx(video_tools=video_tools) as transformer:
+            logger.info(
+                "Running denoising loop (%d steps, %dx%d %d frames @ %.1f fps)",
+                len(sigmas) - 1,
+                width,
+                height,
+                frames,
+                fps,
+            )
             return self.run(
                 transformer,
                 denoiser,
@@ -387,6 +462,8 @@ class PromptEncoder:
         offload_mode: OffloadMode = OffloadMode.NONE,
         text_encoder_builder: BuilderProtocol | None = None,
     ) -> None:
+        self._gemma_root = gemma_root
+        self._checkpoint_path = checkpoint_path
         self._dtype = dtype
         self._device = device
         self._offload_mode = offload_mode
@@ -432,7 +509,7 @@ class PromptEncoder:
 
     def _build_embeddings_processor(self) -> EmbeddingsProcessor:
         """Build the embeddings processor on the target device."""
-        return self._embeddings_processor_builder.build(device=self._device, dtype=self._dtype).to(self._device).eval()
+        return self._embeddings_processor_builder.build(device=self._device, dtype=self._dtype).eval()
 
     def _text_encoder_ctx(self) -> AbstractContextManager:
         if self._offload_mode != OffloadMode.NONE:
@@ -448,6 +525,7 @@ class PromptEncoder:
         enhance_prompt_seed: int = 42,
     ) -> list[EmbeddingsProcessorOutput]:
         """Encode *prompts* through Gemma -> embeddings processor, freeing each model after use."""
+        logger.info("Building text encoder from %s", self._gemma_root)
         with self._text_encoder_ctx() as text_encoder:
             if enhance_first_prompt:
                 prompts = list(prompts)
@@ -455,9 +533,12 @@ class PromptEncoder:
                     text_encoder, prompts[0], enhance_prompt_image, seed=enhance_prompt_seed
                 )
             raw_outputs = [text_encoder.encode(p) for p in prompts]
+        logger.info("Text encoder done, building embeddings processor from %s", self._checkpoint_path)
 
         with gpu_model(self._build_embeddings_processor()) as embeddings_processor:
-            return [embeddings_processor.process_hidden_states(hs, mask) for hs, mask in raw_outputs]
+            result = [embeddings_processor.process_hidden_states(hs, mask) for hs, mask in raw_outputs]
+        logger.info("Prompt encoding complete")
+        return result
 
 
 # ---------------------------------------------------------------------------
@@ -487,7 +568,7 @@ class ImageConditioner:
         )
 
     def _build_encoder(self) -> VideoEncoder:
-        return self._encoder_builder.build(device=self._device, dtype=self._dtype).to(self._device).eval()
+        return self._encoder_builder.build(device=self._device, dtype=self._dtype).eval()
 
     def __call__(self, fn: Callable[[VideoEncoder], T]) -> T:
         """Build video encoder → call *fn(encoder)* → free encoder."""
@@ -511,6 +592,7 @@ class VideoUpsampler:
         device: torch.device,
         registry: Registry | None = None,
     ) -> None:
+        self._upsampler_path = upsampler_path
         self._dtype = dtype
         self._device = device
         self._encoder_builder = Builder(
@@ -527,13 +609,10 @@ class VideoUpsampler:
 
     def __call__(self, latent: torch.Tensor) -> torch.Tensor:
         """Upsample *latent* using video encoder + spatial upsampler, then free both."""
+        logger.info("Building video encoder + spatial upsampler from %s", self._upsampler_path)
         with (
-            gpu_model(
+            gpu_model(self._encoder_builder.build(device=self._device, dtype=self._dtype).eval()) as encoder,
-                self._encoder_builder.build(device=self._device, dtype=self._dtype).to(self._device).eval()
+            gpu_model(self._upsampler_builder.build(device=self._device, dtype=self._dtype).eval()) as upsampler,
-            ) as encoder,
-            gpu_model(
-                self._upsampler_builder.build(device=self._device, dtype=self._dtype).to(self._device).eval()
-            ) as upsampler,
         ):
             return upsample_video(latent=latent, video_encoder=encoder, upsampler=upsampler)
 
@@ -557,6 +636,7 @@ class VideoDecoder:
         memory_efficient: bool = True,
         decoder_builder: BuilderProtocol | None = None,
     ) -> None:
+        self._checkpoint_path = checkpoint_path
         self._dtype = dtype
         self._device = device
         if decoder_builder is not None:
@@ -577,7 +657,8 @@ class VideoDecoder:
         generator: torch.Generator | None = None,
     ) -> Iterator[torch.Tensor]:
         """Decode *latent* to pixel-space video chunks. Decoder freed after exhaustion."""
-        decoder = self._decoder_builder.build(device=self._device, dtype=self._dtype).to(self._device).eval()
+        logger.info("Building video decoder from %s", self._checkpoint_path)
+        decoder = self._decoder_builder.build(device=self._device, dtype=self._dtype).eval()
         return _cleanup_iter(decoder.decode_video(latent, tiling_config, generator), decoder)
 
 
@@ -596,6 +677,7 @@ class AudioDecoder:
         device: torch.device,
         registry: Registry | None = None,
     ) -> None:
+        self._checkpoint_path = checkpoint_path
         self._dtype = dtype
         self._device = device
         self._decoder_builder = Builder(
@@ -613,13 +695,10 @@ class AudioDecoder:
 
     def __call__(self, latent: torch.Tensor) -> Audio:
         """Decode audio *latent* through VAE decoder + vocoder, then free both."""
+        logger.info("Building audio decoder + vocoder from %s", self._checkpoint_path)
         with (
-            gpu_model(
+            gpu_model(self._decoder_builder.build(device=self._device, dtype=self._dtype).eval()) as decoder,
-                self._decoder_builder.build(device=self._device, dtype=self._dtype).to(self._device).eval()
+            gpu_model(self._vocoder_builder.build(device=self._device, dtype=self._dtype).eval()) as vocoder,
-            ) as decoder,
-            gpu_model(
-                self._vocoder_builder.build(device=self._device, dtype=self._dtype).to(self._device).eval()
-            ) as vocoder,
         ):
             return vae_decode_audio(latent, decoder, vocoder)
 
@@ -653,7 +732,5 @@ class AudioConditioner:
 
     def __call__(self, fn: Callable[[torch.nn.Module], T]) -> T:
         """Build audio encoder → call *fn(encoder)* → free encoder."""
-        with gpu_model(
+        with gpu_model(self._encoder_builder.build(device=self._device, dtype=self._dtype).eval()) as encoder:
-            self._encoder_builder.build(device=self._device, dtype=self._dtype).to(self._device).eval()
-        ) as encoder:
             return fn(encoder)

packages/ltx-pipelines/src/ltx_pipelines/utils/denoisers.py

@@ -138,6 +138,8 @@ def _guided_denoise(  # noqa: PLR0913,PLR0915
     ptb_configs = [ptb for _, _, _, ptb in passes]
     n = len(passes)
 
+    orig_b = (video_state or audio_state).latent.shape[0]
+
     def _batched_sigma(state: LatentState) -> torch.Tensor:
         """Expand scalar sigma to (n * B,) matching the repeated state."""
         return sigma.expand(state.latent.shape[0] * n)
@@ -162,8 +164,16 @@ def _guided_denoise(  # noqa: PLR0913,PLR0915
             enabled=not a_skip,
         )
 
+    # Replicate each pass's PerturbationConfig to all `orig_b` samples it
+    # carries, so `BatchedPerturbationConfig.mask_like` returns a per-sample
+    # mask (length n*orig_b) instead of a per-pass mask (length n). Without
+    # this expansion the mask is broadcast against a (n*orig_b, T, D) tensor
+    # and the multiplication fails with a batch-dim mismatch whenever
+    # `orig_b > 1` (e.g. multi-prompt benchmark panels).
+    batched_ptb_configs = [ptb for ptb in ptb_configs for _ in range(orig_b)]
+
     all_v, all_a = transformer(
-        video=batched_video, audio=batched_audio, perturbations=BatchedPerturbationConfig(ptb_configs)
+        video=batched_video, audio=batched_audio, perturbations=BatchedPerturbationConfig(batched_ptb_configs)
     )
 
     # Split results back and combine via guiders.

packages/ltx-pipelines/src/ltx_pipelines/utils/quantization_factory.py

@@ -0,0 +1,36 @@
+"""User-facing quantization-policy dispatch.
+``ltx-core`` exposes one ``build_policy`` factory per backend. This module
+provides the user-facing string-keyed dispatch used by CLI args and pipeline
+defaults — keeping the enum out of ``ltx-core`` so adding/removing backends is
+a single-file change here.
+"""
+
+from enum import Enum
+
+from typing_extensions import assert_never
+
+from ltx_core.quantization import QuantizationPolicy
+from ltx_core.quantization.fp8_cast import build_policy as _build_fp8_cast_policy
+from ltx_core.quantization.fp8_scaled_mm import build_policy as _build_fp8_scaled_mm_policy
+
+
+class QuantizationKind(str, Enum):
+    FP8_CAST = "fp8-cast"
+    FP8_SCALED_MM = "fp8-scaled-mm"
+
+    def to_policy(self, checkpoint_path: str | None = None) -> QuantizationPolicy:
+        """Build the :class:`QuantizationPolicy` for this kind.
+        ``checkpoint_path`` is required for both backends: ``FP8_SCALED_MM``
+        uses it to discover the layer set from ``.weight_scale`` tensors,
+        and ``FP8_CAST`` uses it to fold any prequant scales into the fp8
+        weight at load time.
+        """
+        if checkpoint_path is None:
+            raise ValueError(f"{self.value} quantization requires checkpoint_path.")
+        match self:
+            case QuantizationKind.FP8_CAST:
+                return _build_fp8_cast_policy(checkpoint_path)
+            case QuantizationKind.FP8_SCALED_MM:
+                return _build_fp8_scaled_mm_policy(checkpoint_path)
+            case _:
+                assert_never(self)

packages/ltx-trainer/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "ltx-trainer"
-version = "1.1.3"
+version = "1.1.4"
 description = "LTX-2 training, democratized."
 readme = "README.md"
 authors = [
@@ -48,7 +48,7 @@ build-backend = "hatchling.build"
 
 
 [tool.ruff]
-target-version = "1.1.3"
+target-version = "1.1.4"
 line-length = 120
 
 [tool.ruff.lint]

uv.lock

@@ -2063,7 +2063,7 @@ wheels = [
 
 [[package]]
 name = "ltx-core"
-version = "1.1.3"
+version = "1.1.4"
 source = { editable = "packages/ltx-core" }
 dependencies = [
     { name = "accelerate" },
@@ -2121,7 +2121,7 @@ dev = [{ name = "scikit-image", specifier = ">=0.25.2" }]
 
 [[package]]
 name = "ltx-pipelines"
-version = "1.1.3"
+version = "1.1.4"
 source = { editable = "packages/ltx-pipelines" }
 dependencies = [
     { name = "av" },
@@ -2144,7 +2144,7 @@ requires-dist = [
 
 [[package]]
 name = "ltx-trainer"
-version = "1.1.3"
+version = "1.1.4"
 source = { editable = "packages/ltx-trainer" }
 dependencies = [
     { name = "accelerate" },