Automated PR - 2026-07-07

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

@@ -25,8 +25,12 @@ from ltx_core.model.transformer.modality import Modality
 
 def _split_perturbations(config: BatchedPerturbationConfig, sizes: list[int]) -> list[BatchedPerturbationConfig]:
     """Split a ``BatchedPerturbationConfig`` along the batch dimension."""
-    it = iter(config.perturbations)
-    return [BatchedPerturbationConfig([next(it) for _ in range(s)]) for s in sizes]
+    chunks = []
+    offset = 0
+    for size in sizes:
+        chunks.append(config.batch_slice(offset, offset + size))
+        offset += size
+    return chunks
 
 
 def _merge_tensors(tensors: list[torch.Tensor | None]) -> torch.Tensor | None:
@@ -54,6 +58,10 @@ class BatchSplitAdapter(nn.Module):
         self._model = model
         self._max_batch_size = max_batch_size
 
+    @property
+    def num_blocks(self) -> int:
+        return self._model.num_blocks
+
     def _get_chunk_sizes(self, batch_size: int) -> list[int]:
         full, remainder = divmod(batch_size, self._max_batch_size)
         sizes = [self._max_batch_size] * full
@@ -65,7 +73,7 @@ class BatchSplitAdapter(nn.Module):
         self,
         video: Modality | None,
         audio: Modality | None,
-        perturbations: BatchedPerturbationConfig,
+        perturbations: BatchedPerturbationConfig | None,
     ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
         batch_size = (video or audio).latent.shape[0]
 
@@ -77,7 +85,9 @@ class BatchSplitAdapter(nn.Module):
 
         v_chunks = video.split(sizes) if video is not None else [None] * n
         a_chunks = audio.split(sizes) if audio is not None else [None] * n
-        p_chunks = _split_perturbations(perturbations, sizes)
+        # A None config means "perturb nothing"; forward it per chunk so the inner model
+        # builds a per-chunk all-keep mask (splitting None has nothing to slice).
+        p_chunks = _split_perturbations(perturbations, sizes) if perturbations is not None else [None] * n
 
         chunk_results = [
             self._model(video=vc, audio=ac, perturbations=pc)

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

@@ -17,6 +17,7 @@ from ltx_core.block_streaming.disk import DiskBlockReader, DiskTensorReader, Lor
 from ltx_core.block_streaming.pool import BufferPool
 from ltx_core.block_streaming.provider import WeightsProvider
 from ltx_core.block_streaming.source import DiskWeightSource, PinnedBlock, PinnedWeightSource, WeightSource
+from ltx_core.block_streaming.stream_sync import create_stream_sync
 from ltx_core.block_streaming.utils import (
     carve_buffer,
     derive_layout,
@@ -25,6 +26,7 @@ from ltx_core.block_streaming.utils import (
     resolve_attr,
 )
 from ltx_core.block_streaming.wrapper import BlockStreamingWrapper
+from ltx_core.devices import synchronize_device
 from ltx_core.loader.fuse_loras import FuseRule, bf16_fuse_rule, fuse_lora_weights
 from ltx_core.loader.helpers import create_meta_model, load_state_dict, read_model_config
 from ltx_core.loader.module_ops import ModuleOps
@@ -72,9 +74,14 @@ class StreamingModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType])
             ``"transformer_blocks"``).
         blocks_prefix: State-dict key prefix for block weights
             (e.g. ``"transformer_blocks"``).
+        cpu_slots_count: Default number of pinned CPU buffer slots used by
+            :meth:`build` when it is not given an explicit ``cpu_slots_count``.
+            ``None`` = RAM streaming (all blocks pinned); a small value (e.g.
+            ``DISK_CPU_SLOTS``) selects disk streaming. Lets a builder fully
+            encode its offload behaviour so callers need not re-specify it.
     """
 
-    def __init__(
+    def __init__(  # noqa: PLR0913
         self,
         model_class_configurator: type[ModelConfigurator[ModelType]],
         model_path: str | tuple[str, ...],
@@ -86,6 +93,7 @@ class StreamingModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType])
         fuse_rule: FuseRule = bf16_fuse_rule,
         blocks_attr: str = "",
         blocks_prefix: str = "",
+        cpu_slots_count: int | None = None,
     ) -> None:
         # Read-only: typed with the covariant ModelType, so it must not be a mutable attribute.
         self._model_class_configurator: Final = model_class_configurator
@@ -98,6 +106,7 @@ class StreamingModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType])
         self._fuse_rule = fuse_rule
         self._blocks_attr = blocks_attr
         self._blocks_prefix = blocks_prefix
+        self._cpu_slots_count = cpu_slots_count
 
     @property
     def model_class_configurator(self) -> type[ModelConfigurator[ModelType]]:
@@ -107,6 +116,10 @@ class StreamingModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType])
     def model_path(self) -> str | tuple[str, ...]:
         return self._model_path
 
+    @property
+    def checkpoint(self) -> str | tuple[str, ...]:
+        return self._model_path
+
     @property
     def model_sd_ops(self) -> SDOps | None:
         return self._model_sd_ops
@@ -139,6 +152,10 @@ class StreamingModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType])
     def blocks_prefix(self) -> str:
         return self._blocks_prefix
 
+    @property
+    def cpu_slots_count(self) -> int | None:
+        return self._cpu_slots_count
+
     def with_sd_ops(self, sd_ops: SDOps | None) -> Self:
         clone = copy.copy(self)
         clone._model_sd_ops = sd_ops
@@ -188,8 +205,10 @@ class StreamingModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType])
         Args:
             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).
+            cpu_slots_count: Number of pinned CPU buffer slots. ``None`` falls
+                back to the builder's configured ``cpu_slots_count``, and if that
+                is also ``None``, to RAM streaming (all blocks pre-loaded with
+                LoRA fusion).
             gpu_slots_count: Number of GPU buffer slots.
                 ``None`` = ``_DEFAULT_GPU_SLOTS`` (2).
         """
@@ -216,6 +235,7 @@ class StreamingModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType])
                 f"missing indices {missing}, unexpected indices {extra}"
             )
 
+        cpu_slots_count = cpu_slots_count if cpu_slots_count is not None else self._cpu_slots_count
         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
 
@@ -234,16 +254,11 @@ class StreamingModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType])
 
         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,
-            device,
-            reuse_barrier=lambda event: copy_stream.wait_event(event),
-        )
+        sync = create_stream_sync(device)
+        gpu_pool = BufferPool(source.slot_nbytes, gpu_slots_count, device, reuse_barrier=sync.reuse_barrier)
         provider = WeightsProvider(
             gpu_pool,
-            copy_stream,
+            sync,
             device,
             source,
             lora_sources,
@@ -327,7 +342,7 @@ class StreamingModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType])
             block_sd.sd[key] = None
             should_sync = True
         if should_sync:
-            torch.cuda.synchronize()
+            synchronize_device()
 
         # Fill remaining pinned keys from the source state dict.
         for key, view in fill_views.items():

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

@@ -153,9 +153,11 @@ class LoraSource:
         if pair is None:
             return None
         a, b = pair
-        if device is not None and device.type == "cuda":
-            a = a.to(device=device, non_blocking=True)
-            b = b.to(device=device, non_blocking=True)
+        # Move A/B to a GPU-class target (CUDA/MPS) so the B@A aggregation runs on
+        # the device; on a CPU target they stay put. non_blocking only helps CUDA.
+        if device is not None and device.type in ("cuda", "mps"):
+            a = a.to(device=device, non_blocking=device.type == "cuda")
+            b = b.to(device=device, non_blocking=device.type == "cuda")
         if dtype is not None:
             a = a.to(dtype=dtype)
             b = b.to(dtype=dtype)

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

@@ -8,6 +8,7 @@ from typing import Callable
 import torch
 
 from ltx_core.block_streaming import utils
+from ltx_core.block_streaming.stream_sync import StreamEvent
 
 
 class BufferPool:
@@ -27,13 +28,13 @@ class BufferPool:
         slot_nbytes: int,
         capacity: int,
         device: torch.device,
-        reuse_barrier: Callable[[torch.cuda.Event], None],
+        reuse_barrier: Callable[[StreamEvent], None],
         pin_memory: bool = False,
     ) -> None:
         self._slot_nbytes = slot_nbytes
         self._capacity = capacity
         self._free: deque[torch.Tensor] = deque()
-        self._events: dict[int, torch.cuda.Event] = {}
+        self._events: dict[int, StreamEvent] = {}
         self._reuse_barrier = reuse_barrier
         buffer = utils.alloc_buffer(max(slot_nbytes * capacity, 1), device, pin_memory)
         for slot in range(capacity):
@@ -59,7 +60,7 @@ class BufferPool:
             self._reuse_barrier(event)
         return buffer
 
-    def release(self, buffer: torch.Tensor, event: torch.cuda.Event | None = None) -> None:
+    def release(self, buffer: torch.Tensor, event: StreamEvent | 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

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

@@ -10,8 +10,9 @@ import torch
 from ltx_core.block_streaming.disk import LoraSource
 from ltx_core.block_streaming.pool import BufferPool
 from ltx_core.block_streaming.source import WeightSource
+from ltx_core.block_streaming.stream_sync import StreamEvent, StreamSync
 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.fuse_loras import FuseRule, aggregate_lora_products, bf16_fuse_rule, device_fuse_rule
 from ltx_core.loader.primitives import StateDict
 
 _EMPTY_STATE_DICT = StateDict(sd={}, device=torch.device("cpu"), size=0, dtype=set())
@@ -31,8 +32,9 @@ class WeightsProvider:
     """Provides GPU-ready block weights via H2D copy from a pinned CPU weight source.
     Args:
         pool: Pre-allocated GPU weight buffer pool.
-        copy_stream: Dedicated CUDA stream for async H2D copies.
-        target_device: GPU device for compute.
+        sync: Coordinates copy-vs-compute ordering for the backend
+            (see :class:`StreamSync`).
+        target_device: device for compute.
         source: Pinned CPU weight source.
         lora_sources: LoRA adapters fused on H2D copy.
         blocks_prefix: State-dict prefix for LoRA key matching.
@@ -43,17 +45,17 @@ class WeightsProvider:
     def __init__(
         self,
         pool: BufferPool,
-        copy_stream: torch.cuda.Stream,
+        sync: StreamSync,
         target_device: torch.device,
         source: WeightSource,
         lora_sources: list[LoraSource] | None = None,
         blocks_prefix: str = "",
         fuse_rule: FuseRule = bf16_fuse_rule,
     ) -> None:
-        self._copy_stream = copy_stream
+        self._sync = sync
         self._pool = pool
         self._cache: OrderedDict[int, CachedBlock] = OrderedDict()
-        self._events: dict[int, torch.cuda.Event] = {}
+        self._events: dict[int, StreamEvent | None] = {}
         self._target_device = target_device
         self._source = source
         self._lora_sources = lora_sources or []
@@ -88,36 +90,43 @@ class WeightsProvider:
         gpu_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.
+    ) -> StreamEvent | None:
+        """Copy block weights to the target device and fuse LoRAs.
         *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.
+        in *gpu_weights*.
+        The copy + fusion run under :meth:`StreamSync.copy_scope`, then
+        :meth:`StreamSync.commit_copy` orders the copy before compute and returns
+        a guard event for the source to reuse (the ordering is committed 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):
-            raw[:nbytes].copy_(cpu_buffer[:nbytes], non_blocking=True)
+        with self._sync.copy_scope():
+            raw[:nbytes].copy_(cpu_buffer[:nbytes], non_blocking=self._sync.is_async_copy)
             if self._lora_sources:
                 self._fuse_block_loras(idx, gpu_weights)
-            h2d_event = torch.cuda.Event()
-            h2d_event.record(self._copy_stream)
 
-        torch.cuda.current_stream(self._target_device).wait_event(h2d_event)
-        return h2d_event
+        return self._sync.commit_copy()
 
-    def release(self, idx: int, event: torch.cuda.Event) -> None:
-        """Attach a compute-done event -- waited before this buffer is recycled."""
+    def release(self, idx: int, event: StreamEvent | None) -> None:
+        """Attach a compute-done guard, waited before this buffer is recycled
+        (``None`` when the backend needs no guard)."""
         self._events[idx] = event
 
+    def mark_block_done(self, idx: int) -> None:
+        """Record a compute-done guard for block *idx* and queue it for slot reuse.
+        Called once the block's forward pass has been enqueued, so the buffer is
+        not overwritten by a later copy until this compute completes."""
+        self.release(idx, self._sync.record_compute_done())
+
     def cleanup(self) -> None:
-        """Synchronize streams and release all resources."""
-        self._copy_stream.synchronize()
-        torch.cuda.current_stream(self._target_device).synchronize()
+        """Drain outstanding copy/compute work and release all resources."""
+        self._sync.synchronize()
         self._cache.clear()
         self._events.clear()
         self._source.cleanup()
@@ -128,19 +137,29 @@ class WeightsProvider:
         return len(self._cache)
 
     def _fuse_block_loras(self, idx: int, weights: dict[str, torch.Tensor]) -> None:
-        """Fuse LoRA deltas directly into GPU block weights via ``fuse_rule``."""
-        agg_dtype = self._fuse_rule.aggregation_dtype
+        """Fuse LoRA deltas directly into GPU block weights via ``fuse_rule``.
+        The fusion device+dtype come from :func:`device_fuse_rule`: on MPS it
+        aggregates the ``B@A`` on the GPU in fp32 (fast, and fp32 avoids the
+        bf16-on-MPS unreliability), with the rule casting back to the weight
+        dtype; CUDA/CPU keep the rule's dtype. ``get_ab`` places A/B on the
+        target device for CUDA/MPS, so aggregation and the in-place fuse stay
+        co-located there.
+        """
+        rule = device_fuse_rule(self._target_device, self._fuse_rule)
         for name, tensor in weights.items():
             if not name.endswith(".weight"):
                 continue
             prefix = f"{self._blocks_prefix}.{idx}.{name}".removesuffix(".weight")
             products = (
                 ab
-                for ab in (s.get_ab(prefix, device=self._target_device, dtype=agg_dtype) for s in self._lora_sources)
+                for ab in (
+                    s.get_ab(prefix, device=self._target_device, dtype=rule.aggregation_dtype)
+                    for s in self._lora_sources
+                )
                 if ab is not None
             )
-            deltas = aggregate_lora_products(products, agg_dtype)
+            deltas = aggregate_lora_products(products, rule.aggregation_dtype)
             if deltas is None:
                 continue
-            fused = self._fuse_rule(name, tensor, deltas, _EMPTY_STATE_DICT)
+            fused = rule(name, tensor, deltas, _EMPTY_STATE_DICT)
             tensor.copy_(fused[name])

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

@@ -8,6 +8,7 @@ import torch
 
 from ltx_core.block_streaming.block_fetcher import BlockFetcher, FetchHandle
 from ltx_core.block_streaming.pool import BufferPool
+from ltx_core.block_streaming.stream_sync import StreamEvent
 from ltx_core.block_streaming.utils import carve_buffer, layout_nbytes
 from ltx_core.loader.primitives import TensorLayout
 
@@ -33,7 +34,7 @@ class WeightSource(Protocol):
         """Return one contiguous CPU buffer for block *idx*."""
         ...
 
-    def release(self, idx: int, event: torch.cuda.Event | None) -> None:
+    def release(self, idx: int, event: StreamEvent | None) -> None:
         """Signal that an async operation using these weights is guarded by *event*."""
         ...
 
@@ -108,7 +109,7 @@ class DiskWeightSource(WeightSource):
             self._ensure_scheduled((idx + k) % self._blocks_number)
         return scheduled.raw
 
-    def release(self, idx: int, event: torch.cuda.Event | None) -> None:
+    def release(self, idx: int, event: StreamEvent | None) -> None:
         raw_buffer = self._in_flight.pop(idx)
         self._pool.release(raw_buffer, event=event)
 
@@ -164,7 +165,7 @@ class PinnedWeightSource(WeightSource):
     def get(self, idx: int) -> torch.Tensor:
         return self._blocks[idx].buffer
 
-    def release(self, idx: int, event: torch.cuda.Event | None) -> None:
+    def release(self, idx: int, event: StreamEvent | None) -> None:
         pass
 
     def cleanup(self) -> None:

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

@@ -0,0 +1,174 @@
+"""Copy/compute synchronization for block streaming, abstracted across backends.
+Weight streaming overlaps an H2D weight copy with block compute. The two
+operations must be ordered both ways:
+* copy -> compute: a block must not be read before its weights have landed.
+* compute -> reuse: a GPU buffer slot must not be overwritten by the next
+  copy until the compute that read it has finished.
+On CUDA these are expressed with a dedicated copy stream and cross-stream
+events. MPS exposes no user-facing streams (only ``torch.mps.Event`` on a single
+implicit queue), and CPU is fully synchronous. :class:`StreamSync` hides those
+differences behind one protocol; :func:`create_stream_sync` picks the backend
+implementation. The event types (``torch.cuda.Event`` / ``torch.mps.Event``)
+share the small :class:`StreamEvent` surface the pool and source rely on.
+Kept internal to the streaming module -- nothing else needs stream coordination.
+"""
+
+from __future__ import annotations
+
+import contextlib
+from typing import Protocol, runtime_checkable
+
+import torch
+
+from ltx_core.devices import is_mps_available
+
+
+@runtime_checkable
+class StreamEvent(Protocol):
+    """A device synchronization marker (``torch.cuda.Event`` / ``torch.mps.Event``)."""
+
+    def wait(self) -> None:
+        """Device-side: make subsequently queued work wait for this event."""
+        ...
+
+    def synchronize(self) -> None:
+        """Host-side: block the calling thread until this event completes."""
+        ...
+
+
+class StreamSync(Protocol):
+    """Coordinates the H2D copy against block compute for one streaming model."""
+
+    @property
+    def is_async_copy(self) -> bool:
+        """Whether H2D copies may be enqueued asynchronously."""
+        ...
+
+    def copy_scope(self) -> contextlib.AbstractContextManager[None]:
+        """Context to enqueue the H2D copy under (the copy stream on CUDA)."""
+        ...
+
+    def commit_copy(self) -> StreamEvent | None:
+        """Record a copy-done event and make compute wait on it.
+        Returns the event so the source can guard reuse of the CPU buffer (the
+        disk path host-synchronizes on it), or ``None`` when copies are
+        synchronous and no guard is needed.
+        """
+        ...
+
+    def record_compute_done(self) -> StreamEvent | None:
+        """Record an event marking the end of a block's compute, for slot reuse."""
+        ...
+
+    def reuse_barrier(self, event: StreamEvent | None) -> None:
+        """Before a slot is overwritten by a new copy, wait for *event* (prior compute)."""
+        ...
+
+    def synchronize(self) -> None:
+        """Drain all outstanding copy and compute work."""
+        ...
+
+
+class CudaStreamSync:
+    """CUDA: a dedicated copy stream plus cross-stream events.
+    H2D copies run on ``copy_stream`` so they overlap compute on the default
+    stream; events order the two directions explicitly.
+    """
+
+    def __init__(self, device: torch.device) -> None:
+        self._device = device
+        self._copy_stream = torch.cuda.Stream(device=device)
+
+    @property
+    def is_async_copy(self) -> bool:
+        return True
+
+    def copy_scope(self) -> contextlib.AbstractContextManager[None]:
+        return torch.cuda.stream(self._copy_stream)
+
+    def commit_copy(self) -> StreamEvent:
+        event = torch.cuda.Event()
+        event.record(self._copy_stream)
+        torch.cuda.current_stream(self._device).wait_event(event)
+        return event
+
+    def record_compute_done(self) -> StreamEvent:
+        event = torch.cuda.Event()
+        event.record(torch.cuda.current_stream(self._device))
+        return event
+
+    def reuse_barrier(self, event: StreamEvent | None) -> None:
+        if event is not None:
+            self._copy_stream.wait_event(event)
+
+    def synchronize(self) -> None:
+        self._copy_stream.synchronize()
+        torch.cuda.current_stream(self._device).synchronize()
+
+
+class MpsStreamSync:
+    """MPS: one implicit queue with ``torch.mps.Event`` markers.
+    There is no user-facing copy stream, so copy and compute already serialize
+    on the single default queue. The events make that ordering explicit -- and,
+    crucially, let the buffer pool guard slot reuse on the compute-done event
+    rather than relying on the implicit single-queue ordering. ``Event.wait``
+    enqueues a device-side wait on the default queue (it does not block the
+    host); ``Event.synchronize`` is the host-blocking variant.
+    """
+
+    @property
+    def is_async_copy(self) -> bool:
+        return False
+
+    def copy_scope(self) -> contextlib.AbstractContextManager[None]:
+        return contextlib.nullcontext()
+
+    def commit_copy(self) -> StreamEvent:
+        event = torch.mps.Event()
+        event.record()
+        event.wait()
+        return event
+
+    def record_compute_done(self) -> StreamEvent:
+        event = torch.mps.Event()
+        event.record()
+        return event
+
+    def reuse_barrier(self, event: StreamEvent | None) -> None:
+        if event is not None:
+            event.wait()
+
+    def synchronize(self) -> None:
+        torch.mps.synchronize()
+
+
+class SynchronousStreamSync:
+    """CPU (and any non-accelerator backend): copies are synchronous, no events."""
+
+    @property
+    def is_async_copy(self) -> bool:
+        return False
+
+    def copy_scope(self) -> contextlib.AbstractContextManager[None]:
+        return contextlib.nullcontext()
+
+    def commit_copy(self) -> None:
+        return None
+
+    def record_compute_done(self) -> None:
+        return None
+
+    def reuse_barrier(self, event: StreamEvent | None) -> None:  # noqa: ARG002
+        return None
+
+    def synchronize(self) -> None:
+        return None
+
+
+def create_stream_sync(device: torch.device) -> StreamSync:
+    """Return the :class:`StreamSync` implementation for *device*'s backend."""
+    if device.type == "cuda":
+        return CudaStreamSync(device)
+    if device.type == "mps" and is_mps_available():
+        return MpsStreamSync()
+    return SynchronousStreamSync()

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

@@ -88,12 +88,13 @@ def alloc_buffer(nbytes: int, device: torch.device | None, pin_memory: bool) ->
     """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
-    registration fails. Raises if pinning is requested without a CUDA runtime,
-    since pinning is fundamentally a CUDA driver operation.
+    registration fails. Pinning is fundamentally a CUDA driver operation; when
+    requested without a CUDA runtime (e.g. on MPS/CPU, where H2D copies are
+    synchronous and pinning is meaningless) it degrades to a normal allocation.
     """
+    if pin_memory and not torch.cuda.is_available():
+        pin_memory = False  # pinning is CUDA-only; degrade gracefully off-CUDA
     if pin_memory and (device is None or torch.device(device).type == "cpu"):
-        if not torch.cuda.is_available():
-            raise RuntimeError("pin_memory=True requires CUDA, which is not available")
         buf = _alloc_pinned_exact(nbytes)
         if buf is not None:
             return buf

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

@@ -42,6 +42,10 @@ class BlockStreamingWrapper(nn.Module):
         self._hooks: list[torch.utils.hooks.RemovableHandle] = []
         self._register_hooks()
 
+    @property
+    def num_blocks(self) -> int:
+        return self._model.num_blocks
+
     # ------------------------------------------------------------------
     # Hook registration
     # ------------------------------------------------------------------
@@ -55,10 +59,10 @@ class BlockStreamingWrapper(nn.Module):
             assign_tensor_to_module(block, name, gpu_weights[name])
 
     def _post_hook(self, block_idx: int) -> None:
-        """Record a compute-done event and release the block weights."""
-        compute_done = torch.cuda.Event()
-        compute_done.record(torch.cuda.current_stream(self._target_device))
-        self._provider.release(block_idx, event=compute_done)
+        """Release the block weights once its forward pass has been enqueued.
+        The provider guards the buffer against reuse until this block's compute
+        completes."""
+        self._provider.mark_block_done(block_idx)
 
     def _register_hooks(self) -> None:
         for idx, block in enumerate(self._blocks):

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

@@ -0,0 +1,92 @@
+"""Device abstraction for CUDA, Apple Silicon (MPS), and CPU backends.
+Centralizes backend detection and the handful of APIs that genuinely differ
+across accelerators (synchronization, allocator cache, memory queries, RNG
+state). Selection order is CUDA -> MPS -> CPU.
+CUDA-only optimizations (FlashAttention, Triton blockwise FP8/FP6,
+bitsandbytes, NCCL) are gated at their call sites, not here. MPS in particular
+has no ``float64`` support and no fp8 dtype support; use
+:func:`highest_precision_float` to stay within what the backend can represent.
+"""
+
+from __future__ import annotations
+
+import gc
+import logging
+
+import torch
+
+logger = logging.getLogger(__name__)
+
+DeviceSpec = torch.device | None
+
+
+def is_mps_available() -> bool:
+    """Return whether PyTorch can use the Apple Metal/MPS backend."""
+    mps_backend = getattr(torch.backends, "mps", None)
+    return bool(mps_backend is not None and mps_backend.is_available())
+
+
+def get_preferred_device(local_rank: int | None = None) -> torch.device:
+    """Prefer CUDA, then MPS, then CPU.
+    ``local_rank`` is only meaningful for CUDA multi-process launches. MPS exposes
+    a single logical device in PyTorch, so rank-based indexing is not used there.
+    """
+    if torch.cuda.is_available():
+        index = torch.cuda.current_device() if local_rank is None else local_rank
+        return torch.device("cuda", index)
+    if is_mps_available():
+        return torch.device("mps")
+    return torch.device("cpu")
+
+
+def resolve_device(device: DeviceSpec = None, *, local_rank: int | None = None) -> torch.device:
+    """Return *device*, or the best available accelerator when it is ``None``."""
+    if device is None:
+        return get_preferred_device(local_rank=local_rank)
+    return device
+
+
+def supports_float64(device: DeviceSpec) -> bool:
+    """Return whether *device* can represent ``torch.float64``.
+    MPS has no double-precision support; CUDA and CPU do.
+    """
+    return resolve_device(device).type != "mps"
+
+
+def highest_precision_float(device: DeviceSpec) -> torch.dtype:
+    """Return the widest float the backend supports: ``float64`` on CUDA/CPU,
+    ``float32`` on MPS.
+    Use for numerically sensitive accumulators (e.g. sampler ODE math) that
+    request double precision but must degrade gracefully on MPS.
+    """
+    return torch.float64 if supports_float64(device) else torch.float32
+
+
+def synchronize_device(device: DeviceSpec = None) -> None:
+    """Synchronize CUDA or MPS work if the selected backend supports it."""
+    resolved = resolve_device(device)
+    if resolved.type == "cuda" and torch.cuda.is_available():
+        torch.cuda.synchronize(resolved)
+    elif resolved.type == "mps" and is_mps_available():
+        torch.mps.synchronize()
+
+
+def empty_device_cache(device: DeviceSpec = None) -> None:
+    """Release cached allocator memory for CUDA or MPS."""
+    resolved = resolve_device(device)
+    if resolved.type == "cuda" and torch.cuda.is_available():
+        torch.cuda.empty_cache()
+    elif resolved.type == "mps" and is_mps_available():
+        torch.mps.empty_cache()
+
+
+def cleanup_accelerator_memory(device: DeviceSpec = None) -> None:
+    """Run Python GC and release CUDA/MPS allocator caches."""
+    gc.collect()
+    empty_device_cache(device)
+    synchronize_device(device)
+    try:
+        if hasattr(torch._C, "_host_emptyCache"):
+            torch._C._host_emptyCache()
+    except Exception:
+        logger.warning("Host empty cache cleanup failed; ignoring.", exc_info=True)

packages/ltx-core/src/ltx_core/guidance/perturbations.py

@@ -1,17 +1,19 @@
 from dataclasses import dataclass
-from enum import Enum
+from enum import IntEnum
 
 import torch
 from torch._prims_common import DeviceLikeType
 
 
-class PerturbationType(Enum):
-    """Types of attention perturbations for STG (Spatio-Temporal Guidance)."""
+class PerturbationType(IntEnum):
+    """Types of attention perturbations for STG (Spatio-Temporal Guidance).
+    The integer value is the row index into ``BatchedPerturbationConfig._block_masks`` dim 0.
+    """
 
-    SKIP_A2V_CROSS_ATTN = "skip_a2v_cross_attn"
-    SKIP_V2A_CROSS_ATTN = "skip_v2a_cross_attn"
-    SKIP_VIDEO_SELF_ATTN = "skip_video_self_attn"
-    SKIP_AUDIO_SELF_ATTN = "skip_audio_self_attn"
+    SKIP_VIDEO_SELF_ATTN = 0
+    SKIP_AUDIO_SELF_ATTN = 1
+    SKIP_A2V_CROSS_ATTN = 2
+    SKIP_V2A_CROSS_ATTN = 3
 
 
 @dataclass(frozen=True)
@@ -48,32 +50,84 @@ class PerturbationConfig:
         return PerturbationConfig([])
 
 
-@dataclass(frozen=True)
 class BatchedPerturbationConfig:
-    """Perturbation configurations for a batch, with utilities for generating attention masks."""
+    """Per-block attention keep-masks for a batch, built once from a list of per-sample configs.
+    Construction materializes ``_block_masks`` -- a ``(len(PerturbationType), num_blocks, B)`` tensor
+    (1 = keep, 0 = perturbed) whose dim-0 row index is the ``PerturbationType`` value -- from the
+    perturbation structure. The per-sample config list is NOT retained: every consumer reads the
+    tensor (``mask`` indexes it; ``any_in_batch`` / ``all_in_batch`` read the host mirror).
+    The host build (reading the Python structure) happens here, in ``__init__``, so it MUST be run
+    eagerly OUTSIDE any ``torch.compile`` / CUDA-graph-capture region. The compiled block then reads
+    perturbation purely as the runtime ``_block_masks`` tensor and never recompiles per config.
+    """
 
-    perturbations: list[PerturbationConfig]
+    _block_masks: torch.Tensor  # keep-mask on the compute device, indexed [PerturbationType, block, sample]
+    # Host mirror so any_in_batch / all_in_batch stay sync-free and graph-break-free. Present for
+    # configs that may hit the eager skip shortcuts; None for compiled-only configs built via
+    # ``from_masks`` (the compiled processor reads only ``_block_masks``).
+    _block_masks_cpu: torch.Tensor | None
 
-    def mask(
-        self, perturbation_type: PerturbationType, block: int, device: DeviceLikeType, dtype: torch.dtype
-    ) -> torch.Tensor:
-        mask = torch.ones((len(self.perturbations),), device=device, dtype=dtype)
-        for batch_idx, perturbation in enumerate(self.perturbations):
-            if perturbation.is_perturbed(perturbation_type, block):
-                mask[batch_idx] = 0
+    def __init__(
+        self,
+        perturbations: list[PerturbationConfig],
+        num_blocks: int,
+        device: DeviceLikeType | None = None,
+        dtype: torch.dtype | None = None,
+    ) -> None:
+        keep = [
+            [
+                [not pc.is_perturbed(PerturbationType(direction), block) for pc in perturbations]
+                for block in range(num_blocks)
+            ]
+            for direction in range(len(PerturbationType))
+        ]
+        self._block_masks_cpu = torch.tensor(keep, dtype=dtype, device="cpu")
+        self._block_masks = self._block_masks_cpu if device is None else self._block_masks_cpu.to(device)
 
-        return mask
+    @classmethod
+    def from_masks(
+        cls, block_masks: torch.Tensor, block_masks_cpu: torch.Tensor | None = None
+    ) -> "BatchedPerturbationConfig":
+        """Construct from prebuilt mask tensors (e.g. a batch-dim slice), bypassing the host build.
+        ``block_masks_cpu`` is only consumed by ``any_in_batch`` / ``all_in_batch`` (the eager
+        processor's skip shortcuts); pass it when the result may take that path. The compiled
+        processor reads only ``_block_masks``, so callers on that path may omit the mirror.
+        """
+        obj = cls.__new__(cls)
+        obj._block_masks = block_masks
+        obj._block_masks_cpu = block_masks_cpu
+        return obj
 
-    def mask_like(self, perturbation_type: PerturbationType, block: int, values: torch.Tensor) -> torch.Tensor:
-        mask = self.mask(perturbation_type, block, values.device, values.dtype)
-        return mask.view(mask.numel(), *([1] * len(values.shape[1:])))
+    def batch_slice(self, start: int, end: int) -> "BatchedPerturbationConfig":
+        """A view over samples ``[start:end]`` of the batch, by slicing the mask tensors.
+        Slicing (never rebuilding) keeps the host mask build outside any compiled / capture region.
+        """
+        cpu_mask = self._block_masks_cpu[:, :, start:end] if self._block_masks_cpu is not None else None
+        return BatchedPerturbationConfig.from_masks(self._block_masks[:, :, start:end], cpu_mask)
+
+    def mask(self, perturbation_type: PerturbationType, block: int) -> torch.Tensor:
+        """This block's ``(B, 1, 1)`` keep-mask for one perturbation type, as an OWNED tensor.
+        A ``clone`` (not a view into ``_block_masks``) so the masks attached to a block
+        (e.g. self- and cross-attention) don't alias the same storage -- aliased graph inputs are
+        fragile under ``torch.compile``.
+        """
+        return self._block_masks[perturbation_type, block].reshape(-1, 1, 1).clone()
 
     def any_in_batch(self, perturbation_type: PerturbationType, block: int) -> bool:
-        return any(perturbation.is_perturbed(perturbation_type, block) for perturbation in self.perturbations)
+        assert self._block_masks_cpu is not None, "host mirror required by the skip-shortcut processor path"
+        return bool((self._block_masks_cpu[perturbation_type, block] == 0).any())
 
     def all_in_batch(self, perturbation_type: PerturbationType, block: int) -> bool:
-        return all(perturbation.is_perturbed(perturbation_type, block) for perturbation in self.perturbations)
+        assert self._block_masks_cpu is not None, "host mirror required by the skip-shortcut processor path"
+        return bool((self._block_masks_cpu[perturbation_type, block] == 0).all())
 
     @staticmethod
-    def empty(batch_size: int) -> "BatchedPerturbationConfig":
-        return BatchedPerturbationConfig([PerturbationConfig.empty() for _ in range(batch_size)])
+    def empty(
+        batch_size: int,
+        num_blocks: int,
+        device: DeviceLikeType | None = None,
+        dtype: torch.dtype | None = None,
+    ) -> "BatchedPerturbationConfig":
+        return BatchedPerturbationConfig(
+            [PerturbationConfig.empty() for _ in range(batch_size)], num_blocks, device, dtype
+        )

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

@@ -1,5 +1,5 @@
 from collections.abc import Callable, Iterable, Iterator
-from dataclasses import dataclass
+from dataclasses import dataclass, replace
 from typing import NamedTuple
 
 import torch
@@ -71,7 +71,26 @@ def _bf16_fuse(
 bf16_fuse_rule = FuseRule(aggregation_dtype=torch.bfloat16, fuse_fn=_bf16_fuse)
 
 
-def _get_device() -> torch.device:
+def device_fuse_rule(target_device: torch.device, base_rule: FuseRule) -> FuseRule:
+    """Return the fuse rule to use when fusing onto *target_device*.
+    On MPS, swap the rule's aggregation dtype to fp32: the LoRA ``B@A`` then runs
+    on the GPU (far faster than fusing on CPU) and fp32 sidesteps the bf16-on-MPS
+    numerical unreliability that would otherwise force the slow CPU path. The
+    rule's ``fuse_fn`` still casts the fused result back to the weight dtype.
+    CUDA/CPU keep *base_rule* unchanged.
+    """
+    if target_device.type == "mps":
+        return replace(base_rule, aggregation_dtype=torch.float32)
+    return base_rule
+
+
+def _fusion_device(target_device: torch.device) -> torch.device:
+    """Device to run the fusion on: the target's own accelerator (CUDA/MPS), else
+    CUDA when present (accelerating a CPU-resident fuse), else CPU. The caller
+    moves the fused result back to the weight's device afterwards.
+    """
+    if target_device.type in ("cuda", "mps"):
+        return target_device
     if torch.cuda.is_available():
         return torch.device("cuda", torch.cuda.current_device())
     return torch.device("cpu")
@@ -110,21 +129,22 @@ def fuse_lora_weights(
     used for fusion; caller is responsible for moving them to their final
     destination.
     """
-    fusion_device = _get_device()
+    rule = device_fuse_rule(model_sd.device, fuse_rule)
+    fusion_device = _fusion_device(model_sd.device)
     for key in _affected_weight_keys(lora_sd_and_strengths):
         original_weight = model_sd.sd.get(key)
         if original_weight is None:
             continue
 
-        products = _products_for_sd_key(lora_sd_and_strengths, key, fuse_rule.aggregation_dtype, fusion_device)
-        deltas = aggregate_lora_products(products, fuse_rule.aggregation_dtype)
+        products = _products_for_sd_key(lora_sd_and_strengths, key, rule.aggregation_dtype, fusion_device)
+        deltas = aggregate_lora_products(products, rule.aggregation_dtype)
         if deltas is None:
             continue
 
         original_device = original_weight.device
         weight = original_weight.to(device=fusion_device)
 
-        fused = fuse_rule(key, weight, deltas, model_sd)
+        fused = rule(key, weight, deltas, model_sd)
 
         for k, v in fused.items():
             yield k, v.to(device=original_device) if preserve_input_device else v

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

@@ -87,6 +87,11 @@ class ModelBuilderProtocol(BuilderProtocol[BuiltType], Protocol[BuiltType]):
     - build: Create and initialize a model from state dictionary and apply dtype transformations
     """
 
+    @property
+    def checkpoint(self) -> str | tuple[str, ...]:
+        """Path(s) to the checkpoint this builder loads from (for logging/diagnostics)."""
+        ...
+
     @property
     def model_sd_ops(self) -> SDOps | None: ...
 

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

@@ -148,6 +148,10 @@ class SingleGPUModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType],
     def model_path(self) -> str | tuple[str, ...]:
         return self._model_path
 
+    @property
+    def checkpoint(self) -> str | tuple[str, ...]:
+        return self._model_path
+
     @property
     def model_loader(self) -> StateDictLoader:
         return self._model_loader

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

@@ -93,7 +93,7 @@ class VideoModalityTilingHelper:
             keep_per_tile_cond = self._all_tiles_cond_keep(modality)  # (num_tiles, num_cond) bool
             tile_idx = next((i for i, t in enumerate(self._tiles) if t.in_coords == tile.in_coords), None)
             if tile_idx is None:
-                raise ValueError(
+                raise RuntimeError(
                     f"Tile with in_coords={tile.in_coords} is not in this helper's tile set; "
                     f"pass a tile obtained from `helper.tiles`."
                 )
@@ -164,7 +164,7 @@ class VideoModalityTilingHelper:
 
         if output is not None:
             if output.shape != expected_shape:
-                raise ValueError(f"Expected output shape {expected_shape}, got {output.shape}")
+                raise RuntimeError(f"Expected output shape {expected_shape}, got {output.shape}")
             result = output
         else:
             result = torch.zeros(*expected_shape, device=tile_to_blend.device, dtype=tile_to_blend.dtype)

packages/ltx-core/src/ltx_core/model/audio_vae/vocoder.py

@@ -1,4 +1,6 @@
+import contextlib
 import math
+from collections.abc import Iterator
 from typing import List
 
 import einops
@@ -13,6 +15,25 @@ def get_padding(kernel_size: int, dilation: int = 1) -> int:
     return int((kernel_size * dilation - dilation) / 2)
 
 
+@contextlib.contextmanager
+def _module_in_fp32(module: nn.Module, *, enabled: bool) -> Iterator[None]:
+    """Temporarily cast *module* to float32, restoring its original dtype on exit.
+    Used for the MPS vocoder path where fp32 autocast is unavailable, so the
+    weights must be materialized in float32 for the forward pass. Restores to the
+    module's original weight dtype (captured here), not the input dtype. When
+    *enabled* is False this is a no-op, so callers can wrap unconditionally.
+    """
+    if not enabled:
+        yield
+        return
+    module_dtype = next(module.parameters()).dtype
+    module.float()
+    try:
+        yield
+    finally:
+        module.to(module_dtype)
+
+
 # ---------------------------------------------------------------------------
 # Anti-aliased resampling helpers (kaiser-sinc filters) for BigVGAN v2
 # Adopted from https://github.com/NVIDIA/BigVGAN
@@ -564,15 +585,29 @@ class VocoderWithBWE(nn.Module):
         # compound through 108 sequential convolutions and degrade spectral
         # metrics (mel_l1, MRSTFT) by 40-90% while perceptual quality (CDPAM)
         # is unaffected.  fp32 eliminates this degradation.
-        # We use autocast(dtype=float32) rather than self.float() because it
-        # upcasts bf16 weights per-op at kernel level, avoiding the temporary
-        # memory spike of self.float() / self.to(original_dtype).
+        # On CUDA/CPU we use autocast(dtype=float32) rather than self.float()
+        # because it upcasts bf16 weights per-op at kernel level, avoiding the
+        # temporary memory spike of self.float() / self.to(original_dtype).
         # Benchmarked on H100 (128.5M-param model):
         #   autocast fp32: +70 MB peak VRAM, 123 ms  (vs 482 MB / 95 ms for bf16)
         #   model.float(): +324 MB peak VRAM, 149 ms
         # Tested: both approaches produce bit-identical output.
+        # MPS autocast does not upcast conv weights to fp32 (it only supports
+        # lower-precision autocast dtypes), which would leave the float32 input
+        # running against bf16 conv weights and raise a dtype mismatch. There we
+        # fall back to materializing the weights in fp32 for the pass (bit-identical
+        # per the note above; the memory spike is negligible for this small model).
+        # The vocoder is normally built in fp32 on MPS, so this fallback is then a
+        # no-op -- it only triggers if a bf16 module is run on MPS directly.
+        device_type = mel_spec.device.type
+        module_dtype = next(self.parameters()).dtype
+        fp32_ctx = (
+            _module_in_fp32(self, enabled=module_dtype != torch.float32)
+            if device_type == "mps"
+            else torch.autocast(device_type=device_type, dtype=torch.float32)
+        )
 
-        with torch.autocast(device_type=mel_spec.device.type, dtype=torch.float32):
+        with fp32_ctx:
             x = self.vocoder(mel_spec.float())
             _, _, length_low_rate = x.shape
             output_length = length_low_rate * self.output_sampling_rate // self.input_sampling_rate

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

@@ -24,16 +24,21 @@ class LTXModelProtocol(Protocol):
     so protocol-typed values stay callable via ``model(...)``.
     """
 
+    @property
+    def num_blocks(self) -> int:
+        """Number of transformer blocks, delegated through any wrappers to the ``LTXModel``."""
+        ...
+
     def forward(
         self,
         video: Modality | None,
         audio: Modality | None,
-        perturbations: BatchedPerturbationConfig,
+        perturbations: BatchedPerturbationConfig | None,
     ) -> tuple[torch.Tensor | None, torch.Tensor | None]: ...
 
     def __call__(
         self,
         video: Modality | None,
         audio: Modality | None,
-        perturbations: BatchedPerturbationConfig,
+        perturbations: BatchedPerturbationConfig | None,
     ) -> tuple[torch.Tensor | None, torch.Tensor | None]: ...

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

@@ -1,4 +1,5 @@
 import functools
+import sys
 from dataclasses import dataclass, field
 from enum import Enum
 from typing import Protocol
@@ -27,28 +28,26 @@ def _torch_default_sdpa_priority() -> list[SDPBackend]:
     return [SDPBackend(p) for p in torch._C._get_sdp_priority_order()]
 
 
-memory_efficient_attention = None
 flash_attn_interface = None
 flash_attn_4_func = None
 try:
-    from xformers.ops import memory_efficient_attention
-except ImportError:
-    memory_efficient_attention = None
-try:
-    # FlashAttention3 and XFormersAttention cannot be used together
-    if memory_efficient_attention is None:
-        import flash_attn_interface
+    import flash_attn_interface
 except ImportError:
     flash_attn_interface = None
 try:
     from flash_attn.cute import flash_attn_func as flash_attn_4_func
 except ImportError:
     flash_attn_4_func = None
+try:
+    # macOS only: routes SDPA to Apple's prebuilt MPSGraph attention kernel.
+    from mps_sdpa import sdpa_opt as _mps_sdpa_opt
+except ImportError:
+    _mps_sdpa_opt = None
 
 
 class AttentionCallable(Protocol):
     """Unmasked attention. Backends without a mask kernel (FA3/FA4) implement only
-    this protocol; backends that support masks too (Pytorch/SDPA, xFormers) are
+    this protocol; backends that support masks too (Pytorch/SDPA) are
     structurally usable here and as :class:`MaskedAttentionCallable`."""
 
     def __call__(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, heads: int) -> torch.Tensor: ...
@@ -73,9 +72,8 @@ class PytorchAttention(AttentionCallable):
 
     @property
     def label(self) -> str:
-        """Human-readable identifier for this backend. Encodes the SDPA priority
-        list so a single-backend pin reads differently from the full-priority
-        dispatcher walk."""
+        """Human-readable identifier. 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__(
@@ -101,50 +99,48 @@ class PytorchAttention(AttentionCallable):
         return out
 
 
-class XFormersAttention(AttentionCallable):
-    label = "xFormers"
+class MPSSdpaAttention(AttentionCallable):
+    """Apple-fused scaled-dot-product attention on MPS.
+    Routes to ``mps_sdpa.sdpa_opt``, which calls Apple's prebuilt
+    ``MPSGraph.scaledDotProductAttention`` kernel (via a zero-copy bridge)
+    instead of torch's ``sdpa_general_mps`` graph. The Apple kernel does not
+    materialize the ``[B, H, Nq, Nk]`` score matrix, so it avoids the
+    long-sequence memory wall that makes torch's materializing MPS SDPA
+    unusable on video latents (~32x faster at a 14k-token latent on an M4 Pro).
+    It is a hard dependency on Apple Silicon (the ``mps-sdpa`` platform-marked
+    requirement), so AUTOMATIC always has it on MPS. Unlike a JIT-compiled Metal
+    flash kernel it needs no runtime shader compilation, so it is robust
+    across macOS / Metal revisions.
+    Accepts an optional additive-float or boolean ``mask`` broadcastable to
+    ``[B, H, Nq, Nk]``, so it serves both the unmasked and masked protocols.
+    """
+
+    @property
+    def label(self) -> str:
+        return "MPS-SDPA"
 
     def __call__(
-        self,
-        q: torch.Tensor,
-        k: torch.Tensor,
-        v: torch.Tensor,
-        heads: int,
-        mask: torch.Tensor | None = None,
+        self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, heads: int, mask: torch.Tensor | None = None
     ) -> torch.Tensor:
-        if memory_efficient_attention is None:
-            raise RuntimeError("XFormersAttention was selected but `xformers` is not installed.")
+        if _mps_sdpa_opt is None:
+            raise RuntimeError("MPSSdpaAttention was selected but `mps-sdpa` is not installed.")
+        if q.device.type != "mps":
+            raise RuntimeError("MPSSdpaAttention requires MPS. Use PyTorch SDPA on CPU or CUDA.")
 
         b, _, dim_head = q.shape
         dim_head //= heads
-
-        # xformers expects [B, M, H, K]
-        q, k, v = (t.view(b, -1, heads, dim_head) for t in (q, k, v))
+        q, k, v = (t.view(b, -1, heads, dim_head).transpose(1, 2) for t in (q, k, v))
 
         if mask is not None:
-            # add a singleton batch dimension
+            # add a batch dimension if there isn't already one
             if mask.ndim == 2:
                 mask = mask.unsqueeze(0)
-            # add a singleton heads dimension
+            # add a heads dimension if there isn't already one
             if mask.ndim == 3:
                 mask = mask.unsqueeze(1)
-            # pad to a multiple of 8
-            pad = 8 - mask.shape[-1] % 8
-            # the xformers docs says that it's allowed to have a mask of shape (1, Nq, Nk)
-            # but when using separated heads, the shape has to be (B, H, Nq, Nk)
-            # in flux, this matrix ends up being over 1GB
-            # here, we create a mask with the same batch/head size as the input mask (potentially singleton or full)
-            mask_out = torch.empty(
-                [mask.shape[0], mask.shape[1], q.shape[1], mask.shape[-1] + pad], dtype=q.dtype, device=q.device
-            )
 
-            mask_out[..., : mask.shape[-1]] = mask
-            # doesn't this remove the padding again??
-            mask = mask_out[..., : mask.shape[-1]]
-            mask = mask.expand(b, heads, -1, -1)
-
-        out = memory_efficient_attention(q.to(v.dtype), k.to(v.dtype), v, attn_bias=mask, p=0.0)
-        out = out.reshape(b, -1, heads * dim_head)
+        out = _mps_sdpa_opt(q, k, v, attn_mask=mask)
+        out = out.transpose(1, 2).reshape(b, -1, heads * dim_head)
         return out
 
 
@@ -160,6 +156,8 @@ class FlashAttention3(AttentionCallable):
     ) -> torch.Tensor:
         if flash_attn_interface is None:
             raise RuntimeError("FlashAttention3 was selected but `FlashAttention3` is not installed.")
+        if q.device.type != "cuda":
+            raise RuntimeError("FlashAttention3 requires CUDA. Use PyTorch SDPA on CPU or MPS.")
 
         b, _, dim_head = q.shape
         dim_head //= heads
@@ -183,6 +181,8 @@ class FlashAttention4(AttentionCallable):
     ) -> torch.Tensor:
         if flash_attn_4_func is None:
             raise RuntimeError("FlashAttention4 was selected but `flash-attn-4` is not installed.")
+        if q.device.type != "cuda":
+            raise RuntimeError("FlashAttention4 requires CUDA. Use PyTorch SDPA on CPU or MPS.")
 
         b, _, dim_head = q.shape
         dim_head //= heads
@@ -198,7 +198,7 @@ class FlashAttention4(AttentionCallable):
 # AUTOMATIC inspects installed extras and the GPU arch and returns the fastest
 # 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).
+# may end up on different backends (e.g. FA3 unmasked + SDPA masked on H100).
 
 
 def _sdpa_can_use(backend: SDPBackend, *, with_mask: bool) -> bool:
@@ -236,6 +236,20 @@ _SDPA_FULL_PRIORITY: tuple[SDPBackend, ...] = (
 )
 
 
+def _on_macos() -> bool:
+    """True on macOS, where torch's native SDPA materializes the score matrix and
+    AUTOMATIC routes to Apple's fused ``mps-sdpa`` kernel instead."""
+    return sys.platform == "darwin"
+
+
+def _mps_sdpa_available() -> bool:
+    """True when the ``mps-sdpa`` package is importable. It is a platform-marked
+    hard dependency on Apple Silicon, so this is always True there; it is only
+    False on non-Apple-Silicon macs (e.g. Intel/CPU), where AUTOMATIC falls back
+    to torch's SDPA (acceptable on CPU, which has no MPS memory wall)."""
+    return _mps_sdpa_opt is not None
+
+
 def _sdpa_full_priority() -> PytorchAttention:
     """Hand SDPA the full backend priority order; let torch's dispatcher pick at call time.
     ``sdpa_kernel(_SDPA_FULL_PRIORITY, set_priority=True)`` enables all four
@@ -253,10 +267,14 @@ def _sdpa_full_priority() -> PytorchAttention:
 def _select_primary_attention() -> AttentionCallable:
     """Pick the fastest unmasked attention based on installed extras and GPU arch.
     Priority by arch:
-    - Hopper (sm_90, H100): FA3 / xFormers (mutually exclusive at import) > FA4 > SDPA.
+    - Hopper (sm_90, H100): FA3 > FA4 > SDPA.
     - Datacenter Blackwell (sm_100, B200): FA4 > SDPA. FA4 is intentionally *not*
       picked on consumer Blackwell (sm_120) -- known regressions in newer
       FA4 betas; users who want it on sm_120 must opt in explicitly.
+    - macOS (Apple Silicon / MPS): Apple's fused MPSGraph kernel via ``mps-sdpa``
+      (a platform-marked hard dependency on Apple Silicon) -- it avoids the
+      full-score-matrix memory wall on long video sequences. On a non-Apple-Silicon
+      mac (Intel/CPU) it falls back to torch's SDPA.
     - Everywhere else (Ada, Ampere, CPU): SDPA with the full backend priority
       list -- torch's runtime dispatcher picks the best fit at call time.
     """
@@ -265,21 +283,23 @@ def _select_primary_attention() -> AttentionCallable:
         if major == 9:
             if flash_attn_interface is not None:
                 return FlashAttention3()
-            if memory_efficient_attention is not None:
-                return XFormersAttention()
             if flash_attn_4_func is not None:
                 return FlashAttention4()
         if major == 10 and flash_attn_4_func is not None:
             return FlashAttention4()
+    if _on_macos():
+        return MPSSdpaAttention() if _mps_sdpa_available() else _sdpa_full_priority()
     return _sdpa_full_priority()
 
 
 def _select_masked_attention() -> MaskedAttentionCallable:
-    """Pick a mask-aware attention. Prefers xFormers when installed; else SDPA with
-    the full priority list (the dispatcher rejects FLASH automatically when a
-    mask is present and walks past it)."""
-    if memory_efficient_attention is not None:
-        return XFormersAttention()
+    """Pick a mask-aware attention. On macOS, Apple's fused MPSGraph kernel via
+    ``mps-sdpa`` (a hard dependency on Apple Silicon, else torch's SDPA on
+    Intel/CPU macs); else SDPA with the full priority list (the dispatcher
+    rejects FLASH automatically when a mask is present and walks past it --
+    torch SDPA handles the additive mask directly)."""
+    if _on_macos():
+        return MPSSdpaAttention() if _mps_sdpa_available() else _sdpa_full_priority()
     return _sdpa_full_priority()
 
 
@@ -323,18 +343,21 @@ def _resolve_sdpa_variant(backend: SDPBackend, name: str, *, with_mask: bool) ->
 
 class AttentionFunction(Enum):
     PYTORCH = "pytorch"
-    XFORMERS = "xformers"
     FLASH_ATTENTION_3 = "flash_attention_3"
     FLASH_ATTENTION_4 = "flash_attention_4"
     SDPA_CUDNN = "sdpa_cudnn"
     SDPA_FLASH = "sdpa_flash"
     SDPA_EFFICIENT = "sdpa_efficient"
     SDPA_MATH = "sdpa_math"
+    # Apple's fused MPSGraph SDPA via the `mps-sdpa` package (macOS/MPS only, a
+    # platform-marked hard dependency on Apple Silicon). The AUTOMATIC default on
+    # MPS; never materializes the score matrix.
+    MPS_SDPA = "mps_sdpa"
     # Pick the fastest unmasked backend for the current GPU/extras combo; see
     # :func:`automatic_attention`. Default for :class:`AttentionOps`.
     AUTOMATIC = "automatic"
 
-    def to_callable(self) -> AttentionCallable:  # noqa: PLR0911
+    def to_callable(self) -> AttentionCallable:  # noqa: PLR0911, PLR0912
         """Resolve to a concrete callable. Use this at module init time so that
         torch.compile can trace through the attention call without graph breaks.
         Every non-AUTOMATIC variant raises :class:`RuntimeError` when the backend
@@ -348,24 +371,32 @@ class AttentionFunction(Enum):
                 return automatic_attention()
             case AttentionFunction.PYTORCH:
                 return PytorchAttention()
-            case AttentionFunction.XFORMERS:
-                if memory_efficient_attention is None:
-                    raise RuntimeError("AttentionFunction.XFORMERS selected but `xformers` is not installed.")
-                return XFormersAttention()
             case AttentionFunction.FLASH_ATTENTION_3:
                 if flash_attn_interface is None:
                     raise RuntimeError(
                         "AttentionFunction.FLASH_ATTENTION_3 selected but `flash-attn-3` is not installed."
                     )
+                if not torch.cuda.is_available():
+                    raise RuntimeError(
+                        "AttentionFunction.FLASH_ATTENTION_3 requires CUDA. Use PyTorch SDPA on CPU or MPS."
+                    )
                 return FlashAttention3()
             case AttentionFunction.FLASH_ATTENTION_4:
                 if flash_attn_4_func is None:
                     raise RuntimeError(
                         "AttentionFunction.FLASH_ATTENTION_4 selected but `flash-attn-4` is not installed."
                     )
+                if not torch.cuda.is_available():
+                    raise RuntimeError(
+                        "AttentionFunction.FLASH_ATTENTION_4 requires CUDA. Use PyTorch SDPA on CPU or MPS."
+                    )
                 return FlashAttention4()
             case AttentionFunction.SDPA_MATH:
                 return PytorchAttention(priority=[SDPBackend.MATH])
+            case AttentionFunction.MPS_SDPA:
+                if _mps_sdpa_opt is None:
+                    raise RuntimeError("AttentionFunction.MPS_SDPA selected but `mps-sdpa` is not installed.")
+                return MPSSdpaAttention()
             case AttentionFunction.SDPA_CUDNN:
                 return _resolve_sdpa_variant(
                     SDPBackend.CUDNN_ATTENTION, "AttentionFunction.SDPA_CUDNN", with_mask=False
@@ -388,10 +419,13 @@ class MaskedAttentionFunction(Enum):
     Keeping them out makes "this backend cannot mask" a type error, not a runtime one."""
 
     PYTORCH = "pytorch"
-    XFORMERS = "xformers"
     SDPA_CUDNN = "sdpa_cudnn"
     SDPA_EFFICIENT = "sdpa_efficient"
     SDPA_MATH = "sdpa_math"
+    # Apple's fused MPSGraph SDPA via the `mps-sdpa` package (macOS/MPS only, a
+    # platform-marked hard dependency on Apple Silicon); the AUTOMATIC default on
+    # MPS. Mask-aware.
+    MPS_SDPA = "mps_sdpa"
     # Pick the fastest mask-capable backend for the current extras combo; see
     # :func:`automatic_masked_attention`. Default for the masked slot of
     # :class:`AttentionOps`.
@@ -410,12 +444,12 @@ class MaskedAttentionFunction(Enum):
                 return automatic_masked_attention()
             case MaskedAttentionFunction.PYTORCH:
                 return PytorchAttention()
-            case MaskedAttentionFunction.XFORMERS:
-                if memory_efficient_attention is None:
-                    raise RuntimeError("MaskedAttentionFunction.XFORMERS selected but `xformers` is not installed.")
-                return XFormersAttention()
             case MaskedAttentionFunction.SDPA_MATH:
                 return PytorchAttention(priority=[SDPBackend.MATH])
+            case MaskedAttentionFunction.MPS_SDPA:
+                if _mps_sdpa_opt is None:
+                    raise RuntimeError("MaskedAttentionFunction.MPS_SDPA selected but `mps-sdpa` is not installed.")
+                return MPSSdpaAttention()
             case MaskedAttentionFunction.SDPA_CUDNN:
                 return _resolve_sdpa_variant(
                     SDPBackend.CUDNN_ATTENTION, "MaskedAttentionFunction.SDPA_CUDNN", with_mask=True
@@ -499,7 +533,7 @@ class Attention(torch.nn.Module):
             context: Key/value context tensor of shape ``(B, S, context_dim)``.
                 Falls back to ``x`` (self-attention) when *None*.
             mask: Optional attention mask. Interpretation depends on the attention
-                backend (additive bias for xformers/PyTorch SDPA). A non-None
+                backend (additive bias for PyTorch SDPA). A non-None
                 ``mask`` routes to ``masked_attention_function``; ``None`` keeps
                 the unmasked path.
             pe: Rotary positional embeddings applied to both ``q`` and ``k``.

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

@@ -1,4 +1,4 @@
-from dataclasses import dataclass, field
+from dataclasses import dataclass, field, replace
 from typing import Any
 
 import torch
@@ -27,19 +27,19 @@ class CompilationConfig:
     dynamo_config: dict[str, Any] = field(default_factory=lambda: dict(_DEFAULT_DYNAMO_CONFIG))
 
 
-class _SeqDynamicMarkingProcessor:
-    """Marks the per-block seq dim dynamic, then delegates to an inner processor.
-    Installed by ``compile_transformer`` so the per-block compile artifact stays
-    shape-polymorphic. Wraps whatever ``block_input_processor`` was already on
-    the model -- callers that customised the processor keep their customisation;
-    only the seq-dim marking is layered on top. Lives outside the compiled
-    region, so ``mark_dynamic`` runs in eager mode on the tensors that are
-    about to cross into the trace.
+class CompiledBlockPerturbationsProcessor(BlockPerturbationsProcessor):
+    """Per-block input prep for compiled blocks: mark the seq dim dynamic, then attach perturbation
+    as config-independent runtime masks so the block traces ONCE.
+    The ``mark_dynamic`` calls keep the per-block compile artifact shape-polymorphic; they run in
+    eager mode (this processor lives outside the compiled region) on the tensors about to cross into
+    the trace. Both keep-masks are then attached UNCONDITIONALLY and the skip flags pinned False, so
+    the trace is identical for every pass (cond / uncond / STG): the block never sees a flipped
+    Python bool (``self_attn_all_perturbed``) or a None-vs-tensor mask that Dynamo would specialise
+    on, so the STG pass no longer triggers a recompile. An all-keep mask blends to a no-op
+    (``out*1 + v*0``); an all-zero mask reproduces the skip. Reads only ``mask`` (runtime tensor
+    indexing), never the host-side ``all_in_batch`` / ``any_in_batch``.
     """
 
-    def __init__(self, inner: BlockPerturbationsProcessor) -> None:
-        self.inner = inner
-
     def __call__(
         self,
         args: TransformerArgs,
@@ -84,7 +84,14 @@ class _SeqDynamicMarkingProcessor:
         # and broadcasts, leave it static. Same guard pattern as `timesteps`.
         if args.cross_scale_shift_timestep is not None and args.cross_scale_shift_timestep.shape[1] > 1:
             torch._dynamo.mark_dynamic(args.cross_scale_shift_timestep, 1)
-        return self.inner(args, perturbations, block_idx, self_attn_type, cross_attn_type)
+        # Perturbation as config-independent runtime masks (skip flags pinned False -> no recompile).
+        return replace(
+            args,
+            self_attn_perturbation_mask=perturbations.mask(self_attn_type, block_idx),
+            self_attn_all_perturbed=False,
+            cross_attn_perturbation_mask=perturbations.mask(cross_attn_type, block_idx),
+            cross_attn_skip_all=False,
+        )
 
 
 def compile_transformer(model: LTXModel, config: CompilationConfig) -> LTXModel:
@@ -100,7 +107,7 @@ def compile_transformer(model: LTXModel, config: CompilationConfig) -> LTXModel:
         torch.compile(m, mode=config.mode, backend=config.backend, fullgraph=config.fullgraph, dynamic=config.dynamic)
         for m in model.transformer_blocks
     )
-    model.block_input_processor = _SeqDynamicMarkingProcessor(inner=model.block_input_processor)
+    model.block_input_processor = CompiledBlockPerturbationsProcessor()
 
     def patched_dynamo_forward(*args, **kwargs) -> tuple[torch.Tensor, torch.Tensor]:
         torch.compiler.cudagraph_mark_step_begin()

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

@@ -77,7 +77,7 @@ class LTXModel(torch.nn.Module):
         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
+        # (PYTORCH/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",
@@ -358,21 +358,18 @@ class LTXModel(torch.nn.Module):
         """
         self._enable_gradient_checkpointing = enable
 
+    @property
+    def num_blocks(self) -> int:
+        """Number of transformer blocks."""
+        return len(self.transformer_blocks)
+
     def _process_transformer_blocks(
         self,
         video: TransformerArgs | None,
         audio: TransformerArgs | None,
-        perturbations: BatchedPerturbationConfig | None,
+        perturbations: BatchedPerturbationConfig,
     ) -> tuple[TransformerArgs | None, TransformerArgs | None]:
-        """Process transformer blocks for LTXAV.
-        Per-block perturbation masks are precomputed here and attached to each
-        modality's ``TransformerArgs`` so the block forward has no per-block
-        identity to specialise on — all blocks share a single Dynamo cache slot.
-        """
-        if perturbations is None:
-            batch_size = (video or audio).x.shape[0]
-            perturbations = BatchedPerturbationConfig.empty(batch_size)
-
+        """Process transformer blocks for LTX."""
         for block_idx, block in enumerate(self.transformer_blocks):
             if video is not None:
                 video = self.block_input_processor(
@@ -424,7 +421,7 @@ class LTXModel(torch.nn.Module):
         return x
 
     def forward(
-        self, video: Modality | None, audio: Modality | None, perturbations: BatchedPerturbationConfig
+        self, video: Modality | None, audio: Modality | None, perturbations: BatchedPerturbationConfig | None
     ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
         """
         Forward pass for LTX models.
@@ -438,6 +435,11 @@ class LTXModel(torch.nn.Module):
 
         video_args = self.video_args_preprocessor.prepare(video, audio) if video is not None else None
         audio_args = self.audio_args_preprocessor.prepare(audio, video) if audio is not None else None
+        # Materialize the no-perturbation mask here (eager); a None config means "perturb nothing"
+        # -> all-keep masks. The block loop never builds masks.
+        if perturbations is None:
+            ref = (video_args or audio_args).x
+            perturbations = BatchedPerturbationConfig.empty(ref.shape[0], self.num_blocks, ref.device, ref.dtype)
         # Process transformer blocks
         video_out, audio_out = self._process_transformer_blocks(
             video=video_args,
@@ -477,6 +479,11 @@ class LegacyX0Model(torch.nn.Module):
         super().__init__()
         self.velocity_model = velocity_model
 
+    @property
+    def num_blocks(self) -> int:
+        """Number of transformer blocks."""
+        return self.velocity_model.num_blocks
+
     def forward(
         self,
         video: Modality | None,
@@ -506,11 +513,16 @@ class X0Model(torch.nn.Module):
         super().__init__()
         self.velocity_model = velocity_model
 
+    @property
+    def num_blocks(self) -> int:
+        """Number of transformer blocks."""
+        return self.velocity_model.num_blocks
+
     def forward(
         self,
         video: Modality | None,
         audio: Modality | None,
-        perturbations: BatchedPerturbationConfig,
+        perturbations: BatchedPerturbationConfig | None,
     ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
         """
         Denoise the video and audio according to the sigma.

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

@@ -62,18 +62,16 @@ class BlockPerturbationsProcessor:
         self_attn_type: PerturbationType,
         cross_attn_type: PerturbationType,
     ) -> "TransformerArgs":
-        device, dtype = args.x.device, args.x.dtype
-
         all_self = perturbations.all_in_batch(self_attn_type, block_idx)
         any_self = perturbations.any_in_batch(self_attn_type, block_idx)
         self_mask: torch.Tensor | None = None
         if any_self and not all_self:
-            self_mask = perturbations.mask(self_attn_type, block_idx, device, dtype).view(-1, 1, 1)
+            self_mask = perturbations.mask(self_attn_type, block_idx)
 
         all_cross = perturbations.all_in_batch(cross_attn_type, block_idx)
         cross_mask: torch.Tensor | None = None
         if not all_cross:
-            cross_mask = perturbations.mask(cross_attn_type, block_idx, device, dtype).view(-1, 1, 1)
+            cross_mask = perturbations.mask(cross_attn_type, block_idx)
 
         return replace(
             args,

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

@@ -0,0 +1,11 @@
+"""
+Multi-GPU utilities for LTX models.
+This package provides utilities for running LTX models across multiple GPUs
+using tiled data-parallel techniques and sharded state-dict utilities.
+"""
+
+from ltx_core.multigpu import transformer, vae
+from ltx_core.multigpu.sharded_sd import ShardedSD
+from ltx_core.tiling import DimensionTilingConfig, TileCountConfig
+
+__all__ = ["DimensionTilingConfig", "ShardedSD", "TileCountConfig", "transformer", "vae"]

packages/ltx-core/src/ltx_core/multigpu/gemma/__init__.py

@@ -0,0 +1,6 @@
+"""Multi-GPU utilities for the Gemma text encoder."""
+
+from ltx_core.multigpu.gemma.accelerate_wrapper import AccelerateGemmaWrapper
+from ltx_core.multigpu.gemma.loader import load_gemma_with_device_map
+
+__all__ = ["AccelerateGemmaWrapper", "load_gemma_with_device_map"]

packages/ltx-core/src/ltx_core/multigpu/gemma/accelerate_wrapper.py

@@ -0,0 +1,30 @@
+"""Accelerate-based Gemma text encoder wrapper for multi-GPU inference.
+One rank (``src_rank``) holds the real ``GemmaTextEncoder`` loaded with
+``device_map="auto"``; other ranks hold a lightweight stub. Every public
+method runs on the source rank and broadcasts results to all ranks via the
+provided NCCL process group.
+The ``broadcast_group`` should cover **all ranks that need the
+embeddings** (typically the transformer group or world group).
+"""
+
+from __future__ import annotations
+
+import torch
+
+from ltx_core.multigpu.gemma.broadcast_wrapper import BroadcastGemmaWrapper
+
+
+class AccelerateGemmaWrapper(BroadcastGemmaWrapper):
+    """Source-rank encode + NCCL broadcast around a sharded ``GemmaTextEncoder``."""
+
+    def encode(
+        self,
+        prompts: list[str],
+        padding_side: str = "left",
+    ) -> list[tuple[tuple[torch.Tensor, ...], torch.Tensor]]:
+        """Fuse all prompts into one Gemma call on the source rank, broadcast each output."""
+        if self._rank == self._src_rank:
+            local_outputs = self._encoder.encode(prompts, padding_side)
+        else:
+            local_outputs = [(None, None)] * len(prompts)
+        return [self._broadcast_encoder_output(hs, mask, self._src_rank) for hs, mask in local_outputs]

packages/ltx-core/src/ltx_core/multigpu/gemma/batch_parallel_wrapper.py

@@ -0,0 +1,75 @@
+"""Batch-parallel Gemma text encoder wrapper for multi-GPU inference.
+Each rank holds a full :class:`GemmaTextEncoder` replica resident on its
+own GPU.  ``encode`` partitions the prompt list across ranks (each rank
+encodes a disjoint slice) and broadcasts every prompt's outputs from its
+encoding rank to all other ranks, so all ranks end up with the full list
+in the original order.
+``enhance_t2v`` / ``enhance_i2v`` (inherited) involve sampling, so they
+execute on ``src_rank`` only and the generated string is broadcast.
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.distributed as dist
+
+from ltx_core.multigpu.gemma.broadcast_wrapper import BroadcastGemmaWrapper
+from ltx_core.text_encoders.gemma.encoders.base_encoder import GemmaTextEncoder
+
+
+def _partition(total: int, world_size: int) -> list[int]:
+    """Spread ``total`` items across ``world_size`` ranks; remainder lands on the first ranks."""
+    base, rem = divmod(total, world_size)
+    return [base + (1 if i < rem else 0) for i in range(world_size)]
+
+
+class BatchParallelGemmaWrapper(BroadcastGemmaWrapper):
+    """Per-rank Gemma replica; ``encode`` parallelises a batch across ranks."""
+
+    _encoder: GemmaTextEncoder  # always resident on every rank, unlike the base's optional encoder
+
+    def __init__(
+        self,
+        encoder: GemmaTextEncoder,
+        broadcast_group: dist.ProcessGroup | None,
+        src_rank: int,
+        dtype: torch.dtype = torch.bfloat16,
+        device: torch.device | None = None,
+    ) -> None:
+        """Wrap a per-rank Gemma replica for batch-parallel encoding.
+        Args:
+            encoder: Full Gemma replica; required and resident on every rank (unlike
+                the base, where it is optional and real only on ``src_rank``).
+            broadcast_group: NCCL group spanning the ranks that share the encode work.
+            src_rank: Rank within ``broadcast_group`` that runs the inherited sampling
+                methods (``enhance_t2v`` / ``enhance_i2v``); ``encode`` uses every rank.
+            dtype: Target dtype for output tensors.
+            device: Target device for output tensors; defaults to the current CUDA device.
+        """
+        super().__init__(encoder, broadcast_group, src_rank, dtype, device)
+        self._world_size = dist.get_world_size(broadcast_group)
+
+    def encode(
+        self,
+        prompts: list[str],
+        padding_side: str = "left",
+    ) -> list[tuple[tuple[torch.Tensor, ...], torch.Tensor]]:
+        """Partition prompts across ranks, encode in parallel, broadcast per-prompt outputs.
+        With B prompts on W ranks, each rank gets ``ceil(B/W)`` or ``floor(B/W)``
+        prompts; the typical pos+neg case (B=2, W=2) gives one prompt per rank,
+        running both Gemma forwards concurrently on different GPUs.
+        """
+        n = len(prompts)
+        if n == 0:
+            return []
+        counts = _partition(n, self._world_size)
+        start = sum(counts[: self._rank])
+        local_prompts = prompts[start : start + counts[self._rank]]
+        local_outputs = self._encoder.encode(local_prompts, padding_side) if local_prompts else []
+
+        all_outputs: list[tuple[tuple[torch.Tensor, ...], torch.Tensor]] = []
+        for owner_rank, owner_count in enumerate(counts):
+            for slot in range(owner_count):
+                hs, mask = local_outputs[slot] if owner_rank == self._rank else (None, None)
+                all_outputs.append(self._broadcast_encoder_output(hs, mask, owner_rank))
+        return all_outputs

packages/ltx-core/src/ltx_core/multigpu/gemma/broadcast_wrapper.py

@@ -0,0 +1,108 @@
+"""Shared base for the multi-GPU Gemma wrappers: src-rank gating + NCCL broadcast."""
+
+from __future__ import annotations
+
+import torch
+import torch.distributed as dist
+
+from ltx_core.text_encoders.gemma.encoders.base_encoder import GemmaTextEncoder
+
+
+class BroadcastGemmaWrapper(torch.nn.Module):
+    """Encoder/group plumbing, prompt enhancement, and result broadcast.
+    Subclasses implement ``encode`` (the stub below raises).
+    Args:
+        encoder: The encoder; real on ``src_rank``, may be ``None`` elsewhere.
+        broadcast_group: NCCL group covering ranks that need the embeddings.
+        src_rank: Rank *within* ``broadcast_group`` that holds the real encoder and runs
+            the sampling-based ``enhance_*`` methods; its results are broadcast to every
+            other rank in the group. Builders derive it from a global driver rank via
+            ``dist.get_group_rank``.
+        dtype: Target dtype for output tensors.
+        device: Target device for output tensors.
+    """
+
+    def __init__(
+        self,
+        encoder: GemmaTextEncoder | None,
+        broadcast_group: dist.ProcessGroup | None,
+        src_rank: int,
+        dtype: torch.dtype = torch.bfloat16,
+        device: torch.device | None = None,
+    ) -> None:
+        super().__init__()
+        if device is None and torch.cuda.is_available():
+            device = torch.device("cuda", torch.cuda.current_device())
+        self._encoder = encoder
+        self._group = broadcast_group
+        self._src_rank = src_rank
+        self._rank = dist.get_rank(broadcast_group)
+        self._dtype = dtype
+        self._device = device
+
+    def encode(
+        self,
+        prompts: list[str],
+        padding_side: str = "left",
+    ) -> list[tuple[tuple[torch.Tensor, ...], torch.Tensor]]:
+        """Encode a batch of prompts to per-prompt hidden states; implemented by subclasses."""
+        raise NotImplementedError
+
+    def enhance_t2v(
+        self,
+        prompt: str,
+        max_new_tokens: int = 512,
+        system_prompt: str | None = None,
+        seed: int = 10,
+    ) -> str:
+        result = None
+        if self._rank == self._src_rank:
+            result = self._encoder.enhance_t2v(prompt, max_new_tokens, system_prompt, seed)
+        return self._broadcast_str(result)
+
+    def enhance_i2v(
+        self,
+        prompt: str,
+        image: torch.Tensor,
+        max_new_tokens: int = 512,
+        system_prompt: str | None = None,
+        seed: int = 10,
+    ) -> str:
+        result = None
+        if self._rank == self._src_rank:
+            result = self._encoder.enhance_i2v(prompt, image, max_new_tokens, system_prompt, seed)
+        return self._broadcast_str(result)
+
+    def _broadcast_str(self, value: str | None) -> str:
+        obj_list: list[str | None] = [value]
+        dist.broadcast_object_list(obj_list, src=self._src_rank, group=self._group)
+        result = obj_list[0]
+        assert result is not None, "broadcast returned None; check src_rank/broadcast_group"
+        return result
+
+    def _broadcast_encoder_output(
+        self,
+        hidden_states: tuple[torch.Tensor, ...] | None,
+        attention_mask: torch.Tensor | None,
+        src_rank: int,
+    ) -> tuple[tuple[torch.Tensor, ...], torch.Tensor]:
+        """Broadcast hidden states + attention mask via NCCL from ``src_rank``."""
+        if self._rank == src_rank:
+            meta = [{"hs_shapes": [h.shape for h in hidden_states], "mask_shape": attention_mask.shape}]
+        else:
+            meta = [None]
+        dist.broadcast_object_list(meta, src=src_rank, group=self._group)
+        info = meta[0]
+
+        if self._rank != src_rank:
+            hidden_states = tuple(torch.empty(s, device=self._device, dtype=self._dtype) for s in info["hs_shapes"])
+            attention_mask = torch.empty(info["mask_shape"], device=self._device, dtype=torch.long)
+        else:
+            hidden_states = tuple(h.to(device=self._device, dtype=self._dtype) for h in hidden_states)
+            attention_mask = attention_mask.to(device=self._device)
+
+        for h in hidden_states:
+            dist.broadcast(h, src=src_rank, group=self._group)
+        dist.broadcast(attention_mask, src=src_rank, group=self._group)
+
+        return hidden_states, attention_mask

packages/ltx-core/src/ltx_core/multigpu/gemma/loader.py

@@ -0,0 +1,55 @@
+"""Load GemmaTextEncoder with Accelerate ``device_map="auto"``.
+The Gemma LLM backbone is spread across available CUDA devices using
+HuggingFace Accelerate's automatic device placement.
+Mirrors the ``PromptEncoder`` text-encoder loading in
+``ltx_pipelines.utils.blocks`` but uses ``device_map="auto"`` instead of
+placing the entire model on a single GPU.
+"""
+
+from __future__ import annotations
+
+import logging
+
+import torch
+from transformers import AutoImageProcessor, Gemma3ForConditionalGeneration, Gemma3Processor
+
+from ltx_core.text_encoders.gemma.encoders.base_encoder import GemmaTextEncoder
+from ltx_core.text_encoders.gemma.tokenizer import LTXVGemmaTokenizer
+from ltx_core.utils import find_matching_file
+
+logger = logging.getLogger(__name__)
+
+
+def load_gemma_with_device_map(
+    gemma_root_path: str,
+    dtype: torch.dtype = torch.bfloat16,
+) -> GemmaTextEncoder:
+    """Load GemmaTextEncoder with the LLM backbone spread across GPUs.
+    Uses ``Gemma3ForConditionalGeneration.from_pretrained(device_map="auto")``
+    to distribute layers across available CUDA devices.
+    Args:
+        gemma_root_path: Path to Gemma model directory.
+        dtype: Data type for model weights.
+    """
+    model_folder = str(find_matching_file(gemma_root_path, "model*.safetensors").parent)
+    tokenizer_path = str(find_matching_file(gemma_root_path, "tokenizer.model").parent)
+    processor_path = str(find_matching_file(gemma_root_path, "preprocessor_config.json").parent)
+
+    logger.info("Loading Gemma LLM with device_map='auto'...")
+    gemma_model = Gemma3ForConditionalGeneration.from_pretrained(
+        model_folder,
+        dtype=dtype,
+        device_map="auto",
+        local_files_only=True,
+    )
+
+    tokenizer = LTXVGemmaTokenizer(tokenizer_path, 1024)
+    image_processor = AutoImageProcessor.from_pretrained(processor_path, local_files_only=True, use_fast=False)
+    processor = Gemma3Processor(image_processor=image_processor, tokenizer=tokenizer.tokenizer)
+
+    return GemmaTextEncoder(
+        model=gemma_model,
+        tokenizer=tokenizer,
+        processor=processor,
+        dtype=dtype,
+    )

packages/ltx-core/src/ltx_core/multigpu/sharded_sd.py

@@ -0,0 +1,191 @@
+"""Sharded state dict with distributed weight backup.
+Each rank stores ~1/N of the model weights. Bucketed broadcasts
+restore weights into a target state dict using only a small,
+caller-provided staging buffer.
+"""
+
+from __future__ import annotations
+
+import hashlib
+from dataclasses import dataclass
+
+import torch
+import torch.distributed as dist
+
+
+def _stable_owner(key: str, world: int) -> int:
+    """Deterministic rank assignment (same across all processes)."""
+    h = hashlib.md5(key.encode("utf-8")).digest()
+    return int.from_bytes(h[:8], "little") % world
+
+
+def _nbytes(t: torch.Tensor) -> int:
+    return t.numel() * t.element_size()
+
+
+@dataclass
+class ShardedSD:
+    """Sharded state dict with distributed weight backup.
+    Distributes model weights across ranks for memory-efficient backup
+    and restoration.  Can be used for any scenario where a full state dict
+    needs to be restored from sharded storage (e.g. LoRA hot-swap, weight
+    rollback, checkpoint recovery).
+    - Deterministic ownership: ``MD5(key) % world_size``
+    - Local storage only for owned keys (VRAM ≈ 1/world_size of model)
+    - Bucketed broadcast using a single small staging buffer
+    Usage::
+        backup = ShardedSD.from_state_dict(model.state_dict(), group)
+        staging = torch.empty(64 * 1024 * 1024, dtype=torch.uint8, device=device)
+        backup.broadcast_shards_into(target_sd, staging)    # cooperative: all ranks must call
+    """
+
+    keys: tuple[str, ...]
+    """All parameter keys in the original state dict, in insertion order."""
+    key_sizes: dict[str, int]
+    """Byte size of each parameter tensor (numel * element_size)."""
+    owner_of: dict[str, int]
+    """Maps each key to the rank that stores it."""
+    local_shard: dict[str, torch.Tensor]
+    """Tensors owned by this rank (subset of the full state dict)."""
+    rank: int
+    """This process's rank within the group."""
+    world: int
+    """Total number of ranks in the group."""
+    group: dist.ProcessGroup
+    """NCCL process group used for broadcast operations."""
+    _owner_groups: dict[int, list[str]]
+    """Keys grouped by owning rank, sorted by descending tensor size."""
+
+    @classmethod
+    def from_state_dict(
+        cls,
+        sd: dict[str, torch.Tensor],
+        group: dist.ProcessGroup,
+        clone: bool = True,
+    ) -> ShardedSD:
+        rank = dist.get_rank(group)
+        world = dist.get_world_size(group)
+
+        keys = tuple(sd.keys())
+        # Co-locate .weight_scale with its .weight on the same rank.
+        owner_of: dict[str, int] = {}
+        for k in keys:
+            if k.endswith(".weight_scale"):
+                parent = k.replace(".weight_scale", ".weight")
+                if parent in sd:
+                    owner_of[k] = _stable_owner(parent, world)
+                    continue
+            owner_of[k] = _stable_owner(k, world)
+
+        key_sizes = {k: _nbytes(v) for k, v in sd.items()}
+
+        local_shard: dict[str, torch.Tensor] = {}
+        for k, v in sd.items():
+            if owner_of[k] == rank:
+                local_shard[k] = v.clone() if clone else v
+
+        owner_groups: dict[int, list[str]] = {r: [] for r in range(world)}
+        for k in keys:
+            owner_groups[owner_of[k]].append(k)
+        for r in range(world):
+            owner_groups[r].sort(key=lambda kk: key_sizes[kk], reverse=True)
+
+        return cls(
+            keys=keys,
+            key_sizes=key_sizes,
+            owner_of=owner_of,
+            local_shard=local_shard,
+            rank=rank,
+            world=world,
+            group=group,
+            _owner_groups=owner_groups,
+        )
+
+    def broadcast_shards_into(
+        self,
+        target_sd: dict[str, torch.Tensor],
+        staging: torch.Tensor,
+    ) -> None:
+        """Broadcast stored weights from sharded backup into *target_sd*.
+        This is a **cooperative operation** — all ranks in the process group
+        must call it simultaneously.
+        *staging* is a caller-owned ``uint8`` scratch buffer; its size sets the
+        broadcast granularity (tensors larger than it split across rounds). It
+        may be shared by instances that never broadcast at the same time. Writes
+        directly into existing tensors in *target_sd*.
+        """
+        if staging.dtype != torch.uint8 or staging.numel() == 0:
+            raise ValueError("staging must be a non-empty uint8 buffer")
+        for owner, klist in self._owner_groups.items():
+            if klist:
+                self._broadcast_group(owner, klist, target_sd, staging)
+
+    def _broadcast_group(
+        self,
+        owner: int,
+        keys: list[str],
+        target_sd: dict[str, torch.Tensor],
+        staging: torch.Tensor,
+    ) -> None:
+        """Pack & broadcast params from *owner*, splitting tensors across rounds."""
+        rounds = self._plan_rounds(keys, staging.numel())
+
+        for round_chunks in rounds:
+            filled = 0
+            if self.rank == owner:
+                for k, offset, chunk_size in round_chunks:
+                    src = self.local_shard[k]
+                    if not src.is_contiguous():
+                        raise RuntimeError(f"ShardedSD: local shard tensor '{k}' is not contiguous")
+                    src_bytes = src.view(torch.uint8).view(-1)
+                    staging[filled : filled + chunk_size].copy_(
+                        src_bytes[offset : offset + chunk_size], non_blocking=True
+                    )
+                    filled += chunk_size
+            else:
+                filled = sum(chunk_size for (_, _, chunk_size) in round_chunks)
+
+            if filled == 0:
+                continue
+            view = staging[:filled]
+            dist.broadcast(view, src=owner, group=self.group)
+
+            cursor = 0
+            for k, offset, chunk_size in round_chunks:
+                dst = target_sd[k]
+                if not dst.is_contiguous():
+                    raise RuntimeError(f"ShardedSD: target tensor '{k}' is not contiguous")
+                dst_bytes = dst.view(torch.uint8).view(-1)
+                dst_bytes[offset : offset + chunk_size].copy_(staging[cursor : cursor + chunk_size], non_blocking=True)
+                cursor += chunk_size
+
+    def _plan_rounds(self, keys: list[str], capacity: int) -> list[list[tuple[str, int, int]]]:
+        """Build rounds that pack a *capacity*-byte buffer, splitting tensors if needed.
+        Returns a list of rounds, each containing ``(key, byte_offset, chunk_bytes)`` tuples.
+        """
+        rounds: list[list[tuple[str, int, int]]] = []
+        current: list[tuple[str, int, int]] = []
+        used = 0
+
+        for k in keys:
+            remaining = self.key_sizes[k]
+            offset = 0
+
+            while remaining > 0:
+                space = capacity - used
+                if space == 0:
+                    rounds.append(current)
+                    current = []
+                    used = 0
+                    space = capacity
+
+                chunk = min(remaining, space)
+                current.append((k, offset, chunk))
+                used += chunk
+                offset += chunk
+                remaining -= chunk
+
+        if current:
+            rounds.append(current)
+
+        return rounds

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

@@ -0,0 +1,13 @@
+"""
+Multi-GPU transformer utilities for LTX models.
+This module provides utilities for running LTX transformer models across multiple GPUs
+using tiled data parallelism.
+"""
+
+from ltx_core.multigpu.transformer.tiled_data_parallel import (
+    TiledDataParallelModelWrapper,
+)
+
+__all__ = [
+    "TiledDataParallelModelWrapper",
+]

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

@@ -0,0 +1,270 @@
+import torch
+import torch.distributed as dist
+
+from ltx_core.model.transformer.attention import AttentionCallable, MaskedAttentionCallable
+
+# Mirrors the kernel's DEFAULT_BARRIER_TIMEOUT_SECONDS (configs.cuh), which All2All converts to
+# cycles via the device peak SM clock. Stored so the timeout can be read back to reset after a raise.
+_DEFAULT_ALL2ALL_TIMEOUT_SECONDS = 10.0
+
+
+class AttentionManager:
+    def __init__(
+        self,
+        max_tokens: int,
+        num_heads: int,
+        head_dim: int,
+        tensor_dtype: torch.dtype,
+        group: torch.distributed.ProcessGroup,
+        copy_out_: bool = False,
+    ) -> None:
+        # Lazy: ltx_kernels is an optional GPU-only dep, and this constructor already
+        # requires CUDA -- so importing it here (not at module scope) keeps the multigpu
+        # modules importable without the kernels installed (e.g. CPU CI test collection).
+        from ltx_kernels import All2All  # noqa: PLC0415
+
+        self.rank = dist.get_rank(group)
+        self.world_size = dist.get_world_size(group)
+        self.max_tokens = max_tokens
+        hidden_dim = num_heads * head_dim
+        num_sms = torch.cuda.get_device_properties(self.rank).multi_processor_count
+        self.copy_out = copy_out_
+        buffer_seqlen = (max_tokens + self.world_size - 1) // self.world_size
+        self.all2all_heads, self.all2all_q = (
+            All2All(
+                rank=self.rank,
+                world_size=self.world_size,
+                seqlen=buffer_seqlen,
+                hidden_dim=hidden_dim,
+                num_sms=num_sms,
+                tensor_dtype=tensor_dtype,
+                group=group,
+            )
+            for _ in range(2)
+        )
+        self.all2all_k, self.all2all_v = (
+            All2All(
+                rank=self.rank,
+                world_size=self.world_size,
+                seqlen=buffer_seqlen,
+                hidden_dim=hidden_dim,
+                num_sms=num_sms,
+                tensor_dtype=tensor_dtype,
+                group=group,
+            )
+            if not self.copy_out
+            else self.all2all_q
+            for _ in range(2)
+        )
+        self.group = group
+        self._all2all_timeout_seconds = _DEFAULT_ALL2ALL_TIMEOUT_SECONDS
+
+    def set_seqlen_all2all(self, seqlens: list[int]) -> None:
+        # Route through the wrappers so the registered custom ops' fake-impl
+        # shape info gets updated alongside the C++ runtime's rank_tokens.
+        self.all2all_q.set_rank_tokens(seqlens)
+        self.all2all_k.set_rank_tokens(seqlens)
+        self.all2all_v.set_rank_tokens(seqlens)
+        self.all2all_heads.set_rank_tokens(seqlens)
+
+    @property
+    def all2all_timeout_seconds(self) -> float:
+        """The all2all barrier (deadlock-detection) timeout, in seconds, applied to every instance."""
+        return self._all2all_timeout_seconds
+
+    @all2all_timeout_seconds.setter
+    def all2all_timeout_seconds(self, seconds: float) -> None:
+        # Raise it for the first ``torch.compile`` forward -- where one rank's recompile can delay its
+        # all2all kernel launch past the steady-state timeout, tripping the barrier -- then reset to
+        # the prior value. ``all2all_k``/``all2all_v`` may alias ``all2all_q`` (copy-out path);
+        # setting twice is idempotent. Fan out first (it validates) so a rejected value leaves the
+        # stored steady-state value untouched.
+        for a2a in (self.all2all_q, self.all2all_k, self.all2all_v, self.all2all_heads):
+            a2a.set_timeout_seconds(seconds)
+        self._all2all_timeout_seconds = seconds
+
+    def send_recv_qkv(
+        self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        t_q = self.all2all_q.send_recv_heads(q, copy_out=self.copy_out)
+        t_k = self.all2all_k.send_recv_heads(k, copy_out=self.copy_out)
+        t_v = self.all2all_v.send_recv_heads(v, copy_out=self.copy_out)
+        return t_q, t_k, t_v
+
+    def gather_heads(self, heads_local: torch.Tensor) -> torch.Tensor:
+        out = self.all2all_heads.gather_heads(heads_local, copy_out=self.copy_out)
+        return out
+
+
+class _All2AllRedistribute:
+    """Shared redistribute/gather pipeline for self-attention SP wrappers.
+    Folds the head dim view-and-shuffle so the masked and unmasked variants only
+    have to choose how to invoke the inner attention (with or without the mask
+    kwarg) -- the rest of the SP plumbing is identical.
+    """
+
+    def __init__(self, manager: AttentionManager) -> None:
+        self.manager = manager
+
+    def redistribute(
+        self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, heads: int
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int, int]:
+        if heads % self.manager.world_size != 0:
+            raise ValueError(f"heads ({heads}) must be divisible by world_size ({self.manager.world_size})")
+
+        head_dim = q.shape[-1] // heads
+        q = q.view(q.shape[0], q.shape[1], heads, head_dim)
+        k = k.view(k.shape[0], k.shape[1], heads, head_dim)
+        v = v.view(v.shape[0], v.shape[1], heads, head_dim)
+
+        t_q, t_k, t_v = self.manager.send_recv_qkv(q, k, v)
+        local_heads = heads // self.manager.world_size
+
+        # `flatten` / `unflatten` collapse only the head dims, avoiding a `-1` in the
+        # seq position -- that would otherwise be ambiguous if the seq is 0 for a
+        # zero-token modality.
+        t_q = t_q.flatten(-2)
+        t_k = t_k.flatten(-2)
+        t_v = t_v.flatten(-2)
+        return t_q, t_k, t_v, local_heads, head_dim
+
+    def gather(self, hidden_states: torch.Tensor, local_heads: int, head_dim: int) -> torch.Tensor:
+        hidden_states = hidden_states.unflatten(-1, (local_heads, head_dim))
+        hidden_states = self.manager.gather_heads(hidden_states)
+        return hidden_states.flatten(-2)
+
+
+class All2AllAttention(AttentionCallable):
+    def __init__(self, manager: AttentionManager, original_attention: AttentionCallable):
+        self._sp = _All2AllRedistribute(manager)
+        self.original_attention = original_attention
+
+    def __call__(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        heads: int,
+    ) -> torch.Tensor:
+        t_q, t_k, t_v, local_heads, head_dim = self._sp.redistribute(q, k, v, heads)
+        hidden_states = self.original_attention(q=t_q, k=t_k, v=t_v, heads=local_heads)
+        return self._sp.gather(hidden_states, local_heads, head_dim)
+
+
+class MaskedAll2AllAttention(MaskedAttentionCallable):
+    def __init__(self, manager: AttentionManager, original_attention: MaskedAttentionCallable):
+        self._sp = _All2AllRedistribute(manager)
+        self.original_attention = original_attention
+
+    def __call__(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        heads: int,
+        mask: torch.Tensor,
+    ) -> torch.Tensor:
+        t_q, t_k, t_v, local_heads, head_dim = self._sp.redistribute(q, k, v, heads)
+        hidden_states = self.original_attention(q=t_q, k=t_k, v=t_v, heads=local_heads, mask=mask)
+        return self._sp.gather(hidden_states, local_heads, head_dim)
+
+
+class _AudioAll2AllRedistribute:
+    """Shared redistribute/gather pipeline for audio cross-attention SP wrappers.
+    Q is sliced locally per rank (no cross-rank shuffle on Q because the audio
+    sequence length is small enough to replicate); K/V are redistributed across
+    ranks via ``send_recv_heads``; outputs are gathered via
+    ``all_gather_into_tensor`` along the head dimension. The masked and unmasked
+    variants share this plumbing and only differ in how they invoke the inner
+    attention.
+    """
+
+    def __init__(self, manager: AttentionManager) -> None:
+        self.manager = manager
+
+    def redistribute(
+        self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, heads: int
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int, int]:
+        if heads % self.manager.world_size != 0:
+            raise ValueError(f"heads ({heads}) must be divisible by world_size ({self.manager.world_size})")
+
+        head_dim = q.shape[-1] // heads
+        heads_per_rank = heads // self.manager.world_size
+        rank = self.manager.rank
+
+        q = q.view(q.shape[0], q.shape[1], heads, head_dim)
+        k = k.view(k.shape[0], k.shape[1], heads, head_dim)
+        v = v.view(v.shape[0], v.shape[1], heads, head_dim)
+
+        t_q = q[:, :, heads_per_rank * rank : heads_per_rank * (rank + 1), :].clone()
+        t_k = self.manager.all2all_k.send_recv_heads(k, copy_out=self.manager.copy_out)
+        t_v = self.manager.all2all_v.send_recv_heads(v, copy_out=self.manager.copy_out)
+
+        # `flatten` / `unflatten` collapse only the head dims, avoiding a `-1` in the
+        # seq position -- that would otherwise be ambiguous if the seq is 0 for a
+        # zero-token modality.
+        t_q = t_q.flatten(-2)
+        t_k = t_k.flatten(-2)
+        t_v = t_v.flatten(-2)
+        return t_q, t_k, t_v, heads_per_rank, head_dim
+
+    def gather(self, hidden_states: torch.Tensor, heads_per_rank: int, head_dim: int) -> torch.Tensor:
+        # (B, S, heads_per_rank, head_dim). Move head dim to dim 0 so all_gather_into_tensor
+        # gathers along it; permute back after the collective.
+        hidden_states = hidden_states.unflatten(-1, (heads_per_rank, head_dim)).permute(2, 0, 1, 3).contiguous()
+        gathered = torch.empty(
+            (heads_per_rank * self.manager.world_size, *hidden_states.shape[1:]),
+            dtype=hidden_states.dtype,
+            device=hidden_states.device,
+        )
+        dist.all_gather_into_tensor(gathered, hidden_states, group=self.manager.group)
+        # (heads, B, S, head_dim) -> (B, S, heads, head_dim) -> (B, S, heads * head_dim)
+        return gathered.permute(1, 2, 0, 3).flatten(-2)
+
+
+class AudioAll2AllAttention(AttentionCallable):
+    """All2All attention for audio cross-attention (video_to_audio).
+    Q is sliced locally per rank, K/V are redistributed via send_recv_heads,
+    then outputs are gathered via all_gather across the head dimension.
+    """
+
+    def __init__(self, manager: AttentionManager, original_attention: AttentionCallable):
+        self._sp = _AudioAll2AllRedistribute(manager)
+        self.original_attention = original_attention
+
+    def __call__(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        heads: int,
+    ) -> torch.Tensor:
+        t_q, t_k, t_v, heads_per_rank, head_dim = self._sp.redistribute(q, k, v, heads)
+        hidden_states = self.original_attention(q=t_q, k=t_k, v=t_v, heads=heads_per_rank)
+        return self._sp.gather(hidden_states, heads_per_rank, head_dim)
+
+
+class MaskedAudioAll2AllAttention(MaskedAttentionCallable):
+    """Masked counterpart to :class:`AudioAll2AllAttention`.
+    No current caller invokes A2V / V2A cross-attention with a mask, so the SP
+    mutator pre-installs an unmasked-only :class:`AudioAll2AllAttention` and the
+    masked slot stays at the model default. Defined now so adding masked audio
+    cross-attention later is just an SP-mutator change, not a missing-piece
+    discovery.
+    """
+
+    def __init__(self, manager: AttentionManager, original_attention: MaskedAttentionCallable):
+        self._sp = _AudioAll2AllRedistribute(manager)
+        self.original_attention = original_attention
+
+    def __call__(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        heads: int,
+        mask: torch.Tensor,
+    ) -> torch.Tensor:
+        t_q, t_k, t_v, heads_per_rank, head_dim = self._sp.redistribute(q, k, v, heads)
+        hidden_states = self.original_attention(q=t_q, k=t_k, v=t_v, heads=heads_per_rank, mask=mask)
+        return self._sp.gather(hidden_states, heads_per_rank, head_dim)

packages/ltx-core/src/ltx_core/multigpu/transformer/sequence_parallel.py

@@ -0,0 +1,300 @@
+"""
+Multi-GPU inference wrapper for LTX transformer models.
+This module provides utilities for running LTX model inference across multiple GPUs
+using sequence parallelism. It:
+- Tiles the video inputs across GPUs in the sequence (token) dimension
+- Patches video self-attention operations with all2all attention
+- Runs the model forward pass on each GPU with its local tile
+- Gathers all tokens back to all GPUs after the forward pass
+"""
+
+from dataclasses import replace
+from itertools import accumulate
+
+import torch
+
+from ltx_core.guidance.perturbations import BatchedPerturbationConfig
+from ltx_core.loader.module_ops import ModuleOps
+from ltx_core.model.transformer.attention import Attention
+from ltx_core.model.transformer.modality import Modality
+from ltx_core.model.transformer.model import LTXModel
+from ltx_core.model.transformer.transformer import BasicAVTransformerBlock
+from ltx_core.multigpu.transformer.attention import (
+    All2AllAttention,
+    AttentionManager,
+    AudioAll2AllAttention,
+    MaskedAll2AllAttention,
+    MaskedAudioAll2AllAttention,
+)
+
+
+def compute_sequence_partition(
+    total_tokens: int,
+    world_size: int,
+) -> list[int]:
+    """
+    Compute uniform per-rank token counts.
+    Requires ``total_tokens % world_size == 0`` — callers must pad up-front via
+    :func:`pad_modality_for_uniform_sharding`. Uniform sharding lets the
+    All2All custom-op fakes derive output shapes symbolically from input shapes
+    (``x.shape[1] * world_size`` / ``x.shape[1] // world_size``) instead of from
+    Python int args.
+    """
+    if total_tokens % world_size != 0:
+        raise ValueError(
+            f"compute_sequence_partition expects uniform sharding: total_tokens "
+            f"({total_tokens}) must be divisible by world_size ({world_size}). "
+            f"Pad the modality up-front."
+        )
+    per_rank = total_tokens // world_size
+    return [per_rank] * world_size
+
+
+def pad_modality_for_uniform_sharding(
+    modality: Modality,
+    world_size: int,
+) -> tuple[Modality, int]:
+    """Pad the seq dim up to the next multiple of ``world_size`` and attach a
+    padding-aware attention bias so the padded keys are ignored.
+    Returns ``(padded_modality, original_seq_len)``. If no padding is needed
+    the original modality is returned unchanged.
+    """
+    t_orig = modality.latent.shape[1]
+    pad = (-t_orig) % world_size
+    if pad == 0:
+        return modality, t_orig
+
+    t_padded = t_orig + pad
+    b = modality.latent.shape[0]
+    device = modality.latent.device
+    dtype = modality.latent.dtype
+
+    latent_pad = torch.zeros(b, pad, modality.latent.shape[2], dtype=dtype, device=device)
+    latent = torch.cat([modality.latent, latent_pad], dim=1)
+
+    timesteps_pad_shape = list(modality.timesteps.shape)
+    timesteps_pad_shape[1] = pad
+    timesteps_pad = torch.zeros(timesteps_pad_shape, dtype=modality.timesteps.dtype, device=modality.timesteps.device)
+    timesteps = torch.cat([modality.timesteps, timesteps_pad], dim=1)
+
+    positions_pad_shape = list(modality.positions.shape)
+    positions_pad_shape[2] = pad
+    positions_pad = torch.zeros(positions_pad_shape, dtype=modality.positions.dtype, device=modality.positions.device)
+    positions = torch.cat([modality.positions, positions_pad], dim=2)
+
+    if modality.attention_mask is None:
+        # Key-only padding mask in the canonical [0, 1] form: 1 on valid keys,
+        # 0 on padded keys. Shape (1, 1, T_padded) broadcasts across batch and
+        # queries -- O(T) memory instead of materialising a dense (B, T, T)
+        # matrix just to mask `pad` (< world_size) keys.
+        # `_prepare_self_attention_mask` does the standard 3D -> 4D log-space
+        # conversion and produces a (1, 1, 1, T_padded) bias.
+        attention_mask = torch.ones(1, 1, t_padded, dtype=torch.float32, device=device)
+        attention_mask[:, :, t_orig:] = 0.0
+    else:
+        # User-supplied (B, T, T) [0, 1] mask: extend with padded rows/cols.
+        # Padded query rows attend to all valid keys so their softmax stays
+        # well-defined (the outputs are sliced off after the gather, but a
+        # fully-masked row would produce NaN).
+        old = modality.attention_mask
+        attention_mask = torch.zeros(b, t_padded, t_padded, dtype=old.dtype, device=old.device)
+        attention_mask[:, :t_orig, :t_orig] = old
+        attention_mask[:, t_orig:, :t_orig] = 1.0
+
+    padded = replace(
+        modality,
+        latent=latent,
+        timesteps=timesteps,
+        positions=positions,
+        attention_mask=attention_mask,
+    )
+    return padded, t_orig
+
+
+def compute_sequence_offsets(token_counts: list[int]) -> list[int]:
+    """
+    Compute the starting offset for each rank's token partition.
+    Args:
+        token_counts: List of token counts per rank.
+    Returns:
+        List of starting offsets for each rank.
+    """
+    return [0, *accumulate(token_counts[:-1])]
+
+
+def tile_modality_for_rank(
+    modality: Modality,
+    rank: int,
+    world_size: int,
+) -> tuple[Modality, list[int]]:
+    """
+    Tile a modality's tensors for a specific GPU rank.
+    Splits the sequence dimension (dim 1 for latent/timesteps, dim 2 for positions)
+    across GPUs, returning the local tile for the given rank.
+    Args:
+        modality: The modality to tile.
+        rank: Current GPU rank.
+        world_size: Total number of GPUs.
+    Returns:
+        Tuple of (tiled_modality, token_counts_per_rank).
+    """
+    total_tokens = modality.latent.shape[1]
+    token_counts = compute_sequence_partition(total_tokens, world_size)
+    offsets = compute_sequence_offsets(token_counts)
+
+    start = offsets[rank]
+    end = start + token_counts[rank]
+
+    # Tile latent: (B, T, D) -> (B, T_local, D)
+    tiled_latent = modality.latent[:, start:end, :]
+
+    # Tile timesteps: (B, T) -> (B, T_local)
+    tiled_timesteps = modality.timesteps[:, start:end]
+
+    # Tile positions: (B, 3, T, 2) -> (B, 3, T_local, 2)
+    tiled_positions = modality.positions[:, :, start:end, :]
+
+    tiled_modality = replace(
+        modality,
+        latent=tiled_latent,
+        timesteps=tiled_timesteps,
+        positions=tiled_positions,
+    )
+
+    return tiled_modality, token_counts
+
+
+def gather_output_tokens(
+    local_output: torch.Tensor,
+    token_counts: list[int],
+    group: torch.distributed.ProcessGroup | None = None,
+) -> torch.Tensor:
+    """
+    Gather output tokens from all GPUs back into a single tensor.
+    Args:
+        local_output: Local output tensor of shape (B, T_local, D).
+        token_counts: Number of tokens on each rank.
+        group: Process group for communication. If None, uses default group.
+    Returns:
+        Gathered tensor of shape (B, T_total, D) on all ranks.
+    """
+    world_size = len(token_counts)
+    batch_size = local_output.shape[0]
+    hidden_dim = local_output.shape[2]
+
+    # Prepare output tensors for all_gather
+    max_tokens = max(token_counts)
+
+    # Pad local output to max size for uniform all_gather
+    padded_local = torch.zeros(
+        batch_size,
+        max_tokens,
+        hidden_dim,
+        dtype=local_output.dtype,
+        device=local_output.device,
+    )
+    padded_local[:, : local_output.shape[1], :] = local_output
+
+    # All gather padded outputs
+    gathered_list = [torch.zeros_like(padded_local) for _ in range(world_size)]
+    torch.distributed.all_gather(gathered_list, padded_local, group=group)
+
+    # Extract actual tokens (remove padding) and concatenate
+    outputs = []
+    for i, count in enumerate(token_counts):
+        outputs.append(gathered_list[i][:, :count, :])
+
+    return torch.cat(outputs, dim=1)
+
+
+def create_video_self_attention_module_ops(
+    attention_manager: AttentionManager,
+) -> ModuleOps:
+    """
+    Create ModuleOps for patching video self-attention with all2all attention.
+    This patches the `attn1` attribute on BasicAVTransformerBlock instances,
+    which is the video self-attention module.
+    Args:
+        attention_manager: The AttentionManager instance for all2all communication.
+    Returns:
+        ModuleOps that can be used to patch the model.
+    """
+
+    def mutator(module: torch.nn.Module) -> torch.nn.Module:
+        for block in module.transformer_blocks:
+            if not isinstance(block, BasicAVTransformerBlock):
+                continue
+
+            # Video self-attention: ``Attention.forward`` may receive a non-None
+            # ``mask`` (``video.self_attention_mask``), so wrap both slots; the
+            # branch in ``Attention.forward`` then routes to whichever wrapper
+            # corresponds to the actual call.
+            if hasattr(block, "attn1"):
+                attn1 = block.attn1
+                if isinstance(attn1, Attention):
+                    attn1.attention_function = All2AllAttention(attention_manager, attn1.attention_function)
+                    attn1.masked_attention_function = MaskedAll2AllAttention(
+                        attention_manager, attn1.masked_attention_function
+                    )
+            # video_to_audio cross-attention: no current caller passes a mask
+            # (see ``BasicAVTransformerBlock.forward``), so the masked branch
+            # is dead code today. Wrap both slots anyway so that if a future
+            # caller adds a mask, the SP plumbing is already in place rather
+            # than silently bypassing All2All on that path.
+            if hasattr(block, "video_to_audio_attn"):
+                video_to_audio_attn = block.video_to_audio_attn
+                if isinstance(video_to_audio_attn, Attention):
+                    video_to_audio_attn.attention_function = AudioAll2AllAttention(
+                        attention_manager, video_to_audio_attn.attention_function
+                    )
+                    video_to_audio_attn.masked_attention_function = MaskedAudioAll2AllAttention(
+                        attention_manager, video_to_audio_attn.masked_attention_function
+                    )
+        return module
+
+    return ModuleOps(
+        name="video_self_attention_all2all",
+        matcher=lambda module: isinstance(module, LTXModel),
+        mutator=mutator,
+    )
+
+
+class SequenceParallelModelWrapper(torch.nn.Module):
+    def __init__(self, model: torch.nn.Module, attention_manager: AttentionManager):
+        super().__init__()
+        self.model = model
+        self.attention_manager = attention_manager
+
+    @property
+    def num_blocks(self) -> int:
+        return self.model.num_blocks
+
+    def forward(
+        self, video: Modality | None, audio: Modality | None, perturbations: BatchedPerturbationConfig | None
+    ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
+        if video is None:
+            return self.model(video, audio, perturbations)
+
+        # Pad the video seq dim up to a multiple of world_size so all ranks get
+        # equal shards. The attention mask we attach makes the padded keys
+        # invisible to attention; padded rows are sliced off after the gather.
+        video, t_orig = pad_modality_for_uniform_sharding(video, self.attention_manager.world_size)
+
+        video_tile, token_counts = tile_modality_for_rank(
+            video, self.attention_manager.rank, self.attention_manager.world_size
+        )
+        total_tokens = sum(token_counts)
+        if total_tokens > self.attention_manager.max_tokens:
+            raise ValueError(
+                f"Total video token count ({total_tokens}) exceeds attention_manager max_tokens "
+                f"({self.attention_manager.max_tokens}). Use a smaller resolution or fewer frames."
+            )
+        self.attention_manager.set_seqlen_all2all(token_counts)
+        torch.distributed.barrier(self.attention_manager.group)
+        video, audio = self.model(video_tile, audio, perturbations)
+        video = gather_output_tokens(video, token_counts, self.attention_manager.group)
+        # Unpad: drop the rows we added in `pad_modality_for_uniform_sharding` to make
+        # the seq dim divisible by world_size, restoring the caller's original length.
+        if video.shape[1] != t_orig:
+            video = video[:, :t_orig, :]
+        return video, audio

packages/ltx-core/src/ltx_core/multigpu/transformer/tiled_data_parallel.py

@@ -0,0 +1,99 @@
+"""Tiled Data Parallel model wrapper for the LTX transformer.
+Each GPU processes one or more tiles of the patchified
+``(frames, height, width)`` latent.  Tiles are assigned to ranks via
+round-robin, so the number of tiles may exceed the number of GPUs.
+Tiles may overlap; overlapping regions are blended with trapezoidal
+masks so that seam artefacts are suppressed.  Each rank accumulates
+its assigned tiles locally, then a single ``all_reduce`` synchronises
+the blended output across all ranks.
+Conditioning tokens (appended after the generated tokens) are filtered
+per tile: only tokens whose positions overlap with the tile's spatial
+extent (or that have negative time coordinates) are included.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+import torch.distributed as dist
+
+from ltx_core.modality_tiling import VideoModalityTilingHelper
+from ltx_core.model.transformer.modality import Modality
+from ltx_core.tiling import TileCountConfig
+from ltx_core.tools import VideoLatentTools
+
+if TYPE_CHECKING:
+    from ltx_core.guidance.perturbations import BatchedPerturbationConfig
+
+
+class TiledDataParallelModelWrapper(torch.nn.Module):
+    """Wraps an ``X0Model`` for tiled data parallelism.
+    Tiles are distributed across ranks via round-robin, allowing more
+    tiles than GPUs (e.g. 16 tiles on 4 GPUs = 4 tiles per rank).
+    Each rank processes its assigned tiles sequentially, blending each
+    into a full-size accumulator.  A single ``all_reduce(SUM)`` after
+    all local tiles produces the final result (blend masks sum to 1
+    globally across all tiles).
+    Audio is processed untiled on every tile forward; the outputs are
+    summed via ``all_reduce`` and divided by the total tile count so
+    that all ranks stay in sync.
+    """
+
+    def __init__(
+        self,
+        model: torch.nn.Module,
+        *,
+        video_tools: VideoLatentTools,
+        tiling: TileCountConfig,
+        group: dist.ProcessGroup,
+        normalize_positions: bool = True,
+    ) -> None:
+        super().__init__()
+        self.model = model
+        self.group = group
+        self.world_size = dist.get_world_size(group)
+        self._normalize_positions = normalize_positions
+        self._helper = VideoModalityTilingHelper(tiling, video_tools)
+        all_tiles = self._helper.tiles
+        rank = dist.get_rank(group)
+        self._tiles = [t for i, t in enumerate(all_tiles) if i % self.world_size == rank]
+
+    @property
+    def num_blocks(self) -> int:
+        return self.model.num_blocks
+
+    def forward(
+        self,
+        video: Modality | None,
+        audio: Modality | None,
+        perturbations: BatchedPerturbationConfig | None,
+    ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
+        if video is None:
+            return self.model(video, audio, perturbations)
+
+        # Each rank processes its assigned tiles and accumulates locally.
+        denoised_video: torch.Tensor | None = None
+        denoised_audio: torch.Tensor | None = None
+        for tile in self._tiles:
+            tiled_video, ctx = self._helper.tile_modality(video, tile, normalize_positions=self._normalize_positions)
+            tile_out, audio_out = self.model(tiled_video, audio, perturbations)
+            blended = self._helper.blend(tile_out, tile, ctx)
+            denoised_video = blended if denoised_video is None else denoised_video + blended
+            if audio_out is not None:
+                denoised_audio = audio_out if denoised_audio is None else denoised_audio + audio_out
+
+        assert denoised_video is not None
+
+        # All-reduce: sum blended tiles across ranks (masks sum to 1 globally).
+        denoised_video = denoised_video.contiguous()
+        dist.all_reduce(denoised_video, op=dist.ReduceOp.SUM, group=self.group)
+
+        # Average audio across all tile forwards (each saw different video context).
+        if denoised_audio is not None:
+            total_tiles = len(self._helper.tiles)
+            denoised_audio = denoised_audio.contiguous()
+            dist.all_reduce(denoised_audio, op=dist.ReduceOp.SUM, group=self.group)
+            denoised_audio = denoised_audio / total_tiles
+
+        return denoised_video, denoised_audio

packages/ltx-core/src/ltx_core/multigpu/vae/__init__.py

@@ -0,0 +1,5 @@
+"""Multi-GPU utilities for VAE decoding."""
+
+from ltx_core.multigpu.vae.distributed_decoder import DistributedVideoDecoder
+
+__all__ = ["DistributedVideoDecoder"]

packages/ltx-core/src/ltx_core/multigpu/vae/distributed_decoder.py

@@ -0,0 +1,307 @@
+"""Distributed video decoder that partitions the latent across ranks.
+Tiles are assigned to ranks via round-robin, so the number of tiles
+may exceed the number of GPUs (e.g. 16 tiles on 4 GPUs = 4 tiles per
+rank).  Each rank decodes its assigned tiles sequentially.  Workers
+put their list of decoded tiles into a ``mp.Queue`` (CUDA IPC —
+zero-copy handle sharing).  The driver collects all tiles, blends
+overlap zones, and returns temporal batches distributed across devices.
+The tiling configuration comes from ``MGPUConfig.vae_tiling`` (set at
+construction time), NOT from the pipeline's SGPU tiling kwarg.  MGPU
+tiling controls parallelism; SGPU tiling controls single-GPU VRAM
+management — they are independent concerns.
+"""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Callable, Iterator
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+import torch
+import torch.distributed as dist
+from einops import rearrange
+from torch.multiprocessing import Queue
+
+from ltx_core.model.video_vae.tiling import TilingConfig
+from ltx_core.model.video_vae.video_vae import (
+    VideoDecoder,
+    map_spatial_slice,
+    map_temporal_slice,
+    to_mapping_operation,
+)
+from ltx_core.tiling import (
+    Tile,
+    create_tiles,
+    split_by_count,
+    split_by_count_temporal_causal,
+)
+from ltx_core.types import SpatioTemporalScaleFactors, VideoLatentShape
+
+if TYPE_CHECKING:
+    from ltx_core.tiling import TileCountConfig
+
+logger = logging.getLogger(__name__)
+
+
+# ------------------------------------------------------------------
+# Data structures
+# ------------------------------------------------------------------
+
+
+@dataclass(frozen=True)
+class DecodedTile:
+    """A VAE-decoded tile with pixel-space placement.
+    Attributes:
+        pixels: ``[F_tile, H_tile, W_tile, C]`` in the decoder's native dtype.
+        pixel_tile: Carries ``out_coords`` (f, h, w slices) and ``blend_mask``.
+    """
+
+    pixels: torch.Tensor
+    pixel_tile: Tile
+
+
+# ------------------------------------------------------------------
+# Tile construction helpers
+# ------------------------------------------------------------------
+
+
+def _to_decoded_tile(
+    raw: torch.Tensor,
+    tile: Tile,
+) -> DecodedTile:
+    """Convert raw decoder output ``[B, C, F, H, W]`` to a :class:`DecodedTile`.
+    Rearranges to ``[F, H, W, C]`` and normalises ``[-1, 1] → [0, 1]``.
+    """
+    pixels = rearrange(raw[0], "c f h w -> f h w c")
+    pixels = ((pixels + 1.0) / 2.0).clamp(0.0, 1.0)
+    return DecodedTile(pixels=pixels, pixel_tile=tile)
+
+
+# ------------------------------------------------------------------
+# Tile assembly
+# ------------------------------------------------------------------
+
+
+def compute_summed_weights(
+    tiles: list[DecodedTile],
+    total_frames: int,
+    output_height: int,
+    output_width: int,
+) -> torch.Tensor:
+    """Build the ``[F, H, W]`` denominator for weighted blending."""
+    weights = torch.zeros(total_frames, output_height, output_width)
+    for tile in tiles:
+        f_slice, h_slice, w_slice = tile.pixel_tile.out_coords
+        weights[f_slice, h_slice, w_slice] += tile.pixel_tile.blend_mask
+    return weights.clamp(min=1e-8)
+
+
+def gather_frames(
+    tiles: list[DecodedTile],
+    total_frames: int,
+    output_height: int,
+    output_width: int,
+    num_temporal_batches: int,
+    world_size: int,
+    weights: torch.Tensor,
+    device_fn: Callable[[int], str | torch.device] | None = None,
+) -> Iterator[torch.Tensor]:
+    """Assemble decoded tiles into temporal batches distributed across GPUs.
+    Each temporal batch is allocated on the device returned by *device_fn(batch_index)*.
+    By default batches are placed round-robin on ``cuda:0`` … ``cuda:<world_size-1>``.
+    """
+    if device_fn is None:
+        device_fn = lambda b: f"cuda:{b % world_size}"  # noqa: E731
+
+    batch_size = (total_frames + num_temporal_batches - 1) // num_temporal_batches
+
+    for b in range(num_temporal_batches):
+        batch_range = slice(b * batch_size, min((b + 1) * batch_size, total_frames))
+        batch_len = batch_range.stop - batch_range.start
+        if batch_len <= 0:
+            break
+
+        device = device_fn(b)
+        dtype = tiles[0].pixels.dtype
+        output = torch.zeros(batch_len, output_height, output_width, 3, device=device, dtype=dtype)
+
+        for tile in tiles:
+            f_slice, h_slice, w_slice = tile.pixel_tile.out_coords
+
+            overlap = slice(max(batch_range.start, f_slice.start), min(batch_range.stop, f_slice.stop))
+            if overlap.start >= overlap.stop:
+                continue
+
+            tile_frames = slice(overlap.start - f_slice.start, overlap.stop - f_slice.start)
+            out_frames = slice(overlap.start - batch_range.start, overlap.stop - batch_range.start)
+
+            blend = tile.pixel_tile.blend_mask[tile_frames].to(device=device)
+            output[out_frames, h_slice, w_slice, :] += tile.pixels[tile_frames].to(device=device) * blend[:, :, :, None]
+
+        batch_weights = weights[batch_range.start : batch_range.stop].to(device=device)
+        output.div_(batch_weights[:, :, :, None])
+        yield output
+
+
+# ------------------------------------------------------------------
+# Main class
+# ------------------------------------------------------------------
+
+
+class DistributedVideoDecoder(torch.nn.Module):
+    """Distributed VAE decoder with queue-based tile collection.
+    All ranks decode their latent tile in parallel.  Workers send
+    their :class:`DecodedTile` to the driver rank via the shared
+    ``mp.Queue`` (CUDA IPC — zero-copy).  The driver collects all
+    tiles, blends overlapping regions, and returns temporal batches
+    as an iterator.
+    Parameters
+    ----------
+    decoder:
+        The real (local) ``VideoDecoder`` instance.
+    queue:
+        ``mp.Queue`` shared across all ranks for CUDA IPC tile transfer.
+    vae_group:
+        NCCL process group for the VAE ranks. Used to derive
+        ``rank`` and ``world_size`` within the group.
+    vae_tiling:
+        MGPU tiling config that determines how the latent is split.
+    driver_rank:
+        Group-local rank of the driver process (the rank that collects
+        and assembles tiles).
+    """
+
+    def __init__(
+        self,
+        decoder: VideoDecoder,
+        queue: Queue,  # type: ignore[type-arg]
+        vae_group: dist.ProcessGroup,
+        vae_tiling: TileCountConfig,
+        driver_rank: int = 0,
+    ) -> None:
+        super().__init__()
+        self.decoder = decoder
+        self.queue = queue
+        self.vae_group = vae_group
+        self.rank = dist.get_rank(vae_group)
+        self.world_size = dist.get_world_size(vae_group)
+        self.vae_tiling = vae_tiling
+        self.driver_rank = driver_rank
+
+    def forward(
+        self,
+        sample: torch.Tensor,
+        timestep: torch.Tensor | None = None,
+        generator: torch.Generator | None = None,
+    ) -> torch.Tensor:
+        """Non-tiled path: fall back to local decode."""
+        return self.decoder(sample, timestep, generator)
+
+    def decode_video(
+        self,
+        latent: torch.Tensor,
+        tiling_config: TilingConfig | None = None,
+        generator: torch.Generator | None = None,
+        device_fn: Callable[[int], str | torch.device] | None = None,
+    ) -> Iterator[torch.Tensor]:
+        """Distributed decode — all ranks decode, driver assembles.
+        Not a generator so that worker side-effects (decode + queue.put)
+        execute eagerly regardless of whether the caller iterates.
+        1. Each rank decodes its latent tile (with optional intra-GPU tiling).
+        2. Workers send their :class:`DecodedTile` to the driver via the queue.
+        3. The driver collects all tiles, blends overlaps, and returns
+           temporal batches distributed across GPUs.
+        """
+        if (
+            self.vae_tiling.frames.num_tiles > 1
+            and tiling_config is not None
+            and tiling_config.temporal_config is not None
+        ):
+            raise ValueError(
+                "Cannot combine multi-GPU temporal tiling (vae_tiling.frames.num_tiles > 1) "
+                "with single-GPU temporal tiling (tiling_config.temporal_config). "
+                "Use only one to avoid causal decoding conflicts."
+            )
+
+        latent_shape = VideoLatentShape.from_torch_shape(latent.shape)
+        scale = self.decoder.video_downscale_factors
+        full_shape = latent_shape.upscale(scale)
+
+        # Phase 1: each rank decodes its assigned tiles.
+        my_tiles = self._decode_tiles(latent, latent_shape, scale, generator, tiling_config)
+
+        # Phase 2: workers send tiles to driver.
+        if self.rank != self.driver_rank:
+            self.queue.put((self.rank, my_tiles))
+            return iter([])
+
+        # Phase 3: driver collects and assembles.
+        all_tiles = self._collect_tiles(my_tiles)
+        weights = compute_summed_weights(all_tiles, full_shape.frames, full_shape.height, full_shape.width)
+        batches = gather_frames(
+            all_tiles,
+            full_shape.frames,
+            full_shape.height,
+            full_shape.width,
+            self.world_size,
+            self.world_size,
+            weights,
+            device_fn=device_fn,
+        )
+        return batches
+
+    # ------------------------------------------------------------------
+    # Private helpers
+    # ------------------------------------------------------------------
+
+    def _decode_tiles(
+        self,
+        latent: torch.Tensor,
+        latent_shape: VideoLatentShape,
+        scale: SpatioTemporalScaleFactors,
+        generator: torch.Generator | None,
+        tiling_config: TilingConfig | None = None,
+    ) -> list[DecodedTile]:
+        """Decode this rank's assigned latent tiles and convert to :class:`DecodedTile` list."""
+        all_tiles = create_tiles(
+            torch.Size([latent_shape.frames, latent_shape.height, latent_shape.width]),
+            splitters=[
+                split_by_count_temporal_causal(self.vae_tiling.frames.num_tiles, self.vae_tiling.frames.overlap),
+                split_by_count(self.vae_tiling.height.num_tiles, self.vae_tiling.height.overlap),
+                split_by_count(self.vae_tiling.width.num_tiles, self.vae_tiling.width.overlap),
+            ],
+            mappers=[
+                to_mapping_operation(map_temporal_slice, scale.time),
+                to_mapping_operation(map_spatial_slice, scale.height),
+                to_mapping_operation(map_spatial_slice, scale.width),
+            ],
+        )
+        my_tiles = [t for i, t in enumerate(all_tiles) if i % self.world_size == self.rank]
+        decoded = []
+        for tile in my_tiles:
+            latent_slice = latent[:, :, tile.in_coords[0], tile.in_coords[1], tile.in_coords[2]]
+            if tiling_config is not None:
+                chunks = list(self.decoder.tiled_decode(latent_slice, tiling_config, generator=generator))
+                raw = torch.cat(chunks, dim=2)
+            else:
+                raw = self.decoder.forward(latent_slice, generator=generator)
+            decoded.append(_to_decoded_tile(raw, tile))
+        return decoded
+
+    def _collect_tiles(self, driver_tiles: list[DecodedTile]) -> list[DecodedTile]:
+        """Collect tiles from all workers via the queue. Returns flat list of all tiles.
+        Sorted by rank so the downstream reduction in ``gather_frames`` /
+        ``compute_summed_weights`` (in-place ``+=`` over overlapping pixel
+        regions) processes tiles in a fixed order. Queue-arrival order would
+        otherwise vary run-to-run and yield 1-ulp bf16 drift from
+        non-associative floating-point summation.
+        """
+        per_rank: dict[int, list[DecodedTile]] = {self.driver_rank: driver_tiles}
+        for _ in range(self.world_size - 1):
+            worker_rank, worker_tiles = self.queue.get()
+            per_rank[worker_rank] = worker_tiles
+        result: list[DecodedTile] = []
+        for rank in sorted(per_rank):
+            result.extend(per_rank[rank])
+        return result

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

@@ -0,0 +1,47 @@
+"""Public API for blockwise FP8/FP6 quantization.
+The implementation lives in :mod:`._impl`, which imports the compiled
+``ltx_kernels.blockwise`` kernels at top level. This module deliberately defers
+that import so that ``ltx_core.quantization.blockwise`` remains importable
+without those kernels built; the gate fires only when one of the policy builders
+is actually called.
+"""
+
+from ltx_core.quantization.policy import QuantizationPolicy
+
+__all__ = ["build_fp6_policy", "build_fp8_policy"]
+
+
+def _import_impl():  # noqa: ANN202 - internal helper
+    try:
+        from ltx_core.quantization.blockwise import _impl  # noqa: PLC0415
+
+        return _impl
+    except ImportError as e:
+        raise RuntimeError(
+            "ltx-kernels not built; blockwise FP8/FP6 quantization requires it. "
+            "Build it on a CUDA host with `uv sync --group kernels` (or "
+            "`uv pip install -e packages/ltx-kernels --no-build-isolation`) before "
+            "calling build_fp8_policy() / build_fp6_policy()."
+        ) from e
+
+
+def build_fp8_policy() -> QuantizationPolicy:
+    """Build a blockwise FP8 quantization policy. Raises ``RuntimeError`` if ``ltx-kernels`` is not built."""
+    impl = _import_impl()
+    return QuantizationPolicy(
+        sd_ops=impl.build_sd_ops_fp8(),
+        module_ops=(impl.build_module_ops_fp8(),),
+        model_configurator=impl.BlockwiseFP8LTXModelConfigurator,
+        fuse_rule=impl.fuse_rule_fp8,
+    )
+
+
+def build_fp6_policy() -> QuantizationPolicy:
+    """Build a blockwise FP6 quantization policy. Raises ``RuntimeError`` if ``ltx-kernels`` is not built."""
+    impl = _import_impl()
+    return QuantizationPolicy(
+        sd_ops=impl.build_sd_ops_fp6(),
+        module_ops=(impl.build_module_ops_fp6(),),
+        model_configurator=impl.BlockwiseFP6LTXModelConfigurator,
+        fuse_rule=impl.fuse_rule_fp6,
+    )

packages/ltx-core/src/ltx_core/quantization/blockwise/_impl.py

@@ -0,0 +1,431 @@
+"""Implementation of blockwise FP8/FP6 quantization. Depends on ``ltx_kernels``.
+This module imports the compiled ``ltx_kernels.blockwise`` kernels at top level
+— without them built, simply importing this file raises :class:`ImportError`.
+The intended access path is through ``ltx_core.quantization.blockwise.__init__``
+which catches that and re-raises as a clean :class:`RuntimeError`. Do not import
+this module directly from non-quantization code.
+"""
+
+from typing import Callable, ClassVar, List, NamedTuple, Protocol, Type
+
+import torch
+from ltx_kernels.blockwise.functional import (
+    blockwise_dequantize,
+    blockwise_quantize_adanorm_triton,
+    blockwise_quantize_rms_fma_triton,
+    fp6_blockwise_quantize_weights_torch,
+    fp6_pack_tensor,
+    fp6_unpack_tensor,
+    fp8_blockwise_quantize_weights_torch,
+    gated_attention_triton,
+    rms_norm_rope,
+    rms_norm_split_rope,
+)
+from ltx_kernels.blockwise.linear import BlockwiseFP6Linear, BlockwiseFP8Linear
+from torch import nn
+
+from ltx_core.loader.fuse_loras import FuseRule, bf16_fuse_rule
+from ltx_core.loader.module_ops import ModuleOps
+from ltx_core.loader.primitives import StateDict
+from ltx_core.loader.sd_ops import KeyValueOperationResult, SDOps
+from ltx_core.model.model_protocol import ModelConfigurator
+from ltx_core.model.transformer import LTXModel
+from ltx_core.model.transformer.model_configurator import LTXModelConfigurator, LTXVideoOnlyModelConfigurator
+from ltx_core.model.transformer.ops import (
+    AdaZeroCallable,
+    GatedAttentionCallable,
+    PostSACallable,
+    PreAttentionCallable,
+)
+from ltx_core.model.transformer.rope import LTXRopeType
+from ltx_core.model.transformer.transformer import TransformerOpsConfig
+
+
+class FromLinearProtocol(Protocol):
+    """Protocol for nn.Module subclasses that can be constructed from an nn.Linear."""
+
+    @classmethod
+    def from_linear(cls, linear: nn.Linear, transform_weights: bool = True) -> nn.Module: ...
+
+
+class BlockwiseQuantizedWeight(NamedTuple):
+    """Result of blockwise quantization: a quantized weight tensor and its per-block scale.
+    For FP8: ``weight`` is ``float8_e4m3fn``, ``scale`` is ``float32`` shaped
+    ``[out // 128, in // 128]``.
+    For FP6: ``weight`` is packed ``uint8`` shaped ``[out, (in // 4) * 3]``,
+    ``scale`` is ``float32`` shaped ``[out // 128, in // 128]``.
+    """
+
+    weight: torch.Tensor
+    scale: torch.Tensor
+
+
+EXCLUDED_LAYER_SUBSTRINGS = (
+    "patchify_proj",
+    "adaln_single",
+    "av_ca_video_scale_shift_adaln_single",
+    "av_ca_a2v_gate_adaln_single",
+    "caption_projection",
+    "proj_out",
+    "audio_patchify_proj",
+    "audio_adaln_single",
+    "av_ca_audio_scale_shift_adaln_single",
+    "av_ca_v2a_gate_adaln_single",
+    "audio_caption_projection",
+    "audio_proj_out",
+    "to_gate_logits",
+    "scale_shift_table",
+)
+
+
+_QUANTIZABLE_FLOAT_DTYPES = (torch.bfloat16, torch.float16, torch.float32)
+
+
+def _is_quantizable_float(x: torch.Tensor | torch.dtype) -> bool:
+    """Whether ``x`` is an unquantized high-precision float (bf16 / fp16 / fp32).
+    FP8 / FP6 weights are floats too but they're already in a quantized layout
+    and must not be re-quantized.
+    """
+    dtype = x.dtype if isinstance(x, torch.Tensor) else x
+    return dtype in _QUANTIZABLE_FLOAT_DTYPES
+
+
+def _should_skip_layer(layer_name: str, excluded_layer_substrings: tuple[str, ...]) -> bool:
+    return any(substring in layer_name for substring in excluded_layer_substrings)
+
+
+def _replace_linear_modules(model: torch.nn.Module, linear_cls: Type[FromLinearProtocol]) -> torch.nn.Module:
+    skip_list = ["to_gate_logits", "scale_shift_table"]
+    for name, module in model.named_modules():
+        if "transformer_block" in name and isinstance(module, torch.nn.Linear):
+            if _should_skip_layer(name, skip_list):
+                continue
+            *parent_path, child_name = name.split(".")
+            parent = model
+            for part in parent_path:
+                parent = getattr(parent, part)
+            setattr(
+                parent,
+                child_name,
+                linear_cls.from_linear(module, False),
+            )
+            del module.weight
+            del module.bias
+            torch.cuda.empty_cache()
+    return model
+
+
+# ---------------------------------------------------------------------------
+# Weight quantization helpers
+# ---------------------------------------------------------------------------
+
+
+def _blockwise_quantize_weight_helper(
+    value: torch.Tensor,
+    quant_fn: Callable[[torch.Tensor, int], tuple[torch.Tensor, torch.Tensor]],
+    pack_fn: Callable[[torch.Tensor], torch.Tensor],
+) -> BlockwiseQuantizedWeight:
+    orig_device = value.device
+    w_quant, w_scales = quant_fn(value.cuda())
+    return BlockwiseQuantizedWeight(
+        weight=pack_fn(w_quant).to(device=orig_device),
+        scale=w_scales.to(device=orig_device),
+    )
+
+
+def _fp8_blockwise_quantize_weight(value: torch.Tensor) -> BlockwiseQuantizedWeight:
+    return _blockwise_quantize_weight_helper(value, fp8_blockwise_quantize_weights_torch, lambda x: x)
+
+
+def _fp6_blockwise_quantize_weight(value: torch.Tensor) -> BlockwiseQuantizedWeight:
+    return _blockwise_quantize_weight_helper(value, fp6_blockwise_quantize_weights_torch, fp6_pack_tensor)
+
+
+def _create_weight_quantize_op(
+    excluded_layer_substrings: tuple[str, ...],
+    quantization_func: Callable[[torch.Tensor], BlockwiseQuantizedWeight],
+) -> Callable[[str, torch.Tensor], list[KeyValueOperationResult]]:
+    """KeyValueOperation that blockwise-quantizes a 2D BF16 ``.weight`` and emits ``.weight_scale``."""
+
+    def quantize_weight(key: str, value: torch.Tensor) -> list[KeyValueOperationResult]:
+        if _should_skip_layer(key, excluded_layer_substrings):
+            return [KeyValueOperationResult(key, value)]
+        if value.dim() != 2 or not _is_quantizable_float(value):
+            return [KeyValueOperationResult(key, value)]
+        quantized = quantization_func(value)
+        scale_key = key.replace(".weight", ".weight_scale")
+        return [
+            KeyValueOperationResult(key, quantized.weight),
+            KeyValueOperationResult(scale_key, quantized.scale),
+        ]
+
+    return quantize_weight
+
+
+def _create_bias_to_fp32_op(
+    excluded_layer_substrings: tuple[str, ...],
+) -> Callable[[str, torch.Tensor], list[KeyValueOperationResult]]:
+    """KeyValueOperation that casts a ``.bias`` tensor to FP32.
+    ``BlockwiseFP{8,6}Linear`` registers ``.bias`` as float32; the load-time
+    cast keeps the checkpoint's BF16 bias compatible with that param dtype.
+    """
+
+    def bias_to_fp32(key: str, value: torch.Tensor) -> list[KeyValueOperationResult]:
+        if _should_skip_layer(key, excluded_layer_substrings):
+            return [KeyValueOperationResult(key, value)]
+        return [KeyValueOperationResult(key, value.float())]
+
+    return bias_to_fp32
+
+
+# ---------------------------------------------------------------------------
+# Q8 activation callables (formerly in model.transformer.ops)
+# ---------------------------------------------------------------------------
+
+
+class Q8KernelsPreAttention(PreAttentionCallable):
+    def __call__(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        attn_module: nn.Module,
+        mask: torch.Tensor | None,  # noqa: ARG002
+        pe: torch.Tensor | None,
+        k_pe: torch.Tensor | None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if attn_module.rope_type == LTXRopeType.INTERLEAVED:
+            rope_func = rms_norm_rope
+        elif attn_module.rope_type == LTXRopeType.SPLIT:
+            rope_func = rms_norm_split_rope
+        else:
+            raise ValueError(f"Invalid rope type: {attn_module.rope_type}")
+
+        if pe is not None:
+            k_pe = k_pe if k_pe is not None else pe
+            q = rope_func(q, pe[0], pe[1], attn_module.q_norm.weight, False)
+            k = rope_func(k, k_pe[0], k_pe[1], attn_module.k_norm.weight, False)
+        else:
+            q = attn_module.q_norm(q)
+            k = attn_module.k_norm(k)
+        return q, k
+
+
+class Q8KernelsAdaZeroFunction(AdaZeroCallable):
+    def __call__(
+        self,
+        x: torch.Tensor,
+        eps: float,  # noqa: ARG002
+        scale: torch.Tensor,
+        shift: torch.Tensor,
+    ) -> torch.Tensor:
+        return blockwise_quantize_adanorm_triton(x, None, scale, shift, torch.float8_e4m3fn, 1.0)
+
+
+class Q8KernelsPostSAFunction(PostSACallable):
+    def __call__(
+        self,
+        x: torch.Tensor,
+        y: torch.Tensor,
+        norm_weights: torch.Tensor | None,  # noqa: ARG002
+        eps: float,  # noqa: ARG002
+        gate: torch.Tensor,
+    ) -> List[torch.Tensor]:
+        # Dequantize the fused result: the cross-attention AdaLN path applies a BF16
+        # scale/shift, which cannot operate on the (fp8, scales) payload.
+        normed_fp8 = blockwise_quantize_rms_fma_triton(x, y, gate)
+        return x, blockwise_dequantize(normed_fp8)
+
+
+class Q8KernelsGatedAttention(GatedAttentionCallable):
+    def __call__(
+        self,
+        x: torch.Tensor,
+        attn_out: torch.Tensor,
+        attn_module: nn.Module,
+    ) -> torch.Tensor:
+        # Self-attention path: ``x`` arrives as the ``(fp8, scales)`` tuple
+        # produced by Q8KernelsAdaZeroFunction. Cross-attention path
+        # (apply_cross_attention_adaln) feeds plain BF16, so dequantize only
+        # when needed.
+        if isinstance(x, tuple):
+            x = blockwise_dequantize(x)
+        gate_logits = attn_module.to_gate_logits(x)
+        return gated_attention_triton(attn_out, gate_logits)
+
+
+# ---------------------------------------------------------------------------
+# Fuse rules
+# ---------------------------------------------------------------------------
+
+
+_BLOCK = 128
+
+
+def _blockwise_dequantize_2d(weight_fp8: torch.Tensor, weight_scale: torch.Tensor) -> torch.Tensor:
+    """Dequantize a 2D blockwise-FP8 weight ``[out, in]`` with per-block scale
+    ``[out//128, in//128]`` to BF16.
+    ``ltx_kernels.blockwise.blockwise_dequantize`` is built for 3D activations where
+    scales are ``[b*s, in//128]`` — one row per token. Weights are block-
+    quantized along the row dim too, so we expand the row axis 128x via
+    ``repeat_interleave`` and reuse the kernel.
+    """
+    out_features, in_features = weight_fp8.shape
+    scales_per_row = weight_scale.repeat_interleave(_BLOCK, dim=0)
+    return blockwise_dequantize((weight_fp8.unsqueeze(0), scales_per_row)).view(out_features, in_features)
+
+
+def _blockwise_fp8_fuse(
+    key: str,
+    weight: torch.Tensor,
+    deltas: torch.Tensor,
+    model_sd: StateDict,
+) -> dict[str, torch.Tensor]:
+    """Dequantize the FP8 weight + per-block scale to BF16, add the BF16 delta,
+    and re-quantize blockwise. Both ``.weight`` and the companion
+    ``.weight_scale`` are emitted so the loaded layer matches what
+    ``BlockwiseFP8Linear`` expects.
+    Excluded layers (see ``EXCLUDED_LAYER_SUBSTRINGS``) stay BF16 and have no
+    ``.weight_scale`` companion — for those, fall back to a plain bf16 fuse.
+    """
+    scale_key = key.replace(".weight", ".weight_scale")
+    if scale_key not in model_sd.sd:
+        return bf16_fuse_rule(key, weight, deltas, model_sd)
+    weight_scale = model_sd.sd[scale_key]
+    bf16_weight = _blockwise_dequantize_2d(weight, weight_scale)
+    merged = bf16_weight + deltas.to(dtype=bf16_weight.dtype)
+    new_fp8_weight, new_weight_scale = fp8_blockwise_quantize_weights_torch(merged.cuda())
+    return {
+        key: new_fp8_weight.to(device=weight.device),
+        scale_key: new_weight_scale.to(device=weight.device),
+    }
+
+
+def _blockwise_fp6_fuse(
+    key: str,
+    weight: torch.Tensor,
+    deltas: torch.Tensor,
+    model_sd: StateDict,
+) -> dict[str, torch.Tensor]:
+    """Mirror ``BlockwiseFP6Linear.fp8weight`` for the dequant side: unpack the
+    packed ``uint8`` weight to ``float8_e4m3fn``, dequantize via the per-block
+    scale to BF16, add the BF16 delta, re-quantize to FP6, and pack back to
+    uint8. Both ``.weight`` (packed uint8) and ``.weight_scale`` are emitted.
+    Note: ``fp6_unpack_tensor`` restores the dropped e_1/e_2 exponent bits as 0,
+    so the dequant->add->requant round-trip is lossy on those bits even when no
+    LoRA delta is applied. This matches what ``BlockwiseFP6Linear`` already does
+    at inference time via its ``fp8weight`` property, so the fused weight is
+    numerically consistent with the unfused inference path.
+    Excluded layers (see ``EXCLUDED_LAYER_SUBSTRINGS``) stay BF16 and have no
+    ``.weight_scale`` companion — for those, fall back to a plain bf16 fuse.
+    """
+    scale_key = key.replace(".weight", ".weight_scale")
+    if scale_key not in model_sd.sd:
+        return bf16_fuse_rule(key, weight, deltas, model_sd)
+    weight_scale = model_sd.sd[scale_key]
+    # Packed shape is [out, (in // 4) * 3]; recover in_features.
+    original_n = weight.shape[-1] * 4 // 3
+    fp8_view = fp6_unpack_tensor(weight, original_n).view(torch.float8_e4m3fn)
+    bf16_weight = _blockwise_dequantize_2d(fp8_view, weight_scale)
+    merged = bf16_weight + deltas.to(dtype=bf16_weight.dtype)
+    new_fp8, new_scale = fp6_blockwise_quantize_weights_torch(merged)
+    new_packed = fp6_pack_tensor(new_fp8.view(torch.uint8))
+    return {
+        key: new_packed.to(device=weight.device),
+        scale_key: new_scale.to(device=weight.device),
+    }
+
+
+# ---------------------------------------------------------------------------
+# Configurators (TransformerOpsConfig with Q8 activation callables)
+# ---------------------------------------------------------------------------
+
+
+def _build_blockwise_ops_config() -> TransformerOpsConfig:
+    return TransformerOpsConfig.from_functions(
+        preattention=Q8KernelsPreAttention(),
+        gated_attention=Q8KernelsGatedAttention(),
+        ada_zero=Q8KernelsAdaZeroFunction(),
+        post_sa=Q8KernelsPostSAFunction(),
+    )
+
+
+# FP6 is weight-only; activation ops match FP8.
+_BLOCKWISE_OPS = _build_blockwise_ops_config()
+
+
+class BlockwiseFP8LTXModelConfigurator(ModelConfigurator[LTXModel]):
+    BASE: ClassVar[type[ModelConfigurator[LTXModel]]] = LTXModelConfigurator
+    OPS: ClassVar[TransformerOpsConfig] = _BLOCKWISE_OPS
+
+    @classmethod
+    def from_config(cls, config: dict) -> LTXModel:
+        return cls.BASE.from_config(config, ops=cls.OPS)
+
+
+class BlockwiseFP8LTXVideoOnlyModelConfigurator(BlockwiseFP8LTXModelConfigurator):
+    BASE = LTXVideoOnlyModelConfigurator
+
+
+class BlockwiseFP6LTXModelConfigurator(BlockwiseFP8LTXModelConfigurator):
+    pass
+
+
+class BlockwiseFP6LTXVideoOnlyModelConfigurator(BlockwiseFP8LTXVideoOnlyModelConfigurator):
+    pass
+
+
+# ---------------------------------------------------------------------------
+# SDOps / ModuleOps / FuseRule assembly
+# ---------------------------------------------------------------------------
+
+
+def build_sd_ops_fp8() -> SDOps:
+    return (
+        SDOps("blockwise_fp8_weights")
+        .with_kv_operation(
+            _create_weight_quantize_op(EXCLUDED_LAYER_SUBSTRINGS, _fp8_blockwise_quantize_weight),
+            key_prefix="transformer_blocks.",
+            key_suffix=".weight",
+        )
+        .with_kv_operation(
+            _create_bias_to_fp32_op(EXCLUDED_LAYER_SUBSTRINGS),
+            key_prefix="transformer_blocks.",
+            key_suffix=".bias",
+        )
+    )
+
+
+def build_sd_ops_fp6() -> SDOps:
+    return (
+        SDOps("blockwise_fp6_weights")
+        .with_kv_operation(
+            _create_weight_quantize_op(EXCLUDED_LAYER_SUBSTRINGS, _fp6_blockwise_quantize_weight),
+            key_prefix="transformer_blocks.",
+            key_suffix=".weight",
+        )
+        .with_kv_operation(
+            _create_bias_to_fp32_op(EXCLUDED_LAYER_SUBSTRINGS),
+            key_prefix="transformer_blocks.",
+            key_suffix=".bias",
+        )
+    )
+
+
+def build_module_ops_fp8() -> ModuleOps:
+    return ModuleOps(
+        name="blockwise_fp8_prepare_for_loading",
+        matcher=lambda model: isinstance(model, LTXModel),
+        mutator=lambda model: _replace_linear_modules(model, BlockwiseFP8Linear),
+    )
+
+
+def build_module_ops_fp6() -> ModuleOps:
+    return ModuleOps(
+        name="blockwise_fp6_prepare_for_loading",
+        matcher=lambda model: isinstance(model, LTXModel),
+        mutator=lambda model: _replace_linear_modules(model, BlockwiseFP6Linear),
+    )
+
+
+fuse_rule_fp8 = FuseRule(aggregation_dtype=torch.bfloat16, fuse_fn=_blockwise_fp8_fuse)
+fuse_rule_fp6 = FuseRule(aggregation_dtype=torch.bfloat16, fuse_fn=_blockwise_fp6_fuse)

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

@@ -10,7 +10,6 @@ from ltx_core.loader.module_ops import ModuleOps
 from ltx_core.loader.primitives import StateDict
 from ltx_core.model.transformer import LTXModel
 from ltx_core.quantization.policy import QuantizationPolicy
-from ltx_core.quantization.trtllm_scaled_usable import trtllm_scaled_mm_usable
 
 
 def _read_safetensors_dtypes(path: str) -> dict[str, str]:
@@ -50,34 +49,21 @@ class FP8Linear(nn.Module):
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         origin_shape = x.shape
 
-        if trtllm_scaled_mm_usable():
-            qinput, cur_input_scale = torch.ops.tensorrt_llm.static_quantize_e4m3_per_tensor(x, self.input_scale)
-            if qinput.dim() == 3:
-                qinput = qinput.reshape(-1, qinput.shape[-1])
-            output = torch.ops.trtllm.cublas_scaled_mm(
-                qinput,
-                self.weight.t(),
-                scale_a=cur_input_scale,
-                scale_b=self.weight_scale,
-                bias=None,
-                out_dtype=x.dtype,
-            )
-        else:
-            # Clamp before cast: out-of-range values cast to NaN/saturated FP8, which
-            # produces black-screen output on some checkpoints (e.g. ltx-2-19b-dev-fp8).
-            fp8_min = torch.finfo(torch.float8_e4m3fn).min
-            fp8_max = torch.finfo(torch.float8_e4m3fn).max
-            qinput = torch.clamp(x * self.input_scale.reciprocal(), fp8_min, fp8_max).to(torch.float8_e4m3fn)
-            if qinput.dim() == 3:
-                qinput = qinput.reshape(-1, qinput.shape[-1])
-            output = torch._scaled_mm(
-                qinput,
-                self.weight.t(),
-                scale_a=self.input_scale,
-                scale_b=self.weight_scale,
-                out_dtype=x.dtype,
-                use_fast_accum=True,
-            )
+        # Clamp before cast: out-of-range values cast to NaN/saturated FP8, which
+        # produces black-screen output on some checkpoints (e.g. ltx-2-19b-dev-fp8).
+        fp8_min = torch.finfo(torch.float8_e4m3fn).min
+        fp8_max = torch.finfo(torch.float8_e4m3fn).max
+        qinput = torch.clamp(x * self.input_scale.reciprocal(), fp8_min, fp8_max).to(torch.float8_e4m3fn)
+        if qinput.dim() == 3:
+            qinput = qinput.reshape(-1, qinput.shape[-1])
+        output = torch._scaled_mm(
+            qinput,
+            self.weight.t(),
+            scale_a=self.input_scale,
+            scale_b=self.weight_scale,
+            out_dtype=x.dtype,
+            use_fast_accum=True,
+        )
 
         if self.bias is not None:
             output = output + self.bias.to(output.dtype)

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

@@ -1,37 +0,0 @@
-"""Runtime detection of TensorRT-LLM FP8 scaled-matmul availability.
-When the TRT-LLM ops are usable on the current host (Linux + Hopper-class CUDA
-+ tensorrt_llm wheel installed) we use them since they outperform the PyTorch-native
-``torch._scaled_mm`` path. Otherwise we fall back to the native implementation,
-which is portable across platforms (Windows, macOS, AMD GPUs).
-The check runs once and is cached.
-"""
-
-from __future__ import annotations
-
-import platform
-from functools import cache
-
-import torch
-
-
-@cache
-def trtllm_scaled_mm_usable() -> bool:
-    if platform.system() != "Linux":
-        return False
-
-    if not torch.cuda.is_available():
-        return False
-
-    major, minor = torch.cuda.get_device_capability()
-    sm = major * 10 + minor
-
-    if sm < 90 or sm >= 120:
-        return False
-
-    # The import is load-bearing — registers the trtllm torch ops as a side effect.
-    try:
-        import tensorrt_llm  # noqa: F401, PLC0415
-    except Exception:
-        return False
-
-    return True

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

@@ -75,7 +75,9 @@ class GemmaTextEncoder(torch.nn.Module):
         pad_token_id = self.processor.tokenizer.pad_token_id if self.processor.tokenizer.pad_token_id is not None else 0
         model_inputs = _pad_inputs_for_attention_alignment(model_inputs, pad_token_id=pad_token_id)
 
-        with torch.inference_mode(), torch.random.fork_rng(devices=[self.model.device]):
+        # fork_rng device pinning is only supported for CUDA; MPS/CPU fork CPU RNG only.
+        fork_devices = [self.model.device] if self.model.device.type == "cuda" else []
+        with torch.inference_mode(), torch.random.fork_rng(devices=fork_devices):
             torch.manual_seed(seed)
             outputs = self.model.generate(
                 **model_inputs,

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

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import itertools
+import math
 from dataclasses import dataclass, replace
 from typing import Callable, NamedTuple
 
@@ -462,3 +463,16 @@ class TileCountConfig:
     frames: DimensionTilingConfig = DimensionTilingConfig(num_tiles=1, overlap=0)
     height: DimensionTilingConfig = DimensionTilingConfig(num_tiles=1, overlap=0)
     width: DimensionTilingConfig = DimensionTilingConfig(num_tiles=1, overlap=0)
+
+
+def balanced_tile_split(num_tiles: int) -> tuple[int, int]:
+    """Factor ``num_tiles`` into ``(small, large)`` as square as possible.
+    ``small`` is the largest divisor not exceeding the square root, so
+    ``small * large == num_tiles`` and ``small <= large``. E.g. 2 -> (1, 2),
+    4 -> (2, 2), 8 -> (2, 4), 16 -> (4, 4). The caller decides which tiled
+    dimension gets which factor.
+    """
+    if num_tiles < 1:
+        raise ValueError(f"num_tiles must be >= 1, got {num_tiles}")
+    small = next(d for d in range(math.isqrt(num_tiles), 0, -1) if num_tiles % d == 0)
+    return small, num_tiles // small