Merge pull request #248 from Lightricks/pr-2026-07-07-be3b401

Public sync - 2026-07-07

.gitignore

@@ -44,12 +44,18 @@ wandb/
 *.wav
 *.webp
 
-# HDR IC-LoRA e2e test baseline (checked in via Git LFS)
-!packages/ltx-pipelines/tests/assets/expected_hdr_ic_lora_exr/frame_*.exr
+# Full-params expected results for the --full-e2e quality lane (checked in via Git LFS)
+!packages/ltx-pipelines/tests/assets/full_e2e/
+!packages/ltx-pipelines/tests/assets/full_e2e/*.mp4
+!packages/ltx-pipelines/tests/assets/full_e2e/*.wav
+!packages/ltx-pipelines/tests/assets/full_e2e/expected_hdr_ic_lora_exr/frame_*.exr
+
+# HDR IC-LoRA e2e test input clip (checked in via Git LFS)
 !packages/ltx-pipelines/tests/assets/hdr_ic_lora_test_input.mp4
 
-# Text-to-audio (T2A) e2e test baseline (checked in via Git LFS)
-!packages/ltx-pipelines/tests/assets/expected_t2a_one_stage_ltx2_3.wav
+# Fast integration-profile goldens (bit-exact decoded video/audio; checked in via Git LFS)
+!packages/ltx-pipelines/tests/assets/integration/
+!packages/ltx-pipelines/tests/assets/integration/*.safetensors
 
 # ltx-bench Grafana dashboards (source of truth in the repo)
 !packages/ltx-bench/grafana/*.json

README.md

@@ -14,19 +14,43 @@
 
 ## 🚀 Quick Start
 
+Clone the repo
+
 ```bash
-# Clone the repository
 git clone https://github.com/Lightricks/LTX-2.git
 cd LTX-2
-
-# Set up the environment
-uv sync --frozen
-source .venv/bin/activate
 ```
 
-### Required Models
+Download the relevant [models](https://huggingface.co/Lightricks/LTX-2.3) or use the [Hugging Face CLI](https://huggingface.co/docs/huggingface_hub/guides/cli)
 
-Download the following models from the [LTX-2.3 HuggingFace repository](https://huggingface.co/Lightricks/LTX-2.3):
+```bash
+hf auth login
+hf download Lightricks/LTX-2.3 \
+    ltx-2.3-22b-distilled-1.1.safetensors ltx-2.3-spatial-upscaler-x2-1.1.safetensors --local-dir models/ltx-2.3
+hf download google/gemma-3-12b-it-qat-q4_0-unquantized --local-dir models/gemma-3-12b
+```
+
+If you get a 401/403, accept the model terms on Hugging Face and log in with a **Read** token (fine-grained tokens need the "read gated repos" scope enabled).
+
+Generate
+
+```bash
+uv run python -m ltx_pipelines.distilled \
+    --distilled-checkpoint-path models/ltx-2.3/ltx-2.3-22b-distilled-1.1.safetensors \
+    --spatial-upsampler-path    models/ltx-2.3/ltx-2.3-spatial-upscaler-x2-1.1.safetensors \
+    --gemma-root models/gemma-3-12b \
+    --seed 42 \
+    --output-path output.mp4 \
+    --prompt "A medium close-up shot features a Caucasian man with a beard, wearing a green and white baseball cap without any letters on the front, and a light blue shirt over a white t-shirt. He is positioned in the center of the frame, looking intently directly at the camera, his eyes focused on camera. His facial expression is one of deep concentration, with his brow slightly raised. As he looks straight at the camera, a quick sniff sound is heard, and then he speaks with a deep male voice and a satisfied tone, saying, 'I think it's so good.' The camera remains static throughout, maintaining a shallow depth of field, which keeps the man in sharp focus while the background is softly blurred, showing a beige wall behind him. After a brief pause, another short, audible sniff is heard. The man then continues to speak, his voice maintaining the same quality, as he states, 'So good. So good.' He elaborates further, emphasizing his point with a final statement, 'This got to be, it's got to be the best tool I've ever seen.'"
+```
+
+In cases of GPU memory constraints, consider `--quantization fp8-cast --offload {cpu, disk}`. See [additional flags](packages/ltx-pipelines/docs/installation.md#common-cli-flags).
+
+This uses the distilled model and pipeline for fast results. For better quality or other capabilities, see [Models](#full-model-list) and [Pipelines](#available-pipelines).
+
+### Full Model List
+
+For pipelines beyond the quickstart, download the relevant models from the [LTX-2.3 HuggingFace repository](https://huggingface.co/Lightricks/LTX-2.3):
 
 **LTX-2.3 Model Checkpoint** (choose and download one of the following)
   * [`ltx-2.3-22b-dev.safetensors`](https://huggingface.co/Lightricks/LTX-2.3/blob/main/ltx-2.3-22b-dev.safetensors) - [Download](https://huggingface.co/Lightricks/LTX-2.3/resolve/main/ltx-2.3-22b-dev.safetensors)
@@ -76,9 +100,9 @@ Download the following models from the [LTX-2.3 HuggingFace repository](https://
 ### ⚡ Optimization Tips
 
 * **Use DistilledPipeline** - Fastest inference with only 8 predefined sigmas (8 steps stage 1, 4 steps stage 2)
-* **Enable FP8 quantization** - Enables lower memory footprint: `--quantization fp8-cast` (CLI) or `quantization=QuantizationPolicy.fp8_cast()` (Python). Fp8-cast should be used with bf16 checkpoints, it shall downcast them on the fly. For Hopper GPUs with TensorRT-LLM, use `--quantization fp8-scaled-mm` for FP8 scaled matrix multiplication. Fp8-scaled-mm should be used with fp8 checkpoints.
-* **Install attention optimizations** - On datacenter Blackwell GPUs (B200), install FlashAttention 4 manually: `uv pip install 'flash-attn-4==4.0.0b9'` (this specific revision is the one we have verified against torch 2.9.1+cu128; newer betas have known issues on consumer Blackwell). On other CUDA GPUs (including Hopper), use xFormers (`uv sync --extra xformers`).
-* **Use gradient estimation** - Reduce inference steps from 40 to 20-30 while maintaining quality (see [pipeline documentation](packages/ltx-pipelines/README.md#denoising-loop-optimization))
+* **Enable FP8 quantization** - Enables lower memory footprint: `--quantization fp8-cast` (CLI) or `quantization=QuantizationPolicy.fp8_cast()` (Python). Fp8-cast should be used with bf16 checkpoints, it shall downcast them on the fly. On Hopper+ GPUs with native FP8 support, use `--quantization fp8-scaled-mm` for FP8 scaled matrix multiplication. Fp8-scaled-mm should be used with fp8 checkpoints.
+* **Install attention optimizations** - On datacenter Blackwell GPUs (B200), install FlashAttention 4 manually: `uv pip install 'flash-attn-4==4.0.0b9'` (this specific revision is the one we have verified against torch 2.9.1+cu128; newer betas have known issues on consumer Blackwell). On Hopper GPUs, install the FlashAttention 3 wheel. On other CUDA GPUs, PyTorch SDPA is used automatically. An installed backend is selected automatically at runtime; forcing a specific one is a Python-API option (`AttentionFunction.FLASH_ATTENTION_3`/`FLASH_ATTENTION_4`), not a CLI flag.
+* **Use gradient estimation** - Reduce inference steps from 40 to 20-30 while maintaining quality (see [pipeline documentation](packages/ltx-pipelines/docs/optimization.md#denoising-loop-optimization))
 * **Skip memory cleanup** - If you have sufficient VRAM, disable automatic memory cleanup between stages for faster processing
 * **Choose single-stage pipeline** - Use `TI2VidOneStagePipeline` for faster generation when high resolution isn't required
 
@@ -94,7 +118,7 @@ When writing prompts, focus on detailed, chronological descriptions of actions a
 - Describe lighting and colors
 - Note any changes or sudden events
 
-For additional guidance on writing a prompt please refer to <https://ltx.video/blog/how-to-prompt-for-ltx-2>
+For additional guidance on writing a prompt please refer to <https://ltx.io/blog/prompting-guide-for-ltx-2>
 
 ### Automatic Prompt Enhancement
 

packages/ltx-core/README.md

@@ -11,7 +11,7 @@ The foundational library for the LTX-2 Audio-Video generation model. This packag
 - **`block_streaming/`**: Memory-efficient inference that streams transformer blocks through the GPU one at a time (from pinned CPU buffers or directly from disk)
 - **`model/`**: PyTorch implementations of the LTX-2 Transformer, Video VAE, Audio VAE, Vocoder and Upscaler
 - **`text_encoders/gemma`**: Gemma text encoder implementation with tokenizers, feature extractors, and separate encoders for audio-video and video-only generation
-- **`quantization/`**: FP8 quantization backends (FP8-TensorRT-LLM scaled MM, FP8 cast) for reduced memory footprint.
+- **`quantization/`**: FP8 quantization backends (FP8 scaled MM, FP8 cast) for reduced memory footprint.
 
 ## 🚀 Quick Start
 
@@ -118,11 +118,9 @@ model = builder.build(device=torch.device("cuda"))
 
 The `quantization/` module provides FP8 quantization support for the LTX-2 transformer, significantly reducing memory usage while maintaining quality. Two backends are available:
 
-#### FP8 Scaled MM (TensorRT-LLM)
+#### FP8 Scaled MM
 
-Uses NVIDIA TensorRT-LLM's `cublas_scaled_mm` for efficient FP8 matrix multiplication. Weights are stored in FP8 format with per-tensor scaling, and inputs are quantized dynamically (or statically with calibration data).
-
-**Requirements**: `uv sync --frozen --extra fp8-trtllm`
+Uses PyTorch's `torch._scaled_mm` for efficient FP8 matrix multiplication. Weights are stored in FP8 format with per-tensor scaling, and inputs are quantized dynamically.
 
 **Usage with QuantizationPolicy:**
 

packages/ltx-core/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "ltx-core"
-version = "1.1.6"
+version = "1.1.7"
 description = "Core implementation of Lightricks' LTX-2 model"
 readme = "README.md"
 requires-python = ">=3.10"
@@ -13,38 +13,14 @@ dependencies = [
     "safetensors",
     "accelerate",
     "scipy>=1.14",
+    # Apple Silicon only: Apple's fused MPSGraph SDPA, the AUTOMATIC attention
+    # backend on MPS. The marker installs it on Apple Silicon and prunes it
+    # everywhere else (Linux/CUDA), so it is a hard requirement exactly where it
+    # is the only viable attention kernel. Requires torch>=2.11 (within the
+    # torch~=2.7 floor); the resolver forks torch to >=2.11 on macOS.
+    "mps-sdpa>=0.2.0; sys_platform == 'darwin' and platform_machine == 'arm64'",
 ]
 
-[project.optional-dependencies]
-xformers = ["xformers"]
-fp8-trtllm = [
-    "tensorrt-llm==1.0.0",
-    "onnx>=1.16.0,<1.20.0",
-    "openmpi",
-]
-
-[tool.uv]
-conflicts = [
-  [
-    { extra = "xformers" },
-    { extra = "fp8-trtllm" },
-  ],
-]
-
-[tool.uv.sources]
-xformers = { index = "pytorch" }
-tensorrt-llm = { index = "nvidia" }
-
-[[tool.uv.index]]
-name = "pytorch"
-url = "https://download.pytorch.org/whl/cu129"
-explicit = true
-
-[[tool.uv.index]]
-name = "nvidia"
-url = "https://pypi.nvidia.com/"
-explicit = true
-
 [build-system]
 requires = ["uv_build>=0.9.8,<0.10.0"]
 build-backend = "uv_build"

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

packages/ltx-kernels/MANIFEST.in

@@ -0,0 +1 @@
+recursive-include csrc *.h *.cuh *.hpp *.cpp *.cu

packages/ltx-kernels/README.md

@@ -0,0 +1,83 @@
+# ltx-kernels
+
+Custom CUDA/C++ kernels for `ltx-core`. Three compiled extensions:
+
+- **`all2all_cpp`** -- All2All communication kernels for multi-GPU tensor
+  parallelism, used by the sequence-parallel inference path.
+- **`ops_cpp`** -- Fused element ops for blockwise quantization: `rms_norm_rope`,
+  `rms_norm_split_rope`, and FP6 pack/unpack.
+- **`blockwise_cpp`** -- Blockwise FP8 GEMM. SM89 (GeForce/Ada) kernel always;
+  the SM90 (Hopper, `deep_gemm`) kernel is added when a `9.0` architecture is
+  requested.
+
+The Python surface for blockwise quantization lives in
+`ltx_kernels.blockwise` (`functional`, `linear`, `triton_ops`).
+
+## Requirements
+
+- CUDA toolkit (nvcc) matching your GPU architecture
+- PyTorch with CUDA support
+- Linux
+
+## Building
+
+`ltx-kernels` is excluded from the uv workspace, so a plain `uv sync` does not
+build it. From the repository root, build it via the opt-in `kernels` group
+(editable, no build isolation -- torch must already be installed):
+
+```bash
+uv sync --group kernels
+```
+
+Equivalently, install it directly:
+
+```bash
+uv pip install -e packages/ltx-kernels --no-build-isolation
+```
+
+Set `TORCH_CUDA_ARCH_LIST` to target specific architectures (speeds up compilation):
+
+```bash
+# H100 only
+TORCH_CUDA_ARCH_LIST="9.0" uv pip install -e packages/ltx-kernels --no-build-isolation
+
+# Multiple architectures
+TORCH_CUDA_ARCH_LIST="9.0 9.0a 10.0 12.0" uv pip install -e packages/ltx-kernels --no-build-isolation
+```
+
+When `TORCH_CUDA_ARCH_LIST` is unset the build targets every supported
+architecture (so `uv pip install` "just works" on a dev box); pin it on build
+hosts to cut compile time. Any `9.0` entry enables the SM90 GEMM kernel, which
+is compiled for `sm_90a` (the deep_gemm kernel uses wgmma/TMA).
+
+### cutlass headers
+
+`blockwise_cpp` includes cute/cutlass headers (header-only; compiled into the
+extension, with no runtime dependency). The build fetches them automatically on
+first use: a blobless, `include/`-only sparse clone of cutlass pinned to commit
+`afa17722` (v3.8.0), cached under `~/.cache/ltx-kernels/` (~25 MB) and reused
+across builds.
+
+- Set `CUTLASS_DIR=/path/to/cutlass` to use an existing checkout (uses
+  `$CUTLASS_DIR/include` and skips the fetch).
+- Set `LTX_KERNELS_CACHE_DIR` to override the cache location.
+
+To bump cutlass, change `CUTLASS_REF` in `setup.py`.
+
+## Testing
+
+Tests require a CUDA GPU:
+
+```bash
+uv run pytest packages/ltx-kernels/tests/ -v
+```
+
+## Operations
+
+`all2all_cpp`:
+
+- **send_recv_heads** -- Redistributes attention heads across GPUs (All2All)
+- **gather_heads** -- Inverse of send_recv_heads
+- **allgather** -- Gathers sequence tokens from all ranks
+
+All operations support BFloat16 and Float8 (e4m3fn) data types.

packages/ltx-kernels/csrc/all2all/all2all.cpp

@@ -0,0 +1,424 @@
+/**
+ * @file all2all.cpp
+ * @brief Implementation of All2All communication primitives for multi-GPU tensor parallelism.
+ *
+ * This file implements the All2All class which provides efficient inter-GPU communication
+ * using CUDA IPC (Inter-Process Communication). The implementation supports:
+ * - Head redistribution for tensor-parallel attention (send_recv_heads, gather_heads)
+ * - Sequence gathering for cross-rank aggregation (allgather)
+ *
+ * All operations use a barrier-based synchronization protocol where each GPU writes
+ * directly to remote GPU memory via IPC, then signals completion through atomic
+ * operations on barrier counters.
+ */
+
+#include <ATen/cuda/CUDAContext.h>
+#include <ATen/cuda/CUDADataType.h>
+#include <c10/cuda/CUDAGuard.h>
+
+#include <chrono>
+#include <cuda_runtime.h>
+#include <memory>
+#include <pybind11/functional.h>
+#include <torch/python.h>
+
+#include "all2all.hpp"
+#include "cuda/api.cuh"
+#include "cuda/configs.cuh"
+
+namespace ltx_kernels {
+namespace all2all {
+
+/**
+ * Constructs the All2All communication manager.
+ *
+ * Memory Allocation Strategy:
+ * The constructor allocates a single contiguous GPU memory block that contains:
+ *   1. Data buffer (tensor_bytes): Space for tensor data exchange
+ *   2. Barrier signals (MAX_NUM_PEERS * sizeof(int)): Per-rank completion counters
+ *   3. Buffer pointers (MAX_NUM_PEERS * sizeof(void*)): GPU-accessible pointer array
+ *   4. Barrier pointer array (MAX_NUM_PEERS * sizeof(int*)): GPU-accessible signal pointers
+ *
+ * This layout minimizes memory allocations and allows the entire region to be
+ * shared via a single IPC handle.
+ */
+All2All::All2All(int rank, int world_size, int num_tokens, int hidden_dim, int num_sms, at::ScalarType tensor_dtype,
+                 double timeout_seconds)
+    : rank(rank), world_size(world_size), num_sms(num_sms), max_tokens(num_tokens), num_elems(0), tensor_bytes(0),
+      tensor_dtype(tensor_dtype) {
+  num_elems = int64_t(num_tokens) * int64_t(hidden_dim);
+  tensor_bytes = num_elems * elementSize(tensor_dtype);
+
+  // Derive the barrier timeout from the device's peak SM clock so the wall-clock guard is
+  // correct on any GPU (the kernel counts SM cycles via clock64). Use cudaDeviceGetAttribute,
+  // not cudaDeviceProp::clockRate, which was removed in CUDA 13. The attribute is in kHz.
+  int device = 0;
+  CUDA_CHECK(cudaGetDevice(&device));
+  int sm_clock_khz = 0;
+  CUDA_CHECK(cudaDeviceGetAttribute(&sm_clock_khz, cudaDevAttrClockRate, device));
+  sm_clock_hz_ = static_cast<double>(sm_clock_khz) * 1e3;
+  set_timeout_seconds(timeout_seconds);
+
+  // Calculate sizes for each region of the shared memory block
+  int64_t ptrs_bytes = MAX_NUM_PEERS * sizeof(void *);
+  int64_t barrier_signal_bytes = MAX_NUM_PEERS * sizeof(int);
+  int64_t barrier_signal_ptrs_bytes = MAX_NUM_PEERS * sizeof(int *);
+
+  // Allocate GPU memory for token count arrays (used by kernels)
+  CUDA_CHECK(cudaMalloc(reinterpret_cast<void **>(&rank_tokens_gpu), sizeof(int) * MAX_NUM_PEERS));
+  CUDA_CHECK(cudaMalloc(reinterpret_cast<void **>(&prefix_rank_tokens_gpu), sizeof(int) * MAX_NUM_PEERS));
+
+  // Allocate the main shared memory block and create IPC handle
+  // Layout: [data_buffer | barrier_signals | buffer_ptrs | barrier_signal_ptrs]
+  CUDA_CHECK(
+      cudaMalloc(&buffer_ptrs[rank], tensor_bytes + barrier_signal_bytes + ptrs_bytes + barrier_signal_ptrs_bytes));
+  CUDA_CHECK(cudaIpcGetMemHandle(&ipc_handlers[rank], buffer_ptrs[rank]));
+
+  // Set up pointers to each region within the allocated block
+  buffer_ptrs_gpu =
+      reinterpret_cast<void **>(static_cast<uint8_t *>(buffer_ptrs[rank]) + tensor_bytes + barrier_signal_bytes);
+  barrier_signal_ptrs[rank] = reinterpret_cast<int *>(static_cast<uint8_t *>(buffer_ptrs[rank]) + tensor_bytes);
+  barrier_signal_ptrs_gpu = reinterpret_cast<int **>(static_cast<uint8_t *>(buffer_ptrs[rank]) + tensor_bytes +
+                                                     barrier_signal_bytes + ptrs_bytes);
+
+  // Initialize barrier signals to zero
+  CUDA_CHECK(cudaMemset(barrier_signal_ptrs[rank], 0, barrier_signal_bytes));
+}
+
+All2All::~All2All() noexcept(false) {
+  if (!destroyed) {
+    printf("WARNING: destroy() was not called, which can leak resources.\n");
+    fflush(stdout);
+    destroy();
+  }
+}
+
+/**
+ * Releases all allocated resources.
+ *
+ * This must be called explicitly before destruction to ensure proper cleanup of:
+ * - IPC memory mappings to remote GPUs
+ * - Local GPU memory allocations
+ *
+ * The method synchronizes the device to ensure all pending operations complete
+ * before releasing resources.
+ */
+void All2All::destroy() {
+  if (destroyed) {
+    return;
+  }
+  CUDA_CHECK(cudaDeviceSynchronize());
+
+  // Close IPC mappings to remote GPU memory (skip our own rank)
+  // Only close handles that were actually opened via sync()
+  for (int i = 0; i < world_size; i++) {
+    if (i != rank && buffer_ptrs[i] != nullptr) {
+      CUDA_CHECK(cudaIpcCloseMemHandle(buffer_ptrs[i]));
+    }
+  }
+
+  // Free local GPU memory allocations
+  CUDA_CHECK(cudaFree(buffer_ptrs[rank]));
+  CUDA_CHECK(cudaFree(rank_tokens_gpu));
+  CUDA_CHECK(cudaFree(prefix_rank_tokens_gpu));
+  destroyed = true;
+}
+
+/**
+ * Opens IPC memory mappings to all peer GPUs.
+ *
+ * This method processes IPC handles gathered from all ranks and opens memory
+ * mappings to enable direct GPU-to-GPU memory access. After calling this method,
+ * each GPU can read/write directly to any other GPU's buffer via buffer_ptrs.
+ *
+ * The barrier_signal_ptrs are also set up to point to the correct offset within
+ * each peer's shared memory block.
+ */
+void All2All::sync(const std::vector<std::optional<pybind11::bytearray>> &all_gathered_handles) {
+  for (int i = 0; i < world_size; i++) {
+    auto handle_str = std::string(all_gathered_handles[i].value());
+    EP_HOST_ASSERT(handle_str.size() == CUDA_IPC_HANDLE_SIZE);
+
+    if (i != rank) {
+      // Open IPC mapping to remote GPU's memory
+      std::memcpy(ipc_handlers[i].reserved, handle_str.c_str(), CUDA_IPC_HANDLE_SIZE);
+      CUDA_CHECK(cudaIpcOpenMemHandle(&buffer_ptrs[i], ipc_handlers[i], cudaIpcMemLazyEnablePeerAccess));
+      // Calculate offset to barrier signals in remote buffer
+      barrier_signal_ptrs[i] = reinterpret_cast<int *>(static_cast<uint8_t *>(buffer_ptrs[i]) + tensor_bytes);
+    } else {
+      // Verify our own handle matches what we sent
+      EP_HOST_ASSERT(std::memcmp(ipc_handlers[i].reserved, handle_str.c_str(), CUDA_IPC_HANDLE_SIZE) == 0);
+    }
+  }
+
+  // Copy pointer arrays to GPU for kernel access
+  CUDA_CHECK(cudaMemcpy(buffer_ptrs_gpu, buffer_ptrs, sizeof(void *) * world_size, cudaMemcpyHostToDevice));
+  CUDA_CHECK(
+      cudaMemcpy(barrier_signal_ptrs_gpu, barrier_signal_ptrs, sizeof(int *) * world_size, cudaMemcpyHostToDevice));
+  CUDA_CHECK(cudaDeviceSynchronize());
+}
+
+pybind11::bytearray All2All::get_local_ipc_handle() const {
+  return {ipc_handlers[rank].reserved, CUDA_IPC_HANDLE_SIZE};
+}
+
+/**
+ * Configures token distribution across ranks for the current batch.
+ *
+ * This method computes prefix sums needed by the kernels to calculate source
+ * and destination offsets. It must be called before any communication operation
+ * when the token distribution changes between batches.
+ *
+ * Example: For rank_num_tokens = {128, 96, 128, 64}
+ *   - rank_tokens = {128, 96, 128, 64}
+ *   - prefix_rank_tokens = {0, 128, 224, 352}
+ *   - total_tokens = 416
+ */
+void All2All::set_rank_tokens(const std::vector<int> &rank_num_tokens) {
+  EP_HOST_ASSERT(static_cast<int>(rank_num_tokens.size()) == world_size);
+
+  // Initialize prefix sums to zero
+  for (int i = 0; i < world_size; i++) {
+    prefix_rank_tokens[i] = 0;
+  }
+
+  // Compute prefix sums (exclusive scan)
+  for (int i = 0; i < world_size; i++) {
+    rank_tokens[i] = rank_num_tokens[i];
+    if (i > 0) {
+      prefix_rank_tokens[i] = prefix_rank_tokens[i - 1] + rank_tokens[i - 1];
+    }
+  }
+
+  // Total tokens is the sum of all rank tokens
+  total_tokens = prefix_rank_tokens[world_size - 1] + rank_tokens[world_size - 1];
+
+  // Copy to GPU for kernel access
+  CUDA_CHECK(cudaMemcpy(rank_tokens_gpu, rank_tokens, sizeof(int) * MAX_NUM_PEERS, cudaMemcpyHostToDevice));
+  CUDA_CHECK(
+      cudaMemcpy(prefix_rank_tokens_gpu, prefix_rank_tokens, sizeof(int) * MAX_NUM_PEERS, cudaMemcpyHostToDevice));
+  CUDA_CHECK(cudaDeviceSynchronize());
+}
+
+/**
+ * Creates a tensor from the local IPC buffer.
+ *
+ * This helper method returns either a zero-copy view of the IPC buffer or
+ * a newly allocated tensor with the data copied. The zero-copy mode is more
+ * efficient but the tensor lifetime is tied to the All2All instance.
+ *
+ * @note The buffer pointer is cast to the template type T for proper interpretation.
+ */
+at::Tensor All2All::get_local_buffer_tensor(at::Tensor &x, int batch_size, int out_tokens, int out_heads, int head_size,
+                                            bool should_copy, cudaStream_t stream) {
+  auto ptr = buffer_ptrs[rank];
+  if (should_copy) {
+    // Allocate new tensor and copy data from IPC buffer
+    auto out_tensor = torch::empty({batch_size, out_tokens, out_heads, head_size}, x.options());
+    CUDA_CHECK(cudaMemcpyAsync(out_tensor.data_ptr(), ptr,
+                               int64_t(batch_size) * int64_t(out_tokens) * int64_t(out_heads) * int64_t(head_size) *
+                                   int64_t(elementSize(x.scalar_type())),
+                               cudaMemcpyDeviceToDevice, stream));
+    return out_tensor;
+  } else {
+    // Return a view directly into the IPC buffer (zero-copy)
+    auto out_tensor = torch::from_blob(ptr, {batch_size, out_tokens, out_heads, head_size}, x.options());
+    return out_tensor;
+  }
+}
+
+/**
+ * All2All communication to redistribute attention heads across GPUs.
+ *
+ * This operation is used in tensor-parallel transformers to exchange attention heads:
+ *   - Before: Each GPU has all tokens but only a subset of heads
+ *   - After: Each GPU has all tokens with heads redistributed
+ *
+ * Tensor Layout Transformation:
+ *   Input:  [batch, local_tokens, all_heads, head_size]  per GPU
+ *   Output: [batch, all_tokens, heads_per_rank, head_size]  per GPU
+ *
+ * The operation partitions heads evenly: heads_per_rank = all_heads / world_size
+ * GPU i receives heads [i*heads_per_rank : (i+1)*heads_per_rank] from all GPUs.
+ */
+at::Tensor All2All::send_recv_heads(at::Tensor &x, bool copy_output) {
+  // Validate input tensor properties
+  EP_HOST_ASSERT(x.dim() == 4 and x.is_contiguous());
+  EP_HOST_ASSERT(x.dtype() == tensor_dtype);
+  EP_HOST_ASSERT(x.device().is_cuda());
+  EP_HOST_ASSERT(x.device().index() == rank);
+
+  int batch_size = x.size(0);
+  int num_tokens = x.size(1);
+  int num_heads = x.size(2);
+  int head_size = x.size(3);
+
+  // Output dimensions after redistribution
+  int out_tokens = total_tokens;          // All tokens from all ranks
+  int out_heads = num_heads / world_size; // Each rank gets 1/world_size of heads
+
+  EP_HOST_ASSERT(int64_t(batch_size) * int64_t(out_tokens) * int64_t(out_heads) * int64_t(head_size) *
+                     int64_t(elementSize(x.scalar_type())) <=
+                 tensor_bytes);
+
+  at::cuda::CUDAGuard device_guard{x.device()};
+  auto stream = at::cuda::getCurrentCUDAStream().stream();
+
+  // Launch the All2All kernel
+  all2all_cuda::all2all_head_launch(buffer_ptrs_gpu, barrier_signal_ptrs_gpu, x.data_ptr(), prefix_rank_tokens_gpu,
+                                    rank, world_size, batch_size, total_tokens, num_tokens, num_heads, head_size,
+                                    stream, num_sms, tensor_dtype, timeout_cycles_);
+
+  return get_local_buffer_tensor(x, batch_size, out_tokens, out_heads, head_size, copy_output, stream);
+}
+
+/**
+ * Inverse All2All to gather heads back to original distribution.
+ *
+ * This is the inverse operation of send_recv_heads(). It redistributes data
+ * so each GPU gets back its original tokens with all attention heads.
+ *
+ * Tensor Layout Transformation:
+ *   Input:  [batch, all_tokens, heads_per_rank, head_size]  per GPU
+ *   Output: [batch, local_tokens, all_heads, head_size]  per GPU
+ *
+ * Each GPU sends its portion of tokens to the originating rank, reconstructing
+ * the original head distribution.
+ */
+at::Tensor All2All::gather_heads(at::Tensor &x, bool copy_output) {
+  // Validate input tensor properties
+  EP_HOST_ASSERT(x.dim() == 4 and x.is_contiguous());
+  EP_HOST_ASSERT(x.dtype() == tensor_dtype);
+  EP_HOST_ASSERT(x.device().is_cuda());
+  EP_HOST_ASSERT(x.device().index() == rank);
+
+  at::cuda::CUDAGuard device_guard{x.device()};
+  auto stream = at::cuda::getCurrentCUDAStream().stream();
+
+  int batch_size = x.size(0);
+  int num_heads = x.size(2) * world_size; // Reconstruct total head count
+  int head_size = x.size(3);
+
+  // Output dimensions: this rank's tokens with all heads
+  int out_tokens = rank_tokens[rank];
+  int out_heads = num_heads;
+
+  EP_HOST_ASSERT(int64_t(batch_size) * int64_t(out_tokens) * int64_t(out_heads) * int64_t(head_size) *
+                     int64_t(elementSize(x.scalar_type())) <=
+                 tensor_bytes);
+
+  // Launch the gather kernel
+  all2all_cuda::all2all_head_gather_launch(buffer_ptrs_gpu, barrier_signal_ptrs_gpu, x.data_ptr(), rank_tokens_gpu,
+                                           prefix_rank_tokens_gpu, rank, world_size, batch_size, total_tokens,
+                                           num_heads, head_size, stream, num_sms, tensor_dtype, timeout_cycles_);
+
+  return get_local_buffer_tensor(x, batch_size, out_tokens, out_heads, head_size, copy_output, stream);
+}
+
+/**
+ * AllGather operation to collect sequence tokens from all ranks.
+ *
+ * Each GPU contributes its local sequence tokens, which are gathered into
+ * a complete sequence replicated on all GPUs. This is typically used after
+ * tensor-parallel operations to reconstruct the full sequence.
+ *
+ * Tensor Layout Transformation:
+ *   Input:  [batch, local_seqlen, heads, head_size]  per GPU
+ *   Output: [batch, total_seqlen, heads, head_size]  per GPU (identical on all GPUs)
+ *
+ * Each GPU's tokens are placed at offset prefix_rank_tokens[rank] in the output.
+ */
+at::Tensor All2All::allgather(at::Tensor &x, bool copy_output) {
+  // Validate input tensor properties
+  EP_HOST_ASSERT(x.dim() == 4 and x.is_contiguous());
+  EP_HOST_ASSERT(x.dtype() == tensor_dtype);
+  EP_HOST_ASSERT(x.device().is_cuda());
+  EP_HOST_ASSERT(x.device().index() == rank);
+
+  at::cuda::CUDAGuard device_guard{x.device()};
+  auto stream = at::cuda::getCurrentCUDAStream().stream();
+
+  int batch_size = x.size(0);
+  int seqlen = x.size(1);
+  int num_heads = x.size(2);
+  int head_size = x.size(3);
+
+  // Output contains all tokens from all ranks
+  int out_tokens = total_tokens;
+  int out_heads = num_heads;
+  int hidden_dim = num_heads * head_size;
+
+  EP_HOST_ASSERT(int64_t(batch_size) * int64_t(out_tokens) * int64_t(out_heads) * int64_t(head_size) *
+                     int64_t(elementSize(x.scalar_type())) <=
+                 tensor_bytes);
+
+  // Launch the allgather kernel
+  all2all_cuda::allgather_launch(buffer_ptrs_gpu, barrier_signal_ptrs_gpu, x.data_ptr(), prefix_rank_tokens_gpu, rank,
+                                 world_size, batch_size, seqlen, hidden_dim, total_tokens, stream, num_sms,
+                                 tensor_dtype, timeout_cycles_);
+
+  return get_local_buffer_tensor(x, batch_size, out_tokens, out_heads, head_size, copy_output, stream);
+}
+
+} // namespace all2all
+} // namespace ltx_kernels
+
+/**
+ * Python bindings for the All2All communication library.
+ *
+ * Usage from Python:
+ *   import all2all_cpp
+ *
+ *   # Create instance (one per GPU)
+ *   comm = all2all_cpp.All2All(rank, world_size, max_tokens, hidden_dim, num_sms, dtype)
+ *
+ *   # Exchange IPC handles and synchronize
+ *   handle = comm.get_local_ipc_handle()
+ *   # ... gather handles via NCCL ...
+ *   comm.sync(all_handles)
+ *
+ *   # Set token distribution
+ *   comm.set_rank_tokens([128, 128, 128, 128])
+ *
+ *   # Perform operations
+ *   output = comm.send_recv_heads(input_tensor, copy_output=False)
+ *
+ *   # Cleanup
+ *   comm.destroy()
+ */
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+  m.doc() = "High-performance All2All communication library for multi-GPU tensor parallelism.\n\n"
+            "This library provides IPC-based All2All operations optimized for transformer models.\n"
+            "Supported operations:\n"
+            "  - send_recv_heads: Redistribute attention heads across GPUs\n"
+            "  - gather_heads: Inverse of send_recv_heads\n"
+            "  - allgather: Gather sequence tokens from all ranks\n";
+
+  pybind11::class_<ltx_kernels::all2all::All2All>(
+      m, "All2All",
+      "Manages All2All communication state for multi-GPU operations.\n\n"
+      "Args:\n"
+      "    rank: This GPU's rank (0 to world_size-1)\n"
+      "    world_size: Total number of GPUs\n"
+      "    num_tokens: Maximum tokens per rank\n"
+      "    hidden_dim: Hidden dimension (heads * head_size)\n"
+      "    num_sms: Number of SMs for kernel launches\n"
+      "    tensor_dtype: Tensor data type (torch.bfloat16 or torch.float8_e4m3fn)\n"
+      "    timeout_seconds: Optional initial barrier timeout in seconds (defaults to the kernel default)")
+      .def(pybind11::init<int, int, int, int, int, at::ScalarType>())
+      .def(pybind11::init<int, int, int, int, int, at::ScalarType, double>())
+      .def("get_local_ipc_handle", &ltx_kernels::all2all::All2All::get_local_ipc_handle,
+           "Returns the IPC handle for this rank's buffer.")
+      .def("sync", &ltx_kernels::all2all::All2All::sync, "Opens IPC mappings to all peer GPUs using gathered handles.")
+      .def("destroy", &ltx_kernels::all2all::All2All::destroy,
+           "Releases all GPU resources. Must be called before destruction.")
+      .def("send_recv_heads", &ltx_kernels::all2all::All2All::send_recv_heads,
+           "All2All operation to redistribute attention heads.")
+      .def("gather_heads", &ltx_kernels::all2all::All2All::gather_heads,
+           "Inverse All2All to gather heads back to original distribution.")
+      .def("allgather", &ltx_kernels::all2all::All2All::allgather, "Gathers sequence tokens from all ranks.")
+      .def("set_rank_tokens", &ltx_kernels::all2all::All2All::set_rank_tokens,
+           "Sets token counts per rank for the current batch.")
+      .def("set_timeout_seconds", &ltx_kernels::all2all::All2All::set_timeout_seconds,
+           "Sets the barrier timeout in seconds (converted to cycles via the device peak SM clock).");
+}

packages/ltx-kernels/csrc/all2all/all2all.hpp

@@ -0,0 +1,265 @@
+/**
+ * @file all2all.hpp
+ * @brief High-performance All2All communication primitives for multi-GPU tensor parallelism.
+ *
+ * This library provides efficient All2All communication operations optimized for transformer
+ * models using tensor parallelism. It uses CUDA IPC (Inter-Process Communication) for
+ * zero-copy data transfer between GPUs in the same node.
+ *
+ * ## Architecture Overview
+ *
+ * The All2All class manages shared memory buffers accessible by all GPUs via IPC handles.
+ * Each GPU allocates a contiguous memory region containing:
+ *   - Data buffer: Stores tensor data for exchange
+ *   - Barrier signals: Synchronization counters for coordination
+ *   - GPU pointer arrays: Device-accessible pointers to all peer buffers
+ *
+ * Memory Layout (per GPU):
+ * ```
+ * |<---- tensor_bytes ---->|<-- barrier signals -->|<-- buffer_ptrs_gpu -->|<-- barrier_signal_ptrs_gpu -->|
+ * |      Data Buffer       |    MAX_PEERS * int    |   MAX_PEERS * void*   |      MAX_PEERS * int*         |
+ * ```
+ *
+ * ## Supported Operations
+ *
+ * 1. **send_recv_heads**: Redistributes attention heads across GPUs (All2All)
+ *    - Input: [batch, tokens, heads, head_size] on each GPU
+ *    - Output: [batch, total_tokens, heads/world_size, head_size] on each GPU
+ *
+ * 2. **gather_heads**: Inverse of send_recv_heads
+ *    - Gathers distributed heads back to original distribution
+ *
+ * 3. **allgather**: Gathers sequence data from all ranks
+ *    - Each GPU contributes its local tokens to form the complete sequence
+ *
+ * ## Thread Safety
+ *
+ * - The class is NOT thread-safe. Each thread/process should have its own instance.
+ * - Multiple CUDA streams may use the same instance sequentially.
+ * - The `destroy()` method MUST be called before destruction to properly release IPC handles.
+ *
+ * ## Usage Example
+ *
+ * ```cpp
+ * // Initialize on each GPU
+ * auto comm = All2All(rank, world_size, max_tokens, hidden_dim, num_sms, dtype);
+ *
+ * // Exchange IPC handles (via NCCL or other collective)
+ * auto my_handle = comm.get_local_ipc_handle();
+ * // ... gather all handles ...
+ * comm.sync(all_handles);
+ *
+ * // Set token distribution for current batch
+ * comm.set_rank_tokens({128, 128, 128, 128});  // tokens per rank
+ *
+ * // Perform All2All on attention heads
+ * auto result = comm.send_recv_heads(input_tensor, copy_output=false);
+ *
+ * // Clean up
+ * comm.destroy();
+ * ```
+ */
+
+#pragma once
+
+#include "cuda/configs.cuh"
+#include "event.hpp"
+#include <cmath>
+#include <limits>
+#include <pybind11/pybind11.h>
+#include <pybind11/pytypes.h>
+#include <stdexcept>
+#include <torch/types.h>
+#include <tuple>
+#include <vector>
+
+namespace ltx_kernels {
+namespace all2all {
+
+/**
+ * @class All2All
+ * @brief Manages All2All communication state and operations for multi-GPU tensor parallelism.
+ *
+ * This class encapsulates the IPC-based communication infrastructure needed for
+ * efficient All2All operations. It maintains shared memory buffers, barrier signals,
+ * and provides methods for head-parallel tensor redistribution.
+ */
+struct All2All {
+private:
+  int rank;             ///< This GPU's rank (0 to world_size-1)
+  int world_size;       ///< Total number of GPUs in the communication group
+  int num_sms;          ///< Number of SMs to use for kernel launches
+  int max_tokens;       ///< Maximum number of tokens the buffer was allocated for
+  int64_t num_elems;    ///< Number of elements in the data buffer (tokens * hidden_dim)
+  int64_t tensor_bytes; ///< Size of the data buffer in bytes
+
+  /// Host array of pointers to each rank's data buffer (GPU memory)
+  void *buffer_ptrs[MAX_NUM_PEERS] = {nullptr};
+  /// Device-accessible array of buffer pointers (copied to GPU)
+  void **buffer_ptrs_gpu = nullptr;
+
+  /// Host array of pointers to each rank's barrier signal buffer
+  int *barrier_signal_ptrs[MAX_NUM_PEERS] = {nullptr};
+  /// Device-accessible array of barrier signal pointers
+  int **barrier_signal_ptrs_gpu = nullptr;
+
+  /// IPC handles for sharing memory between processes
+  cudaIpcMemHandle_t ipc_handlers[MAX_NUM_PEERS];
+
+  at::ScalarType tensor_dtype; ///< Data type of tensors (BFloat16 or Float8_e4m3fn)
+  bool destroyed = false;      ///< Flag to track if resources have been released
+
+  int total_tokens;                      ///< Sum of tokens across all ranks for current batch
+  int rank_tokens[MAX_NUM_PEERS];        ///< Number of tokens on each rank
+  int prefix_rank_tokens[MAX_NUM_PEERS]; ///< Cumulative sum of tokens (for offset calculation)
+  int *rank_tokens_gpu = nullptr;        ///< Device copy of rank_tokens
+  int *prefix_rank_tokens_gpu = nullptr; ///< Device copy of prefix_rank_tokens
+
+  /// Device peak SM clock in Hz (from cudaDeviceGetAttribute(cudaDevAttrClockRate)), queried
+  /// once at construction. Used to convert a wall-clock timeout in seconds to barrier cycles.
+  double sm_clock_hz_ = 0.0;
+
+  /// All2All barrier timeout in GPU clock cycles. The constructor sets it from
+  /// DEFAULT_BARRIER_TIMEOUT_SECONDS and the queried SM clock; raise it (set_timeout_seconds)
+  /// to tolerate large cross-rank kernel-launch skew during the first torch.compile forward,
+  /// where one rank's recompile can delay its launch past the steady-state timeout.
+  uint64_t timeout_cycles_ = 0;
+
+public:
+  /**
+   * @brief Constructs an All2All communication manager.
+   *
+   * Allocates GPU memory for the local data buffer, barrier signals, and pointer arrays.
+   * The IPC handle for the local buffer is created and can be retrieved via get_local_ipc_handle().
+   *
+   * @param rank This GPU's rank in the communication group (0-indexed)
+   * @param world_size Total number of GPUs/ranks
+   * @param num_tokens Maximum number of tokens this rank will handle
+   * @param hidden_dim Hidden dimension size (heads * head_size)
+   * @param num_sms Number of CUDA SMs to use for kernel execution
+   * @param tensor_dtype Data type for tensors (BFloat16 or Float8_e4m3fn)
+   * @param timeout_seconds Initial barrier timeout in seconds (see set_timeout_seconds); may be
+   *        raised/reset at runtime for the first torch.compile forward
+   */
+  All2All(int rank, int world_size, int num_tokens, int hidden_dim, int num_sms, at::ScalarType tensor_dtype,
+          double timeout_seconds = DEFAULT_BARRIER_TIMEOUT_SECONDS);
+
+  /**
+   * @brief Destructor - warns if destroy() was not called.
+   *
+   * @warning Always call destroy() explicitly before the destructor to properly
+   *          release IPC handles. Failing to do so may leak resources.
+   */
+  ~All2All() noexcept(false);
+
+  /**
+   * @brief Synchronizes IPC handles from all ranks and opens remote memory mappings.
+   *
+   * This method must be called after all ranks have created their All2All instances
+   * and exchanged IPC handles via an external collective (e.g., NCCL allgather).
+   *
+   * @param all_gathered_handles Vector of IPC handles from all ranks (indexed by rank)
+   */
+  void sync(const std::vector<std::optional<pybind11::bytearray>> &all_gathered_handles);
+
+  /**
+   * @brief Returns the IPC handle for this rank's shared buffer.
+   *
+   * The returned handle should be gathered across all ranks and passed to sync().
+   *
+   * @return pybind11::bytearray containing the CUDA IPC handle (CUDA_IPC_HANDLE_SIZE bytes)
+   */
+  pybind11::bytearray get_local_ipc_handle() const;
+
+  /**
+   * @brief Creates a tensor view or copy of the local output buffer.
+   *
+   * @param x Reference tensor for options (dtype, device)
+   * @param batch_size Batch dimension size
+   * @param out_tokens Output token dimension size
+   * @param out_heads Output heads dimension size
+   * @param head_size Head dimension size
+   * @param should_copy If true, copies data to a new tensor; if false, returns a view
+   * @param stream CUDA stream for async copy
+   * @return Tensor with shape [batch_size, out_tokens, out_heads, head_size]
+   */
+  at::Tensor get_local_buffer_tensor(at::Tensor &x, int batch_size, int out_tokens, int out_heads, int head_size,
+                                     bool should_copy, cudaStream_t stream);
+
+  /**
+   * @brief Releases all GPU resources and closes IPC handles.
+   *
+   * This method MUST be called before the object is destroyed. It synchronizes
+   * the device, closes remote IPC mappings, and frees local GPU memory.
+   */
+  void destroy();
+
+  /**
+   * @brief Performs All2All communication to redistribute attention heads.
+   *
+   * Redistributes tensor from [batch, local_tokens, all_heads, head_size] to
+   * [batch, all_tokens, local_heads, head_size]. Each rank sends its portion
+   * of heads to the corresponding target rank.
+   *
+   * @param x Input tensor with shape [batch, num_tokens, num_heads, head_size]
+   * @param copy_output If true, returns a copy; if false, returns a view of the IPC buffer
+   * @return Tensor with shape [batch, total_tokens, num_heads/world_size, head_size]
+   */
+  at::Tensor send_recv_heads(at::Tensor &x, bool copy_output);
+
+  /**
+   * @brief Performs inverse All2All to gather heads back to original distribution.
+   *
+   * Inverse of send_recv_heads(). Redistributes from [batch, all_tokens, local_heads, head_size]
+   * back to [batch, local_tokens, all_heads, head_size].
+   *
+   * @param x Input tensor with shape [batch, total_tokens, heads_per_rank, head_size]
+   * @param copy_output If true, returns a copy; if false, returns a view of the IPC buffer
+   * @return Tensor with shape [batch, rank_tokens[rank], num_heads, head_size]
+   */
+  at::Tensor gather_heads(at::Tensor &x, bool copy_output);
+
+  /**
+   * @brief Gathers sequence tokens from all ranks.
+   *
+   * Each rank contributes its local sequence tokens, which are gathered into
+   * a complete sequence on all ranks.
+   *
+   * @param x Input tensor with shape [batch, seqlen, num_heads, head_size]
+   * @param copy_output If true, returns a copy; if false, returns a view of the IPC buffer
+   * @return Tensor with shape [batch, total_tokens, num_heads, head_size]
+   */
+  at::Tensor allgather(at::Tensor &x, bool copy_output);
+
+  /**
+   * @brief Sets the token count for each rank in the current batch.
+   *
+   * Must be called before send_recv_heads(), gather_heads(), or allgather()
+   * to configure the token distribution. This allows variable-length sequences
+   * across ranks.
+   *
+   * @param rank_num_tokens Vector of token counts, one per rank (must have world_size elements)
+   */
+  void set_rank_tokens(const std::vector<int> &rank_num_tokens);
+
+  /**
+   * @brief Sets the all2all barrier timeout in seconds.
+   *
+   * Converted to GPU clock cycles using the device's peak SM clock (queried at construction).
+   * Relaxes deadlock detection during the first torch.compile forward, where asymmetric
+   * per-rank recompilation can delay a rank's kernel launch beyond the steady-state timeout.
+   * Reset to the default for steady-state replay.
+   */
+  void set_timeout_seconds(double seconds) {
+    if (!std::isfinite(seconds) || seconds < 0.0) {
+      throw std::invalid_argument("All2All timeout (seconds) must be finite and non-negative");
+    }
+    // Saturate rather than overflow the float->uint64 cast (out-of-range conversion is UB).
+    const double cycles = seconds * sm_clock_hz_;
+    const double max_cycles = static_cast<double>(std::numeric_limits<uint64_t>::max());
+    timeout_cycles_ = cycles >= max_cycles ? std::numeric_limits<uint64_t>::max() : static_cast<uint64_t>(cycles);
+  }
+};
+
+} // namespace all2all
+} // namespace ltx_kernels

packages/ltx-kernels/csrc/all2all/cuda/all2all_heads.cu

@@ -0,0 +1,372 @@
+/**
+ * @file all2all_heads.cu
+ * @brief CUDA kernels for All2All attention head redistribution.
+ *
+ * This file implements the GPU kernels for redistributing attention heads across
+ * multiple GPUs using IPC-based direct memory access. The kernels are designed
+ * for tensor-parallel transformer models where attention heads need to be
+ * exchanged between GPUs.
+ *
+ * ## Algorithm Overview
+ *
+ * The kernels use a direct-write approach where each GPU writes its data directly
+ * to the target GPU's memory buffer via IPC. This avoids intermediate copies and
+ * achieves near-peak memory bandwidth utilization.
+ *
+ * ## SM Work Distribution (Round-Robin)
+ *
+ * SMs are distributed round-robin among target ranks to handle non-divisible SM counts:
+ *   - SM i writes to rank (i % world_size)
+ *   - With 132 SMs and 8 GPUs: ranks 0-3 get 17 SMs, ranks 4-7 get 16 SMs
+ *   - Each SM group processes all tokens for its assigned target rank
+ *   - Within each group, SMs cooperate to cover all tokens in strided fashion
+ *
+ * ## Synchronization Protocol
+ *
+ * After data transfer, a barrier synchronization ensures all ranks have completed:
+ *   1. Each SM atomically increments the target rank's barrier counter for this rank
+ *   2. SM 0 waits until it has received signals from all ranks
+ *   3. Barrier counters are reset for the next operation
+ */
+
+#include "cuda/configs.cuh"
+#include "cuda/exceptions.cuh"
+#include "cuda/utils.cuh"
+#include <ATen/cuda/CUDADataType.h>
+
+namespace ltx_kernels {
+namespace all2all {
+namespace all2all_cuda {
+
+/**
+ * @brief All2All kernel for redistributing attention heads across GPUs.
+ *
+ * This kernel performs the "send" phase of All2All: each GPU writes its assigned
+ * subset of attention heads to all other GPUs. The data layout transformation is:
+ *
+ *   Source: [batch, num_tokens, num_heads, head_size]
+ *   Dest:   [batch, total_tokens, heads_per_rank, head_size]
+ *
+ * Each GPU writes heads [target_rank * heads_per_rank : (target_rank+1) * heads_per_rank]
+ * to target_rank's buffer at token offset prefix_rank_tokens[rank].
+ *
+ * ## Memory Layout
+ *
+ * Input tensor x (row-major, contiguous):
+ *   - Batch dimension: outermost
+ *   - Token dimension: batch_stride = num_tokens * num_heads * head_size
+ *   - Head dimension: token_stride = num_heads * head_size
+ *   - Head element: head_stride = head_size
+ *
+ * Output buffer (per target rank):
+ *   - Similar layout but with heads_per_rank instead of num_heads
+ *   - Tokens from this rank placed at offset prefix_rank_tokens[rank]
+ *
+ * ## Thread Block Organization
+ *
+ * Each thread block handles multiple tokens cooperatively:
+ *   - Threads are organized in a 2D logical grid (rows=tokens, cols=elements)
+ *   - Each thread copies 16 bytes (int4) per iteration
+ *   - num_threads_per_token = (heads_per_rank * head_size) / elements_per_thread
+ *   - num_tokens_per_copy = num_threads / num_threads_per_token
+ *
+ * @tparam ELEM_T Element type (at::BFloat16 or at::Float8_e4m3fn)
+ * @param buffer_ptrs Device array of pointers to each rank's data buffer
+ * @param barrier_signal_ptrs Device array of pointers to each rank's barrier signals
+ * @param x Source tensor data pointer
+ * @param rank This GPU's rank
+ * @param world_size Total number of GPUs
+ * @param batch_size Number of batches
+ * @param num_tokens Number of tokens on this rank
+ * @param num_heads Total number of attention heads
+ * @param head_size Size of each attention head
+ * @param total_tokens Sum of tokens across all ranks
+ * @param prefix_rank_tokens Cumulative token counts for offset calculation
+ */
+template <typename ELEM_T>
+__global__ void send_recv_all2all(void **buffer_ptrs, int **barrier_signal_ptrs, void *x, int rank, int world_size,
+                                  int batch_size, int num_tokens, int num_heads, int head_size, int total_tokens,
+                                  int *prefix_rank_tokens, uint64_t timeout_cycles) {
+  // Grid dimensions
+  int num_sms = gridDim.x;
+  int sm_id = blockIdx.x;
+  int num_threads = blockDim.x;
+
+  // === SM Work Distribution (Round-Robin) ===
+  // Use modular assignment to handle num_sms not divisible by world_size.
+  // This ensures all SMs are utilized: some ranks get ceil(num_sms/world_size)
+  // SMs, others get floor(num_sms/world_size) SMs.
+  int64_t target_rank = get_target_rank(sm_id, world_size);
+  int64_t rank_local_sm_id = get_rank_local_sm_id(sm_id, world_size);
+  int64_t num_sms_for_this_rank = get_num_sms_for_rank(target_rank, num_sms, world_size);
+
+  // === Head Assignment ===
+  // Heads are partitioned evenly: rank i gets heads [i*hpr : (i+1)*hpr]
+  int64_t heads_per_rank = num_heads / world_size;
+  int64_t head_id = target_rank * heads_per_rank; // Starting head for target rank
+
+  // === Thread Mapping ===
+  // Each thread copies an int4 (16 bytes) per memory operation
+  // Threads form a 2D grid: (tokens_per_copy, threads_per_token)
+  int64_t num_elems_per_thread = sizeof(int4) / sizeof(ELEM_T);
+  int64_t num_threads_per_token = heads_per_rank * head_size / num_elems_per_thread;
+  int64_t num_tokens_per_copy = num_threads / num_threads_per_token;
+
+  // 2D thread coordinates within the logical grid
+  int64_t copy_thr_col_idx = threadIdx.x % num_threads_per_token; // Element offset
+  int64_t copy_thr_row_idx = threadIdx.x / num_threads_per_token; // Token offset
+
+  // Get target rank's buffer pointer
+  auto ptr = reinterpret_cast<void *>(static_cast<int8_t *>(buffer_ptrs[target_rank]));
+
+  // Use 64-bit arithmetic to avoid overflow for large tensors
+  int64_t num_tokens_64b = int64_t(num_tokens);
+  int64_t num_heads_64b = int64_t(num_heads);
+  int64_t head_size_64b = int64_t(head_size);
+
+  // === Main Copy Loop ===
+  // Iterate over batches and tokens, with SMs in the same group
+  // working on different token ranges in strided fashion
+  for (int64_t batch_ind = 0; batch_ind < batch_size; batch_ind++) {
+    // Strided token iteration: each SM in the group handles different token ranges
+    for (int64_t token_idx = rank_local_sm_id * num_tokens_per_copy; token_idx < num_tokens;
+         token_idx += num_tokens_per_copy * num_sms_for_this_rank) {
+      int64_t copy_token_idx = token_idx + copy_thr_row_idx;
+      // Destination token index accounts for this rank's offset in the global sequence
+      int64_t dst_token_idx = prefix_rank_tokens[rank] + copy_token_idx;
+
+      if (copy_token_idx >= num_tokens)
+        break;
+
+      // === Pointer Arithmetic ===
+      // Source: Read from this rank's input tensor at [batch, token, head_id:head_id+hpr, :]
+      // Note: We read a contiguous chunk of heads starting at head_id
+      int4 *shuffled_x_ptr =
+          reinterpret_cast<int4 *>(reinterpret_cast<uint8_t *>(x) +
+                                   batch_ind * num_tokens_64b * num_heads_64b * head_size_64b * sizeof(ELEM_T) +
+                                   copy_token_idx * num_heads_64b * head_size_64b * sizeof(ELEM_T) +
+                                   head_id * head_size_64b * sizeof(ELEM_T)) +
+          copy_thr_col_idx;
+
+      // Destination: Write to target rank's buffer at [batch, dst_token, :, :]
+      // The buffer has layout [batch, total_tokens, heads_per_rank, head_size]
+      int4 *shuffled_buffer_ptr =
+          reinterpret_cast<int4 *>(reinterpret_cast<uint8_t *>(ptr) +
+                                   batch_ind * total_tokens * heads_per_rank * head_size_64b * sizeof(ELEM_T) +
+                                   dst_token_idx * heads_per_rank * head_size_64b * sizeof(ELEM_T)) +
+          copy_thr_col_idx;
+
+      // Non-allocating store to avoid polluting L1 cache
+      st_na_global(shuffled_buffer_ptr, __ldg(shuffled_x_ptr));
+    }
+  }
+
+  // === Barrier Synchronization ===
+  // Signal completion to target rank and wait for all ranks to finish
+  barrier_wait_and_reset_roundrobin(barrier_signal_ptrs, target_rank, rank, world_size, num_sms, sm_id, threadIdx.x,
+                                    timeout_cycles);
+}
+
+/**
+ * @brief All2All kernel for gathering attention heads back to original distribution.
+ *
+ * This kernel performs the inverse of send_recv_all2all: it gathers heads from
+ * all ranks back to reconstruct the original tensor layout. Each GPU reads from
+ * its local buffer and writes its portion of heads to all target ranks.
+ *
+ * Data layout transformation:
+ *   Source: [batch, total_tokens, heads_per_rank, head_size]  (per GPU)
+ *   Dest:   [batch, rank_tokens[target], num_heads, head_size]  (per target GPU)
+ *
+ * ## Memory Layout
+ *
+ * Input tensor x (this rank's portion after send_recv_all2all):
+ *   - Contains all tokens but only heads_per_rank heads
+ *   - Layout: [batch, total_tokens, heads_per_rank, head_size]
+ *
+ * Output buffer (per target rank):
+ *   - Contains only that rank's tokens but all heads
+ *   - Layout: [batch, rank_tokens[target], num_heads, head_size]
+ *   - This rank writes heads [rank * heads_per_rank : (rank+1) * heads_per_rank]
+ *
+ * @tparam ELEM_T Element type (at::BFloat16 or at::Float8_e4m3fn)
+ * @param buffer_ptrs Device array of pointers to each rank's data buffer
+ * @param barrier_signal_ptrs Device array of pointers to barrier signals
+ * @param x Source tensor data (this rank's buffer after send_recv)
+ * @param rank This GPU's rank
+ * @param world_size Total number of GPUs
+ * @param batch_size Number of batches
+ * @param num_heads Total number of heads (reconstructed)
+ * @param head_size Size of each attention head
+ * @param rank_tokens Number of tokens for each rank
+ * @param total_tokens Sum of tokens across all ranks
+ * @param prefix_rank_tokens Cumulative token counts for offset calculation
+ */
+template <typename ELEM_T>
+__global__ void gather_heads(void **buffer_ptrs, int **barrier_signal_ptrs, void *x, int rank, int world_size,
+                             int batch_size, int num_heads, int head_size, const int *__restrict__ rank_tokens,
+                             int total_tokens, int *prefix_rank_tokens, uint64_t timeout_cycles) {
+  // Grid dimensions
+  int num_sms = gridDim.x;
+  int sm_id = blockIdx.x;
+  int num_threads = blockDim.x;
+
+  // === SM Work Distribution (Round-Robin) ===
+  // Same partitioning as send_recv_all2all
+  int64_t target_rank = get_target_rank(sm_id, world_size);
+  int64_t rank_local_sm_id = get_rank_local_sm_id(sm_id, world_size);
+  int64_t num_sms_for_this_rank = get_num_sms_for_rank(target_rank, num_sms, world_size);
+  int64_t heads_per_rank = num_heads / world_size;
+
+  // === Thread Mapping ===
+  int64_t num_elems_per_thread = sizeof(int4) / sizeof(ELEM_T);
+  int64_t num_threads_per_token = heads_per_rank * head_size / num_elems_per_thread;
+  int64_t num_tokens_per_copy = num_threads / num_threads_per_token;
+
+  int64_t copy_thr_col_idx = threadIdx.x % num_threads_per_token;
+  int64_t copy_thr_row_idx = threadIdx.x / num_threads_per_token;
+
+  // Number of tokens owned by target rank
+  const int64_t tgt_tokens = int64_t(rank_tokens[target_rank]);
+
+  // This rank writes its heads at offset [rank * heads_per_rank] in the output
+  int64_t head_idx = rank * heads_per_rank;
+  int64_t num_heads_64b = int64_t(num_heads);
+  int64_t head_size_64b = int64_t(head_size);
+  int64_t total_tokens_64b = int64_t(total_tokens);
+
+  // Get target rank's buffer pointer
+  auto ptr = reinterpret_cast<void *>(static_cast<int8_t *>(buffer_ptrs[target_rank]));
+
+  // === Main Copy Loop ===
+  // Process target rank's tokens: read from global position, write to local position
+  for (int64_t batch_idx = 0; batch_idx < batch_size; batch_idx++) {
+    for (int64_t token_idx = rank_local_sm_id * num_tokens_per_copy; token_idx < tgt_tokens;
+         token_idx += num_tokens_per_copy * num_sms_for_this_rank) {
+      int64_t copy_token = token_idx + copy_thr_row_idx;
+      if (copy_token >= tgt_tokens)
+        break;
+
+      // Source: Read from global token position (target rank's tokens in our buffer)
+      int64_t src_token_idx = prefix_rank_tokens[target_rank] + copy_token;
+      // Destination: Write to local token position in target's buffer
+      int64_t dst_token_idx = copy_token;
+
+      // Source pointer: our input tensor at [batch, src_token, :, :]
+      int4 *shuffled_x_ptr =
+          reinterpret_cast<int4 *>(reinterpret_cast<uint8_t *>(x) +
+                                   batch_idx * total_tokens_64b * heads_per_rank * head_size_64b * sizeof(ELEM_T) +
+                                   src_token_idx * heads_per_rank * head_size_64b * sizeof(ELEM_T)) +
+          copy_thr_col_idx;
+
+      // Destination pointer: target's buffer at [batch, dst_token, head_idx:head_idx+hpr, :]
+      int4 *shuffled_buffer_ptr =
+          reinterpret_cast<int4 *>(reinterpret_cast<uint8_t *>(ptr) +
+                                   batch_idx * tgt_tokens * num_heads_64b * head_size_64b * sizeof(ELEM_T) +
+                                   dst_token_idx * num_heads_64b * head_size_64b * sizeof(ELEM_T) +
+                                   head_idx * head_size_64b * sizeof(ELEM_T)) +
+          copy_thr_col_idx;
+
+      st_na_global(shuffled_buffer_ptr, __ldg(shuffled_x_ptr));
+    }
+  }
+
+  // === Barrier Synchronization ===
+  barrier_wait_and_reset_roundrobin(barrier_signal_ptrs, target_rank, rank, world_size, num_sms, sm_id, threadIdx.x,
+                                    timeout_cycles);
+}
+
+/**
+ * @brief Host function to launch the gather_heads kernel.
+ *
+ * Selects the appropriate template instantiation based on tensor data type
+ * and launches the kernel with the specified number of SMs.
+ *
+ * @param buffer_ptrs Device array of buffer pointers
+ * @param barrier_signal_ptrs Device array of barrier signal pointers
+ * @param x Input tensor data pointer
+ * @param rank_tokens Token count per rank (device memory)
+ * @param prefix_rank_tokens Cumulative token counts (device memory)
+ * @param rank This GPU's rank
+ * @param world_size Total number of GPUs
+ * @param batch_size Number of batches
+ * @param total_tokens Sum of tokens across all ranks
+ * @param num_heads Total number of attention heads
+ * @param head_size Size of each attention head
+ * @param stream CUDA stream for async execution
+ * @param num_sms Number of SMs to launch
+ * @param tensor_dtype Data type (BFloat16 or Float8_e4m3fn)
+ */
+void all2all_head_gather_launch(void **buffer_ptrs, int **barrier_signal_ptrs, void *x, const int *rank_tokens,
+                                int *prefix_rank_tokens, int rank, int world_size, int batch_size, int total_tokens,
+                                int num_heads, int head_size, cudaStream_t stream, int num_sms,
+                                at::ScalarType tensor_dtype, uint64_t timeout_cycles) {
+  do {
+    if (tensor_dtype == at::ScalarType::BFloat16) {
+      gather_heads<at::BFloat16><<<num_sms, DEFAULT_KERNEL_THREADS, 0, stream>>>(
+          buffer_ptrs, barrier_signal_ptrs, x, rank, world_size, batch_size, num_heads, head_size, rank_tokens,
+          total_tokens, prefix_rank_tokens, timeout_cycles);
+    } else if (tensor_dtype == at::ScalarType::Float8_e4m3fn) {
+      gather_heads<at::Float8_e4m3fn><<<num_sms, DEFAULT_KERNEL_THREADS, 0, stream>>>(
+          buffer_ptrs, barrier_signal_ptrs, x, rank, world_size, batch_size, num_heads, head_size, rank_tokens,
+          total_tokens, prefix_rank_tokens, timeout_cycles);
+    }
+
+    // Check for kernel launch errors
+    cudaError_t e = cudaGetLastError();
+    if (e != cudaSuccess) {
+      EPException cuda_exception("CUDA", __FILE__, __LINE__, cudaGetErrorString(e));
+      fprintf(stderr, "%s\n", cuda_exception.what());
+      throw cuda_exception;
+    }
+  } while (0);
+}
+
+/**
+ * @brief Host function to launch the send_recv_all2all kernel.
+ *
+ * Selects the appropriate template instantiation based on tensor data type
+ * and launches the kernel with the specified number of SMs.
+ *
+ * @param buffer_ptrs Device array of buffer pointers
+ * @param barrier_signal_ptrs Device array of barrier signal pointers
+ * @param x Input tensor data pointer
+ * @param prefix_rank_tokens Cumulative token counts (device memory)
+ * @param rank This GPU's rank
+ * @param world_size Total number of GPUs
+ * @param batch_size Number of batches
+ * @param total_tokens Sum of tokens across all ranks
+ * @param num_tokens Number of tokens on this rank
+ * @param num_heads Total number of attention heads
+ * @param head_size Size of each attention head
+ * @param stream CUDA stream for async execution
+ * @param num_sms Number of SMs to launch
+ * @param tensor_dtype Data type (BFloat16 or Float8_e4m3fn)
+ */
+void all2all_head_launch(void **buffer_ptrs, int **barrier_signal_ptrs, void *x, int *prefix_rank_tokens, int rank,
+                         int world_size, int batch_size, int total_tokens, int num_tokens, int num_heads, int head_size,
+                         cudaStream_t stream, int num_sms, at::ScalarType tensor_dtype, uint64_t timeout_cycles) {
+  do {
+    if (tensor_dtype == at::ScalarType::BFloat16) {
+      send_recv_all2all<at::BFloat16><<<num_sms, DEFAULT_KERNEL_THREADS, 0, stream>>>(
+          buffer_ptrs, barrier_signal_ptrs, x, rank, world_size, batch_size, num_tokens, num_heads, head_size,
+          total_tokens, prefix_rank_tokens, timeout_cycles);
+    } else if (tensor_dtype == at::ScalarType::Float8_e4m3fn) {
+      send_recv_all2all<at::Float8_e4m3fn><<<num_sms, DEFAULT_KERNEL_THREADS, 0, stream>>>(
+          buffer_ptrs, barrier_signal_ptrs, x, rank, world_size, batch_size, num_tokens, num_heads, head_size,
+          total_tokens, prefix_rank_tokens, timeout_cycles);
+    }
+
+    // Check for kernel launch errors
+    cudaError_t e = cudaGetLastError();
+    if (e != cudaSuccess) {
+      EPException cuda_exception("CUDA", __FILE__, __LINE__, cudaGetErrorString(e));
+      fprintf(stderr, "%s\n", cuda_exception.what());
+      throw cuda_exception;
+    }
+  } while (0);
+}
+
+} // namespace all2all_cuda
+} // namespace all2all
+} // namespace ltx_kernels

packages/ltx-kernels/csrc/all2all/cuda/allgather.cu

@@ -0,0 +1,198 @@
+/**
+ * @file allgather.cu
+ * @brief CUDA kernel for AllGather operation using IPC-based direct memory access.
+ *
+ * This file implements the GPU kernel for gathering sequence tokens from all GPUs
+ * into a complete sequence on each GPU. Unlike the head redistribution kernels,
+ * this kernel preserves the head dimension and only gathers across the token
+ * (sequence) dimension.
+ *
+ * ## Algorithm Overview
+ *
+ * Each GPU broadcasts its local tokens to all other GPUs' buffers:
+ *   - GPU i writes its tokens to position [prefix_rank_tokens[i]] in each buffer
+ *   - After completion, all buffers contain the full sequence [0:total_tokens]
+ *
+ * ## Use Case
+ *
+ * This is typically used after tensor-parallel computation to reconstruct the
+ * full sequence for operations that require global context (e.g., output projection).
+ */
+
+#include "cuda/configs.cuh"
+#include "cuda/exceptions.cuh"
+#include "cuda/utils.cuh"
+#include <ATen/cuda/CUDADataType.h>
+
+namespace ltx_kernels {
+namespace all2all {
+namespace all2all_cuda {
+
+/**
+ * @brief AllGather kernel to collect sequence tokens from all ranks.
+ *
+ * Each GPU writes its local sequence tokens to all other GPUs' buffers at the
+ * appropriate offset. After synchronization, all GPUs have the complete sequence.
+ *
+ * Data layout transformation:
+ *   Input per GPU:  [batch, seqlen, hidden_dim]
+ *   Output per GPU: [batch, total_tokens, hidden_dim]  (identical on all GPUs)
+ *
+ * ## Memory Layout
+ *
+ * Input tensor x (contiguous):
+ *   - Shape: [batch, seqlen, hidden_dim]
+ *   - hidden_dim = num_heads * head_size (flattened)
+ *
+ * Output buffer (per target rank, after gather):
+ *   - Shape: [batch, total_tokens, hidden_dim]
+ *   - This rank's tokens placed at offset rank_tokens_prefix[rank]
+ *
+ * ## Thread Mapping
+ *
+ * Similar to all2all_heads, threads cooperate to copy tokens:
+ *   - Each thread copies 16 bytes (int4)
+ *   - Threads per token = hidden_dim * sizeof(ELEM_T) / sizeof(int4)
+ *   - Multiple tokens processed per thread block
+ *
+ * @tparam ELEM_T Element type (__nv_bfloat16 or at::Float8_e4m3fn)
+ * @param x Source tensor data pointer (this rank's tokens)
+ * @param buffer_ptrs Device array of pointers to each rank's data buffer
+ * @param barrier_signal_ptrs Device array of pointers to barrier signals
+ * @param batch_size Number of batches
+ * @param seqlen Number of tokens on this rank
+ * @param hidden_dim Hidden dimension size (num_heads * head_size)
+ * @param world_size Total number of GPUs
+ * @param rank This GPU's rank
+ * @param total_tokens Sum of tokens across all ranks
+ * @param rank_tokens_prefix Cumulative token counts (device memory)
+ */
+template <typename ELEM_T>
+__global__ void allgather(void *x, void **buffer_ptrs, int **barrier_signal_ptrs, int batch_size, int seqlen,
+                          int hidden_dim, int world_size, int rank, int total_tokens, int *rank_tokens_prefix,
+                          uint64_t timeout_cycles) {
+
+  // Grid dimensions
+  int num_sms = gridDim.x;
+  int sm_id = blockIdx.x;
+  int num_threads = blockDim.x;
+
+  // === SM Work Distribution (Round-Robin) ===
+  // Use modular assignment to handle num_sms not divisible by world_size.
+  // This ensures all SMs are utilized: some ranks get ceil(num_sms/world_size)
+  // SMs, others get floor(num_sms/world_size) SMs.
+  int tgt_rank = get_target_rank(sm_id, world_size);
+  int rank_local_sm_id = get_rank_local_sm_id(sm_id, world_size);
+  int num_sms_for_this_rank = get_num_sms_for_rank(tgt_rank, num_sms, world_size);
+
+  // Get target rank's buffer pointer
+  auto ptr = reinterpret_cast<void *>(static_cast<int8_t *>(buffer_ptrs[tgt_rank]));
+
+  // === Thread Mapping ===
+  // Each thread copies one int4 (16 bytes)
+  int64_t num_elems_per_thread = sizeof(int4) / sizeof(ELEM_T);
+  int64_t num_threads_per_token = hidden_dim / num_elems_per_thread;
+  int64_t num_tokens_per_copy = num_threads / num_threads_per_token;
+
+  // 2D thread coordinates
+  int64_t copy_thr_col_idx = threadIdx.x % num_threads_per_token; // Element offset
+  int64_t copy_thr_row_idx = threadIdx.x / num_threads_per_token; // Token offset
+
+  // Use 64-bit arithmetic to avoid overflow
+  int64_t hidden_dim_64b = int64_t(hidden_dim);
+  int64_t total_tokens_64b = int64_t(total_tokens);
+  int64_t seqlen_64b = int64_t(seqlen);
+
+  // === Main Copy Loop ===
+  // Broadcast this rank's tokens to all target ranks' buffers
+  for (int64_t batch_idx = 0; batch_idx < batch_size; batch_idx++) {
+    // Strided token iteration within SM group for this target rank
+    for (int64_t token_idx = rank_local_sm_id * num_tokens_per_copy; token_idx < seqlen;
+         token_idx += num_tokens_per_copy * num_sms_for_this_rank) {
+      int64_t copy_token = token_idx + copy_thr_row_idx;
+      if (copy_token >= seqlen)
+        break;
+
+      // Source: local token index in input tensor
+      int64_t src_token_idx = copy_token;
+      // Destination: global token index in output buffer
+      // This rank's tokens start at prefix_rank_tokens[rank]
+      int64_t dst_token_idx = copy_token + rank_tokens_prefix[rank];
+
+      // Source pointer: input tensor at [batch, src_token, :]
+      int4 *shuffled_x_ptr = reinterpret_cast<int4 *>(reinterpret_cast<uint8_t *>(x) +
+                                                      batch_idx * seqlen_64b * hidden_dim_64b * sizeof(ELEM_T) +
+                                                      src_token_idx * hidden_dim_64b * sizeof(ELEM_T)) +
+                             copy_thr_col_idx;
+
+      // Destination pointer: target buffer at [batch, dst_token, :]
+      int4 *shuffled_buffer_ptr =
+          reinterpret_cast<int4 *>(reinterpret_cast<uint8_t *>(ptr) +
+                                   batch_idx * total_tokens_64b * hidden_dim_64b * sizeof(ELEM_T) +
+                                   dst_token_idx * hidden_dim_64b * sizeof(ELEM_T)) +
+          copy_thr_col_idx;
+
+      // Non-allocating store for better cache behavior
+      st_na_global(shuffled_buffer_ptr, __ldg(shuffled_x_ptr));
+    }
+  }
+
+  // === Barrier Synchronization ===
+  // Signal completion to target rank and wait for all ranks
+  // Use round-robin variant since SM counts per rank may differ
+  barrier_wait_and_reset_roundrobin(barrier_signal_ptrs, tgt_rank, rank, world_size, num_sms, sm_id, threadIdx.x,
+                                    timeout_cycles);
+}
+
+/**
+ * @brief Host function to launch the allgather kernel.
+ *
+ * Launches the AllGather kernel with the specified configuration.
+ * Uses ALLGATHER_KERNEL_THREADS (1024) threads per block for higher
+ * occupancy than the All2All kernels.
+ *
+ * @param buffer_ptrs Device array of buffer pointers
+ * @param barrier_signal_ptrs Device array of barrier signal pointers
+ * @param x Input tensor data pointer
+ * @param prefix_rank_tokens Cumulative token counts (device memory)
+ * @param rank This GPU's rank
+ * @param world_size Total number of GPUs
+ * @param batch_size Number of batches
+ * @param seqlen Number of tokens on this rank
+ * @param hidden_dim Hidden dimension size
+ * @param total_tokens Sum of tokens across all ranks
+ * @param stream CUDA stream for async execution
+ * @param num_sms Number of SMs to launch
+ * @param tensor_dtype Data type (BFloat16 or Float8_e4m3fn)
+ */
+void allgather_launch(void **buffer_ptrs, int **barrier_signal_ptrs, void *x, int *prefix_rank_tokens, int rank,
+                      int world_size, int batch_size, int seqlen, int hidden_dim, int total_tokens, cudaStream_t stream,
+                      int num_sms, at::ScalarType tensor_dtype, uint64_t timeout_cycles) {
+  do {
+    if (tensor_dtype == at::ScalarType::BFloat16) {
+      allgather<at::BFloat16><<<num_sms, ALLGATHER_KERNEL_THREADS, 0, stream>>>(
+          x, buffer_ptrs, barrier_signal_ptrs, batch_size, seqlen, hidden_dim, world_size, rank, total_tokens,
+          prefix_rank_tokens, timeout_cycles);
+    } else if (tensor_dtype == at::ScalarType::Float8_e4m3fn) {
+      allgather<at::Float8_e4m3fn><<<num_sms, ALLGATHER_KERNEL_THREADS, 0, stream>>>(
+          x, buffer_ptrs, barrier_signal_ptrs, batch_size, seqlen, hidden_dim, world_size, rank, total_tokens,
+          prefix_rank_tokens, timeout_cycles);
+    } else {
+      EPException dtype_exception("allgather_launch", __FILE__, __LINE__, "Unsupported dtype");
+      fprintf(stderr, "%s\n", dtype_exception.what());
+      throw dtype_exception;
+    }
+
+    // Check for kernel launch errors
+    cudaError_t e = cudaGetLastError();
+    if (e != cudaSuccess) {
+      EPException cuda_exception("CUDA", __FILE__, __LINE__, cudaGetErrorString(e));
+      fprintf(stderr, "%s\n", cuda_exception.what());
+      throw cuda_exception;
+    }
+  } while (0);
+}
+
+} // namespace all2all_cuda
+} // namespace all2all
+} // namespace ltx_kernels

packages/ltx-kernels/csrc/all2all/cuda/api.cuh

@@ -0,0 +1,99 @@
+/**
+ * @file api.cuh
+ * @brief CUDA kernel launch function declarations for All2All operations.
+ *
+ * This header provides the host-callable interface for launching the All2All
+ * CUDA kernels. These functions handle template instantiation and kernel
+ * configuration based on the tensor data type.
+ */
+
+#pragma once
+
+#include <ATen/cuda/CUDADataType.h>
+#include <vector>
+
+namespace ltx_kernels {
+namespace all2all {
+namespace all2all_cuda {
+
+/**
+ * @brief Launches the All2All head redistribution kernel.
+ *
+ * Redistributes attention heads across GPUs:
+ *   Input:  [batch, num_tokens, num_heads, head_size] per GPU
+ *   Output: [batch, total_tokens, num_heads/world_size, head_size] per GPU
+ *
+ * @param buffer_ptrs Device array of pointers to each rank's data buffer
+ * @param barrier_signal_ptrs Device array of pointers to barrier signals
+ * @param x Source tensor data pointer
+ * @param prefix_rank_tokens Cumulative token counts per rank (device memory)
+ * @param rank This GPU's rank (0 to world_size-1)
+ * @param world_size Total number of GPUs
+ * @param batch_size Batch dimension size
+ * @param total_tokens Sum of tokens across all ranks
+ * @param num_tokens Number of tokens on this rank
+ * @param num_heads Total number of attention heads
+ * @param head_size Size of each attention head
+ * @param stream CUDA stream for async execution
+ * @param num_sms Number of SMs to use for the kernel
+ * @param tensor_dtype Data type (BFloat16 or Float8_e4m3fn)
+ */
+void all2all_head_launch(void **buffer_ptrs, int **barrier_signal_ptrs, void *x, int *prefix_rank_tokens, int rank,
+                         int world_size, int batch_size, int total_tokens, int num_tokens, int num_heads, int head_size,
+                         cudaStream_t stream, int num_sms, at::ScalarType tensor_dtype, uint64_t timeout_cycles);
+
+/**
+ * @brief Launches the gather heads kernel (inverse of all2all_head_launch).
+ *
+ * Redistributes tokens back to original head distribution:
+ *   Input:  [batch, total_tokens, heads_per_rank, head_size] per GPU
+ *   Output: [batch, rank_tokens[rank], num_heads, head_size] per GPU
+ *
+ * @param buffer_ptrs Device array of pointers to each rank's data buffer
+ * @param barrier_signal_ptrs Device array of pointers to barrier signals
+ * @param x Source tensor data pointer
+ * @param rank_tokens Token count for each rank (device memory)
+ * @param prefix_rank_tokens Cumulative token counts (device memory)
+ * @param rank This GPU's rank
+ * @param world_size Total number of GPUs
+ * @param batch_size Batch dimension size
+ * @param total_tokens Sum of tokens across all ranks
+ * @param num_heads Total number of attention heads (reconstructed)
+ * @param head_size Size of each attention head
+ * @param stream CUDA stream for async execution
+ * @param num_sms Number of SMs to use for the kernel
+ * @param tensor_dtype Data type (BFloat16 or Float8_e4m3fn)
+ */
+void all2all_head_gather_launch(void **buffer_ptrs, int **barrier_signal_ptrs, void *x, const int *rank_tokens,
+                                int *prefix_rank_tokens, int rank, int world_size, int batch_size, int total_tokens,
+                                int num_heads, int head_size, cudaStream_t stream, int num_sms,
+                                at::ScalarType tensor_dtype, uint64_t timeout_cycles);
+
+/**
+ * @brief Launches the AllGather kernel for sequence tokens.
+ *
+ * Gathers sequence tokens from all ranks:
+ *   Input:  [batch, seqlen, hidden_dim] per GPU
+ *   Output: [batch, total_tokens, hidden_dim] per GPU (identical on all)
+ *
+ * @param buffer_ptrs Device array of pointers to each rank's data buffer
+ * @param barrier_signal_ptrs Device array of pointers to barrier signals
+ * @param x Source tensor data pointer
+ * @param prefix_rank_tokens Cumulative token counts (device memory)
+ * @param rank This GPU's rank
+ * @param world_size Total number of GPUs
+ * @param batch_size Batch dimension size
+ * @param seqlen Number of tokens on this rank
+ * @param hidden_dim Hidden dimension size (num_heads * head_size)
+ * @param total_tokens Sum of tokens across all ranks
+ * @param stream CUDA stream for async execution
+ * @param num_sms Number of SMs to use for the kernel
+ * @param tensor_dtype Data type (BFloat16 or Float8_e4m3fn)
+ */
+void allgather_launch(void **buffer_ptrs, int **barrier_signal_ptrs, void *x, int *prefix_rank_tokens, int rank,
+                      int world_size, int batch_size, int seqlen, int hidden_dim, int total_tokens, cudaStream_t stream,
+                      int num_sms, at::ScalarType tensor_dtype, uint64_t timeout_cycles);
+
+} // namespace all2all_cuda
+} // namespace all2all
+} // namespace ltx_kernels

packages/ltx-kernels/csrc/blockwise/api.cpp

@@ -0,0 +1,89 @@
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <torch/extension.h>
+#include <vector>
+#include <stdio.h>
+
+#ifdef __SM90__
+#include "sm90_fp8_gemm_1d2d_bias.hpp"
+#endif
+
+#include "sm89_fp8_gemm_1d2d.hpp"
+
+namespace blockwise{
+template <int N>
+static auto get_shape(const torch::Tensor& t) {
+    return [&t] <size_t... Is> (std::index_sequence<Is...>) {
+        return std::make_tuple(static_cast<int>(t.sizes()[Is])...);
+    }(std::make_index_sequence<N>());
+}
+
+#ifdef __SM90__
+static void fp8_gemm_nt_sm90(const std::pair<torch::Tensor, torch::Tensor>& a,
+                        const std::pair<torch::Tensor, torch::Tensor>& b,
+                        const torch::Tensor& d,
+                        const std::optional<torch::Tensor>& bias,
+                        const std::optional<torch::Tensor>& c, const int num_sms) {
+
+    // Type and shape checks
+    const auto& [m , k ] = get_shape<2>(a.first);
+    const auto& [n , k_] = get_shape<2>(b.first);
+    const auto& [m_, n_] = get_shape<2>(d);
+
+    // The SM90 kernel always adds bias; synthesize a zero bias when the layer is
+    // bias-less (e.g. the no-bias video FFN of v3 checkpoints), mirroring SM89 below.
+    torch::Tensor bias_tensor = bias.has_value()
+        ? bias.value()
+        : torch::zeros({n}, d.options().dtype(torch::kFloat32));
+    sm90_fp8_gemm_1d2d_bias(a.first, a.second, b.first, b.second, bias_tensor, c, d, m, n, k, num_sms);
+}
+#endif
+
+static void fp8_gemm_nt_sm89(const std::pair<torch::Tensor, torch::Tensor>& a,
+                             const std::pair<torch::Tensor, torch::Tensor>& b,
+                             const torch::Tensor& d,
+                             const std::optional<torch::Tensor>& bias,
+                             const bool use_fast_accum = true) {
+
+    const auto& [m, k] = get_shape<2>(a.first);
+    const auto& [n, k_] = get_shape<2>(b.first);
+    const auto& [m_, n_] = get_shape<2>(d);
+
+    // The SM89 kernel always adds bias; synthesize a zero bias when the layer is
+    // bias-less so we add 0 rather than uninitialized memory (mirrors SM90 above).
+    torch::Tensor bias_tensor = bias.has_value()
+        ? bias.value()
+        : torch::zeros({n}, d.options().dtype(torch::kFloat32));
+
+    blockwise::sm89_fp8_gemm_1d2d_bias(
+        a.first, a.second,    // a data, sfa scales
+        b.first, b.second,    // b data, sfb scales
+        bias_tensor,          // bias (or empty tensor)
+        d,                    // output
+        m, n, k,
+        use_fast_accum);      // pass through accumulation mode
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+    // m.def("package_name", &function_name, "function_docstring"")
+#ifdef __SM90__
+    m.def("fp8_gemm_nt_sm90", &fp8_gemm_nt_sm90,
+          py::arg("a"), py::arg("b"), py::arg("d"),
+          py::arg("bias") = std::nullopt,
+          py::arg("c") = std::nullopt,
+          py::arg("num_sms") = 132
+    );
+#endif
+    m.def("fp8_gemm_nt_sm89", &fp8_gemm_nt_sm89,
+          py::arg("a"), py::arg("b"), py::arg("d"),
+          py::arg("bias") = std::nullopt,
+          py::arg("use_fast_accum") = true
+    );
+}
+};
+
+
+
+
+
+

packages/ltx-kernels/csrc/blockwise/config.hpp

@@ -0,0 +1,92 @@
+#pragma once
+#include <torch/python.h>
+#include <cute/arch/mma_sm100_umma.hpp>
+#include "utils.hpp"
+#include "exceptions.hpp"
+
+namespace blockwise{
+struct MulticastConfig {
+    int num_multicast;
+    bool is_multicast_on_a;
+
+    MulticastConfig(const int& num_multicast, const bool& is_multicast_on_a):
+        num_multicast(num_multicast), is_multicast_on_a(is_multicast_on_a) {
+        DG_HOST_ASSERT(1 <= num_multicast and num_multicast <= 2);
+    }
+};
+
+struct SharedMemoryConfig {
+    int smem_size;
+    int swizzle_a_mode;
+    int swizzle_b_mode;
+    int swizzle_cd_mode;
+};
+
+struct ThreadConfig {
+    int num_threads;
+
+    // SM90
+    int num_tma_threads;
+    int num_math_threads;
+
+    // SM100
+    int num_non_epilogue_threads;
+    int num_epilogue_threads;
+
+    static ThreadConfig sm90(const int& num_tma_threads,
+                             const int& num_math_threads) {
+        auto config = ThreadConfig();
+        config.num_threads = num_tma_threads + num_math_threads;
+        config.num_tma_threads = num_tma_threads;
+        config.num_math_threads = num_math_threads;
+        return config;
+    }
+
+    static ThreadConfig sm100(const int& num_non_epilogue_threads,
+                              const int& num_epilogue_threads) {
+        auto config = ThreadConfig();
+        config.num_threads = num_non_epilogue_threads + num_epilogue_threads;
+        config.num_non_epilogue_threads = num_non_epilogue_threads;
+        config.num_epilogue_threads = num_epilogue_threads;
+        return config;
+    }
+};
+
+template<int SM>
+struct GemmConfig{};
+// {
+//     // Templated configs
+    
+//     at::ScalarType ab_dtype, cd_dtype;
+//     bool with_accumulation;
+//     int block_m, block_n, block_k;
+//     int num_stages, num_last_stages;
+
+//     // Templated device configs
+//     int num_sms;
+
+//     // Structured configs
+//     MulticastConfig multicast_config;
+//     SharedMemoryConfig smem_config;
+//     ThreadConfig thread_config;
+// };
+
+
+template <> 
+struct GemmConfig<90>
+{
+    at::ScalarType ab_dtype = torch::kFloat8_e4m3fn;
+    at::ScalarType cd_dtype = torch::kBFloat16;
+    bool with_accumulation = false; 
+    int block_m = 256;
+    int block_n = 128;
+    int block_k = 128;
+    int num_stages = 3;
+    int num_last_stages = 2;
+    int num_sms = 132;
+    MulticastConfig multicast_config{2, true};
+    SharedMemoryConfig smem_config{216240, 128, 128, 128};
+    ThreadConfig thread_config = ThreadConfig::sm90(128, 256);
+};
+
+};

packages/ltx-kernels/csrc/blockwise/exceptions.hpp

@@ -0,0 +1,65 @@
+#pragma once
+
+#include <exception>
+#include <string>
+#include <sstream>
+
+namespace blockwise {
+
+class DGException final : public std::exception {
+    std::string message = {};
+
+public:
+    explicit DGException(const char *name, const char* file, const int line, const std::string& error) {
+        message = std::string(name) + " error (" + file + ":" + std::to_string(line) + "): " + error;
+    }
+
+    const char *what() const noexcept override {
+        return message.c_str();
+    }
+};
+
+#ifndef DG_STATIC_ASSERT
+#define DG_STATIC_ASSERT(cond, ...) static_assert(cond, __VA_ARGS__)
+#endif
+
+#ifndef DG_HOST_ASSERT
+#define DG_HOST_ASSERT(cond) \
+do { \
+    if (not (cond)) { \
+        throw DGException("Assertion", __FILE__, __LINE__, #cond); \
+    } \
+} while (0)
+#endif
+
+#ifndef DG_HOST_UNREACHABLE
+#define DG_HOST_UNREACHABLE(reason) (throw DGException("Assertion", __FILE__, __LINE__, reason))
+#endif
+
+// #ifndef DG_CUDA_DRIVER_CHECK
+// #define DG_CUDA_DRIVER_CHECK(cmd) \
+// do { \
+//     const auto& e = (cmd); \
+//     if (e != CUDA_SUCCESS) { \
+//         std::stringstream ss; \
+//         const char *name, *info; \
+//         cuGetErrorName(e, &name), cuGetErrorString(e, &info); \
+//         ss << static_cast<int>(e) << " (" << name << ", " << info << ")"; \
+//         throw DGException("CUDA driver", __FILE__, __LINE__, ss.str()); \
+//     } \
+// } while (0)
+// #endif
+
+#ifndef DG_CUDA_RUNTIME_CHECK
+#define DG_CUDA_RUNTIME_CHECK(cmd) \
+do { \
+    const auto& e = (cmd); \
+    if (e != cudaSuccess) { \
+        std::stringstream ss; \
+        ss << static_cast<int>(e) << " (" << cudaGetErrorName(e) << ", " << cudaGetErrorString(e) << ")"; \
+        throw DGException("CUDA runtime", __FILE__, __LINE__, ss.str()); \
+    } \
+} while (0)
+#endif
+
+} // namespace deep_gemm

packages/ltx-kernels/csrc/blockwise/kernels/deep_gemm/include/deep_gemm/common/cute_tie.cuh

@@ -0,0 +1,48 @@
+#pragma once
+
+namespace cute {
+
+struct ignore_t {
+    template <typename T>
+    constexpr const ignore_t& operator=(T&&) const noexcept {
+        return *this;
+    }
+};
+
+inline constexpr ignore_t ignore{};
+
+} // namespace cute
+
+#define CUTE_TIE_CONCAT_IMPL(A, B) A##B
+#define CUTE_TIE_CONCAT(A, B) CUTE_TIE_CONCAT_IMPL(A, B)
+
+#define CUTE_TIE_GET_NTH_ARG(_1, _2, _3, _4, _5, _6, _7, _8, _9, _10, N, ...) N
+#define CUTE_TIE_COUNT_ARGS(...) \
+    CUTE_TIE_GET_NTH_ARG(__VA_ARGS__, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0)
+
+#define CUTE_TIE_OP_DECL(I, TUPLE, VAR) auto VAR = ::cute::get<I>(TUPLE)
+#define CUTE_TIE_OP_ASSIGN(I, TUPLE, VAR) VAR = ::cute::get<I>(TUPLE)
+
+#define CUTE_TIE_APPLY_OP_1(OP, T, V1) OP(0, T, V1);
+#define CUTE_TIE_APPLY_OP_2(OP, T, V1, V2) OP(0, T, V1); OP(1, T, V2);
+#define CUTE_TIE_APPLY_OP_3(OP, T, V1, V2, V3) OP(0, T, V1); OP(1, T, V2); OP(2, T, V3);
+#define CUTE_TIE_APPLY_OP_4(OP, T, V1, V2, V3, V4) OP(0, T, V1); OP(1, T, V2); OP(2, T, V3); OP(3, T, V4);
+#define CUTE_TIE_APPLY_OP_5(OP, T, V1, V2, V3, V4, V5) OP(0, T, V1); OP(1, T, V2); OP(2, T, V3); OP(3, T, V4); OP(4, T, V5);
+
+#define CUTE_TIE_DECL(TUPLE_EXPR, ...) \
+    auto&& CUTE_TIE_CONCAT(cute_tie__temp_tuple_, __LINE__) = (TUPLE_EXPR); \
+    CUTE_TIE_CONCAT(CUTE_TIE_APPLY_OP_, CUTE_TIE_COUNT_ARGS(__VA_ARGS__)) ( \
+        CUTE_TIE_OP_DECL, \
+        CUTE_TIE_CONCAT(cute_tie__temp_tuple_, __LINE__), \
+        __VA_ARGS__ \
+    )
+
+#define CUTE_TIE(TUPLE_EXPR, ...) \
+    do { \
+        auto&& CUTE_TIE_CONCAT(cute_tie__temp_tuple_, __LINE__) = (TUPLE_EXPR); \
+        CUTE_TIE_CONCAT(CUTE_TIE_APPLY_OP_, CUTE_TIE_COUNT_ARGS(__VA_ARGS__)) ( \
+            CUTE_TIE_OP_ASSIGN, \
+            CUTE_TIE_CONCAT(cute_tie__temp_tuple_, __LINE__), \
+            __VA_ARGS__ \
+        ); \
+    } while (0)

packages/ltx-kernels/csrc/blockwise/kernels/deep_gemm/include/deep_gemm/common/epilogue_utils.cuh

@@ -0,0 +1,27 @@
+#pragma once
+
+#include <deep_gemm/common/types.hpp>
+#include <deep_gemm/common/utils.cuh>
+
+namespace deep_gemm {
+
+struct EpilogueIdentity {
+    template <uint32_t STORE_BLOCK_N>
+    __device__ __forceinline__ static uint32_t apply_index_n(const uint32_t &n_idx) {
+        return n_idx;
+    }
+};
+
+template <uint32_t kLeft, uint32_t kMid, uint32_t kRight>
+struct EpilogueHeadSplits: EpilogueIdentity {
+    template <uint32_t STORE_BLOCK_N>
+    __device__ __forceinline__ static uint32_t apply_index_n(const uint32_t &n_idx) {
+        DG_STATIC_ASSERT(kLeft % STORE_BLOCK_N == 0 and kMid % STORE_BLOCK_N == 0 
+                         and kRight % STORE_BLOCK_N == 0, "Invalid head splits config");
+        return n_idx + (n_idx + kRight) / (kLeft + kRight) * kMid;
+    }
+};
+
+#pragma clang diagnostic pop
+
+} // namespace deep_gemm

packages/ltx-kernels/csrc/blockwise/kernels/deep_gemm/include/deep_gemm/common/reduction.cuh

@@ -0,0 +1,44 @@
+#pragma once
+
+#include <cuda_bf16.h>
+#include <cuda_fp8.h>
+#include <cuda/std/cstdint>
+#include <cuda/std/utility>
+
+#include <deep_gemm/common/utils.cuh>
+
+// Operation functors
+template <typename T> struct ReduceSum { __device__ T operator()(T a, T b) const { return a + b; } };
+template <typename T> struct ReduceMax { __device__ T operator()(T a, T b) const { return a > b ? a : b; } };
+template <typename T> struct ReduceMin { __device__ T operator()(T a, T b) const { return a < b ? a : b; } };
+template <typename T> struct ReduceAnd { __device__ T operator()(T a, T b) const { return a & b; } };
+template <typename T> struct ReduceOr  { __device__ T operator()(T a, T b) const { return a | b; } };
+
+// Unified reduction function
+template <int kNumLanesPerGroup, bool kIntergroupReduce, typename T, typename Op>
+__forceinline__ __device__ T warp_reduce(T value, Op op) {
+    DG_STATIC_ASSERT(kNumLanesPerGroup == 32 or kNumLanesPerGroup == 16 or kNumLanesPerGroup == 8 or
+                     kNumLanesPerGroup ==  4 or kNumLanesPerGroup == 2  or kNumLanesPerGroup == 1,
+                     "Invalid number of lanes");
+    constexpr uint32_t mask = 0xffffffff;
+    if constexpr (kIntergroupReduce) {
+        if constexpr (kNumLanesPerGroup <=  1) value = op(value, __shfl_xor_sync(mask, value,  1));
+        if constexpr (kNumLanesPerGroup <=  2) value = op(value, __shfl_xor_sync(mask, value,  2));
+        if constexpr (kNumLanesPerGroup <=  4) value = op(value, __shfl_xor_sync(mask, value,  4));
+        if constexpr (kNumLanesPerGroup <=  8) value = op(value, __shfl_xor_sync(mask, value,  8));
+        if constexpr (kNumLanesPerGroup <= 16) value = op(value, __shfl_xor_sync(mask, value, 16));
+    } else {
+        if constexpr (kNumLanesPerGroup >= 32) value = op(value, __shfl_xor_sync(mask, value, 16));
+        if constexpr (kNumLanesPerGroup >= 16) value = op(value, __shfl_xor_sync(mask, value,  8));
+        if constexpr (kNumLanesPerGroup >=  8) value = op(value, __shfl_xor_sync(mask, value,  4));
+        if constexpr (kNumLanesPerGroup >=  4) value = op(value, __shfl_xor_sync(mask, value,  2));
+        if constexpr (kNumLanesPerGroup >=  2) value = op(value, __shfl_xor_sync(mask, value,  1));
+    }
+    return value;
+}
+
+// Convenience aliases
+template <int kNumLanesPerGroup = 32, bool kIntergroupReduce = false, typename T>
+__forceinline__ __device__ T warp_reduce_sum(T value) {
+    return warp_reduce<kNumLanesPerGroup, kIntergroupReduce, T>(value, ReduceSum<T>{});
+}

packages/ltx-kernels/csrc/blockwise/kernels/deep_gemm/include/deep_gemm/common/scheduler.cuh

@@ -0,0 +1,239 @@
+#pragma once
+
+#include <deep_gemm/common/types.hpp>
+#include <deep_gemm/common/utils.cuh>
+
+namespace deep_gemm {
+
+enum class KGroupedIndexType {
+    MN,
+    K,
+    SF_K,
+};
+
+template <GemmType kGemmType, uint32_t BLOCK_M, uint32_t BLOCK_N, uint32_t kNumSMs, bool kIsMulticastOnA>
+static constexpr uint32_t get_num_1d_blocks_per_group() {
+    // Select the best from candidates
+    uint32_t num_best_blocks = 0, min_usage = cute::numeric_limits<uint32_t>::max();
+    for (const auto& candidate: {8u, 16u}) {
+        const auto& usage = kIsMulticastOnA ?
+                    candidate * BLOCK_N + constexpr_ceil_div(kNumSMs, candidate) * BLOCK_M: // Grouping on N
+                    candidate * BLOCK_M + constexpr_ceil_div(kNumSMs, candidate) * BLOCK_N; // Grouping on M
+        if (usage < min_usage)
+            min_usage = usage, num_best_blocks = candidate;
+    }
+    return num_best_blocks;
+}
+
+#pragma clang diagnostic push
+#pragma ide diagnostic ignored "cppcoreguidelines-pro-type-member-init"
+template <GemmType kGemmType,
+          uint32_t BLOCK_M, uint32_t BLOCK_N,
+          uint32_t kNumGroups,
+          uint32_t kNumMulticast, bool kIsMulticastOnA,
+          uint32_t kNumSMs,
+          uint32_t SF_K_ALIGNMENT = 512u,  // for k-grouped GEMM only: 128 (SM90 float SF) or 512 (SM100 UE8M0 SF)
+          uint32_t kNum1DBlocksPerGroup = get_num_1d_blocks_per_group<kGemmType, BLOCK_M, BLOCK_N, kNumSMs, kIsMulticastOnA>()>
+struct Scheduler {
+    int current_iter = -1;
+
+    // Block configs
+    uint32_t num_blocks;
+    uint32_t num_m_blocks;
+    uint32_t num_n_blocks;
+
+    // For SM90 multicast checks
+    uint32_t num_blocks_in_group;
+    bool is_peer_cta_alive = true;
+
+    // For grouped GEMM
+    int* grouped_layout;
+    uint32_t current_group_idx = 0;
+    // Only used for masked layout
+    uint32_t current_m_cumsum = 0;
+    // Only used for k-grouped layout
+    uint32_t current_shape_k, current_num_valid_groups = 0, current_k_cumsum = 0, current_sf_k_cumsum = 0;
+    uint32_t next_group_idx, next_shape_k;
+
+    // Only used for k-grouped gemm
+    __device__ __forceinline__ void get_next_k_group(uint32_t &group_idx, uint32_t &shape_k) const {
+        for (; group_idx < kNumGroups; ++ group_idx) {
+            shape_k = __ldg(grouped_layout + group_idx);
+            if (shape_k > 0)
+                break;
+        }
+    }
+
+    // ReSharper disable once CppPossiblyUninitializedMember
+    __device__ __forceinline__ explicit Scheduler(const uint32_t& shape_m, const uint32_t& shape_n, const uint32_t& shape_k,
+                                                  int* grouped_layout = nullptr) {
+        num_m_blocks = ceil_div(shape_m, BLOCK_M);
+        num_n_blocks = ceil_div(shape_n, BLOCK_N);
+        current_shape_k = shape_k;
+        if constexpr (kGemmType == GemmType::Normal) {
+            num_blocks = num_m_blocks * num_n_blocks;
+        } else if (kGemmType == GemmType::MGroupedContiguous) {
+            num_blocks = num_m_blocks * num_n_blocks;
+            this->grouped_layout = grouped_layout;
+        } else if (kGemmType == GemmType::MGroupedMasked) {
+            this->grouped_layout = grouped_layout;
+        } else if (kGemmType == GemmType::KGroupedContiguous) {
+            this->grouped_layout = grouped_layout;
+            get_next_k_group(current_group_idx, current_shape_k);
+            next_group_idx = current_group_idx + 1;
+            get_next_k_group(next_group_idx, next_shape_k);
+        }
+    }
+
+    __device__ __forceinline__ void get_swizzled_block_idx(const uint32_t& block_idx, uint32_t& m_block_idx, uint32_t& n_block_idx) {
+        DG_STATIC_ASSERT(kNum1DBlocksPerGroup % kNumMulticast == 0, "Invalid group size");
+
+        // Swizzle for better L2 usages
+        const auto& primary_num_blocks = kIsMulticastOnA ? num_n_blocks : num_m_blocks;
+        const auto& secondary_num_blocks = kIsMulticastOnA ? num_m_blocks : num_n_blocks;
+        const auto& num_blocks_per_group = secondary_num_blocks * kNum1DBlocksPerGroup;
+        const auto& group_idx = block_idx / num_blocks_per_group;
+        auto first_block_idx = group_idx * kNum1DBlocksPerGroup;
+        auto in_group_idx = block_idx % num_blocks_per_group;
+        num_blocks_in_group = min(kNum1DBlocksPerGroup, primary_num_blocks - first_block_idx);
+
+        // Fix unaligned TMA multicast
+        // NOTES: for SM90 only, as SM90 can dynamically disable TMA multicast
+        // while SM100 uses 2-CTA, which can not be dynamically disabled
+#if __CUDA_ARCH__ < 1000
+        if (kNumMulticast > 1 and num_blocks_in_group % 2 != 0) {
+            if (in_group_idx < (num_blocks_in_group ^ 1) * secondary_num_blocks) {
+                num_blocks_in_group = num_blocks_in_group ^ 1;
+            } else {
+                in_group_idx = in_group_idx - (num_blocks_in_group ^ 1) * secondary_num_blocks;
+                first_block_idx += num_blocks_in_group ^ 1;
+                num_blocks_in_group = 1;
+            }
+        }
+#endif
+
+        // Convert to final M/N block indices
+        // `kIsMulticastOnA == true` leads to groups on N
+        if constexpr (kIsMulticastOnA) {
+            m_block_idx = in_group_idx / num_blocks_in_group;
+            n_block_idx = first_block_idx + in_group_idx % num_blocks_in_group;
+        } else {
+            m_block_idx = first_block_idx + in_group_idx % num_blocks_in_group;
+            n_block_idx = in_group_idx / num_blocks_in_group;
+        }
+    }
+
+    template <bool kWithGroupOffset, KGroupedIndexType kIndexType = KGroupedIndexType::MN>
+    __device__ __forceinline__ uint32_t get_global_idx(const uint32_t shape_dim, const uint32_t block_size,
+                                                       const uint32_t& block_idx, const uint32_t& m_block_idx = 0) {
+        if constexpr (kGemmType == GemmType::Normal) {
+            return block_idx * block_size;
+        } else if constexpr (kGemmType == GemmType::MGroupedContiguous) {
+            const auto offset = kWithGroupOffset ? cute::max(0, __ldg(grouped_layout + m_block_idx * BLOCK_M)) : 0;
+            return offset * shape_dim + block_idx * block_size;
+        } else if constexpr (kGemmType == GemmType::MGroupedMasked) {
+            const auto offset = kWithGroupOffset ? current_group_idx : 0;
+            return offset * shape_dim + block_idx * block_size;
+        } else if constexpr (kGemmType == GemmType::KGroupedContiguous) {
+            auto offset = 0;
+            if constexpr (kWithGroupOffset) {
+                if constexpr (kIndexType == KGroupedIndexType::MN)
+                    offset = current_group_idx * shape_dim;
+                else if constexpr (kIndexType == KGroupedIndexType::K)
+                    offset = current_k_cumsum;
+                else if constexpr (kIndexType == KGroupedIndexType::SF_K)
+                    offset = current_sf_k_cumsum;
+            }
+            return offset + block_idx * block_size;
+        }
+    }
+
+    __device__ __forceinline__ bool get_next_block(uint32_t& m_block_idx, uint32_t& n_block_idx) {
+        const auto next_block_idx = (++ current_iter) * kNumSMs + blockIdx.x;
+
+        if constexpr (kGemmType == GemmType::MGroupedMasked) {
+            while (true) {
+                // End of the task
+                if (current_group_idx == kNumGroups)
+                    return false;
+
+                // Within current group
+                num_m_blocks = ceil_div(static_cast<uint32_t>(__ldg(grouped_layout + current_group_idx)), BLOCK_M);
+                const auto current_m_block_cumsum = current_m_cumsum + num_m_blocks;
+                if (next_block_idx < current_m_block_cumsum * num_n_blocks)
+                    break;
+
+                // Move to check the next group
+                current_group_idx ++, current_m_cumsum = current_m_block_cumsum;
+            }
+
+            get_swizzled_block_idx(next_block_idx - current_m_cumsum * num_n_blocks, m_block_idx, n_block_idx);
+        } else if (kGemmType == GemmType::KGroupedContiguous) {
+            while (true) {
+                // End of the task
+                if (current_group_idx == kNumGroups)
+                    return false;
+
+                // Within current group
+                if (next_block_idx < (current_num_valid_groups + 1) * num_m_blocks * num_n_blocks)
+                    break;
+
+                // Move to check the next group
+                current_k_cumsum += current_shape_k;
+                current_sf_k_cumsum += ceil_div(current_shape_k, SF_K_ALIGNMENT);
+                current_num_valid_groups ++;
+
+                current_group_idx = next_group_idx ++;
+                current_shape_k = next_shape_k;
+                get_next_k_group(next_group_idx, next_shape_k);
+            }
+
+            get_swizzled_block_idx(next_block_idx - current_num_valid_groups * num_m_blocks * num_n_blocks, m_block_idx, n_block_idx);
+        } else {
+            if (next_block_idx >= num_blocks)
+                return false;
+
+            // For SM90 only
+            // NOTES: we don't have to set `is_peer_cta_alive` for masked grouped GEMM, as it must be aligned
+            is_peer_cta_alive = num_n_blocks % kNumMulticast == 0 or                  // Always aligned on N (constant bypass)
+                                num_m_blocks % kNumMulticast == 0 or                  // Always aligned on M (constant bypass)
+                                (next_block_idx ^ 1) < num_blocks;                    // Peer CTA in bound
+            get_swizzled_block_idx(next_block_idx, m_block_idx, n_block_idx);
+        }
+        return true;
+    }
+
+    // For SM90 only
+    __device__ __forceinline__ bool is_tma_multicast_valid(const uint32_t& m_block_idx) const {
+        if (num_blocks_in_group == 1)
+            return false;
+        if constexpr (kGemmType == GemmType::Normal or kGemmType == GemmType::MGroupedMasked or kGemmType == GemmType::KGroupedContiguous) {
+            return true;
+        } else {
+            DG_STATIC_ASSERT(kGemmType == GemmType::MGroupedContiguous, "Invalid Gemm type");
+            if constexpr (kIsMulticastOnA) {
+                return true;
+            } else {
+                const auto& group_idx = __ldg(grouped_layout + m_block_idx * BLOCK_M);
+                const auto& peer_group_idx = __ldg(grouped_layout + (m_block_idx ^ 1) * BLOCK_M);
+                return group_idx == peer_group_idx;
+            }
+        }
+    }
+
+    // For SM90 only
+    // ReSharper disable once CppNotAllPathsReturnValue
+    __device__ __forceinline__ bool is_computation_valid(const uint32_t& m_block_idx, const uint32_t& m_offset) const {
+        if constexpr (kGemmType == GemmType::Normal) {
+            return true;
+        } else if constexpr (kGemmType == GemmType::MGroupedContiguous) {
+            return __ldg(grouped_layout + m_offset + m_block_idx * BLOCK_M) >= 0;
+        } else if constexpr (kGemmType == GemmType::MGroupedMasked) {
+            return m_offset + m_block_idx * BLOCK_M < __ldg(grouped_layout + current_group_idx);
+        }
+    }
+};
+
+#pragma clang diagnostic pop
+
+} // namespace deep_gemm

packages/ltx-kernels/csrc/blockwise/kernels/deep_gemm/include/deep_gemm/common/sm100_utils.cuh

@@ -0,0 +1,260 @@
+#pragma once
+
+#include <cute/atom/mma_traits_sm100.hpp>
+#include <cute/arch/mma_sm100_umma.hpp>
+#include <cute/arch/tmem_allocator_sm100.hpp>
+
+#include <deep_gemm/common/utils.cuh>
+
+namespace deep_gemm::sm100 {
+
+template <uint32_t BLOCK_INNER, uint32_t kSwizzleMode, typename dtype_t>
+constexpr uint32_t get_inner_block_atom_size() {
+    return kSwizzleMode == 0 ? BLOCK_INNER : kSwizzleMode / sizeof(dtype_t);
+}
+
+template <uint32_t BLOCK_INNER, uint32_t BLOCK_OUTER,
+          uint32_t kSwizzleMode, uint32_t kNumMulticast,
+          typename dtype_t>
+__device__ __forceinline__ void
+tma_copy(void const* desc_ptr, cutlass::arch::ClusterTransactionBarrier* barrier_ptr,
+         dtype_t* smem_ptr, const uint32_t& inner_idx, const int32_t& outer_idx) {
+    DG_STATIC_ASSERT(1 <= kNumMulticast and kNumMulticast <= 2, "Invalid multicast config");
+    DG_STATIC_ASSERT(static_cast<uint64_t>(cute::TMA::CacheHintSm90::EVICT_NORMAL) ==
+                     static_cast<uint64_t>(cute::TMA::CacheHintSm100::EVICT_NORMAL), "Invalid cache hint");
+
+    // 2-CTA function will send signals to the leader CTA only
+    const auto copy_func = kNumMulticast == 1 ? cute::SM90_TMA_LOAD_2D::copy : cute::SM100_TMA_2SM_LOAD_2D::copy;
+
+    // Issue multiple TMAs
+    constexpr uint32_t BLOCK_INNER_ATOM = get_inner_block_atom_size<BLOCK_INNER, kSwizzleMode, dtype_t>();
+    #pragma unroll
+    for (uint32_t i = 0; i < BLOCK_INNER / BLOCK_INNER_ATOM; ++ i) {
+        copy_func(desc_ptr, reinterpret_cast<uint64_t*>(barrier_ptr),
+                  static_cast<uint64_t>(cute::TMA::CacheHintSm100::EVICT_NORMAL),
+                  smem_ptr + i * BLOCK_OUTER * BLOCK_INNER_ATOM, inner_idx + i * BLOCK_INNER_ATOM, outer_idx);
+    }
+}
+
+__device__ __forceinline__
+cute::UMMA::SmemDescriptor make_smem_desc(cute::UMMA::LayoutType layout, void* smem_ptr,
+                                          uint32_t stride_byte_offset, uint32_t leading_byte_offset) {
+    cute::UMMA::SmemDescriptor desc;
+
+    // Set the version for SM100
+    desc.version_ = 1;
+
+    // Legacy mode
+    desc.lbo_mode_ = 0;
+
+    // Layout
+    desc.layout_type_ = static_cast<uint8_t>(layout);
+
+    // Start address
+    const auto uint_ptr = cute::cast_smem_ptr_to_uint(smem_ptr);
+    desc.start_address_ = static_cast<uint16_t>(uint_ptr >> 4);
+
+    // Base offset
+    desc.base_offset_ = 0;
+
+    // SBO and LBO
+    desc.stride_byte_offset_ = stride_byte_offset >> 4;
+    desc.leading_byte_offset_ = leading_byte_offset >> 4;
+
+    return desc;
+}
+
+__device__ __forceinline__
+cute::UMMA::SmemDescriptor make_sf_desc(void* smem_ptr) {
+    // NOTES: the UTCCP layout is K-major by default
+    // Atom size: 8 x 128 bits
+    // {SBO, LBO} means the byte stride between atoms on {MN, K}
+    // Since the UTCCP we used is 128b-wide (only 1 atom on K), so LBO can be zero
+    return make_smem_desc(cute::UMMA::LayoutType::SWIZZLE_NONE, smem_ptr, 8 * 16, 0);
+}
+
+__device__ __forceinline__
+void replace_smem_desc_addr(cute::UMMA::SmemDescriptor& desc, const void* smem_ptr) {
+    const auto uint_ptr = cute::cast_smem_ptr_to_uint(smem_ptr);
+    desc.start_address_ = static_cast<uint16_t>(uint_ptr >> 4);
+}
+
+__device__ __forceinline__
+static uint32_t get_atom_base(const cute::UMMA::LayoutType& layout_type) {
+    return layout_type == cute::UMMA::LayoutType::SWIZZLE_128B_BASE32B ? 32 : 16;
+}
+
+// ReSharper disable once CppNotAllPathsReturnValue
+template <cute::UMMA::Major kMajorMode, uint32_t kSwizzleMode, bool kUseBase32, typename dtype_t>
+constexpr static cute::UMMA::LayoutType to_umma_layout_type() {
+    DG_STATIC_ASSERT(kSwizzleMode == 0 or kSwizzleMode == 16 or
+                     kSwizzleMode == 32 or kSwizzleMode == 64 or
+                     kSwizzleMode == 128, "Invalid swizzling mode");
+    // A special case
+    if constexpr ((cute::is_same_v<dtype_t, float> and kMajorMode == cute::UMMA::Major::MN) or kUseBase32) {
+        DG_STATIC_ASSERT(kUseBase32, "Invalid swizzling base");
+        return cute::UMMA::LayoutType::SWIZZLE_128B_BASE32B;
+    }
+
+    // Normal cases
+    if constexpr (kSwizzleMode == 0)   return cute::UMMA::LayoutType::SWIZZLE_NONE;
+    if constexpr (kSwizzleMode == 16)  return cute::UMMA::LayoutType::SWIZZLE_NONE;
+    if constexpr (kSwizzleMode == 32)  return cute::UMMA::LayoutType::SWIZZLE_32B;
+    if constexpr (kSwizzleMode == 64)  return cute::UMMA::LayoutType::SWIZZLE_64B;
+    if constexpr (kSwizzleMode == 128) return cute::UMMA::LayoutType::SWIZZLE_128B;
+}
+
+template <cute::UMMA::Major kMajorMode, uint32_t BLOCK_MN, uint32_t kSwizzleMode, typename dtype_t>
+__device__ __forceinline__
+constexpr uint32_t get_umma_desc_stride_k() {
+    return kMajorMode == cute::UMMA::Major::K ? 1 : get_inner_block_atom_size<BLOCK_MN, kSwizzleMode, dtype_t>();
+}
+
+template <cute::UMMA::Major kMajorMode, uint32_t BLOCK_MN, uint32_t kSwizzleMode, typename dtype_t>
+__device__ __forceinline__
+uint32_t advance_umma_desc_lo(const uint32_t& base, const uint32_t& offset, const uint32_t& k_idx) {
+    return base + (((offset + k_idx * get_umma_desc_stride_k<kMajorMode, BLOCK_MN, kSwizzleMode, dtype_t>()) * static_cast<uint32_t>(sizeof(dtype_t))) >> 4u);
+}
+
+template <cute::UMMA::Major kMajorMode, uint32_t BLOCK_MN, uint32_t BLOCK_K, uint32_t kSwizzleMode, bool kUseBase32 = false, typename dtype_t>
+__device__ __forceinline__
+cute::UMMA::SmemDescriptor make_umma_desc(dtype_t* base_smem_ptr, uint32_t mn_idx, uint32_t k_idx) {
+    const uint32_t stride_k = get_umma_desc_stride_k<kMajorMode, BLOCK_MN, kSwizzleMode, dtype_t>();
+    const auto& layout_type = to_umma_layout_type<kMajorMode, kSwizzleMode, kUseBase32, dtype_t>();
+    const auto& num_non_contiguous = 128 / get_atom_base(layout_type);
+    if constexpr (kMajorMode == cute::UMMA::Major::K) {
+        // NOTES: for K-major layout, the swizzle must be 128B (also, atom index must be 0), as `BLOCK_K` is always 128
+        DG_STATIC_ASSERT(kSwizzleMode == BLOCK_K * sizeof(dtype_t), "Unexpected value");
+
+        // Atom size: 8 x `kSwizzleMode` (in bytes, on K)
+        // {SBO, LBO} means the byte stride between atoms on {MN, K}
+        // NOTES: on K, there is only 1 atom as asserted previously, so LBO can be 0
+        const uint32_t stride_byte_offset = num_non_contiguous * BLOCK_K * sizeof(dtype_t);
+        const uint32_t leading_byte_offset = 0;
+        return make_smem_desc(layout_type,
+                              base_smem_ptr + mn_idx * BLOCK_K + k_idx * stride_k,
+                              stride_byte_offset, leading_byte_offset);
+    } else {
+        constexpr uint32_t BLOCK_MN_ATOM = get_inner_block_atom_size<BLOCK_MN, kSwizzleMode, dtype_t>();
+
+        // Must have no in-atom MN-idx
+        // NOTES: no worries for the runtime assert, the `mn_idx` are constants at compilation time
+        DG_DEVICE_ASSERT(mn_idx % BLOCK_MN_ATOM == 0);
+        DG_STATIC_ASSERT(kSwizzleMode > 0, "Invalid swizzling");
+
+        // Atom size: `kSwizzleMode` (in bytes, on MN) x 8
+        // NOTES: `kSwizzleMode == 16` mean non-swizzling but interleaving
+        // {SBO, LBO} means the byte stride between atoms on {K, MN} for swizzling
+        // {SBO, LBO} means the byte stride between atoms on {MN, K} for non-swizzling
+        uint32_t stride_byte_offset = num_non_contiguous * BLOCK_MN_ATOM * sizeof(dtype_t);
+        uint32_t leading_byte_offset = BLOCK_K * BLOCK_MN_ATOM * sizeof(dtype_t);
+        if constexpr (kSwizzleMode == 16)
+            swap(stride_byte_offset, leading_byte_offset);
+        return make_smem_desc(layout_type,
+                              base_smem_ptr + mn_idx * BLOCK_K + k_idx * stride_k,
+                              stride_byte_offset, leading_byte_offset);
+    }
+}
+
+__device__  __forceinline__
+uint64_t make_runtime_instr_desc_with_sf_id(cute::UMMA::InstrDescriptorBlockScaled desc, const uint32_t& sf_id) {
+    desc.a_sf_id_ = sf_id, desc.b_sf_id_ = sf_id;
+    return static_cast<uint64_t>(static_cast<uint32_t>(desc)) << 32;
+}
+
+template <uint32_t kNumCols>
+__device__ constexpr uint32_t get_num_aligned_tmem_cols() {
+    DG_STATIC_ASSERT(kNumCols <= 512, "Too many tensor memory columns");
+    if (kNumCols <=  32) return  32;
+    if (kNumCols <=  64) return  64;
+    if (kNumCols <= 128) return 128;
+    if (kNumCols <= 256) return 256;
+    return 512;
+}
+
+__device__ __forceinline__ void tcgen05_before_thread_sync() {
+    asm volatile("tcgen05.fence::before_thread_sync;");
+}
+
+__device__ __forceinline__ void tcgen05_after_thread_sync() {
+    asm volatile("tcgen05.fence::after_thread_sync;");
+}
+
+// UMMA versions with relaxed assertions
+struct SM100_MMA_F16BF16_SS {
+    __device__ static void
+    fma(uint64_t const& desc_a,
+        uint64_t const& desc_b,
+        uint32_t const& tmem_c,
+        uint32_t const& scale_c,
+        uint64_t const& desc) {
+        asm volatile(
+            "{\n\t"
+            ".reg .pred p;\n\t"
+            "setp.ne.b32 p, %4, 0;\n\t"
+            "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p; \n\t"
+            "}\n"
+            :: "r"(tmem_c), "l"(desc_a), "l"(desc_b), "r"(static_cast<uint32_t>(desc >> 32)), "r"(scale_c));
+    }
+};
+
+struct SM100_MMA_F16BF16_2x1SM_SS {
+    __device__ static void
+    fma(uint64_t const& desc_a,
+        uint64_t const& desc_b,
+        uint32_t const& tmem_c,
+        uint32_t const& scale_c,
+        uint64_t const& desc) {
+        asm volatile(
+            "{\n\t"
+            ".reg .pred p;\n\t"
+            "setp.ne.b32 p, %4, 0;\n\t"
+            "tcgen05.mma.cta_group::2.kind::f16 [%0], %1, %2, %3, p; \n\t"
+            "}\n"
+            :: "r"(tmem_c), "l"(desc_a), "l"(desc_b), "r"(static_cast<uint32_t>(desc >> 32)), "r"(scale_c));
+    }
+};
+
+struct SM100_MMA_MXF8F6F4_SS {
+    __device__ static void
+    fma(uint64_t const& desc_a,
+        uint64_t const& desc_b,
+        uint32_t const& tmem_c,
+        uint32_t const& scale_c,
+        uint64_t const& desc,
+        uint32_t const& tmem_sfa,
+        uint32_t const& tmem_sfb) {
+        asm volatile(
+          "{\n\t"
+          ".reg .pred p;\n\t"
+          "setp.ne.b32 p, %4, 0;\n\t"
+          "tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale [%0], %1, %2, %3, [%5], [%6], p; \n\t"
+          "}\n"
+          :
+          : "r"(tmem_c), "l"(desc_a), "l"(desc_b), "r"(static_cast<uint32_t>(desc >> 32)), "r"(scale_c),
+            "r"(tmem_sfa), "r"(tmem_sfb));
+    }
+};
+
+struct SM100_MMA_MXF8F6F4_2x1SM_SS {
+    __device__ static void
+    fma(uint64_t const& desc_a,
+        uint64_t const& desc_b,
+        uint32_t const& tmem_c,
+        uint32_t const& scale_c,
+        uint64_t const& desc,
+        uint32_t const& tmem_sfa,
+        uint32_t const& tmem_sfb) {
+        asm volatile(
+          "{\n\t"
+          ".reg .pred p;\n\t"
+          "setp.ne.b32 p, %4, 0;\n\t"
+          "tcgen05.mma.cta_group::2.kind::mxf8f6f4.block_scale [%0], %1, %2, %3, [%5], [%6], p; \n\t"
+          "}\n"
+          :
+          : "r"(tmem_c), "l"(desc_a), "l"(desc_b), "r"(static_cast<uint32_t>(desc >> 32)), "r"(scale_c),
+            "r"(tmem_sfa), "r"(tmem_sfb));
+    }
+};
+
+} // namespace `deep_gemm::sm100`

packages/ltx-kernels/csrc/blockwise/kernels/deep_gemm/include/deep_gemm/common/sm90_utils.cuh

@@ -0,0 +1,283 @@
+#pragma once
+
+#include <cute/arch/copy_sm90_tma.hpp>
+#include <cute/arch/cluster_sm90.hpp>
+#include <cute/arch/mma_sm90_gmma.hpp>
+#include <cute/arch/mma_sm90_gmma_ext.hpp>
+
+#include <deep_gemm/common/utils.cuh>
+
+namespace deep_gemm::sm90 {
+
+template <int N_, typename MMA>
+struct FP8MMA {
+
+    template <size_t ...Idx>
+    __forceinline__ __device__ static void call_fma_impl(uint64_t const& desc_a, uint64_t const& desc_b, float* d, bool scale_d, cute::index_sequence<Idx...>) {
+        using namespace cute::SM90::GMMA;
+        MMA::fma(desc_a, desc_b, d[Idx]..., (scale_d ? ScaleOut::One : ScaleOut::Zero));
+    }
+
+    __forceinline__ __device__ static void wgmma(uint64_t const& desc_a, uint64_t const& desc_b, float* d, bool scale_d) {
+        call_fma_impl(desc_a, desc_b, d, scale_d, cute::make_index_sequence<N_/2>{});
+    }
+
+    static constexpr int M = 64;
+    static constexpr int N = N_;
+    static constexpr int K = 32;
+    static constexpr int kNumAccum = M * N / 128;
+};
+
+template <int N>
+struct FP8MMASelector {
+
+    static constexpr auto select_mma() {
+        using namespace cute::SM90::GMMA;
+        if constexpr (N == 8) return MMA_64x8x32_F32E4M3E4M3_SS_TN();
+        if constexpr (N == 16) return MMA_64x16x32_F32E4M3E4M3_SS_TN();
+        if constexpr (N == 24) return MMA_64x24x32_F32E4M3E4M3_SS_TN();
+        if constexpr (N == 32) return MMA_64x32x32_F32E4M3E4M3_SS_TN();
+        if constexpr (N == 40) return MMA_64x40x32_F32E4M3E4M3_SS_TN();
+        if constexpr (N == 48) return MMA_64x48x32_F32E4M3E4M3_SS_TN();
+        if constexpr (N == 56) return MMA_64x56x32_F32E4M3E4M3_SS_TN();
+        if constexpr (N == 64) return MMA_64x64x32_F32E4M3E4M3_SS_TN();
+        if constexpr (N == 72) return MMA_64x72x32_F32E4M3E4M3_SS_TN();
+        if constexpr (N == 80) return MMA_64x80x32_F32E4M3E4M3_SS_TN();
+        if constexpr (N == 88) return MMA_64x88x32_F32E4M3E4M3_SS_TN();
+        if constexpr (N == 96) return MMA_64x96x32_F32E4M3E4M3_SS_TN();
+        if constexpr (N == 104) return MMA_64x104x32_F32E4M3E4M3_SS_TN();
+        if constexpr (N == 112) return MMA_64x112x32_F32E4M3E4M3_SS_TN();
+        if constexpr (N == 120) return MMA_64x120x32_F32E4M3E4M3_SS_TN();
+        if constexpr (N == 128) return MMA_64x128x32_F32E4M3E4M3_SS_TN();
+        if constexpr (N == 136) return MMA_64x136x32_F32E4M3E4M3_SS_TN();
+        if constexpr (N == 144) return MMA_64x144x32_F32E4M3E4M3_SS_TN();
+        if constexpr (N == 152) return MMA_64x152x32_F32E4M3E4M3_SS_TN();
+        if constexpr (N == 160) return MMA_64x160x32_F32E4M3E4M3_SS_TN();
+        if constexpr (N == 168) return MMA_64x168x32_F32E4M3E4M3_SS_TN();
+        if constexpr (N == 176) return MMA_64x176x32_F32E4M3E4M3_SS_TN();
+        if constexpr (N == 184) return MMA_64x184x32_F32E4M3E4M3_SS_TN();
+        if constexpr (N == 192) return MMA_64x192x32_F32E4M3E4M3_SS_TN();
+        if constexpr (N == 200) return MMA_64x200x32_F32E4M3E4M3_SS_TN();
+        if constexpr (N == 208) return MMA_64x208x32_F32E4M3E4M3_SS_TN();
+        if constexpr (N == 216) return MMA_64x216x32_F32E4M3E4M3_SS_TN();
+        if constexpr (N == 224) return MMA_64x224x32_F32E4M3E4M3_SS_TN();
+        if constexpr (N == 232) return MMA_64x232x32_F32E4M3E4M3_SS_TN();
+        if constexpr (N == 240) return MMA_64x240x32_F32E4M3E4M3_SS_TN();
+        if constexpr (N == 248) return MMA_64x248x32_F32E4M3E4M3_SS_TN();
+        if constexpr (N == 256) return MMA_64x256x32_F32E4M3E4M3_SS_TN();
+    }
+
+    static constexpr auto select_type() {
+        return FP8MMA<N, decltype(select_mma())>();
+    }
+
+    using type = decltype(select_type());
+};
+
+template <int N_, typename MMA>
+struct BF16MMA {
+
+    template <size_t ...Idx>
+    __forceinline__ __device__ static void call_fma_impl(uint64_t const& desc_a, uint64_t const& desc_b, float* d, bool scale_d, cute::index_sequence<Idx...>) {
+        using namespace cute::SM90::GMMA;
+        MMA::fma(desc_a, desc_b, d[Idx]..., (scale_d ? ScaleOut::One : ScaleOut::Zero));
+    }
+
+    __forceinline__ __device__ static void wgmma(uint64_t const& desc_a, uint64_t const& desc_b, float* d, bool scale_d) {
+        call_fma_impl(desc_a, desc_b, d, scale_d, cute::make_index_sequence<N_/2>{});
+    }
+
+    static constexpr int M = 64;
+    static constexpr int N = N_;
+    static constexpr int K = 16;
+    static constexpr int kNumAccum = M * N / 128;
+};
+
+template <int N>
+struct BF16MMASelector {
+
+    static constexpr auto select_mma() {
+        using namespace cute::SM90::GMMA;
+        if constexpr (N == 8) return MMA_64x8x16_F32BF16BF16_SS<Major::K, Major::K>();
+        if constexpr (N == 16) return MMA_64x16x16_F32BF16BF16_SS<Major::K, Major::K>();
+        if constexpr (N == 24) return MMA_64x24x16_F32BF16BF16_SS<Major::K, Major::K>();
+        if constexpr (N == 32) return MMA_64x32x16_F32BF16BF16_SS<Major::K, Major::K>();
+        if constexpr (N == 40) return MMA_64x40x16_F32BF16BF16_SS<Major::K, Major::K>();
+        if constexpr (N == 48) return MMA_64x48x16_F32BF16BF16_SS<Major::K, Major::K>();
+        if constexpr (N == 56) return MMA_64x56x16_F32BF16BF16_SS<Major::K, Major::K>();
+        if constexpr (N == 64) return MMA_64x64x16_F32BF16BF16_SS<Major::K, Major::K>();
+        if constexpr (N == 72) return MMA_64x72x16_F32BF16BF16_SS<Major::K, Major::K>();
+        if constexpr (N == 80) return MMA_64x80x16_F32BF16BF16_SS<Major::K, Major::K>();
+        if constexpr (N == 88) return MMA_64x88x16_F32BF16BF16_SS<Major::K, Major::K>();
+        if constexpr (N == 96) return MMA_64x96x16_F32BF16BF16_SS<Major::K, Major::K>();
+        if constexpr (N == 104) return MMA_64x104x16_F32BF16BF16_SS<Major::K, Major::K>();
+        if constexpr (N == 112) return MMA_64x112x16_F32BF16BF16_SS<Major::K, Major::K>();
+        if constexpr (N == 120) return MMA_64x120x16_F32BF16BF16_SS<Major::K, Major::K>();
+        if constexpr (N == 128) return MMA_64x128x16_F32BF16BF16_SS<Major::K, Major::K>();
+        if constexpr (N == 136) return MMA_64x136x16_F32BF16BF16_SS<Major::K, Major::K>();
+        if constexpr (N == 144) return MMA_64x144x16_F32BF16BF16_SS<Major::K, Major::K>();
+        if constexpr (N == 152) return MMA_64x152x16_F32BF16BF16_SS<Major::K, Major::K>();
+        if constexpr (N == 160) return MMA_64x160x16_F32BF16BF16_SS<Major::K, Major::K>();
+        if constexpr (N == 168) return MMA_64x168x16_F32BF16BF16_SS<Major::K, Major::K>();
+        if constexpr (N == 176) return MMA_64x176x16_F32BF16BF16_SS<Major::K, Major::K>();
+        if constexpr (N == 184) return MMA_64x184x16_F32BF16BF16_SS<Major::K, Major::K>();
+        if constexpr (N == 192) return MMA_64x192x16_F32BF16BF16_SS<Major::K, Major::K>();
+        if constexpr (N == 200) return MMA_64x200x16_F32BF16BF16_SS<Major::K, Major::K>();
+        if constexpr (N == 208) return MMA_64x208x16_F32BF16BF16_SS<Major::K, Major::K>();
+        if constexpr (N == 216) return MMA_64x216x16_F32BF16BF16_SS<Major::K, Major::K>();
+        if constexpr (N == 224) return MMA_64x224x16_F32BF16BF16_SS<Major::K, Major::K>();
+        if constexpr (N == 232) return MMA_64x232x16_F32BF16BF16_SS<Major::K, Major::K>();
+        if constexpr (N == 240) return MMA_64x240x16_F32BF16BF16_SS<Major::K, Major::K>();
+        if constexpr (N == 248) return MMA_64x248x16_F32BF16BF16_SS<Major::K, Major::K>();
+        if constexpr (N == 256) return MMA_64x256x16_F32BF16BF16_SS<Major::K, Major::K>();
+    }
+
+    static constexpr auto select_type() {
+        return BF16MMA<N, decltype(select_mma())>();
+    }
+
+    using type = decltype(select_type());
+};
+
+
+template <typename dtype_t>
+struct SM90_U32x2_STSM_N {
+    __device__ __forceinline__ static void
+    copy(dtype_t src_0, dtype_t src_1, void* smem_dst) {
+        const uint32_t src[2] = {*reinterpret_cast<uint32_t*>(&src_0), *reinterpret_cast<uint32_t*>(&src_1)};
+        asm volatile("stmatrix.sync.aligned.x2.m8n8.shared.b16 [%0], {%1, %2};\n"
+                     :: "l"(smem_dst), "r"(src[0]), "r"(src[1]));
+    }
+};
+
+struct SM90_U32x2_LDSM_N {
+    __device__ __forceinline__ static void
+    copy(uint32_t& dst_0, uint32_t& dst_1, void* smem_src) {
+        asm volatile("ldmatrix.sync.aligned.x2.m8n8.shared.b16 {%0, %1}, [%2];\n"
+                     : "=r"(dst_0), "=r"(dst_1)
+                     : "l"(smem_src));
+    }
+};
+
+struct SM90_U32x4_LDSM_N {
+    __device__ __forceinline__ static void
+    copy(uint32_t& dst_0, uint32_t& dst_1, uint32_t& dst_2, uint32_t& dst_3, void* smem_src) {
+        asm volatile("ldmatrix.sync.aligned.x4.m8n8.shared.b16 {%0, %1, %2, %3}, [%4];\n"
+                     : "=r"(dst_0), "=r"(dst_1), "=r"(dst_2), "=r"(dst_3)
+                     : "l"(smem_src));
+    }
+};
+
+__forceinline__ __device__ void warpgroup_arrive() {
+    asm volatile("wgmma.fence.sync.aligned;\n" ::: "memory");
+}
+
+__forceinline__ __device__ void warpgroup_commit_batch() {
+    asm volatile("wgmma.commit_group.sync.aligned;\n" ::: "memory");
+}
+
+__forceinline__ __device__ void warpgroup_fence_operand(float& reg) {
+    asm volatile("" : "+f"(reg) :: "memory");
+}
+
+template <int N>
+__forceinline__ __device__ void warpgroup_wait() {
+    DG_STATIC_ASSERT(N >= 0 and N <= 7, "WGMMA wait: N must be in range [0, 7]");
+    asm volatile("wgmma.wait_group.sync.aligned %0;\n" :: "n"(N) : "memory");
+}
+
+// TODO: replace with CUTLASS solution
+union GmmaDescriptor {
+    __host__ __device__ constexpr GmmaDescriptor() noexcept: desc_(0) {}
+
+    __host__ __device__ constexpr GmmaDescriptor(uint64_t desc) noexcept: desc_(desc) {}
+
+    __host__ __device__ constexpr GmmaDescriptor(GmmaDescriptor const &t) noexcept: desc_(t.desc_) {}
+
+    __host__ __device__ constexpr GmmaDescriptor(GmmaDescriptor &&t) noexcept: desc_(t.desc_) {}
+
+    __host__ __device__ constexpr GmmaDescriptor &operator=(GmmaDescriptor const &t) noexcept {
+        desc_ = t.desc_;
+        return *this;
+    }
+
+    __host__ __device__ constexpr GmmaDescriptor &operator=(GmmaDescriptor &&t) noexcept {
+        desc_ = t.desc_;
+        return *this;
+    }
+
+    uint64_t desc_;
+    uint32_t reg32_[2];
+    uint16_t reg16_[4];
+
+    struct {
+        uint16_t start_address_: 14, : 2;
+        uint16_t leading_byte_offset_: 14, : 2;
+        uint16_t stride_byte_offset_: 14, : 2;
+        uint8_t : 1, base_offset_: 3, : 4;
+        uint8_t : 6, layout_type_: 2;
+    } bitfield;
+
+    // Decay to an `uint64_t`
+    __host__ __device__ constexpr operator uint64_t() const noexcept { return desc_; }
+};
+
+template <class PointerType>
+__device__ GmmaDescriptor make_smem_desc(PointerType smem_ptr, const int& layout_type,
+                                         const int& leading_byte_offset = 0,
+                                         const int& stride_byte_offset = 1024) {
+    GmmaDescriptor desc;
+    const auto& uint_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
+    desc.bitfield.start_address_ = uint_ptr >> 4;
+    desc.bitfield.layout_type_ = layout_type;
+    desc.bitfield.leading_byte_offset_ = leading_byte_offset >> 4;
+    desc.bitfield.stride_byte_offset_ = stride_byte_offset >> 4;
+    desc.bitfield.base_offset_ = 0;
+    return desc;
+}
+
+__device__ __forceinline__ void
+tma_copy(void const* desc_ptr, uint64_t* barrier_ptr, void* smem_ptr,
+         const uint32_t& crd_0, const uint32_t& crd_1, const uint32_t& num_tma_multicast = 1) {
+    constexpr auto cache_hint = static_cast<uint64_t>(cute::TMA::CacheHintSm90::EVICT_NORMAL);
+    if (num_tma_multicast == 1) {
+        cute::SM90_TMA_LOAD_2D::copy(desc_ptr, barrier_ptr, cache_hint, smem_ptr, crd_0, crd_1);
+    } else if (cute::block_rank_in_cluster() == 0) {
+        cute::SM90_TMA_LOAD_MULTICAST_2D::copy(desc_ptr, barrier_ptr, (1 << num_tma_multicast) - 1, cache_hint, smem_ptr, crd_0, crd_1);
+    }
+}
+
+__device__ __forceinline__ void
+tma_3d_copy(void const* desc_ptr, uint64_t* barrier_ptr, void* smem_ptr,
+            const uint32_t& crd_0, const uint32_t& crd_1, const uint32_t& crd_2) {
+    constexpr auto cache_hint = static_cast<uint64_t>(cute::TMA::CacheHintSm90::EVICT_NORMAL);
+    cute::SM90_TMA_LOAD_3D::copy(desc_ptr, barrier_ptr, cache_hint, smem_ptr, crd_0, crd_1, crd_2);
+}
+
+// Tensormap related
+__device__ __forceinline__ void tensor_map_release_cta() {
+    asm volatile ("fence.proxy.tensormap::generic.release.cta;");
+}
+
+__device__ __forceinline__ void tensor_map_acquire_cta(const cute::TmaDescriptor* gmem_desc_ptr) {
+    auto gmem_int_desc = reinterpret_cast<uint64_t>(gmem_desc_ptr);
+    asm volatile ("fence.proxy.tensormap::generic.acquire.cta [%0], 128;" :: "l"(gmem_int_desc) : "memory");
+}
+
+__device__ __forceinline__ void tensor_map_replace_global_addr_in_smem(cute::TmaDescriptor* smem_desc, const void* new_addr) {
+    auto smem_int_desc = static_cast<uint32_t>(__cvta_generic_to_shared(smem_desc));
+    const auto new_int64_addr = reinterpret_cast<uint64_t>(new_addr);
+    asm volatile ("tensormap.replace.tile.global_address.shared::cta.b1024.b64 [%0], %1;" :: "r"(smem_int_desc), "l"(new_int64_addr));
+}
+
+__device__ __forceinline__ void tensor_map_replace_global_inner_dim_stride_in_smem(cute::TmaDescriptor* smem_desc, const uint32_t& new_dim, const uint64_t& new_stride) {
+    auto smem_int_desc = __cvta_generic_to_shared(smem_desc);
+    asm volatile ("tensormap.replace.tile.global_dim.shared::cta.b1024.b32 [%0], 0, %1;" :: "l"(smem_int_desc), "r"(new_dim));
+#if ((__CUDACC_VER_MAJOR__ > 12) or ((__CUDACC_VER_MAJOR__ == 12) and (__CUDACC_VER_MINOR__ >= 3)))
+    asm volatile("tensormap.replace.tile.global_stride.shared::cta.b1024.b64 [%0], 0, %1;" :: "l"(smem_int_desc), "l"(new_stride));
+#else
+    DG_STATIC_ASSERT(false, "Invalid CUDA version");
+#endif
+}
+
+} // namespace `deep_gemm::sm90`

packages/ltx-kernels/csrc/blockwise/kernels/deep_gemm/include/deep_gemm/common/types.hpp

@@ -0,0 +1,18 @@
+#pragma once
+
+namespace deep_gemm {
+
+enum class GemmType {
+    Normal              = 0,
+    MGroupedContiguous  = 1,
+    MGroupedMasked      = 2,
+    KGroupedContiguous  = 3,
+};
+
+enum class KernelType {
+    Kernel1D1D = 0,
+    Kernel1D2D = 1,
+    KernelNoSF = 2
+};
+
+} // namespace deep_gemm

packages/ltx-kernels/csrc/blockwise/kernels/deep_gemm/include/deep_gemm/common/utils.cuh

@@ -0,0 +1,179 @@
+#pragma once
+
+#include <cuda_bf16.h>
+#include <cuda_fp8.h>
+#include <cuda/std/cstdint>
+#include <cuda/std/utility>
+#include <cute/container/tuple.hpp>
+
+#include "cute_tie.cuh"
+
+#ifdef __CLION_IDE__
+
+__host__ __device__ __forceinline__ void host_device_printf(const char* format, ...) {
+    asm volatile("trap;");
+}
+
+#define printf host_device_printf
+#endif
+
+#ifndef DG_DEVICE_ASSERT
+#define DG_DEVICE_ASSERT(cond) \
+do { \
+    if (not (cond)) { \
+        printf("Assertion failed: %s:%d, condition: %s\n", __FILE__, __LINE__, #cond); \
+        asm("trap;"); \
+    } \
+} while (0)
+#endif
+
+#ifndef DG_TRAP_ONLY_DEVICE_ASSERT
+#define DG_TRAP_ONLY_DEVICE_ASSERT(cond) \
+do { \
+    if (not (cond)) \
+        asm("trap;"); \
+} while (0)
+#endif
+
+#ifndef DG_STATIC_ASSERT
+#define DG_STATIC_ASSERT(cond, ...) static_assert(cond, __VA_ARGS__)
+#endif
+
+namespace deep_gemm {
+
+template <typename FuncT>
+struct PatternVisitor {
+    FuncT func;
+
+    __device__ __host__
+    explicit PatternVisitor(FuncT&& func): func(std::forward<FuncT>(func)) {}
+
+    __device__ __host__
+    auto operator [](const uint32_t& i) {
+        return func(i);
+    }
+};
+
+template <typename T>
+__device__ __host__ T ceil_div(T a, T b) {
+    return (a + b - 1) / b;
+}
+
+template <typename T>
+__device__ __host__ constexpr T constexpr_ceil_div(T a, T b) {
+    return (a + b - 1) / b;
+}
+
+template <typename T>
+__device__ __host__ T align(T a, T b) {
+    return ceil_div(a, b) * b;
+}
+
+template <typename T>
+__device__ __host__ constexpr T constexpr_align(T a, T b) {
+    return constexpr_ceil_div(a, b) * b;
+}
+
+template <typename T>
+__device__ __host__ constexpr T constexpr_gcd(T a, T b) {
+    return b == 0 ? a : constexpr_gcd(b, a % b);
+}
+
+template<typename T>
+__forceinline__ __device__ void swap(T& a, T& b) {
+    T temp = a;
+    a = b;
+    b = temp;
+}
+
+__forceinline__ __device__ uint32_t get_sm_idx() {
+    uint32_t sm_idx;
+    asm ("mov.u32 %0, %%smid;" : "=r"(sm_idx));
+    return sm_idx;
+}
+
+__forceinline__ __device__ uint32_t get_lane_idx() {
+    uint32_t lane_id;
+    asm ("mov.u32 %0, %laneid;" : "=r"(lane_id));
+    return lane_id;
+}
+
+__device__  __forceinline__ uint32_t ld_shared(const uint32_t* ptr) {
+    uint32_t ret;
+    asm volatile("ld.shared.u32 %0, [%1];" : "=r"(ret) : "l"(ptr));
+    return ret;
+}
+
+__device__  __forceinline__ float2 ld_shared(const float2* ptr) {
+    float2 ret;
+    asm volatile("ld.shared.v2.f32 {%0, %1}, [%2];" : "=f"(ret.x), "=f"(ret.y) : "l"(ptr));
+    return ret;
+}
+
+__device__  __forceinline__ float4 ld_shared(const float4* ptr) {
+    float4 ret;
+    asm volatile("ld.shared.v4.f32 {%0, %1, %2, %3}, [%4];" : "=f"(ret.x), "=f"(ret.y), "=f"(ret.z), "=f"(ret.w) : "l"(ptr));
+    return ret;
+}
+
+__device__  __forceinline__ uint4 ld_shared(const uint4* ptr) {
+    uint4 ret;
+    asm volatile("ld.shared.v4.u32 {%0, %1, %2, %3}, [%4];" : "=r"(ret.x), "=r"(ret.y), "=r"(ret.z), "=r"(ret.w) : "l"(ptr));
+    return ret;
+}
+
+__device__  __forceinline__ float ld_shared(const float* ptr) {
+    float ret;
+    asm volatile("ld.shared.f32 %0, [%1];" : "=f"(ret) : "l"(ptr));
+    return ret;
+}
+
+__device__ __forceinline__ void st_shared(const float* ptr, float val) {
+    asm volatile("st.shared.f32 [%0], %1;" :: "l"(ptr), "f"(val));
+}
+
+__device__ __forceinline__ void st_shared(const float2* ptr, float2 val) {
+    asm volatile("st.shared.v2.f32 [%0], {%1, %2};" :: "l"(ptr), "f"(val.x), "f"(val.y));
+}
+
+__device__ __forceinline__ void st_shared(const uint32_t* ptr, uint32_t val) {
+    asm volatile("st.shared.u32 [%0], %1;" :: "l"(ptr), "r"(val));
+}
+
+__device__  __forceinline__ void st_shared(const void* ptr, uint32_t x, uint32_t y) {
+    asm volatile("st.shared.v2.u32 [%0], {%1, %2};" :: "l"(ptr), "r"(x), "r"(y));
+}
+
+__device__  __forceinline__ void st_shared(const void* ptr, uint32_t x, uint32_t y, uint32_t z, uint32_t w) {
+    asm volatile("st.shared.v4.u32 [%0], {%1, %2, %3, %4};" :: "l"(ptr), "r"(x), "r"(y), "r"(z), "r"(w));
+}
+
+template <typename old_t>
+__device__ __forceinline__ int cast_into_bf16_and_pack(old_t& x, old_t& y) {
+    auto bf16x2 = __float22bfloat162_rn({*reinterpret_cast<float*>(&x), *reinterpret_cast<float*>(&y)});
+    return *reinterpret_cast<int*>(&bf16x2);
+}
+
+__device__ __forceinline__ void prefetch_l1(void *ptr) {
+    asm volatile("prefetch.global.L1 [%0];" :: "l"(ptr));
+}
+
+template <uint32_t kNumBytes>
+struct Vectorized {
+    static auto zeros() {
+        // TODO: add `ulonglong4` for SM100 once `__ldg` support this
+        if constexpr (kNumBytes > 0 and kNumBytes % 16 == 0) {
+            return make_uint4(0, 0, 0, 0);
+        } else if constexpr (kNumBytes > 0 and kNumBytes % 8 == 0) {
+            return make_uint2(0, 0);
+        } else if constexpr (kNumBytes > 0 and kNumBytes % 4 == 0) {
+            return 0;
+        } else {
+            DG_STATIC_ASSERT(kNumBytes > 0 and kNumBytes % 4 == 0, "Invalid vectorization");
+        }
+    }
+
+    using vec_t = decltype(zeros());
+};
+
+} // namespace `deep_gemm`

packages/ltx-kernels/csrc/blockwise/kernels/deep_gemm/include/deep_gemm/impls/sm90_fp8_gemm_1d2d.cuh

@@ -0,0 +1,408 @@
+#pragma once
+
+#pragma clang diagnostic push
+#pragma clang diagnostic ignored "-Wunknown-attributes"
+
+#include <cutlass/arch/barrier.h>
+#include <cutlass/arch/reg_reconfig.h>
+
+#include <cute/arch/cluster_sm90.hpp>
+#include <cute/arch/copy_sm90_desc.hpp>
+#include <cute/arch/copy_sm90_tma.hpp>
+
+#include <deep_gemm/common/epilogue_utils.cuh>
+#include <deep_gemm/common/utils.cuh>
+#include <deep_gemm/common/scheduler.cuh>
+#include <deep_gemm/common/sm90_utils.cuh>
+
+namespace deep_gemm {
+
+using namespace deep_gemm::sm90;
+
+template <uint32_t kNumFormerIters, uint32_t kGap, uint32_t kEnd, typename func_t>
+__device__ void dispatch_num_former_iters(uint32_t num_former_iters, const func_t& func) {
+    if (num_former_iters == kNumFormerIters) {
+        func(cute::Int<kNumFormerIters>{});
+        return;
+    }
+
+    if constexpr (kNumFormerIters + kGap <= kEnd)
+        dispatch_num_former_iters<kNumFormerIters + kGap, kGap, kEnd>(num_former_iters, func);
+}
+
+template <uint32_t SHAPE_M, uint32_t SHAPE_N, uint32_t SHAPE_K,
+          uint32_t kNumGroups,
+          uint32_t BLOCK_M, uint32_t BLOCK_N, uint32_t BLOCK_K,
+          uint32_t kSwizzleDMode,
+          uint32_t kNumStages, uint32_t kNumLastStages,
+          uint32_t kNumTMAThreads, uint32_t kNumMathThreads,
+          uint32_t kNumTMAMulticast, bool kIsTMAMulticastOnA,
+          uint32_t kNumSMs, GemmType kGemmType,
+          typename epilogue_type_t>
+__global__ __launch_bounds__(kNumTMAThreads + kNumMathThreads, 1) void
+sm90_fp8_gemm_1d2d_impl(float* sfb, int* grouped_layout,
+                        uint32_t shape_m, uint32_t shape_n, uint32_t shape_k,
+                        const __grid_constant__ cute::TmaDescriptor tensor_map_a,
+                        const __grid_constant__ cute::TmaDescriptor tensor_map_b,
+                        const __grid_constant__ cute::TmaDescriptor tensor_map_d,
+                        const __grid_constant__ cute::TmaDescriptor tensor_map_sfa) {
+#if (defined(__CUDA_ARCH__) and (__CUDA_ARCH__ >= 900)) or defined(__CLION_IDE__)
+    // Scaling checks
+    DG_STATIC_ASSERT(BLOCK_K == 128, "Only support per-128-channel FP8 scaling");
+    DG_STATIC_ASSERT(constexpr_ceil_div(BLOCK_N, BLOCK_K) == 1 or (constexpr_gcd(BLOCK_N, BLOCK_K) == BLOCK_N - BLOCK_K), "Too much B scales in a single block");
+
+    // Types
+    using WGMMA = typename FP8MMASelector<BLOCK_N>::type;
+    using Barrier = cutlass::arch::ClusterTransactionBarrier;
+    DG_STATIC_ASSERT(BLOCK_M % WGMMA::M == 0, "Invalid block size");
+
+    // Overwrite shape constants if the compiler gives
+    shape_m = SHAPE_M != 0 ? SHAPE_M : shape_m;
+    shape_n = SHAPE_N != 0 ? SHAPE_N : shape_n;
+    shape_k = SHAPE_K != 0 ? SHAPE_K : shape_k;
+
+    // Shared memory
+    static constexpr bool kMustUseUniformedScaleB = (BLOCK_K % BLOCK_N == 0);
+    static constexpr uint32_t SMEM_D_SIZE = BLOCK_M * BLOCK_N * sizeof(__nv_bfloat16);
+    static constexpr uint32_t SMEM_A_SIZE_PER_STAGE = BLOCK_M * BLOCK_K * sizeof(__nv_fp8_e4m3);
+    static constexpr uint32_t SMEM_B_SIZE_PER_STAGE = BLOCK_N * BLOCK_K * sizeof(__nv_fp8_e4m3);
+    static constexpr uint32_t SMEM_SFA_SIZE_PER_STAGE = BLOCK_M * sizeof(float);
+    const uint32_t& shape_k_scales = ceil_div(shape_k, BLOCK_K);
+    const uint32_t& smem_sfb_size = align<uint32_t>(shape_k_scales * (kMustUseUniformedScaleB ? 1 : 2) * sizeof(float), sizeof(Barrier));
+
+    // Configs
+    const uint32_t num_total_k_blocks = ceil_div(shape_k, BLOCK_K);
+    const uint32_t warp_idx = __shfl_sync(0xffffffff, threadIdx.x / 32, 0);
+    const uint32_t lane_idx = get_lane_idx();
+
+    // Prefetch TMA descriptors at the very beginning
+    if (warp_idx == kNumMathThreads / 32 and cute::elect_one_sync()) {
+        cute::prefetch_tma_descriptor(&tensor_map_a);
+        cute::prefetch_tma_descriptor(&tensor_map_b);
+        cute::prefetch_tma_descriptor(&tensor_map_sfa);
+        cute::prefetch_tma_descriptor(&tensor_map_d);
+    }
+    __syncwarp();
+
+    // Align to 1024 bytes for swizzle-128B
+    extern __shared__ __align__(1024) uint8_t smem_buffer[];
+    DG_STATIC_ASSERT(SMEM_D_SIZE % 1024 == 0, "Shared memory of A/B must be aligned to 1024 bytes");
+
+    // Data on shared memory
+    auto smem_d = reinterpret_cast<__nv_bfloat16*>(smem_buffer);
+    auto smem_a = PatternVisitor([&](const uint32_t& i) {
+        return reinterpret_cast<__nv_fp8_e4m3*>(smem_buffer + SMEM_D_SIZE + i * SMEM_A_SIZE_PER_STAGE);
+    });
+    auto smem_b = PatternVisitor([&](const uint32_t& i) {
+        return reinterpret_cast<__nv_fp8_e4m3*>(smem_buffer + SMEM_D_SIZE + kNumStages * SMEM_A_SIZE_PER_STAGE + i * SMEM_B_SIZE_PER_STAGE);
+    });
+    constexpr uint32_t SMEM_SF_OFFSET = SMEM_D_SIZE + kNumStages * (SMEM_A_SIZE_PER_STAGE + SMEM_B_SIZE_PER_STAGE);
+    auto smem_sfa = PatternVisitor([&](const uint32_t& i) {
+        return reinterpret_cast<float*>(smem_buffer + SMEM_SF_OFFSET + i * SMEM_SFA_SIZE_PER_STAGE);
+    });
+    auto smem_sfb = reinterpret_cast<float*>(smem_buffer + SMEM_SF_OFFSET + kNumStages * SMEM_SFA_SIZE_PER_STAGE);
+
+    // Fill barriers
+    auto barrier_start_ptr = reinterpret_cast<Barrier*>(reinterpret_cast<uint8_t*>(smem_sfb) + smem_sfb_size);
+    auto full_barriers     = PatternVisitor([&](const uint32_t& i) { return barrier_start_ptr + i; });
+    auto empty_barriers    = PatternVisitor([&](const uint32_t& i) { return barrier_start_ptr + kNumStages + i; });
+
+    // Initialize barriers
+    DG_STATIC_ASSERT(kNumTMAMulticast <= 32, "Too many TMA multicast");
+    if (warp_idx == kNumMathThreads / 32 + 1 and cute::elect_one_sync()) {
+        // NOTES: we always use `lane_idx` to arrive for the `lane_idx`-th CTA in the cluster,
+        // even with TMA multicast disabled, we want to make the behavior aligned
+        #pragma unroll
+        for (uint32_t i = 0; i < kNumStages; ++ i) {
+            full_barriers[i]->init(1);
+            empty_barriers[i]->init(kNumTMAMulticast * kNumMathThreads / 32);
+        }
+
+        // Make initialized barrier visible in async proxy
+        cutlass::arch::fence_barrier_init();
+    }
+
+    // Synchronize all threads to make barrier visible in normal memory model
+    (kNumTMAMulticast > 1) ? cute::cluster_sync() : __syncthreads();
+
+    // Register reconfigurations
+    constexpr uint32_t kNumTMARegisters = 40;
+    constexpr uint32_t kNumMathRegisters = 232;
+
+    // Block scheduler
+    uint32_t m_block_idx, n_block_idx;
+    auto scheduler = Scheduler<kGemmType, BLOCK_M, BLOCK_N, kNumGroups, kNumTMAMulticast, kIsTMAMulticastOnA, kNumSMs>(shape_m, shape_n, shape_k, grouped_layout);
+
+    // Pipeline and TMA phases
+    uint32_t stage_idx = 0, phase = 0;
+    auto advance_pipeline = [&](uint32_t& k_block_idx) {
+        ++ k_block_idx;
+
+        // Flip phases only if reach the next first stage
+        stage_idx = stage_idx == kNumStages - 1 ? 0 : stage_idx + 1;
+        phase ^= stage_idx == 0;
+    };
+
+    if (warp_idx >= kNumMathThreads / 32) {
+        // TMA warp-group for loading data
+        cutlass::arch::warpgroup_reg_dealloc<kNumTMARegisters>();
+
+        // NOTES: only one thread (or warp) will be used
+        if (warp_idx == kNumMathThreads / 32 and cute::elect_one_sync()) {
+            // Persistently schedule over blocks
+            while (scheduler.get_next_block(m_block_idx, n_block_idx)) {
+                // Assign TMA multicast number into A and B
+                // NOTES: there may be additional odd rows/columns or cases where multicast is not possible.
+                const bool is_tma_multicast_valid = scheduler.is_tma_multicast_valid(m_block_idx);
+                const uint32_t num_tma_multicast_a = (kIsTMAMulticastOnA and is_tma_multicast_valid) ? kNumTMAMulticast : 1;
+                const uint32_t num_tma_multicast_b = (not kIsTMAMulticastOnA and is_tma_multicast_valid) ? kNumTMAMulticast : 1;
+                DG_STATIC_ASSERT(kNumTMAMulticast <= 2, "Scheduler does not support > 2 TMA multicast");
+
+                for (uint32_t k_block_idx = 0; k_block_idx < num_total_k_blocks; advance_pipeline(k_block_idx)) {
+                    // Wait consumer release
+                    empty_barriers[stage_idx]->wait(phase ^ 1);
+
+                    // Issue TMA A
+                    constexpr bool kWithGroupOffsetA = kGemmType == GemmType::MGroupedMasked;
+                    auto& full_barrier = *full_barriers[stage_idx];
+                    const uint32_t k_idx = k_block_idx * BLOCK_K;
+                    tma_copy(&tensor_map_a, reinterpret_cast<uint64_t*>(&full_barrier),
+                             smem_a[stage_idx], k_idx, scheduler.get_global_idx<kWithGroupOffsetA>(shape_m, BLOCK_M, m_block_idx),
+                             num_tma_multicast_a);
+                    tma_copy(&tensor_map_sfa, reinterpret_cast<uint64_t*>(&full_barrier),
+                             smem_sfa[stage_idx], m_block_idx * BLOCK_M, scheduler.get_global_idx<kWithGroupOffsetA>(shape_k_scales, 1, k_block_idx),
+                             num_tma_multicast_a);
+
+                    // Issue TMA B
+                    tma_copy(&tensor_map_b, reinterpret_cast<uint64_t*>(&full_barrier),
+                             smem_b[stage_idx], k_idx, scheduler.get_global_idx<true>(shape_n, BLOCK_N, n_block_idx, m_block_idx),
+                             num_tma_multicast_b);
+                    full_barrier.arrive_and_expect_tx(SMEM_A_SIZE_PER_STAGE + SMEM_B_SIZE_PER_STAGE + SMEM_SFA_SIZE_PER_STAGE);
+                }
+            }
+
+            // To safely deconstruct distributed shared barriers, we need another round of empty waits
+            if constexpr (kNumTMAMulticast > 1) {
+                for (uint32_t i = 0; i < kNumStages; advance_pipeline(i))
+                    empty_barriers[stage_idx]->wait(phase ^ 1);
+            }
+        }
+    } else {
+        // Math warp-groups for WGMMA
+        cutlass::arch::warpgroup_reg_alloc<kNumMathRegisters>();
+
+        // NOTES: use `__shfl_sync` to encourage NVCC to use unified registers
+        const auto math_wg_idx = __shfl_sync(0xffffffff, threadIdx.x / 128, 0);
+        const auto r_0 = warp_idx * 16 + lane_idx / 4, r_1 = r_0 + 8;
+
+        auto a_desc = make_smem_desc(smem_a[0] + math_wg_idx * WGMMA::M * BLOCK_K, 1);
+        auto b_desc = make_smem_desc(smem_b[0], 1);
+        const uint32_t a_desc_lo = __shfl_sync(0xffffffff, a_desc.reg32_[0], 0);
+        const uint32_t b_desc_lo = __shfl_sync(0xffffffff, b_desc.reg32_[0], 0);
+
+        // Persistently schedule over blocks
+        while (scheduler.get_next_block(m_block_idx, n_block_idx)) {
+            // Decide the number of scales B to load
+            DG_TRAP_ONLY_DEVICE_ASSERT(shape_n % 8 == 0);
+            uint32_t num_former_iters = BLOCK_N / 8, num_full_iters = num_former_iters;
+            if constexpr (not kMustUseUniformedScaleB) {
+                num_former_iters = min(BLOCK_N, BLOCK_K - n_block_idx * BLOCK_N % BLOCK_K) / 8;
+                num_full_iters = min(shape_n - n_block_idx * BLOCK_N, BLOCK_N) / 8;
+            }
+            uint32_t num_sfb = shape_k_scales * (num_former_iters >= num_full_iters ? 1 : 2);
+
+            // Load B scales with math warp-groups
+            // NOTES: except the first warp, we want to overlap loading B scales with TMA stores between tasks
+            if (threadIdx.x >= 32) {
+                auto num_previous_lines = scheduler.get_global_idx<true>(ceil_div(shape_n, BLOCK_K), 0, 0, m_block_idx);
+                auto local_sfb = sfb + (num_previous_lines + ((n_block_idx * BLOCK_N) / BLOCK_K)) * shape_k_scales;
+                #pragma unroll
+                for (uint32_t i = threadIdx.x - 32; i < num_sfb; i += kNumMathThreads - 32)
+                    st_shared(smem_sfb + i, __ldg(local_sfb + i));
+            }
+            cutlass::arch::NamedBarrier::sync(kNumMathThreads, 0);
+
+            // Accumulation for WGMMA or CUDA promotion
+            constexpr uint32_t WAVE_BLOCK_M = WGMMA::M * (BLOCK_M <= 64 ? 1 : 2);
+            DG_STATIC_ASSERT(BLOCK_M % WAVE_BLOCK_M == 0, "Invalid block sizes");
+            float accum[WGMMA::kNumAccum], final_accum[WGMMA::kNumAccum * (BLOCK_M / WAVE_BLOCK_M)] = {0};
+
+            // Empty barrier arrival
+            auto empty_barrier_arrive = [&]() {
+                if constexpr (kNumTMAMulticast == 1) {
+                    lane_idx == 0 ? empty_barriers[stage_idx]->arrive() : void();
+                } else {
+                    auto target_cta = scheduler.is_peer_cta_alive ? lane_idx : cute::block_rank_in_cluster();
+                    lane_idx < kNumTMAMulticast ? empty_barriers[stage_idx]->arrive(target_cta) : void();
+                }
+            };
+
+            // Skip useless computations
+            if (scheduler.is_computation_valid(m_block_idx, math_wg_idx * WGMMA::M)) {
+                // The compiler must know the dynamic variable `num_former_iters`'s real value
+                constexpr bool kShouldOptimize = BLOCK_K / constexpr_gcd(BLOCK_K, BLOCK_N) <= 4 and not kMustUseUniformedScaleB;
+                constexpr uint32_t kGap = constexpr_gcd(BLOCK_K, BLOCK_N) / 8;
+                constexpr uint32_t kEnd = kShouldOptimize ? BLOCK_K / 8 : 0;
+
+                // Dispatch `num_former_iters` and launch MMAs
+                dispatch_num_former_iters<0, kGap, kEnd>(kShouldOptimize ? num_former_iters : 0, [&](auto _) {
+                    #pragma unroll 8
+                    for (uint32_t k_block_idx = 0; k_block_idx < num_total_k_blocks; advance_pipeline(k_block_idx)) {
+                        const auto& a_desc_base_lo = a_desc_lo + stage_idx * (SMEM_A_SIZE_PER_STAGE / 16);
+                        const auto& b_desc_base_lo = b_desc_lo + stage_idx * (SMEM_B_SIZE_PER_STAGE / 16);
+
+                        // Read B scales
+                        float scale_b_0 = ld_shared(smem_sfb + k_block_idx), scale_b_1;
+                        // NOTES: even some blocks do not need to read the second row, but we still load one to align with other blocks
+                        if constexpr (not kMustUseUniformedScaleB)
+                            scale_b_1 = ld_shared(smem_sfb + k_block_idx + shape_k_scales);
+
+                        // Wait TMA arrivals
+                        full_barriers[stage_idx]->wait(phase);
+
+                        // TODO: remove some useless computation for unaligned Ms
+                        #pragma unroll
+                        for (uint32_t local_idx = 0; local_idx < BLOCK_M / WAVE_BLOCK_M; ++ local_idx) {
+                            auto m_offset = local_idx * WAVE_BLOCK_M;
+
+                            // Read A scales
+                            // NOTES: all shared memory read must be prior to `warpgroup_arrive` to avoid next scheduled block polluting the results
+                            auto scale_a_0 = ld_shared(smem_sfa[stage_idx] + r_0 + m_offset);
+                            auto scale_a_1 = ld_shared(smem_sfa[stage_idx] + r_1 + m_offset);
+
+                            // Commit WGMMA instructions
+                            #pragma unroll
+                            for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i)
+                                warpgroup_fence_operand(accum[i]);
+                            warpgroup_arrive();
+                            #pragma unroll
+                            for (uint32_t k = 0; k < BLOCK_K / WGMMA::K; ++ k) {
+                                a_desc.reg32_[0] = a_desc_base_lo + (m_offset * BLOCK_K + k * WGMMA::K) / 16;
+                                b_desc.reg32_[0] = b_desc_base_lo + k * WGMMA::K / 16;
+                                WGMMA::wgmma(a_desc, b_desc, accum, k);
+                            }
+                            warpgroup_commit_batch();
+                            #pragma unroll
+                            for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i)
+                                warpgroup_fence_operand(accum[i]);
+                            warpgroup_wait<0>();
+
+                            // Notify barrier arrival at the last warpgroup wave
+                            if (local_idx == BLOCK_M / WAVE_BLOCK_M - 1)
+                                empty_barrier_arrive();
+
+                            // Promote with scales
+                            // NOTES: making it as predicates is very important for performance, comparing to two loops
+                            float scale_0_0 = scale_a_0 * scale_b_0, scale_1_0 = scale_a_1 * scale_b_0;
+                            float scale_0_1, scale_1_1;
+                            if constexpr (not kMustUseUniformedScaleB)
+                                scale_0_1 = scale_a_0 * scale_b_1, scale_1_1 = scale_a_1 * scale_b_1;
+
+                            auto shifted_accum = final_accum + WGMMA::kNumAccum * local_idx;
+                            #pragma unroll
+                            for (uint32_t i = 0; i < WGMMA::kNumAccum / 4; ++ i) {
+                                // NOTES: for unrolled `num_former_iters` cases, we expect the compiler to automatically make it a constant
+                                bool predicate = kMustUseUniformedScaleB or i < num_former_iters;
+                                shifted_accum[i * 4 + 0] += (predicate ? scale_0_0 : scale_0_1) * accum[i * 4 + 0];
+                                shifted_accum[i * 4 + 1] += (predicate ? scale_0_0 : scale_0_1) * accum[i * 4 + 1];
+                                shifted_accum[i * 4 + 2] += (predicate ? scale_1_0 : scale_1_1) * accum[i * 4 + 2];
+                                shifted_accum[i * 4 + 3] += (predicate ? scale_1_0 : scale_1_1) * accum[i * 4 + 3];
+                            }
+                        }
+                    }
+                });
+            } else {
+                #pragma unroll
+                for (uint32_t k_block_idx = 0; k_block_idx < num_total_k_blocks; advance_pipeline(k_block_idx)) {
+                    full_barriers[stage_idx]->wait(phase);
+                    empty_barrier_arrive();
+                }
+            }
+
+            // TMA checks
+            constexpr uint32_t kNumElemBytes = sizeof(nv_bfloat16);
+            constexpr uint32_t TMA_D_BLOCK_N = kSwizzleDMode == 0 ? BLOCK_N : (kSwizzleDMode / kNumElemBytes);
+            constexpr uint32_t WGMMA_M_PER_WARP = WGMMA::M / 4;
+            DG_STATIC_ASSERT(BLOCK_M % 8 == 0, "Invalid swizzling atom");
+            DG_STATIC_ASSERT(BLOCK_N % TMA_D_BLOCK_N == 0 and BLOCK_N / TMA_D_BLOCK_N <= 32,
+                            "Unaligned TMA store or too many TMA store instructions");
+            DG_STATIC_ASSERT(TMA_D_BLOCK_N % 8 == 0, "Invalid TMA block N");
+
+            // Wait last TMA store to be finished
+            if (threadIdx.x < BLOCK_N / TMA_D_BLOCK_N)
+                cute::tma_store_wait<0>();
+            cutlass::arch::NamedBarrier::sync(kNumMathThreads, 0);
+
+            // Write back to shared memory using STSM and issue TMA stores
+            DG_STATIC_ASSERT(WGMMA::kNumAccum % 4 == 0, "Invalid STSM x2 vectorization");
+            #pragma unroll
+            for (uint32_t local_idx = 0; local_idx < BLOCK_M / WAVE_BLOCK_M; ++ local_idx) {
+                auto m_offset = local_idx * WAVE_BLOCK_M;
+                auto shifted_accum = final_accum + WGMMA::kNumAccum * local_idx;
+                #pragma unroll
+                for (auto i = 0; i < WGMMA::kNumAccum / 4; ++ i) {
+                    // Swizzle or padding into the correct address
+                    uint8_t* smem_ptr = nullptr;
+                    if constexpr (kSwizzleDMode > 0) {
+                        // Calculate the swizzling atom offset and in-atom offset
+                        constexpr uint32_t kNumBankGroupBytes = 16;
+                        auto atom_offset = i / (TMA_D_BLOCK_N / 8), in_atom_offset = i % (TMA_D_BLOCK_N / 8);
+
+                        // Calculate the index of the bank group to be written in the atom
+                        auto bank_group_index = in_atom_offset + lane_idx * (kSwizzleDMode / kNumBankGroupBytes);
+
+                        // Reshape the atom in another view and swizzle
+                        //  - original: `(BLOCK_M, kSwizzleDMode / kNumBankGroupBytes)`
+                        //  - new: `(BLOCK_M * kSwizzleDMode / kNumBankGroupBytes / 8, 8)`
+                        constexpr bool kHasShortcut = (kSwizzleDMode / kNumBankGroupBytes) == 8;
+                        auto row = kHasShortcut ? (in_atom_offset / 8 + lane_idx) : (bank_group_index / 8);
+                        auto col = kHasShortcut ? (in_atom_offset) : (bank_group_index % 8);
+                        col ^= row % (kSwizzleDMode / 16);
+
+                        // Add back into the base pointer
+                        // NOTES: think twice before modifying this, as changes may affect the number of instructions
+                        smem_ptr = reinterpret_cast<uint8_t*>(smem_d) +                // Base pointer
+                            warp_idx * (WGMMA_M_PER_WARP * kSwizzleDMode) +            // Warp offset
+                            m_offset * kSwizzleDMode +                                 // Wave offset
+                            atom_offset * BLOCK_M * kSwizzleDMode +                    // Swizzle atom offset (constants)
+                            row * (kNumBankGroupBytes * 8) + col * kNumBankGroupBytes; // In-atom offset
+                    } else {
+                        // No swizzling, just padding
+                        smem_ptr = reinterpret_cast<uint8_t*>(smem_d + (m_offset + warp_idx * WGMMA_M_PER_WARP + lane_idx) * BLOCK_N + i * 8);
+                    }
+
+                    // NOTES: only 16 lanes' addresses are used
+                    SM90_U32x2_STSM_N<nv_bfloat162>::copy(
+                        __float22bfloat162_rn({shifted_accum[i * 4 + 0], shifted_accum[i * 4 + 1]}),
+                        __float22bfloat162_rn({shifted_accum[i * 4 + 2], shifted_accum[i * 4 + 3]}),
+                        smem_ptr
+                    );
+                }
+            }
+            cute::tma_store_fence();
+            cutlass::arch::NamedBarrier::sync(kNumMathThreads, 0);
+
+            // Use TMA store to write back to global memory
+            // TODO: compatible with FP32 output
+            constexpr bool kWithGroupOffsetD = kGemmType == GemmType::MGroupedMasked;
+            DG_STATIC_ASSERT(kNumMathThreads >= BLOCK_N / TMA_D_BLOCK_N, "Too many TMA blocks");
+            if (threadIdx.x < BLOCK_N / TMA_D_BLOCK_N) {
+                auto in_block_n_offset = threadIdx.x * TMA_D_BLOCK_N;
+                auto smem_ptr = smem_d + in_block_n_offset * BLOCK_M;
+                cute::SM90_TMA_STORE_2D::copy(&tensor_map_d, smem_ptr,
+                                              epilogue_type_t::apply_index_n<TMA_D_BLOCK_N>(n_block_idx * BLOCK_N + in_block_n_offset),
+                                              scheduler.get_global_idx<kWithGroupOffsetD>(shape_m, BLOCK_M, m_block_idx));
+                cute::tma_store_arrive();
+            }
+            __syncwarp();
+        }
+    }
+#else
+    if (blockIdx.x == 0 and threadIdx.x == 0)
+        DG_DEVICE_ASSERT(false and "This kernel only support sm_90a");
+#endif
+}
+
+};  // namespace deep_gemm
+
+#pragma clang diagnostic pop

packages/ltx-kernels/csrc/blockwise/kernels/deep_gemm/include/deep_gemm/impls/sm90_fp8_gemm_1d2d_bias.cu

@@ -0,0 +1,590 @@
+#include <cutlass/arch/barrier.h>
+#include <cutlass/arch/reg_reconfig.h>
+
+#include <cute/arch/cluster_sm90.hpp>
+#include <cute/arch/copy_sm90_desc.hpp>
+#include <cute/arch/copy_sm90_tma.hpp>
+
+#include <deep_gemm/common/epilogue_utils.cuh>
+#include <deep_gemm/common/utils.cuh>
+#include <deep_gemm/common/scheduler.cuh>
+#include <deep_gemm/common/sm90_utils.cuh>
+
+#include <stdexcept>
+#include <string>
+
+// LT-PATCH: upstream hard-#defines `__CUDA_ARCH__ 900` here, which forces the wgmma
+// kernel body on every compile pass and makes this source impossible to place in a
+// multi-arch fat binary (it emits sm_90-only instructions during e.g. the sm_89 pass
+// -> ptxas error). Removed so the existing `#if __CUDA_ARCH__ >= 900 ... #else assert
+// #endif` guard takes effect per-arch: the real body is built only into the sm_90a
+// cubin, other arches get a host-visible assert stub. The sm_90a pass is unchanged
+// (nvcc defines __CUDA_ARCH__=900 there regardless).
+
+namespace deep_gemm {
+
+using namespace deep_gemm::sm90;
+
+template <uint32_t kNumFormerIters, uint32_t kGap, uint32_t kEnd, typename func_t>
+__device__ void dispatch_num_former_iters(uint32_t num_former_iters, const func_t& func) {
+    if (num_former_iters == kNumFormerIters) {
+        func(cute::Int<kNumFormerIters>{});
+        return;
+    }
+
+    if constexpr (kNumFormerIters + kGap <= kEnd)
+        dispatch_num_former_iters<kNumFormerIters + kGap, kGap, kEnd>(num_former_iters, func);
+}
+
+template <uint32_t SHAPE_M, uint32_t SHAPE_N, uint32_t SHAPE_K,
+          uint32_t kNumGroups,
+          uint32_t BLOCK_M, uint32_t BLOCK_N, uint32_t BLOCK_K,
+          uint32_t kSwizzleDMode,
+          uint32_t kNumStages, uint32_t kNumLastStages,
+          uint32_t kNumTMAThreads, uint32_t kNumMathThreads,
+          uint32_t kNumTMAMulticast, bool kIsTMAMulticastOnA,
+          uint32_t kNumSMs, GemmType kGemmType,
+          typename epilogue_type_t>
+__global__ __launch_bounds__(kNumTMAThreads + kNumMathThreads, 1) void
+sm90_fp8_gemm_1d2d_bias_impl(float* sfb, float* bias, int* grouped_layout,
+                        uint32_t shape_m, uint32_t shape_n, uint32_t shape_k,
+                        const __grid_constant__ cute::TmaDescriptor tensor_map_a,
+                        const __grid_constant__ cute::TmaDescriptor tensor_map_b,
+                        const __grid_constant__ cute::TmaDescriptor tensor_map_d,
+                        const __grid_constant__ cute::TmaDescriptor tensor_map_sfa) {
+// LT-PATCH: was `__CUDA_ARCH__ >= 900`. Tightened to Hopper-only (< 1000) so that in a
+// multi-arch fat binary that also targets Blackwell (sm_100/sm_120), this wgmma body is
+// NOT emitted for those passes (wgmma is sm_90a-only) -- they get the `#else` assert stub
+// instead. Blackwell dispatches to the SM89 kernel at runtime, so the stub is never run.
+#if (defined(__CUDA_ARCH__) and (__CUDA_ARCH__ >= 900) and (__CUDA_ARCH__ < 1000)) or defined(__CLION_IDE__)
+    // Scaling checks
+    DG_STATIC_ASSERT(BLOCK_K == 128, "Only support per-128-channel FP8 scaling");
+    DG_STATIC_ASSERT(constexpr_ceil_div(BLOCK_N, BLOCK_K) == 1 or (constexpr_gcd(BLOCK_N, BLOCK_K) == BLOCK_N - BLOCK_K), "Too much B scales in a single block");
+
+    // Types
+    using WGMMA = typename FP8MMASelector<BLOCK_N>::type;
+    using Barrier = cutlass::arch::ClusterTransactionBarrier;
+    DG_STATIC_ASSERT(BLOCK_M % WGMMA::M == 0, "Invalid block size");
+
+    // Overwrite shape constants if the compiler gives
+    shape_m = SHAPE_M != 0 ? SHAPE_M : shape_m;
+    shape_n = SHAPE_N != 0 ? SHAPE_N : shape_n;
+    shape_k = SHAPE_K != 0 ? SHAPE_K : shape_k;
+
+    // Shared memory
+    static constexpr bool kMustUseUniformedScaleB = (BLOCK_K % BLOCK_N == 0);
+    static constexpr uint32_t SMEM_D_SIZE = BLOCK_M * BLOCK_N * sizeof(__nv_bfloat16);
+    static constexpr uint32_t SMEM_A_SIZE_PER_STAGE = BLOCK_M * BLOCK_K * sizeof(__nv_fp8_e4m3);
+    static constexpr uint32_t SMEM_B_SIZE_PER_STAGE = BLOCK_N * BLOCK_K * sizeof(__nv_fp8_e4m3);
+    static constexpr uint32_t SMEM_SFA_SIZE_PER_STAGE = BLOCK_M * sizeof(float);
+    const uint32_t& shape_k_scales = ceil_div(shape_k, BLOCK_K);
+    const uint32_t& smem_sfb_size = align<uint32_t>(shape_k_scales * (kMustUseUniformedScaleB ? 1 : 2) * sizeof(float), sizeof(Barrier));
+
+    // Configs
+    const uint32_t num_total_k_blocks = ceil_div(shape_k, BLOCK_K);
+    const uint32_t warp_idx = __shfl_sync(0xffffffff, threadIdx.x / 32, 0);
+    const uint32_t lane_idx = get_lane_idx();
+
+    // Prefetch TMA descriptors at the very beginning
+    if (warp_idx == kNumMathThreads / 32 and cute::elect_one_sync()) {
+        cute::prefetch_tma_descriptor(&tensor_map_a);
+        cute::prefetch_tma_descriptor(&tensor_map_b);
+        cute::prefetch_tma_descriptor(&tensor_map_sfa);
+        cute::prefetch_tma_descriptor(&tensor_map_d);
+    }
+    __syncwarp();
+
+    // Align to 1024 bytes for swizzle-128B
+    extern __shared__ __align__(1024) uint8_t smem_buffer[];
+    DG_STATIC_ASSERT(SMEM_D_SIZE % 1024 == 0, "Shared memory of A/B must be aligned to 1024 bytes");
+
+    // Data on shared memory
+    auto smem_d = reinterpret_cast<__nv_bfloat16*>(smem_buffer);
+    auto smem_a = PatternVisitor([&](const uint32_t& i) {
+        return reinterpret_cast<__nv_fp8_e4m3*>(smem_buffer + SMEM_D_SIZE + i * SMEM_A_SIZE_PER_STAGE);
+    });
+    auto smem_b = PatternVisitor([&](const uint32_t& i) {
+        return reinterpret_cast<__nv_fp8_e4m3*>(smem_buffer + SMEM_D_SIZE + kNumStages * SMEM_A_SIZE_PER_STAGE + i * SMEM_B_SIZE_PER_STAGE);
+    });
+    constexpr uint32_t SMEM_SF_OFFSET = SMEM_D_SIZE + kNumStages * (SMEM_A_SIZE_PER_STAGE + SMEM_B_SIZE_PER_STAGE);
+    auto smem_sfa = PatternVisitor([&](const uint32_t& i) {
+        return reinterpret_cast<float*>(smem_buffer + SMEM_SF_OFFSET + i * SMEM_SFA_SIZE_PER_STAGE);
+    });
+    auto smem_sfb = reinterpret_cast<float*>(smem_buffer + SMEM_SF_OFFSET + kNumStages * SMEM_SFA_SIZE_PER_STAGE);
+
+    // Fill barriers
+    auto barrier_start_ptr = reinterpret_cast<Barrier*>(reinterpret_cast<uint8_t*>(smem_sfb) + smem_sfb_size);
+    auto full_barriers     = PatternVisitor([&](const uint32_t& i) { return barrier_start_ptr + i; });
+    auto empty_barriers    = PatternVisitor([&](const uint32_t& i) { return barrier_start_ptr + kNumStages + i; });
+
+    // Initialize barriers
+    DG_STATIC_ASSERT(kNumTMAMulticast <= 32, "Too many TMA multicast");
+    if (warp_idx == kNumMathThreads / 32 + 1 and cute::elect_one_sync()) {
+        // NOTES: we always use `lane_idx` to arrive for the `lane_idx`-th CTA in the cluster,
+        // even with TMA multicast disabled, we want to make the behavior aligned
+        #pragma unroll
+        for (uint32_t i = 0; i < kNumStages; ++ i) {
+            full_barriers[i]->init(1);
+            empty_barriers[i]->init(kNumTMAMulticast * kNumMathThreads / 32);
+        }
+
+        // Make initialized barrier visible in async proxy
+        cutlass::arch::fence_barrier_init();
+    }
+
+    // Synchronize all threads to make barrier visible in normal memory model
+    (kNumTMAMulticast > 1) ? cute::cluster_sync() : __syncthreads();
+
+    // Register reconfigurations
+    constexpr uint32_t kNumTMARegisters = 40;
+    constexpr uint32_t kNumMathRegisters = 232;
+
+    // Block scheduler
+    uint32_t m_block_idx, n_block_idx;
+    auto scheduler = Scheduler<kGemmType, BLOCK_M, BLOCK_N, kNumGroups, kNumTMAMulticast, kIsTMAMulticastOnA, kNumSMs>(shape_m, shape_n, shape_k, grouped_layout);
+
+    // Pipeline and TMA phases
+    uint32_t stage_idx = 0, phase = 0;
+    auto advance_pipeline = [&](uint32_t& k_block_idx) {
+        ++ k_block_idx;
+
+        // Flip phases only if reach the next first stage
+        stage_idx = stage_idx == kNumStages - 1 ? 0 : stage_idx + 1;
+        phase ^= stage_idx == 0;
+    };
+
+    if (warp_idx >= kNumMathThreads / 32) {
+        // TMA warp-group for loading data
+        cutlass::arch::warpgroup_reg_dealloc<kNumTMARegisters>();
+
+        // NOTES: only one thread (or warp) will be used
+        if (warp_idx == kNumMathThreads / 32 and cute::elect_one_sync()) {
+            // Persistently schedule over blocks
+            while (scheduler.get_next_block(m_block_idx, n_block_idx)) {
+                // Assign TMA multicast number into A and B
+                // NOTES: there may be additional odd rows/columns or cases where multicast is not possible.
+                const bool is_tma_multicast_valid = scheduler.is_tma_multicast_valid(m_block_idx);
+                const uint32_t num_tma_multicast_a = (kIsTMAMulticastOnA and is_tma_multicast_valid) ? kNumTMAMulticast : 1;
+                const uint32_t num_tma_multicast_b = (not kIsTMAMulticastOnA and is_tma_multicast_valid) ? kNumTMAMulticast : 1;
+                DG_STATIC_ASSERT(kNumTMAMulticast <= 2, "Scheduler does not support > 2 TMA multicast");
+
+                for (uint32_t k_block_idx = 0; k_block_idx < num_total_k_blocks; advance_pipeline(k_block_idx)) {
+                    // Wait consumer release
+                    empty_barriers[stage_idx]->wait(phase ^ 1);
+
+                    // Issue TMA A
+                    constexpr bool kWithGroupOffsetA = kGemmType == GemmType::MGroupedMasked;
+                    auto& full_barrier = *full_barriers[stage_idx];
+                    const uint32_t k_idx = k_block_idx * BLOCK_K;
+                    tma_copy(&tensor_map_a, reinterpret_cast<uint64_t*>(&full_barrier),
+                             smem_a[stage_idx], k_idx, scheduler.get_global_idx<kWithGroupOffsetA>(shape_m, BLOCK_M, m_block_idx),
+                             num_tma_multicast_a);
+                    tma_copy(&tensor_map_sfa, reinterpret_cast<uint64_t*>(&full_barrier),
+                             smem_sfa[stage_idx], m_block_idx * BLOCK_M, scheduler.get_global_idx<kWithGroupOffsetA>(shape_k_scales, 1, k_block_idx),
+                             num_tma_multicast_a);
+
+                    // Issue TMA B
+                    tma_copy(&tensor_map_b, reinterpret_cast<uint64_t*>(&full_barrier),
+                             smem_b[stage_idx], k_idx, scheduler.get_global_idx<true>(shape_n, BLOCK_N, n_block_idx, m_block_idx),
+                             num_tma_multicast_b);
+                    full_barrier.arrive_and_expect_tx(SMEM_A_SIZE_PER_STAGE + SMEM_B_SIZE_PER_STAGE + SMEM_SFA_SIZE_PER_STAGE);
+                }
+            }
+
+            // To safely deconstruct distributed shared barriers, we need another round of empty waits
+            if constexpr (kNumTMAMulticast > 1) {
+                for (uint32_t i = 0; i < kNumStages; advance_pipeline(i))
+                    empty_barriers[stage_idx]->wait(phase ^ 1);
+            }
+        }
+    } else {
+        // Math warp-groups for WGMMA
+        cutlass::arch::warpgroup_reg_alloc<kNumMathRegisters>();
+
+        // NOTES: use `__shfl_sync` to encourage NVCC to use unified registers
+        const auto math_wg_idx = __shfl_sync(0xffffffff, threadIdx.x / 128, 0);
+        const auto r_0 = warp_idx * 16 + lane_idx / 4, r_1 = r_0 + 8;
+
+        auto a_desc = make_smem_desc(smem_a[0] + math_wg_idx * WGMMA::M * BLOCK_K, 1);
+        auto b_desc = make_smem_desc(smem_b[0], 1);
+        const uint32_t a_desc_lo = __shfl_sync(0xffffffff, a_desc.reg32_[0], 0);
+        const uint32_t b_desc_lo = __shfl_sync(0xffffffff, b_desc.reg32_[0], 0);
+
+        // Persistently schedule over blocks
+        while (scheduler.get_next_block(m_block_idx, n_block_idx)) {
+            // Decide the number of scales B to load
+            DG_TRAP_ONLY_DEVICE_ASSERT(shape_n % 8 == 0);
+            uint32_t num_former_iters = BLOCK_N / 8, num_full_iters = num_former_iters;
+            if constexpr (not kMustUseUniformedScaleB) {
+                num_former_iters = min(BLOCK_N, BLOCK_K - n_block_idx * BLOCK_N % BLOCK_K) / 8;
+                num_full_iters = min(shape_n - n_block_idx * BLOCK_N, BLOCK_N) / 8;
+            }
+            uint32_t num_sfb = shape_k_scales * (num_former_iters >= num_full_iters ? 1 : 2);
+
+            // Load B scales with math warp-groups
+            // NOTES: except the first warp, we want to overlap loading B scales with TMA stores between tasks
+            if (threadIdx.x >= 32) {
+                auto num_previous_lines = scheduler.get_global_idx<true>(ceil_div(shape_n, BLOCK_K), 0, 0, m_block_idx);
+                auto local_sfb = sfb + (num_previous_lines + ((n_block_idx * BLOCK_N) / BLOCK_K)) * shape_k_scales;
+                #pragma unroll
+                for (uint32_t i = threadIdx.x - 32; i < num_sfb; i += kNumMathThreads - 32)
+                    st_shared(smem_sfb + i, __ldg(local_sfb + i));
+            }
+            cutlass::arch::NamedBarrier::sync(kNumMathThreads, 0);
+
+            // Accumulation for WGMMA or CUDA promotion
+            constexpr uint32_t WAVE_BLOCK_M = WGMMA::M * (BLOCK_M <= 64 ? 1 : 2);
+            DG_STATIC_ASSERT(BLOCK_M % WAVE_BLOCK_M == 0, "Invalid block sizes");
+            float accum[WGMMA::kNumAccum], final_accum[WGMMA::kNumAccum * (BLOCK_M / WAVE_BLOCK_M)] = {0};
+
+            // Empty barrier arrival
+            auto empty_barrier_arrive = [&]() {
+                if constexpr (kNumTMAMulticast == 1) {
+                    lane_idx == 0 ? empty_barriers[stage_idx]->arrive() : void();
+                } else {
+                    auto target_cta = scheduler.is_peer_cta_alive ? lane_idx : cute::block_rank_in_cluster();
+                    lane_idx < kNumTMAMulticast ? empty_barriers[stage_idx]->arrive(target_cta) : void();
+                }
+            };
+
+            // Skip useless computations
+            if (scheduler.is_computation_valid(m_block_idx, math_wg_idx * WGMMA::M)) {
+                // The compiler must know the dynamic variable `num_former_iters`'s real value
+                constexpr bool kShouldOptimize = BLOCK_K / constexpr_gcd(BLOCK_K, BLOCK_N) <= 4 and not kMustUseUniformedScaleB;
+                constexpr uint32_t kGap = constexpr_gcd(BLOCK_K, BLOCK_N) / 8;
+                constexpr uint32_t kEnd = kShouldOptimize ? BLOCK_K / 8 : 0;
+
+                // Dispatch `num_former_iters` and launch MMAs
+                dispatch_num_former_iters<0, kGap, kEnd>(kShouldOptimize ? num_former_iters : 0, [&](auto _) {
+                    #pragma unroll 8
+                    for (uint32_t k_block_idx = 0; k_block_idx < num_total_k_blocks; advance_pipeline(k_block_idx)) {
+                        const auto& a_desc_base_lo = a_desc_lo + stage_idx * (SMEM_A_SIZE_PER_STAGE / 16);
+                        const auto& b_desc_base_lo = b_desc_lo + stage_idx * (SMEM_B_SIZE_PER_STAGE / 16);
+
+                        // Read B scales
+                        float scale_b_0 = ld_shared(smem_sfb + k_block_idx), scale_b_1;
+                        // NOTES: even some blocks do not need to read the second row, but we still load one to align with other blocks
+                        if constexpr (not kMustUseUniformedScaleB)
+                            scale_b_1 = ld_shared(smem_sfb + k_block_idx + shape_k_scales);
+
+                        // Wait TMA arrivals
+                        full_barriers[stage_idx]->wait(phase);
+
+                        // TODO: remove some useless computation for unaligned Ms
+                        #pragma unroll
+                        for (uint32_t local_idx = 0; local_idx < BLOCK_M / WAVE_BLOCK_M; ++ local_idx) {
+                            auto m_offset = local_idx * WAVE_BLOCK_M;
+
+                            // Read A scales
+                            // NOTES: all shared memory read must be prior to `warpgroup_arrive` to avoid next scheduled block polluting the results
+                            auto scale_a_0 = ld_shared(smem_sfa[stage_idx] + r_0 + m_offset);
+                            auto scale_a_1 = ld_shared(smem_sfa[stage_idx] + r_1 + m_offset);
+
+                            // Commit WGMMA instructions
+                            #pragma unroll
+                            for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i)
+                                warpgroup_fence_operand(accum[i]);
+                            warpgroup_arrive();
+                            #pragma unroll
+                            for (uint32_t k = 0; k < BLOCK_K / WGMMA::K; ++ k) {
+                                a_desc.reg32_[0] = a_desc_base_lo + (m_offset * BLOCK_K + k * WGMMA::K) / 16;
+                                b_desc.reg32_[0] = b_desc_base_lo + k * WGMMA::K / 16;
+                                WGMMA::wgmma(a_desc, b_desc, accum, k);
+                            }
+                            warpgroup_commit_batch();
+                            #pragma unroll
+                            for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i)
+                                warpgroup_fence_operand(accum[i]);
+                            warpgroup_wait<0>();
+
+                            // Notify barrier arrival at the last warpgroup wave
+                            if (local_idx == BLOCK_M / WAVE_BLOCK_M - 1)
+                                empty_barrier_arrive();
+
+                            // Promote with scales
+                            // NOTES: making it as predicates is very important for performance, comparing to two loops
+                            float scale_0_0 = scale_a_0 * scale_b_0, scale_1_0 = scale_a_1 * scale_b_0;
+                            float scale_0_1, scale_1_1;
+                            if constexpr (not kMustUseUniformedScaleB)
+                                scale_0_1 = scale_a_0 * scale_b_1, scale_1_1 = scale_a_1 * scale_b_1;
+
+                            auto shifted_accum = final_accum + WGMMA::kNumAccum * local_idx;
+                            #pragma unroll
+                            for (uint32_t i = 0; i < WGMMA::kNumAccum / 4; ++ i) {
+                                // NOTES: for unrolled `num_former_iters` cases, we expect the compiler to automatically make it a constant
+                                bool predicate = kMustUseUniformedScaleB or i < num_former_iters;
+                                shifted_accum[i * 4 + 0] += (predicate ? scale_0_0 : scale_0_1) * accum[i * 4 + 0];
+                                shifted_accum[i * 4 + 1] += (predicate ? scale_0_0 : scale_0_1) * accum[i * 4 + 1];
+                                shifted_accum[i * 4 + 2] += (predicate ? scale_1_0 : scale_1_1) * accum[i * 4 + 2];
+                                shifted_accum[i * 4 + 3] += (predicate ? scale_1_0 : scale_1_1) * accum[i * 4 + 3];
+                            }
+                        }
+                    }
+                });
+            } else {
+                #pragma unroll
+                for (uint32_t k_block_idx = 0; k_block_idx < num_total_k_blocks; advance_pipeline(k_block_idx)) {
+                    full_barriers[stage_idx]->wait(phase);
+                    empty_barrier_arrive();
+                }
+            }
+
+            // TMA checks
+            constexpr uint32_t kNumElemBytes = sizeof(nv_bfloat16);
+            constexpr uint32_t TMA_D_BLOCK_N = kSwizzleDMode == 0 ? BLOCK_N : (kSwizzleDMode / kNumElemBytes);
+            constexpr uint32_t WGMMA_M_PER_WARP = WGMMA::M / 4;
+            DG_STATIC_ASSERT(BLOCK_M % 8 == 0, "Invalid swizzling atom");
+            DG_STATIC_ASSERT(BLOCK_N % TMA_D_BLOCK_N == 0 and BLOCK_N / TMA_D_BLOCK_N <= 32,
+                            "Unaligned TMA store or too many TMA store instructions");
+            DG_STATIC_ASSERT(TMA_D_BLOCK_N % 8 == 0, "Invalid TMA block N");
+            // Wait last TMA store to be finished
+            float* bias_ptr = bias + n_block_idx*BLOCK_N + (lane_idx % 4) * 2;
+            #pragma unroll
+            for(uint32_t local_idx=0; local_idx < BLOCK_M / WAVE_BLOCK_M; ++ local_idx){
+                auto shifted_accum = final_accum + WGMMA::kNumAccum * local_idx;
+                #pragma unroll
+                for (auto i = 0; i < WGMMA::kNumAccum / 4; ++ i) {
+                    shifted_accum[4*i + 0] += bias_ptr[8*i + 0];
+                    shifted_accum[4*i + 1] += bias_ptr[8*i + 1];
+                    shifted_accum[4*i + 2] += bias_ptr[8*i + 0];
+                    shifted_accum[4*i + 3] += bias_ptr[8*i + 1];
+                }
+            }
+            
+            if (threadIdx.x < BLOCK_N / TMA_D_BLOCK_N)
+                cute::tma_store_wait<0>();
+            cutlass::arch::NamedBarrier::sync(kNumMathThreads, 0);
+            
+            // Write back to shared memory using STSM and issue TMA stores
+            DG_STATIC_ASSERT(WGMMA::kNumAccum % 4 == 0, "Invalid STSM x2 vectorization");
+            #pragma unroll
+            for (uint32_t local_idx = 0; local_idx < BLOCK_M / WAVE_BLOCK_M; ++ local_idx) {
+                auto m_offset = local_idx * WAVE_BLOCK_M;
+                auto shifted_accum = final_accum + WGMMA::kNumAccum * local_idx;
+                #pragma unroll
+                for (auto i = 0; i < WGMMA::kNumAccum / 4; ++ i) {
+                    // Swizzle or padding into the correct address
+                    uint8_t* smem_ptr = nullptr;
+                    if constexpr (kSwizzleDMode > 0) {
+                        // Calculate the swizzling atom offset and in-atom offset
+                        constexpr uint32_t kNumBankGroupBytes = 16;
+                        auto atom_offset = i / (TMA_D_BLOCK_N / 8), in_atom_offset = i % (TMA_D_BLOCK_N / 8);
+
+                        // Calculate the index of the bank group to be written in the atom
+                        auto bank_group_index = in_atom_offset + lane_idx * (kSwizzleDMode / kNumBankGroupBytes);
+
+                        // Reshape the atom in another view and swizzle
+                        //  - original: `(BLOCK_M, kSwizzleDMode / kNumBankGroupBytes)`
+                        //  - new: `(BLOCK_M * kSwizzleDMode / kNumBankGroupBytes / 8, 8)`
+                        constexpr bool kHasShortcut = (kSwizzleDMode / kNumBankGroupBytes) == 8;
+                        auto row = kHasShortcut ? (in_atom_offset / 8 + lane_idx) : (bank_group_index / 8);
+                        auto col = kHasShortcut ? (in_atom_offset) : (bank_group_index % 8);
+                        col ^= row % (kSwizzleDMode / 16);
+
+                        // Add back into the base pointer
+                        // NOTES: think twice before modifying this, as changes may affect the number of instructions
+                        smem_ptr = reinterpret_cast<uint8_t*>(smem_d) +                // Base pointer
+                            warp_idx * (WGMMA_M_PER_WARP * kSwizzleDMode) +            // Warp offset
+                            m_offset * kSwizzleDMode +                                 // Wave offset
+                            atom_offset * BLOCK_M * kSwizzleDMode +                    // Swizzle atom offset (constants)
+                            row * (kNumBankGroupBytes * 8) + col * kNumBankGroupBytes; // In-atom offset
+                    } else {
+                        // No swizzling, just padding
+                        smem_ptr = reinterpret_cast<uint8_t*>(smem_d + (m_offset + warp_idx * WGMMA_M_PER_WARP + lane_idx) * BLOCK_N + i * 8);
+                    }
+
+                    // NOTES: only 16 lanes' addresses are used
+                    SM90_U32x2_STSM_N<nv_bfloat162>::copy(
+                        __float22bfloat162_rn({shifted_accum[i * 4 + 0], shifted_accum[i * 4 + 1]}),
+                        __float22bfloat162_rn({shifted_accum[i * 4 + 2], shifted_accum[i * 4 + 3]}),
+                        smem_ptr
+                    );
+                }
+            }
+            cute::tma_store_fence();
+            cutlass::arch::NamedBarrier::sync(kNumMathThreads, 0);
+
+            // Use TMA store to write back to global memory
+            // TODO: compatible with FP32 output
+            constexpr bool kWithGroupOffsetD = kGemmType == GemmType::MGroupedMasked;
+            DG_STATIC_ASSERT(kNumMathThreads >= BLOCK_N / TMA_D_BLOCK_N, "Too many TMA blocks");
+            if (threadIdx.x < BLOCK_N / TMA_D_BLOCK_N) {
+                auto in_block_n_offset = threadIdx.x * TMA_D_BLOCK_N;
+                auto smem_ptr = smem_d + in_block_n_offset * BLOCK_M;
+                cute::SM90_TMA_STORE_2D::copy(&tensor_map_d, smem_ptr,
+                                              epilogue_type_t::apply_index_n<TMA_D_BLOCK_N>(n_block_idx * BLOCK_N + in_block_n_offset),
+                                              scheduler.get_global_idx<kWithGroupOffsetD>(shape_m, BLOCK_M, m_block_idx));
+                cute::tma_store_arrive();
+            }
+            __syncwarp();
+        }
+    }
+#else
+    if (blockIdx.x == 0 and threadIdx.x == 0)
+        DG_DEVICE_ASSERT(false and "This kernel only support sm_90a");
+#endif
+}
+
+static cudaLaunchConfig_t construct_launch_config(const cudaStream_t& stream, const int& smem_size,
+                                                  const dim3& grid_dim, const dim3& block_dim, const int& cluster_dim) {
+    
+    cudaLaunchConfig_t config;
+    config.gridDim = grid_dim;
+    config.blockDim = block_dim;
+    config.dynamicSmemBytes = smem_size;
+    config.stream = stream;
+    config.numAttrs = 0;
+    config.attrs = nullptr;
+
+    // NOTES: must use `static` or the `attr` will be deconstructed
+    static cudaLaunchAttribute attr;
+    if (cluster_dim > 1) {
+        attr.id = cudaLaunchAttributeClusterDimension;
+        attr.val.clusterDim = {static_cast<unsigned>(cluster_dim), 1, 1};
+        config.attrs = &attr;
+        config.numAttrs = 1;
+    }
+    return config;
+}
+
+
+// static auto launch_kernel(auto kernel, const cudaLaunchConfig_t& config, float* sfb, float* bias, int* grouped_layout,
+//                         uint32_t shape_m, uint32_t shape_n, uint32_t shape_k,
+//                         const  CUtensorMap  tensor_map_a,
+//                         const  CUtensorMap  tensor_map_b,
+//                         const  CUtensorMap  tensor_map_d,
+//                         const  CUtensorMap  tensor_map_sfa) {
+//     // void* ptr_args[] = {&sfb, &bias, &grouped_layout, &shape_m, &shape_n, &shape_k, &tensor_map_a, &tensor_map_b, &tensor_map_d, &tensor_map_sfa};
+//     return 
+// }
+
+
+template<int N, int K>
+void sm90_fp8_gemm_1d2d_bias_launch(int num_sms, int num_threads, int cluster_dim, int smem_size, cudaStream_t stream, float* sfb, float* bias, int* grouped_layout,
+                        uint32_t shape_m, uint32_t shape_n, uint32_t shape_k,
+                        const  CUtensorMap  tensor_map_a,
+                        const  CUtensorMap  tensor_map_b,
+                        const  CUtensorMap  tensor_map_d,
+                        const  CUtensorMap  tensor_map_sfa){
+    dim3 grid{num_sms, 1, 1};
+    dim3 block{num_threads, 1, 1};
+    const auto config = construct_launch_config(stream, smem_size, grid, block, cluster_dim);
+    if(num_sms == 132){
+        auto kernel = &sm90_fp8_gemm_1d2d_bias_impl<0, N, K, 1, 256, 128, 128, 128, 3, (K / 128) % 3, 128, 256, 2, true, 132, GemmType::Normal, EpilogueIdentity>;
+        cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);
+        cudaLaunchKernelEx(&config, kernel, sfb, bias, grouped_layout, shape_m, shape_n, shape_k, tensor_map_a, tensor_map_b, tensor_map_d, tensor_map_sfa);
+    } else if(num_sms == 116) {
+        auto kernel = &sm90_fp8_gemm_1d2d_bias_impl<0, N, K, 1, 256, 128, 128, 128, 3, (K / 128) % 3, 128, 256, 2, true, 116, GemmType::Normal, EpilogueIdentity>;
+        cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);
+        cudaLaunchKernelEx(&config, kernel, sfb, bias, grouped_layout, shape_m, shape_n, shape_k, tensor_map_a, tensor_map_b, tensor_map_d, tensor_map_sfa);
+    } else if (num_sms == 100) {
+        auto kernel = &sm90_fp8_gemm_1d2d_bias_impl<0, N, K, 1, 256, 128, 128, 128, 3, (K / 128) % 3, 128, 256, 2, true, 100, GemmType::Normal, EpilogueIdentity>;
+        cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);
+        cudaLaunchKernelEx(&config, kernel, sfb, bias, grouped_layout, shape_m, shape_n, shape_k, tensor_map_a, tensor_map_b, tensor_map_d, tensor_map_sfa);
+    } else {
+        // The supported SM counts are exactly the branches above (the only kernels
+        // instantiated). Fail loudly instead of falling through with no launch,
+        // which would leave the output buffer uninitialized.
+        throw std::runtime_error("Unsupported num_sms=" + std::to_string(num_sms)
+                                 + " (blockwise SM90 GEMM is built for 132, 116, and 100 SMs)");
+    }
+
+    // launch_kernel(kernel, config, sfb, bias, grouped_layout, shape_m, shape_n, shape_k, tensor_map_a, tensor_map_b, tensor_map_d, tensor_map_sfa);
+}
+template void sm90_fp8_gemm_1d2d_bias_launch<2048, 2048>(int num_sms, int num_threads, int cluster_dim, int smem_size, cudaStream_t stream, float* sfb, float* bias, int* grouped_layout,
+                        uint32_t shape_m, uint32_t shape_n, uint32_t shape_k,
+                        const  CUtensorMap  tensor_map_a,
+                        const  CUtensorMap  tensor_map_b,
+                        const  CUtensorMap  tensor_map_d,
+                        const  CUtensorMap  tensor_map_sfa);
+template void sm90_fp8_gemm_1d2d_bias_launch<4096, 2048>(int num_sms, int num_threads, int cluster_dim, int smem_size, cudaStream_t stream, float* sfb, float* bias, int* grouped_layout,
+                        uint32_t shape_m, uint32_t shape_n, uint32_t shape_k,
+                        const  CUtensorMap  tensor_map_a,
+                        const  CUtensorMap  tensor_map_b,
+                        const  CUtensorMap  tensor_map_d,
+                        const  CUtensorMap  tensor_map_sfa);
+template void sm90_fp8_gemm_1d2d_bias_launch<8192, 2048>(int num_sms, int num_threads, int cluster_dim, int smem_size, cudaStream_t stream, float* sfb, float* bias, int* grouped_layout,
+                        uint32_t shape_m, uint32_t shape_n, uint32_t shape_k,
+                        const  CUtensorMap  tensor_map_a,
+                        const  CUtensorMap  tensor_map_b,
+                        const  CUtensorMap  tensor_map_d,
+                        const  CUtensorMap  tensor_map_sfa);
+template void sm90_fp8_gemm_1d2d_bias_launch<16384, 2048>(int num_sms, int num_threads, int cluster_dim, int smem_size, cudaStream_t stream, float* sfb, float* bias, int* grouped_layout,
+                        uint32_t shape_m, uint32_t shape_n, uint32_t shape_k,
+                        const  CUtensorMap  tensor_map_a,
+                        const  CUtensorMap  tensor_map_b,
+                        const  CUtensorMap  tensor_map_d,
+                        const  CUtensorMap  tensor_map_sfa);
+template void sm90_fp8_gemm_1d2d_bias_launch<2048, 4096>(int num_sms, int num_threads, int cluster_dim, int smem_size, cudaStream_t stream, float* sfb, float* bias, int* grouped_layout,
+                        uint32_t shape_m, uint32_t shape_n, uint32_t shape_k,
+                        const  CUtensorMap  tensor_map_a,
+                        const  CUtensorMap  tensor_map_b,
+                        const  CUtensorMap  tensor_map_d,
+                        const  CUtensorMap  tensor_map_sfa);
+template void sm90_fp8_gemm_1d2d_bias_launch<4096, 4096>(int num_sms, int num_threads, int cluster_dim, int smem_size, cudaStream_t stream, float* sfb, float* bias, int* grouped_layout,
+                        uint32_t shape_m, uint32_t shape_n, uint32_t shape_k,
+                        const  CUtensorMap  tensor_map_a,
+                        const  CUtensorMap  tensor_map_b,
+                        const  CUtensorMap  tensor_map_d,
+                        const  CUtensorMap  tensor_map_sfa);
+template void sm90_fp8_gemm_1d2d_bias_launch<8192, 4096>(int num_sms, int num_threads, int cluster_dim, int smem_size, cudaStream_t stream, float* sfb, float* bias, int* grouped_layout,
+                        uint32_t shape_m, uint32_t shape_n, uint32_t shape_k,
+                        const  CUtensorMap  tensor_map_a,
+                        const  CUtensorMap  tensor_map_b,
+                        const  CUtensorMap  tensor_map_d,
+                        const  CUtensorMap  tensor_map_sfa);
+template void sm90_fp8_gemm_1d2d_bias_launch<16384, 4096>(int num_sms, int num_threads, int cluster_dim, int smem_size, cudaStream_t stream, float* sfb, float* bias, int* grouped_layout,
+                        uint32_t shape_m, uint32_t shape_n, uint32_t shape_k,
+                        const  CUtensorMap  tensor_map_a,
+                        const  CUtensorMap  tensor_map_b,
+                        const  CUtensorMap  tensor_map_d,
+                        const  CUtensorMap  tensor_map_sfa);
+template void sm90_fp8_gemm_1d2d_bias_launch<2048, 8192>(int num_sms, int num_threads, int cluster_dim, int smem_size, cudaStream_t stream, float* sfb, float* bias, int* grouped_layout,
+                        uint32_t shape_m, uint32_t shape_n, uint32_t shape_k,
+                        const  CUtensorMap  tensor_map_a,
+                        const  CUtensorMap  tensor_map_b,
+                        const  CUtensorMap  tensor_map_d,
+                        const  CUtensorMap  tensor_map_sfa);
+template void sm90_fp8_gemm_1d2d_bias_launch<4096, 8192>(int num_sms, int num_threads, int cluster_dim, int smem_size, cudaStream_t stream, float* sfb, float* bias, int* grouped_layout,
+                        uint32_t shape_m, uint32_t shape_n, uint32_t shape_k,
+                        const  CUtensorMap  tensor_map_a,
+                        const  CUtensorMap  tensor_map_b,
+                        const  CUtensorMap  tensor_map_d,
+                        const  CUtensorMap  tensor_map_sfa);
+template void sm90_fp8_gemm_1d2d_bias_launch<8192, 8192>(int num_sms, int num_threads, int cluster_dim, int smem_size, cudaStream_t stream, float* sfb, float* bias, int* grouped_layout,
+                        uint32_t shape_m, uint32_t shape_n, uint32_t shape_k,
+                        const  CUtensorMap  tensor_map_a,
+                        const  CUtensorMap  tensor_map_b,
+                        const  CUtensorMap  tensor_map_d,
+                        const  CUtensorMap  tensor_map_sfa);
+template void sm90_fp8_gemm_1d2d_bias_launch<16384, 8192>(int num_sms, int num_threads, int cluster_dim, int smem_size, cudaStream_t stream, float* sfb, float* bias, int* grouped_layout,
+                        uint32_t shape_m, uint32_t shape_n, uint32_t shape_k,
+                        const  CUtensorMap  tensor_map_a,
+                        const  CUtensorMap  tensor_map_b,
+                        const  CUtensorMap  tensor_map_d,
+                        const  CUtensorMap  tensor_map_sfa);
+template void sm90_fp8_gemm_1d2d_bias_launch<2048, 16384>(int num_sms, int num_threads, int cluster_dim, int smem_size, cudaStream_t stream, float* sfb, float* bias, int* grouped_layout,
+                        uint32_t shape_m, uint32_t shape_n, uint32_t shape_k,
+                        const  CUtensorMap  tensor_map_a,
+                        const  CUtensorMap  tensor_map_b,
+                        const  CUtensorMap  tensor_map_d,
+                        const  CUtensorMap  tensor_map_sfa);
+template void sm90_fp8_gemm_1d2d_bias_launch<4096, 16384>(int num_sms, int num_threads, int cluster_dim, int smem_size, cudaStream_t stream, float* sfb, float* bias, int* grouped_layout,
+                        uint32_t shape_m, uint32_t shape_n, uint32_t shape_k,
+                        const  CUtensorMap  tensor_map_a,
+                        const  CUtensorMap  tensor_map_b,
+                        const  CUtensorMap  tensor_map_d,
+                        const  CUtensorMap  tensor_map_sfa);
+template void sm90_fp8_gemm_1d2d_bias_launch<8192, 16384>(int num_sms, int num_threads, int cluster_dim, int smem_size, cudaStream_t stream, float* sfb, float* bias, int* grouped_layout,
+                        uint32_t shape_m, uint32_t shape_n, uint32_t shape_k,
+                        const  CUtensorMap  tensor_map_a,
+                        const  CUtensorMap  tensor_map_b,
+                        const  CUtensorMap  tensor_map_d,
+                        const  CUtensorMap  tensor_map_sfa);
+template void sm90_fp8_gemm_1d2d_bias_launch<16384, 16384>(int num_sms, int num_threads, int cluster_dim, int smem_size, cudaStream_t stream, float* sfb, float* bias, int* grouped_layout,
+                        uint32_t shape_m, uint32_t shape_n, uint32_t shape_k,
+                        const  CUtensorMap  tensor_map_a,
+                        const  CUtensorMap  tensor_map_b,
+                        const  CUtensorMap  tensor_map_d,
+                        const  CUtensorMap  tensor_map_sfa);
+};  // namespace deep_gemm

packages/ltx-kernels/csrc/blockwise/kernels/geforce/gemm.cu

@@ -0,0 +1,287 @@
+#include "cutlass/cutlass.h"
+#include "cutlass/layout/layout.h"
+#include <cute/tensor.hpp>
+#include <c10/cuda/CUDAException.h> 
+#include <torch/extension.h>
+#include <torch/python.h>
+#include <cuda_runtime.h>
+#include <iostream>
+
+#include "kernel_traits.cuh"
+#include "static_switch.h"
+
+namespace sm89{
+using namespace cute;
+
+__device__ static void copy_1d(float* gmem_src, float* smem_dst)
+{
+
+    uint32_t smem_int_ptr = cast_smem_ptr_to_uint((void*)smem_dst);
+    asm volatile("cp.async.ca.shared.global.L2::128B [%0], [%1], %2;\n"
+        :: "r"(smem_int_ptr),
+           "l"(gmem_src),
+           "n"(sizeof(float)));
+}
+
+template <typename KernelTraits=gemm_traits<128, 256, 2, 4096, 2, 4, true, half_t, bfloat16_t>>
+__global__ void gemm_fp8_kernel(float_e4m3_t* Aptr, float* sfa, float_e4m3_t* Bptr, float* sfb, float* bias_ptr, void* out, int M, int N, int K, int TMA_ALIGNED_M){
+    using output_t = typename KernelTraits::out_t;
+    using SmemLayoutA = typename KernelTraits::SmemLayoutA;
+    using SmemLayoutB = typename KernelTraits::SmemLayoutB;
+    using SmemLayoutC = typename KernelTraits::SmemLayoutC;
+    
+    constexpr int BM = KernelTraits::BM;
+    constexpr int BN = KernelTraits::BN;
+    constexpr int BK = KernelTraits::BK;
+    constexpr int Ksfa = KernelTraits::KSF;
+    constexpr bool has_bias = KernelTraits::HasBias;
+    extern __shared__ float smem_[];
+    float *bias_shm = smem_;
+    float *sfa_shm = reinterpret_cast<float*>(bias_shm + cosize(typename KernelTraits::SmemLayoutBias{}));
+    
+    output_t* C_shm = reinterpret_cast<output_t*>(sfa_shm + cosize(typename KernelTraits::SmemLayoutSFA{}));
+    float_e4m3_t* A_shm = reinterpret_cast<float_e4m3_t*>(sfa_shm + cosize(typename KernelTraits::SmemLayoutSFA{}));
+    float_e4m3_t* B_shm = reinterpret_cast<float_e4m3_t*>(A_shm + cosize(SmemLayoutA{}));
+  
+    int idx = threadIdx.x;
+    int ix = blockIdx.x;
+    int iy = blockIdx.y;
+    // sfa += BM * iy;
+    sfb += KernelTraits::NUM_SFB_PER_STEP * ix * Ksfa;
+    
+
+    output_t* Cptr = reinterpret_cast<output_t*>(out);
+
+    Tensor A = make_tensor(make_gmem_ptr(Aptr), make_shape(M, K), make_stride(K, Int<1>{}));
+    Tensor B = make_tensor(make_gmem_ptr(Bptr), make_shape(N, K), make_stride(K, Int<1>{}));
+    Tensor D = make_tensor(make_gmem_ptr(Cptr), make_shape(M, N), make_stride(N, Int<1>{}));
+    Tensor SFA = make_tensor(make_gmem_ptr(sfa), make_shape(M, Ksfa), make_stride(Int<1>{}, TMA_ALIGNED_M));
+    
+    Tensor gA = local_tile(A, make_tile(Int<BM>{}, Int<BK>{}), make_coord(iy, _)); 
+    Tensor gB = local_tile(B, make_tile(Int<BN>{}, Int<BK>{}), make_coord(ix, _)); 
+    Tensor gD = local_tile(D, make_tile(Int<BM>{}, Int<BN>{}), make_coord(iy, ix)); 
+    Tensor gSFA = local_tile(SFA, make_tile(Int<BM>{}, Int<1>{}), make_coord(iy, _));
+
+    auto sBias = make_tensor(make_smem_ptr(bias_shm), typename KernelTraits::SmemLayoutBias{});
+    if constexpr (has_bias){
+        Tensor Bias = make_tensor(make_gmem_ptr(bias_ptr), make_shape(_1{}, N), make_stride(N, Int<1>{}));
+        Tensor gBias = local_tile(Bias, make_tile(Int<1>{}, Int<BN>{}), make_coord(_, ix));
+        typename KernelTraits::G2SBiasCopy g2s_bias_copy;
+        auto g2s_bias_thr_copy = g2s_bias_copy.get_slice(idx);
+        auto tCBiasgBias = g2s_bias_thr_copy.partition_S(gBias);
+        auto tCBiassBias = g2s_bias_thr_copy.partition_D(sBias);
+        if(idx < BN){
+            copy_1d((float*)&gBias(0) + idx, (float*)&sBias(0) + idx);
+        }
+    }
+
+    auto sSFA = make_tensor(make_smem_ptr(sfa_shm), typename KernelTraits::SmemLayoutSFA{});
+    auto sA = make_tensor(make_smem_ptr(A_shm), SmemLayoutA{}); 
+    auto sB = make_tensor(make_smem_ptr(B_shm), SmemLayoutB{});
+
+    typename KernelTraits::MMATile tiled_mma;
+    auto thr_mma = tiled_mma.get_slice(threadIdx.x);
+
+    auto tCrA = thr_mma.partition_fragment_A(gA(_, _, 0));  
+    auto tCrB = thr_mma.partition_fragment_B(gB(_, _, 0)); 
+    auto tCrD = thr_mma.partition_fragment_C(gD);           
+    clear(tCrD);
+    auto tCrD_fp32 = make_tensor_like<float>(tCrD);
+    clear(tCrD_fp32);
+
+    typename KernelTraits::G2STiledCopy g2s_tiled_copy;
+    auto g2s_thr_copy = g2s_tiled_copy.get_slice(idx);
+    auto tAgA_copy = g2s_thr_copy.partition_S(gA); 
+    auto tAsA_copy = g2s_thr_copy.partition_D(sA); 
+    auto tBgB_copy = g2s_thr_copy.partition_S(gB); 
+    auto tBsB_copy = g2s_thr_copy.partition_D(sB); 
+
+    auto s2r_tiled_copy_a = make_tiled_copy_A(typename KernelTraits::S2RCopyAtomA{}, tiled_mma);
+    auto s2r_thr_copy_a = s2r_tiled_copy_a.get_slice(idx);
+    auto tAsA = s2r_thr_copy_a.partition_S(sA);     
+    auto tCrA_view = s2r_thr_copy_a.retile_D(tCrA); 
+
+
+    auto s2r_tiled_copy_b = make_tiled_copy_B(typename KernelTraits::S2RCopyAtomB{}, tiled_mma);
+    auto s2r_thr_copy_b = s2r_tiled_copy_b.get_slice(idx);
+    auto tBsB = s2r_thr_copy_b.partition_S(sB);    
+    auto tCrB_view = s2r_thr_copy_b.retile_D(tCrB);
+
+    auto cA = make_identity_tensor(make_shape(size<0>(sA), size<1>(sA)));
+    auto tAcA = g2s_thr_copy.partition_S(cA);
+    int residual = M - iy*BM;
+
+    int itile_to_read = 0;
+    int ismem_read = 0;
+    int ismem_write = 0;
+    int ismem_read_sfa = 0;
+    constexpr int kStages = KernelTraits::KStages;
+
+    #pragma unroll
+    for(int istage=0; istage<kStages - 1; ++istage){
+        for (size_t m = 0; m < size<1>(tAsA_copy); m++)
+        {
+            for (size_t k = 0; k < size<2>(tAsA_copy); k++)
+            {
+                if(get<0>(tAcA(0, m, k)) < residual){
+                    cute::copy(g2s_tiled_copy, tAgA_copy(_, m, k, istage), tAsA_copy(_, m, k, istage)); 
+                }
+            }  
+        }
+        if(idx < KernelTraits::THREADS_SFA_COPY && (BM * iy + idx * KernelTraits::SFA_ELEMS_PER_COPY < M)) {
+            copy_1d((float*)&gSFA(0, 0, istage) + idx*KernelTraits::SFA_ELEMS_PER_COPY,  (float*)&sSFA(0, istage) + idx*KernelTraits::SFA_ELEMS_PER_COPY);
+        }
+        cute::copy(g2s_tiled_copy, tBgB_copy(_, _, _, istage), tBsB_copy(_, _, _, istage));
+        cp_async_fence();
+        ++itile_to_read;
+        ++ismem_write;
+    }
+
+    cp_async_wait<kStages - 2>();
+    __syncthreads();
+
+    cute::copy(s2r_tiled_copy_a, tAsA(_, _, 0, ismem_read), tCrA_view(_, _, 0)); 
+    cute::copy(s2r_tiled_copy_b, tBsB(_, _, 0, ismem_read), tCrB_view(_, _, 0));
+
+    static constexpr int nk = size<2>(tCrA);
+    auto sfa_tv = typename KernelTraits::SFAThreadLayout{};
+    static constexpr int NTILES = KernelTraits::NTiles;
+    #pragma unroll
+    for(int itile = 0; itile < NTILES; itile++){
+        clear(tCrD);
+        #pragma unroll
+        for(int ik = 0; ik < nk; ik++){
+            int ik_next = (ik + 1) % nk;
+            if(ik == nk - 1) {
+                cp_async_wait<kStages - 2>();
+                __syncthreads();
+                ismem_read = (ismem_read + 1) % kStages;
+            }
+            cute::copy(s2r_tiled_copy_a, tAsA(_, _, ik_next, ismem_read), tCrA_view(_, _, ik_next));
+            cute::copy(s2r_tiled_copy_b, tBsB(_, _, ik_next, ismem_read), tCrB_view(_, _, ik_next));
+            if(ik == 0){
+                if(itile_to_read < NTILES){
+                    for (size_t m = 0; m < size<1>(tAsA_copy); m++)
+                    {
+                        for (size_t k = 0; k < size<2>(tAsA_copy); k++)
+                        {
+                            if(get<0>(tAcA(0, m, k)) < residual){
+                                cute::copy(g2s_tiled_copy, tAgA_copy(_, m, k, itile_to_read), tAsA_copy(_, m, k, ismem_write));
+                            }
+                        }
+                    }
+                    cute::copy(g2s_tiled_copy, tBgB_copy(_, _, _, itile_to_read), tBsB_copy(_, _, _, ismem_write));
+                    if(idx < KernelTraits::THREADS_SFA_COPY && (BM * iy + idx * KernelTraits::SFA_ELEMS_PER_COPY < M)) {
+                        copy_1d((float*)&gSFA(0, 0, itile_to_read) + idx * KernelTraits::SFA_ELEMS_PER_COPY,  (float*)&sSFA(0, ismem_write) + idx*KernelTraits::SFA_ELEMS_PER_COPY);
+                    }
+                    ++itile_to_read;
+                    ismem_write = (ismem_write + 1) % kStages;
+                }
+                cp_async_fence();
+            }
+            cute::gemm(tiled_mma, tCrD, tCrA(_, _, ik), tCrB(_, _, ik), tCrD);
+        }
+       
+        int sf_ind = itile / KernelTraits::TILES_PER_BLOCK;
+        float sfb_val = sfb[sf_ind];
+        #pragma unroll
+        for(int i = 0; i < size<1>(tCrD); i++){ // (MMA, MMA_M, MMA_N) = (4, 4, 4)
+            float sfa_val_1 = sSFA(sfa_tv(idx) + i * KernelTraits::MMA_WARP_M, ismem_read_sfa);
+            float sfa_val_2  = sSFA(sfa_tv(idx) + 8 + i * KernelTraits::MMA_WARP_M, ismem_read_sfa);
+            #pragma unroll
+            for(int j = 0; j < size<2>(tCrD); j++){
+                tCrD_fp32(0, i, j) += sfa_val_1 * sfb_val * float(tCrD(0, i, j));
+                tCrD_fp32(1, i, j) += sfa_val_1 * sfb_val * float(tCrD(1, i, j));
+                tCrD_fp32(2, i, j) += sfa_val_2 * sfb_val * float(tCrD(2, i, j));
+                tCrD_fp32(3, i, j) += sfa_val_2 * sfb_val * float(tCrD(3, i, j));
+            }
+        } 
+        ismem_read_sfa = (ismem_read_sfa + 1) % kStages;
+        __syncthreads();
+    }
+
+    auto tCrBias = make_tensor<float>(Layout<Shape<_2, Int<size<2>(tCrD_fp32)>>>{});
+    auto bias_threads = typename KernelTraits::BiasThreadLayout{};
+    if constexpr (has_bias){
+        #pragma unroll
+        for(int i = 0; i<size<2>(tCrD_fp32); i++){
+            tCrBias(0, i) = sBias(bias_threads(idx) + i * KernelTraits::MMA_WARP_N);
+            tCrBias(1, i) = sBias(1 + bias_threads(idx) + i * KernelTraits::MMA_WARP_N);
+        }
+        #pragma unroll
+        for(int i = 0; i<size<1>(tCrD_fp32); i++){
+            #pragma unroll
+            for (int j = 0; j < size<2>(tCrD_fp32) ; j++)
+            {
+                tCrD_fp32(0, i, j) += tCrBias(0, j);
+                tCrD_fp32(1, i, j) += tCrBias(1, j);
+                tCrD_fp32(2, i, j) += tCrBias(0, j);
+                tCrD_fp32(3, i, j) += tCrBias(1, j);
+            }
+        }
+    }
+
+    auto sC = make_tensor(make_smem_ptr(C_shm), SmemLayoutC{});
+    auto r2s_tiled_copy_c = make_tiled_copy_C(typename KernelTraits::R2SCopyAtomC{}, tiled_mma);
+    auto r2s_thr_copy_c = r2s_tiled_copy_c.get_slice(idx);
+    auto tCrC_r2s = r2s_thr_copy_c.retile_S(tCrD_fp32);  
+    auto tCsC_r2s = r2s_thr_copy_c.partition_D(sC);  
+
+    typename KernelTraits::S2GCopyC s2g_tiled_copy_c;
+    auto s2g_thr_copy_c = s2g_tiled_copy_c.get_thread_slice(idx);
+    auto tCsC_s2g = s2g_thr_copy_c.partition_S(sC);  
+    auto tCgC_s2g = s2g_thr_copy_c.partition_D(gD);  
+
+    int pipe = size<2>(tCsC_r2s);
+
+    auto cC = make_identity_tensor(make_shape(size<0>(gD), size<1>(gD)));
+    auto tCcC = s2g_thr_copy_c.partition_D(cC);
+
+    for(int i = 0; i< size<1>(tCrC_r2s); i++){
+        for(int j = 0; j < size<2>(tCrC_r2s); j+=pipe){
+            for(int step = 0; step < pipe; ++step){
+                auto fragment = make_tensor_like<output_t>(tCrC_r2s(_, i, j + step));
+                cute::copy(tCrC_r2s(_, i, j + step), fragment);
+                cute::copy(r2s_tiled_copy_c, fragment, tCsC_r2s(_, 0, step)); 
+            }
+            __syncthreads();
+            if (get<0>(tCcC(0, i, j / pipe)) < residual){
+                cute::copy(s2g_tiled_copy_c, tCsC_s2g(_, 0, 0), tCgC_s2g(_, i, j / pipe));                
+            }
+            __syncthreads();
+        }
+    }
+}
+
+template <bool has_bias, typename accum_type>
+void fp8_kernel_launch(void* Aptr, void* sfa, void* Bptr, void* sfb, void* bias_ptr, void* out, int M, int N, int K, cudaStream_t stream) {
+    int TMA_ALIGNED_M = ((M + sizeof(float) - 1) / sizeof(float)) * sizeof(float); // SIZEOF(float) = 4 
+    BLOCK_K_SWITCH(K_, M_SWITCH( 
+        using KernelTraits = gemm_traits<BM, BN, 3, K_, WARP_ROW, WARP_COL, has_bias, accum_type, bfloat16_t>;
+        auto kernel = &gemm_fp8_kernel<KernelTraits>;
+        int BX = (N + KernelTraits::BN - 1) / KernelTraits::BN;
+        int BY = (M + KernelTraits::BM - 1) / KernelTraits::BM;
+        dim3 block(KernelTraits::NUM_THREADS);
+        dim3 gridDim(BX, BY);
+        cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, KernelTraits::SmemSize);
+        kernel<<<gridDim, KernelTraits::NUM_THREADS, KernelTraits::SmemSize, stream>>>((float_e4m3_t*)Aptr, (float*)sfa, (float_e4m3_t*)Bptr, (float*)sfb, (float*)bias_ptr, out, M, N, K, TMA_ALIGNED_M);
+        C10_CUDA_KERNEL_LAUNCH_CHECK();))
+}
+
+template<bool use_fast_accum>
+void fp8_bias_gemm_cuda(void* Aptr, void* SFA, void* Bptr, void* SFB, void* bias_ptr, void* out, int M, int N, int K, cudaStream_t stream){
+    using accum_type = std::conditional_t<use_fast_accum, half_t, float>;
+    fp8_kernel_launch<true, accum_type>(Aptr, SFA, Bptr, SFB, bias_ptr, out, M, N, K, stream);
+}
+// template<bool use_fast_accum>
+// void fp8_gemm_cuda(void* Aptr, void* SFA, void* Bptr, void* SFB, void* out, int M, int N, int K, cudaStream_t stream){
+//     using accum_type = std::conditional_t<use_fast_accum, half_t, float>;
+//     BLOCK_K_SWITCH(num_acc_upcast_steps, fp8_kernel_launch<false, num_acc_upcast_steps, accum_type>(Aptr, SFA, Bptr, SFB, nullptr, out, M, N, K, stream);) 
+// }
+
+// template void fp8_gemm_cuda<true>(void* Aptr, void* SFA, void* Bptr, void* SFB, void* out, int M, int N, int K, cudaStream_t stream);
+// template void fp8_gemm_cuda<false>(void* Aptr, void* SFA, void* Bptr, void* SFB, void* out, int M, int N, int K, cudaStream_t stream);
+
+template void fp8_bias_gemm_cuda<true>(void* Aptr, void* SFA, void* Bptr, void* SFB, void* bias_ptr, void* out, int M, int N, int K, cudaStream_t stream);
+template void fp8_bias_gemm_cuda<false>(void* Aptr, void* SFA, void* Bptr, void* SFB, void* bias_ptr, void* out, int M, int N, int K, cudaStream_t stream);
+}; // namespace sm89

packages/ltx-kernels/csrc/blockwise/kernels/geforce/kernel_traits.cuh

@@ -0,0 +1,121 @@
+#pragma once
+
+#include <cute/tensor.hpp>
+
+#include <cutlass/cutlass.h>
+#include <cutlass/layout/layout.h>
+#include <cutlass/numeric_types.h>
+
+#include "mma_sm89_fp16.hpp"
+#include "mma_traits_sm89_fp16.hpp"
+
+using namespace cute;
+
+template<int BYTES> struct BytesToType {};
+template<> struct BytesToType<4> {
+    using Type = uint32_t;
+    static_assert(sizeof(Type) == 4);
+};
+template<> struct BytesToType<2> {
+    using Type = uint16_t;
+    static_assert(sizeof(Type) == 2);
+};
+
+template<int BM_, int BN_, int KStages_, int K_, int WARP_ROW_=2, int WARP_COL_=2, bool HasBias_=false, typename accum_t_=cutlass::half_t, typename out_t_=cutlass::bfloat16_t>
+struct gemm_traits {
+    static constexpr int BLOCK_SIZE = 128;
+    static constexpr int K = K_;
+    static constexpr int BM = BM_;
+    static constexpr int BN = BN_;
+    static constexpr int BK = 128;
+    static constexpr int TILES_PER_BLOCK = BLOCK_SIZE / BK;
+    static constexpr int NUM_SFB_PER_STEP = BN / BLOCK_SIZE;
+    static constexpr int NTiles = K / BK;
+    static constexpr int KSF = K / BLOCK_SIZE; 
+    static constexpr int KStages = KStages_;
+    static constexpr int WARP_ROW = WARP_ROW_;
+    static constexpr int WARP_COL = WARP_COL_;
+    static constexpr int NUM_WARPS = WARP_ROW * WARP_COL;
+    static constexpr int NUM_THREADS = NUM_WARPS * 32;
+    static constexpr int MMA_WARP_M = WARP_ROW * 16;
+    static constexpr int MMA_WARP_N = WARP_COL * 8;
+    static constexpr int MMA_WARP_K = 32;
+    using accum_t = accum_t_;
+    using out_t = out_t_;
+    using SwizzleLayoutO = std::conditional_t<
+        std::is_same_v<out_t_, cutlass::bfloat16_t>,
+        Swizzle<3, 3, 3>,
+        Swizzle<2, 4, 3>
+    >;
+    using SwizzleLayoutAB = Swizzle<2, 4, 3>;
+    using MMA_Atom_SM89 = std::conditional_t<
+        std::is_same_v<accum_t, cutlass::half_t>,
+        MMA_Atom<SM89_16x8x32_F16E4M3E4M3F16_TN>,
+        MMA_Atom<SM89_16x8x32_F32E4M3E4M3F32_TN>
+    >;
+    static constexpr int INPUT_ELEMS_PER_COPY = sizeof(uint128_t) / sizeof(float_e4m3_t);
+    static constexpr int OUTPUT_ELEMS_PER_COPY = sizeof(uint128_t) / sizeof(out_t_); 
+    static constexpr int THREADS_PER_ROW = BK / INPUT_ELEMS_PER_COPY;
+    using GMEMLayout = Layout< Shape <Int<NUM_THREADS / THREADS_PER_ROW>, Int<THREADS_PER_ROW>>, Stride<Int<THREADS_PER_ROW>, _1>>;
+    using G2SCopyAtom = Copy_Atom<SM80_CP_ASYNC_CACHEGLOBAL<cute::uint128_t>, float_e4m3_t>;
+    using G2STiledCopy = decltype(
+        make_tiled_copy(
+            G2SCopyAtom{},
+            GMEMLayout{},
+            Layout<Shape<_1, Int<INPUT_ELEMS_PER_COPY>>>{}
+        )
+    );
+    using S2RCopyAtomA = Copy_Atom<SM75_U32x4_LDSM_N, float_e4m3_t>;
+    using S2RCopyAtomB = Copy_Atom<SM75_U32x2_LDSM_N, float_e4m3_t>;
+    using SmemLayoutAtom = decltype(composition(
+        Swizzle<2, 4, 3>{},
+        make_layout(make_shape(Int<8>{}, Int<BK>{}),
+                    make_stride(Int<BK>{}, Int<1>{}))));
+    using SmemLayoutA = decltype(
+        tile_to_shape(SmemLayoutAtom{}, make_shape(Int<BM>{}, Int<BK>{}, Int<KStages>{}))
+    );
+    using SmemLayoutB = decltype(
+        tile_to_shape(SmemLayoutAtom{}, make_shape(Int<BN>{}, Int<BK>{}, Int<KStages>{}))
+    );
+    using MMATile = decltype(
+        make_tiled_mma(
+            MMA_Atom_SM89{},
+            Layout<Shape<Int<WARP_ROW>, Int<WARP_COL>, _1>>{},
+            Tile<Int<MMA_WARP_M>, Int<MMA_WARP_N>, Int<MMA_WARP_K>>{}
+        )
+    );
+
+    static constexpr int ELEMS_PER_TILE = MMA_WARP_M * MMA_WARP_N;
+    static constexpr int NUM_ELEMS_PER_WRITE = NUM_THREADS * sizeof(cute::uint128_t) / sizeof(out_t_);
+    static constexpr int OUT_PIPE = NUM_ELEMS_PER_WRITE / ELEMS_PER_TILE;
+    // using SmemLayoutC = Layout<Shape<Int<BM>, Int<BN>>, Stride<Int<BN>, Int<1>>>;
+
+    using SmemLayoutC = decltype(
+            make_layout( 
+                make_shape(Int<MMA_WARP_M>{}, Int<MMA_WARP_N*OUT_PIPE>{}), 
+                make_stride(Int<MMA_WARP_N*OUT_PIPE>{}, Int<1>{})
+            )
+    );  
+    static constexpr int THREADS_PER_ROW_WRITE = MMA_WARP_N * OUT_PIPE / OUTPUT_ELEMS_PER_COPY;
+    using R2SCopyAtomC = Copy_Atom<UniversalCopy<typename BytesToType<2*sizeof(out_t)>::Type>, out_t>;
+    using S2GCopyAtomC = Copy_Atom<UniversalCopy<cute::uint128_t>, out_t>;
+    using S2GCopyC = decltype(make_tiled_copy(S2GCopyAtomC{},
+                            make_layout(make_shape(Int<NUM_THREADS / THREADS_PER_ROW_WRITE>{}, Int<THREADS_PER_ROW_WRITE>{}),
+                                        make_stride(Int<THREADS_PER_ROW_WRITE>{}, Int<1>{})),
+                            make_layout(make_shape(Int<1>{}, Int<OUTPUT_ELEMS_PER_COPY>{}))));
+    
+   using G2SBiasCopyAtom = Copy_Atom<SM80_CP_ASYNC_CACHEALWAYS<float>, float>;
+   using G2SBiasCopy = decltype(make_tiled_copy(G2SBiasCopyAtom{}, make_layout(
+                    make_shape(Int<1>{},Int<BN>{}), make_stride(Int<BN>{}, Int<1>{})),
+                    make_layout(make_shape(Int<1>{},Int<1>{}), make_stride(Int<1>{}, Int<1>{}))));
+    using sfa_copy_vtype = float;
+    static constexpr int SFA_ELEMS_PER_COPY = sizeof(sfa_copy_vtype)/sizeof(float);
+    static constexpr int THREADS_SFA_COPY = BM * sizeof(float) / sizeof(sfa_copy_vtype); 
+    // using G2SSFACopyAtom = Copy_Atom<SM80_CP_ASYNC_CACHEALWAYS<cute::uint128_t>, float>;
+    static constexpr bool HasBias = HasBias_;
+    using SmemLayoutBias = Layout<Shape<Int<1>, Int<BN>>, Stride<Int<BN>, Int<1>>>;
+    using SmemLayoutSFA = Layout<Shape<Int<BM>, Int<KStages>>, Stride<Int<1>, Int<BM>>>;
+    using BiasThreadLayout = Layout<Shape<Shape<_4, _8>, Shape<Int<WARP_ROW>, Int<WARP_COL>>>, Stride<Stride<_2, _0>, Stride<_0, _8>>>;
+    using SFAThreadLayout = Layout<Shape<Shape<_4, _8>, Shape<Int<WARP_ROW>, Int<WARP_COL>>>, Stride<Stride<_0, _1>, Stride<_16, _0>>>;
+    static constexpr int SmemSize = cute::max(cute::cosize(SmemLayoutA{})+cute::cosize(SmemLayoutB{}), cute::cosize(SmemLayoutC{})*sizeof(out_t)) + cute::cosize(SmemLayoutBias{}) * sizeof(float) + cute::cosize(SmemLayoutSFA{})*sizeof(float);
+};

packages/ltx-kernels/csrc/blockwise/kernels/geforce/mma_sm89_fp16.hpp

@@ -0,0 +1,84 @@
+#pragma once
+
+#include <cute/config.hpp>
+#include <cute/arch/mma.hpp>
+////////////////////////////////////////////////////////////////////////////////
+
+#if (__CUDACC_VER_MAJOR__ > 12) || (__CUDACC_VER_MAJOR__ == 12 && __CUDACC_VER_MINOR__ >= 4)
+#  define CUTE_ARCH_MMA_F32_SM89_SUPPORTED
+#endif
+
+#if (__CUDACC_VER_MAJOR__ > 12) || (__CUDACC_VER_MAJOR__ == 12 && __CUDACC_VER_MINOR__ >= 8)
+#  define CUTE_ARCH_MMA_F16_SM89_SUPPORTED
+#endif
+
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 890)
+#  if defined(CUTE_ARCH_MMA_F32_SM89_SUPPORTED)
+#    define CUTE_ARCH_MMA_F32_SM89_ENABLED
+#  endif
+
+#  if defined(CUTE_ARCH_MMA_F16_SM89_SUPPORTED)
+#    define CUTE_ARCH_MMA_F16_SM89_ENABLED
+#  endif
+#endif
+
+namespace cute {
+struct SM89_16x8x32_F32E4M3E4M3F32_TN
+{
+  using DRegisters = float[4];
+  using ARegisters = uint32_t[4];
+  using BRegisters = uint32_t[2];
+  using CRegisters = float[4];
+
+  CUTE_HOST_DEVICE static void
+  fma(float      & d0, float      & d1, float      & d2, float      & d3,
+      uint32_t const& a0, uint32_t const& a1, uint32_t const& a2, uint32_t const& a3,
+      uint32_t const& b0, uint32_t const& b1,
+      float const& c0, float const& c1, float const& c2, float const& c3)
+  {
+#if defined(CUTE_ARCH_MMA_F32_SM89_ENABLED)
+    asm(
+      "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 "
+      "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n"
+      : "=f"(d0), "=f"(d1), "=f"(d2), "=f"(d3)
+      :
+        "r"(a0), "r"(a1), "r"(a2), "r"(a3),
+        "r"(b0), "r"(b1),
+        "f"(c0), "f"(c1), "f"(c2), "f"(c3)
+  );
+#else
+    CUTE_INVALID_CONTROL_PATH("Attempting to use SM89_16x8x32_F32E4M3E4M3F32_TN without CUTE_ARCH_MMA_F32_SM89_ENABLED");
+#endif
+  }
+};
+
+  // MMA 16x8x32 TN
+struct SM89_16x8x32_F16E4M3E4M3F16_TN
+{
+  using DRegisters = uint32_t[2];
+  using ARegisters = uint32_t[4];
+  using BRegisters = uint32_t[2];
+  using CRegisters = uint32_t[2];
+
+  CUTE_HOST_DEVICE static void
+  fma(uint32_t      & d0, uint32_t      & d1,
+      uint32_t const& a0, uint32_t const& a1, uint32_t const& a2, uint32_t const& a3,
+      uint32_t const& b0, uint32_t const& b1,
+      uint32_t const& c0, uint32_t const& c1)
+  {
+#if defined(CUTE_ARCH_MMA_F16_SM89_ENABLED)
+    asm(
+      "mma.sync.aligned.m16n8k32.row.col.f16.e4m3.e4m3.f16 "
+      "{%0,%1}, {%2,%3,%4,%5}, {%6,%7}, {%8,%9};\n"
+      : "=r"(d0), "=r"(d1)
+      :
+        "r"(a0), "r"(a1), "r"(a2), "r"(a3),
+        "r"(b0), "r"(b1),
+        "r"(c0), "r"(c1)
+  );
+#else
+    CUTE_INVALID_CONTROL_PATH("Attempting to use SM89_16x8x32_F32E4M3E4M3F32_TN without CUTE_ARCH_MMA_F16_SM89_ENABLED");
+#endif
+  }
+};
+}

packages/ltx-kernels/csrc/blockwise/kernels/geforce/mma_traits_sm89_fp16.hpp

@@ -0,0 +1,50 @@
+#pragma once
+
+#include <cute/atom/mma_traits.hpp>
+#include <cute/layout.hpp>
+#include <cute/numeric/numeric_types.hpp>
+#include "mma_sm89_fp16.hpp"
+
+namespace cute 
+{
+
+namespace {
+
+// (T32,V4) -> (M16,N8)
+using SM80_16x8_Row = Layout<Shape <Shape < _4,_8>,Shape < _2,_2>>,
+                             Stride<Stride<_32,_1>,Stride<_16,_8>>>;
+
+}
+
+template <>
+struct MMA_Traits<SM89_16x8x32_F32E4M3E4M3F32_TN> {
+  using ValTypeD = float;
+  using ValTypeA = float_e4m3_t;
+  using ValTypeB = float_e4m3_t;
+  using ValTypeC = float;
+
+  using Shape_MNK = Shape<_16,_8,_32>;
+  using ThrID   = Layout<_32>;
+  using ALayout = Layout<Shape <Shape < _4,_8>,Shape < _4,_2,  _2>>,
+                         Stride<Stride<_64,_1>,Stride<_16,_8,_256>>>;
+  using BLayout = Layout<Shape <Shape < _4,_8>,Shape <_4,  _2>>,
+                         Stride<Stride<_32,_1>,Stride<_8,_128>>>;
+  using CLayout = SM80_16x8_Row;
+};
+
+template <>
+struct MMA_Traits<SM89_16x8x32_F16E4M3E4M3F16_TN> {
+  using ValTypeD = half_t;
+  using ValTypeA = float_e4m3_t;
+  using ValTypeB = float_e4m3_t;
+  using ValTypeC = half_t;
+
+  using Shape_MNK = Shape<_16,_8,_32>;
+  using ThrID   = Layout<_32>;
+  using ALayout = Layout<Shape <Shape < _4,_8>,Shape < _4,_2,  _2>>,
+                         Stride<Stride<_64,_1>,Stride<_16,_8,_256>>>;
+  using BLayout = Layout<Shape <Shape < _4,_8>,Shape <_4,  _2>>,
+                         Stride<Stride<_32,_1>,Stride<_8,_128>>>;
+  using CLayout = SM80_16x8_Row;
+};
+}

packages/ltx-kernels/csrc/blockwise/kernels/geforce/static_switch.h

@@ -0,0 +1,35 @@
+#pragma once
+
+#define BOOL_SWITCH(COND, CONST_NAME, ...)      \
+    if (COND) {                                 \
+      constexpr static bool CONST_NAME = true;  \
+      __VA_ARGS__                               \
+    } else {                                    \
+      constexpr static bool CONST_NAME = false; \
+      __VA_ARGS__                               \
+    }                                           
+//K/128
+#define BLOCK_K_SWITCH(COSNT_NAME, ...)             \
+    if (K == 2048) {                          \
+      constexpr static int COSNT_NAME = 2048;   \
+      __VA_ARGS__                             \
+    }                                         \
+    else if (K == 4096) {                     \
+      constexpr static int COSNT_NAME = 4096;    \
+      __VA_ARGS__                               \
+    } else if (K == 8192) {                      \
+      constexpr static int COSNT_NAME = 8192;   \
+      __VA_ARGS__                               \
+    } else if (K == 16384) {                      \
+      constexpr static int COSNT_NAME = 16384;    \
+      __VA_ARGS__                               \
+    } else {                                    \
+      TORCH_CHECK(false, "Unsupported K value: ", K);   \
+    }
+
+#define M_SWITCH(...) \
+      constexpr static int BM = 64;    \
+      constexpr static int BN = 128;    \
+      constexpr static int WARP_ROW = 2;    \
+      constexpr static int WARP_COL = 4;    \
+    __VA_ARGS__                                                          

packages/ltx-kernels/csrc/blockwise/runtime_utils.hpp

@@ -0,0 +1,206 @@
+#pragma once
+
+#include <cuda.h>
+#include <cuda_runtime.h>
+
+#include <torch/python.h>
+#include "exceptions.hpp"
+
+namespace blockwise {
+template <typename T>
+static T ceil_div(const T& a, const T& b) {
+    return (a + b - 1) / b;
+}
+template <typename T>
+static constexpr T align(const T& a, const T& b) {
+    return ceil_div(a, b) * b;
+}
+
+static int get_tma_aligned_size(const int& x, const int& element_size) {
+    constexpr int kNumTMAAlignmentBytes = 16;
+    DG_HOST_ASSERT(kNumTMAAlignmentBytes % element_size == 0);
+    return align(x, kNumTMAAlignmentBytes / element_size);
+}
+static std::pair<int, int> get_inner_outer_dims(const cute::UMMA::Major& major, const int& k, const int& mn) {
+    return major == cute::UMMA::Major::K ? std::make_pair(k, mn) : std::make_pair(mn, k);
+}
+
+static int get_non_contiguous_dim(const cute::UMMA::Major& major) {
+    return major == cute::UMMA::Major::K ? -2 : -1;
+}
+
+static int get_compiled_dim(const int& dim, const char& name, const std::string& compiled_dims) {
+    for (const char& c: compiled_dims) {
+        if (name == c)
+            return dim;
+    }
+    return 0;
+}
+static CUtensorMapDataType aten_dtype_to_tensor_map_dtype(const at::ScalarType& dtype,
+                                                          const bool& allow_tf32) {
+    if (allow_tf32 and dtype == torch::kFloat)
+        return CU_TENSOR_MAP_DATA_TYPE_TFLOAT32;
+
+    switch (dtype) {
+        case torch::kInt:           return CU_TENSOR_MAP_DATA_TYPE_INT32;
+        case torch::kFloat:         return CU_TENSOR_MAP_DATA_TYPE_FLOAT32;
+        case torch::kBFloat16:      return CU_TENSOR_MAP_DATA_TYPE_BFLOAT16;
+        case torch::kFloat8_e4m3fn: return CU_TENSOR_MAP_DATA_TYPE_UINT8;
+        default: DG_HOST_UNREACHABLE("Unsupported dtype");
+    }
+}
+
+static CUtensorMapSwizzle mode_into_tensor_map_swizzle(const int& mode, const int& base) {
+#if CUDA_VERSION >= 12080
+    if (base != 0) {
+        DG_HOST_ASSERT(base == 32 and mode == 128);
+        return CU_TENSOR_MAP_SWIZZLE_128B_ATOM_32B;
+    }
+#endif
+
+    DG_HOST_ASSERT(base == 0);
+    switch (mode) {
+        case   0:
+        case  16: return CU_TENSOR_MAP_SWIZZLE_NONE;
+        case  32: return CU_TENSOR_MAP_SWIZZLE_32B;
+        case  64: return CU_TENSOR_MAP_SWIZZLE_64B;
+        case 128: return CU_TENSOR_MAP_SWIZZLE_128B;
+        default: DG_HOST_UNREACHABLE("Unsupported swizzling mode");
+    }
+}
+
+static CUtensorMap make_tma_2d_desc(const torch::Tensor& t,
+                                    int gmem_inner_dim, int gmem_outer_dim,
+                                    int smem_inner_dim, int smem_outer_dim,
+                                    const int& gmem_outer_stride,
+                                    const int& swizzle_mode, const int& swizzle_base = 0,
+                                    const bool& allow_tf32 = false) {
+    const auto& elem_size = static_cast<int>(t.element_size());
+    if (swizzle_mode != 0)
+        smem_inner_dim = swizzle_mode / elem_size;
+
+    CUtensorMap tensor_map;
+    const cuuint64_t gmem_dims[2] = {static_cast<cuuint64_t>(gmem_inner_dim), static_cast<cuuint64_t>(gmem_outer_dim)};
+    const cuuint32_t smem_dims[2] = {static_cast<cuuint32_t>(smem_inner_dim), static_cast<cuuint32_t>(smem_outer_dim)};
+    const cuuint64_t gmem_strides[1] = {static_cast<cuuint64_t>(gmem_outer_stride * elem_size), };
+    const cuuint32_t elem_strides[2] = {1, 1};
+    // if (get_env<int>("DG_JIT_DEBUG")) {
+    //     printf("Making TMA desc: global memory: %d %d, shared memory: %d %d, outer stride: %d, swizzle: %d (base: %d), elem size: %d\n",
+    //            gmem_inner_dim, gmem_outer_dim, smem_inner_dim, smem_outer_dim,
+    //            gmem_outer_stride, swizzle_mode, swizzle_base, elem_size);
+    // }
+    cuTensorMapEncodeTiled(
+        &tensor_map, aten_dtype_to_tensor_map_dtype(t.scalar_type(), allow_tf32),
+        2, t.data_ptr(), gmem_dims, gmem_strides, smem_dims, elem_strides,
+        CU_TENSOR_MAP_INTERLEAVE_NONE, mode_into_tensor_map_swizzle(swizzle_mode, swizzle_base),
+        CU_TENSOR_MAP_L2_PROMOTION_L2_256B, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE);
+    return tensor_map;
+}
+
+static CUtensorMap make_tma_3d_desc(const torch::Tensor& t,
+                                    const int& gmem_dim_0, const int& gmem_dim_1, const int& gmem_dim_2,
+                                    const int& smem_dim_0, const int& smem_dim_1, const int& smem_dim_2,
+                                    const int& gmem_stride_0, const int& gmem_stride_1,
+                                    const int& swizzle_mode, const int& swizzle_base = 0,
+                                    const bool& allow_tf32 = false) {
+    const auto& elem_size = static_cast<int>(t.element_size());
+    if (swizzle_mode != 0)
+        DG_HOST_ASSERT(smem_dim_0 == swizzle_mode / elem_size);
+
+    CUtensorMap tensor_map;
+    const cuuint64_t gmem_dims[3] = {static_cast<cuuint64_t>(gmem_dim_0), static_cast<cuuint64_t>(gmem_dim_1), static_cast<cuuint64_t>(gmem_dim_2),};
+    const cuuint32_t smem_dims[3] = {static_cast<cuuint32_t>(smem_dim_0), static_cast<cuuint32_t>(smem_dim_1), static_cast<cuuint32_t>(smem_dim_2)};
+    const cuuint64_t gmem_strides[2] = {static_cast<cuuint64_t>(gmem_stride_0 * elem_size), static_cast<cuuint64_t>(gmem_stride_1 * elem_size)};
+    const cuuint32_t elem_strides[3] = {1, 1, 1};
+    // if (get_env<int>("DG_JIT_DEBUG")) {
+    //     printf("Making 3D TMA desc: global memory: %d %d %d, shared memory: %d %d %d, outer stride: %d %d, swizzle: %d, elem size: %d\n",
+    //            gmem_dim_0, gmem_dim_1, gmem_dim_2, smem_dim_0, smem_dim_1, smem_dim_2,
+    //            gmem_stride_0, gmem_stride_1, swizzle_mode, elem_size);
+    // }
+    cuTensorMapEncodeTiled(
+        &tensor_map, aten_dtype_to_tensor_map_dtype(t.scalar_type(), allow_tf32),
+        3, t.data_ptr(), gmem_dims, gmem_strides, smem_dims, elem_strides,
+        CU_TENSOR_MAP_INTERLEAVE_NONE, mode_into_tensor_map_swizzle(swizzle_mode, swizzle_base),
+        CU_TENSOR_MAP_L2_PROMOTION_L2_256B, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE);
+    return tensor_map;
+}
+
+static CUtensorMap make_tma_a_desc(const cute::UMMA::Major& major,
+                                   const torch::Tensor& t,
+                                   const int& shape_m, const int& shape_k,
+                                   const int& block_m, const int& block_k,
+                                   const int& outer_stride,
+                                   const int& num_groups,
+                                   const int& swizzle_mode, const int& swizzle_base = 0,
+                                   const bool& allow_tf32 = false) {
+    if (num_groups > 1)
+        DG_HOST_ASSERT(major == cute::UMMA::Major::K);
+    const auto& [gmem_inner_dim, gmem_outer_dim] = get_inner_outer_dims(major, shape_k, shape_m * num_groups);
+    const auto& [smem_inner_dim, smem_outer_dim] = get_inner_outer_dims(major, block_k, block_m);
+    return make_tma_2d_desc(t,
+                            gmem_inner_dim, gmem_outer_dim,
+                            smem_inner_dim, smem_outer_dim,
+                            outer_stride,
+                            swizzle_mode, swizzle_base,
+                            allow_tf32);
+}
+
+static CUtensorMap make_tma_b_desc(const cute::UMMA::Major& major,
+                                   const torch::Tensor& t,
+                                   const int& shape_n, const int& shape_k,
+                                   const int& block_n, const int& block_k,
+                                   const int& outer_stride,
+                                   const int& num_groups,
+                                   const int& swizzle_mode, const int& swizzle_base = 0,
+                                   const bool& allow_tf32 = false) {
+    const auto& [gmem_inner_dim, gmem_outer_dim] = get_inner_outer_dims(major, shape_k, shape_n);
+    const auto& [smem_inner_dim, smem_outer_dim] = get_inner_outer_dims(major, block_k, block_n);
+
+    // `num_groups` is always applied into the outer dimensions
+    return make_tma_2d_desc(t,
+                            gmem_inner_dim, gmem_outer_dim * num_groups,
+                            smem_inner_dim, smem_outer_dim,
+                            outer_stride,
+                            swizzle_mode, swizzle_base,
+                            allow_tf32);
+}
+
+static CUtensorMap make_tma_cd_desc(const torch::Tensor& t,
+                                    const int& shape_m, const int& shape_n,
+                                    const int& block_m, const int& block_n,
+                                    const int& outer_stride,
+                                    const int& num_groups,
+                                    const int& swizzle_mode, const int& swizzle_base = 0,
+                                    const bool& allow_tf32 = false) {
+    // Swizzling requires the inner box dim to be less or equal than `kSwizzleCDMode`
+    // bytes, so `BLOCK_N * sizeof(T) / kSwizzleCDMode` TMA stores are required
+    return make_tma_2d_desc(t,
+                            shape_n, shape_m * num_groups,
+                            block_n, block_m,
+                            outer_stride,
+                            swizzle_mode, swizzle_base,
+                            allow_tf32);
+}
+
+static CUtensorMap make_tma_sf_desc(const cute::UMMA::Major& major,
+                                    const torch::Tensor& t,
+                                    int shape_mn, int shape_k,
+                                    const int& block_mn, const int& block_k,
+                                    const int& num_groups,
+                                    const int& swizzle_mode, const int& swizzle_base = 0,
+                                    const bool& allow_tf32 = false) {
+    DG_HOST_ASSERT(major == cute::UMMA::Major::MN);
+
+    // TODO: maybe swizzle SF as well
+    DG_HOST_ASSERT(swizzle_mode == 0);
+
+    shape_mn = get_tma_aligned_size(shape_mn, static_cast<int>(t.element_size()));
+    return make_tma_2d_desc(t,
+                            shape_mn, ceil_div(shape_k, block_k * (t.scalar_type() == torch::kFloat ? 1 : 4)) * num_groups,
+                            block_mn, 1,
+                            shape_mn,
+                            swizzle_mode, swizzle_base,
+                            allow_tf32);
+}
+
+} // namespace deep_gemm

packages/ltx-kernels/csrc/blockwise/sm89_fp8_gemm_1d2d.hpp

@@ -0,0 +1,39 @@
+#pragma once
+#include <cuda.h>
+#include <cuda_runtime.h>
+#include <nvrtc.h>
+
+#include <torch/python.h>
+#include <ATen/cuda/CUDAContext.h>
+
+#include "kernels/geforce/static_switch.h"
+
+namespace sm89 {
+template<bool use_fast_accum>
+void fp8_bias_gemm_cuda(void* Aptr, void* SFA, void* Bptr, void* SFB, void* bias_ptr, void* out, int M, int N, int K, cudaStream_t stream);
+}
+
+namespace blockwise {
+static void sm89_fp8_gemm_1d2d_bias(const torch::Tensor& a, const torch::Tensor& sfa,
+                               const torch::Tensor& b, const torch::Tensor& sfb,
+                               const torch::Tensor& bias,
+                               const torch::Tensor& d,
+                               const int& m, const int& n, const int& k,
+                               const bool use_fast_accum) {
+    auto stream = at::cuda::getCurrentCUDAStream().stream();
+
+    if (use_fast_accum) {
+        sm89::fp8_bias_gemm_cuda<true>(
+                a.data_ptr(), sfa.data_ptr(),
+                b.data_ptr(), sfb.data_ptr(),
+                bias.data_ptr(), d.data_ptr(),
+                m, n, k, stream);
+    } else {
+        sm89::fp8_bias_gemm_cuda<false>(
+                a.data_ptr(), sfa.data_ptr(),
+                b.data_ptr(), sfb.data_ptr(),
+                bias.data_ptr(), d.data_ptr(),
+                m, n, k, stream);
+    }
+}
+}

packages/ltx-kernels/csrc/blockwise/sm90_fp8_gemm_1d2d_bias.hpp

@@ -0,0 +1,66 @@
+
+
+#pragma once
+#include <cuda.h>
+#include <cuda_runtime.h>
+#include <nvrtc.h>
+
+#include <torch/python.h>
+#include <ATen/cuda/CUDAContext.h>
+
+#include <cute/arch/mma_sm100_desc.hpp>
+#include "runtime_utils.hpp"
+
+#include "config.hpp"
+#include "static_switch.hpp"
+
+namespace deep_gemm{
+template<int N, int K>
+void sm90_fp8_gemm_1d2d_bias_launch(int num_sms, int num_threads, int cluster_dim, int smem_size, cudaStream_t stream, float* sfb, float* bias, int* grouped_layout,
+                        uint32_t shape_m, uint32_t shape_n, uint32_t shape_k,
+                        const CUtensorMap tensor_map_a,
+                        const CUtensorMap tensor_map_b,
+                        const CUtensorMap tensor_map_d,
+                        const CUtensorMap tensor_map_sfa);
+};
+
+namespace blockwise{
+
+static void sm90_fp8_gemm_1d2d_bias(const torch::Tensor& a, const torch::Tensor& sfa,
+                               const torch::Tensor& b, const torch::Tensor& sfb,
+                               const torch::Tensor& bias,
+                               const std::optional<torch::Tensor>& c,
+                               const torch::Tensor& d,
+                               const int& m, const int& n, const int& k, const int num_sms) {
+    // DG_HOST_ASSERT(not c.has_value() and d.scalar_type() == torch::kBFloat16);
+    const auto& config = GemmConfig<90>();
+
+    // Requires no TMA splits
+    // DG_HOST_ASSERT(config.smem_config.swizzle_a_mode == config.block_k);
+    // DG_HOST_ASSERT(config.smem_config.swizzle_b_mode == config.block_k);
+    int smem_size = k == 16384 || k == 8192 ? 216624 : config.smem_config.smem_size;     
+    const auto& tensor_map_a = make_tma_a_desc(cute::UMMA::Major::K, a, m, k,
+                                               config.block_m,
+                                               config.block_k,
+                                               static_cast<int>(a.stride(-2)), 1,
+                                               config.smem_config.swizzle_a_mode);
+    const auto& tensor_map_b = make_tma_b_desc(cute::UMMA::Major::K, b, n, k,
+                                               config.block_n,
+                                               config.block_k,
+                                               static_cast<int>(b.stride(-2)), 1,
+                                               config.smem_config.swizzle_b_mode);
+    const auto& tensor_map_d = make_tma_cd_desc(d, m, static_cast<int>(d.size(-1)),
+                                                config.block_m,
+                                                config.block_n,
+                                                static_cast<int>(d.stride(-2)), 1,
+                                                config.smem_config.swizzle_cd_mode);
+    const auto& tensor_map_sfa = make_tma_sf_desc(cute::UMMA::Major::MN, sfa, m, k,
+                                                  config.block_m, config.block_k, 1, 0);
+    auto stream = at::cuda::getCurrentCUDAStream().stream();
+    // Launch    
+    DIM_SWITCH(k, K, 
+        DIM_SWITCH(n, N, 
+            deep_gemm::sm90_fp8_gemm_1d2d_bias_launch<N, K>(num_sms, config.thread_config.num_threads, config.multicast_config.num_multicast, smem_size, stream, (float*)sfb.data_ptr(), (float*)bias.data_ptr(), nullptr, m, n, k, tensor_map_a, tensor_map_b, tensor_map_d, tensor_map_sfa);)
+    )
+}
+};

packages/ltx-kernels/csrc/blockwise/static_switch.hpp

@@ -0,0 +1,60 @@
+#pragma once
+
+#define DIM_SWITCH(VAR_NAME, CONST_NAME, ...)    \
+    if (VAR_NAME == 4096) {                         \
+      constexpr static int CONST_NAME = 4096;       \
+      __VA_ARGS__                                 \
+    } else if (VAR_NAME == 2048){                  \
+        constexpr static int CONST_NAME = 2048;    \
+        __VA_ARGS__                                 \
+    } else if (VAR_NAME == 8192){                   \
+        constexpr static int CONST_NAME = 8192;    \
+        __VA_ARGS__                                \
+    } else if(VAR_NAME == 16384) {                  \
+        constexpr static int CONST_NAME = 16384;     \
+        __VA_ARGS__                                  \
+    } else {                                         \
+        TORCH_CHECK(false, "Unsupported DIM_SWITCH value: ", VAR_NAME); \
+    }
+
+#define BOOL_SWITCH(COND, CONST_NAME, ...)      \
+    if (COND) {                                 \
+      constexpr static bool CONST_NAME = true;  \
+      __VA_ARGS__                               \
+    } else {                                    \
+      constexpr static bool CONST_NAME = false; \
+      __VA_ARGS__                               \
+    }                                           \
+//K/128
+#define BLOCK_K_SWITCH(COSNT_NAME, ...)             \
+    if (K == 2048) {                          \
+      constexpr static int COSNT_NAME = 16;   \
+      __VA_ARGS__                             \
+    }                                         \
+    else if (K == 4096) {                     \
+      constexpr static int COSNT_NAME = 32;    \
+      __VA_ARGS__                               \
+    } else if (K == 8192) {                      \
+      constexpr static int COSNT_NAME = 64;   \
+      __VA_ARGS__                               \
+    } else if (K == 16384) {                      \
+      constexpr static int COSNT_NAME = 128;    \
+      __VA_ARGS__                               \
+    } else {                                    \
+      TORCH_CHECK(false, "Unsupported K value: ", K);   \
+    }
+
+#define M_SWITCH(...) \
+    if (M <= 1024) {                             \
+        constexpr static int BM = 128;    \
+        constexpr static int BN = 128;    \
+        constexpr static int WARP_ROW = 2;    \
+        constexpr static int WARP_COL = 2;    \
+      __VA_ARGS__                               \
+    } else {                                    \
+        constexpr static int BM = 128;    \
+        constexpr static int BN = 256;    \
+        constexpr static int WARP_ROW = 2;    \
+        constexpr static int WARP_COL = 4;    \
+        __VA_ARGS__                           \
+    }                                          

packages/ltx-kernels/csrc/blockwise/utils.hpp

@@ -0,0 +1,38 @@
+#pragma once
+
+#include <cuda_bf16.h>
+#include <cuda_fp8.h>
+#include <cuda/std/cstdint>
+#include <cuda/std/utility>
+#include <cute/container/tuple.hpp>
+
+#ifdef __CLION_IDE__
+
+__host__ __device__ __forceinline__ void host_device_printf(const char* format, ...) {
+    asm volatile("trap;");
+}
+
+#define printf host_device_printf
+#endif
+
+#ifndef DG_DEVICE_ASSERT
+#define DG_DEVICE_ASSERT(cond) \
+do { \
+    if (not (cond)) { \
+        printf("Assertion failed: %s:%d, condition: %s\n", __FILE__, __LINE__, #cond); \
+        asm("trap;"); \
+    } \
+} while (0)
+#endif
+
+#ifndef DG_TRAP_ONLY_DEVICE_ASSERT
+#define DG_TRAP_ONLY_DEVICE_ASSERT(cond) \
+do { \
+    if (not (cond)) \
+        asm("trap;"); \
+} while (0)
+#endif
+
+#ifndef DG_STATIC_ASSERT
+#define DG_STATIC_ASSERT(cond, ...) static_assert(cond, __VA_ARGS__)
+#endif

packages/ltx-kernels/csrc/include/cuda/configs.cuh

@@ -0,0 +1,88 @@
+/**
+ * @file configs.cuh
+ * @brief Configuration constants and compile-time settings for ltx-kernels.
+ *
+ * This header defines the tunable parameters and constants used throughout
+ * the ltx-kernels communication library. These values are chosen to balance
+ * performance across different GPU architectures.
+ */
+
+#pragma once
+
+#include <cstdint>
+#include <cuda_bf16.h>
+#include <cuda_runtime.h>
+
+namespace ltx_kernels {
+// =============================================================================
+// Synchronization Configuration
+// =============================================================================
+
+/**
+ * @brief Default barrier timeout in seconds.
+ *
+ * If a barrier wait exceeds this timeout, the kernel traps to indicate a deadlock or
+ * communication failure. All2All converts it to clock cycles using the device's peak SM clock
+ * (cudaDeviceGetAttribute(cudaDevAttrClockRate)), so the wall-clock guard holds regardless of GPU.
+ */
+constexpr double DEFAULT_BARRIER_TIMEOUT_SECONDS = 10.0;
+
+// =============================================================================
+// Hardware Limits
+// =============================================================================
+
+/**
+ * @brief Maximum number of peer GPUs supported for IPC communication.
+ *
+ * This limits the size of static arrays for buffer pointers and barrier signals.
+ * Set to 8 to support up to 8-way tensor parallelism (common for DGX systems).
+ */
+constexpr int MAX_NUM_PEERS = 8;
+
+// =============================================================================
+// Kernel Configuration
+// =============================================================================
+
+/**
+ * @brief Default number of threads per block for All2All kernels.
+ *
+ * Used by send_recv_all2all and gather_heads kernels. The value 512 provides
+ * good occupancy while leaving registers for complex pointer arithmetic.
+ */
+constexpr int DEFAULT_KERNEL_THREADS = 512;
+
+/**
+ * @brief Number of threads per block for the AllGather kernel.
+ *
+ * AllGather uses more threads (1024) because its memory access pattern
+ * is simpler (no head selection), allowing higher thread-level parallelism.
+ */
+constexpr int ALLGATHER_KERNEL_THREADS = 1024;
+
+} // namespace ltx_kernels
+
+// =============================================================================
+// Torch/CUDA Compatibility Fixes
+// =============================================================================
+
+/*
+ * PyTorch sometimes disables CUDA half/bfloat16 operators and conversions
+ * to avoid ambiguity in template resolution. We re-enable them here since
+ * our kernels explicitly handle these types.
+ */
+
+#ifdef __CUDA_NO_HALF_CONVERSIONS__
+#undef __CUDA_NO_HALF_CONVERSIONS__
+#endif
+#ifdef __CUDA_NO_HALF_OPERATORS__
+#undef __CUDA_NO_HALF_OPERATORS__
+#endif
+#ifdef __CUDA_NO_HALF2_OPERATORS__
+#undef __CUDA_NO_HALF2_OPERATORS__
+#endif
+#ifdef __CUDA_NO_BFLOAT16_CONVERSIONS__
+#undef __CUDA_NO_BFLOAT16_CONVERSIONS__
+#endif
+#ifdef __CUDA_NO_BFLOAT162_OPERATORS__
+#undef __CUDA_NO_BFLOAT162_OPERATORS__
+#endif

packages/ltx-kernels/csrc/include/cuda/exceptions.cuh

@@ -0,0 +1,170 @@
+/**
+ * @file exceptions.cuh
+ * @brief Exception handling and assertion macros for CUDA/C++ code.
+ *
+ * This header provides a unified exception type and assertion macros for
+ * both host and device code. The macros capture file and line information
+ * for easier debugging of errors.
+ *
+ * ## Usage Examples
+ *
+ * ```cpp
+ * // Check CUDA API call
+ * CUDA_CHECK(cudaMalloc(&ptr, size));
+ *
+ * // Host-side assertion
+ * EP_HOST_ASSERT(tensor.is_contiguous());
+ *
+ * // Device-side assertion (inside kernel)
+ * EP_DEVICE_ASSERT(threadIdx.x < MAX_THREADS);
+ *
+ * // Compile-time assertion
+ * EP_STATIC_ASSERT(sizeof(int4) == 16, "int4 must be 16 bytes");
+ * ```
+ */
+
+#pragma once
+
+#include <exception>
+#include <string>
+
+#include "configs.cuh"
+
+// =============================================================================
+// Static Assertions
+// =============================================================================
+
+/**
+ * @brief Compile-time assertion macro.
+ *
+ * @param cond Condition that must be true at compile time
+ * @param reason Human-readable error message if condition fails
+ */
+#ifndef EP_STATIC_ASSERT
+#define EP_STATIC_ASSERT(cond, reason) static_assert(cond, reason)
+#endif
+
+// =============================================================================
+// Exception Type
+// =============================================================================
+
+/**
+ * @class EPException
+ * @brief Custom exception type with file/line information.
+ *
+ * EPException captures the location (file, line) and context (name, error)
+ * of the error for debugging. It inherits from std::exception for
+ * compatibility with standard C++ exception handling.
+ *
+ * ## Message Format
+ *
+ * The what() message has the format:
+ * "Failed: <name> error <file>:<line> '<error message>'"
+ */
+class EPException : public std::exception {
+private:
+  std::string message = {}; ///< Formatted error message
+
+public:
+  /**
+   * @brief Constructs an EPException with location and error information.
+   *
+   * @param name Category of error (e.g., "CUDA", "Assertion")
+   * @param file Source file where error occurred (__FILE__)
+   * @param line Line number where error occurred (__LINE__)
+   * @param error Description of the error
+   */
+  explicit EPException(const char *name, const char *file, const int line, const std::string &error) {
+    message = std::string("Failed: ") + name + " error " + file + ":" + std::to_string(line) + " '" + error + "'";
+  }
+
+  /**
+   * @brief Returns the formatted error message.
+   * @return C-string containing the error message
+   */
+  const char *what() const noexcept override { return message.c_str(); }
+};
+
+// =============================================================================
+// Runtime Assertion Macros
+// =============================================================================
+
+/**
+ * @brief Checks CUDA API return value and throws on error.
+ *
+ * Use this macro to wrap all CUDA runtime API calls. If the call fails,
+ * an EPException is thrown with the CUDA error string.
+ *
+ * @param cmd CUDA API call expression
+ * @throws EPException if the CUDA call returns an error
+ *
+ * Example:
+ * ```cpp
+ * CUDA_CHECK(cudaMalloc(&ptr, size));
+ * CUDA_CHECK(cudaMemcpy(dst, src, size, cudaMemcpyDeviceToDevice));
+ * ```
+ */
+#ifndef CUDA_CHECK
+#define CUDA_CHECK(cmd)                                                                                                \
+  do {                                                                                                                 \
+    cudaError_t e = (cmd);                                                                                             \
+    if (e != cudaSuccess) {                                                                                            \
+      throw EPException("CUDA", __FILE__, __LINE__, cudaGetErrorString(e));                                            \
+    }                                                                                                                  \
+  } while (0)
+#endif
+
+/**
+ * @brief Host-side assertion that throws on failure.
+ *
+ * Use this for runtime checks in host code. If the condition is false,
+ * an EPException is thrown with the condition as the error message.
+ *
+ * @param cond Condition to check (must be true)
+ * @throws EPException if condition is false
+ *
+ * Example:
+ * ```cpp
+ * EP_HOST_ASSERT(tensor.dim() == 4);
+ * EP_HOST_ASSERT(rank >= 0 && rank < world_size);
+ * ```
+ */
+#ifndef EP_HOST_ASSERT
+#define EP_HOST_ASSERT(cond)                                                                                           \
+  do {                                                                                                                 \
+    if (not(cond)) {                                                                                                   \
+      throw EPException("Assertion", __FILE__, __LINE__, #cond);                                                       \
+    }                                                                                                                  \
+  } while (0)
+#endif
+
+/**
+ * @brief Device-side assertion that traps on failure.
+ *
+ * Use this for runtime checks inside CUDA kernels. If the condition is
+ * false, prints an error message and executes a trap instruction to
+ * halt the GPU.
+ *
+ * @warning This causes the entire kernel to abort. Use sparingly and
+ *          consider removing from release builds for performance.
+ *
+ * @param cond Condition to check (must be true)
+ *
+ * Example:
+ * ```cpp
+ * __global__ void my_kernel(int* data, int size) {
+ *     int idx = threadIdx.x + blockIdx.x * blockDim.x;
+ *     EP_DEVICE_ASSERT(idx < size);
+ *     data[idx] = 42;
+ * }
+ * ```
+ */
+#ifndef EP_DEVICE_ASSERT
+#define EP_DEVICE_ASSERT(cond)                                                                                         \
+  do {                                                                                                                 \
+    if (not(cond)) {                                                                                                   \
+      printf("Assertion failed: %s:%d, condition: %s\n", __FILE__, __LINE__, #cond);                                   \
+      asm("trap;");                                                                                                    \
+    }                                                                                                                  \
+  } while (0)
+#endif

packages/ltx-kernels/csrc/include/cuda/utils.cuh

@@ -0,0 +1,360 @@
+/**
+ * @file utils.cuh
+ * @brief Low-level CUDA utility functions for memory operations and synchronization.
+ *
+ * This header provides optimized PTX assembly wrappers for memory operations
+ * that bypass cache hierarchy or use specific memory ordering semantics.
+ * These are critical for achieving peak bandwidth in multi-GPU communication.
+ *
+ * ## Memory Operation Types
+ *
+ * - **Non-allocating stores (st_na)**: Bypass L1 cache to avoid polluting it
+ *   with data that won't be reused locally
+ * - **Non-caching loads (ld_nc)**: Bypass L1 cache for streaming reads
+ * - **Acquire/Release**: Memory ordering for synchronization
+ * - **System scope (sys)**: Visibility across all GPUs, not just this one
+ *
+ * ## Cache Hints
+ *
+ * - L1::no_allocate: Don't allocate in L1 on miss (streaming pattern)
+ * - L2::256B: Use 256-byte L2 cache lines
+ * - volatile: Bypass all caches, always go to memory
+ */
+
+#pragma once
+#include <stdint.h>
+
+// =============================================================================
+// PTX Instruction Selection
+// =============================================================================
+
+/**
+ * Store instruction macro. When DISABLE_AGGRESSIVE_PTX_INSTRS is not defined,
+ * uses non-allocating stores to avoid polluting L1 cache with write-only data.
+ */
+#ifndef DISABLE_AGGRESSIVE_PTX_INSTRS
+#define ST_NA_FUNC "st.global.L1::no_allocate"
+#else
+#define ST_NA_FUNC "st.global"
+#endif
+
+/**
+ * Load instruction macro. When DISABLE_AGGRESSIVE_PTX_INSTRS is not defined,
+ * uses non-caching loads optimized for streaming access patterns.
+ */
+#ifndef DISABLE_AGGRESSIVE_PTX_INSTRS
+#define LD_NC_FUNC "ld.global.nc.L1::no_allocate.L2::256B"
+#else
+#define LD_NC_FUNC "ld.volatile.global.L2::256B"
+#endif
+
+namespace ltx_kernels {
+
+// =============================================================================
+// Round-Robin SM Distribution Helpers
+// =============================================================================
+
+/**
+ * @brief Compute target rank for a given SM using round-robin distribution.
+ *
+ * Round-robin assignment ensures all SMs are utilized even when num_sms
+ * is not evenly divisible by world_size.
+ *
+ * @param sm_id The SM/block ID (blockIdx.x)
+ * @param world_size Total number of ranks
+ * @return Target rank for this SM
+ */
+__device__ __forceinline__ int get_target_rank(int sm_id, int world_size) { return sm_id % world_size; }
+
+/**
+ * @brief Compute local SM index within a rank's SM group.
+ *
+ * With round-robin, SM i is the (i / world_size)-th SM assigned to its rank.
+ *
+ * @param sm_id The SM/block ID (blockIdx.x)
+ * @param world_size Total number of ranks
+ * @return Local index of this SM within its assigned rank's group
+ */
+__device__ __forceinline__ int get_rank_local_sm_id(int sm_id, int world_size) { return sm_id / world_size; }
+
+/**
+ * @brief Compute number of SMs assigned to a specific rank.
+ *
+ * With round-robin distribution:
+ *   - Ranks [0, extra) get (base + 1) SMs each
+ *   - Ranks [extra, world_size) get base SMs each
+ * where base = num_sms / world_size, extra = num_sms % world_size
+ *
+ * @param target_rank The rank to query
+ * @param num_sms Total number of SMs launched
+ * @param world_size Total number of ranks
+ * @return Number of SMs assigned to target_rank
+ */
+__device__ __forceinline__ int get_num_sms_for_rank(int target_rank, int num_sms, int world_size) {
+  int base_sms = num_sms / world_size;
+  int extra_sms = num_sms % world_size;
+  return base_sms + (target_rank < extra_sms ? 1 : 0);
+}
+
+// =============================================================================
+// Control Flow
+// =============================================================================
+
+/**
+ * @brief Triggers a GPU trap (fatal error).
+ *
+ * Used for unrecoverable errors like synchronization timeout.
+ * Causes the kernel to abort and report an error to the host.
+ */
+__device__ __forceinline__ void trap() { asm("trap;"); }
+
+// =============================================================================
+// Memory Ordering Operations (for synchronization)
+// =============================================================================
+
+/**
+ * @brief System-scope store with release ordering.
+ *
+ * Ensures all prior memory operations are visible before this store.
+ * System scope means visibility across all GPUs (for IPC communication).
+ *
+ * @param ptr Pointer to global memory
+ * @param val Value to store
+ */
+__device__ __forceinline__ void st_release_sys_global(const int *ptr, int val) {
+  asm volatile("st.release.sys.global.s32 [%0], %1;" ::"l"(ptr), "r"(val) : "memory");
+}
+
+/**
+ * @brief System-scope store with relaxed ordering.
+ *
+ * No ordering guarantees - fastest store but requires external synchronization.
+ *
+ * @param ptr Pointer to global memory
+ * @param val Value to store
+ */
+__device__ __forceinline__ void st_relaxed_sys_global(const int *ptr, int val) {
+  asm volatile("st.relaxed.sys.global.s32 [%0], %1;" ::"l"(ptr), "r"(val) : "memory");
+}
+
+/**
+ * @brief CTA-scope store with release ordering.
+ *
+ * Ensures visibility within the thread block (CTA = Cooperative Thread Array).
+ *
+ * @param ptr Pointer to global memory
+ * @param val Value to store
+ */
+__device__ __forceinline__ void st_release_cta(const int *ptr, int val) {
+  asm volatile("st.release.cta.s32 [%0], %1;" ::"l"(ptr), "r"(val) : "memory");
+}
+
+/**
+ * @brief System-scope load with acquire ordering (32-bit).
+ *
+ * Ensures subsequent memory operations are ordered after this load.
+ * System scope for IPC visibility across GPUs.
+ *
+ * @param ptr Pointer to global memory
+ * @return Loaded value
+ */
+__device__ __forceinline__ int ld_acquire_sys_global(const int *ptr) {
+  int ret;
+  asm volatile("ld.acquire.sys.global.s32 %0, [%1];" : "=r"(ret) : "l"(ptr));
+  return ret;
+}
+
+/**
+ * @brief System-scope load with acquire ordering (64-bit).
+ *
+ * @param ptr Pointer to global memory
+ * @return Loaded value
+ */
+__device__ __forceinline__ uint64_t ld_acquire_sys_global(const uint64_t *ptr) {
+  uint64_t ret;
+  asm volatile("ld.acquire.sys.global.u64 %0, [%1];" : "=l"(ret) : "l"(ptr));
+  return ret;
+}
+
+/**
+ * @brief GPU-scope load with acquire ordering.
+ *
+ * Visibility limited to this GPU (not for IPC).
+ *
+ * @param ptr Pointer to global memory
+ * @return Loaded value
+ */
+__device__ __forceinline__ int ld_acquire_global(const int *ptr) {
+  int ret;
+  asm volatile("ld.acquire.gpu.global.s32 %0, [%1];" : "=r"(ret) : "l"(ptr));
+  return ret;
+}
+
+/**
+ * @brief Volatile load bypassing all caches.
+ *
+ * Always reads from memory, never from cache. Used for polling
+ * synchronization variables that may be updated by other GPUs.
+ *
+ * @param ptr Pointer to global memory
+ * @return Loaded value
+ */
+__device__ __forceinline__ int ld_volatile_global(const int *ptr) {
+  int ret;
+  asm volatile("ld.volatile.global.s32 %0, [%1];" : "=r"(ret) : "l"(ptr));
+  return ret;
+}
+
+// =============================================================================
+// Optimized Bulk Memory Operations
+// =============================================================================
+
+/**
+ * @brief Non-allocating 128-bit store.
+ *
+ * Stores an int4 (128 bits / 16 bytes) without allocating in L1 cache.
+ * Optimal for write-streaming patterns where data won't be read locally.
+ *
+ * @param ptr Destination pointer (must be 16-byte aligned)
+ * @param value Data to store
+ */
+__device__ __forceinline__ void st_na_global(const int4 *ptr, const int4 &value) {
+  asm volatile(ST_NA_FUNC ".v4.s32 [%0], {%1, %2, %3, %4};" ::"l"(ptr), "r"(value.x), "r"(value.y), "r"(value.z),
+               "r"(value.w));
+}
+
+/**
+ * @brief Non-caching 128-bit load.
+ *
+ * Loads an int4 bypassing L1 cache with optimized L2 caching (256B lines).
+ * Optimal for read-streaming patterns.
+ *
+ * @param ptr Source pointer (must be 16-byte aligned)
+ * @return Loaded int4 value
+ */
+__device__ __forceinline__ int4 ld_nc_global(const int4 *ptr) {
+  int4 ret;
+  asm volatile(LD_NC_FUNC ".v4.s32 {%0, %1, %2, %3}, [%4];"
+               : "=r"(ret.x), "=r"(ret.y), "=r"(ret.z), "=r"(ret.w)
+               : "l"(ptr));
+  return ret;
+}
+
+/**
+ * @brief Barrier synchronization pattern for multi-GPU communication.
+ *
+ * This function implements a barrier synchronization protocol used in All2All
+ * and AllGather operations. It signals completion to target ranks and waits
+ * for all expected signals to arrive before resetting the barrier.
+ *
+ * Protocol:
+ * 1. Thread 0 of each block signals completion to the target rank
+ * 2. Block 0 waits for all ranks to signal (with timeout protection)
+ * 3. Once all signals received, reset the barrier counters
+ *
+ * @param barrier_signal_ptrs Array of pointers to barrier signal buffers for each rank
+ * @param target_rank The rank this block is sending data to
+ * @param rank This GPU's rank
+ * @param world_size Total number of GPUs/ranks
+ * @param expected_count Number of signals expected (typically num_sms_per_rank)
+ * @param sm_id The SM/block ID (blockIdx.x)
+ * @param thread_id The thread ID within the block (threadIdx.x)
+ * @param timeout_cycles Number of cycles to wait before timeout
+ */
+__device__ __forceinline__ void barrier_wait_and_reset(int **barrier_signal_ptrs, int target_rank, int rank,
+                                                       int world_size, int expected_count, int sm_id, int thread_id,
+                                                       uint64_t timeout_cycles) {
+  // Release: fence so peers see our data writes, then sync before signaling.
+  __threadfence_system();
+  __syncthreads();
+
+  // Thread 0 signals completion to target rank
+  if (thread_id == 0) {
+    atomicAdd_system(barrier_signal_ptrs[target_rank] + rank, 1);
+  }
+
+  // Synchronize before checking signals
+  __syncthreads();
+
+  // Only block 0 waits for all signals and resets the barrier
+  if (sm_id == 0 && thread_id < world_size) {
+    auto start_time = clock64();
+    while (true) {
+      // Acquire: seeing the signal guarantees the peer's data is visible.
+      int recv_count = ld_acquire_sys_global(barrier_signal_ptrs[rank] + thread_id);
+      if (recv_count == expected_count) {
+        break;
+      }
+      if (clock64() - start_time >= timeout_cycles) {
+        printf("All2All barrier timeout: rank=%d, waiting_for_source=%d, expected=%d, got=%d\n", rank, thread_id,
+               expected_count, recv_count);
+        trap();
+      }
+    }
+    // Reset barrier for next use
+    atomicSub_system(barrier_signal_ptrs[rank] + thread_id, expected_count);
+  }
+}
+
+/**
+ * @brief Barrier synchronization for round-robin SM distribution.
+ *
+ * Similar to barrier_wait_and_reset, but handles the case where SMs are
+ * distributed round-robin across ranks, resulting in different target ranks
+ * receiving different numbers of signals.
+ *
+ * With round-robin: target ranks [0, extra) receive (base + 1) signals from
+ * each source, and target ranks [extra, world_size) receive base signals
+ * from each source. Note that ALL sources send the same count to a given
+ * receiver - the count depends on the receiver's rank position.
+ *
+ * @param barrier_signal_ptrs Array of pointers to barrier signal buffers for each rank
+ * @param target_rank The rank this block is sending data to
+ * @param rank This GPU's rank
+ * @param world_size Total number of GPUs/ranks
+ * @param num_sms Total number of SMs launched (used to compute expected counts)
+ * @param sm_id The SM/block ID (blockIdx.x)
+ * @param thread_id The thread ID within the block (threadIdx.x)
+ * @param timeout_cycles Number of cycles to wait before timeout
+ */
+__device__ __forceinline__ void barrier_wait_and_reset_roundrobin(int **barrier_signal_ptrs, int target_rank, int rank,
+                                                                  int world_size, int num_sms, int sm_id, int thread_id,
+                                                                  uint64_t timeout_cycles) {
+  // Release: fence so peers see our data writes, then sync before signaling.
+  __threadfence_system();
+  __syncthreads();
+
+  // Thread 0 signals completion to target rank
+  if (thread_id == 0) {
+    atomicAdd_system(barrier_signal_ptrs[target_rank] + rank, 1);
+  }
+
+  // Synchronize before checking signals
+  __syncthreads();
+
+  // Only block 0 waits for all signals and resets the barrier
+  // Each thread handles one source rank
+  if (sm_id == 0 && thread_id < world_size) {
+    // All sources send the same number of signals to THIS receiver.
+    // The count depends on how many SMs target this rank (the receiver).
+    int expected_from_each_source = get_num_sms_for_rank(rank, num_sms, world_size);
+
+    auto start_time = clock64();
+    while (true) {
+      // Acquire: seeing the signal guarantees the peer's data is visible.
+      int recv_count = ld_acquire_sys_global(barrier_signal_ptrs[rank] + thread_id);
+      if (recv_count == expected_from_each_source) {
+        break;
+      }
+      if (clock64() - start_time >= timeout_cycles) {
+        printf("All2All barrier timeout (roundrobin): rank=%d, waiting_for_source=%d, expected=%d, got=%d\n", rank,
+               thread_id, expected_from_each_source, recv_count);
+        trap();
+      }
+    }
+    // Reset barrier for next use
+    atomicSub_system(barrier_signal_ptrs[rank] + thread_id, expected_from_each_source);
+  }
+}
+
+} // namespace ltx_kernels

packages/ltx-kernels/csrc/include/event.hpp

@@ -0,0 +1,114 @@
+/**
+ * @file event.hpp
+ * @brief CUDA stream and event synchronization utilities.
+ *
+ * This header provides wrapper types and helper functions for managing
+ * CUDA events and stream synchronization in PyTorch/ATen environment.
+ * These utilities are used to coordinate asynchronous operations across
+ * multiple CUDA streams.
+ */
+
+#pragma once
+
+#include <ATen/cuda/CUDAContext.h>
+#include <memory>
+
+#include "cuda/exceptions.cuh"
+
+namespace ltx_kernels {
+
+/**
+ * @struct EventHandle
+ * @brief RAII wrapper for a CUDA event with automatic recording.
+ *
+ * EventHandle encapsulates a torch::Event and automatically records it
+ * on the specified (or current) CUDA stream upon construction. This
+ * provides a convenient way to capture the completion point of stream
+ * operations for synchronization purposes.
+ *
+ * ## Usage Example
+ *
+ * ```cpp
+ * // Record event on current stream
+ * EventHandle ev1;
+ *
+ * // Record event on specific stream
+ * EventHandle ev2(my_stream);
+ *
+ * // Make current stream wait for the event
+ * ev1.current_stream_wait();
+ * ```
+ */
+struct EventHandle {
+  /// Shared pointer to the underlying torch::Event
+  std::shared_ptr<torch::Event> event;
+
+  /**
+   * @brief Constructs an EventHandle and records on the current CUDA stream.
+   *
+   * The event captures the completion point of all operations submitted
+   * to the current stream before this constructor is called.
+   */
+  EventHandle() {
+    event = std::make_shared<torch::Event>(torch::kCUDA);
+    event->record(at::cuda::getCurrentCUDAStream());
+  }
+
+  /**
+   * @brief Constructs an EventHandle and records on the specified stream.
+   *
+   * @param stream The CUDA stream to record the event on
+   */
+  explicit EventHandle(const at::cuda::CUDAStream &stream) {
+    event = std::make_shared<torch::Event>(torch::kCUDA);
+    event->record(stream);
+  }
+
+  /// Copy constructor (shares the underlying event)
+  EventHandle(const EventHandle &other) = default;
+
+  /**
+   * @brief Makes the current CUDA stream wait for this event.
+   *
+   * After this call returns, operations submitted to the current stream
+   * will not execute until the event has been reached on its recording stream.
+   */
+  void current_stream_wait() const { at::cuda::getCurrentCUDAStream().unwrap().wait(*event); }
+};
+
+/**
+ * @brief Creates and records a CUDA event on the specified stream.
+ *
+ * @param s The CUDA stream to record on
+ * @return A torch::Event that has been recorded on stream s
+ */
+inline torch::Event create_event(const at::cuda::CUDAStream &s) {
+  auto event = torch::Event(torch::kCUDA);
+  event.record(s);
+  return event;
+}
+
+/**
+ * @brief Makes stream s_0 wait for stream s_1's current position.
+ *
+ * After this call, operations on s_0 will not execute until all operations
+ * currently queued on s_1 have completed.
+ *
+ * @param s_0 The stream that will wait
+ * @param s_1 The stream to wait for
+ * @pre s_0 and s_1 must be different streams
+ */
+inline void stream_wait(const at::cuda::CUDAStream &s_0, const at::cuda::CUDAStream &s_1) {
+  EP_HOST_ASSERT(s_0.id() != s_1.id());
+  s_0.unwrap().wait(create_event(s_1));
+}
+
+/**
+ * @brief Makes a stream wait for a previously recorded event.
+ *
+ * @param s The stream that will wait
+ * @param event The event to wait for
+ */
+inline void stream_wait(const at::cuda::CUDAStream &s, const EventHandle &event) { s.unwrap().wait(*event.event); }
+
+} // namespace ltx_kernels

packages/ltx-kernels/csrc/ops/fp6_bitpack.cpp

@@ -0,0 +1,73 @@
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <torch/extension.h>
+#include <torch/python.h>
+
+#include <vector>
+
+void fp6_pack_cuda(
+    at::Tensor& x,
+    at::Tensor& out,
+    cudaStream_t stream
+);
+
+void fp6_unpack_cuda(
+    at::Tensor& x,
+    at::Tensor& out,
+    cudaStream_t stream
+);
+
+at::Tensor fp6_pack(at::Tensor &x) {
+    // TORCH_CHECK(x.dtype() == torch::kUInt8, "Input tensor must be uint8");
+    TORCH_CHECK(x.is_cuda(), "Input tensor must be on CUDA");
+    TORCH_CHECK(x.is_contiguous(), "Input tensor must be contiguous");
+    TORCH_CHECK(x.dim() == 2, "Input tensor must be 2D [m, n]");
+
+    int64_t m = x.size(0);
+    int64_t n = x.size(1);
+
+    TORCH_CHECK(n % 8 == 0, "n must be divisible by 8, got ", n);
+
+    // Output shape: [m, n*3/4] since 4 elements of 8-bit = 32 bits, 4 elements of 6-bit = 24 bits = 3 bytes
+    int64_t n_packed = n * 3 / 4;
+
+    auto options = torch::TensorOptions()
+        .dtype(torch::kUInt8)
+        .device(x.device());
+
+    at::Tensor out = torch::empty({m, n_packed}, options);
+
+    at::cuda::CUDAGuard device_guard{x.get_device()};
+    auto stream = at::cuda::getCurrentCUDAStream().stream();
+
+    fp6_pack_cuda(x, out, stream);
+
+    return out;
+}
+
+at::Tensor fp6_unpack(at::Tensor &x, int64_t original_n) {
+    TORCH_CHECK(x.dtype() == torch::kUInt8, "Input tensor must be uint8");
+    TORCH_CHECK(x.is_cuda(), "Input tensor must be on CUDA");
+    TORCH_CHECK(x.is_contiguous(), "Input tensor must be contiguous");
+    TORCH_CHECK(x.dim() == 2, "Input tensor must be 2D [m, n_packed]");
+    TORCH_CHECK(original_n % 8 == 0, "original_n must be divisible by 8, got ", original_n);
+
+    int64_t m = x.size(0);
+    int64_t n_packed = x.size(1);
+
+    TORCH_CHECK(n_packed == original_n * 3 / 4,
+                "Packed size mismatch: expected ", original_n * 3 / 4, " got ", n_packed);
+
+    auto options = torch::TensorOptions()
+        .dtype(torch::kUInt8)
+        .device(x.device());
+
+    at::Tensor out = torch::empty({m, original_n}, options);
+
+    at::cuda::CUDAGuard device_guard{x.get_device()};
+    auto stream = at::cuda::getCurrentCUDAStream().stream();
+
+    fp6_unpack_cuda(x, out, stream);
+
+    return out;
+}

packages/ltx-kernels/csrc/ops/fp6_pack.cu

@@ -0,0 +1,211 @@
+#include <c10/cuda/CUDAException.h>
+#include <cuda_runtime.h>
+#include <cuda.h>
+
+#include <ATen/ATen.h>
+#include <torch/types.h>
+
+// Device function to pack 8-bit to 6-bit
+// 8-bit layout: s e_1 e_2 e_3 m_1 m_2 m_3 m_4 (bits 7-0)
+// 6-bit layout: s e_3 m_1 m_2 m_3 m_4 (bits 5-0)
+// Drop e_1 (bit 6) and e_2 (bit 5)
+__device__ __forceinline__ uint8_t pack_8bit_to_6bit(uint8_t input) {
+    // Extract the sign bit (bit 7)
+    uint8_t sign = (input >> 7) & 0x1;
+
+    // Extract e_3 (bit 4)
+    uint8_t e_3 = (input >> 4) & 0x1;
+
+    // Extract mantissa bits (bits 3-0)
+    uint8_t mantissa = input & 0x0F;
+
+    // Pack into 6-bit format: s e_3 m_1 m_2 m_3 m_4
+    uint8_t result = (sign << 5) | (e_3 << 4) | mantissa;
+
+    return result & 0x3F; // Mask to 6 bits
+}
+
+// Device function to pack 4 x 6-bit values into 3 bytes
+__device__ __forceinline__ void pack_4x6bit_to_3bytes(const uint8_t* input_6bit, uint8_t* output_3bytes) {
+    uint8_t v0 = input_6bit[0] & 0x3F;
+    uint8_t v1 = input_6bit[1] & 0x3F;
+    uint8_t v2 = input_6bit[2] & 0x3F;
+    uint8_t v3 = input_6bit[3] & 0x3F;
+
+    // Pack: [v0: 6 bits][v1: 6 bits][v2: 6 bits][v3: 6 bits] = 24 bits = 3 bytes
+    output_3bytes[0] = (v0 << 2) | (v1 >> 4);
+    output_3bytes[1] = (v1 << 4) | (v2 >> 2);
+    output_3bytes[2] = (v2 << 6) | v3;
+}
+
+// CUDA kernel for packing 2D tensor
+// Input: [m, n] uint8 tensor
+// Output: [m, n*3/4] uint8 tensor
+__global__ void fp6_pack_kernel(
+    const uint8_t* __restrict__ input,
+    uint8_t* __restrict__ output,
+    int m,
+    int n,
+    int n_packed
+) {
+    // Each thread processes one row and 4 elements at a time
+    int row = blockIdx.x;
+    int col_group = blockIdx.y * blockDim.x + threadIdx.x;
+
+    if (row >= m) return;
+
+    // Calculate input and output positions
+    int input_col = col_group * 4;
+    if (input_col >= n) return;
+
+    int output_col = col_group * 3;
+
+    const uint8_t* input_row = input + row * n;
+    uint8_t* output_row = output + row * n_packed;
+
+    uint8_t temp_6bit[4];
+
+    // Pack 4 elements
+    #pragma unroll
+    for (int i = 0; i < 4; i++) {
+        if (input_col + i < n) {
+            temp_6bit[i] = pack_8bit_to_6bit(input_row[input_col + i]);
+        } else {
+            temp_6bit[i] = 0;
+        }
+    }
+
+    // Write 3 bytes to output
+    uint8_t temp_3bytes[3];
+    pack_4x6bit_to_3bytes(temp_6bit, temp_3bytes);
+
+    if (output_col < n_packed) output_row[output_col] = temp_3bytes[0];
+    if (output_col + 1 < n_packed) output_row[output_col + 1] = temp_3bytes[1];
+    if (output_col + 2 < n_packed) output_row[output_col + 2] = temp_3bytes[2];
+}
+
+// Device function to unpack 6-bit to 8-bit
+__device__ __forceinline__ uint8_t unpack_6bit_to_8bit(uint8_t input) {
+    input = input & 0x3F; // Ensure only 6 bits
+
+    uint8_t sign = (input >> 5) & 0x1;
+    uint8_t e_3 = (input >> 4) & 0x1;
+    uint8_t mantissa = input & 0x0F;
+
+    // Reconstruct 8-bit with e_1 and e_2 set to 0
+    uint8_t result = (sign << 7) | (e_3 << 4) | mantissa;
+
+    return result;
+}
+
+// Device function to unpack 3 bytes into 4 x 6-bit values
+__device__ __forceinline__ void unpack_3bytes_to_4x6bit(const uint8_t* input_3bytes, uint8_t* output_6bit) {
+    output_6bit[0] = (input_3bytes[0] >> 2) & 0x3F;
+    output_6bit[1] = ((input_3bytes[0] << 4) | (input_3bytes[1] >> 4)) & 0x3F;
+    output_6bit[2] = ((input_3bytes[1] << 2) | (input_3bytes[2] >> 6)) & 0x3F;
+    output_6bit[3] = input_3bytes[2] & 0x3F;
+}
+
+// CUDA kernel for unpacking 2D tensor
+// Input: [m, n_packed] uint8 tensor
+// Output: [m, n] uint8 tensor
+__global__ void fp6_unpack_kernel(
+    const uint8_t* __restrict__ input,
+    uint8_t* __restrict__ output,
+    int m,
+    int n_packed,
+    int n
+) {
+    // Each thread processes one row and 4 elements at a time
+    int row = blockIdx.x;
+    int col_group = blockIdx.y * blockDim.x + threadIdx.x;
+
+    if (row >= m) return;
+
+    // Calculate input and output positions
+    int input_col = col_group * 3;
+    if (input_col >= n_packed) return;
+
+    int output_col = col_group * 4;
+
+    const uint8_t* input_row = input + row * n_packed;
+    uint8_t* output_row = output + row * n;
+
+    // Read 3 bytes
+    uint8_t temp_3bytes[3];
+    temp_3bytes[0] = (input_col < n_packed) ? input_row[input_col] : 0;
+    temp_3bytes[1] = (input_col + 1 < n_packed) ? input_row[input_col + 1] : 0;
+    temp_3bytes[2] = (input_col + 2 < n_packed) ? input_row[input_col + 2] : 0;
+
+    // Unpack to 4 x 6-bit values
+    uint8_t temp_6bit[4];
+    unpack_3bytes_to_4x6bit(temp_3bytes, temp_6bit);
+
+    // Convert to 8-bit and write
+    #pragma unroll
+    for (int i = 0; i < 4; i++) {
+        if (output_col + i < n) {
+            output_row[output_col + i] = unpack_6bit_to_8bit(temp_6bit[i]);
+        }
+    }
+}
+
+// Host function to launch pack kernel
+void fp6_pack_cuda(
+    at::Tensor& x,
+    at::Tensor& out,
+    cudaStream_t stream
+) {
+    int m = x.size(0);
+    int n = x.size(1);
+    int n_packed = out.size(1);
+
+    const uint8_t* input_ptr = (uint8_t*)x.data_ptr();
+    uint8_t* output_ptr = (uint8_t*)out.data_ptr();
+
+    // Each thread handles 4 input elements -> 3 output bytes
+    int num_groups = (n + 3) / 4;
+
+    int threads = 256;
+    dim3 blocks(m, (num_groups + threads - 1) / threads);
+
+    fp6_pack_kernel<<<blocks, threads, 0, stream>>>(
+        input_ptr,
+        output_ptr,
+        m,
+        n,
+        n_packed
+    );
+
+    C10_CUDA_KERNEL_LAUNCH_CHECK();
+}
+
+// Host function to launch unpack kernel
+void fp6_unpack_cuda(
+    at::Tensor& x,
+    at::Tensor& out,
+    cudaStream_t stream
+) {
+    int m = x.size(0);
+    int n_packed = x.size(1);
+    int n = out.size(1);
+
+    const uint8_t* input_ptr = (uint8_t*)x.data_ptr();
+    uint8_t* output_ptr = (uint8_t*)out.data_ptr();
+
+    // Each thread handles 3 input bytes -> 4 output elements
+    int num_groups = (n + 3) / 4;
+
+    int threads = 256;
+    dim3 blocks(m, (num_groups + threads - 1) / threads);
+
+    fp6_unpack_kernel<<<blocks, threads, 0, stream>>>(
+        input_ptr,
+        output_ptr,
+        m,
+        n_packed,
+        n
+    );
+
+    C10_CUDA_KERNEL_LAUNCH_CHECK();
+}

packages/ltx-kernels/csrc/ops/include/fast_hadamard_transform.h

@@ -0,0 +1,163 @@
+/******************************************************************************
+ * Copyright (c) 2023, Tri Dao.
+ ******************************************************************************/
+
+#pragma once
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+struct HadamardParamsBase {
+    using index_t = int64_t;
+
+    int batch, dim, log_N;
+
+    index_t x_batch_stride;
+    index_t out_batch_stride;
+
+    float scale;
+
+    // Common data pointers.
+    void *__restrict__ x_ptr;
+    void *__restrict__ out_ptr;
+};
+
+struct UnifiedHadamardParamsBase{
+    using index_t = int64_t;
+
+    int batch, dim, log_N;
+    int batch_fma_change;
+
+    index_t x_batch_stride;
+    index_t out_batch_stride;
+    index_t fma_batch_stride;
+    index_t cos_freq_batch_stride;
+    index_t sin_freq_batch_stride;
+    
+    float scale;
+
+    // Common data pointers.
+    void *__restrict__ x_ptr;
+    void *__restrict__ out_ptr;
+    void *__restrict__ out_scales_ptr;
+
+    void *__restrict__ y_scale_ptr;
+    void *__restrict__ z_shift_ptr;
+
+    void *__restrict__ weights_ptr;
+
+    void *__restrict__ cos_freq_ptr;
+    void *__restrict__ sin_freq_ptr;
+
+};
+
+
+struct DequantHadamardParamsBase {
+    using index_t = int64_t;
+
+    int batch, dim, log_N;
+
+    index_t x_batch_stride;
+    index_t out_batch_stride;
+
+    float scale;
+
+    // Common data pointers.
+    void *__restrict__ x_ptr;
+    void *__restrict__ scales_ptr;
+    void *__restrict__ out_ptr;
+};
+
+
+struct QuantHadamardParamsBase {
+    using index_t = int64_t;
+
+    int batch, dim, log_N;
+
+    index_t x_batch_stride;
+    index_t out_batch_stride;
+
+    float scale;
+
+    // Common data pointers.
+    void *__restrict__ x_ptr;
+    void *__restrict__ out_ptr;
+    void *__restrict__ out_scales_ptr;
+};
+
+struct NormFMAHadamardParamsBase {
+    using index_t = int64_t;
+
+    int batch, dim, log_N;
+    int seqlen;
+
+    index_t x_batch_stride;
+    index_t out_batch_stride;
+    index_t fma_batch_stride;
+
+    float scale;
+
+    // Common data pointers.
+    void *__restrict__ x_ptr;
+    void *__restrict__ out_ptr;
+    void *__restrict__ y_scale_ptr;
+    void *__restrict__ z_shift_ptr;
+    void *__restrict__ weights_ptr;  
+};
+
+
+struct NormRopeHadamardParamsBase {
+    using index_t = int64_t;
+
+    int batch, dim, log_N;
+    
+    index_t x_batch_stride;
+    index_t out_batch_stride;
+    index_t cos_freq_batch_stride;
+    index_t sin_freq_batch_stride;
+    
+    float scale;
+
+    // Common data pointers.
+    void *__restrict__ x_ptr;
+    void *__restrict__ out_ptr;
+    void *__restrict__ cos_freq_ptr;
+    void *__restrict__ sin_freq_ptr;
+    void *__restrict__ weights_ptr;  
+};
+
+
+struct NormHadamardParamsBase {
+    using index_t = int64_t;
+
+    int batch, dim, log_N;
+    
+    index_t x_batch_stride;
+    index_t out_batch_stride;
+    
+    float scale;
+
+    // Common data pointers.
+    void *__restrict__ x_ptr;
+    void *__restrict__ out_ptr;
+    void *__restrict__ weights_ptr;  
+};
+
+
+struct RopeHadamardParamsBase {
+    using index_t = int64_t;
+
+    int batch, dim, log_N;
+    
+    index_t x_batch_stride;
+    index_t out_batch_stride;
+    index_t cos_freq_batch_stride;
+    index_t sin_freq_batch_stride;
+    
+    float scale;
+
+    // Common data pointers.
+    void *__restrict__ x_ptr;
+    void *__restrict__ out_ptr;
+    void *__restrict__ cos_freq_ptr;
+    void *__restrict__ sin_freq_ptr;  
+};

packages/ltx-kernels/csrc/ops/include/fast_hadamard_transform_common.h

@@ -0,0 +1,319 @@
+/******************************************************************************
+ * Copyright (c) 2023, Tri Dao.
+ ******************************************************************************/
+
+#pragma once
+
+#include <cuda_bf16.h>
+#include <cuda_fp16.h>
+
+#define FULL_MASK 0xffffffff
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+template<typename TYPE> struct QuantMax {};
+template<> struct QuantMax<int8_t> { static constexpr float value = 127.0; };
+template<> struct QuantMax<at::Float8_e4m3fn> { static constexpr float value = 256.0; };
+
+struct uint8 {
+    uint4 u;
+    uint4 v;
+};
+
+template<int BYTES> struct BytesToType {};
+
+template<>
+struct BytesToType<32> {
+    using Type = uint8;
+    static_assert(sizeof(Type) == 32);
+};
+
+template<> struct BytesToType<16> {
+    using Type = uint4;
+    static_assert(sizeof(Type) == 16);
+};
+
+template<> struct BytesToType<8> {
+    using Type = uint64_t;
+    static_assert(sizeof(Type) == 8);
+};
+
+template<> struct BytesToType<4> {
+    using Type = uint32_t;
+    static_assert(sizeof(Type) == 4);
+};
+
+template<> struct BytesToType<2> {
+    using Type = uint16_t;
+    static_assert(sizeof(Type) == 2);
+};
+
+template<> struct BytesToType<1> {
+    using Type = uint8_t;
+    static_assert(sizeof(Type) == 1);
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<typename T>
+struct SumOp {
+__device__ inline T operator()(T const & x, T const & y) { return x + y; }
+};
+
+template<typename T>
+struct MaxOp {
+__device__ inline T operator()(T const & x, T const & y) { return max(x, y); }
+};
+
+template <>
+struct MaxOp<float> {
+// This is slightly faster
+__device__ inline float operator()(float const &x, float const &y) { return max(x, y); }
+};
+
+
+template<int THREADS>
+struct Allreduce {
+    static_assert(THREADS == 32 || THREADS == 16 || THREADS == 8 || THREADS == 4);
+    template<typename T, typename Operator>
+    static __device__ inline T run(T x, Operator &op) {
+        constexpr int OFFSET = THREADS / 2;
+        x = op(x, __shfl_xor_sync(uint32_t(-1), x, OFFSET));
+        return Allreduce<OFFSET>::run(x, op);
+    }
+};
+
+template<>
+struct Allreduce<2> {
+template<typename T, typename Operator>
+static __device__ inline T run(T x, Operator &op) {
+    x = op(x, __shfl_xor_sync(uint32_t(-1), x, 1));
+    return x;
+}
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// https://stackoverflow.com/questions/35311711/whats-the-right-way-to-compute-integral-base-2-logarithms-at-compile-time
+constexpr int cilog2(int val) { return val > 0 ? 1 + cilog2(val >> 1) : -1; }
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<int kLogN, int kNChunks>
+__device__ __forceinline__ void hadamard_mult_thread(float x[kNChunks][1 << kLogN]) {
+    constexpr int N = 1 << kLogN;
+    #pragma unroll
+    for (int i = 0; i < kLogN; ++i) {
+        const int stride = 1 << i;
+        #pragma unroll
+        for (int j = 0; j < N / 2; ++j) {
+            const int lo = j & (stride - 1);
+            const int idx = (j - lo) * 2 + lo;
+            #pragma unroll
+            for (int c = 0; c < kNChunks; ++c) {
+                const float a = x[c][idx];
+                const float b = x[c][idx + stride];
+                x[c][idx] = a + b;
+                x[c][idx + stride] = a - b;
+            }
+        }
+    }
+}
+
+template<int kLogWarpSize, int kStepStart, int kNChunks, int kNItems>
+__device__ __forceinline__ void hadamard_mult_warp(float x[kNChunks][kNItems]) {
+    constexpr int N = 1 << kLogWarpSize;
+    int lane_id = threadIdx.x % N;
+    #pragma unroll
+    for (int step = kStepStart; step < kLogWarpSize; ++step) {
+        const int lane_mask = 1 << step;
+        const float sign = (lane_id & lane_mask) ? -1.f : 1.f;
+        #pragma unroll
+        for (int c = 0; c < kNChunks; ++c) {
+            #pragma unroll
+            for (int i = 0; i < kNItems; ++i) {
+                float x_val_other = __shfl_xor_sync(FULL_MASK, x[c][i], lane_mask);
+                x[c][i] = sign * x[c][i] + x_val_other;
+            }
+        }
+    }
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <int kNChunks, int kNElts, typename input_t>
+inline __device__ void load_input(input_t *x, float x_vals[kNChunks][kNElts], int dim) {
+    using vec_t = typename BytesToType<sizeof(input_t) * kNElts>::Type;
+    input_t x_vals_load[kNChunks][kNElts] = {0};
+    #pragma unroll
+    for (int c = 0; c < kNChunks; ++c) {
+        if ((c * blockDim.x + threadIdx.x) * kNElts < dim) {
+            reinterpret_cast<vec_t*>(x_vals_load)[c] = reinterpret_cast<const vec_t*>(x)[c * blockDim.x + threadIdx.x];
+        }
+    }
+    #pragma unroll
+    for (int c = 0; c < kNChunks; ++c) {
+        #pragma unroll
+        for (int i = 0; i < kNElts; ++i) { x_vals[c][i] = float(x_vals_load[c][i]); }
+    }
+}
+
+
+template <int kNChunks, int kNElts, typename output_t, bool do_round>
+inline __device__ void store_output(output_t *out, float out_vals[kNChunks][kNElts], int dim, float scale=1.f) {
+    using vec_t = typename BytesToType<sizeof(output_t) * kNElts>::Type;
+    output_t out_vals_store[kNChunks][kNElts];
+    #pragma unroll
+    for (int c = 0; c < kNChunks; ++c) {
+        #pragma unroll
+        for (int i = 0; i < kNElts; ++i) { 
+            if constexpr (do_round){
+                out_vals_store[c][i] = round(out_vals[c][i] * scale);
+            } else {
+                out_vals_store[c][i] = out_vals[c][i] * scale;
+            }
+             
+        }
+    }
+    #pragma unroll
+    for (int c = 0; c < kNChunks; ++c) {
+        if ((c * blockDim.x + threadIdx.x) * kNElts < dim) {
+            reinterpret_cast<vec_t*>(out)[c * blockDim.x + threadIdx.x] = reinterpret_cast<const vec_t*>(out_vals_store)[c];
+        }
+    }
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// Pre=true means the exchange before the hadamard_mult_warp, Pre=false means after.
+template <int kNChunks, int kChunksPerExchange, int kNElts, int kWarpSize, int kNWarps, bool Pre, typename vec_t>
+inline __device__ void exchange_smem_pre(float x_vals[kNChunks][kNElts], vec_t *smem) {
+    constexpr int kNThreads = kWarpSize * kNWarps;
+    constexpr int kNExchangePerVec = kNElts / (sizeof(vec_t) / sizeof(float));
+    const int warp_id = threadIdx.x / kWarpSize;
+    const int lane_id = threadIdx.x % kWarpSize;
+    const int row_t = threadIdx.x % kNWarps;
+    const int col_t = threadIdx.x / kNWarps;
+    // We use the XOR swizzle trick (new_col = col ^ row) to avoid / reduce smem bank conflicts.
+    #pragma unroll
+    for (int c0 = 0; c0 < kNChunks / kChunksPerExchange; ++c0) {
+        __syncthreads();
+        #pragma unroll
+        for (int c1 = 0; c1 < kChunksPerExchange; ++c1) {
+            #pragma unroll
+            for (int r = 0; r < kNExchangePerVec; ++r) {
+                smem[(c1 * kNExchangePerVec + r) * kNThreads + (Pre ? warp_id * kWarpSize + lane_id ^ warp_id : row_t * kWarpSize + col_t ^ row_t)] = reinterpret_cast<vec_t*>(x_vals[c0 * kChunksPerExchange + c1])[r];
+            }
+        }
+        __syncthreads();
+        #pragma unroll
+        for (int c1 = 0; c1 < kChunksPerExchange; ++c1) {
+            #pragma unroll
+            for (int r = 0; r < kNExchangePerVec; ++r) {
+                reinterpret_cast<vec_t*>(x_vals[c0 * kChunksPerExchange + c1])[r] = smem[(c1 * kNExchangePerVec + r) * kNThreads + (Pre ? row_t * kWarpSize + col_t ^ row_t : warp_id * kWarpSize + lane_id ^ warp_id)];
+            }
+        }
+    }
+}
+
+inline __device__ float gelu_approximate(float x){
+    constexpr float sqrthalfpi2 = 0.7978845608028653558798921198687637369517172623298693153318516593f;
+    constexpr float factor = 0.044715f;
+    return 0.5f*x*(1.0f + tanhf(sqrthalfpi2*(x + factor*x*x*x)));
+}
+
+template <int kNChunks, int kNElts>
+inline __device__ void fused_gelu(float x_vals[kNChunks][kNElts]){
+    #pragma unroll
+    for (size_t c = 0; c < kNChunks; c++)
+    {
+        #pragma unroll
+        for (size_t i = 0; i < kNElts; i++)
+        {
+            x_vals[c][i] = gelu_approximate(x_vals[c][i]);
+        }
+        
+    }
+}
+
+template <int kNChunks, int kNElts, int kNWarps, bool norm_affine>
+inline __device__ void fused_rms_norm(float x_vals[kNChunks][kNElts], float weights_vals[kNChunks][kNElts], float* smem_sum, float dim){
+    float thread_squared_sum = 0.0f;
+    const int warp_id = threadIdx.x / 32;
+
+    #pragma unroll
+    for (size_t c = 0; c < kNChunks; c++)
+    {
+        #pragma unroll
+        for (size_t i = 0; i < kNElts; i++)
+        {
+            thread_squared_sum += x_vals[c][i] * x_vals[c][i]; 
+        }
+        
+    }
+    SumOp<float> sum_op;
+    float warp_sum = Allreduce<32>::run(thread_squared_sum, sum_op);
+
+    if(threadIdx.x % 32 == 0){
+        smem_sum[warp_id] = warp_sum;
+    }
+    __syncthreads();
+    float norm = 0.0f;
+    #pragma unroll  
+    for (size_t i = 0; i < kNWarps; i++)
+    {
+        norm += smem_sum[i];
+    }
+    
+    norm *= 1.0f/dim;
+    norm = rsqrtf(norm + 0.0000001f);
+
+    #pragma unroll
+    for (size_t c = 0; c < kNChunks; c++)
+    {
+        #pragma unroll
+        for (size_t i = 0; i < kNElts; i++)
+        {
+            if constexpr (norm_affine){
+                x_vals[c][i] *= (norm * weights_vals[c][i]);
+            } else {
+                x_vals[c][i] *= norm;
+            }
+        }
+    }
+}
+
+
+template <int kNChunks, int kNElts>
+inline __device__ void fused_rope(float x_vals[kNChunks][kNElts], float sin_freqs_vals[kNChunks][kNElts], float cos_freqs_vals[kNChunks][kNElts]){
+    #pragma unroll
+    for (size_t c = 0; c < kNChunks; c++)
+    {
+        #pragma unroll
+        for (size_t i = 0; i < kNElts; i+=2)
+        {
+            float x_1 = x_vals[c][i];
+            float x_2 = x_vals[c][i+1];
+            x_vals[c][i] = -x_2*sin_freqs_vals[c][i] + x_1*cos_freqs_vals[c][i];
+            x_vals[c][i+1] = x_1*sin_freqs_vals[c][i+1] + x_2*cos_freqs_vals[c][i+1];
+        }
+    }
+}
+
+
+template <int kNChunks, int kNElts, bool add_one_scale>
+inline __device__ void fused_multiply_add(float x_vals[kNChunks][kNElts], float y_scale_vals[kNChunks][kNElts], float z_shift_vals[kNChunks][kNElts]) {
+
+    #pragma unroll
+    for (size_t c = 0; c < kNChunks; c++)
+    {
+        #pragma unroll
+        for (size_t i = 0; i < kNElts; i++)
+        {
+            if constexpr (add_one_scale){
+                x_vals[c][i] = x_vals[c][i] * (1.0f + y_scale_vals[c][i]) + z_shift_vals[c][i];
+            } else {
+                x_vals[c][i] = x_vals[c][i] * y_scale_vals[c][i] + z_shift_vals[c][i];
+            }    
+        }
+    }
+}

packages/ltx-kernels/csrc/ops/include/fast_hadamard_transform_special.h

@@ -0,0 +1,138 @@
+
+/******************************************************************************
+ * Copyright (c) 2023, Tri Dao.
+ ******************************************************************************/
+
+// This file is auto-generated. See "code_gen.py"
+
+
+#pragma once
+
+
+__device__ __forceinline__ void hadamard_mult_thread_12(float x[12]) {
+    float out[12];
+    out[0] = + x[0] - x[1] + x[2] + x[3] + x[4] + x[5] + x[6] + x[7] + x[8] + x[9] + x[10] + x[11];
+    out[1] = - x[0] - x[1] + x[2] - x[3] + x[4] - x[5] + x[6] - x[7] + x[8] - x[9] + x[10] - x[11];
+    out[2] = + x[0] + x[1] + x[2] - x[3] + x[4] + x[5] - x[6] - x[7] - x[8] - x[9] + x[10] + x[11];
+    out[3] = + x[0] - x[1] - x[2] - x[3] + x[4] - x[5] - x[6] + x[7] - x[8] + x[9] + x[10] - x[11];
+    out[4] = + x[0] + x[1] + x[2] + x[3] + x[4] - x[5] + x[6] + x[7] - x[8] - x[9] - x[10] - x[11];
+    out[5] = + x[0] - x[1] + x[2] - x[3] - x[4] - x[5] + x[6] - x[7] - x[8] + x[9] - x[10] + x[11];
+    out[6] = + x[0] + x[1] - x[2] - x[3] + x[4] + x[5] + x[6] - x[7] + x[8] + x[9] - x[10] - x[11];
+    out[7] = + x[0] - x[1] - x[2] + x[3] + x[4] - x[5] - x[6] - x[7] + x[8] - x[9] - x[10] + x[11];
+    out[8] = + x[0] + x[1] - x[2] - x[3] - x[4] - x[5] + x[6] + x[7] + x[8] - x[9] + x[10] + x[11];
+    out[9] = + x[0] - x[1] - x[2] + x[3] - x[4] + x[5] + x[6] - x[7] - x[8] - x[9] + x[10] - x[11];
+    out[10] = + x[0] + x[1] + x[2] + x[3] - x[4] - x[5] - x[6] - x[7] + x[8] + x[9] + x[10] - x[11];
+    out[11] = + x[0] - x[1] + x[2] - x[3] - x[4] + x[5] - x[6] + x[7] + x[8] - x[9] - x[10] - x[11];
+    #pragma unroll
+    for (int i = 0; i < 12; i++) { x[i] = out[i]; }
+}
+
+
+__device__ __forceinline__ void hadamard_mult_thread_20(float x[20]) {
+    float out[20];
+    out[0] = + x[0] - x[1] - x[2] - x[3] - x[4] + x[5] - x[6] - x[7] - x[8] - x[9] + x[10] + x[11] - x[12] - x[13] + x[14] + x[15] - x[16] + x[17] + x[18] - x[19];
+    out[1] = - x[0] + x[1] - x[2] - x[3] - x[4] - x[5] + x[6] - x[7] - x[8] - x[9] + x[10] + x[11] + x[12] - x[13] - x[14] - x[15] + x[16] - x[17] + x[18] + x[19];
+    out[2] = - x[0] - x[1] + x[2] - x[3] - x[4] - x[5] - x[6] + x[7] - x[8] - x[9] - x[10] + x[11] + x[12] + x[13] - x[14] + x[15] - x[16] + x[17] - x[18] + x[19];
+    out[3] = - x[0] - x[1] - x[2] + x[3] - x[4] - x[5] - x[6] - x[7] + x[8] - x[9] - x[10] - x[11] + x[12] + x[13] + x[14] + x[15] + x[16] - x[17] + x[18] - x[19];
+    out[4] = - x[0] - x[1] - x[2] - x[3] + x[4] - x[5] - x[6] - x[7] - x[8] + x[9] + x[10] - x[11] - x[12] + x[13] + x[14] - x[15] + x[16] + x[17] - x[18] + x[19];
+    out[5] = - x[0] + x[1] + x[2] + x[3] + x[4] + x[5] - x[6] - x[7] - x[8] - x[9] - x[10] + x[11] - x[12] - x[13] + x[14] + x[15] + x[16] - x[17] - x[18] + x[19];
+    out[6] = + x[0] - x[1] + x[2] + x[3] + x[4] - x[5] + x[6] - x[7] - x[8] - x[9] + x[10] - x[11] + x[12] - x[13] - x[14] + x[15] + x[16] + x[17] - x[18] - x[19];
+    out[7] = + x[0] + x[1] - x[2] + x[3] + x[4] - x[5] - x[6] + x[7] - x[8] - x[9] - x[10] + x[11] - x[12] + x[13] - x[14] - x[15] + x[16] + x[17] + x[18] - x[19];
+    out[8] = + x[0] + x[1] + x[2] - x[3] + x[4] - x[5] - x[6] - x[7] + x[8] - x[9] - x[10] - x[11] + x[12] - x[13] + x[14] - x[15] - x[16] + x[17] + x[18] + x[19];
+    out[9] = + x[0] + x[1] + x[2] + x[3] - x[4] - x[5] - x[6] - x[7] - x[8] + x[9] + x[10] - x[11] - x[12] + x[13] - x[14] + x[15] - x[16] - x[17] + x[18] + x[19];
+    out[10] = - x[0] - x[1] + x[2] + x[3] - x[4] + x[5] - x[6] + x[7] + x[8] - x[9] + x[10] - x[11] - x[12] - x[13] - x[14] - x[15] + x[16] + x[17] + x[18] + x[19];
+    out[11] = - x[0] - x[1] - x[2] + x[3] + x[4] - x[5] + x[6] - x[7] + x[8] + x[9] - x[10] + x[11] - x[12] - x[13] - x[14] + x[15] - x[16] + x[17] + x[18] + x[19];
+    out[12] = + x[0] - x[1] - x[2] - x[3] + x[4] + x[5] - x[6] + x[7] - x[8] + x[9] - x[10] - x[11] + x[12] - x[13] - x[14] + x[15] + x[16] - x[17] + x[18] + x[19];
+    out[13] = + x[0] + x[1] - x[2] - x[3] - x[4] + x[5] + x[6] - x[7] + x[8] - x[9] - x[10] - x[11] - x[12] + x[13] - x[14] + x[15] + x[16] + x[17] - x[18] + x[19];
+    out[14] = - x[0] + x[1] + x[2] - x[3] - x[4] - x[5] + x[6] + x[7] - x[8] + x[9] - x[10] - x[11] - x[12] - x[13] + x[14] + x[15] + x[16] + x[17] + x[18] - x[19];
+    out[15] = - x[0] + x[1] - x[2] - x[3] + x[4] - x[5] - x[6] + x[7] + x[8] - x[9] + x[10] - x[11] - x[12] - x[13] - x[14] + x[15] - x[16] - x[17] - x[18] - x[19];
+    out[16] = + x[0] - x[1] + x[2] - x[3] - x[4] - x[5] - x[6] - x[7] + x[8] + x[9] - x[10] + x[11] - x[12] - x[13] - x[14] - x[15] + x[16] - x[17] - x[18] - x[19];
+    out[17] = - x[0] + x[1] - x[2] + x[3] - x[4] + x[5] - x[6] - x[7] - x[8] + x[9] - x[10] - x[11] + x[12] - x[13] - x[14] - x[15] - x[16] + x[17] - x[18] - x[19];
+    out[18] = - x[0] - x[1] + x[2] - x[3] + x[4] + x[5] + x[6] - x[7] - x[8] - x[9] - x[10] - x[11] - x[12] + x[13] - x[14] - x[15] - x[16] - x[17] + x[18] - x[19];
+    out[19] = + x[0] - x[1] - x[2] + x[3] - x[4] - x[5] + x[6] + x[7] - x[8] - x[9] - x[10] - x[11] - x[12] - x[13] + x[14] - x[15] - x[16] - x[17] - x[18] + x[19];
+    #pragma unroll
+    for (int i = 0; i < 20; i++) { x[i] = out[i]; }
+}
+
+
+__device__ __forceinline__ void hadamard_mult_thread_28(float x[28]) {
+    float out[28];
+    out[0] = + x[0] - x[1] - x[2] - x[3] - x[4] - x[5] - x[6] + x[7] + x[8] - x[9] - x[10] - x[11] - x[12] + x[13] + x[14] - x[15] + x[16] - x[17] - x[18] + x[19] - x[20] + x[21] - x[22] - x[23] + x[24] + x[25] - x[26] - x[27];
+    out[1] = - x[0] + x[1] - x[2] - x[3] - x[4] - x[5] - x[6] + x[7] + x[8] + x[9] - x[10] - x[11] - x[12] - x[13] - x[14] + x[15] - x[16] + x[17] - x[18] - x[19] + x[20] - x[21] + x[22] - x[23] - x[24] + x[25] + x[26] - x[27];
+    out[2] = - x[0] - x[1] + x[2] - x[3] - x[4] - x[5] - x[6] - x[7] + x[8] + x[9] + x[10] - x[11] - x[12] - x[13] + x[14] - x[15] + x[16] - x[17] + x[18] - x[19] - x[20] - x[21] - x[22] + x[23] - x[24] - x[25] + x[26] + x[27];
+    out[3] = - x[0] - x[1] - x[2] + x[3] - x[4] - x[5] - x[6] - x[7] - x[8] + x[9] + x[10] + x[11] - x[12] - x[13] - x[14] + x[15] - x[16] + x[17] - x[18] + x[19] - x[20] + x[21] - x[22] - x[23] + x[24] - x[25] - x[26] + x[27];
+    out[4] = - x[0] - x[1] - x[2] - x[3] + x[4] - x[5] - x[6] - x[7] - x[8] - x[9] + x[10] + x[11] + x[12] - x[13] - x[14] - x[15] + x[16] - x[17] + x[18] - x[19] + x[20] + x[21] + x[22] - x[23] - x[24] + x[25] - x[26] - x[27];
+    out[5] = - x[0] - x[1] - x[2] - x[3] - x[4] + x[5] - x[6] - x[7] - x[8] - x[9] - x[10] + x[11] + x[12] + x[13] + x[14] - x[15] - x[16] + x[17] - x[18] + x[19] - x[20] - x[21] + x[22] + x[23] - x[24] - x[25] + x[26] - x[27];
+    out[6] = - x[0] - x[1] - x[2] - x[3] - x[4] - x[5] + x[6] + x[7] - x[8] - x[9] - x[10] - x[11] + x[12] + x[13] - x[14] + x[15] - x[16] - x[17] + x[18] - x[19] + x[20] - x[21] - x[22] + x[23] + x[24] - x[25] - x[26] + x[27];
+    out[7] = - x[0] - x[1] + x[2] + x[3] + x[4] + x[5] - x[6] + x[7] - x[8] - x[9] - x[10] - x[11] - x[12] - x[13] - x[14] + x[15] + x[16] - x[17] - x[18] + x[19] + x[20] + x[21] - x[22] + x[23] - x[24] - x[25] + x[26] - x[27];
+    out[8] = - x[0] - x[1] - x[2] + x[3] + x[4] + x[5] + x[6] - x[7] + x[8] - x[9] - x[10] - x[11] - x[12] - x[13] + x[14] - x[15] + x[16] + x[17] - x[18] - x[19] + x[20] - x[21] + x[22] - x[23] + x[24] - x[25] - x[26] + x[27];
+    out[9] = + x[0] - x[1] - x[2] - x[3] + x[4] + x[5] + x[6] - x[7] - x[8] + x[9] - x[10] - x[11] - x[12] - x[13] + x[14] + x[15] - x[16] + x[17] + x[18] - x[19] - x[20] + x[21] - x[22] + x[23] - x[24] + x[25] - x[26] - x[27];
+    out[10] = + x[0] + x[1] - x[2] - x[3] - x[4] + x[5] + x[6] - x[7] - x[8] - x[9] + x[10] - x[11] - x[12] - x[13] - x[14] + x[15] + x[16] - x[17] + x[18] + x[19] - x[20] - x[21] + x[22] - x[23] + x[24] - x[25] + x[26] - x[27];
+    out[11] = + x[0] + x[1] + x[2] - x[3] - x[4] - x[5] + x[6] - x[7] - x[8] - x[9] - x[10] + x[11] - x[12] - x[13] - x[14] - x[15] + x[16] + x[17] - x[18] + x[19] + x[20] - x[21] - x[22] + x[23] - x[24] + x[25] - x[26] + x[27];
+    out[12] = + x[0] + x[1] + x[2] + x[3] - x[4] - x[5] - x[6] - x[7] - x[8] - x[9] - x[10] - x[11] + x[12] - x[13] + x[14] - x[15] - x[16] + x[17] + x[18] - x[19] + x[20] + x[21] - x[22] - x[23] + x[24] - x[25] + x[26] - x[27];
+    out[13] = - x[0] + x[1] + x[2] + x[3] + x[4] - x[5] - x[6] - x[7] - x[8] - x[9] - x[10] - x[11] - x[12] + x[13] + x[14] + x[15] - x[16] - x[17] + x[18] + x[19] - x[20] - x[21] + x[22] - x[23] - x[24] + x[25] - x[26] + x[27];
+    out[14] = - x[0] + x[1] - x[2] + x[3] + x[4] - x[5] + x[6] + x[7] - x[8] - x[9] + x[10] + x[11] - x[12] - x[13] + x[14] - x[15] - x[16] - x[17] - x[18] - x[19] - x[20] - x[21] - x[22] + x[23] + x[24] + x[25] + x[26] - x[27];
+    out[15] = + x[0] - x[1] + x[2] - x[3] + x[4] + x[5] - x[6] - x[7] + x[8] - x[9] - x[10] + x[11] + x[12] - x[13] - x[14] + x[15] - x[16] - x[17] - x[18] - x[19] - x[20] - x[21] - x[22] - x[23] + x[24] + x[25] + x[26] + x[27];
+    out[16] = - x[0] + x[1] - x[2] + x[3] - x[4] + x[5] + x[6] - x[7] - x[8] + x[9] - x[10] - x[11] + x[12] + x[13] - x[14] - x[15] + x[16] - x[17] - x[18] - x[19] - x[20] + x[21] - x[22] - x[23] - x[24] + x[25] + x[26] + x[27];
+    out[17] = + x[0] - x[1] + x[2] - x[3] + x[4] - x[5] + x[6] + x[7] - x[8] - x[9] + x[10] - x[11] - x[12] + x[13] - x[14] - x[15] - x[16] + x[17] - x[18] - x[19] - x[20] + x[21] + x[22] - x[23] - x[24] - x[25] + x[26] + x[27];
+    out[18] = + x[0] + x[1] - x[2] + x[3] - x[4] + x[5] - x[6] + x[7] + x[8] - x[9] - x[10] + x[11] - x[12] - x[13] - x[14] - x[15] - x[16] - x[17] + x[18] - x[19] - x[20] + x[21] + x[22] + x[23] - x[24] - x[25] - x[26] + x[27];
+    out[19] = - x[0] + x[1] + x[2] - x[3] + x[4] - x[5] + x[6] - x[7] + x[8] + x[9] - x[10] - x[11] + x[12] - x[13] - x[14] - x[15] - x[16] - x[17] - x[18] + x[19] - x[20] + x[21] + x[22] + x[23] + x[24] - x[25] - x[26] - x[27];
+    out[20] = + x[0] - x[1] + x[2] + x[3] - x[4] + x[5] - x[6] - x[7] - x[8] + x[9] + x[10] - x[11] - x[12] + x[13] - x[14] - x[15] - x[16] - x[17] - x[18] - x[19] + x[20] - x[21] + x[22] + x[23] + x[24] + x[25] - x[26] - x[27];
+    out[21] = - x[0] + x[1] + x[2] - x[3] - x[4] + x[5] + x[6] - x[7] + x[8] - x[9] + x[10] + x[11] - x[12] + x[13] + x[14] + x[15] - x[16] - x[17] - x[18] - x[19] + x[20] + x[21] - x[22] - x[23] - x[24] - x[25] - x[26] - x[27];
+    out[22] = + x[0] - x[1] + x[2] + x[3] - x[4] - x[5] + x[6] + x[7] - x[8] + x[9] - x[10] + x[11] + x[12] - x[13] + x[14] + x[15] + x[16] - x[17] - x[18] - x[19] - x[20] - x[21] + x[22] - x[23] - x[24] - x[25] - x[26] - x[27];
+    out[23] = + x[0] + x[1] - x[2] + x[3] + x[4] - x[5] - x[6] - x[7] + x[8] - x[9] + x[10] - x[11] + x[12] + x[13] - x[14] + x[15] + x[16] + x[17] - x[18] - x[19] - x[20] - x[21] - x[22] + x[23] - x[24] - x[25] - x[26] - x[27];
+    out[24] = - x[0] + x[1] + x[2] - x[3] + x[4] + x[5] - x[6] + x[7] - x[8] + x[9] - x[10] + x[11] - x[12] + x[13] - x[14] - x[15] + x[16] + x[17] + x[18] - x[19] - x[20] - x[21] - x[22] - x[23] + x[24] - x[25] - x[26] - x[27];
+    out[25] = - x[0] - x[1] + x[2] + x[3] - x[4] + x[5] + x[6] + x[7] + x[8] - x[9] + x[10] - x[11] + x[12] - x[13] - x[14] - x[15] - x[16] + x[17] + x[18] + x[19] - x[20] - x[21] - x[22] - x[23] - x[24] + x[25] - x[26] - x[27];
+    out[26] = + x[0] - x[1] - x[2] + x[3] + x[4] - x[5] + x[6] - x[7] + x[8] + x[9] - x[10] + x[11] - x[12] + x[13] - x[14] - x[15] - x[16] - x[17] + x[18] + x[19] + x[20] - x[21] - x[22] - x[23] - x[24] - x[25] + x[26] - x[27];
+    out[27] = + x[0] + x[1] - x[2] - x[3] + x[4] + x[5] - x[6] + x[7] - x[8] + x[9] + x[10] - x[11] + x[12] - x[13] + x[14] - x[15] - x[16] - x[17] - x[18] + x[19] + x[20] - x[21] - x[22] - x[23] - x[24] - x[25] - x[26] + x[27];
+    #pragma unroll
+    for (int i = 0; i < 28; i++) { x[i] = out[i]; }
+}
+
+
+__device__ __forceinline__ void hadamard_mult_thread_40(float x[40]) {
+    float out[40];
+    out[0] = + x[0] - x[1] - x[2] - x[3] - x[4] - x[5] - x[6] - x[7] - x[8] - x[9] - x[10] - x[11] - x[12] - x[13] - x[14] - x[15] - x[16] - x[17] - x[18] - x[19] + x[20] - x[21] - x[22] - x[23] - x[24] - x[25] - x[26] - x[27] - x[28] - x[29] - x[30] - x[31] - x[32] - x[33] - x[34] - x[35] - x[36] - x[37] - x[38] - x[39];
+    out[1] = + x[0] + x[1] - x[2] + x[3] + x[4] - x[5] - x[6] - x[7] - x[8] + x[9] - x[10] + x[11] - x[12] + x[13] + x[14] + x[15] + x[16] - x[17] - x[18] + x[19] + x[20] + x[21] - x[22] + x[23] + x[24] - x[25] - x[26] - x[27] - x[28] + x[29] - x[30] + x[31] - x[32] + x[33] + x[34] + x[35] + x[36] - x[37] - x[38] + x[39];
+    out[2] = + x[0] + x[1] + x[2] - x[3] + x[4] + x[5] - x[6] - x[7] - x[8] - x[9] + x[10] - x[11] + x[12] - x[13] + x[14] + x[15] + x[16] + x[17] - x[18] - x[19] + x[20] + x[21] + x[22] - x[23] + x[24] + x[25] - x[26] - x[27] - x[28] - x[29] + x[30] - x[31] + x[32] - x[33] + x[34] + x[35] + x[36] + x[37] - x[38] - x[39];
+    out[3] = + x[0] - x[1] + x[2] + x[3] - x[4] + x[5] + x[6] - x[7] - x[8] - x[9] - x[10] + x[11] - x[12] + x[13] - x[14] + x[15] + x[16] + x[17] + x[18] - x[19] + x[20] - x[21] + x[22] + x[23] - x[24] + x[25] + x[26] - x[27] - x[28] - x[29] - x[30] + x[31] - x[32] + x[33] - x[34] + x[35] + x[36] + x[37] + x[38] - x[39];
+    out[4] = + x[0] - x[1] - x[2] + x[3] + x[4] - x[5] + x[6] + x[7] - x[8] - x[9] - x[10] - x[11] + x[12] - x[13] + x[14] - x[15] + x[16] + x[17] + x[18] + x[19] + x[20] - x[21] - x[22] + x[23] + x[24] - x[25] + x[26] + x[27] - x[28] - x[29] - x[30] - x[31] + x[32] - x[33] + x[34] - x[35] + x[36] + x[37] + x[38] + x[39];
+    out[5] = + x[0] + x[1] - x[2] - x[3] + x[4] + x[5] - x[6] + x[7] + x[8] - x[9] - x[10] - x[11] - x[12] + x[13] - x[14] + x[15] - x[16] + x[17] + x[18] + x[19] + x[20] + x[21] - x[22] - x[23] + x[24] + x[25] - x[26] + x[27] + x[28] - x[29] - x[30] - x[31] - x[32] + x[33] - x[34] + x[35] - x[36] + x[37] + x[38] + x[39];
+    out[6] = + x[0] + x[1] + x[2] - x[3] - x[4] + x[5] + x[6] - x[7] + x[8] + x[9] - x[10] - x[11] - x[12] - x[13] + x[14] - x[15] + x[16] - x[17] + x[18] + x[19] + x[20] + x[21] + x[22] - x[23] - x[24] + x[25] + x[26] - x[27] + x[28] + x[29] - x[30] - x[31] - x[32] - x[33] + x[34] - x[35] + x[36] - x[37] + x[38] + x[39];
+    out[7] = + x[0] + x[1] + x[2] + x[3] - x[4] - x[5] + x[6] + x[7] - x[8] + x[9] + x[10] - x[11] - x[12] - x[13] - x[14] + x[15] - x[16] + x[17] - x[18] + x[19] + x[20] + x[21] + x[22] + x[23] - x[24] - x[25] + x[26] + x[27] - x[28] + x[29] + x[30] - x[31] - x[32] - x[33] - x[34] + x[35] - x[36] + x[37] - x[38] + x[39];
+    out[8] = + x[0] + x[1] + x[2] + x[3] + x[4] - x[5] - x[6] + x[7] + x[8] - x[9] + x[10] + x[11] - x[12] - x[13] - x[14] - x[15] + x[16] - x[17] + x[18] - x[19] + x[20] + x[21] + x[22] + x[23] + x[24] - x[25] - x[26] + x[27] + x[28] - x[29] + x[30] + x[31] - x[32] - x[33] - x[34] - x[35] + x[36] - x[37] + x[38] - x[39];
+    out[9] = + x[0] - x[1] + x[2] + x[3] + x[4] + x[5] - x[6] - x[7] + x[8] + x[9] - x[10] + x[11] + x[12] - x[13] - x[14] - x[15] - x[16] + x[17] - x[18] + x[19] + x[20] - x[21] + x[22] + x[23] + x[24] + x[25] - x[26] - x[27] + x[28] + x[29] - x[30] + x[31] + x[32] - x[33] - x[34] - x[35] - x[36] + x[37] - x[38] + x[39];
+    out[10] = + x[0] + x[1] - x[2] + x[3] + x[4] + x[5] + x[6] - x[7] - x[8] + x[9] + x[10] - x[11] + x[12] + x[13] - x[14] - x[15] - x[16] - x[17] + x[18] - x[19] + x[20] + x[21] - x[22] + x[23] + x[24] + x[25] + x[26] - x[27] - x[28] + x[29] + x[30] - x[31] + x[32] + x[33] - x[34] - x[35] - x[36] - x[37] + x[38] - x[39];
+    out[11] = + x[0] - x[1] + x[2] - x[3] + x[4] + x[5] + x[6] + x[7] - x[8] - x[9] + x[10] + x[11] - x[12] + x[13] + x[14] - x[15] - x[16] - x[17] - x[18] + x[19] + x[20] - x[21] + x[22] - x[23] + x[24] + x[25] + x[26] + x[27] - x[28] - x[29] + x[30] + x[31] - x[32] + x[33] + x[34] - x[35] - x[36] - x[37] - x[38] + x[39];
+    out[12] = + x[0] + x[1] - x[2] + x[3] - x[4] + x[5] + x[6] + x[7] + x[8] - x[9] - x[10] + x[11] + x[12] - x[13] + x[14] + x[15] - x[16] - x[17] - x[18] - x[19] + x[20] + x[21] - x[22] + x[23] - x[24] + x[25] + x[26] + x[27] + x[28] - x[29] - x[30] + x[31] + x[32] - x[33] + x[34] + x[35] - x[36] - x[37] - x[38] - x[39];
+    out[13] = + x[0] - x[1] + x[2] - x[3] + x[4] - x[5] + x[6] + x[7] + x[8] + x[9] - x[10] - x[11] + x[12] + x[13] - x[14] + x[15] + x[16] - x[17] - x[18] - x[19] + x[20] - x[21] + x[22] - x[23] + x[24] - x[25] + x[26] + x[27] + x[28] + x[29] - x[30] - x[31] + x[32] + x[33] - x[34] + x[35] + x[36] - x[37] - x[38] - x[39];
+    out[14] = + x[0] - x[1] - x[2] + x[3] - x[4] + x[5] - x[6] + x[7] + x[8] + x[9] + x[10] - x[11] - x[12] + x[13] + x[14] - x[15] + x[16] + x[17] - x[18] - x[19] + x[20] - x[21] - x[22] + x[23] - x[24] + x[25] - x[26] + x[27] + x[28] + x[29] + x[30] - x[31] - x[32] + x[33] + x[34] - x[35] + x[36] + x[37] - x[38] - x[39];
+    out[15] = + x[0] - x[1] - x[2] - x[3] + x[4] - x[5] + x[6] - x[7] + x[8] + x[9] + x[10] + x[11] - x[12] - x[13] + x[14] + x[15] - x[16] + x[17] + x[18] - x[19] + x[20] - x[21] - x[22] - x[23] + x[24] - x[25] + x[26] - x[27] + x[28] + x[29] + x[30] + x[31] - x[32] - x[33] + x[34] + x[35] - x[36] + x[37] + x[38] - x[39];
+    out[16] = + x[0] - x[1] - x[2] - x[3] - x[4] + x[5] - x[6] + x[7] - x[8] + x[9] + x[10] + x[11] + x[12] - x[13] - x[14] + x[15] + x[16] - x[17] + x[18] + x[19] + x[20] - x[21] - x[22] - x[23] - x[24] + x[25] - x[26] + x[27] - x[28] + x[29] + x[30] + x[31] + x[32] - x[33] - x[34] + x[35] + x[36] - x[37] + x[38] + x[39];
+    out[17] = + x[0] + x[1] - x[2] - x[3] - x[4] - x[5] + x[6] - x[7] + x[8] - x[9] + x[10] + x[11] + x[12] + x[13] - x[14] - x[15] + x[16] + x[17] - x[18] + x[19] + x[20] + x[21] - x[22] - x[23] - x[24] - x[25] + x[26] - x[27] + x[28] - x[29] + x[30] + x[31] + x[32] + x[33] - x[34] - x[35] + x[36] + x[37] - x[38] + x[39];
+    out[18] = + x[0] + x[1] + x[2] - x[3] - x[4] - x[5] - x[6] + x[7] - x[8] + x[9] - x[10] + x[11] + x[12] + x[13] + x[14] - x[15] - x[16] + x[17] + x[18] - x[19] + x[20] + x[21] + x[22] - x[23] - x[24] - x[25] - x[26] + x[27] - x[28] + x[29] - x[30] + x[31] + x[32] + x[33] + x[34] - x[35] - x[36] + x[37] + x[38] - x[39];
+    out[19] = + x[0] - x[1] + x[2] + x[3] - x[4] - x[5] - x[6] - x[7] + x[8] - x[9] + x[10] - x[11] + x[12] + x[13] + x[14] + x[15] - x[16] - x[17] + x[18] + x[19] + x[20] - x[21] + x[22] + x[23] - x[24] - x[25] - x[26] - x[27] + x[28] - x[29] + x[30] - x[31] + x[32] + x[33] + x[34] + x[35] - x[36] - x[37] + x[38] + x[39];
+    out[20] = + x[0] - x[1] - x[2] - x[3] - x[4] - x[5] - x[6] - x[7] - x[8] - x[9] - x[10] - x[11] - x[12] - x[13] - x[14] - x[15] - x[16] - x[17] - x[18] - x[19] - x[20] + x[21] + x[22] + x[23] + x[24] + x[25] + x[26] + x[27] + x[28] + x[29] + x[30] + x[31] + x[32] + x[33] + x[34] + x[35] + x[36] + x[37] + x[38] + x[39];
+    out[21] = + x[0] + x[1] - x[2] + x[3] + x[4] - x[5] - x[6] - x[7] - x[8] + x[9] - x[10] + x[11] - x[12] + x[13] + x[14] + x[15] + x[16] - x[17] - x[18] + x[19] - x[20] - x[21] + x[22] - x[23] - x[24] + x[25] + x[26] + x[27] + x[28] - x[29] + x[30] - x[31] + x[32] - x[33] - x[34] - x[35] - x[36] + x[37] + x[38] - x[39];
+    out[22] = + x[0] + x[1] + x[2] - x[3] + x[4] + x[5] - x[6] - x[7] - x[8] - x[9] + x[10] - x[11] + x[12] - x[13] + x[14] + x[15] + x[16] + x[17] - x[18] - x[19] - x[20] - x[21] - x[22] + x[23] - x[24] - x[25] + x[26] + x[27] + x[28] + x[29] - x[30] + x[31] - x[32] + x[33] - x[34] - x[35] - x[36] - x[37] + x[38] + x[39];
+    out[23] = + x[0] - x[1] + x[2] + x[3] - x[4] + x[5] + x[6] - x[7] - x[8] - x[9] - x[10] + x[11] - x[12] + x[13] - x[14] + x[15] + x[16] + x[17] + x[18] - x[19] - x[20] + x[21] - x[22] - x[23] + x[24] - x[25] - x[26] + x[27] + x[28] + x[29] + x[30] - x[31] + x[32] - x[33] + x[34] - x[35] - x[36] - x[37] - x[38] + x[39];
+    out[24] = + x[0] - x[1] - x[2] + x[3] + x[4] - x[5] + x[6] + x[7] - x[8] - x[9] - x[10] - x[11] + x[12] - x[13] + x[14] - x[15] + x[16] + x[17] + x[18] + x[19] - x[20] + x[21] + x[22] - x[23] - x[24] + x[25] - x[26] - x[27] + x[28] + x[29] + x[30] + x[31] - x[32] + x[33] - x[34] + x[35] - x[36] - x[37] - x[38] - x[39];
+    out[25] = + x[0] + x[1] - x[2] - x[3] + x[4] + x[5] - x[6] + x[7] + x[8] - x[9] - x[10] - x[11] - x[12] + x[13] - x[14] + x[15] - x[16] + x[17] + x[18] + x[19] - x[20] - x[21] + x[22] + x[23] - x[24] - x[25] + x[26] - x[27] - x[28] + x[29] + x[30] + x[31] + x[32] - x[33] + x[34] - x[35] + x[36] - x[37] - x[38] - x[39];
+    out[26] = + x[0] + x[1] + x[2] - x[3] - x[4] + x[5] + x[6] - x[7] + x[8] + x[9] - x[10] - x[11] - x[12] - x[13] + x[14] - x[15] + x[16] - x[17] + x[18] + x[19] - x[20] - x[21] - x[22] + x[23] + x[24] - x[25] - x[26] + x[27] - x[28] - x[29] + x[30] + x[31] + x[32] + x[33] - x[34] + x[35] - x[36] + x[37] - x[38] - x[39];
+    out[27] = + x[0] + x[1] + x[2] + x[3] - x[4] - x[5] + x[6] + x[7] - x[8] + x[9] + x[10] - x[11] - x[12] - x[13] - x[14] + x[15] - x[16] + x[17] - x[18] + x[19] - x[20] - x[21] - x[22] - x[23] + x[24] + x[25] - x[26] - x[27] + x[28] - x[29] - x[30] + x[31] + x[32] + x[33] + x[34] - x[35] + x[36] - x[37] + x[38] - x[39];
+    out[28] = + x[0] + x[1] + x[2] + x[3] + x[4] - x[5] - x[6] + x[7] + x[8] - x[9] + x[10] + x[11] - x[12] - x[13] - x[14] - x[15] + x[16] - x[17] + x[18] - x[19] - x[20] - x[21] - x[22] - x[23] - x[24] + x[25] + x[26] - x[27] - x[28] + x[29] - x[30] - x[31] + x[32] + x[33] + x[34] + x[35] - x[36] + x[37] - x[38] + x[39];
+    out[29] = + x[0] - x[1] + x[2] + x[3] + x[4] + x[5] - x[6] - x[7] + x[8] + x[9] - x[10] + x[11] + x[12] - x[13] - x[14] - x[15] - x[16] + x[17] - x[18] + x[19] - x[20] + x[21] - x[22] - x[23] - x[24] - x[25] + x[26] + x[27] - x[28] - x[29] + x[30] - x[31] - x[32] + x[33] + x[34] + x[35] + x[36] - x[37] + x[38] - x[39];
+    out[30] = + x[0] + x[1] - x[2] + x[3] + x[4] + x[5] + x[6] - x[7] - x[8] + x[9] + x[10] - x[11] + x[12] + x[13] - x[14] - x[15] - x[16] - x[17] + x[18] - x[19] - x[20] - x[21] + x[22] - x[23] - x[24] - x[25] - x[26] + x[27] + x[28] - x[29] - x[30] + x[31] - x[32] - x[33] + x[34] + x[35] + x[36] + x[37] - x[38] + x[39];
+    out[31] = + x[0] - x[1] + x[2] - x[3] + x[4] + x[5] + x[6] + x[7] - x[8] - x[9] + x[10] + x[11] - x[12] + x[13] + x[14] - x[15] - x[16] - x[17] - x[18] + x[19] - x[20] + x[21] - x[22] + x[23] - x[24] - x[25] - x[26] - x[27] + x[28] + x[29] - x[30] - x[31] + x[32] - x[33] - x[34] + x[35] + x[36] + x[37] + x[38] - x[39];
+    out[32] = + x[0] + x[1] - x[2] + x[3] - x[4] + x[5] + x[6] + x[7] + x[8] - x[9] - x[10] + x[11] + x[12] - x[13] + x[14] + x[15] - x[16] - x[17] - x[18] - x[19] - x[20] - x[21] + x[22] - x[23] + x[24] - x[25] - x[26] - x[27] - x[28] + x[29] + x[30] - x[31] - x[32] + x[33] - x[34] - x[35] + x[36] + x[37] + x[38] + x[39];
+    out[33] = + x[0] - x[1] + x[2] - x[3] + x[4] - x[5] + x[6] + x[7] + x[8] + x[9] - x[10] - x[11] + x[12] + x[13] - x[14] + x[15] + x[16] - x[17] - x[18] - x[19] - x[20] + x[21] - x[22] + x[23] - x[24] + x[25] - x[26] - x[27] - x[28] - x[29] + x[30] + x[31] - x[32] - x[33] + x[34] - x[35] - x[36] + x[37] + x[38] + x[39];
+    out[34] = + x[0] - x[1] - x[2] + x[3] - x[4] + x[5] - x[6] + x[7] + x[8] + x[9] + x[10] - x[11] - x[12] + x[13] + x[14] - x[15] + x[16] + x[17] - x[18] - x[19] - x[20] + x[21] + x[22] - x[23] + x[24] - x[25] + x[26] - x[27] - x[28] - x[29] - x[30] + x[31] + x[32] - x[33] - x[34] + x[35] - x[36] - x[37] + x[38] + x[39];
+    out[35] = + x[0] - x[1] - x[2] - x[3] + x[4] - x[5] + x[6] - x[7] + x[8] + x[9] + x[10] + x[11] - x[12] - x[13] + x[14] + x[15] - x[16] + x[17] + x[18] - x[19] - x[20] + x[21] + x[22] + x[23] - x[24] + x[25] - x[26] + x[27] - x[28] - x[29] - x[30] - x[31] + x[32] + x[33] - x[34] - x[35] + x[36] - x[37] - x[38] + x[39];
+    out[36] = + x[0] - x[1] - x[2] - x[3] - x[4] + x[5] - x[6] + x[7] - x[8] + x[9] + x[10] + x[11] + x[12] - x[13] - x[14] + x[15] + x[16] - x[17] + x[18] + x[19] - x[20] + x[21] + x[22] + x[23] + x[24] - x[25] + x[26] - x[27] + x[28] - x[29] - x[30] - x[31] - x[32] + x[33] + x[34] - x[35] - x[36] + x[37] - x[38] - x[39];
+    out[37] = + x[0] + x[1] - x[2] - x[3] - x[4] - x[5] + x[6] - x[7] + x[8] - x[9] + x[10] + x[11] + x[12] + x[13] - x[14] - x[15] + x[16] + x[17] - x[18] + x[19] - x[20] - x[21] + x[22] + x[23] + x[24] + x[25] - x[26] + x[27] - x[28] + x[29] - x[30] - x[31] - x[32] - x[33] + x[34] + x[35] - x[36] - x[37] + x[38] - x[39];
+    out[38] = + x[0] + x[1] + x[2] - x[3] - x[4] - x[5] - x[6] + x[7] - x[8] + x[9] - x[10] + x[11] + x[12] + x[13] + x[14] - x[15] - x[16] + x[17] + x[18] - x[19] - x[20] - x[21] - x[22] + x[23] + x[24] + x[25] + x[26] - x[27] + x[28] - x[29] + x[30] - x[31] - x[32] - x[33] - x[34] + x[35] + x[36] - x[37] - x[38] + x[39];
+    out[39] = + x[0] - x[1] + x[2] + x[3] - x[4] - x[5] - x[6] - x[7] + x[8] - x[9] + x[10] - x[11] + x[12] + x[13] + x[14] + x[15] - x[16] - x[17] + x[18] + x[19] - x[20] + x[21] - x[22] - x[23] + x[24] + x[25] + x[26] + x[27] - x[28] + x[29] - x[30] + x[31] - x[32] - x[33] - x[34] - x[35] + x[36] + x[37] - x[38] - x[39];
+    #pragma unroll
+    for (int i = 0; i < 40; i++) { x[i] = out[i]; }
+}
+

packages/ltx-kernels/csrc/ops/include/static_switch.h

@@ -0,0 +1,23 @@
+// Inspired by https://github.com/NVIDIA/DALI/blob/main/include/dali/core/static_switch.h
+// and https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/Dispatch.h
+
+#pragma once
+
+/// @param COND       - a boolean expression to switch by
+/// @param CONST_NAME - a name given for the constexpr bool variable.
+/// @param ...       - code to execute for true and false
+///
+/// Usage:
+/// ```
+/// BOOL_SWITCH(flag, BoolConst, [&] {
+///     some_function<BoolConst>(...);
+/// });
+/// ```
+#define BOOL_SWITCH(COND, CONST_NAME, ...)      \
+    if (COND) {                                 \
+      constexpr static bool CONST_NAME = true;  \
+      __VA_ARGS__                               \
+    } else {                                    \
+      constexpr static bool CONST_NAME = false; \
+      __VA_ARGS__                               \
+    }                                           \

packages/ltx-kernels/csrc/ops/ops_api.cpp

@@ -0,0 +1,31 @@
+
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
+
+#include <torch/extension.h>
+#include <torch/python.h>
+
+
+at::Tensor rms_norm_rope(at::Tensor &x, c10::optional<at::Tensor>& weights_, at::Tensor &cos_freqs, at::Tensor &sin_freqs,  bool out_16bit);
+
+at::Tensor fp6_pack(at::Tensor &x);
+at::Tensor fp6_unpack(at::Tensor &x, int64_t original_n);
+
+at::Tensor rms_norm_split_rope(
+    at::Tensor &x,
+    at::Tensor &sin_freqs,
+    at::Tensor &cos_freqs,
+    at::Tensor &weights,
+    bool out_fp8
+);
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+    m.def("rms_norm_rope", &rms_norm_rope,
+          "fused  norm + rope + cvt");
+    m.def("fp6_pack", &fp6_pack,
+          "Pack 8-bit to 6-bit by dropping e_1 and e_2 bits");
+    m.def("fp6_unpack", &fp6_unpack,
+          "Unpack 6-bit to 8-bit (with e_1 and e_2 set to 0)");
+    m.def("rms_norm_split_rope", &rms_norm_split_rope,
+          "RMS norm + split RoPE with optional FP8 output");
+}

packages/ltx-kernels/csrc/ops/rms_norm_rope.cpp

@@ -0,0 +1,114 @@
+/******************************************************************************
+ * Copyright (c) 2023, Tri Dao.
+ ******************************************************************************/
+
+// Host entry point for the fused RMS-norm + RoPE kernel.
+
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <torch/extension.h>
+#include <vector>
+
+#include "fast_hadamard_transform.h"
+
+#define CHECK_SHAPE(x, ...) TORCH_CHECK(x.sizes() == torch::IntArrayRef({__VA_ARGS__}), #x " must have shape (" #__VA_ARGS__ ")")
+
+template<typename input_t, typename output_t, bool norm_affine>
+void rms_norm_rope_cuda(NormRopeHadamardParamsBase &params, cudaStream_t stream);
+
+void set_norm_rope_hadamard_params(NormRopeHadamardParamsBase &params,
+                         // sizes
+                         const size_t batch,
+                         const size_t dim,
+                         const size_t multiple,
+                         // device pointers
+                         const at::Tensor x,
+                         const at::Tensor cos_freqs,
+                         const at::Tensor sin_freqs,
+                         const at::Tensor weights,
+                         const at::Tensor out,
+
+                         bool norm_affine,
+                         float scale
+                         ) {
+
+    // Reset the parameters
+    memset(&params, 0, sizeof(params));
+
+    params.batch = batch;
+    params.dim = dim;
+    params.log_N = int(ceil(std::log2(dim / multiple)));
+
+    // Set the pointers and strides.
+    params.x_ptr = x.data_ptr();
+    params.out_ptr = out.data_ptr();
+    params.cos_freq_ptr = cos_freqs.data_ptr();
+    params.sin_freq_ptr = sin_freqs.data_ptr();
+    if (norm_affine){
+        params.weights_ptr = weights.data_ptr();
+    } else {
+        params.weights_ptr = nullptr;
+    }
+    // All stride are in elements, not bytes.
+    params.x_batch_stride = x.stride(0);
+    params.out_batch_stride = out.stride(0);
+    params.cos_freq_batch_stride = cos_freqs.stride(0);
+    params.sin_freq_batch_stride = sin_freqs.stride(0);
+
+    params.scale = scale;
+
+}
+
+at::Tensor rms_norm_rope(at::Tensor &x, c10::optional<at::Tensor>& weights_, at::Tensor &cos_freqs, at::Tensor &sin_freqs, bool out_16bit) {
+    auto input_type = x.scalar_type();
+    float scale = 1.0f; // :D
+    TORCH_CHECK(input_type == at::ScalarType::BFloat16);
+    TORCH_CHECK(x.is_cuda());
+    const auto shapes_og = x.sizes();
+    const int dim_og = x.size(-1);
+    x = x.reshape({-1, dim_og});
+    if (x.stride(-1) != 1) { x = x.contiguous(); }
+    const auto sizes = x.sizes();
+    const int batch_size = sizes[0];
+    cos_freqs = cos_freqs.reshape({-1, dim_og});
+    sin_freqs = sin_freqs.reshape({-1, dim_og});
+    at::Tensor weights;
+    bool norm_affine = false;
+    if(weights_.has_value()){
+        weights = weights_.value();
+        norm_affine = true;
+    }
+    CHECK_SHAPE(x, batch_size, dim_og);
+    TORCH_CHECK(x.stride(1) == 1);
+    if (dim_og % 8 != 0) {
+        x = torch::nn::functional::pad(x, torch::nn::functional::PadFuncOptions({0, 8 - dim_og % 8}));
+    }
+    const int dim = x.size(1);
+    at::Tensor out;
+    if (out_16bit){
+        out = torch::empty(x.sizes(), x.options().dtype(torch::kBFloat16));
+    } else {
+        out = torch::empty(x.sizes(), x.options().dtype(torch::kFloat8_e4m3fn));
+    }
+    at::cuda::CUDAGuard device_guard{(char)x.get_device()};
+    auto stream = at::cuda::getCurrentCUDAStream().stream();
+    NormRopeHadamardParamsBase params;
+    set_norm_rope_hadamard_params(params, batch_size, dim, 1, x, cos_freqs, sin_freqs, weights, out, norm_affine, scale);
+    TORCH_CHECK(dim % 8 == 0, "fast_hadamard_transform only supports hidden dimension divisible by 8 for now");
+    TORCH_CHECK(dim <= 32768, "fast_hadamard_transform only supports hidden dimension at most 32768 for now");
+    if (norm_affine){
+        if (out_16bit){
+            rms_norm_rope_cuda<at::BFloat16, at::BFloat16, true>(params, stream);
+        } else {
+            rms_norm_rope_cuda<at::BFloat16, at::Float8_e4m3fn, true>(params, stream);
+        }
+
+    } else {
+        if (out_16bit){
+            rms_norm_rope_cuda<at::BFloat16, at::BFloat16, false>(params, stream);
+        } else {
+            rms_norm_rope_cuda<at::BFloat16, at::Float8_e4m3fn, false>(params, stream);
+        }
+    }
+    return out.reshape(shapes_og);
+}

packages/ltx-kernels/csrc/ops/rms_norm_rope_cuda.cu

@@ -0,0 +1,228 @@
+/******************************************************************************
+ * Copyright (c) 2023, Tri Dao.
+ ******************************************************************************/
+
+// #pragma once
+
+#include <c10/util/BFloat16.h>
+#include <c10/util/Half.h>
+#include <c10/util/Float8_e4m3fn.h>
+#include <c10/cuda/CUDAException.h>  // For C10_CUDA_CHECK and C10_CUDA_KERNEL_LAUNCH_CHECK
+
+#include "fast_hadamard_transform.h"
+#include "fast_hadamard_transform_common.h"
+#include "fast_hadamard_transform_special.h"
+#include "static_switch.h"
+
+
+template<int kNThreads_, int kLogN_, typename input_t_, typename output_t_, bool norm_affine_>
+struct norm_rope_kernel_traits {
+    using input_t = input_t_;
+    using output_t = output_t_;
+
+    static constexpr int kNThreads = kNThreads_;
+    static constexpr int kLogN = kLogN_;
+    static constexpr int N = 1 << kLogN;
+    static constexpr int kNBytes = sizeof(input_t);
+    static constexpr int OutkNBytes = sizeof(output_t);
+    
+    static constexpr bool norm_affine = norm_affine_;
+    
+    static_assert(kNBytes == 1 || kNBytes == 2 || kNBytes == 4);
+    static constexpr int kNElts = kNBytes == 4 ? 4 : kNBytes == 2 ? 8 : 8;
+    // It's possible that we need to do 2 rounds of exchange if input_t is 16 bits
+    // (since then we'd have 8 values of float, and each round we can exchange 4 floats).
+    static constexpr int kNExchangePerVec = sizeof(float) / sizeof(input_t);
+    
+    using vec_t = typename BytesToType<kNBytes * kNElts>::Type;
+    using vec_t_out = typename BytesToType<OutkNBytes * kNElts>::Type;
+
+    static constexpr int kNChunks = N / (kNElts * kNThreads);
+    // We don't want to use more than 32 KB of shared memory.
+    static constexpr int kSmemExchangeSize = std::min(N * 4, 32 * 1024);
+    static constexpr int kNExchangeRounds = N * 4 / kSmemExchangeSize;
+    static_assert(kNExchangeRounds * kSmemExchangeSize == N * 4);
+    static constexpr int kSmemSize = kSmemExchangeSize;
+};
+
+
+template<typename Ktraits>
+__global__ __launch_bounds__(Ktraits::kNThreads)
+void norm_rope_cvt_kernel(NormRopeHadamardParamsBase params) {
+    constexpr int kNThreads = Ktraits::kNThreads;
+    constexpr int kNElts = Ktraits::kNElts;
+    constexpr int kNExchangePerVec = Ktraits::kNExchangePerVec;
+    constexpr int kNExchangeRounds = Ktraits::kNExchangeRounds;
+    constexpr int kNChunks = Ktraits::kNChunks;
+    constexpr bool norm_affine = Ktraits::norm_affine;
+
+    using input_t = typename Ktraits::input_t;
+    using output_t = typename Ktraits::output_t;
+    using vec_t = typename Ktraits::vec_t;
+    using out_vec_t = typename Ktraits::vec_t_out;
+    using weights_t = typename Ktraits::input_t;
+    using freqs_t = typename Ktraits::input_t;
+
+    constexpr int kLogNElts = cilog2(Ktraits::kNElts);
+    static_assert(1 << kLogNElts == kNElts, "kNElts must be a power of 2");
+    constexpr int kWarpSize = std::min(kNThreads, 32);
+    constexpr int kLogWarpSize = cilog2(kWarpSize);
+    static_assert(1 << kLogWarpSize == kWarpSize, "Warp size must be a power of 2");
+    constexpr int kNWarps = kNThreads / kWarpSize;
+    constexpr int kLogNWarps = cilog2(kNWarps);
+    static_assert(1 << kLogNWarps == kNWarps, "kNWarps must be a power of 2");
+    constexpr int kLoadsPerExchange = Ktraits::kSmemExchangeSize / (sizeof(vec_t) * kNThreads);
+    static_assert(kLoadsPerExchange * sizeof(vec_t) * kNThreads == Ktraits::kSmemExchangeSize, "kSmemExchangeSize should be a power of 2");
+    static_assert(kNExchangeRounds * kLoadsPerExchange * sizeof(vec_t) == kNChunks * kNElts * sizeof(float));
+
+    constexpr int kChunksPerExchange = Ktraits::kSmemExchangeSize / (sizeof(vec_t) * kNExchangePerVec * kNThreads);
+    static_assert(kChunksPerExchange * sizeof(vec_t) * kNExchangePerVec * kNThreads == Ktraits::kSmemExchangeSize);
+    constexpr int kNExchanges = kNChunks / kChunksPerExchange;
+    static_assert(kNExchanges * kChunksPerExchange == kNChunks);
+
+    // Shared memory.
+    extern __shared__ char smem_[];
+    vec_t *smem_exchange = reinterpret_cast<vec_t *>(smem_);
+
+    const int batch_id = blockIdx.x;
+    const int warp_id = threadIdx.x / 32;
+
+    input_t *x = reinterpret_cast<input_t *>(params.x_ptr) + batch_id * params.x_batch_stride;
+    output_t *out = reinterpret_cast<output_t *>(params.out_ptr) + batch_id * params.out_batch_stride;
+    weights_t *weights = norm_affine ? reinterpret_cast<weights_t*>(params.weights_ptr) : nullptr;
+
+    float x_vals[kNChunks][kNElts];
+    float weights_vals[kNChunks][kNElts];
+
+    load_input<kNChunks, kNElts, input_t>(x, x_vals, params.dim);
+
+    //RMS Norm START
+    float thread_squared_sum = 0.0f;
+    #pragma unroll
+    for (size_t c = 0; c < kNChunks; c++)
+    {
+        #pragma unroll
+        for (size_t i = 0; i < kNElts; i++)
+        {
+            thread_squared_sum += x_vals[c][i] * x_vals[c][i]; 
+        }
+        
+    }
+    SumOp<float> sum_op;
+    float warp_sum = Allreduce<32>::run(thread_squared_sum, sum_op);
+    float *smem_sum = reinterpret_cast<float*>(smem_);
+    if(threadIdx.x % 32 == 0){
+        smem_sum[warp_id] = warp_sum;
+    }
+    __syncthreads();
+    float norm = 0.0f;
+     #pragma unroll  
+    for (size_t i = 0; i < kNWarps; i++)
+    {
+        norm += smem_sum[i];
+    }
+    
+    norm *= 1.0f/params.dim;
+    norm = rsqrtf(norm);
+    if constexpr (norm_affine){
+        load_input<kNChunks, kNElts, weights_t>(weights, weights_vals, params.dim);
+    }
+    #pragma unroll
+    for (size_t c = 0; c < kNChunks; c++)
+    {
+        #pragma unroll
+        for (size_t i = 0; i < kNElts; i++)
+        {
+            if constexpr (norm_affine){
+                x_vals[c][i] *= (norm * weights_vals[c][i]);
+            } else {
+                x_vals[c][i] *= norm;
+            }
+        }
+        
+    }
+    //RMS NORM END
+    
+    //ROPE START
+    float sin_freqs_vals[kNChunks][kNElts];
+    float cos_freqs_vals[kNChunks][kNElts];
+
+    freqs_t *cos_freqs = reinterpret_cast<freqs_t*>(params.cos_freq_ptr) + batch_id * params.cos_freq_batch_stride;
+    freqs_t *sin_freqs = reinterpret_cast<freqs_t*>(params.sin_freq_ptr) + batch_id * params.sin_freq_batch_stride;
+
+    load_input<kNChunks, kNElts, freqs_t>(cos_freqs, cos_freqs_vals, params.dim);
+    load_input<kNChunks, kNElts, freqs_t>(sin_freqs, sin_freqs_vals, params.dim);
+
+    #pragma unroll
+    for (size_t c = 0; c < kNChunks; c++)
+    {
+        #pragma unroll
+        for (size_t i = 0; i < kNElts; i+=2)
+        {
+            float x_1 = x_vals[c][i];
+            float x_2 = x_vals[c][i+1];
+            x_vals[c][i] = -x_2*sin_freqs_vals[c][i] + x_1*cos_freqs_vals[c][i];
+            x_vals[c][i+1] = x_1*sin_freqs_vals[c][i+1] + x_2*cos_freqs_vals[c][i+1];
+        }
+    }
+    //ROPE END
+    
+    store_output<kNChunks, kNElts, output_t, false>(out, x_vals, params.dim, params.scale);
+}
+
+template<int kNThreads, int kLogN, typename input_t, typename output_t, bool norm_affine>
+void norm_rope_cvt_launch(NormRopeHadamardParamsBase &params, cudaStream_t stream) {
+    using Ktraits = norm_rope_kernel_traits<kNThreads, kLogN, input_t, output_t, norm_affine>;
+    constexpr int kSmemSize = Ktraits::kSmemSize;
+    dim3 grid(params.batch);
+    auto kernel = &norm_rope_cvt_kernel<Ktraits>;
+    if (kSmemSize >= 48 * 1024) {
+        C10_CUDA_CHECK(cudaFuncSetAttribute(
+            kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemSize));
+        }
+    kernel<<<grid, Ktraits::kNThreads, kSmemSize, stream>>>(params);
+    C10_CUDA_KERNEL_LAUNCH_CHECK();
+}
+
+template<typename input_t, typename output_t, bool norm_affine>
+void rms_norm_rope_cuda(NormRopeHadamardParamsBase &params, cudaStream_t stream) {
+    if (params.log_N == 3) {
+        norm_rope_cvt_launch<1, 3, input_t, output_t, norm_affine>(params, stream);
+    } else if (params.log_N == 4) {
+        norm_rope_cvt_launch<2, 4, input_t, output_t, norm_affine>(params, stream);
+    } else if (params.log_N == 5) {
+        norm_rope_cvt_launch<4, 5, input_t, output_t, norm_affine>(params, stream);
+    } else if (params.log_N == 6) {
+        norm_rope_cvt_launch<8, 6, input_t, output_t, norm_affine>(params, stream);
+    } else if (params.log_N == 7) {
+        norm_rope_cvt_launch<16, 7, input_t, output_t, norm_affine>(params, stream);
+    } else if (params.log_N == 8) {
+        norm_rope_cvt_launch<32, 8, input_t, output_t, norm_affine>(params, stream);
+    } else if (params.log_N == 9) {
+        norm_rope_cvt_launch<32, 9, input_t, output_t, norm_affine>(params, stream);
+    } else if (params.log_N == 10) {
+        norm_rope_cvt_launch<128, 10, input_t, output_t, norm_affine>(params, stream);
+    } else if (params.log_N == 11) {
+        norm_rope_cvt_launch<256, 11, input_t, output_t, norm_affine>(params, stream);
+    } else if (params.log_N == 12) {
+        norm_rope_cvt_launch<256, 12, input_t, output_t, norm_affine>(params, stream);
+    } else if (params.log_N == 13) {
+        norm_rope_cvt_launch<256, 13, input_t, output_t, norm_affine>(params, stream);
+    } else if (params.log_N == 14) {
+        norm_rope_cvt_launch<256, 14, input_t, output_t, norm_affine>(params, stream);
+    } else if (params.log_N == 15) {
+        norm_rope_cvt_launch<256, 15, input_t, output_t, norm_affine>(params, stream);
+    }
+}
+
+
+template void rms_norm_rope_cuda<at::BFloat16, at::Float8_e4m3fn, false>(NormRopeHadamardParamsBase &params, cudaStream_t stream);
+template void rms_norm_rope_cuda<at::BFloat16, at::Float8_e4m3fn, true>(NormRopeHadamardParamsBase &params, cudaStream_t stream);
+
+template void rms_norm_rope_cuda<at::BFloat16, at::BFloat16, false>(NormRopeHadamardParamsBase &params, cudaStream_t stream);
+template void rms_norm_rope_cuda<at::BFloat16, at::BFloat16, true>(NormRopeHadamardParamsBase &params, cudaStream_t stream);
+
+// template void fast_hadamard_transform_cuda<at::BFloat16, at::BFloat16>(HadamardParamsBase &params, cudaStream_t stream);
+
+// template void fast_hadamard_transform_cuda<at::Float8_e4m3fn, at::BFloat16>(HadamardParamsBase &params, cudaStream_t stream);
+// template void fast_hadamard_transform_cuda<at::Float8_e4m3fn, at::Float8_e4m3fn>(HadamardParamsBase &params, cudaStream_t stream);

packages/ltx-kernels/csrc/ops/rms_norm_split_rope.cpp

@@ -0,0 +1,111 @@
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <torch/extension.h>
+#include <torch/python.h>
+
+#include <vector>
+
+// Forward declaration of CUDA kernel template
+template<typename out_t>
+void rms_norm_split_rope_cuda(
+    void* x,
+    void* sin_freqs,
+    void* cos_freqs,
+    void* weights,
+    int b,
+    int s,
+    int n,
+    int h,
+    long cos_sb, long cos_sn, long cos_ss,
+    long sin_sb, long sin_sn, long sin_ss,
+    void* out,
+    cudaStream_t stream
+);
+
+at::Tensor rms_norm_split_rope(
+    at::Tensor &x,
+    at::Tensor &sin_freqs,
+    at::Tensor &cos_freqs,
+    at::Tensor &weights,
+    bool out_fp8
+) {
+    TORCH_CHECK(x.scalar_type() == at::ScalarType::BFloat16, "Input must be BFloat16");
+    TORCH_CHECK(sin_freqs.scalar_type() == at::ScalarType::BFloat16, "sin_freqs must be BFloat16");
+    TORCH_CHECK(cos_freqs.scalar_type() == at::ScalarType::BFloat16, "cos_freqs must be BFloat16");
+    TORCH_CHECK(x.is_cuda(), "Input must be on CUDA");
+    TORCH_CHECK(sin_freqs.is_cuda(), "sin_freqs must be on CUDA");
+    TORCH_CHECK(cos_freqs.is_cuda(), "cos_freqs must be on CUDA");
+
+    // Get dimensions
+    // x: [b, s, h]
+    // cos, sin: [b, n, s, d] where n*d = h/2
+    int b = x.size(0);
+    int s = x.size(1);
+    int h = x.size(2);
+
+    TORCH_CHECK(cos_freqs.dim() == 4, "cos_freqs must be 4D");
+    TORCH_CHECK(sin_freqs.dim() == 4, "sin_freqs must be 4D");
+
+    int n = cos_freqs.size(1);
+    int d = h / n;
+
+    
+    // Require a contiguous innermost (d/2) dim for the vectorized int4 freq load,
+    // but keep the outer (b, n, s) strides: apply_split_rotary_emb hands us a
+    // swapaxes view (logical [b, n, s, d/2], physical [b, s, n, d/2]) whose inner
+    // stride is already 1, so this never copies it. The strides are forwarded to
+    // the kernel so the read is correct regardless of the physical layout.
+    if (x.stride(-1) != 1) { x = x.contiguous(); }
+    if (cos_freqs.stride(-1) != 1) { cos_freqs = cos_freqs.contiguous(); }
+    if (sin_freqs.stride(-1) != 1) { sin_freqs = sin_freqs.contiguous(); }
+
+    long cos_sb = cos_freqs.stride(0), cos_sn = cos_freqs.stride(1), cos_ss = cos_freqs.stride(2);
+    long sin_sb = sin_freqs.stride(0), sin_sn = sin_freqs.stride(1), sin_ss = sin_freqs.stride(2);
+
+    // Create output tensor
+    at::Tensor out;
+    if (out_fp8) {
+        out = torch::empty(x.sizes(), x.options().dtype(torch::kFloat8_e4m3fn));
+    } else {
+        out = torch::empty(x.sizes(), x.options().dtype(torch::kBFloat16));
+    }
+    
+    // Setup CUDA
+    at::cuda::CUDAGuard device_guard{(char)x.get_device()};
+    auto stream = at::cuda::getCurrentCUDAStream().stream();
+
+    // Launch kernel
+    if (out_fp8) {
+        rms_norm_split_rope_cuda<at::Float8_e4m3fn>(
+            x.data_ptr(),
+            sin_freqs.data_ptr(),
+            cos_freqs.data_ptr(),
+            weights.data_ptr(),  // weights (optional, not used yet)
+            b,
+            s,
+            n,
+            h,
+            cos_sb, cos_sn, cos_ss,
+            sin_sb, sin_sn, sin_ss,
+            (void*)out.data_ptr(),
+            stream
+        );
+    } else {
+        rms_norm_split_rope_cuda<at::BFloat16>(
+            x.data_ptr(),
+            sin_freqs.data_ptr(),
+            cos_freqs.data_ptr(),
+            weights.data_ptr(),  // weights (optional, not used yet)
+            b,
+            s,
+            n,
+            h,
+            cos_sb, cos_sn, cos_ss,
+            sin_sb, sin_sn, sin_ss,
+            (void*)out.data_ptr(),
+            stream
+        );
+    }
+
+    return out;
+}

packages/ltx-kernels/csrc/ops/rms_norm_split_rope_cuda.cu

@@ -0,0 +1,202 @@
+#include <c10/util/BFloat16.h>
+#include <c10/util/Float8_e4m3fn.h>
+#include <c10/cuda/CUDAException.h>
+#include <cuda_runtime.h>
+#include <cuda_bf16.h>
+#include <cuda_fp8.h>
+
+// CUDA kernel template for RMS norm + split RoPE
+// out_t can be at::Float8_e4m3fn or at::BFloat16
+
+using bf16 = __nv_bfloat16;
+using fp8 = __nv_fp8_e4m3;
+__device__ __forceinline__ void _load_x(bf16* x, float x_vals[8], int h){
+    bf16 x_tmp[8];
+    *reinterpret_cast<int4*>(x_tmp) = reinterpret_cast<int4*>(x + blockIdx.x * h)[threadIdx.x];
+    #pragma unroll
+    for(int i = 0; i < 8; i++){
+        x_vals[i] = float(x_tmp[i]);
+    }
+}
+// Load 8 freq values for this thread from a table laid out logically as
+// [b, n, s, d/2] (what apply_split_rotary_emb produces -- a swapaxes view whose
+// physical layout is [b, s, n, d/2]). The strides (sb, sn, ss; inner d/2 stride
+// is 1) are forwarded from the host so the read is correct for both that
+// non-contiguous view and a genuinely contiguous [b, n, s, d/2] tensor. Both
+// head-halves map to the same freq element (mirrors the eager cos.unsqueeze(-2)).
+__device__ __forceinline__ void _load_freqs(
+    const bf16* freqs, float x_vals[8], int s, int d, long sb, long sn, long ss
+){
+    bf16 x_tmp[8];
+    int threads_per_head = d / 8;
+    int head_idx = threadIdx.x / threads_per_head;
+    int lane = threadIdx.x % (threads_per_head / 2);
+    int b_idx = blockIdx.x / s;
+    int t_idx = blockIdx.x % s;
+    long off = b_idx * sb + head_idx * sn + t_idx * ss + (long)lane * 8;
+    *reinterpret_cast<int4*>(x_tmp) = *reinterpret_cast<const int4*>(freqs + off);
+    #pragma unroll
+    for(int i = 0; i < 8; i++){
+        x_vals[i] = float(x_tmp[i]);
+    }
+}
+
+template<typename out_t>
+__global__ void _rms_norm_split_rope_kernel(bf16* x, bf16* sin_freqs, bf16* cos_freqs, void* out, bf16* weights, int b, int s, int n, int h,
+    long cos_sb, long cos_sn, long cos_ss, long sin_sb, long sin_sn, long sin_ss){
+    int token_idx = blockIdx.x;
+    int tid = threadIdx.x;
+    int lane_id = tid % 32;
+    // freqs have shape [b, s, h/2]
+    // each thread block calculate one row
+    // there are h/8 threads in thread block, each thread processes 8 values
+    // num_of_rows = b * s
+    // freqs have h/2 dim
+    // gridDim is (num_of_rows, 1, 1)
+    // calculate rms norm x_normed = x/x_norm * weights. x_norm is calculated across row, it means thread block wide sum reduction
+
+    extern __shared__ float smem[];
+
+    // Step 1: Load input values (8 per thread)
+    float x_vals[8];
+    _load_x(x, x_vals, h);
+
+    float sum_sq = 0.0f;
+    #pragma unroll
+    for(int i = 0; i < 8; i++){
+        sum_sq += x_vals[i] * x_vals[i];
+    }
+
+    // Warp-level reduction
+    #pragma unroll
+    for(int offset = 16; offset > 0; offset >>= 1){
+        sum_sq += __shfl_xor_sync(0xffffffff, sum_sq, offset);
+    }
+
+    if(tid % 32 == 0){
+        smem[tid / 32] = sum_sq;
+    }
+    __syncthreads();
+
+    // Final reduction across warps
+    if(tid == 0){
+        float total_sum = 0.0f;
+        int num_warps = blockDim.x / 32;
+        for(int i = 0; i < num_warps; i++){
+            total_sum += smem[i];
+        }
+        // RMS: sqrt(mean(x^2))
+        float rms = rsqrtf(total_sum / h + 1e-6f);  // Add epsilon for numerical stability
+        smem[0] = rms;
+    }
+    __syncthreads();
+
+    float inv_rms = smem[0];
+
+    // Step 3: Apply RMS normalization (and weights if provided)
+    #pragma unroll
+    for(int i = 0; i < 8; i++){
+        x_vals[i] *= inv_rms;
+        // TODO: Apply weights if provided
+        if(weights != nullptr) x_vals[i] *= float(weights[tid * 8 + i]);
+    }
+
+    // Step 4: Calculate dimensions for split RoPE
+    // Conceptually: [b, s, h] -> [b, s, n, 2*d] -> [b, s, n, 2, d]
+    // where h = n * 2 * d
+    int d = h / n;
+    float x_other_vals[8];
+
+    int threads_per_head = d / 8;
+    int head_idx = tid / threads_per_head;
+    int idx_in_head = tid % threads_per_head;
+    bool is_first_half = idx_in_head < (threads_per_head / 2);
+
+    // LT-PATCH: full-warp mask. The original (1u << threads_per_head) - 1 only marks
+    // the first head's lanes active, so lanes belonging to heads beyond the first are
+    // not in the mask -> __shfl_xor_sync result is undefined and can corrupt RoPE. The
+    // XOR pattern keeps data within each power-of-two head group, so a full-warp mask
+    // is correct for every lane.
+    const unsigned mask = 0xffffffffu;
+    const int laneMask = threads_per_head / 2;           // 4, 8, or 16
+    #pragma unroll
+    for (int i = 0; i < 8; i++) {
+        x_other_vals[i] = __shfl_xor_sync(mask, x_vals[i], laneMask);
+    }
+
+    float cos_vals[8], sin_vals[8];
+    _load_freqs(cos_freqs, cos_vals, s, d, cos_sb, cos_sn, cos_ss);
+    _load_freqs(sin_freqs, sin_vals, s, d, sin_sb, sin_sn, sin_ss);
+    #pragma unroll
+    for(int i = 0; i < 8; i++){
+        x_vals[i] = cos_vals[i]*x_vals[i];
+    }
+
+
+    float sign = is_first_half ? -1.0f : 1.0f;
+    for(int i = 0; i < 8; i++){
+        x_vals[i] += sign*sin_vals[i]*x_other_vals[i];
+    }
+
+    // Step 6: Convert and store output
+    if constexpr (std::is_same_v<out_t, at::Float8_e4m3fn>){
+        fp8 out_tmp[8];
+        #pragma unroll
+        for(int i = 0; i < 8; i++){
+            out_tmp[i] = fp8(x_vals[i]);
+        }
+        *reinterpret_cast<int64_t*>((fp8*)out + token_idx * h + tid * 8) = *reinterpret_cast<int64_t*>(out_tmp);
+    } else {
+        bf16 out_tmp[8];
+        #pragma unroll
+        for(int i = 0; i < 8; i++){
+            out_tmp[i] = __float2bfloat16(x_vals[i]);
+        }
+        *reinterpret_cast<int4*>((bf16*)out + token_idx * h + tid * 8) = *reinterpret_cast<int4*>(out_tmp);
+    }
+}
+
+template<typename out_t>
+void rms_norm_split_rope_cuda(
+    void* x,           // Input: [b, s, h]
+    void* sin_freqs,   // Sin frequencies: [b, n, s, d]
+    void* cos_freqs,   // Cos frequencies: [b, n, s, d]
+    void* weights,
+    int b,                      // Batch size
+    int s,                      // Sequence length
+    int n,                      // Number of heads (32)
+    int h,                      // Hidden dimension (2048, 4096, or 8192)
+    long cos_sb, long cos_sn, long cos_ss,  // cos_freqs strides (b, n, s)
+    long sin_sb, long sin_sn, long sin_ss,  // sin_freqs strides (b, n, s)
+    void* out,                // Output: [b, s, h]
+    cudaStream_t stream
+) {
+    int num_tokens = b * s;
+    int num_threads = h / 8;  // Each thread processes 8 elements
+    int smem_size = (num_threads / 32 + 1) * sizeof(float);  // Shared memory for reductions
+
+    dim3 grid(num_tokens);
+    dim3 block(num_threads);
+
+    _rms_norm_split_rope_kernel<out_t><<<grid, block, smem_size, stream>>>(
+        reinterpret_cast<bf16*>(x),
+        reinterpret_cast<bf16*>(sin_freqs),
+        reinterpret_cast<bf16*>(cos_freqs),
+        out,
+        reinterpret_cast<bf16*>(weights),
+        b, s, n, h,
+        cos_sb, cos_sn, cos_ss,
+        sin_sb, sin_sn, sin_ss
+    );
+
+    C10_CUDA_KERNEL_LAUNCH_CHECK();
+}
+
+// Explicit template instantiations
+template void rms_norm_split_rope_cuda<at::BFloat16>(
+    void*, void*, void*, void*, int, int, int, int, long, long, long, long, long, long, void*, cudaStream_t
+);
+
+template void rms_norm_split_rope_cuda<at::Float8_e4m3fn>(
+    void*, void*, void*, void*, int, int, int, int, long, long, long, long, long, long, void*, cudaStream_t
+);

packages/ltx-kernels/pyproject.toml

@@ -0,0 +1,18 @@
+[build-system]
+requires = [
+  "setuptools>=61",
+  "wheel",
+  "torch",
+]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "ltx-kernels"
+version = "0.1.0"
+dependencies = ["torch"]
+
+[tool.setuptools]
+package-dir = {"" = "src"}
+
+[tool.setuptools.packages.find]
+where = ["src"]

packages/ltx-kernels/setup.py

@@ -0,0 +1,228 @@
+import os
+import re
+import subprocess
+from pathlib import Path
+
+import setuptools
+import torch
+from torch.utils.cpp_extension import CUDA_HOME, BuildExtension, CUDAExtension
+
+ROOT = Path(__file__).resolve().parent
+
+# cutlass headers for the blockwise GEMM kernels. Pinned to the upstream commit the
+# build was validated against and fetched into a cache dir, rather than carried as a
+# git submodule (keeps the public-repo sync clean and the clone CI-cacheable).
+CUTLASS_REPO = "https://github.com/NVIDIA/cutlass.git"
+CUTLASS_REF = "afa1772203677c5118fcd82537a9c8fefbcc7008"  # v3.8.0
+
+
+def _nvidia_include_dirs() -> list[str]:
+    """Include dirs from pip-installed nvidia packages (e.g. cusparse headers)."""
+    try:
+        import nvidia  # noqa: PLC0415
+
+        return [str(p) for pkg in Path(nvidia.__path__[0]).iterdir() if (p := pkg / "include").is_dir()]
+    except ImportError:
+        return []
+
+
+def _arch_tokens() -> list[str]:
+    """Normalized entries from TORCH_CUDA_ARCH_LIST (e.g. ['8.9', '9.0'])."""
+    raw = os.environ.get("TORCH_CUDA_ARCH_LIST", "")
+    return [re.sub(r"\+PTX$", "", tok).strip() for tok in raw.replace(",", " ").split() if tok.strip()]
+
+
+# Arch codes the blockwise FP8 GEMM kernels support, as nvcc `sm_<code>` targets.
+# SM89 ("geforce") is the generic fp8 kernel: it runs on Ada and is also the kernel
+# dispatched on Blackwell (sm_100a datacenter / sm_120 consumer), so it is compiled for
+# those too. The SM90 ("deep_gemm") kernel is Hopper-only and needs sm_90a (wgmma/TMA);
+# it is declared-always / defined-conditionally, so it stubs out on every non-Hopper
+# pass and can share a multi-arch fat binary. Ampere has no fp8 path. Entries are
+# filtered to what the local nvcc can actually target (see _nvcc_arch_nums), so the
+# Blackwell codes are inert until built with CUDA 12.8+.
+# NOTE: the Blackwell (100a/120) path is implemented but not yet validated on real
+# Blackwell hardware -- needs a B200 + CUDA 12.8 build/run pass.
+_BLOCKWISE_ARCHES = ["89", "90a", "100a", "120"]
+
+
+def _nvcc_arch_nums() -> set[str]:
+    """Architecture numbers this nvcc can target, e.g. {'80', '86', '89', '90'}."""
+    try:
+        out = subprocess.check_output([f"{CUDA_HOME}/bin/nvcc", "--list-gpu-arch"], text=True)
+    except (OSError, subprocess.CalledProcessError):
+        return set()
+    return {m.group(1) for tok in out.split() if (m := re.match(r"compute_(\d+a?)$", tok.strip()))}
+
+
+def _blockwise_gencode() -> tuple[list[str], bool]:
+    """Return (``-gencode`` flags for blockwise_cpp, build_sm90).
+    Honors ``TORCH_CUDA_ARCH_LIST`` when set (mapping 8.9 -> sm_89, 9.0/9.0a -> sm_90a,
+    ignoring arches the kernels do not support); unset builds for every supported arch
+    this nvcc can target. The flags apply uniformly to all sources -- safe because the
+    SM90 source stubs itself out on non-sm_90a passes.
+    """
+    supported = _nvcc_arch_nums()
+    # nvcc may report an arch either plain ("90") or suffixed ("90a"); accept either.
+    base = [a for a in _BLOCKWISE_ARCHES if a in supported or a.rstrip("a") in supported]
+    env = _arch_tokens()
+    if env:
+        sel = []
+        for tok in env:
+            if tok.startswith("8.9"):
+                sel.append("89")
+            elif tok.startswith("9.0"):
+                sel.append("90a")
+            elif tok.startswith("10.0"):
+                sel.append("100a")
+            elif tok.startswith("12.0"):
+                sel.append("120")
+            # Ampere (8.0/8.6) and other arches have no fp8 blockwise kernel.
+        archs = [a for a in dict.fromkeys(sel) if a in base] or base
+    else:
+        archs = base
+    flags = [f"-gencode=arch=compute_{a},code=sm_{a}" for a in archs]
+    return flags, ("90a" in archs)
+
+
+def _cutlass_include() -> str:
+    """Return the cutlass include dir, fetching the pinned commit on first use.
+    Honors ``CUTLASS_DIR`` (a prebuilt cutlass checkout, e.g. a system copy) and
+    otherwise caches a shallow clone of ``CUTLASS_REF`` under
+    ``LTX_KERNELS_CACHE_DIR`` (default ``~/.cache/ltx-kernels``), so it is reused
+    across builds and can be restored from a CI cache. Keeps ``uv sync`` /
+    ``pip install -e`` self-contained without a git submodule.
+    """
+    if env := os.environ.get("CUTLASS_DIR"):
+        return str(Path(env) / "include")
+    cache_root = Path(os.environ.get("LTX_KERNELS_CACHE_DIR", Path.home() / ".cache" / "ltx-kernels"))
+    dest = cache_root / f"cutlass-{CUTLASS_REF}"
+    if not (dest / "include" / "cutlass" / "cutlass.h").is_file():
+        dest.mkdir(parents=True, exist_ok=True)
+        # Blobless partial clone of the exact pinned commit (GitHub allows fetching an
+        # arbitrary SHA), sparse-checked-out to include/ only: cutlass is header-only and
+        # the rest of the repo (tools/test/examples/python, ~85% by size) is unused.
+        subprocess.run(["git", "init", "-q", str(dest)], check=True)
+        subprocess.run(["git", "-C", str(dest), "remote", "add", "origin", CUTLASS_REPO], check=True)
+        # Cone-mode sparse checkout of just include/. Use init + set (not "set --cone",
+        # whose inline flag postdates git 2.35 and is silently parsed as a pattern on older git).
+        subprocess.run(["git", "-C", str(dest), "sparse-checkout", "init", "--cone"], check=True)
+        subprocess.run(["git", "-C", str(dest), "sparse-checkout", "set", "include"], check=True)
+        subprocess.run(
+            ["git", "-C", str(dest), "fetch", "-q", "--depth", "1", "--filter=blob:none", "origin", CUTLASS_REF],
+            check=True,
+        )
+        subprocess.run(["git", "-C", str(dest), "checkout", "-q", "FETCH_HEAD"], check=True)
+    return str(dest / "include")
+
+
+if __name__ == "__main__":
+    if CUDA_HOME is None:
+        raise RuntimeError(
+            "CUDA toolkit not found (CUDA_HOME is None). ltx-kernels compiles CUDA extensions "
+            "and must be built on a host with the CUDA toolkit installed (nvcc on PATH or "
+            "CUDA_HOME set)."
+        )
+
+    ext_modules = []
+
+    # all2all_cpp -- unchanged.
+    all2all_args = ["-O3", "-Wall", "-Wextra", "-Werror", "-Wno-unused-parameter", "-Wno-attributes"]
+    ext_modules.append(
+        CUDAExtension(
+            name="all2all_cpp",
+            include_dirs=[str(ROOT / "csrc/all2all"), str(ROOT / "csrc/include"), *_nvidia_include_dirs()],
+            sources=[
+                "csrc/all2all/all2all.cpp",
+                "csrc/all2all/cuda/all2all_heads.cu",
+                "csrc/all2all/cuda/allgather.cu",
+            ],
+            extra_compile_args={"cxx": all2all_args, "nvcc": ["-O3"]},
+        )
+    )
+
+    # ops_cpp -- arch-independent element ops (rms_norm_rope, rms_norm_split_rope,
+    # fp6 pack/unpack). Arch is driven by TORCH_CUDA_ARCH_LIST / torch defaults.
+    ext_modules.append(
+        CUDAExtension(
+            name="ops_cpp",
+            sources=[
+                "csrc/ops/ops_api.cpp",
+                "csrc/ops/fp6_bitpack.cpp",
+                "csrc/ops/fp6_pack.cu",
+                "csrc/ops/rms_norm_rope.cpp",
+                "csrc/ops/rms_norm_rope_cuda.cu",
+                "csrc/ops/rms_norm_split_rope.cpp",
+                "csrc/ops/rms_norm_split_rope_cuda.cu",
+            ],
+            include_dirs=[str(ROOT / "csrc/ops/include"), *_nvidia_include_dirs()],
+            extra_compile_args={
+                "cxx": ["-O3", "-std=c++17"],
+                "nvcc": [
+                    "-O3",
+                    "-std=c++17",
+                    "-U__CUDA_NO_HALF_OPERATORS__",
+                    "-U__CUDA_NO_HALF_CONVERSIONS__",
+                    "-U__CUDA_NO_HALF2_OPERATORS__",
+                    "-U__CUDA_NO_BFLOAT16_CONVERSIONS__",
+                    "--expt-relaxed-constexpr",
+                    "--expt-extended-lambda",
+                    "--use_fast_math",
+                ],
+            },
+        )
+    )
+
+    # blockwise_cpp -- FP8 GEMM. The SM89 (GeForce) kernel is always built; the SM90
+    # (deep_gemm) kernel + -D__SM90__ are added whenever sm_90a is among the targets.
+    # Arches are an explicit -gencode list (see _blockwise_gencode): TORCH_CUDA_ARCH_LIST
+    # when set, else every supported arch this nvcc can target ("build for everything").
+    # The list is uniform across sources -- the SM90 source declares-always /
+    # defines-conditionally, so it compiles (as a stub) for non-sm_90a arches too. Note
+    # blockwise is unsupported on Ampere and fails at *runtime* there, by design.
+    cutlass_include = _cutlass_include()
+    gencode, build_sm90 = _blockwise_gencode()
+    blockwise_sources = [
+        "csrc/blockwise/api.cpp",
+        "csrc/blockwise/kernels/geforce/gemm.cu",
+    ]
+    abi = f"-D_GLIBCXX_USE_CXX11_ABI={int(torch.compiled_with_cxx11_abi())}"
+    blockwise_cxx = ["-O3", "-std=c++17", "-fPIC", "-Wno-psabi", "-Wno-deprecated-declarations", abi]
+    blockwise_nvcc = [
+        "-O3",
+        "-std=c++17",
+        "--ptxas-options=-O2",
+        "--expt-relaxed-constexpr",
+        "--expt-extended-lambda",
+        "-U__CUDA_NO_HALF_OPERATORS__",
+        "-U__CUDA_NO_HALF_CONVERSIONS__",
+        "-U__CUDA_NO_HALF2_OPERATORS__",
+        "-U__CUDA_NO_BFLOAT16_CONVERSIONS__",
+        *gencode,
+    ]
+    if build_sm90:
+        blockwise_sources.append("csrc/blockwise/kernels/deep_gemm/include/deep_gemm/impls/sm90_fp8_gemm_1d2d_bias.cu")
+        blockwise_cxx.append("-D__SM90__")
+        blockwise_nvcc.append("-D__SM90__")
+
+    ext_modules.append(
+        CUDAExtension(
+            name="blockwise_cpp",
+            sources=blockwise_sources,
+            include_dirs=[
+                f"{CUDA_HOME}/include",
+                f"{CUDA_HOME}/include/cccl",
+                str(ROOT / "csrc/blockwise"),
+                str(ROOT / "csrc/blockwise/kernels/deep_gemm/include"),
+                cutlass_include,
+                *_nvidia_include_dirs(),
+            ],
+            libraries=["cuda", "cudart", "nvrtc"],
+            library_dirs=[f"{CUDA_HOME}/lib64", f"{CUDA_HOME}/lib64/stubs"],
+            extra_compile_args={"cxx": blockwise_cxx, "nvcc": blockwise_nvcc},
+        )
+    )
+
+    setuptools.setup(
+        ext_modules=ext_modules,
+        cmdclass={"build_ext": BuildExtension},
+    )

packages/ltx-kernels/src/ltx_kernels/__init__.py

@@ -0,0 +1,6 @@
+"""ltx-kernels: High-performance CUDA kernels for distributed attention operations."""
+
+from ltx_kernels.all_to_all import All2All
+
+__all__ = ["All2All"]
+__version__ = "0.1.0"

packages/ltx-kernels/src/ltx_kernels/all_to_all.py

@@ -0,0 +1,183 @@
+"""All2All communication primitives for distributed attention.
+The C++ kernels are exposed via ``torch.library.custom_op`` so that
+``torch.compile`` (and CUDA Graph capture under ``mode="reduce-overhead"``)
+can trace through them without a graph break. Each :class:`All2All`
+instance registers itself in a class-level registry indexed by an integer
+``comm_id``; the custom op takes that id plus an input tensor and dispatches
+to the appropriate C++ runtime.
+"""
+
+import math
+import weakref
+from typing import Any, ClassVar
+
+import torch
+import torch.distributed as dist
+from all2all_cpp import All2All as All2AllCpp
+from torch.library import custom_op
+
+# Output shapes are derived symbolically from the input tensor's shape under
+# the assumption of *uniform* sharding (caller pads up-front so
+# `total_tokens % world_size == 0`). `world_size` travels through the op as an int
+# so it enters the traced graph: the fake needs it for the output shape, and since
+# it is constant per process the Dynamo guard it installs never triggers a
+# within-run recompile -- it only keys the compile cache by world size, so a graph
+# compiled under one GPU count is never replayed under another. Per-rank token
+# counts, which do vary per call, stay in the C++ runtime state (`set_rank_tokens`)
+# and never travel through the op.
+
+
+@custom_op("ltx_kernels::send_recv_heads", mutates_args=(), device_types="cuda")
+def _send_recv_heads_op(
+    x: torch.Tensor,
+    comm_id: int,
+    world_size: int,  # noqa: ARG001
+    copy_out: bool,
+) -> torch.Tensor:
+    # copy_out=False returns a view into the IPC buffer (zero-copy). Safe under
+    # cudagraph_trees because the IPC buffer is cudaMalloc'd inside All2All, not
+    # in the static graph pool, so it isn't subject to graph-pool aliasing.
+    return All2All._runtime_registry[comm_id].send_recv_heads(x, copy_out)
+
+
+@_send_recv_heads_op.register_fake
+def _send_recv_heads_fake(
+    x: torch.Tensor,
+    comm_id: int,  # noqa: ARG001
+    world_size: int,
+    copy_out: bool,  # noqa: ARG001
+) -> torch.Tensor:
+    return x.new_empty((x.shape[0], x.shape[1] * world_size, x.shape[2] // world_size, x.shape[3]))
+
+
+@custom_op("ltx_kernels::gather_heads", mutates_args=(), device_types="cuda")
+def _gather_heads_op(
+    x: torch.Tensor,
+    comm_id: int,
+    world_size: int,  # noqa: ARG001
+    copy_out: bool,
+) -> torch.Tensor:
+    return All2All._runtime_registry[comm_id].gather_heads(x, copy_out)
+
+
+@_gather_heads_op.register_fake
+def _gather_heads_fake(
+    x: torch.Tensor,
+    comm_id: int,  # noqa: ARG001
+    world_size: int,
+    copy_out: bool,  # noqa: ARG001
+) -> torch.Tensor:
+    return x.new_empty((x.shape[0], x.shape[1] // world_size, x.shape[2] * world_size, x.shape[3]))
+
+
+class All2All:
+    """IPC-based All2All communication for distributed head-parallel attention.
+    This class manages GPU memory buffers and IPC handles to enable efficient
+    cross-GPU communication for attention head redistribution.
+    Args:
+        rank: Local rank of this process.
+        world_size: Total number of processes in the distributed group.
+        seqlen: Maximum sequence length to allocate buffers for.
+        hidden_dim: Hidden dimension size (num_heads * head_dim).
+        num_sms: Number of SMs to use for kernel execution.
+        tensor_dtype: Data type for tensors (e.g., torch.bfloat16).
+        group: PyTorch distributed process group.
+    """
+
+    _next_id: ClassVar[int] = 0
+    _runtime_registry: ClassVar[dict[int, "All2AllCpp"]] = {}
+
+    @staticmethod
+    def _release_comm(comm_id: int, runtime: "All2AllCpp") -> None:
+        All2All._runtime_registry.pop(comm_id, None)
+        runtime.destroy()
+
+    def __init__(
+        self,
+        rank: int,
+        world_size: int,
+        seqlen: int,
+        hidden_dim: int,
+        num_sms: int,
+        tensor_dtype: torch.dtype,
+        group: torch.distributed.ProcessGroup | None = None,
+        timeout_seconds: float | None = None,
+    ) -> None:
+        self.rank = rank
+        self.world_size = world_size
+        self.num_sms = num_sms
+        self.tensor_dtype = tensor_dtype
+        self.buffer_size = int(seqlen * hidden_dim * tensor_dtype.itemsize)
+
+        # Initialize the C++ runtime. Omit timeout_seconds to keep the kernel's default
+        # (DEFAULT_BARRIER_TIMEOUT_SECONDS); it can still be changed later via set_timeout_seconds.
+        if timeout_seconds is None:
+            self.runtime = All2AllCpp(rank, world_size, seqlen, hidden_dim, num_sms, tensor_dtype)
+        else:
+            if not math.isfinite(timeout_seconds) or timeout_seconds < 0:
+                raise ValueError(f"all2all timeout seconds must be finite and non-negative, got {timeout_seconds}")
+            self.runtime = All2AllCpp(rank, world_size, seqlen, hidden_dim, num_sms, tensor_dtype, timeout_seconds)
+
+        # Register in the class-level table so the custom ops can find us.
+        # Auto-cleanup via weakref.finalize covers the case where destroy()
+        # isn't called explicitly — without it the registry would keep the
+        # runtime (and CUDA/IPC resources) alive for the process lifetime.
+        self._comm_id = All2All._next_id
+        All2All._next_id += 1
+        All2All._runtime_registry[self._comm_id] = self.runtime
+        self._finalizer = weakref.finalize(self, All2All._release_comm, self._comm_id, self.runtime)
+
+        # Exchange IPC handles across all ranks
+        ipc_handles: list[Any] = [None] * world_size
+        local_ipc_handle = self.runtime.get_local_ipc_handle()
+        dist.all_gather_object(ipc_handles, local_ipc_handle, group)
+
+        self.runtime.sync(ipc_handles)
+
+    def set_rank_tokens(self, rank_num_tokens: list[int]) -> None:
+        """Sets per-rank token counts on the C++ runtime."""
+        self.runtime.set_rank_tokens(rank_num_tokens)
+
+    def set_timeout_seconds(self, seconds: float) -> None:
+        """Set the all2all barrier (deadlock-detection) timeout, in seconds.
+        The C++ runtime converts to clock cycles using the device's peak SM clock. Raise it
+        during the first ``torch.compile`` forward, where one rank's recompile can delay its
+        all2all launch past the steady-state timeout; reset for steady state.
+        """
+        if not math.isfinite(seconds) or seconds < 0:
+            raise ValueError(f"all2all timeout seconds must be finite and non-negative, got {seconds}")
+        self.runtime.set_timeout_seconds(seconds)
+
+    def send_recv_heads(self, x: torch.Tensor, *, copy_out: bool = False) -> torch.Tensor:
+        """Exchange attention heads across ranks (All2All pattern).
+        Args:
+            x: Input tensor of shape [batch, tokens, heads, head_dim].
+            copy_out: If True, copy result to a new tensor instead of using buffer.
+        Returns:
+            Output tensor with redistributed heads.
+        """
+        return torch.ops.ltx_kernels.send_recv_heads(x, self._comm_id, self.world_size, copy_out)
+
+    def gather_heads(self, x: torch.Tensor, *, copy_out: bool = False) -> torch.Tensor:
+        """Gather heads back to original distribution (reverse All2All).
+        Args:
+            x: Input tensor with distributed heads.
+            copy_out: If True, copy result to a new tensor instead of using buffer.
+        Returns:
+            Output tensor with gathered heads.
+        """
+        return torch.ops.ltx_kernels.gather_heads(x, self._comm_id, self.world_size, copy_out)
+
+    def allgather(self, x: torch.Tensor, *, copy_out: bool = False) -> torch.Tensor:
+        """Allgather operation across all ranks.
+        Args:
+            x: Input tensor to gather.
+            copy_out: If True, copy result to a new tensor instead of using buffer.
+        Returns:
+            Gathered tensor from all ranks.
+        """
+        return self.runtime.allgather(x, copy_out)
+
+    def destroy(self) -> None:
+        """Release IPC handles and GPU memory buffers."""
+        self._finalizer()

packages/ltx-kernels/src/ltx_kernels/arch.py

@@ -0,0 +1,21 @@
+"""GPU architecture detection for blockwise kernel dispatch."""
+
+import torch
+
+
+def get_device_arch() -> str:
+    """Return a coarse architecture name for the current CUDA device.
+    Used to pick the FP8 GEMM kernel variant (SM89 vs SM90) at runtime.
+    """
+    if not torch.cuda.is_available():
+        raise RuntimeError("CUDA is not available; blockwise kernels require a CUDA device.")
+    major, minor = torch.cuda.get_device_capability(torch.cuda.current_device())
+    if major == 8 and (minor >= 0 and minor < 9):
+        return "ampere"
+    if major == 8 and minor == 9:
+        return "ada"
+    if major == 9 and minor == 0:
+        return "hopper"
+    if major in {10, 12}:
+        return "blackwell"
+    raise NotImplementedError(f"Unsupported GPU compute capability sm_{major}{minor}.")

packages/ltx-kernels/src/ltx_kernels/blockwise/__init__.py

@@ -0,0 +1,36 @@
+"""Blockwise FP8/FP6 quantization kernels.
+Importing this subpackage pulls in the compiled ``ops_cpp`` / ``blockwise_cpp``
+extensions (via :mod:`.functional` and :mod:`.linear`) and ``triton``, so it
+raises :class:`ImportError` on hosts where the kernels were not built. Callers
+that want a soft dependency should guard the import (e.g. ``pytest.importorskip``
+or the lazy gate in ``ltx_core.quantization.blockwise``).
+"""
+
+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
+
+__all__ = [
+    "BlockwiseFP6Linear",
+    "BlockwiseFP8Linear",
+    "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",
+]

packages/ltx-kernels/src/ltx_kernels/blockwise/functional.py

@@ -0,0 +1,234 @@
+"""Blockwise FP8/FP6 functional ops.
+The CUDA kernels live in the ``ops_cpp`` extension.
+"""
+
+from typing import Optional, Tuple
+
+import torch
+from ops_cpp import fp6_pack, fp6_unpack
+from ops_cpp import rms_norm_rope as rms_norm_rope_cuda
+from ops_cpp import rms_norm_split_rope as rms_norm_split_rope_cuda
+
+from ltx_kernels.blockwise.triton_ops import (
+    blockwise_dequantize_triton,
+    blockwise_quantize_adanorm_triton,  # noqa: F401  (re-exported for consumers)
+    blockwise_quantize_rms_fma_triton,  # noqa: F401  (re-exported for consumers)
+    blockwise_quantize_triton,  # noqa: F401  (used by linear.BlockwiseGemmLinearFunc)
+    gated_attention_triton,  # noqa: F401  (re-exported for consumers)
+)
+
+# Quantization scale constants for different precisions
+FP8_SCALE_MAX = 448.0
+FP6_SCALE_MAX = 0.1172
+
+
+@torch.library.custom_op("q8_kernels_ops::rms_norm_rope", mutates_args=())
+def _rms_norm_rope_cuda(
+    x: torch.Tensor,
+    weights: Optional[torch.Tensor],
+    cos_freqs: torch.Tensor,
+    sin_freqs: torch.Tensor,
+    out_16bit: bool,
+) -> torch.Tensor:
+    return rms_norm_rope_cuda(x, weights, cos_freqs, sin_freqs, out_16bit)
+
+
+@torch.library.register_fake("q8_kernels_ops::rms_norm_rope")
+def _rms_norm_rope_cuda_fake(
+    x: torch.Tensor,
+    weights: Optional[torch.Tensor],
+    cos_freqs: torch.Tensor,
+    sin_freqs: torch.Tensor,
+    out_16bit: bool,
+) -> torch.Tensor:
+    out = torch.empty_like(x)
+    if out_16bit:
+        return out.to(torch.bfloat16)
+    else:
+        return out.to(torch.float8_e4m3fn)
+
+
+@torch.library.custom_op("q8_kernels_ops::rms_norm_split_rope", mutates_args=())
+def _rms_norm_split_rope_cuda(
+    x: torch.Tensor,
+    sin_freqs: torch.Tensor,
+    cos_freqs: torch.Tensor,
+    weights: torch.Tensor,
+    out_fp8: bool,
+) -> torch.Tensor:
+    return rms_norm_split_rope_cuda(x, sin_freqs, cos_freqs, weights, out_fp8)
+
+
+@torch.library.register_fake("q8_kernels_ops::rms_norm_split_rope")
+def _rms_norm_split_rope_cuda_fake(
+    x: torch.Tensor,
+    sin_freqs: torch.Tensor,
+    cos_freqs: torch.Tensor,
+    weights: torch.Tensor,
+    out_fp8: bool,
+) -> torch.Tensor:
+    out = torch.empty_like(x)
+    if out_fp8:
+        return out.to(torch.float8_e4m3fn)
+    else:
+        return out.to(torch.bfloat16)
+
+
+class RMSNormRope(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, x, weights, cos_freqs, sin_freqs, out_16bit):
+        return torch.ops.q8_kernels_ops.rms_norm_rope(
+            x, weights, cos_freqs, sin_freqs, out_16bit
+        )
+
+
+class RMSNormSplitRope(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, x, sin_freqs, cos_freqs, weights, out_fp8):
+        return torch.ops.q8_kernels_ops.rms_norm_split_rope(
+            x, sin_freqs, cos_freqs, weights, out_fp8
+        )
+
+
+def rms_norm_rope(
+    x: torch.Tensor,
+    cos_freqs: torch.Tensor,
+    sin_freqs: torch.Tensor,
+    weights: Optional[torch.Tensor] = None,
+    out_fp8: bool = False,
+) -> torch.Tensor:
+    # weights=None flows through to the kernel's no-affine path (nullptr); see
+    # rms_norm_rope.cpp, which dispatches the norm_affine template on has_value().
+    return RMSNormRope.apply(x, weights, cos_freqs, sin_freqs, not out_fp8)
+
+
+def rms_norm_split_rope(
+    x: torch.Tensor,
+    cos_freqs: torch.Tensor,
+    sin_freqs: torch.Tensor,
+    weights: Optional[torch.Tensor] = None,
+    out_fp8: bool = False,
+) -> torch.Tensor:
+    """
+    Apply RMS normalization followed by split RoPE (Rotary Position Embedding).
+    Args:
+        x: Input tensor of shape [b, s, h] where h can be 2048, 4096, or 8192
+        cos_freqs: Cos frequencies of shape [b, n, s, d] where n=32 and n*d = h/2
+        sin_freqs: Sin frequencies of shape [b, n, s, d] where n=32
+        weights: Optional RMS norm weights of shape [h]
+        out_fp8: If True, output in float8_e4m3fn, otherwise bfloat16
+    Returns:
+        Output tensor of shape [b, s, h] with dtype based on out_fp8 parameter
+    """
+    if weights is None:
+        weights = torch.ones(x.shape[-1], dtype=x.dtype, device=x.device)
+    return RMSNormSplitRope.apply(x, sin_freqs, cos_freqs, weights, out_fp8)
+
+
+def blockwise_quantize_weights(w: torch.Tensor, block_size=128, scale_max=FP8_SCALE_MAX) -> tuple[torch.Tensor, torch.Tensor]:
+    w = w.view(w.shape[0]//block_size, block_size, w.shape[1]//block_size, block_size).transpose(1, 2).contiguous()
+    w_absmax = w.float().abs().view(w.shape[0], w.shape[1], block_size*block_size).max(dim=-1, keepdim=False).values
+    w_scales = scale_max/w_absmax
+    w_quant = (w.float()*w_scales[:, :, None, None].float()).to(torch.float8_e4m3fn)
+    w_quant = w_quant.transpose(1, 2).contiguous()
+    w_quant = w_quant.view(w.shape[0]*block_size, -1)
+    return w_quant, (1/w_scales).contiguous()
+
+
+def blockwise_quantize_torch(x: torch.Tensor, block_size=128, scale_max=FP8_SCALE_MAX) -> tuple[torch.Tensor, torch.Tensor]:
+    b, n, h = x.shape
+    x = x.view(-1, x.shape[-1]//block_size, block_size)
+    x_absmax = x.float().abs().max(dim=-1, keepdim=True).values
+    x_scales = scale_max/x_absmax
+    x_scaled = (x * x_scales).to(torch.float8_e4m3fn)
+    return x_scaled.view(b, n, h), 1/x_scales.view(b*n, h//block_size).t().contiguous().float()
+
+
+# Precision-specific convenience functions
+
+def fp8_blockwise_quantize_weights_torch(x: torch.Tensor, block_size=128) -> tuple[torch.Tensor, torch.Tensor]:
+    return blockwise_quantize_weights(x, block_size, FP8_SCALE_MAX)
+
+
+def fp6_blockwise_quantize_weights_torch(x: torch.Tensor, block_size=128) -> tuple[torch.Tensor, torch.Tensor]:
+    return blockwise_quantize_weights(x, block_size, FP6_SCALE_MAX)
+
+
+def blockwise_dequantize(x: Tuple[torch.Tensor, torch.Tensor]) -> torch.Tensor:
+    x_fp8, scales = x
+    return blockwise_dequantize_triton(x_fp8, scales)
+
+
+# FP6 Pack/Unpack operations
+@torch.library.custom_op("q8_kernels_ops::fp6_pack", mutates_args=())
+def _fp6_pack(x: torch.Tensor) -> torch.Tensor:
+    """Pack 8-bit tensor to 6-bit by dropping e_1 and e_2 bits.
+    Args:
+        x: Input tensor of shape [m, n] with dtype uint8 or float8_e4m3fn
+    Returns:
+        Packed tensor of shape [m, n*3/4] with dtype uint8
+    """
+    return fp6_pack(x)
+
+
+@torch.library.register_fake("q8_kernels_ops::fp6_pack")
+def _fp6_pack_fake(x: torch.Tensor) -> torch.Tensor:
+    m, n = x.shape
+    n_packed = (n * 3) // 4
+    return torch.empty(m, n_packed, dtype=torch.uint8, device=x.device)
+
+
+@torch.library.custom_op("q8_kernels_ops::fp6_unpack", mutates_args=())
+def _fp6_unpack(x: torch.Tensor, original_n: int) -> torch.Tensor:
+    """Unpack 6-bit tensor back to 8-bit (with e_1 and e_2 set to 0).
+    Args:
+        x: Packed tensor of shape [m, n_packed] with dtype uint8
+        original_n: Original dimension size (unpacked)
+    Returns:
+        Unpacked tensor of shape [m, original_n] with dtype uint8
+    """
+    return fp6_unpack(x, original_n)
+
+
+@torch.library.register_fake("q8_kernels_ops::fp6_unpack")
+def _fp6_unpack_fake(x: torch.Tensor, original_n: int) -> torch.Tensor:
+    m = x.shape[0]
+    return torch.empty(m, original_n, dtype=torch.uint8, device=x.device)
+
+
+# Convenience wrapper functions
+def fp6_pack_tensor(x: torch.Tensor) -> torch.Tensor:
+    """Pack 8-bit tensor to 6-bit by dropping e_1 and e_2 bits.
+    This function packs FP8 weights into a more memory-efficient FP6 format
+    by dropping the two highest exponent bits (e_1 and e_2). This achieves
+    25% memory reduction (4 bytes -> 3 bytes per 4 elements).
+    Args:
+        x: Input tensor of shape [m, n] with dtype uint8 or float8_e4m3fn.
+           n must be divisible by 8.
+    Returns:
+        Packed tensor of shape [m, n*3/4] with dtype uint8
+    Example:
+        >>> w_fp8 = torch.randn(1024, 4096, dtype=torch.float8_e4m3fn, device='cuda')
+        >>> w_fp8_uint8 = w_fp8.view(torch.uint8)
+        >>> w_packed = fp6_pack_tensor(w_fp8_uint8)
+        >>> print(w_packed.shape)  # [1024, 3072]
+    """
+    return torch.ops.q8_kernels_ops.fp6_pack(x)
+
+
+def fp6_unpack_tensor(x: torch.Tensor, original_n: int) -> torch.Tensor:
+    """Unpack 6-bit tensor back to 8-bit (with e_1 and e_2 set to 0).
+    This function unpacks FP6 weights back to FP8 format. Note that the
+    two dropped exponent bits (e_1 and e_2) are restored as 0.
+    Args:
+        x: Packed tensor of shape [m, n_packed] with dtype uint8
+        original_n: Original dimension size before packing (must be divisible by 8)
+    Returns:
+        Unpacked tensor of shape [m, original_n] with dtype uint8
+    Example:
+        >>> w_packed = torch.randint(0, 256, (1024, 3072), dtype=torch.uint8, device='cuda')
+        >>> w_unpacked = fp6_unpack_tensor(w_packed, 4096)
+        >>> w_fp8 = w_unpacked.view(torch.float8_e4m3fn)
+        >>> print(w_fp8.shape)  # [1024, 4096]
+    """
+    return torch.ops.q8_kernels_ops.fp6_unpack(x, original_n)

packages/ltx-kernels/src/ltx_kernels/blockwise/linear.py

@@ -0,0 +1,279 @@
+"""Blockwise FP8/FP6 quantized linear layers and the FP8 GEMM custom op.
+The FP8 GEMM kernels live in the ``blockwise_cpp`` extension.
+"""
+
+from typing import List, Optional
+
+import torch
+import torch.nn as nn
+from ops_cpp import fp6_pack, fp6_unpack
+
+from ltx_kernels.arch import get_device_arch
+from ltx_kernels.blockwise.functional import (
+    blockwise_quantize_triton,
+    blockwise_quantize_weights,
+    fp6_blockwise_quantize_weights_torch,
+)
+
+# Lazily-initialized, architecture-specific FP8 GEMM callable. Importing this
+# module must not require a CUDA device or the compiled extension; the kernel is
+# only resolved on first use (matching the lazy-import gate in ltx-core's
+# blockwise.__init__).
+_fp8_gemm = None
+
+
+def get_fp8_gemm_nt():
+    """Runtime GPU architecture dispatch for FP8 GEMM."""
+    arch = get_device_arch()
+    if arch in ["ada", "blackwell"]:
+        # Ada/Blackwell use SM89 kernels (GeForce path)
+        from blockwise_cpp import fp8_gemm_nt_sm89
+
+        def _func(a, b, d, bias=None, c=None, num_sms=132, use_fast_accum=True):
+            return fp8_gemm_nt_sm89(a, b, d, bias=bias, use_fast_accum=use_fast_accum)
+
+        return _func
+    elif arch == "hopper":
+        # Hopper uses SM90 kernels (H100 path)
+        from blockwise_cpp import fp8_gemm_nt_sm90
+
+        def _func(a, b, d, bias=None, c=None, num_sms=132, use_fast_accum=True):
+            return fp8_gemm_nt_sm90(a, b, d, bias=bias, c=c, num_sms=num_sms)
+
+        return _func
+    else:
+        raise RuntimeError(f"Unsupported GPU architecture: {arch}")
+
+
+def _fp8_gemm_dispatch(a, b, d, bias=None, c=None, num_sms=132, use_fast_accum=True):
+    global _fp8_gemm
+    if _fp8_gemm is None:
+        _fp8_gemm = get_fp8_gemm_nt()
+    return _fp8_gemm(a, b, d, bias=bias, c=c, num_sms=num_sms, use_fast_accum=use_fast_accum)
+
+
+@torch.library.custom_op("blockwise::fp8_gemm", mutates_args=())
+def blockwise_fp8_gemm(
+    a: List[torch.Tensor],
+    b: List[torch.Tensor],
+    bias: Optional[torch.Tensor],
+    use_fast_accum: bool,
+) -> torch.Tensor:
+    d = torch.empty(a[0].shape[0], b[0].shape[0], dtype=torch.bfloat16, device=b[0].device)
+    _fp8_gemm_dispatch(a, b, d, bias=bias, use_fast_accum=use_fast_accum)
+    return d
+
+
+@blockwise_fp8_gemm.register_fake
+def blockwise_fp8_gemm_fake(
+    a: List[torch.Tensor],
+    b: List[torch.Tensor],
+    bias: Optional[torch.Tensor],
+    use_fast_accum: bool,
+) -> torch.Tensor:
+    o = torch.empty(a[0].shape[0], b[0].shape[0], dtype=torch.bfloat16, device=b[0].device)
+    return o
+
+
+def is_16bit(x) -> bool:
+    return x.dtype == torch.float16 or x.dtype == torch.bfloat16
+
+
+class BlockwiseGemmLinearFunc(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        a: tuple[torch.Tensor, torch.Tensor],
+        w: tuple[torch.Tensor],
+        bias: Optional[torch.Tensor],
+    ) -> torch.Tensor:
+
+        is_16bit_a = is_16bit(a[0])
+        if is_16bit_a:
+            a = blockwise_quantize_triton(a[0])
+
+        b, n, h = a[0].shape
+        out_h, _ = w[0].shape
+        new_a = (a[0].view(-1, h), a[1])
+        d = blockwise_fp8_gemm(new_a, w, bias, use_fast_accum=True)
+        return d.view(b, n, out_h)
+
+
+def blockwise_linear_func(a, b, bias=None):
+    return BlockwiseGemmLinearFunc.apply(a, b, bias)
+
+
+class BlockwiseLinear(nn.Module):
+    """Base class for blockwise quantized linear layers."""
+
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        bias: bool = True,
+        device=None,
+        dtype=None,  # for compatibility with nn.Linear
+    ):
+        super().__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+
+        if bias:
+            # FP32 to match the SM89/SM90 GEMM bias contract (the kernels take float*).
+            self.bias = nn.Parameter(
+                torch.empty(out_features, device=device, dtype=torch.float32),
+                requires_grad=False,
+            )
+        else:
+            self.register_parameter("bias", None)
+
+    @property
+    def fp8weight(self):
+        """Should be implemented by subclasses."""
+        raise NotImplementedError("Subclasses must implement the weight property")
+
+    def forward(self, x):
+        if not isinstance(x, tuple):
+            x = (x, None)
+        return blockwise_linear_func(x, (self.fp8weight, self.weight_scale), self.bias)
+
+    @classmethod
+    def from_linear(cls, linear: nn.Linear, transform_weights=True):
+        """Should be implemented by subclasses."""
+        raise NotImplementedError("Subclasses must implement from_linear method")
+
+
+class BlockwiseFP8Linear(BlockwiseLinear):
+    """Blockwise quantized linear layer using FP8 weights."""
+
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        bias: bool = True,
+        device=None,
+        dtype=None,
+    ):
+        super().__init__(in_features, out_features, bias, device, dtype)
+
+        assert in_features % 128 == 0, f"in_features must be divisible by 128, got {in_features}"
+        assert out_features % 128 == 0, f"out_features must be divisible by 128, got {out_features}"
+
+        self.weight = nn.Parameter(
+            torch.empty(
+                out_features, in_features, device=device, dtype=torch.float8_e4m3fn
+            ),
+            requires_grad=False,
+        )
+        self.weight_scale = nn.Parameter(
+            torch.ones(out_features//128, in_features//128, device=device, dtype=torch.float32),
+            requires_grad=False,
+        )
+
+    @property
+    def fp8weight(self):
+        """Return FP8 weight directly."""
+        return self.weight
+
+    @classmethod
+    def from_linear(cls, linear: nn.Linear, transform_weights=True):
+        layer = cls(
+            linear.in_features,
+            linear.out_features,
+            linear.bias is not None,
+            device=linear.weight.device,
+        )
+
+        if transform_weights:
+            w_fp8, w_scales = blockwise_quantize_weights(linear.weight.data.cuda())
+        else:
+            w_fp8 = torch.empty(linear.out_features, linear.in_features, dtype=torch.float8_e4m3fn, device=linear.weight.device)
+            w_scales = torch.ones(
+                w_fp8.shape[0]//128,
+                w_fp8.shape[1]//128,
+                device=w_fp8.device,
+                dtype=torch.float32,
+            )
+
+        layer.weight.data = w_fp8
+        layer.weight_scale.data = w_scales
+
+        if linear.bias is not None:
+            layer.bias.data = linear.bias.data.to(torch.float32)
+
+        return layer
+
+
+class BlockwiseFP6Linear(BlockwiseLinear):
+    """Blockwise quantized linear layer using FP6 weights with 6-bit packing."""
+
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        bias: bool = True,
+        device=None,
+        dtype=None,
+    ):
+        assert in_features % 4 == 0, f"in_features must be divisible by 4, got {in_features}"
+        assert in_features % 8 == 0, f"in_features must be divisible by 8 for fp6_pack, got {in_features}"
+        assert in_features % 128 == 0, f"in_features must be divisible by 128, got {in_features}"
+        assert out_features % 128 == 0, f"out_features must be divisible by 128, got {out_features}"
+
+        super().__init__(in_features, out_features, bias, device, dtype)
+
+        # Store packed weight: [out_features, (in_features // 4)*3]
+        self.weight = nn.Parameter(
+            torch.empty(
+                out_features, (in_features // 4)*3, device=device, dtype=torch.uint8
+            ),
+            requires_grad=False,
+        )
+
+        # Scales are based on unpacked dimensions
+        self.weight_scale = nn.Parameter(
+            torch.ones(out_features//128, in_features//128, device=device, dtype=torch.float32),
+            requires_grad=False,
+        )
+
+    @property
+    def fp8weight(self):
+        """Unpack the 6-bit weights to 8-bit format on-the-fly."""
+        # Unpack from [out_features, (in_features // 4)*3] to [out_features, in_features]
+        unpacked = fp6_unpack(self.weight, self.in_features)
+        # Convert uint8 back to float8_e4m3fn view
+        return unpacked.view(torch.float8_e4m3fn)
+
+    @classmethod
+    def from_linear(cls, linear: nn.Linear, transform_weights=True):
+        layer = cls(
+            linear.in_features,
+            linear.out_features,
+            linear.bias is not None,
+            device=linear.weight.device,
+        )
+
+        if transform_weights:
+            # Quantize to FP6 (actually FP8 with restricted range)
+            w_fp6, w_scales = fp6_blockwise_quantize_weights_torch(linear.weight.data.cuda())
+
+            # Pack the FP6 weights from [out, in] uint8 to [out, in*3/4] uint8
+            w_fp6_uint8 = w_fp6.view(torch.uint8)
+            w_packed = fp6_pack(w_fp6_uint8)
+
+            layer.weight.data = w_packed
+            layer.weight_scale.data = w_scales
+        else:
+            # If not transforming, assume weights are already in packed format
+            layer.weight.data = torch.empty(linear.out_features, (linear.in_features // 4) * 3, dtype=torch.uint8, device=linear.weight.device)
+            layer.weight_scale.data = torch.ones(
+                linear.weight.shape[0]//128,
+                linear.in_features//128,
+                device=linear.weight.device,
+                dtype=torch.float32,
+            )
+
+        if linear.bias is not None:
+            layer.bias.data = linear.bias.data.to(torch.float32)
+
+        return layer

packages/ltx-kernels/src/ltx_kernels/blockwise/triton_ops.py

@@ -0,0 +1,453 @@
+from typing import *
+
+import torch
+from triton import jit
+from triton import language as tl
+from triton import next_power_of_2
+from torch.library import triton_op, wrap_triton
+
+def get_tma_aligned_size(n, element_size):
+    num_elems_tma = 16 // element_size
+    return ((n + num_elems_tma - 1) // num_elems_tma) * num_elems_tma 
+
+@jit
+def _quantize(x, scale_max: tl.constexpr, NUM_BLOCKS: tl.constexpr, BLOCK_SIZE: tl.constexpr):
+    x = tl.reshape(x, (NUM_BLOCKS, BLOCK_SIZE))
+    x_abs = tl.abs(x)
+    x_abs = tl.broadcast_to(x_abs, (NUM_BLOCKS, BLOCK_SIZE))
+    x_absmax = tl.max(x_abs, axis=1)[:, None]
+    x_scales = scale_max / x_absmax
+    x_quant = (x_scales * x).to(tl.float8e4nv)
+    x_out_scales = 1.0/x_scales
+    return x_quant, x_out_scales
+
+@jit
+def _kernel(X, OUT, SCALES, HDIM, BLOCK_SIZE: tl.constexpr):
+    row_idx = tl.program_id(0)
+    x_ptr = X + row_idx * HDIM
+    out_ptr = OUT + row_idx * HDIM
+
+    h_offset = tl.arange(0, BLOCK_SIZE)
+
+    x = tl.load(x_ptr + h_offset, mask=h_offset < HDIM).to(tl.float32)
+    x_scale = 127.0 / tl.max(tl.abs(x))
+    x_scaled = x * x_scale
+    x_scaled += (0.5 * tl.where(x_scaled >= 0, 1, -1)).to(tl.int8)
+
+    tl.store(out_ptr + h_offset, x_scaled, mask=h_offset < HDIM)
+    tl.store(SCALES + row_idx, 1 / x_scale)
+
+def run_quantize_kernel(x: torch.Tensor, out_dtype: Optional[torch.dtype] = None):
+    x_shape_orig = x.shape
+    x = x.view(-1, x_shape_orig[-1])
+    out = torch.empty(x_shape_orig, dtype=torch.int8, device=x.device)
+    scales = torch.empty(x.shape[0], dtype=torch.float, device=x.device)
+
+    BLOCK_SIZE = next_power_of_2(x_shape_orig[-1])
+    grid = (x.shape[0],)
+    _kernel[grid](x, out, scales, x_shape_orig[-1], BLOCK_SIZE, num_warps=4)
+
+    return out.view(x_shape_orig), scales.view(x_shape_orig[:-1])
+
+@jit
+def _block_quant_norm_kernel(
+    X,
+    Norm_Scale,
+    Norm_Shift,
+    Out_scales,
+    X_out,
+    norm_scale_batch_stride: int,
+    norm_shift_batch_stride: int,
+    norm_scale_token_stride: int,
+    norm_shift_token_stride: int,
+    seqlen: int,
+    H: tl.constexpr,
+    BROADCAST_SEQLEN: tl.constexpr,
+    NUM_BLOCKS: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+    TMA_ALIGNED_MN: tl.constexpr,
+    SCALE_MAX: tl.constexpr,
+):
+    token_idx = tl.program_id(0)
+    batch_id =  token_idx // seqlen
+
+    x_ptr = X + token_idx*H + tl.arange(0, H)
+    if BROADCAST_SEQLEN:
+        norm_scale_ptr = Norm_Scale + batch_id*norm_scale_batch_stride + tl.arange(0, H)
+        norm_shift_ptr = Norm_Shift + batch_id*norm_shift_batch_stride + tl.arange(0, H)
+    else:
+        norm_scale_ptr = Norm_Scale + batch_id*norm_scale_batch_stride + (token_idx % seqlen) * norm_scale_token_stride + tl.arange(0, H)
+        norm_shift_ptr = Norm_Shift + batch_id*norm_shift_batch_stride + (token_idx % seqlen) * norm_shift_token_stride + tl.arange(0, H)
+
+    norm_scales = tl.load(norm_scale_ptr)
+    norm_shift = tl.load(norm_shift_ptr)
+    x = tl.load(x_ptr)
+
+    x_sqr = x*x
+    x_norm = tl.sum(x_sqr) / H
+    x_norm = tl.rsqrt(x_norm + 0.00001)
+    x = (x * x_norm).to(tl.bfloat16)
+    x = x * (1.0+norm_scales) + norm_shift
+
+    o_quant, o_scales = _quantize(x, SCALE_MAX, NUM_BLOCKS, BLOCK_SIZE)
+    x_out_ptr = X_out + token_idx*H + BLOCK_SIZE*tl.arange(0, NUM_BLOCKS)[:, None] + tl.arange(0, BLOCK_SIZE)[None, :]
+    tl.store(x_out_ptr, o_quant)
+    out_scales_ptr = Out_scales + token_idx + TMA_ALIGNED_MN*tl.arange(0, NUM_BLOCKS)[:, None]
+    tl.store(out_scales_ptr, o_scales)
+    
+@jit
+def _gelu(x):
+    return x * tl.sigmoid(1.702*x)
+
+@jit
+def _block_quant_kernel(
+    X,
+    Out_scales,
+    X_out,
+    seqlen: int,
+    H: tl.constexpr,
+    NUM_BLOCKS: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+    TMA_ALIGNED_MN: tl.constexpr,
+    USE_GELU: tl.constexpr,
+    SCALE_MAX: tl.constexpr,
+):
+    batch_id, seq_id = tl.program_id(0), tl.program_id(1)
+    token_idx = batch_id*seqlen + seq_id
+    x_ptr = X + token_idx*H + BLOCK_SIZE*tl.arange(0, NUM_BLOCKS)[:, None] + tl.arange(0, BLOCK_SIZE)[None, :]
+    x = tl.load(x_ptr).to(tl.float32)
+    if USE_GELU:
+        x = _gelu(x)
+
+    o_quant, o_scales = _quantize(x, SCALE_MAX, NUM_BLOCKS, BLOCK_SIZE)
+    x_out_ptr = X_out + token_idx*H + BLOCK_SIZE*tl.arange(0, NUM_BLOCKS)[:, None] + tl.arange(0, BLOCK_SIZE)[None, :]
+    tl.store(x_out_ptr, o_quant)
+    out_scales_ptr = Out_scales + token_idx + TMA_ALIGNED_MN*tl.arange(0, NUM_BLOCKS)[:, None]
+    tl.store(out_scales_ptr, o_scales)
+
+@triton_op("blockwise::quantize", mutates_args=())
+def _quant_blockwise_tma_aligned_func(x: torch.Tensor, use_gelu: bool) -> Tuple[torch.Tensor, torch.Tensor]:
+    b, s, h = x.shape
+    num_rows = b*s
+    block_size = 128
+    num_blocks = h // block_size
+    out = torch.empty((num_rows, h), device=x.device, dtype=torch.float8_e4m3fn)
+    tma_aligned_mn = get_tma_aligned_size(num_rows, 4)
+    scales = torch.empty_strided((num_rows, num_blocks), (1, tma_aligned_mn), dtype=torch.float, device=x.device)
+    scale_max: float = 448.0
+    wrap_triton(_block_quant_kernel)[(b, s)](
+        x,
+        scales,
+        out,
+        H=h,
+        seqlen=s,
+        BLOCK_SIZE=128,
+        NUM_BLOCKS=num_blocks,
+        TMA_ALIGNED_MN=tma_aligned_mn,
+        USE_GELU=use_gelu,
+        SCALE_MAX=scale_max
+    )
+
+    return out.view(b, s, h), scales
+
+def run_quant_blockwise_tma_aligned(x):
+    return _quant_blockwise_tma_aligned_func(x, False)
+
+def run_quant_blockwise_gelu_tma_aligned(x):
+    return _quant_blockwise_tma_aligned_func(x, True)
+
+@triton_op("blockwise::adanorm", mutates_args=())
+def run_quant_blockwise_norm_tma_aligned(x: torch.Tensor, w: Optional[torch.Tensor], norm_scale: torch.Tensor, norm_shift: torch.Tensor, out_dtype: torch.dtype, hd_scale: Optional[float]) -> Tuple[torch.Tensor, torch.Tensor]:
+    b, s, h = x.shape
+    is_broadcast = (norm_scale.shape[1] == 1)
+    num_rows = b*s
+    block_size = 128
+    num_blocks = h // block_size
+    out = torch.empty((num_rows, h), device=x.device, dtype=torch.float8_e4m3fn)
+    tma_aligned_mn = get_tma_aligned_size(num_rows, 4)
+    scales = torch.empty_strided((num_rows, num_blocks), (1, tma_aligned_mn), dtype=torch.float, device=x.device)
+    norm_scale_batch_stride = norm_scale.stride(0)
+    norm_shift_batch_stride = norm_shift.stride(0)
+
+    scale_max: float = 448.0
+    wrap_triton(_block_quant_norm_kernel)[(num_rows, 1, 1)](
+        x,
+        norm_scale,
+        norm_shift,
+        scales,
+        out,
+        norm_scale_batch_stride=norm_scale_batch_stride,
+        norm_shift_batch_stride=norm_shift_batch_stride,
+        norm_scale_token_stride=norm_scale.stride(1),
+        norm_shift_token_stride=norm_shift.stride(1),
+        H=h,
+        seqlen=s,
+        BROADCAST_SEQLEN=is_broadcast,
+        BLOCK_SIZE=128,
+        NUM_BLOCKS=num_blocks,
+        TMA_ALIGNED_MN=tma_aligned_mn,
+        SCALE_MAX=scale_max
+    )
+
+    return out.view(b, s, h), scales
+
+@jit
+def _quant_rms_sum_mult_kernel(
+    X,
+    Y,
+    Z,
+    Out,
+    Out_scales,
+    seqlen: int,
+    z_batch_stride: int,
+    z_token_stride: int,
+    TMA_ALIGNED_MN: int,
+    H: tl.constexpr,
+    NUM_BLOCKS: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+    IS_Z_BROADCAST: tl.constexpr,
+    QUANTIZE: tl.constexpr,
+    SCALE_MAX: tl.constexpr,
+):
+    token_idx = tl.program_id(0)
+    batch_id =  token_idx // seqlen
+
+    x_ptr = X + token_idx*H + tl.arange(0, H)
+    y_ptr = Y + token_idx*H + tl.arange(0, H)
+    if IS_Z_BROADCAST:
+        z_ptr = Z + batch_id*z_batch_stride + tl.arange(0, H)
+    else:
+        z_ptr = Z + batch_id*z_batch_stride + (token_idx % seqlen) * z_token_stride + tl.arange(0, H)
+    x = tl.load(x_ptr)
+    y = tl.load(y_ptr)
+    z = tl.load(z_ptr)
+
+    o = x + y*z
+
+    tl.store(x_ptr, o)
+
+    o_sqr = o * o
+    o_norm = tl.sum(o_sqr, axis=0) / H
+    o_inv = tl.rsqrt(o_norm)
+    o *= o_inv
+
+    if QUANTIZE:
+        o_quant, o_scales = _quantize(o, SCALE_MAX, NUM_BLOCKS, BLOCK_SIZE)
+        x_out_ptr = Out + token_idx*H + BLOCK_SIZE*tl.arange(0, NUM_BLOCKS)[:, None] + tl.arange(0, BLOCK_SIZE)[None, :]
+        tl.store(x_out_ptr, o_quant)
+        out_scales_ptr = Out_scales + token_idx + TMA_ALIGNED_MN*tl.arange(0, NUM_BLOCKS)[:, None]
+        tl.store(out_scales_ptr, o_scales)
+    else:
+        x_out_ptr = Out + token_idx*H + tl.arange(0, H)
+        tl.store(x_out_ptr, o)
+
+@triton_op("blockwise::quant_rms_fma", mutates_args=("x", ))
+def run_quant_blockwise_rms_fma(x: torch.Tensor, y: torch.Tensor, z: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+    b, s, h = x.shape
+    scale_max: float = 448.0
+    is_z_broadcast = (z.shape[1] == 1)
+    num_rows = b*s
+    block_size = 128
+    num_blocks = h // block_size
+    out = torch.empty((num_rows, h), device=x.device, dtype=torch.float8_e4m3fn)
+    tma_aligned_mn = get_tma_aligned_size(num_rows, 4)
+    scales = torch.empty_strided((num_rows, num_blocks), (1, tma_aligned_mn), dtype=torch.float, device=x.device)
+    z_batch_stride = z.stride(0)
+
+    wrap_triton(_quant_rms_sum_mult_kernel)[(num_rows, 1, 1)](
+        x,
+        y,
+        z,
+        out,
+        scales,
+        seqlen=s,
+        z_batch_stride=z_batch_stride,
+        z_token_stride=z.stride(1),
+        TMA_ALIGNED_MN=tma_aligned_mn,
+        H=h,
+        NUM_BLOCKS=num_blocks,
+        BLOCK_SIZE=128,
+        IS_Z_BROADCAST=is_z_broadcast,
+        QUANTIZE=True,
+        SCALE_MAX=scale_max
+    )
+
+    return out.view(b, s, h), scales
+
+
+@triton_op("blockwise::rms_fma", mutates_args=("x", ))
+def run_rms_fma(x: torch.Tensor, y: torch.Tensor, z: torch.Tensor) -> torch.Tensor:
+    b, s, h = x.shape
+    is_z_broadcast = (z.shape[1] == 1)
+    num_rows = b*s
+    block_size = 128
+    num_blocks = h // block_size
+    out = torch.empty((num_rows, h), device=x.device, dtype=torch.bfloat16)
+    z_batch_stride = z.stride(0)
+    
+    wrap_triton(_quant_rms_sum_mult_kernel)[(num_rows, 1, 1)](
+        x,
+        y,
+        z,
+        out,
+        None,
+        seqlen=s,
+        z_batch_stride=z_batch_stride,
+        z_token_stride=z.stride(1),
+        TMA_ALIGNED_MN=0,
+        H=h,
+        NUM_BLOCKS=num_blocks,
+        BLOCK_SIZE=128,
+        IS_Z_BROADCAST=is_z_broadcast,
+        QUANTIZE=False,
+        SCALE_MAX=448.0
+    )
+
+    return out.view(b, s, h)
+
+
+@jit
+def _gated_attention_kernel(
+    X,
+    Gate_Logits,
+    Out,
+    Out_scales,
+    H: tl.constexpr,
+    NUM_HEADS: tl.constexpr,
+    DIM_HEAD: tl.constexpr,
+    NUM_BLOCKS: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+    TMA_ALIGNED_MN: tl.constexpr,
+    QUANTIZE: tl.constexpr,
+    SCALE_MAX: tl.constexpr,
+):
+    token_idx = tl.program_id(0)
+
+    # Load gate logits for all heads: shape [NUM_HEADS]
+    gate_logits = tl.load(Gate_Logits + token_idx * NUM_HEADS + tl.arange(0, NUM_HEADS)).to(tl.float32)
+    # 2*sigmoid so that zero-init gives identity (2 * 0.5 = 1.0)
+    gates = 2.0 * tl.sigmoid(gate_logits)                           # [NUM_HEADS]
+    gates = tl.broadcast_to(gates[:, None], (NUM_HEADS, DIM_HEAD))  # [NUM_HEADS, DIM_HEAD]
+
+    # Load x viewed as [NUM_HEADS, DIM_HEAD]
+    offsets = tl.arange(0, NUM_HEADS)[:, None] * DIM_HEAD + tl.arange(0, DIM_HEAD)[None, :]
+    x = tl.load(X + token_idx * H + offsets).to(tl.float32)
+
+    # Apply per-head gating
+    gated = x * gates  # [NUM_HEADS, DIM_HEAD]
+
+    if QUANTIZE:
+        o_quant, o_scales = _quantize(gated, SCALE_MAX, NUM_BLOCKS, BLOCK_SIZE)
+        out_ptr = Out + token_idx * H + BLOCK_SIZE * tl.arange(0, NUM_BLOCKS)[:, None] + tl.arange(0, BLOCK_SIZE)[None, :]
+        tl.store(out_ptr, o_quant)
+        scales_ptr = Out_scales + token_idx + TMA_ALIGNED_MN * tl.arange(0, NUM_BLOCKS)[:, None]
+        tl.store(scales_ptr, o_scales)
+    else:
+        out_ptr = Out + token_idx * H + offsets
+        tl.store(out_ptr, gated)
+
+
+@triton_op("blockwise::gated_attention", mutates_args=())
+def run_gated_attention(x: torch.Tensor, gate_logits: torch.Tensor, quantize: bool = True) -> Tuple[torch.Tensor, torch.Tensor]:
+    b, t, h = x.shape
+    num_heads = gate_logits.shape[-1]
+    dim_head = h // num_heads
+    block_size = 128
+    num_blocks = h // block_size
+    num_rows = b * t
+    scale_max: float = 448.0
+
+    if quantize:
+        out = torch.empty((num_rows, h), device=x.device, dtype=torch.float8_e4m3fn)
+        tma_aligned_mn = get_tma_aligned_size(num_rows, 4)
+        scales = torch.empty_strided((num_rows, num_blocks), (1, tma_aligned_mn), dtype=torch.float, device=x.device)
+    else:
+        out = torch.empty((num_rows, h), device=x.device, dtype=torch.bfloat16)
+        tma_aligned_mn = 0
+        scales = None
+
+    wrap_triton(_gated_attention_kernel)[(num_rows,)](
+        x,
+        gate_logits,
+        out,
+        scales,
+        H=h,
+        NUM_HEADS=num_heads,
+        DIM_HEAD=dim_head,
+        NUM_BLOCKS=num_blocks,
+        BLOCK_SIZE=128,
+        TMA_ALIGNED_MN=tma_aligned_mn,
+        QUANTIZE=quantize,
+        SCALE_MAX=scale_max,
+    )
+
+    return out.view(b, t, h), scales
+
+
+@jit
+def _blockwise_dequantize_kernel(
+    X,
+    Scales,
+    Out,
+    scales_row_stride: int,
+    scales_col_stride: int,
+    H: tl.constexpr,
+    NUM_BLOCKS: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+):
+    token_idx = tl.program_id(0)
+
+    # Load per-block scales [NUM_BLOCKS] respecting the actual tensor strides
+    scales = tl.load(
+        Scales + token_idx * scales_row_stride + scales_col_stride * tl.arange(0, NUM_BLOCKS)
+    ).to(tl.float32)
+    scales = tl.broadcast_to(scales[:, None], (NUM_BLOCKS, BLOCK_SIZE))
+
+    # Load fp8 input viewed as [NUM_BLOCKS, BLOCK_SIZE]
+    x_offsets = BLOCK_SIZE * tl.arange(0, NUM_BLOCKS)[:, None] + tl.arange(0, BLOCK_SIZE)[None, :]
+    x = tl.load(X + token_idx * H + x_offsets).to(tl.float32)
+
+    # Dequantize and store as bf16
+    tl.store(Out + token_idx * H + x_offsets, (x * scales).to(tl.bfloat16))
+
+
+@triton_op("blockwise::dequantize", mutates_args=())
+def run_blockwise_dequantize(x: torch.Tensor, scales: torch.Tensor) -> torch.Tensor:
+    b, s, h = x.shape
+    block_size = 128
+    num_blocks = h // block_size
+    num_rows = b * s
+    out = torch.empty((num_rows, h), device=x.device, dtype=torch.bfloat16)
+
+    wrap_triton(_blockwise_dequantize_kernel)[(num_rows,)](
+        x,
+        scales,
+        out,
+        scales_row_stride=scales.stride(0),
+        scales_col_stride=scales.stride(1),
+        H=h,
+        NUM_BLOCKS=num_blocks,
+        BLOCK_SIZE=128,
+    )
+
+    return out.view(b, s, h)
+
+
+rowwise_int_quantize_triton = run_quantize_kernel
+blockwise_quantize_adanorm_triton = run_quant_blockwise_norm_tma_aligned
+blockwise_quantize_triton = run_quant_blockwise_tma_aligned
+blockwise_quantize_gelu_triton = run_quant_blockwise_gelu_tma_aligned
+blockwise_quantize_rms_fma_triton = run_quant_blockwise_rms_fma
+rms_fma_triton = run_rms_fma
+gated_attention_triton = run_gated_attention
+blockwise_dequantize_triton = run_blockwise_dequantize
+
+# # Precision-specific convenience functions for Triton kernels
+# def fp8_blockwise_quantize_triton(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+#     return run_quant_blockwise_tma_aligned(x, scale_max=FP8_SCALE_MAX)
+
+# def fp7_blockwise_quantize_triton(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+#     return run_quant_blockwise_tma_aligned(x, scale_max=FP7_SCALE_MAX)
+
+# def fp6_blockwise_quantize_triton(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+#     return run_quant_blockwise_tma_aligned(x, scale_max=FP6_SCALE_MAX)

packages/ltx-pipelines/CLAUDE.md

@@ -49,7 +49,7 @@ Inference pipelines for LTX-2 audio-video generation. Depends on `ltx-core` for
 
 ## Shared building blocks (`utils/blocks.py`)
 
-- `DiffusionStage` -- owns transformer lifecycle; builds model on call, frees on exit via `gpu_model()` context manager (moves params to meta device to release GPU/CPU memory). Accepts optional `stepper` and `loop` overrides.
+- `DiffusionStage` -- owns transformer lifecycle; builds model on call, frees on exit via `gpu_model()` context manager (moves params to meta device to release GPU/CPU memory). Accepts optional `stepper` and `loop` overrides. `__init__` takes a pre-built transformer builder; pipelines construct it via the `DiffusionStage.from_checkpoint(checkpoint_path, ..., loras=...)` classmethod, which builds the standard (and, when offloading, streaming) builders. `with_builder` / `with_loras` return a new stage with a swapped builder / LoRA set without re-specifying config.
 - `PromptEncoder` -- Gemma text encoder + embeddings processor (video 4096-dim, audio 2048-dim).
 - `ImageConditioner` / `AudioConditioner` -- temporary encoder scope; builds encoder, passes to callable, frees.
 - `VideoUpsampler` -- 2x spatial upsampling via encoder + upsampler.

packages/ltx-pipelines/README.md

@@ -1,15 +1,9 @@
 # LTX-2 Pipelines
 
-High-level pipeline implementations for generating audio-video content with Lightricks' **LTX-2** model. This package provides ready-to-use pipelines for text-to-video, image-to-video, video-to-video, and keyframe interpolation tasks.
+High-level pipeline implementations for generating audio-video content with Lightricks' **LTX-2** model. This package provides ready-to-use pipelines for text-to-video, image-to-video, video-to-video, audio-to-video, keyframe interpolation, and retake tasks.
 
 Pipelines are built using building blocks from [`ltx-core`](../ltx-core/) (schedulers, guiders, noisers, patchifiers) and handle the complete inference flow including model loading, encoding, decoding, and file I/O.
 
----
-
-## 📋 Overview
-
-LTX-2 Pipelines provides production-ready implementations that abstract away the complexity of the diffusion process, model loading, and memory management. Each pipeline is optimized for specific use cases and offers different trade-offs between speed, quality, and memory usage.
-
 **Key Features:**
 
 - 🎬 **Multiple Pipeline Types**: Text-to-video, image-to-video, video-to-video, audio-to-video, keyframe interpolation, and retake
@@ -19,27 +13,12 @@ LTX-2 Pipelines provides production-ready implementations that abstract away the
 - 📦 **Self-Contained**: Handles model loading, encoding, decoding, and file I/O
 - 🚀 **CLI Support**: All pipelines can be run as command-line scripts
 
----
-
-## 🚀 Quick Start
-
-`ltx-pipelines` provides ready-made inference pipelines for text-to-video, image-to-video, video-to-video, audio-to-video, keyframe interpolation, and retake. Built using building blocks from [`ltx-core`](../ltx-core/), these pipelines handle the complete inference flow including model loading, encoding, decoding, and file I/O.
-
-## 🔧 Installation
+## Quick Start
 
 ```bash
 # From the repository root
 uv sync --frozen
 
-# Or install as a package
-pip install -e packages/ltx-pipelines
-```
-
-### Running Pipelines
-
-All pipelines can be run directly from the command line. Each pipeline module is executable:
-
-```bash
 # Run a pipeline (example: two-stage text-to-video)
 python -m ltx_pipelines.ti2vid_two_stages \
     --checkpoint-path path/to/checkpoint.safetensors \
@@ -48,544 +27,21 @@ python -m ltx_pipelines.ti2vid_two_stages \
     --gemma-root path/to/gemma \
     --prompt "A beautiful sunset over the ocean" \
     --output-path output.mp4
-
-# View all available options for any pipeline
-python -m ltx_pipelines.ti2vid_two_stages --help
 ```
 
-Available pipeline modules:
-
-- `ltx_pipelines.ti2vid_two_stages` - Two-stage text/image-to-video (recommended).
-- `ltx_pipelines.ti2vid_two_stages_hq` - Two-stage text/image-to-video (different sampler, better quality).
-- `ltx_pipelines.ti2vid_one_stage` - Single-stage text/image-to-video.
-- `ltx_pipelines.t2a_one_stage` - Single-stage text-to-audio (audio-only output).
-- `ltx_pipelines.distilled` - Fast text/image-to-video pipeline using only the distilled model.
-- `ltx_pipelines.ic_lora` - Video-to-video with IC-LoRA.
-- `ltx_pipelines.keyframe_interpolation` - Keyframe interpolation.
-- `ltx_pipelines.a2vid_two_stage` - Audio-to-video generation conditioned on an input audio.
-- `ltx_pipelines.retake` - Regenerate a time region of an existing video.
-- `ltx_pipelines.hdr_ic_lora` - Video-to-video with HDR output (linear float via LogC3 inverse decode).
-- `ltx_pipelines.lipdub` - Lip dubbing / re-voicing with IC-LoRA and audio reference conditioning.
-
-Use `--help` with any pipeline module to see all available options and parameters.
-
----
-
-## 🎯 Pipeline Selection Guide
-
-### Quick Decision Tree
-
-```text
-Do you have an existing video to modify?
-├─ YES → Use RetakePipeline (regenerate a specific time region)
-│
-Do you have an audio file to drive generation?
-├─ YES → Use A2VidPipelineTwoStage (audio-to-video)
-│
-Do you need HDR output (linear float frames for EXR / tonemapping)?
-├─ YES → Use HDRICLoraPipeline (video-to-video with LogC3 inverse decode)
-│
-Do you need to condition on existing images/videos?
-├─ YES → Do you have reference videos for video-to-video?
-│  ├─ YES → Use ICLoraPipeline
-│  └─ NO → Do you have multiple keyframe images to interpolate?
-│     ├─ YES → Use KeyframeInterpolationPipeline
-│     └─ NO → Use TI2VidTwoStagesPipeline (image conditioning only)
-│
-└─ NO → Text-to-video only
-   ├─ Do you need best quality?
-   │  └─ YES → Use TI2VidTwoStagesPipeline (recommended for production)
-   │
-   └─ Do you need fastest inference?
-      └─ YES → Use DistilledPipeline (with 8 predefined sigmas)
-```
-
-> **Note:** [`TI2VidOneStagePipeline`](src/ltx_pipelines/ti2vid_one_stage.py) is primarily for educational purposes. For best quality, use two-stage pipelines ([`TI2VidTwoStagesPipeline`](src/ltx_pipelines/ti2vid_two_stages.py), [`TI2VidTwoStagesHQPipeline`](src/ltx_pipelines/ti2vid_two_stages_hq.py), [`ICLoraPipeline`](src/ltx_pipelines/ic_lora.py), [`KeyframeInterpolationPipeline`](src/ltx_pipelines/keyframe_interpolation.py), [`A2VidPipelineTwoStage`](src/ltx_pipelines/a2vid_two_stage.py), or [`DistilledPipeline`](src/ltx_pipelines/distilled.py)). For editing existing videos, use [`RetakePipeline`](src/ltx_pipelines/retake.py).
-
-### Features Comparison
-
-| Pipeline | Stages | [Multimodal Guidance](#%EF%B8%8F-multimodal-guidance) | Upsampling | Conditioning | Best For |
-| -------- | ------ | --- | ---------- | ------------- | -------- |
-| **TI2VidTwoStagesPipeline** | 2 | ✅ | ✅ | Image | **Production quality** (recommended) |
-| **TI2VidTwoStagesHQPipeline** | 2 | ✅ | ✅ | Image | Same as above, res_2s sampler (higher quality) |
-| **TI2VidOneStagePipeline** | 1 | ✅ | ❌ | Image | Educational, prototyping |
-| **DistilledPipeline** | 2 | ❌ | ✅ | Image | Fastest inference (8 sigmas) |
-| **ICLoraPipeline** | 2 | ✅ | ✅ | Image + Video | Video-to-video transformations |
-| **KeyframeInterpolationPipeline** | 2 | ✅ | ✅ | Keyframes | Animation, interpolation |
-| **A2VidPipelineTwoStage** | 2 | ✅ | ✅ | Audio + Image | Audio-driven video generation |
-| **RetakePipeline** | 1 | ✅ | ❌ | Source Video | Regenerating a time region of a video |
-| **HDRICLoraPipeline** | 2 | ❌ | ✅ | Video | HDR video-to-video (linear float output for EXR) |
-| **LipDubPipeline** | 2 | ✅ | ✅ | Video + Audio | Lip dubbing with audio ref conditioning |
-
----
-
-## 📦 Available Pipelines
-
-### 1. TI2VidTwoStagesPipeline
-
-**Best for:** High-quality text/image-to-video generation with upsampling. **Recommended for production use.**
-
-**Source**: [`src/ltx_pipelines/ti2vid_two_stages.py`](src/ltx_pipelines/ti2vid_two_stages.py)
-
-Two-stage generation: Stage 1 generates low-resolution video with [multimodal guidance](#%EF%B8%8F-multimodal-guidance), Stage 2 upsamples to 2x resolution with distilled LoRA refinement. Supports image conditioning. Highest quality output, slower than one-stage but significantly better quality.
-
-**Use when:** Production-quality video generation, higher resolution needed, quality over speed, text-to-video with image conditioning.
-
----
-
-### 2. TI2VidTwoStagesHQPipeline
-
-**Best for:** Same two-stage text/image-to-video as TI2VidTwoStagesPipeline but with a different sampler and step count.
-
-**Source**: [`src/ltx_pipelines/ti2vid_two_stages_hq.py`](src/ltx_pipelines/ti2vid_two_stages_hq.py)
-
-Uses the **res_2s** second-order sampler instead of Euler. Same stage structure (stage 1 at target resolution with CFG, stage 2 upsampling with distilled LoRA) and image conditioning support. Typically allows fewer steps for comparable quality; trade-offs differ from the default Euler-based pipeline.
-
-**Use when:** You want the same two-stage workflow with fewer steps or prefer the res_2s sampling behavior.
-
----
-
-### 3. TI2VidOneStagePipeline
-
-**Best for:** Educational purposes and quick prototyping.
-
-**Source**: [`src/ltx_pipelines/ti2vid_one_stage.py`](src/ltx_pipelines/ti2vid_one_stage.py)
-
-> **⚠️ Important:** This pipeline is primarily for educational purposes. For production-quality results, use `TI2VidTwoStagesPipeline` or other two-stage pipelines.
-
-Single-stage generation (no upsampling) with [multimodal guidance](#%EF%B8%8F-multimodal-guidance) and image conditioning support. Faster inference but lower resolution output (typically 512x768).
-
-**Use when:** Learning how the pipeline works, quick prototyping, testing, or when high resolution is not needed.
-
----
-
-### 4. DistilledPipeline
-
-**Best for:** Fastest inference with good quality using a distilled model with predefined sigma schedule.
-
-**Source**: [`src/ltx_pipelines/distilled.py`](src/ltx_pipelines/distilled.py)
-
-Two-stage generation with 8 predefined sigmas (8 steps in stage 1, 4 steps in stage 2). No guidance required. Fastest inference among all pipelines. Supports image conditioning. Requires spatial upsampler.
-
-**Use when:** Fastest inference is critical, batch processing many videos, or when you have a distilled model checkpoint.
-
----
-
-### 5. ICLoraPipeline
-
-**Best for:** Video-to-video and image-to-video transformations using IC-LoRA.
-
-**Source**: [`src/ltx_pipelines/ic_lora.py`](src/ltx_pipelines/ic_lora.py)
-
-Two-stage generation with IC-LoRA support. Can condition on reference videos (video-to-video) or images at specific frames. CFG guidance in stage 1, upsampling in stage 2. Requires IC-LoRA trained model.
-
-**Note:** ICLoraPipeline can only be used with a distilled model.
-
-**Use when:** Video-to-video transformations, image-to-video with strong control, or when you have reference videos to guide generation.
-
----
-
-### 6. KeyframeInterpolationPipeline
-
-**Best for:** Generating videos by interpolating between keyframe images.
-
-**Source**: [`src/ltx_pipelines/keyframe_interpolation.py`](src/ltx_pipelines/keyframe_interpolation.py)
-
-Two-stage generation with keyframe interpolation. Uses guiding latents (additive conditioning) instead of replacing latents for smoother transitions. [Multimodal guidance](#%EF%B8%8F-multimodal-guidance) in stage 1, upsampling in stage 2.
-
-**Use when:** You have keyframe images and want to interpolate between them, creating smooth transitions, or animation/motion interpolation tasks.
-
----
-
-### 7. A2VidPipelineTwoStage
-
-**Best for:** Generating video driven by an input audio.
-
-**Source**: [`src/ltx_pipelines/a2vid_two_stage.py`](src/ltx_pipelines/a2vid_two_stage.py)
-
-Two-stage audio-to-video generation. Stage 1 generates video at half resolution with audio conditioning (video-only denoising with the audio frozen), then Stage 2 upsamples by 2x and refines the video while keeping the audio fixed, using a distilled LoRA. The input audio is encoded via the audio VAE and used as the initial audio latent, but the original audio waveform is passed through and returned in the output to preserve fidelity. Supports image conditioning and prompt enhancement.
-
-**Extra CLI arguments:** `--audio-path` (required), `--audio-start-time`, `--audio-max-duration`.
-
-**Use when:** You have an audio clip and want to generate a matching video, audio-reactive video generation, or music visualization.
-
----
-
-### 8. RetakePipeline
-
-**Best for:** Regenerating a specific time region of an existing video while keeping the rest unchanged.
-
-**Source**: [`src/ltx_pipelines/retake.py`](src/ltx_pipelines/retake.py)
-
-Single-stage generation that encodes the source video and audio into latents, applies a temporal region mask to mark `[start_time, end_time]` for regeneration, and denoises only the masked region from a text prompt. Content outside the time window is preserved. Supports independent control over video and audio regeneration (`regenerate_video`, `regenerate_audio` flags), and can use either the full model with CFG guidance or the distilled model with a fixed sigma schedule.
-
-**Extra CLI arguments:** `--video-path` (required), `--start-time` (required), `--end-time` (required).
-
-**Constraints:** Source video frame count must satisfy the 8k+1 format (e.g. 97, 193) and resolution must be multiples of 32.
-
-**Use when:** You want to re-do a specific section of a generated video (e.g. fix a bad segment), selectively regenerate audio or video in a time window, or iterate on part of a result without re-generating the entire clip.
-
----
-
-### 9. HDRICLoraPipeline
-
-**Best for:** Video-to-video generation with HDR output for EXR export and offline tonemapping.
-
-**Source**: [`src/ltx_pipelines/hdr_ic_lora.py`](src/ltx_pipelines/hdr_ic_lora.py)
-
-Two-stage video-to-video on the distilled model with an HDR IC-LoRA. Decoded latents pass through an HDR inverse transform (ARRI LogC3, auto-detected from LoRA metadata) to produce a **linear HDR float** tensor `[f, h, w, c]`. Video-only (audio skipped). Text embeddings are pre-computed externally and loaded from a `.safetensors` file. Tonemapping and EXR saving are the caller's responsibility. LoRA and embeddings: [`Lightricks/LTX-2.3-22b-IC-LoRA-HDR`](https://huggingface.co/Lightricks/LTX-2.3-22b-IC-LoRA-HDR).
-
-**Extra CLI arguments:** `--input` (mp4 or directory, required), `--output-dir` (required), `--hdr-lora` (required), `--text-embeddings` (pre-computed `.safetensors`, required), `--num-frames`, `--spatial-tile` (tiled VAE decode tile size; reduce on lower-VRAM GPUs), `--skip-mp4` (EXR only, no H.264 preview), `--exr-half` (float16 EXR), `--high-quality` (generates 2x frames internally for smoother output, ~2x slower), `--offload {none,cpu,disk}` (weight offloading; disables FP8 quantization when not `none`).
-
-**Use when:** You need linear HDR float output for EXR export, color grading, or custom tonemapping workflows.
-
----
-
-### 10. LipDubPipeline
-
-**Best for:** Lip dubbing, rephrasing while keeping the same speaker identity and matching lip movements to new audio.
-
-**Source**: [`src/ltx_pipelines/lipdub.py`](src/ltx_pipelines/lipdub.py)
-
-Uses IC-LoRA on a **distilled** checkpoint with a **single** lip-dub IC-LoRA applied in **both** stages. The reference clip provides video and audio reference tokens whose VAE latents are appended to the target audio sequence as frozen reference tokens. The frame count and frame rate are derived from the reference video (frame count is silently snapped to the nearest `8k+1`), so the CLI does not accept `--num-frames` or `--frame-rate`. Required: `--reference-video`. Optional: `--reference-strength`. LoRA: [`Lightricks/LTX-2.3-22b-IC-LoRA-LipDub`](https://huggingface.co/Lightricks/LTX-2.3-22b-IC-LoRA-LipDub).
-
-**Note:** Requires a distilled model checkpoint and one lip-dub IC-LoRA (`--lora` exactly once).
-
-**Use when:** Dubbing, rephrasing with matched lips and speaker identity.
-
----
-
-### 11. T2AOneStagePipeline
-
-**Best for:** Text-to-audio — generating speech/audio only (no video) from a text prompt, e.g. driving an audio-style LoRA such as an accent LoRA.
-
-**Source**: [`src/ltx_pipelines/t2a_one_stage.py`](src/ltx_pipelines/t2a_one_stage.py)
-
-Single-stage, **audio-only** generation: the video branch is absent (`video=None`), so only the audio modality is denoised and decoded through the audio VAE + vocoder, producing a wave file. Audio duration is derived from `--num-frames` / `--frame-rate` (the same `8k+1` frame convention as video). Audio guidance (CFG/STG) is optional — the `--audio-*` flags default to the model's values; the video→audio cross-modal guidance is disabled since there is no video modality.
-
-**Extra CLI arguments (all optional, with sensible defaults):** `--num-frames`, `--frame-rate`, `--negative-prompt`, `--audio-cfg-guidance-scale`, `--audio-stg-guidance-scale`, `--audio-stg-blocks`, `--audio-rescale-scale`, `--audio-skip-step`. No `--height/--width/--image` (audio has no spatial dimensions).
-
-**Use when:** You need speech/audio from text alone, or to evaluate an audio-only LoRA (accent, voice style) without generating video.
-
----
-
-## 🎨 Conditioning Types
-
-Pipelines use different conditioning methods from [`ltx-core`](../ltx-core/) for controlling generation. See the [ltx-core conditioning documentation](../ltx-core/README.md#conditioning--control) for details.
-
-### Image Conditioning
-
-All pipelines support image conditioning, but with different methods:
-
-- **Replacing Latents** ([`image_conditionings_by_replacing_latent`](src/ltx_pipelines/utils/helpers.py)):
-  - Used by: `TI2VidOneStagePipeline`, `TI2VidTwoStagesPipeline`, `DistilledPipeline`, `ICLoraPipeline`
-  - Replaces the latent at a specific frame with the encoded image
-  - Strong control over specific frames
-
-- **Guiding Latents** ([`image_conditionings_by_adding_guiding_latent`](src/ltx_pipelines/utils/helpers.py)):
-  - Used by: `KeyframeInterpolationPipeline`
-  - Adds the image as a guiding signal rather than replacing
-  - Better for smooth interpolation between keyframes
-
-### Video Conditioning
-
-- **Video Conditioning** (ICLoraPipeline only):
-  - Conditions on entire reference videos
-  - Useful for video-to-video transformations
-  - Uses `VideoConditionByKeyframeIndex` from [`ltx-core`](../ltx-core/)
-
----
-
-## 🎛️ Multimodal Guidance
-
-LTX-2 pipelines use **multimodal guidance** to steer the diffusion process for both video and audio modalities. Each modality (video, audio) has its own guider with independent parameters, allowing fine-grained control over generation quality and adherence to prompts.
-
-### Guidance Parameters
-
-The `MultiModalGuiderParams` dataclass controls guidance behavior:
-
-| Parameter | Description |
-| --------- | ----------- |
-| `cfg_scale` | **Classifier-Free Guidance** scale. Higher values make the output adhere more strongly to the text prompt. Typical values: 2.0–5.0. Set to **1.0** to disable. |
-| `stg_scale` | **Spatio-Temporal Guidance** scale. Controls perturbation-based guidance for improved temporal coherence. Typical values: 0.5–1.5. Set to **0.0** to disable. |
-| `stg_blocks` | Which transformer blocks to perturb for STG (e.g., `[29]` for the last block). Set to **`[]`** to disable STG. |
-| `rescale_scale` | Rescales the guided prediction to match the variance of the conditional prediction. Helps prevent over-saturation. Typical values: 0.5–0.7. Set to **0.0** to disable. |
-| `modality_scale` | **Modality CFG** scale. Steers the model away from unsynced video and audio results, improving audio-visual coherence. Set to **1.0** to disable. |
-| `skip_step` | Skip guidance every N steps. Can speed up inference with minimal quality loss. Set to **0** to disable (never skip). |
-
-### How It Works
-
-The multimodal guider combines three guidance signals during each denoising step:
-
-1. **CFG (Text Guidance)**: Steers generation toward the text prompt by computing `(cond - uncond_text)`.
-2. **STG (Perturbation Guidance)**: Improves structural coherence by perturbing specific transformer blocks and steering away from the perturbed prediction.
-3. **Modality CFG**: For joint audio-video generation, steers the model away from unsynced video and audio results.
-
-### Example Configuration
-
-```python
-from ltx_core.components.guiders import MultiModalGuiderParams
-
-# Video guider: moderate CFG, STG enabled, modality isolation
-video_guider_params = MultiModalGuiderParams(
-    cfg_scale=3.0,
-    stg_scale=1.0,
-    rescale_scale=0.7,
-    modality_scale=3.0,
-    stg_blocks=[29],
-)
-
-# Audio guider: higher CFG for stronger prompt adherence
-audio_guider_params = MultiModalGuiderParams(
-    cfg_scale=7.0,
-    stg_scale=1.0,
-    rescale_scale=0.7,
-    modality_scale=3.0,
-    stg_blocks=[29],
-)
-```
-
-> **Tip:** Start with the default values from [`constants.py`](src/ltx_pipelines/utils/constants.py) and adjust based on your use case. Higher `cfg_scale` = stronger prompt adherence but potentially less natural motion; higher `stg_scale` = better temporal coherence but slower inference (requires extra forward passes).
->
-> **Tip:** When generating video with audio, set `modality_scale` > 1.0 (e.g., 3.0) to improve audio-visual sync. If generating video-only, set it to 1.0 to disable.
-
----
-
-## ⚡ Optimization Tips
-
-
-### Memory Optimization
-
-**FP8 Quantization (Lower Memory Footprint):**
-
-For smaller GPU memory footprint, use the `--quantization` flag and set `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True`.
-
-Two quantization policies are available:
-
-| Policy | CLI Flag | Description |
-| ------ | -------- | ----------- |
-| **FP8 Cast** | `--quantization fp8-cast` | Downcasts transformer linear weights to FP8 during loading; upcasts on the fly during inference. No extra dependencies. |
-| **FP8 Scaled MM** | `--quantization fp8-scaled-mm` | Uses FP8 scaled matrix multiplication via TensorRT-LLM (`tensorrt_llm` must be installed). Best performance on Hopper GPUs. |
-
-**CLI:**
-
-```bash
-# FP8 Cast (works on any GPU with FP8 support)
-PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True python -m ltx_pipelines.ti2vid_two_stages \
-    --quantization fp8-cast --checkpoint-path=...
-
-# FP8 Scaled MM (requires tensorrt_llm, best on Hopper GPUs)
-PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True python -m ltx_pipelines.ti2vid_two_stages \
-    --quantization fp8-scaled-mm --checkpoint-path=...
-```
-
-**Programmatically:**
-
-When authoring custom scripts, pass a `QuantizationPolicy` to pipeline classes:
-
-```python
-from ltx_core.quantization.fp8_cast import build_policy as build_fp8_cast_policy
-# Alternative:
-# from ltx_core.quantization.fp8_scaled_mm import build_policy as build_fp8_scaled_mm_policy
-
-pipeline = TI2VidTwoStagesPipeline(
-    checkpoint_path=ltx_model_path,
-    distilled_lora=distilled_lora,
-    spatial_upsampler_path=upsampler_path,
-    gemma_root=gemma_root_path,
-    loras=[],
-    quantization=build_fp8_cast_policy(ltx_model_path),
-)
-pipeline(...)
-```
-
-You still need to use `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True` when launching:
-
-```bash
-PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True python my_denoising_pipeline.py
-```
-
-**Memory Cleanup Between Stages:**
-
-By default, pipelines clean GPU memory (especially transformer weights) between stages. If you have enough memory, you can skip this cleanup to reduce running time:
-
-```python
-# In pipeline implementations, memory cleanup happens automatically
-# between stages. For custom pipelines, you can skip:
-# utils.cleanup_memory()  # Comment out if you have enough VRAM
-```
-
-### Compilation (`torch.compile`)
-
-Compiling the transformer blocks with `torch.compile` speeds up inference. It is **opt-in and off by default**. The blocks are compiled shape-polymorphically (the sequence dimension is marked dynamic), so one compiled artifact serves any token count without recompiling.
-
-**CLI** — the `--compile` flag maps directly to `CompilationConfig`:
-
-| Form | Result |
-| ---- | ------ |
-| *(flag absent)* | eager, no compilation |
-| `--compile` | compile with defaults |
-| `--compile KEY=VALUE ...` | compile, overriding individual fields |
-
-```bash
-# Defaults
-python -m ltx_pipelines.ti2vid_two_stages --compile --checkpoint-path=...
-
-# reduce-overhead captures CUDA graphs -- the main latency lever for the denoising loop.
-# Off by default because graph capture reserves static memory pools (extra VRAM), so it
-# trades memory for speed; enable it when you have headroom.
-python -m ltx_pipelines.ti2vid_two_stages --compile mode=reduce-overhead --checkpoint-path=...
-
-# Several overrides at once
-python -m ltx_pipelines.ti2vid_two_stages \
-    --compile mode=max-autotune fullgraph=true dynamic=true --checkpoint-path=...
-```
-
-| Field | Values | Default | Notes |
-| ----- | ------ | ------- | ----- |
-| `mode` | `none`, `reduce-overhead`, `max-autotune`, … | `none` | `reduce-overhead`/`max-autotune` enable CUDA graphs |
-| `backend` | `inductor`, `eager`, … | `inductor` | |
-| `fullgraph` | `true`/`false` | `false` | |
-| `dynamic` | `auto`/`true`/`false` | `auto` | the seq dim is marked dynamic regardless |
-| `inductor_config` | JSON object or path to a `.json` | `{}` | `torch._inductor.config` overrides |
-| `dynamo_config` | JSON object or path to a `.json` | `{"inline_inbuilt_nn_modules": true, "cache_size_limit": 256}` | `torch._dynamo.config` overrides |
-
-**Controlling inductor / dynamo configs.** `inductor_config` and `dynamo_config` take either an inline JSON object or a path to a `.json` file, applied via `torch._inductor.config.patch(...)` / `torch._dynamo.config.patch(...)` around the compiled forward. They **replace the defaults wholesale — they do not merge**, so when overriding `dynamo_config` re-include any defaults you want to keep:
-
-```bash
-python -m ltx_pipelines.ti2vid_two_stages \
-    --compile 'inductor_config={"max_autotune": true}' \
-              'dynamo_config={"inline_inbuilt_nn_modules": true, "cache_size_limit": 256, "recompile_limit": 32}' \
-    --checkpoint-path=...
-```
-
-**Programmatically**, pass a `CompilationConfig` to the pipeline:
-
-```python
-from ltx_core.model.transformer.compiling import CompilationConfig
-
-pipeline = TI2VidTwoStagesPipeline(
-    ...,
-    compilation_config=CompilationConfig(mode="reduce-overhead"),
-)
-```
-
-**Faster cache loads: `unsafe_skip_cache_dynamic_shape_guards` (unsafe, opt-in).** Inductor's FX-graph cache re-checks the dynamic-shape guards stored with each entry on every lookup. Setting this flag skips that re-check (every entry is treated as a guard hit), which speeds up warm and cross-process cache loads. It is **not enabled by default** because it is a correctness hazard: a kernel first compiled at a small sequence length keeps int32 address arithmetic, and reusing it at a larger sequence length (roughly **>58k tokens/rank**) overflows int32 and reads out of bounds — surfacing as a CUDA illegal memory access or silently corrupted output. Only enable it when your token counts stay within the range the cached kernels were compiled for:
-
-```bash
-python -m ltx_pipelines.ti2vid_two_stages \
-    --compile 'inductor_config={"unsafe_skip_cache_dynamic_shape_guards": true}' \
-    --checkpoint-path=...
-```
-
-### Denoising Loop Optimization
-
-**Gradient Estimation Denoising Loop:**
-
-Instead of the standard Euler denoising loop, you can use gradient estimation for fewer steps (~20-30 instead of 40):
-
-```python
-from ltx_pipelines.utils import gradient_estimating_euler_denoising_loop
-
-# Use gradient estimation denoising loop
-def denoising_loop(sigmas, video_state, audio_state, stepper):
-    return gradient_estimating_euler_denoising_loop(
-        sigmas=sigmas,
-        video_state=video_state,
-        audio_state=audio_state,
-        stepper=stepper,
-        transformer=transformer,
-        denoiser=denoiser,
-        ge_gamma=2.0,  # Gradient estimation coefficient
-    )
-```
-
-This allows you to use **20-30 steps instead of 40** while maintaining quality. The gradient estimation function is defined in [`samplers.py`](src/ltx_pipelines/utils/samplers.py).
-
----
-
-## 🔧 Requirements
-
-- **LTX-2 Model Checkpoint** - Local `.safetensors` file
-- **Gemma Text Encoder** - Local Gemma model directory
-- **Spatial Upscaler** - Required for two-stage pipelines (except one-stage)
-- **Distilled LoRA** - Required for two-stage pipelines (except one-stage and distilled)
-
----
-
-## 📖 Example: Image-to-Video
-
-```python
-from ltx_core.components.guiders import MultiModalGuiderParams
-from ltx_core.loader import LTXV_LORA_COMFY_RENAMING_MAP, LoraPathStrengthAndSDOps
-from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
-from ltx_pipelines.ti2vid_two_stages import TI2VidTwoStagesPipeline
-from ltx_pipelines.utils.args import ImageConditioningInput
-from ltx_pipelines.utils.media_io import encode_video
-
-distilled_lora = [
-    LoraPathStrengthAndSDOps(
-        "/path/to/distilled_lora.safetensors",
-        0.6,
-        LTXV_LORA_COMFY_RENAMING_MAP,
-    ),
-]
-
-pipeline = TI2VidTwoStagesPipeline(
-    checkpoint_path="/path/to/checkpoint.safetensors",
-    distilled_lora=distilled_lora,
-    spatial_upsampler_path="/path/to/upsampler.safetensors",
-    gemma_root="/path/to/gemma",
-    loras=[],
-)
-
-video_guider_params = MultiModalGuiderParams(
-    cfg_scale=3.0,
-    stg_scale=1.0,
-    rescale_scale=0.7,
-    modality_scale=3.0,
-    skip_step=0,
-    stg_blocks=[29],
-)
-
-audio_guider_params = MultiModalGuiderParams(
-    cfg_scale=7.0,
-    stg_scale=1.0,
-    rescale_scale=0.7,
-    modality_scale=3.0,
-    skip_step=0,
-    stg_blocks=[29],
-)
-
-# Generate video from image. The pipeline returns (video_iterator, audio);
-# the caller is responsible for encoding to file via encode_video().
-num_frames = 121
-frame_rate = 25.0
-tiling_config = TilingConfig.default()
-video, audio = pipeline(
-    prompt="A serene landscape with mountains in the background",
-    negative_prompt="worst quality, low quality, blurry, distorted",
-    seed=42,
-    height=512,
-    width=768,
-    num_frames=num_frames,
-    frame_rate=frame_rate,
-    num_inference_steps=40,
-    video_guider_params=video_guider_params,
-    audio_guider_params=audio_guider_params,
-    images=[ImageConditioningInput("input_image.jpg", 0, 1.0, 33)],  # path, frame_idx=0, strength=1.0, crf=33
-    tiling_config=tiling_config,
-)
-encode_video(
-    video=video,
-    fps=frame_rate,
-    audio=audio,
-    output_path="output.mp4",
-    video_chunks_number=get_video_chunks_number(num_frames, tiling_config),
-)
-```
-
----
+See [Installation & Usage](docs/installation.md) for full setup, CLI modules, and shared flags.
+
+## 📚 Documentation
+
+| Topic | Description |
+| ----- | ----------- |
+| [Installation & Usage](docs/installation.md) | Install, requirements, running pipelines from the CLI, common flags |
+| [Pipeline Selection Guide](docs/pipeline-selection.md) | Decision tree + feature comparison to pick the right pipeline |
+| [Available Pipelines](docs/pipelines.md) | Full reference for all 11 pipelines |
+| [Conditioning Types](docs/conditioning.md) | Image and video conditioning methods |
+| [Multimodal Guidance](docs/multimodal-guidance.md) | CFG / STG / modality guidance parameters and tuning |
+| [Optimization Tips](docs/optimization.md) | FP8 quantization, `torch.compile`, gradient estimation |
+| [Multi-GPU Inference](docs/multigpu/README.md) | Run a single generation across GPUs for latency (SP, TDP, distributed VAE, distributed Gemma) |
 
 ## 🔗 Related Projects