Automated PR - 2026-06-17

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

@@ -5,8 +5,9 @@ CPU-to-GPU copies, caching, and stream synchronization.  Two weight
 source strategies are available:
 - **RAM streaming** (default): all blocks pre-loaded into pinned CPU
   buffers with LoRA fusion at build time.  Fast, higher CPU memory.
-- **Disk streaming** (``cpu_slots < num_blocks``): blocks read from
-  disk on demand with FIFO eviction.  Slower, lower CPU memory.
+- **Disk streaming** (``cpu_slots < blocks_number``): blocks are read from
+  disk on demand by a :class:`DiskWeightSource`, on a background worker
+  thread.  Slower, lower CPU memory.
 """
 
 from ltx_core.block_streaming.builder import DISK_CPU_SLOTS, StreamingModelBuilder

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

@@ -0,0 +1,83 @@
+"""BlockFetcher: async disk reads on a worker thread."""
+
+from __future__ import annotations
+
+import logging
+import queue
+import threading
+from dataclasses import dataclass
+
+import torch
+
+from ltx_core.block_streaming.disk import DiskBlockReader
+
+logger = logging.getLogger(__name__)
+
+Buffer = dict[str, torch.Tensor]
+
+
+@dataclass(slots=True)
+class FetchHandle:
+    """Caller-facing handle for an outstanding read returned by :meth:`BlockFetcher.submit`.
+    Carries only what the caller needs: a completion event the worker sets and the
+    read's error. The worker updates these once the read finishes.
+    """
+
+    _done: threading.Event
+    _error: BaseException | None = None
+
+    def wait(self) -> BaseException | None:
+        """Block until the read finishes; return its error, or ``None`` on success."""
+        self._done.wait()
+        return self._error
+
+
+@dataclass(slots=True)
+class _ReadRequest:
+    """One outstanding read, internal to :class:`BlockFetcher`.
+    The fetcher's worker reads block ``idx`` into the caller-carved ``buffer`` and
+    updates ``handle`` (its error, then its event) once the read has finished.
+    """
+
+    idx: int
+    buffer: Buffer
+    handle: FetchHandle
+
+
+class BlockFetcher:
+    """Fills caller-supplied buffers on a worker thread."""
+
+    def __init__(self, reader: DiskBlockReader) -> None:
+        self._reader = reader
+        self._request_queue: queue.SimpleQueue[_ReadRequest | None] = queue.SimpleQueue()
+        self._worker = threading.Thread(target=self._run, name="BlockFetcher-IO", daemon=True)
+        self._worker.start()
+
+    def submit(self, idx: int, buffer: Buffer) -> FetchHandle:
+        """Enqueue a read of block *idx* into the caller-carved *buffer*, return its handle."""
+        handle = FetchHandle(_done=threading.Event())
+        request = _ReadRequest(idx=idx, buffer=buffer, handle=handle)
+        self._request_queue.put(request)
+        return handle
+
+    def cleanup(self) -> None:
+        """Drain pending reads, join the worker, close the reader."""
+        self._request_queue.put(None)
+        self._worker.join()
+        self._reader.cleanup()
+
+    def _run(self) -> None:
+        # Pinned buffers are allocated under the caller's inference_mode, so
+        # in-place copy_ from this thread requires inference_mode here too.
+        with torch.inference_mode():
+            while True:
+                request = self._request_queue.get()
+                if request is None:
+                    return
+
+                try:
+                    self._reader.read_into(request.buffer, request.idx)
+                except Exception as exc:
+                    logger.exception("BlockFetcher: fetch failed for item %d", request.idx)
+                    request.handle._error = exc
+                request.handle._done.set()

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

@@ -2,19 +2,28 @@
 
 from __future__ import annotations
 
+import copy
 import logging
-from dataclasses import dataclass, field, replace
-from typing import Generic
+from dataclasses import replace
+from typing import TYPE_CHECKING, Final, Generic
 
 import safetensors
 import torch
 from torch import nn
 
+from ltx_core.block_streaming import utils as bs_utils
+from ltx_core.block_streaming.block_fetcher import BlockFetcher
 from ltx_core.block_streaming.disk import DiskBlockReader, DiskTensorReader, LoraSource
-from ltx_core.block_streaming.pool import WeightPool
+from ltx_core.block_streaming.pool import BufferPool
 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.source import DiskWeightSource, PinnedBlock, PinnedWeightSource, WeightSource
+from ltx_core.block_streaming.utils import (
+    carve_buffer,
+    derive_layout,
+    layout_nbytes,
+    make_block_key,
+    resolve_attr,
+)
 from ltx_core.block_streaming.wrapper import BlockStreamingWrapper
 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
@@ -25,23 +34,29 @@ from ltx_core.loader.primitives import (
     ModelBuilderProtocol,
     StateDict,
     StateDictLoader,
+    TensorLayout,
 )
 from ltx_core.loader.registry import DummyRegistry, Registry
 from ltx_core.loader.sd_ops import SDOps
 from ltx_core.loader.sft_loader import SafetensorsModelStateDictLoader
 from ltx_core.model.model_protocol import ModelConfigurator, ModelType
 
+if TYPE_CHECKING:
+    from typing_extensions import Self
+
 logger = logging.getLogger(__name__)
 
 DISK_CPU_SLOTS = 2
 _DEFAULT_GPU_SLOTS = 2
+_PREFETCH_DEPTH = 2
 
 
-@dataclass(frozen=True)
 class StreamingModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType]):
     """Immutable builder for :class:`BlockStreamingWrapper`.
     Reads block weights from safetensors on demand.  ``cpu_slots`` and
     ``gpu_slots`` control the memory/speed trade-off (see :meth:`build`).
+    The builder is immutable (``with_*`` return modified copies) and exposes
+    its state via read-only properties backed by private attributes.
     Args:
         model_class_configurator: Creates the model from a config dict.
         model_path: One or more ``.safetensors`` checkpoint paths.
@@ -59,30 +74,99 @@ class StreamingModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType])
             (e.g. ``"transformer_blocks"``).
     """
 
-    model_class_configurator: type[ModelConfigurator[ModelType]]
-    model_path: str | tuple[str, ...]
-    model_sd_ops: SDOps | None = None
-    module_ops: tuple[ModuleOps, ...] = field(default_factory=tuple)
-    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
+    def __init__(
+        self,
+        model_class_configurator: type[ModelConfigurator[ModelType]],
+        model_path: str | tuple[str, ...],
+        model_sd_ops: SDOps | None = None,
+        module_ops: tuple[ModuleOps, ...] = (),
+        loras: tuple[LoraPathStrengthAndSDOps, ...] = (),
+        model_loader: StateDictLoader | None = None,
+        registry: Registry | None = None,
+        fuse_rule: FuseRule = bf16_fuse_rule,
+        blocks_attr: str = "",
+        blocks_prefix: str = "",
+    ) -> None:
+        # Read-only: typed with the covariant ModelType, so it must not be a mutable attribute.
+        self._model_class_configurator: Final = model_class_configurator
+        self._model_path = model_path
+        self._model_sd_ops = model_sd_ops
+        self._module_ops = module_ops
+        self._loras = loras
+        self._model_loader = model_loader if model_loader is not None else SafetensorsModelStateDictLoader()
+        self._registry = registry if registry is not None else DummyRegistry()
+        self._fuse_rule = fuse_rule
+        self._blocks_attr = blocks_attr
+        self._blocks_prefix = blocks_prefix
 
-    # Streaming-specific
-    blocks_attr: str = ""
-    blocks_prefix: str = ""
+    @property
+    def model_class_configurator(self) -> type[ModelConfigurator[ModelType]]:
+        return self._model_class_configurator
 
-    def with_sd_ops(self, sd_ops: SDOps | None) -> StreamingModelBuilder:
-        return replace(self, model_sd_ops=sd_ops)
+    @property
+    def model_path(self) -> str | tuple[str, ...]:
+        return self._model_path
 
-    def with_module_ops(self, module_ops: tuple[ModuleOps, ...]) -> StreamingModelBuilder:
-        return replace(self, module_ops=module_ops)
+    @property
+    def model_sd_ops(self) -> SDOps | None:
+        return self._model_sd_ops
 
-    def with_loras(self, loras: tuple[LoraPathStrengthAndSDOps, ...]) -> StreamingModelBuilder:
-        return replace(self, loras=loras)
+    @property
+    def module_ops(self) -> tuple[ModuleOps, ...]:
+        return self._module_ops
 
-    def with_fuse_rule(self, fuse_rule: FuseRule) -> StreamingModelBuilder:
-        return replace(self, fuse_rule=fuse_rule)
+    @property
+    def loras(self) -> tuple[LoraPathStrengthAndSDOps, ...]:
+        return self._loras
+
+    @property
+    def model_loader(self) -> StateDictLoader:
+        return self._model_loader
+
+    @property
+    def registry(self) -> Registry:
+        return self._registry
+
+    @property
+    def fuse_rule(self) -> FuseRule:
+        return self._fuse_rule
+
+    @property
+    def blocks_attr(self) -> str:
+        return self._blocks_attr
+
+    @property
+    def blocks_prefix(self) -> str:
+        return self._blocks_prefix
+
+    def with_sd_ops(self, sd_ops: SDOps | None) -> Self:
+        clone = copy.copy(self)
+        clone._model_sd_ops = sd_ops
+        return clone
+
+    def with_module_ops(self, module_ops: tuple[ModuleOps, ...]) -> Self:
+        clone = copy.copy(self)
+        clone._module_ops = module_ops
+        return clone
+
+    def with_loras(self, loras: tuple[LoraPathStrengthAndSDOps, ...]) -> Self:
+        clone = copy.copy(self)
+        clone._loras = loras
+        return clone
+
+    def with_registry(self, registry: Registry) -> Self:
+        clone = copy.copy(self)
+        clone._registry = registry
+        return clone
+
+    def with_lora_load_device(self, device: torch.device) -> Self:
+        # Streaming fuses LoRAs into pinned CPU buffers; no other staging device is meaningful.
+        raise NotImplementedError("StreamingModelBuilder loads LoRA weights on CPU only.")
+
+    def with_fuse_rule(self, fuse_rule: FuseRule) -> Self:
+        clone = copy.copy(self)
+        clone._fuse_rule = fuse_rule
+        return clone
 
     def model_config(self) -> dict:
         """Read model configuration from the checkpoint metadata."""
@@ -94,16 +178,16 @@ class StreamingModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType])
 
     def build(
         self,
-        target_device: torch.device,
-        dtype: torch.dtype,
+        device: torch.device | None = None,
+        dtype: torch.dtype | None = None,
         cpu_slots_count: int | None = None,
         gpu_slots_count: int | None = None,
         **_kwargs: object,
     ) -> BlockStreamingWrapper:
         """Build and return a ready-to-use :class:`BlockStreamingWrapper`.
         Args:
-            target_device: GPU device for compute.
-            dtype: Weight dtype (e.g. ``torch.bfloat16``).
+            device: GPU device for compute. ``None`` defaults to ``cuda``.
+            dtype: Weight dtype (e.g. ``torch.bfloat16``). Required.
             cpu_slots_count: Number of pinned CPU buffer slots.
                 ``None`` = RAM streaming (all blocks pre-loaded with LoRA fusion).
             gpu_slots_count: Number of GPU buffer slots.
@@ -111,6 +195,9 @@ class StreamingModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType])
         """
         if not self.blocks_prefix:
             raise ValueError("blocks_prefix must be non-empty for streaming")
+        if dtype is None:
+            raise ValueError("StreamingModelBuilder.build requires an explicit dtype")
+        device = device if device is not None else torch.device("cuda")
 
         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)
@@ -120,31 +207,44 @@ class StreamingModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType])
 
         checkpoint_paths = list(self.model_path) if isinstance(self.model_path, tuple) else [self.model_path]
         block_key_map, non_block_keys = _scan_checkpoint_keys(checkpoint_paths, self.model_sd_ops, self.blocks_prefix)
+        expected_indices = set(range(len(blocks)))
+        if set(block_key_map) != expected_indices:
+            missing = sorted(expected_indices - set(block_key_map))
+            extra = sorted(set(block_key_map) - expected_indices)
+            raise ValueError(
+                f"Block weights under prefix '{self.blocks_prefix}.' do not match the {len(blocks)} model blocks: "
+                f"missing indices {missing}, unexpected indices {extra}"
+            )
 
         cpu_slots_count = cpu_slots_count if cpu_slots_count is not None else len(blocks)
         gpu_slots_count = gpu_slots_count if gpu_slots_count is not None else _DEFAULT_GPU_SLOTS
 
         if cpu_slots_count >= len(blocks):
+            lora_sd_and_strengths = self._load_lora_sds()
             source, lora_sources = self._build_pinned_source(
-                meta_model, target_device, dtype, cpu_slots_count, block_key_map, non_block_keys
+                blocks, dtype, cpu_slots_count, block_key_map, lora_sd_and_strengths
             )
+            non_block_loras = lora_sd_and_strengths
         else:
             reader = DiskTensorReader(checkpoint_paths)
             source, lora_sources = self._build_disk_source(
-                meta_model, target_device, dtype, cpu_slots_count, reader, block_key_map, non_block_keys
+                blocks, dtype, cpu_slots_count, reader, block_key_map, prefetch_depth=_PREFETCH_DEPTH
             )
+            non_block_loras = [src.as_state_dict_with_strength() for src in lora_sources]
 
-        copy_stream = torch.cuda.Stream(device=target_device)
-        gpu_pool = WeightPool(
-            source.block_layout,
+        self._load_non_block_weights(meta_model, non_block_keys, device, dtype, non_block_loras)
+
+        copy_stream = torch.cuda.Stream(device=device)
+        gpu_pool = BufferPool(
+            source.slot_nbytes,
             gpu_slots_count,
-            target_device,
+            device,
             reuse_barrier=lambda event: copy_stream.wait_event(event),
         )
         provider = WeightsProvider(
             gpu_pool,
             copy_stream,
-            target_device,
+            device,
             source,
             lora_sources,
             self.blocks_prefix,
@@ -154,24 +254,12 @@ class StreamingModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType])
             model=meta_model,
             blocks=blocks,
             provider=provider,
-            target_device=target_device,
+            target_device=device,
         )
 
-    def _build_pinned_source(
-        self,
-        meta_model: nn.Module,
-        target_device: torch.device,
-        dtype: torch.dtype,
-        cpu_slots_count: int,
-        block_key_map: dict[int, list[tuple[str, str]]],
-        non_block_keys: list[tuple[str, str]],
-    ) -> tuple[WeightSource, list[LoraSource]]:
-        """Pre-load all blocks into pinned CPU buffers with LoRA fusion."""
-        model_sd = load_state_dict(
-            self.model_path, self.model_loader, self.registry, torch.device("cpu"), self.model_sd_ops
-        )
-
-        lora_sd_and_strengths = [
+    def _load_lora_sds(self) -> list[LoraStateDictWithStrength]:
+        """Load each configured LoRA into a state dict for fusion (pinned path)."""
+        return [
             LoraStateDictWithStrength(
                 load_state_dict([lora.path], self.model_loader, self.registry, torch.device("cpu"), lora.sd_ops),
                 lora.strength,
@@ -179,6 +267,25 @@ class StreamingModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType])
             for lora in self.loras
         ]
 
+    def _filtered_sd_ops(self, name_suffix: str, allowed_model_keys: frozenset[str]) -> SDOps:
+        """``model_sd_ops`` restricted to *allowed_model_keys* (post-rename keys).
+        The loader skips keys filtered to None before reading them, so a restricted
+        load never materializes the excluded partition. The distinct ``name`` avoids
+        a registry cache-id collision with the other partition.
+        """
+        base = self.model_sd_ops if self.model_sd_ops is not None else SDOps("streaming").with_matching()
+        allowed = allowed_model_keys if base.allowed_keys is None else (allowed_model_keys & base.allowed_keys)
+        return replace(base, name=f"{base.name}__{name_suffix}", allowed_keys=allowed)
+
+    def _build_pinned_source(
+        self,
+        blocks: nn.ModuleList,
+        dtype: torch.dtype,
+        cpu_slots_count: int,
+        block_key_map: dict[int, list[tuple[str, str]]],
+        lora_sd_and_strengths: list[LoraStateDictWithStrength],
+    ) -> tuple[WeightSource, list[LoraSource]]:
+        """Pre-load each block into its own contiguous pinned CPU buffer with LoRA fusion."""
         for block_idx in block_key_map:
             if block_idx >= cpu_slots_count:
                 raise ValueError(
@@ -186,87 +293,75 @@ class StreamingModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType])
                     f"got block index {block_idx} with only {cpu_slots_count} slots."
                 )
 
-        blocks = resolve_attr(meta_model, self.blocks_attr)
-        block_tensors: dict[str, torch.Tensor] = {}
+        # One contiguous pinned buffer per block, carved into per-param views. The
+        # views (flattened by full key) are filled in place; the source then keeps
+        # only the contiguous buffer and the layout to re-carve it on read.
+        pinned_buffers: dict[int, torch.Tensor] = {}
+        block_layouts: dict[int, TensorLayout] = {}
+        fill_views: dict[str, torch.Tensor] = {}
         for block_idx, entries in block_key_map.items():
-            block_params = dict(blocks[block_idx].named_parameters())
-            for _sft_key, param_name in entries:
-                key = make_block_key(self.blocks_prefix, block_idx, param_name)
-                block_tensors[key] = block_params[param_name]
-        blocks_layout = derive_layout(block_tensors, dtype)
-        pinned_blocks = allocate_layout_views(blocks_layout, pin_memory=True)
+            block_state = _block_state(blocks[block_idx])
+            layout = derive_layout({param_name: block_state[param_name] for _sft_key, param_name in entries}, dtype)
+            buffer = bs_utils.alloc_buffer(layout_nbytes(layout), torch.device("cpu"), pin_memory=True)
+            views = carve_buffer(buffer, layout)
+            pinned_buffers[block_idx] = buffer
+            block_layouts[block_idx] = layout
+            for param_name, view in views.items():
+                fill_views[make_block_key(self.blocks_prefix, block_idx, param_name)] = view
+
+        block_sd = load_state_dict(
+            self.model_path,
+            self.model_loader,
+            self.registry,
+            torch.device("cpu"),
+            self._filtered_sd_ops("blocks", frozenset(fill_views)),
+        )
 
         should_sync = False
         for key, fused in fuse_lora_weights(
-            model_sd, lora_sd_and_strengths, fuse_rule=self.fuse_rule, preserve_input_device=False
+            block_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
-                should_sync = True
-            else:
-                model_sd.sd[key] = fused
+            if key not in fill_views:
+                raise ValueError(f"Block-restricted load produced {key!r}, which is not a pinned block weight")
+            fill_views[key].copy_(fused, non_blocking=True)
+            block_sd.sd[key] = None
+            should_sync = True
         if should_sync:
             torch.cuda.synchronize()
 
         # Fill remaining pinned keys from the source state dict.
-        for key in blocks_layout:
-            if model_sd.sd[key] is None:
+        for key, view in fill_views.items():
+            if block_sd.sd[key] is None:
                 continue
-            pinned_blocks[key].copy_(model_sd.sd[key])
-            model_sd.sd[key] = None
-
-        pinned: dict[int, dict[str, torch.Tensor]] = {
-            block_idx: {
-                param_name: pinned_blocks[make_block_key(self.blocks_prefix, block_idx, param_name)]
-                for _sft_key, param_name in entries
-            }
-            for block_idx, entries in block_key_map.items()
-        }
-
-        non_block_sd: dict[str, torch.Tensor] = {
-            model_key: model_sd.sd[model_key].to(device=target_device, dtype=dtype)
-            for _sft_key, model_key in non_block_keys
-        }
-
-        meta_model.load_state_dict(non_block_sd, strict=False, assign=True)
+            view.copy_(block_sd.sd[key])
+            block_sd.sd[key] = None
 
+        pinned = {idx: PinnedBlock(pinned_buffers[idx], block_layouts[idx]) for idx in pinned_buffers}
         return PinnedWeightSource(pinned), []
 
     def _build_disk_source(
         self,
-        meta_model: nn.Module,
-        target_device: torch.device,
+        blocks: nn.ModuleList,
         dtype: torch.dtype,
         cpu_slots_count: int,
         reader: DiskTensorReader,
         block_key_map: dict[int, list[tuple[str, str]]],
-        non_block_keys: list[tuple[str, str]],
+        prefetch_depth: int,
     ) -> 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
-        relies on module_ops (e.g. fp8_cast) leaving the meta param dtype in
-        sync with the post-sd_ops checkpoint dtype.
+        """Create a DiskWeightSource backed by a DiskBlockReader.
+        Pool slots are sized to the largest block and carved per block on read, so
+        heterogeneous blocks (e.g. layers with differing attention layouts) share
+        one pool. Pool capacity is ``cpu_slots_count + prefetch_depth`` so the
+        lookahead loop in ``DiskWeightSource.get`` never evicts its own target.
+        Layouts come from the meta model; assumes module_ops keep the meta param
+        dtype in sync with the post-sd_ops checkpoint dtype.
         """
-        lora_sources = [LoraSource(lora.path, lora.sd_ops, lora.strength) for lora in self.loras]
+        block_layouts = _block_layouts(blocks, block_key_map, dtype)
+        slot_nbytes = max(layout_nbytes(layout) for layout in block_layouts.values())
 
-        self._load_non_block_weights(
-            reader,
-            non_block_keys,
-            meta_model,
-            target_device,
-            dtype,
-            sd_ops=self.model_sd_ops,
-            lora_sources=lora_sources,
-            fuse_rule=self.fuse_rule,
-        )
-
-        blocks = resolve_attr(meta_model, self.blocks_attr)
-        layout = derive_layout(dict(blocks[0].named_parameters()), dtype)
-
-        cpu_pool = WeightPool(
-            layout,
-            cpu_slots_count,
+        cpu_pool = BufferPool(
+            slot_nbytes,
+            cpu_slots_count + prefetch_depth,
             torch.device("cpu"),
             reuse_barrier=lambda event: event.synchronize(),
             pin_memory=True,
@@ -277,43 +372,42 @@ class StreamingModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType])
             sd_ops=self.model_sd_ops,
             blocks_prefix=self.blocks_prefix,
         )
-        source = DiskWeightSource(cpu_pool, block_reader)
+        fetcher = BlockFetcher(block_reader)
+        source = DiskWeightSource(
+            cpu_pool,
+            fetcher,
+            block_layouts,
+            blocks_number=len(blocks),
+            prefetch_depth=prefetch_depth,
+        )
+        lora_sources = [LoraSource(lora.path, lora.sd_ops, lora.strength) for lora in self.loras]
+
         return source, lora_sources
 
-    @staticmethod
     @torch.inference_mode()
     def _load_non_block_weights(
-        reader: DiskTensorReader,
-        non_block_keys: list[tuple[str, str]],
+        self,
         model: nn.Module,
+        non_block_keys: list[tuple[str, str]],
         device: torch.device,
         dtype: torch.dtype,
-        sd_ops: SDOps | None = None,
-        lora_sources: list[LoraSource] | None = None,
-        fuse_rule: FuseRule = bf16_fuse_rule,
+        lora_sd_and_strengths: list[LoraStateDictWithStrength],
     ) -> None:
-        """Load non-block weights into *model* on *device*, applying ``sd_ops`` and fusing LoRAs.
-        Fusion goes through :func:`fuse_lora_weights` under *fuse_rule* so the
-        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.
+        """Load the non-block weights onto *device* and fuse LoRAs -- both paths.
+        Reads through the loader with ``model_sd_ops`` restricted to the
+        non-block keys, so ``sd_ops`` (incl. kv-ops such as Gemma's ``lm_head``
+        duplication) is applied exactly once and block tensors are never read.
         """
-        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)
-            if sd_ops is not None:
-                for kv in sd_ops.apply_to_key_value(model_key, tensor):
-                    non_block_sd[kv.new_key] = kv.new_value
-            else:
-                non_block_sd[model_key] = tensor
+        non_block_sd_ops = self._filtered_sd_ops("non_block", frozenset(mk for _sft_key, mk in non_block_keys))
+        loaded = load_state_dict(self.model_path, self.model_loader, self.registry, device, non_block_sd_ops)
+        non_block_sd = {key: tensor.to(dtype=dtype) for key, tensor in loaded.sd.items()}
 
-        if lora_sources:
-            lora_sd_and_strengths = [src.as_state_dict_with_strength() for src in lora_sources]
+        if lora_sd_and_strengths:
             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,
+                fuse_rule=self.fuse_rule,
                 preserve_input_device=True,
             ):
                 non_block_sd[key] = fused
@@ -321,6 +415,35 @@ class StreamingModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType])
         model.load_state_dict(non_block_sd, strict=False, assign=True)
 
 
+def _block_state(block: nn.Module) -> dict[str, torch.Tensor]:
+    """Streamed-eligible tensors of a block: parameters then buffers.
+    Block streaming swaps both params and checkpoint-backed buffers (e.g. Gemma4's
+    per-layer ``layer_scalar``), so the layout, pinned packing, and meta-ordering
+    all consult parameters and buffers together. Non-checkpoint (computed) buffers
+    are harmless here -- only keys present in ``block_key_map`` are ever streamed.
+    """
+    return {**dict(block.named_parameters()), **dict(block.named_buffers())}
+
+
+def _block_layouts(
+    blocks: nn.ModuleList,
+    block_key_map: dict[int, list[tuple[str, str]]],
+    dtype: torch.dtype,
+) -> dict[int, TensorLayout]:
+    """Per-block layout of the streamed tensors, taken from the meta model.
+    Blocks may differ in shape and even in which tensors they have (e.g. Gemma4's
+    full-attention layers drop ``v_proj``), so each block gets its own layout in
+    ``block_key_map`` order. The pinned packing, the disk reader, and the GPU carve
+    all key off this same per-block layout, so the provider's contiguous H2D copy
+    is valid for any entry order (no cross-block ordering required).
+    """
+    layouts: dict[int, TensorLayout] = {}
+    for idx, entries in block_key_map.items():
+        state = _block_state(blocks[idx])
+        layouts[idx] = derive_layout({param_name: state[param_name] for _sft_key, param_name in entries}, dtype)
+    return layouts
+
+
 def _scan_checkpoint_keys(
     checkpoint_paths: list[str],
     sd_ops: SDOps | None,

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

@@ -128,8 +128,8 @@ class LoraSource:
 
     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``.
+        Lets non-block fusion consume the already-loaded disk-streaming LoRA
+        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():

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

@@ -1,4 +1,4 @@
-"""Weight buffer pool for block streaming."""
+"""Raw buffer pool for block streaming."""
 
 from __future__ import annotations
 
@@ -7,69 +7,64 @@ from typing import Callable
 
 import torch
 
-from ltx_core.block_streaming.utils import allocate_layout_views
-from ltx_core.loader.primitives import TensorLayout
+from ltx_core.block_streaming import utils
 
 
-class WeightPool:
-    """Fixed pool of pre-allocated weight buffers with event-based reuse.
-    All slots share a single buffer (CPU or GPU); each slot is a
-    contiguous slice carved out of it via :func:`allocate_layout_views`.
+class BufferPool:
+    """Fixed pool of pre-allocated raw buffer slots with event-based reuse.
+    Slots are carved from a single contiguous ``uint8`` buffer; each is
+    ``slot_nbytes`` long and handed out as a raw 1-D ``uint8`` tensor.
     Args:
-        buffer_layout: ``{name: (shape, dtype)}`` for each buffer.
-        capacity: Number of buffers to pre-allocate.
+        slot_nbytes: Byte size of each slot.
+        capacity: Number of slots to pre-allocate.
         device: Device for allocation.
-        reuse_barrier: Called with the pending event before a buffer is reused.
+        reuse_barrier: Called with the pending event before a slot is reused.
         pin_memory: Pin buffers (for async H2D copies from CPU).
     """
 
     def __init__(
         self,
-        buffer_layout: TensorLayout,
+        slot_nbytes: int,
         capacity: int,
         device: torch.device,
         reuse_barrier: Callable[[torch.cuda.Event], None],
         pin_memory: bool = False,
     ) -> None:
-        self._buffer_layout = buffer_layout
+        self._slot_nbytes = slot_nbytes
         self._capacity = capacity
-        self._free: deque[dict[str, torch.Tensor]] = deque()
+        self._free: deque[torch.Tensor] = deque()
         self._events: dict[int, torch.cuda.Event] = {}
         self._reuse_barrier = reuse_barrier
-        memory_layout = {
-            _make_key(slot, name): (shape, dtype)
-            for slot in range(capacity)
-            for name, (shape, dtype) in buffer_layout.items()
-        }
-        all_views = allocate_layout_views(memory_layout, device=device, pin_memory=pin_memory)
+        buffer = utils.alloc_buffer(max(slot_nbytes * capacity, 1), device, pin_memory)
         for slot in range(capacity):
-            self._free.append({name: all_views[_make_key(slot, name)] for name in buffer_layout})
+            self._free.append(buffer[slot * slot_nbytes : (slot + 1) * slot_nbytes])
 
     @property
     def capacity(self) -> int:
         return self._capacity
 
     @property
-    def buffer_layout(self) -> TensorLayout:
-        return self._buffer_layout
+    def slot_nbytes(self) -> int:
+        return self._slot_nbytes
 
-    def acquire(self) -> dict[str, torch.Tensor]:
-        """Take a free buffer, waiting any pending event before returning."""
-        weights = self._free.popleft()
-        event = self._events.pop(id(weights), None)
+    def acquire(self) -> torch.Tensor:
+        """Take a free raw slot, waiting any pending event before returning.
+        Raises :class:`RuntimeError` if every slot is currently in use.
+        """
+        if not self._free:
+            raise RuntimeError(f"BufferPool exhausted: all {self._capacity} buffers are in use")
+        buffer = self._free.popleft()
+        event = self._events.pop(id(buffer), None)
         if event is not None:
             self._reuse_barrier(event)
-        return weights
+        return buffer
 
-    def release(self, weights: dict[str, torch.Tensor], event: torch.cuda.Event | None = None) -> None:
-        """Return a buffer to the free list.
-        If *event* is given it is waited on the next :meth:`acquire`
-        of this buffer, ensuring the prior operation has completed.
+    def release(self, buffer: torch.Tensor, event: torch.cuda.Event | None = None) -> None:
+        """Return a raw slot to the free list.
+        The *buffer* must be the exact tensor object returned by :meth:`acquire`
+        (reuse is keyed on its identity). If *event* is given it is waited on the
+        next :meth:`acquire` of this slot, ensuring the prior operation finished.
         """
         if event is not None:
-            self._events[id(weights)] = event
-        self._free.append(weights)
-
-
-def _make_key(slot: int, name: str) -> str:
-    return f"{slot}/{name}"
+            self._events[id(buffer)] = event
+        self._free.append(buffer)

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

@@ -3,37 +3,28 @@
 from __future__ import annotations
 
 from collections import OrderedDict
+from typing import NamedTuple
 
 import torch
 
 from ltx_core.block_streaming.disk import LoraSource
-from ltx_core.block_streaming.pool import WeightPool
+from ltx_core.block_streaming.pool import BufferPool
 from ltx_core.block_streaming.source import WeightSource
+from ltx_core.block_streaming.utils import carve_buffer, layout_nbytes
 from ltx_core.loader.fuse_loras import FuseRule, aggregate_lora_products, bf16_fuse_rule
 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:
-    """Return a ``uint8`` view spanning every tensor in *weights*, or ``None`` if
-    they don't share one contiguous storage region."""
-    tensors = list(weights.values())
-    if not tensors:
-        return None
-    storage = tensors[0].untyped_storage()
-    storage_ptr = storage.data_ptr()
-    start = end = tensors[0].storage_offset() * tensors[0].element_size()
-    for t in tensors:
-        if t.untyped_storage().data_ptr() != storage_ptr or not t.is_contiguous():
-            return None
-        offset = t.storage_offset() * t.element_size()
-        nbytes = t.numel() * t.element_size()
-        start = min(start, offset)
-        end = max(end, offset + nbytes)
-    view = torch.empty(0, dtype=torch.uint8, device=tensors[0].device)
-    view.set_(storage, start, (end - start,), (1,))
-    return view
+class CachedBlock(NamedTuple):
+    """A cached GPU block: the raw pool slot plus the carved per-key views.
+    The raw slot is what is returned to the pool on eviction; the views are
+    what callers consume.
+    """
+
+    raw: torch.Tensor
+    views: dict[str, torch.Tensor]
 
 
 class WeightsProvider:
@@ -51,7 +42,7 @@ class WeightsProvider:
 
     def __init__(
         self,
-        pool: WeightPool,
+        pool: BufferPool,
         copy_stream: torch.cuda.Stream,
         target_device: torch.device,
         source: WeightSource,
@@ -61,7 +52,7 @@ class WeightsProvider:
     ) -> None:
         self._copy_stream = copy_stream
         self._pool = pool
-        self._cache: OrderedDict[int, dict[str, torch.Tensor]] = OrderedDict()
+        self._cache: OrderedDict[int, CachedBlock] = OrderedDict()
         self._events: dict[int, torch.cuda.Event] = {}
         self._target_device = target_device
         self._source = source
@@ -72,40 +63,45 @@ class WeightsProvider:
     def get(self, idx: int) -> dict[str, torch.Tensor]:
         """Return GPU weights for block *idx*. Does H2D copy on miss."""
         if idx in self._cache:
-            return self._cache[idx]
+            return self._cache[idx].views
 
         # Evict oldest GPU buffer if at capacity.
         if len(self._cache) >= self._pool.capacity:
-            evicted_idx, evicted_weights = self._cache.popitem(last=False)
-            self._pool.release(evicted_weights, event=self._events.pop(evicted_idx, None))
+            evicted_idx, evicted = self._cache.popitem(last=False)
+            self._pool.release(evicted.raw, event=self._events.pop(evicted_idx, None))
 
-        gpu_weights = self._pool.acquire()
-        cpu_weights = self._source.get(idx)
+        layout = self._source.block_layout(idx)
+        raw = self._pool.acquire()
+        gpu_weights = carve_buffer(raw, layout)
+        cpu_buffer = self._source.get(idx)
 
-        h2d_event = self._copy_to_gpu(idx, gpu_weights, cpu_weights)
+        h2d_event = self._copy_to_gpu(idx, raw, gpu_weights, cpu_buffer, layout_nbytes(layout))
         self._source.release(idx, event=h2d_event)
 
-        self._cache[idx] = gpu_weights
+        self._cache[idx] = CachedBlock(raw, gpu_weights)
         return gpu_weights
 
     def _copy_to_gpu(
         self,
         idx: int,
+        raw: torch.Tensor,
         gpu_weights: dict[str, torch.Tensor],
-        cpu_weights: dict[str, torch.Tensor],
+        cpu_buffer: torch.Tensor,
+        nbytes: int,
     ) -> torch.cuda.Event:
         """Enqueue H2D copy + LoRA fusion on the copy stream and wait on compute.
-        The wait is intentionally inside this method so callers -- and
-        instrumentation regions wrapping it -- observe the full transfer time.
+        *cpu_buffer* is one contiguous source buffer carved by the same layout as
+        *raw*, so a single byte copy of its leading *nbytes* reproduces every view
+        in *gpu_weights*. The wait is intentionally inside this method so callers --
+        and instrumentation regions wrapping it -- observe the full transfer time.
         """
+        if not cpu_buffer.is_contiguous() or cpu_buffer.dtype != torch.uint8 or cpu_buffer.numel() < nbytes:
+            raise ValueError(
+                f"source buffer for block {idx} must be a contiguous uint8 buffer of >= {nbytes} bytes, "
+                f"got {cpu_buffer.dim()}-D {cpu_buffer.dtype} with {cpu_buffer.numel()} elements"
+            )
         with torch.cuda.stream(self._copy_stream):
-            gpu_view = _contiguous_byte_view(gpu_weights)
-            cpu_view = _contiguous_byte_view(cpu_weights)
-            if gpu_view is not None and cpu_view is not None and gpu_view.numel() == cpu_view.numel():
-                gpu_view.copy_(cpu_view, non_blocking=True)
-            else:
-                for name, gpu_tensor in gpu_weights.items():
-                    gpu_tensor.copy_(cpu_weights[name], non_blocking=True)
+            raw[:nbytes].copy_(cpu_buffer[:nbytes], non_blocking=True)
             if self._lora_sources:
                 self._fuse_block_loras(idx, gpu_weights)
             h2d_event = torch.cuda.Event()

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

@@ -2,31 +2,38 @@
 
 from __future__ import annotations
 
-from collections import OrderedDict
-from typing import Protocol
+from typing import NamedTuple, Protocol
 
 import torch
 
-from ltx_core.block_streaming.disk import DiskBlockReader
-from ltx_core.block_streaming.pool import WeightPool
+from ltx_core.block_streaming.block_fetcher import BlockFetcher, FetchHandle
+from ltx_core.block_streaming.pool import BufferPool
+from ltx_core.block_streaming.utils import carve_buffer, layout_nbytes
 from ltx_core.loader.primitives import TensorLayout
 
 
 class WeightSource(Protocol):
     """Provides pinned CPU weights for a given block index.
-    Assumes all buffers share an identical layout across all block indices.
+    Blocks may be heterogeneous: each has its own layout, so the source exposes a
+    per-block layout and the byte size of the largest block (which sizes the
+    pool slots -- a smaller block is carved into the front of a max-sized slot).
+    The source is the single source of truth for each block's layout.
     """
 
+    def block_layout(self, idx: int) -> TensorLayout:
+        """Per-block buffer layout (shape + dtype for each param)."""
+        ...
+
     @property
-    def block_layout(self) -> TensorLayout:
-        """Shared per-block buffer layout (shape + dtype for each param)."""
+    def slot_nbytes(self) -> int:
+        """Byte size of the largest block; sizes a pool slot (16-byte aligned)."""
         ...
 
-    def get(self, idx: int) -> dict[str, torch.Tensor]:
-        """Return CPU weights for block *idx*."""
+    def get(self, idx: int) -> torch.Tensor:
+        """Return one contiguous CPU buffer for block *idx*."""
         ...
 
-    def release(self, idx: int, event: torch.cuda.Event) -> None:
+    def release(self, idx: int, event: torch.cuda.Event | None) -> None:
         """Signal that an async operation using these weights is guarded by *event*."""
         ...
 
@@ -35,68 +42,133 @@ class WeightSource(Protocol):
         ...
 
 
-class DiskWeightSource(WeightSource):
-    """Reads block weights from disk into pinned CPU buffers on demand."""
+class _Scheduled(NamedTuple):
+    """A scheduled (possibly in-flight) read: the raw pool slot and its fetch handle."""
+
+    raw: torch.Tensor
+    status: FetchHandle
+
+
+class DiskWeightSource(WeightSource):
+    """WeightSource that streams blocks from disk via a :class:`BlockFetcher`.
+    ``get(idx)`` must be paired with a ``release(idx)`` before *idx* is fetched
+    again; getting a block that is already in flight raises. Each read acquires a
+    raw pool slot, carves it to block *idx*'s layout, and hands the carved views to
+    the fetcher to fill, so one max-sized slot can serve blocks of differing shapes.
+    ``get`` returns that same contiguous slot.
+    """
+
+    def __init__(
+        self,
+        pool: BufferPool,
+        fetcher: BlockFetcher,
+        block_layouts: dict[int, TensorLayout],
+        blocks_number: int,
+        prefetch_depth: int = 0,
+    ) -> None:
+        if blocks_number <= 0:
+            raise ValueError(f"blocks_number must be > 0, got {blocks_number}")
+        if prefetch_depth < 0:
+            raise ValueError(f"prefetch_depth must be >= 0, got {prefetch_depth}")
+        max_layout_nbytes = max((layout_nbytes(layout) for layout in block_layouts.values()), default=0)
+        if pool.slot_nbytes < max_layout_nbytes:
+            raise ValueError(
+                f"pool slot is too small for the largest block: slot {pool.slot_nbytes} bytes < {max_layout_nbytes}"
+            )
 
-    def __init__(self, pool: WeightPool, reader: DiskBlockReader) -> None:
         self._pool = pool
-        self._cache: OrderedDict[int, dict[str, torch.Tensor]] = OrderedDict()
-        self._events: dict[int, torch.cuda.Event] = {}
-        self._reader = reader
+        self._blocks_number = blocks_number
+        self._prefetch_depth = prefetch_depth
+        self._fetcher = fetcher
+        self._block_layouts = block_layouts
+        self._scheduled: dict[int, _Scheduled] = {}
+        self._in_flight: dict[int, torch.Tensor] = {}
+
+    def block_layout(self, idx: int) -> TensorLayout:
+        return self._block_layouts[idx]
 
     @property
-    def block_layout(self) -> TensorLayout:
-        return self._pool.buffer_layout
+    def slot_nbytes(self) -> int:
+        return self._pool.slot_nbytes
 
-    def get(self, idx: int) -> dict[str, torch.Tensor]:
-        """Return CPU weights for block *idx*. Reads from disk on miss."""
-        if idx in self._cache:
-            return self._cache[idx]
+    def get(self, idx: int) -> torch.Tensor:
+        if idx in self._in_flight:
+            raise RuntimeError(f"Block {idx} is already in flight; release it before getting it again")
 
-        if len(self._cache) >= self._pool.capacity:
-            evicted_idx, evicted_weights = self._cache.popitem(last=False)
-            self._pool.release(evicted_weights, event=self._events.pop(evicted_idx, None))
+        scheduled = self._scheduled.pop(idx, None)
+        if scheduled is None:
+            scheduled = self._schedule(idx)
+        error = scheduled.status.wait()
+        if error is not None:
+            self._pool.release(scheduled.raw)
+            raise error
 
-        weights = self._pool.acquire()
-        self._reader.read_into(weights, idx)
-        self._cache[idx] = weights
-        return weights
+        self._in_flight[idx] = scheduled.raw
+        for k in range(1, self._prefetch_depth + 1):
+            self._ensure_scheduled((idx + k) % self._blocks_number)
+        return scheduled.raw
 
-    def release(self, idx: int, event: torch.cuda.Event) -> None:
-        """Attach an H2D event -- waited before this buffer is recycled."""
-        self._events[idx] = event
+    def release(self, idx: int, event: torch.cuda.Event | None) -> None:
+        raw_buffer = self._in_flight.pop(idx)
+        self._pool.release(raw_buffer, event=event)
 
     def cleanup(self) -> None:
-        """Clear cache and close the disk reader."""
-        self._cache.clear()
-        self._events.clear()
-        self._reader.cleanup()
+        self._fetcher.cleanup()
+        while self._in_flight:
+            _, raw_buffer = self._in_flight.popitem()
+            self._pool.release(raw_buffer)
+        while self._scheduled:
+            _, scheduled = self._scheduled.popitem()
+            scheduled.status.wait()
+            self._pool.release(scheduled.raw)
 
-    def __len__(self) -> int:
-        return len(self._cache)
+    def _ensure_scheduled(self, idx: int) -> None:
+        """Schedule a read for *idx* if one is not already pending."""
+        if idx not in self._scheduled:
+            self._scheduled[idx] = self._schedule(idx)
+
+    def _schedule(self, idx: int) -> _Scheduled:
+        """Acquire a raw slot, carve it to block *idx*, enqueue a read, return the handle.
+        The raw slot and its fetch status are returned together so the caller can
+        track both as one unit; the fetcher only receives the carved views to fill.
+        """
+        raw_buffer = self._pool.acquire()
+        carved = carve_buffer(raw_buffer, self._block_layouts[idx])
+        status = self._fetcher.submit(idx, carved)
+        return _Scheduled(raw_buffer, status)
+
+
+class PinnedBlock(NamedTuple):
+    """A pre-loaded pinned block: its single contiguous buffer and the layout to carve it with."""
+
+    buffer: torch.Tensor
+    layout: TensorLayout
 
 
 class PinnedWeightSource(WeightSource):
-    """Pre-loaded pinned CPU weights."""
+    """Pre-loaded pinned CPU weights, one contiguous (possibly heterogeneous) buffer per block."""
 
-    def __init__(self, weights: dict[int, dict[str, torch.Tensor]]) -> None:
-        if not weights:
+    def __init__(self, blocks: dict[int, PinnedBlock]) -> None:
+        if not blocks:
             raise ValueError("PinnedWeightSource requires at least one block")
-        self._weights = weights
+        self._blocks = blocks
+        self._slot_nbytes = max(layout_nbytes(block.layout) for block in blocks.values())
+
+    def block_layout(self, idx: int) -> TensorLayout:
+        return self._blocks[idx].layout
 
     @property
-    def block_layout(self) -> TensorLayout:
-        first_block = self._weights[min(self._weights)]
-        return {name: (t.shape, t.dtype) for name, t in first_block.items()}
+    def slot_nbytes(self) -> int:
+        return self._slot_nbytes
 
-    def get(self, idx: int) -> dict[str, torch.Tensor]:
-        return self._weights[idx]
+    def get(self, idx: int) -> torch.Tensor:
+        return self._blocks[idx].buffer
 
-    def release(self, idx: int, event: torch.cuda.Event) -> None:
+    def release(self, idx: int, event: torch.cuda.Event | None) -> None:
         pass
 
     def cleanup(self) -> None:
-        self._weights.clear()
+        self._blocks.clear()
 
     def __len__(self) -> int:
-        return len(self._weights)
+        return len(self._blocks)

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

@@ -5,7 +5,7 @@ from __future__ import annotations
 import math
 import weakref
 from dataclasses import dataclass
-from typing import Any
+from typing import Any, NamedTuple
 
 import torch
 from torch import nn
@@ -84,7 +84,7 @@ def _alloc_pinned_exact(nbytes: int) -> torch.Tensor | None:
     return buf
 
 
-def _alloc_buffer(nbytes: int, device: torch.device | None, pin_memory: bool) -> torch.Tensor:
+def alloc_buffer(nbytes: int, device: torch.device | None, pin_memory: bool) -> torch.Tensor:
     """Allocate one ``uint8`` buffer for :func:`allocate_layout_views`.
     For pinned host buffers, prefer ``cudaHostRegister`` to dodge the caching
     allocator's power-of-2 rounding. Falls back to the caching allocator if
@@ -112,6 +112,47 @@ class _TensorSlice:
         return math.prod(self.shape) * self.dtype.itemsize
 
 
+class LayoutSlices(NamedTuple):
+    """Per-key tensor slices of a layout plus the total aligned buffer size."""
+
+    slices: dict[str, _TensorSlice]
+    nbytes: int
+
+
+def _layout_slices(layout: TensorLayout) -> LayoutSlices:
+    """Compute the byte offset of each key in *layout* and the total aligned size.
+    The size is at least one byte so empty layouts still produce a valid buffer.
+    """
+    slices: dict[str, _TensorSlice] = {}
+    cursor = 0
+    for key, (shape, dtype) in layout.items():
+        cursor = _align_up(cursor, _BUFFER_ALIGN)
+        slices[key] = _TensorSlice(offset=cursor, shape=shape, dtype=dtype)
+        cursor += slices[key].size()
+    return LayoutSlices(slices, max(_align_up(cursor, _BUFFER_ALIGN), 1))
+
+
+def layout_nbytes(layout: TensorLayout) -> int:
+    """Byte size of one contiguous, 16-byte-aligned buffer holding *layout* (>= 1)."""
+    return _layout_slices(layout).nbytes
+
+
+def carve_buffer(buffer: torch.Tensor, layout: TensorLayout) -> dict[str, torch.Tensor]:
+    """Carve per-key tensor views for *layout* into the front of *buffer*.
+    *buffer* is a 1-D ``uint8`` tensor at least :func:`layout_nbytes` long. Each
+    returned tensor is a non-overlapping slice of its leading bytes reinterpreted
+    at the requested shape and dtype; any trailing bytes are left unused. That
+    slack is what lets one max-sized pool slot hold a smaller (heterogeneous)
+    block. The views keep *buffer*'s storage alive via PyTorch refcounting.
+    """
+    if buffer.dtype != torch.uint8 or buffer.dim() != 1:
+        raise ValueError(f"carve_buffer expects a 1-D uint8 buffer, got {buffer.dim()}-D {buffer.dtype}")
+    slices, nbytes = _layout_slices(layout)
+    if buffer.numel() < nbytes:
+        raise ValueError(f"buffer too small to carve layout: need {nbytes} bytes, got {buffer.numel()}")
+    return {key: buffer[s.offset : s.offset + s.size()].view(s.dtype).view(s.shape) for key, s in slices.items()}
+
+
 def allocate_layout_views(
     layout: TensorLayout,
     device: torch.device | None = None,
@@ -123,12 +164,5 @@ def allocate_layout_views(
     requested shape and dtype. The views keep the underlying storage alive
     via PyTorch refcounting — drop them all to release the memory.
     """
-    slices: dict[str, _TensorSlice] = {}
-    cursor = 0
-    for key, (shape, dtype) in layout.items():
-        cursor = _align_up(cursor, _BUFFER_ALIGN)
-        slices[key] = _TensorSlice(offset=cursor, shape=shape, dtype=dtype)
-        cursor += slices[key].size()
-    # Allocate at least one byte so empty layouts still produce a valid buffer.
-    buffer = _alloc_buffer(max(_align_up(cursor, _BUFFER_ALIGN), 1), device, pin_memory)
-    return {key: buffer[s.offset : s.offset + s.size()].view(s.dtype).view(s.shape) for key, s in slices.items()}
+    buffer = alloc_buffer(layout_nbytes(layout), device, pin_memory)
+    return carve_buffer(buffer, layout)

packages/ltx-core/src/ltx_core/components/guiders.py

@@ -253,6 +253,11 @@ class MultiModalGuider:
         and as scale * (cond - uncond) for stg, steering the denoising process away from the unconditioned
         prediction.
         """
+        dtype = cond.dtype
+        cond = cond.float()
+        uncond_text = uncond_text.float() if isinstance(uncond_text, torch.Tensor) else uncond_text
+        uncond_perturbed = uncond_perturbed.float() if isinstance(uncond_perturbed, torch.Tensor) else uncond_perturbed
+        uncond_modality = uncond_modality.float() if isinstance(uncond_modality, torch.Tensor) else uncond_modality
         pred = (
             cond
             + (self.params.cfg_scale - 1) * (cond - uncond_text)
@@ -265,7 +270,7 @@ class MultiModalGuider:
             factor = self.params.rescale_scale * factor + (1 - self.params.rescale_scale)
             pred = pred * factor
 
-        return pred
+        return pred.to(dtype)
 
     def do_unconditional_generation(self) -> bool:
         """Returns True if the guider is doing unconditional generation."""

packages/ltx-core/src/ltx_core/components/noisers.py

@@ -17,18 +17,20 @@ class GaussianNoiser(Noiser):
 
     def __init__(self, generator: torch.Generator):
         super().__init__()
-
         self.generator = generator
 
-    def __call__(self, latent_state: LatentState, noise_scale: float = 1.0) -> LatentState:
-        noise = torch.randn(
+    def _sample_noise(self, latent_state: LatentState) -> torch.Tensor:
+        return torch.randn(
             *latent_state.latent.shape,
             device=latent_state.latent.device,
             dtype=latent_state.latent.dtype,
             generator=self.generator,
         )
-        scaled_mask = latent_state.denoise_mask * noise_scale
-        latent = noise * scaled_mask + latent_state.latent * (1 - scaled_mask)
+
+    def __call__(self, latent_state: LatentState, noise_scale: float = 1.0) -> LatentState:
+        noise = self._sample_noise(latent_state)
+        latent = torch.lerp(latent_state.latent.float(), noise.float(), noise_scale)
+        latent = torch.lerp(latent_state.clean_latent.float(), latent, latent_state.denoise_mask)
         return replace(
             latent_state,
             latent=latent.to(latent_state.latent.dtype),

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

@@ -7,6 +7,7 @@ from ltx_core.conditioning.types import (
     ConditioningItemAttentionStrengthWrapper,
     VideoConditionByKeyframeIndex,
     VideoConditionByLatentIndex,
+    VideoConditionByMask,
     VideoConditionByReferenceLatent,
 )
 
@@ -17,5 +18,6 @@ __all__ = [
     "ConditioningItemAttentionStrengthWrapper",
     "VideoConditionByKeyframeIndex",
     "VideoConditionByLatentIndex",
+    "VideoConditionByMask",
     "VideoConditionByReferenceLatent",
 ]

packages/ltx-core/src/ltx_core/conditioning/types/__init__.py

@@ -3,6 +3,7 @@
 from ltx_core.conditioning.types.attention_strength_wrapper import ConditioningItemAttentionStrengthWrapper
 from ltx_core.conditioning.types.keyframe_cond import VideoConditionByKeyframeIndex
 from ltx_core.conditioning.types.latent_cond import VideoConditionByLatentIndex
+from ltx_core.conditioning.types.mask_cond import VideoConditionByMask
 from ltx_core.conditioning.types.reference_audio_cond import AudioConditionByReferenceLatent
 from ltx_core.conditioning.types.reference_video_cond import VideoConditionByReferenceLatent
 
@@ -11,5 +12,6 @@ __all__ = [
     "ConditioningItemAttentionStrengthWrapper",
     "VideoConditionByKeyframeIndex",
     "VideoConditionByLatentIndex",
+    "VideoConditionByMask",
     "VideoConditionByReferenceLatent",
 ]

packages/ltx-core/src/ltx_core/conditioning/types/keyframe_cond.py

@@ -10,8 +10,9 @@ from ltx_core.types import LatentState, VideoLatentShape
 class VideoConditionByKeyframeIndex(ConditioningItem):
     """
     Conditions video generation on keyframe latents at a specific frame index.
-    Appends keyframe tokens to the latent state with positions offset by frame_idx,
-    and sets denoise strength according to the strength parameter.
+    Appends keyframe tokens to the sequence with positions offset by frame_idx: the keyframe
+    latents become clean-latent tokens (placeholder zeros in the noisy latent) and the denoise
+    mask is set from the strength parameter.
     To add attention masking, wrap with :class:`ConditioningItemAttentionStrengthWrapper`.
     Args:
         keyframes: Keyframe latents [B, C, F, H, W].
@@ -75,7 +76,7 @@ class VideoConditionByKeyframeIndex(ConditioningItem):
         )
 
         return LatentState(
-            latent=torch.cat([latent_state.latent, tokens], dim=1),
+            latent=torch.cat([latent_state.latent, torch.zeros_like(tokens)], dim=1),
             denoise_mask=torch.cat([latent_state.denoise_mask, denoise_mask], dim=1),
             positions=torch.cat([latent_state.positions, positions], dim=2),
             clean_latent=torch.cat([latent_state.clean_latent, tokens], dim=1),

packages/ltx-core/src/ltx_core/conditioning/types/latent_cond.py

@@ -9,8 +9,8 @@ from ltx_core.types import LatentState
 class VideoConditionByLatentIndex(ConditioningItem):
     """
     Conditions video generation by injecting latents at a specific latent frame index.
-    Replaces tokens in the latent state at positions corresponding to latent_idx,
-    and sets denoise strength according to the strength parameter.
+    Sets the clean latents at positions corresponding to latent_idx to the injected latents,
+    sets denoise strength according to the strength parameter.
     """
 
     def __init__(self, latent: torch.Tensor, strength: float, latent_idx: int):
@@ -37,7 +37,6 @@ class VideoConditionByLatentIndex(ConditioningItem):
 
         latent_state = latent_state.clone()
 
-        latent_state.latent[:, start_token:stop_token] = tokens
         latent_state.clean_latent[:, start_token:stop_token] = tokens
         latent_state.denoise_mask[:, start_token:stop_token] = 1.0 - self.strength
 

packages/ltx-core/src/ltx_core/conditioning/types/mask_cond.py

@@ -0,0 +1,49 @@
+"""Mask-based conditioning for inpainting and spatial conditioning."""
+
+from dataclasses import replace
+
+import torch
+
+from ltx_core.conditioning.item import ConditioningItem
+from ltx_core.tools import LatentTools
+from ltx_core.types import LatentState
+
+
+class VideoConditionByMask(ConditioningItem):
+    """Condition video generation using a binary mask over latent frames.
+    Masked positions (mask=1) receive the provided clean latent values and are
+    excluded from denoising (denoise_mask set to ``1 - strength``). Unmasked
+    positions (mask=0) are left unchanged and denoised normally.
+    The mask operates in **unpatchified latent** space — it should have shape
+    ``[B, F, H, W]`` matching the latent dimensions (after VAE encoding,
+    before patchification). This is consistent with the latent input format
+    used by all other conditioning items.
+    Args:
+        latent: Clean conditioning latents in unpatchified format [B, C, F, H, W].
+            Must match the target shape of the latent tools.
+        mask: Binary mask [B, F, H, W] in unpatchified latent space.
+            1 = conditioning position (clean, excluded from denoising),
+            0 = generated position (noised, denoised normally).
+        strength: Conditioning strength for masked positions. 1.0 = fully clean
+            (no denoising), 0.0 = no conditioning effect. Default 1.0.
+    """
+
+    def __init__(self, latent: torch.Tensor, mask: torch.Tensor, strength: float = 1.0):
+        self.latent = latent
+        self.mask = mask
+        self.strength = strength
+
+    def apply_to(self, latent_state: LatentState, latent_tools: LatentTools) -> LatentState:
+        """Apply mask-based conditioning to the latent state."""
+        tokens = latent_tools.patchifier.patchify(self.latent)
+
+        mask = latent_tools.patchifier.patchify(self.mask.unsqueeze(1))
+
+        m = mask.to(dtype=latent_state.latent.dtype)
+        inv = 1 - m
+
+        return replace(
+            latent_state,
+            clean_latent=latent_state.clean_latent * inv + tokens * m,
+            denoise_mask=latent_state.denoise_mask * inv + (1.0 - self.strength) * m,
+        )

packages/ltx-core/src/ltx_core/conditioning/types/reference_audio_cond.py

@@ -51,7 +51,7 @@ class AudioConditionByReferenceLatent:
         )
 
         return LatentState(
-            latent=torch.cat([latent_state.latent, tokens], dim=1),
+            latent=torch.cat([latent_state.latent, torch.zeros_like(tokens)], dim=1),
             denoise_mask=torch.cat([latent_state.denoise_mask, denoise_mask], dim=1),
             positions=torch.cat([latent_state.positions, self.positions], dim=2),
             clean_latent=torch.cat([latent_state.clean_latent, tokens], dim=1),

packages/ltx-core/src/ltx_core/conditioning/types/reference_video_cond.py

@@ -14,7 +14,8 @@ class VideoConditionByReferenceLatent(ConditioningItem):
     Conditions video generation on a reference video latent for IC-LoRA inference.
     IC-LoRAs are trained by concatenating reference (control signal) and target tokens,
     learning to attend across both. This class replicates that setup at inference by
-    appending reference tokens to the latent sequence.
+    appending the reference tokens to the sequence as clean latents (with placeholder zeros
+    in the noisy latent).
     IC-LoRAs can be trained with lower-resolution references than the target (e.g., 384px
     reference for 768px output) for efficiency and better generalization. The
     `downscale_factor` scales reference positions to match target coordinates, preserving
@@ -22,9 +23,9 @@ class VideoConditionByReferenceLatent(ConditioningItem):
     (stored in LoRA metadata).
     To add attention masking, wrap with :class:`ConditioningItemAttentionStrengthWrapper`.
     Args:
-        latent: Reference video latents [B, C, F, H, W]
-        downscale_factor: Target/reference resolution ratio (e.g., 2 = half-resolution
-            reference). Spatial positions are scaled by this factor.
+        latent: Reference video latents [B, C, F, H, W].
+        downscale_factor: Target/reference spatial ratio (e.g. 2 = half-res ref).
+        temporal_scale_factor: Target/reference temporal ratio S (e.g. 4 = ref at 1/4 fps).
         strength: Conditioning strength. 1.0 = full (reference kept clean),
             0.0 = none (reference denoised). Default 1.0.
     """
@@ -33,10 +34,12 @@ class VideoConditionByReferenceLatent(ConditioningItem):
         self,
         latent: torch.Tensor,
         downscale_factor: int = 1,
+        temporal_scale_factor: int = 1,
         strength: float = 1.0,
     ):
         self.latent = latent
         self.downscale_factor = downscale_factor
+        self.temporal_scale_factor = temporal_scale_factor
         self.strength = strength
 
     def apply_to(
@@ -44,10 +47,9 @@ class VideoConditionByReferenceLatent(ConditioningItem):
         latent_state: LatentState,
         latent_tools: VideoLatentTools,
     ) -> LatentState:
-        """Append reference video tokens with scaled positions."""
+        """Append reference video tokens with positions translated into the target frame."""
         tokens = latent_tools.patchifier.patchify(self.latent)
 
-        # Compute positions for the reference video's actual dimensions
         latent_coords = latent_tools.patchifier.get_patch_grid_bounds(
             output_shape=VideoLatentShape.from_torch_shape(self.latent.shape),
             device=self.latent.device,
@@ -58,9 +60,18 @@ class VideoConditionByReferenceLatent(ConditioningItem):
             causal_fix=latent_tools.causal_fix,
         )
         positions = positions.to(dtype=torch.float32)
-        positions[:, 0, ...] /= latent_tools.fps
 
-        # Scale spatial positions to match target coordinate space
+        # Place ref tokens on their own time spacing (= target_fps / S).
+        positions[:, 0, ...] /= latent_tools.fps / self.temporal_scale_factor
+
+        # Translate into the target's frame so ref's last patch ends with target's last
+        # patch; clamp the causal patch's negative start back to [0, 1/target_fps).
+        if self.temporal_scale_factor != 1:
+            t_target = latent_state.positions[:, 0, 0:1, 1:2].to(dtype=torch.float32)  # = 1/target_fps
+            positions[:, 0, ...] = torch.clamp(
+                positions[:, 0, ...] - (self.temporal_scale_factor - 1) * t_target,
+                min=0,
+            )
         if self.downscale_factor != 1:
             positions[:, 1, ...] *= self.downscale_factor  # height axis
             positions[:, 2, ...] *= self.downscale_factor  # width axis
@@ -83,7 +94,7 @@ class VideoConditionByReferenceLatent(ConditioningItem):
         )
 
         return LatentState(
-            latent=torch.cat([latent_state.latent, tokens], dim=1),
+            latent=torch.cat([latent_state.latent, torch.zeros_like(tokens)], dim=1),
             denoise_mask=torch.cat([latent_state.denoise_mask, denoise_mask], dim=1),
             positions=torch.cat([latent_state.positions, positions], dim=2),
             clean_latent=torch.cat([latent_state.clean_latent, tokens], dim=1),

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

@@ -1,18 +1,21 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, NamedTuple, Protocol
+from typing import TYPE_CHECKING, Any, NamedTuple, Protocol, TypeVar
 
 import torch
 
 from ltx_core.loader.module_ops import ModuleOps
 from ltx_core.loader.sd_ops import SDOps
-from ltx_core.model.model_protocol import ModelType
 
 if TYPE_CHECKING:
+    from typing_extensions import Self
+
     from ltx_core.loader.fuse_loras import FuseRule
     from ltx_core.loader.registry import Registry
 
+BuiltType = TypeVar("BuiltType", covariant=True)  # noqa: PLC0105
+
 
 # Per-key shape and dtype description for a flat collection of tensors.
 TensorLayout = dict[str, tuple[torch.Size, torch.dtype]]
@@ -50,74 +53,75 @@ class StateDictLoader(Protocol):
         """
         Load metadata from path
         """
+        ...
 
     def load(self, path: str | list[str], sd_ops: SDOps | None = None, device: torch.device | None = None) -> StateDict:
         """
         Load state dict from path or paths (for sharded model storage) and apply sd_ops
         """
+        ...
 
 
-class BuilderProtocol(Protocol[ModelType]):
+class BuilderProtocol(Protocol[BuiltType]):
     """Protocol for model builders that produce a model via ``build()``."""
 
     def build(
-        self, device: torch.device | None = None, dtype: torch.dtype | None = None, **kwargs: object
-    ) -> ModelType: ...
+        self,
+        device: torch.device | None = None,
+        dtype: torch.dtype | None = None,
+        **kwargs: Any,  # noqa: ANN401
+    ) -> BuiltType: ...
 
+    @property
+    def registry(self) -> "Registry": ...
 
-class ModelBuilderProtocol(BuilderProtocol[ModelType], Protocol[ModelType]):
-    """
-    Protocol for building PyTorch models from configuration dictionaries.
-    Implementations must provide:
-    - meta_model: Create a model from configuration dictionary and apply module operations
-    - build: Create and initialize a model from state dictionary and apply dtype transformations
-    """
-
-    model_sd_ops: SDOps | None
-    module_ops: tuple[ModuleOps, ...]
-    loras: tuple["LoraPathStrengthAndSDOps", ...]
-    registry: "Registry"
-
-    def meta_model(self, config: dict, module_ops: list[ModuleOps] | None = None) -> ModelType:
-        """
-        Create a model on the meta device from a configuration dictionary.
-        This decouples model creation from weight loading, allowing the model
-        architecture to be instantiated without allocating memory for parameters.
-        Args:
-            config: Model configuration dictionary.
-            module_ops: Optional list of module operations to apply (e.g., quantization).
-        Returns:
-            Model instance on meta device (no actual memory allocated for parameters).
-        """
-        ...
-
-    def with_sd_ops(self, sd_ops: SDOps | None) -> "ModelBuilderProtocol[ModelType]":
-        """Return a copy of this builder with the given state-dict key remapping ops."""
-        ...
-
-    def with_module_ops(self, module_ops: tuple[ModuleOps, ...]) -> "ModelBuilderProtocol[ModelType]":
-        """Return a copy of this builder with the given module operations (e.g. quantization)."""
-        ...
-
-    def with_loras(self, loras: tuple["LoraPathStrengthAndSDOps", ...]) -> "ModelBuilderProtocol[ModelType]":
-        """Return a copy of this builder with the given LoRAs to fuse at build time."""
-        ...
-
-    def with_registry(self, registry: "Registry") -> "ModelBuilderProtocol[ModelType]":
+    def with_registry(self, registry: "Registry") -> "Self":
         """Return a copy of this builder using the given weight registry for allocation."""
         ...
 
-    def with_lora_load_device(self, device: torch.device) -> "ModelBuilderProtocol[ModelType]":
+
+class ModelBuilderProtocol(BuilderProtocol[BuiltType], Protocol[BuiltType]):
+    """
+    Protocol for building PyTorch models from configuration dictionaries.
+    Implementations must provide:
+    - build: Create and initialize a model from state dictionary and apply dtype transformations
+    """
+
+    @property
+    def model_sd_ops(self) -> SDOps | None: ...
+
+    @property
+    def module_ops(self) -> tuple[ModuleOps, ...]: ...
+
+    @property
+    def loras(self) -> tuple["LoraPathStrengthAndSDOps", ...]: ...
+
+    def with_sd_ops(self, sd_ops: SDOps | None) -> "Self":
+        """Return a copy of this builder with the given state-dict key remapping ops."""
+        ...
+
+    def with_module_ops(self, module_ops: tuple[ModuleOps, ...]) -> "Self":
+        """Return a copy of this builder with the given module operations (e.g. quantization)."""
+        ...
+
+    def with_loras(self, loras: tuple["LoraPathStrengthAndSDOps", ...]) -> "Self":
+        """Return a copy of this builder with the given LoRAs to fuse at build time."""
+        ...
+
+    def with_lora_load_device(self, device: torch.device) -> "Self":
         """Return a copy of this builder that loads LoRA weights onto the given device."""
         ...
 
-    def with_fuse_rule(self, fuse_rule: "FuseRule") -> "ModelBuilderProtocol[ModelType]":
+    def with_fuse_rule(self, fuse_rule: "FuseRule") -> "Self":
         """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:
+        self,
+        device: torch.device | None = None,
+        dtype: torch.dtype | None = None,
+        **kwargs: Any,  # noqa: ANN401
+    ) -> BuiltType:
         """
         Build the model
         Args:
@@ -140,8 +144,7 @@ class LoRAAdaptableProtocol(Protocol):
     - lora: Add a LoRA to the model
     """
 
-    def lora(self, lora_path: str, strength: float) -> "LoRAAdaptableProtocol":
-        pass
+    def lora(self, lora_path: str, strength: float, sd_ops: SDOps) -> "LoRAAdaptableProtocol": ...
 
 
 class LoraPathStrengthAndSDOps(NamedTuple):

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

@@ -24,6 +24,7 @@ class ContentMatching:
 
     prefix: str = ""
     suffix: str = ""
+    contains: str = ""
 
 
 class KeyValueOperationResult(NamedTuple):
@@ -72,10 +73,10 @@ class SDOps:
         new_mapping = (*self.mapping, ContentReplacement(content, replacement))
         return replace(self, mapping=new_mapping)
 
-    def with_matching(self, prefix: str = "", suffix: str = "") -> "SDOps":
-        """Create a new SDOps instance with the specified prefix and suffix matching added to the mapping."""
+    def with_matching(self, prefix: str = "", suffix: str = "", contains: str = "") -> "SDOps":
+        """Create a new SDOps instance with the specified prefix, suffix and contains matching added to the mapping."""
 
-        new_mapping = (*self.mapping, ContentMatching(prefix, suffix))
+        new_mapping = (*self.mapping, ContentMatching(prefix, suffix, contains))
         return replace(self, mapping=new_mapping)
 
     def with_additional_allowed_keys(self, keys: frozenset[str]) -> "SDOps":
@@ -100,7 +101,12 @@ class SDOps:
     def apply_to_key(self, key: str) -> str | None:
         """Apply the mapping to the given name."""
         matchers = [content for content in self.mapping if isinstance(content, ContentMatching)]
-        valid = any(key.startswith(f.prefix) and key.endswith(f.suffix) for f in matchers)
+        valid = any(
+            key.startswith(matcher.prefix)
+            and key.endswith(matcher.suffix)
+            and (not matcher.contains or matcher.contains in key)
+            for matcher in matchers
+        )
         if not valid:
             return None
 

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

@@ -1,6 +1,8 @@
+from __future__ import annotations
+
+import copy
 import logging
-from dataclasses import dataclass, field, replace
-from typing import Generic
+from typing import TYPE_CHECKING, Final, Generic
 
 import torch
 from torch import nn
@@ -21,6 +23,9 @@ from ltx_core.loader.sd_ops import SDOps
 from ltx_core.loader.sft_loader import SafetensorsModelStateDictLoader
 from ltx_core.model.model_protocol import ModelConfigurator, ModelType
 
+if TYPE_CHECKING:
+    from typing_extensions import Self
+
 logger: logging.Logger = logging.getLogger(__name__)
 
 
@@ -78,10 +83,12 @@ def _load_model_weights(
     meta_model.load_state_dict(fused_sd, strict=False, assign=True)
 
 
-@dataclass(frozen=True)
 class SingleGPUModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType], LoRAAdaptableProtocol):
     """
     Builder for PyTorch models residing on a single GPU.
+    The builder is immutable: ``with_*``/``lora`` return modified copies. The
+    ``ModelBuilderProtocol`` surface is exposed via read-only properties backed
+    by private attributes.
     Attributes:
         model_class_configurator: Class responsible for constructing the model from a config dict.
         model_path: Path (or tuple of shard paths) to the model's `.safetensors` checkpoint(s).
@@ -98,54 +105,106 @@ class SingleGPUModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType],
         fuse_rule: Per-policy LoRA merge rule. Defaults to ``bf16_fuse_rule``;
     """
 
-    model_class_configurator: type[ModelConfigurator[ModelType]]
-    model_path: str | tuple[str, ...]
-    model_sd_ops: SDOps | None = None
-    module_ops: tuple[ModuleOps, ...] = field(default_factory=tuple)
-    loras: tuple[LoraPathStrengthAndSDOps, ...] = field(default_factory=tuple)
-    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 __init__(
+        self,
+        model_class_configurator: type[ModelConfigurator[ModelType]],
+        model_path: str | tuple[str, ...],
+        model_sd_ops: SDOps | None = None,
+        module_ops: tuple[ModuleOps, ...] = (),
+        loras: tuple[LoraPathStrengthAndSDOps, ...] = (),
+        model_loader: StateDictLoader | None = None,
+        registry: Registry | None = None,
+        lora_load_device: torch.device | None = None,
+        fuse_rule: FuseRule = bf16_fuse_rule,
+    ) -> None:
+        # Read-only: typed with the covariant ModelType, so it must not be a mutable attribute.
+        self._model_class_configurator: Final = model_class_configurator
+        self._model_path = model_path
+        self._model_sd_ops = model_sd_ops
+        self._module_ops = module_ops
+        self._loras = loras
+        self._model_loader = model_loader if model_loader is not None else SafetensorsModelStateDictLoader()
+        self._registry = registry if registry is not None else DummyRegistry()
+        self._lora_load_device = lora_load_device if lora_load_device is not None else torch.device("cpu")
+        self._fuse_rule = 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)))
+    @property
+    def model_sd_ops(self) -> SDOps | None:
+        return self._model_sd_ops
 
-    def with_sd_ops(self, sd_ops: SDOps | None) -> "SingleGPUModelBuilder":
-        return replace(self, model_sd_ops=sd_ops)
+    @property
+    def module_ops(self) -> tuple[ModuleOps, ...]:
+        return self._module_ops
 
-    def with_module_ops(self, module_ops: tuple[ModuleOps, ...]) -> "SingleGPUModelBuilder":
-        return replace(self, module_ops=module_ops)
+    @property
+    def loras(self) -> tuple[LoraPathStrengthAndSDOps, ...]:
+        return self._loras
 
-    def with_loras(self, loras: tuple[LoraPathStrengthAndSDOps, ...]) -> "SingleGPUModelBuilder":
-        return replace(self, loras=loras)
+    @property
+    def registry(self) -> Registry:
+        return self._registry
 
-    def with_registry(self, registry: Registry) -> "SingleGPUModelBuilder":
-        return replace(self, registry=registry)
+    @property
+    def model_path(self) -> str | tuple[str, ...]:
+        return self._model_path
 
-    def with_lora_load_device(self, device: torch.device) -> "SingleGPUModelBuilder":
-        return replace(self, lora_load_device=device)
+    @property
+    def model_loader(self) -> StateDictLoader:
+        return self._model_loader
 
-    def with_fuse_rule(self, fuse_rule: FuseRule) -> "SingleGPUModelBuilder":
-        return replace(self, fuse_rule=fuse_rule)
+    @property
+    def lora_load_device(self) -> torch.device:
+        return self._lora_load_device
+
+    @property
+    def fuse_rule(self) -> FuseRule:
+        return self._fuse_rule
+
+    def lora(self, lora_path: str, strength: float, sd_ops: SDOps) -> Self:
+        clone = copy.copy(self)
+        clone._loras = (*self._loras, LoraPathStrengthAndSDOps(lora_path, strength, sd_ops))
+        return clone
+
+    def with_sd_ops(self, sd_ops: SDOps | None) -> Self:
+        clone = copy.copy(self)
+        clone._model_sd_ops = sd_ops
+        return clone
+
+    def with_module_ops(self, module_ops: tuple[ModuleOps, ...]) -> Self:
+        clone = copy.copy(self)
+        clone._module_ops = module_ops
+        return clone
+
+    def with_loras(self, loras: tuple[LoraPathStrengthAndSDOps, ...]) -> Self:
+        clone = copy.copy(self)
+        clone._loras = loras
+        return clone
+
+    def with_registry(self, registry: Registry) -> Self:
+        clone = copy.copy(self)
+        clone._registry = registry
+        return clone
+
+    def with_lora_load_device(self, device: torch.device) -> Self:
+        clone = copy.copy(self)
+        clone._lora_load_device = device
+        return clone
+
+    def with_fuse_rule(self, fuse_rule: FuseRule) -> Self:
+        clone = copy.copy(self)
+        clone._fuse_rule = fuse_rule
+        return clone
 
     def model_config(self) -> dict:
-        return read_model_config(self.model_path, self.model_loader)
+        return read_model_config(self._model_path, self._model_loader)
 
     def meta_model(self, config: dict, module_ops: tuple[ModuleOps, ...]) -> ModelType:
-        return create_meta_model(self.model_class_configurator, config, module_ops)
+        return create_meta_model(self._model_class_configurator, config, module_ops)
 
     def load_sd(
         self, paths: list[str], registry: Registry, device: torch.device | None, sd_ops: SDOps | None = None
     ) -> StateDict:
-        return load_state_dict(paths, self.model_loader, registry, device, sd_ops)
-
-    def _return_model(self, meta_model: ModelType, device: torch.device) -> ModelType:
-        uninitialized = _check_uninitialized(meta_model)
-        if uninitialized:
-            logger.warning(f"Uninitialized parameters or buffers: {uninitialized}")
-            return meta_model
-        return meta_model.to(device)
+        return load_state_dict(paths, self._model_loader, registry, device, sd_ops)
 
     def build(
         self,
@@ -155,18 +214,23 @@ class SingleGPUModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType],
     ) -> ModelType:
         device = torch.device("cuda") if device is None else device
         config = self.model_config()
-        meta_model = self.meta_model(config, self.module_ops)
+        meta_model = self.meta_model(config, self._module_ops)
 
         _load_model_weights(
             meta_model=meta_model,
-            model_path=self.model_path,
-            loras=self.loras,
-            loader=self.model_loader,
-            registry=self.registry,
+            model_path=self._model_path,
+            loras=self._loras,
+            loader=self._model_loader,
+            registry=self._registry,
             device=device,
             dtype=dtype,
-            model_sd_ops=self.model_sd_ops,
-            lora_load_device=self.lora_load_device,
-            fuse_rule=self.fuse_rule,
+            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)
+
+        uninitialized = _check_uninitialized(meta_model)
+        if uninitialized:
+            logger.warning(f"Uninitialized parameters or buffers: {uninitialized}")
+            return meta_model
+        return meta_model.to(device)

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

@@ -1,6 +1,14 @@
-from typing import Protocol, TypeVar
+from __future__ import annotations
 
-ModelType = TypeVar("ModelType")
+from typing import TYPE_CHECKING, Protocol, TypeVar
+
+import torch
+
+if TYPE_CHECKING:
+    from ltx_core.guidance.perturbations import BatchedPerturbationConfig
+    from ltx_core.model.transformer.modality import Modality
+
+ModelType = TypeVar("ModelType", covariant=True, bound=torch.nn.Module)  # noqa: PLC0105
 
 
 class ModelConfigurator(Protocol[ModelType]):
@@ -8,3 +16,24 @@ class ModelConfigurator(Protocol[ModelType]):
 
     @classmethod
     def from_config(cls, config: dict) -> ModelType: ...
+
+
+class LTXModelProtocol(Protocol):
+    """Velocity-model forward interface shared by ``LTXModel`` and its multi-GPU wrappers.
+    ``forward`` pins the real signature (enforced structurally); ``__call__`` mirrors it
+    so protocol-typed values stay callable via ``model(...)``.
+    """
+
+    def forward(
+        self,
+        video: Modality | None,
+        audio: Modality | None,
+        perturbations: BatchedPerturbationConfig,
+    ) -> tuple[torch.Tensor | None, torch.Tensor | None]: ...
+
+    def __call__(
+        self,
+        video: Modality | None,
+        audio: Modality | None,
+        perturbations: BatchedPerturbationConfig,
+    ) -> tuple[torch.Tensor | None, torch.Tensor | None]: ...

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

@@ -3,13 +3,17 @@
 from ltx_core.model.transformer.modality import Modality
 from ltx_core.model.transformer.model import LTXModel, X0Model
 from ltx_core.model.transformer.model_configurator import (
+    LTXV_AUDIO_ONLY_MODEL_COMFY_RENAMING_MAP,
     LTXV_MODEL_COMFY_RENAMING_MAP,
+    LTXAudioOnlyModelConfigurator,
     LTXModelConfigurator,
     LTXVideoOnlyModelConfigurator,
 )
 
 __all__ = [
+    "LTXV_AUDIO_ONLY_MODEL_COMFY_RENAMING_MAP",
     "LTXV_MODEL_COMFY_RENAMING_MAP",
+    "LTXAudioOnlyModelConfigurator",
     "LTXModel",
     "LTXModelConfigurator",
     "LTXVideoOnlyModelConfigurator",

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

@@ -1,5 +1,4 @@
 import functools
-import logging
 from dataclasses import dataclass, field
 from enum import Enum
 from typing import Protocol
@@ -15,8 +14,6 @@ from ltx_core.model.transformer.ops import (
 )
 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.
@@ -76,9 +73,9 @@ class PytorchAttention(AttentionCallable):
 
     @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."""
+        """Human-readable identifier for this backend. 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__(
@@ -199,10 +196,9 @@ class FlashAttention4(AttentionCallable):
 
 # --- 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).
+# usable callable for each path. The selection runs once per process (cached).
+# 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:
@@ -289,20 +285,23 @@ def _select_masked_attention() -> MaskedAttentionCallable:
 
 @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
+    """Cached AUTOMATIC pick for the unmasked path.
+    Cached so every ``AttentionOps`` in the process shares one instance."""
+    return _select_primary_attention()
 
 
 @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
+    """Cached AUTOMATIC pick for the masked path. See :func:`automatic_attention`."""
+    return _select_masked_attention()
+
+
+def attention_label(fn: AttentionCallable | MaskedAttentionCallable) -> str:
+    """Best-effort human-readable backend name.
+    Built-in callables expose ``.label`` (encoding the SDPA priority list for the
+    Pytorch backends); fall back to the class name for custom or wrapped callables
+    (e.g. the multi-GPU All2All wrappers) that don't define one."""
+    return getattr(fn, "label", type(fn).__name__)
 
 
 def _resolve_sdpa_variant(backend: SDPBackend, name: str, *, with_mask: bool) -> PytorchAttention:
@@ -342,8 +341,7 @@ class AttentionFunction(Enum):
         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.
+        cached :func:`automatic_attention` instance so every build shares one callable.
         """
         match self:
             case AttentionFunction.AUTOMATIC:

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

@@ -11,7 +11,7 @@ from ltx_core.model.transformer.transformer_args import BlockPerturbationsProces
 
 # 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_INDUCTOR_CONFIG: dict[str, Any] = {}
 _DEFAULT_DYNAMO_CONFIG: dict[str, Any] = {"inline_inbuilt_nn_modules": True, "cache_size_limit": 256}
 
 

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

@@ -1,9 +1,12 @@
+import logging
 from enum import Enum
 
 import torch
 
 from ltx_core.guidance.perturbations import BatchedPerturbationConfig, PerturbationType
+from ltx_core.model.model_protocol import LTXModelProtocol
 from ltx_core.model.transformer.adaln import AdaLayerNormSingle, adaln_embedding_coefficient
+from ltx_core.model.transformer.attention import attention_label
 from ltx_core.model.transformer.modality import Modality
 from ltx_core.model.transformer.rope import LTXRopeType
 from ltx_core.model.transformer.transformer import (
@@ -20,6 +23,8 @@ from ltx_core.model.transformer.transformer_args import (
 )
 from ltx_core.utils import to_denoised
 
+logger = logging.getLogger(__name__)
+
 
 class LTXModelType(Enum):
     AudioVideo = "ltx av model"
@@ -70,6 +75,15 @@ class LTXModel(torch.nn.Module):
         cross_attention_adaln: bool = False,
     ):
         super().__init__()
+        # Log the attention backends this transformer is built with. Reading the resolved
+        # ``label`` off the ops reports whatever was selected -- AUTOMATIC, an explicit pin
+        # (PYTORCH/XFORMERS/FA3/FA4/SDPA_*), or a directly supplied callable -- so this is the
+        # single source of truth for which kernel a build uses. Fires once per build.
+        logger.info(
+            "Building transformer with attention backends -- self: %s, masked: %s",
+            attention_label(ops.attention_ops.attention_function),
+            attention_label(ops.attention_ops.masked_attention_function),
+        )
         self._enable_gradient_checkpointing = False
         self.cross_attention_adaln = cross_attention_adaln
         self.use_middle_indices_grid = use_middle_indices_grid
@@ -411,7 +425,7 @@ class LTXModel(torch.nn.Module):
 
     def forward(
         self, video: Modality | None, audio: Modality | None, perturbations: BatchedPerturbationConfig
-    ) -> tuple[torch.Tensor, torch.Tensor]:
+    ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
         """
         Forward pass for LTX models.
         Returns:
@@ -459,7 +473,7 @@ class LegacyX0Model(torch.nn.Module):
     Returns fully denoised output based on the velocities produced by the base model.
     """
 
-    def __init__(self, velocity_model: LTXModel):
+    def __init__(self, velocity_model: LTXModelProtocol):
         super().__init__()
         self.velocity_model = velocity_model
 
@@ -488,7 +502,7 @@ class X0Model(torch.nn.Module):
     Applies scaled denoising to the video and audio according to the timesteps = sigma * denoising_mask.
     """
 
-    def __init__(self, velocity_model: LTXModel):
+    def __init__(self, velocity_model: LTXModelProtocol):
         super().__init__()
         self.velocity_model = velocity_model
 

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

@@ -123,6 +123,58 @@ class LTXVideoOnlyModelConfigurator(ModelConfigurator[LTXModel]):
         )
 
 
+class LTXAudioOnlyModelConfigurator(ModelConfigurator[LTXModel]):
+    """
+    Configurator for LTX audio only model.
+    Builds an audio-only LTX model (``model_type=AudioOnly``) so the video
+    transformer weights are never instantiated or loaded. Useful for
+    text-to-audio inference where the video branch is unused.
+    """
+
+    @classmethod
+    def from_config(cls, config: dict, ops: TransformerOpsConfig = DEFAULT_TRANSFORMER_OPS) -> LTXModel:
+        # Build audio caption projection for 19B models (projection handled in transformer).
+        _, audio_caption_projection = _build_caption_projections(config, is_av=True)
+
+        config = config.get("transformer", {})
+
+        check_config_value(config, "dropout", 0.0)
+        check_config_value(config, "attention_bias", True)
+        check_config_value(config, "num_vector_embeds", None)
+        check_config_value(config, "activation_fn", "gelu-approximate")
+        check_config_value(config, "num_embeds_ada_norm", 1000)
+        check_config_value(config, "use_linear_projection", False)
+        check_config_value(config, "only_cross_attention", False)
+        check_config_value(config, "cross_attention_norm", True)
+        check_config_value(config, "double_self_attention", False)
+        check_config_value(config, "upcast_attention", False)
+        check_config_value(config, "standardization_norm", "rms_norm")
+        check_config_value(config, "norm_elementwise_affine", False)
+        check_config_value(config, "qk_norm", "rms_norm")
+        check_config_value(config, "positional_embedding_type", "rope")
+        check_config_value(config, "use_middle_indices_grid", True)
+
+        return LTXModel(
+            model_type=LTXModelType.AudioOnly,
+            num_layers=config.get("num_layers", 48),
+            norm_eps=config.get("norm_eps", 1e-06),
+            ops=ops,
+            timestep_scale_multiplier=config.get("timestep_scale_multiplier", 1000),
+            use_middle_indices_grid=config.get("use_middle_indices_grid", True),
+            audio_num_attention_heads=config.get("audio_num_attention_heads", 32),
+            audio_attention_head_dim=config.get("audio_attention_head_dim", 64),
+            audio_in_channels=config.get("audio_in_channels", 128),
+            audio_out_channels=config.get("audio_out_channels", 128),
+            audio_cross_attention_dim=config.get("audio_cross_attention_dim", 2048),
+            audio_positional_embedding_max_pos=config.get("audio_positional_embedding_max_pos", [20]),
+            rope_type=LTXRopeType(config.get("rope_type", "split")),
+            double_precision_rope=config.get("frequencies_precision", False) == "float64",
+            apply_gated_attention=config.get("apply_gated_attention", False),
+            audio_caption_projection=audio_caption_projection,
+            cross_attention_adaln=config.get("cross_attention_adaln", False),
+        )
+
+
 def _build_caption_projections(
     config: dict,
     is_av: bool,
@@ -151,3 +203,16 @@ LTXV_MODEL_COMFY_RENAMING_MAP = (
     .with_matching(prefix="model.diffusion_model.")
     .with_replacement("model.diffusion_model.", "")
 )
+
+LTXV_AUDIO_ONLY_MODEL_COMFY_RENAMING_MAP = (
+    SDOps("LTXV_AUDIO_ONLY_MODEL_COMFY_MAP")
+    .with_matching(prefix="model.diffusion_model.", contains="audio_attn1")
+    .with_matching(prefix="model.diffusion_model.", contains="audio_attn2")
+    .with_matching(prefix="model.diffusion_model.", contains="audio_ff")
+    .with_matching(prefix="model.diffusion_model.", contains="audio_patchify")
+    .with_matching(prefix="model.diffusion_model.", contains="audio_proj_out")
+    .with_matching(prefix="model.diffusion_model.", contains="audio_adaln_single")
+    .with_matching(prefix="model.diffusion_model.", contains="audio_prompt")
+    .with_matching(prefix="model.diffusion_model.", contains="audio_scale_shift_table")
+    .with_replacement("model.diffusion_model.", "")
+)

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

@@ -24,7 +24,6 @@ from ltx_core.model.transformer.ops import (
 )
 from ltx_core.model.transformer.rope import LTXRopeType
 from ltx_core.model.transformer.transformer_args import TransformerArgs
-from ltx_core.utils import rms_norm
 
 
 @dataclass
@@ -222,7 +221,7 @@ class BasicAVTransformerBlock(torch.nn.Module):
 
     def _apply_text_cross_attention(
         self,
-        x: torch.Tensor,
+        x_normed: torch.Tensor,
         context: torch.Tensor,
         attn: AttentionCallable,
         scale_shift_table: torch.Tensor,
@@ -232,11 +231,14 @@ class BasicAVTransformerBlock(torch.nn.Module):
         context_mask: torch.Tensor | None,
         cross_attention_adaln: bool = False,
     ) -> torch.Tensor:
-        """Apply text cross-attention, with optional AdaLN modulation."""
+        """Apply text cross-attention, with optional AdaLN modulation.
+        ``x_normed`` is the RMS-normalized self-attention output produced by
+        ``post_sa_function`` -- this method does not normalize again.
+        """
         if cross_attention_adaln:
-            shift_q, scale_q, gate = self.get_ada_values(scale_shift_table, x.shape[0], timestep, slice(6, 9))
+            shift_q, scale_q, gate = self.get_ada_values(scale_shift_table, x_normed.shape[0], timestep, slice(6, 9))
             return apply_cross_attention_adaln(
-                x,
+                x_normed,
                 context,
                 attn,
                 shift_q,
@@ -245,9 +247,8 @@ class BasicAVTransformerBlock(torch.nn.Module):
                 prompt_scale_shift_table,
                 prompt_timestep,
                 context_mask,
-                self.norm_eps,
             )
-        return attn(rms_norm(x, eps=self.norm_eps), context=context, mask=context_mask)
+        return attn(x_normed, context=context, mask=context_mask)
 
     def forward(  # noqa: PLR0915
         self,
@@ -280,10 +281,10 @@ class BasicAVTransformerBlock(torch.nn.Module):
                 perturbation_mask=video.self_attn_perturbation_mask,
                 all_perturbed=video.self_attn_all_perturbed,
             )
-            vx = vx + vx_msa_out * vgate_msa
+            vx, vx_normed = self.post_sa_function(vx, vx_msa_out, None, self.norm_eps, vgate_msa)
             del vgate_msa, norm_vx, vx_msa_out
             vx = vx + self._apply_text_cross_attention(
-                vx,
+                vx_normed,
                 video.context,
                 self.attn2,
                 self.scale_shift_table,
@@ -293,6 +294,7 @@ class BasicAVTransformerBlock(torch.nn.Module):
                 video.context_mask,
                 cross_attention_adaln=self.cross_attention_adaln,
             )
+            del vx_normed
 
         if run_ax:
             ashift_msa, ascale_msa, agate_msa = self.get_ada_values(
@@ -308,10 +310,10 @@ class BasicAVTransformerBlock(torch.nn.Module):
                 perturbation_mask=audio.self_attn_perturbation_mask,
                 all_perturbed=audio.self_attn_all_perturbed,
             )
-            ax = ax + ax_msa_out * agate_msa
+            ax, ax_normed = self.post_sa_function(ax, ax_msa_out, None, self.norm_eps, agate_msa)
             del agate_msa, norm_ax, ax_msa_out
             ax = ax + self._apply_text_cross_attention(
-                ax,
+                ax_normed,
                 audio.context,
                 self.audio_attn2,
                 self.audio_scale_shift_table,
@@ -321,6 +323,7 @@ class BasicAVTransformerBlock(torch.nn.Module):
                 audio.context_mask,
                 cross_attention_adaln=self.cross_attention_adaln,
             )
+            del ax_normed
 
         # Audio - Video cross attention.
         if run_a2v or run_v2a:
@@ -414,7 +417,7 @@ class BasicAVTransformerBlock(torch.nn.Module):
 
 
 def apply_cross_attention_adaln(
-    x: torch.Tensor,
+    x_normed: torch.Tensor,
     context: torch.Tensor,
     attn: AttentionCallable,
     q_shift: torch.Tensor,
@@ -423,13 +426,17 @@ def apply_cross_attention_adaln(
     prompt_scale_shift_table: torch.Tensor,
     prompt_timestep: torch.Tensor,
     context_mask: torch.Tensor | None = None,
-    norm_eps: float = 1e-6,
 ) -> torch.Tensor:
-    batch_size = x.shape[0]
+    """Apply query/key AdaLN modulation then cross-attention.
+    ``x_normed`` is already RMS-normalized by ``post_sa_function``; this only
+    applies the affine (scale/shift) modulation, so the normalization is not
+    repeated here.
+    """
+    batch_size = x_normed.shape[0]
     shift_kv, scale_kv = (
-        prompt_scale_shift_table[None, None].to(device=x.device, dtype=x.dtype)
+        prompt_scale_shift_table[None, None].to(device=x_normed.device, dtype=x_normed.dtype)
         + prompt_timestep.reshape(batch_size, prompt_timestep.shape[1], 2, -1)
     ).unbind(dim=2)
-    attn_input = rms_norm(x, eps=norm_eps) * (1 + q_scale) + q_shift
+    attn_input = x_normed * (1 + q_scale) + q_shift
     encoder_hidden_states = context * (1 + scale_kv) + shift_kv
     return attn(attn_input, context=encoder_hidden_states, mask=context_mask) * q_gate

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

@@ -30,21 +30,31 @@ class GemmaTextEncoder(torch.nn.Module):
 
     def encode(
         self,
-        text: str,
+        prompts: list[str],
         padding_side: str = "left",  # noqa: ARG002
-    ) -> tuple[tuple[torch.Tensor, ...], torch.Tensor]:
-        """Run Gemma LLM and return raw hidden states + attention mask.
-        Calls the inner model (self.model.model) to skip lm_head logits computation (~500 MiB saving).
-        Returns:
-            (hidden_states, attention_mask) where hidden_states is a tuple of per-layer tensors.
+    ) -> list[tuple[tuple[torch.Tensor, ...], torch.Tensor]]:
+        """Run a single fused Gemma forward over a batch of prompts.
+        Calls the inner model (self.model.model) to skip lm_head logits computation
+        (~500 MiB saving). The tokenizer pads every prompt to ``max_length`` (1024),
+        so the inputs stack into a single ``[N, 1024]`` batch with no further padding
+        logic; per-prompt outputs are sliced back to ``[1, 1024, D]`` / ``[1, 1024]``
+        in the original order.
         """
-        token_pairs = self.tokenizer.tokenize_with_weights(text)["gemma"]
-        input_ids = torch.tensor([[t[0] for t in token_pairs]], device=self.model.device)
-        attention_mask = torch.tensor([[w[1] for w in token_pairs]], device=self.model.device)
+        if not prompts:
+            return []
+        tokenized = [self.tokenizer.tokenize_with_weights(t)["gemma"] for t in prompts]
+        input_ids = torch.tensor(
+            [[tok for tok, _ in pairs] for pairs in tokenized],
+            device=self.model.device,
+        )
+        attention_mask = torch.tensor(
+            [[w for _, w in pairs] for pairs in tokenized],
+            device=self.model.device,
+        )
         outputs = self.model.model(input_ids=input_ids, attention_mask=attention_mask, output_hidden_states=True)
         hidden_states = outputs.hidden_states
         del outputs
-        return hidden_states, attention_mask
+        return [(tuple(h[i : i + 1] for h in hidden_states), attention_mask[i : i + 1]) for i in range(len(prompts))]
 
     # --- Prompt enhancement methods ---