Add char_masks preprocessing script for SCAIL-2 training (task 3.7)

scripts/process_char_masks.py turns per-sample character label-map videos/images
(integer pixel labels: 0 = environment, 1..K = characters -> binding slots) into
the pixel-space semantic-mask tensors the SCAIL training/inference path consumes:
{"mask": [K+1, F_pix, H_pix, W_pix]} (ch0 = environment switch, ch1..K = slots).

- Aligns to the target video's latent grid read from the saved latent metadata
  (F_pix=(F-1)*8+1, H*32, W*32), so char_masks/ lines up file-for-file with
  latents/ / driving_latents/ for PrecomputedDataset.
- Nearest-neighbour resize so integer labels are never blended; labels > K are
  dropped with a warning; ch0 filled uniformly with --environment-switch.
- Reuses process_videos.py helpers (naming, atomic save, VAE factors) and matches
  its typer CLI conventions.

Verified on CPU: a synthetic 2-character label map (plus an out-of-range id)
produces mask (7,17,128,128) with ch0 uniform, slots placed correctly, id>K
dropped, and feeds encode_mask_channels to the 8*(K+1)=56 channels. README +
docs/tasks.md 3.7 updated (upstream label-map generation via SAM/tracking is
dataset-specific and still out of scope).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

.gitignore

@@ -44,12 +44,18 @@ wandb/
 *.wav
 *.webp
 
-# HDR IC-LoRA e2e test baseline (checked in via Git LFS)
+# Full-params expected results for the --full-e2e quality lane (checked in via Git LFS)
-!packages/ltx-pipelines/tests/assets/expected_hdr_ic_lora_exr/frame_*.exr
+!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)
+# Fast integration-profile goldens (bit-exact decoded video/audio; checked in via Git LFS)
-!packages/ltx-pipelines/tests/assets/expected_t2a_one_stage_ltx2_3.wav
+!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.
+* **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 other CUDA GPUs (including Hopper), use xFormers (`uv sync --extra xformers`).
+* **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/README.md#denoising-loop-optimization))
+* **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
 

docs/plan.md

@@ -0,0 +1,70 @@
+# SCAIL-2 → LTX-2 移植計畫
+
+將 SCAIL-2（arXiv:2606.10804，`zai-org/SCAIL-2`）的端到端角色動畫手法移植到 LTX-2。
+
+> **注意**：SCAIL-2 原始實作建構於 **Wan 2.1**，非 LTX-2。座標慣例與架構需翻譯到 LTX 的資料流。
+
+## 三個核心機制
+
+1. **Driving latent 直接串接** — 把驅動影片 latent 直接接進 DiT token 序列（不經骨架/pose 中介），用 width 軸座標偏移 ΔW 讓 driving 座標跟主 video 座標分開。
+2. **In-context mask channel** — 疊加額外輸入 channel（1 個環境開關 + K=6 個角色綁定槽，展開為 `4(K+1)=28` channel）到模型輸入，讓模型知道背景/角色對應。
+3. **Mode-specific RoPE** — Animation Mode 與 Replacement Mode 用不同的座標指派規則。
+
+## LTX-2 對應落點
+
+| 機制 | LTX-2 落點 | 改權重？ |
+|---|---|---|
+| 1. Driving 串接 + ΔW | 新 `ConditioningItem`（clone `reference_video_cond.py`），改 `positions` 偏移 | 否 |
+| 3. Mode-specific RoPE | 上述 item 加 `DrivingMode` enum | 否 |
+| 2. In-context mask channel | 加寬 `patchify_proj` 輸入 channel + `LatentState` 帶額外 channel + 投影前 concat | **是**（需微調） |
+
+**關鍵洞察**：LTX 的 `ConditioningItem.apply_to(latent_state, latent_tools) -> LatentState` 就是「串接進序列」的天然注入點，機制 1+3 **完全不用動 transformer / rope.py**。RoPE 由 `LatentState.positions` `[B,3,T,2]`（pixel 座標，axis1=(time,h,w)）驅動；ΔW = 對 `positions[:,2]`（width）加常數。width 正規化上界 `max_pos[2]=2048`，超過會 wrap。
+
+## SCAIL-2 座標規則（論文）
+
+序列 `[z_ref; z_t; z_driv]`，driving 永遠在 width 軸帶固定偏移 ΔW：
+
+| | Animation | Replacement |
+|---|---|---|
+| z_ref | T=0, H=[0,Hv), W=[0,Wv) | T=0, **H=[ΔH_ref, ΔH_ref+Hv)**, W=[0,Wv) |
+| z_t | T=[1,Tv], H=[0,Hv), W=[0,Wv) | T=[0,Tv−1], H=[0,Hv), W=[0,Wv) |
+| z_driv | T=[1,Tv], H=[0,Hv), **W=[ΔW, ΔW+Wv)** | T=[0,Tv−1], H=[0,Hv), **W=[ΔW, ΔW+Wv)** |
+
+## 分階段計畫
+
+### Phase 1 — Driving 串接 item（機制 1+3，推論期 PoC）✅ 已完成
+- 純推論、不改權重、不動 `patchify_proj`、不碰 ltx-trainer。
+- 用現有 checkpoint 驗證資料流正確性（序列長度、座標偏移、attention mask、denoise_mask、不 wrap）。
+- **PoC 僅驗證 plumbing**；LTX-2 未經此訓練，畫面不會是正確動畫。
+
+### Phase 2 — In-context mask channel（機制 2，checkpoint 手術）✅ 已完成（plumbing）
+- **決策**：時間編碼採忠實堆疊，`8×(K+1)=56` channel（LTX 時間因子 8，K=6），非 Wan 的 28。
+- `LTXModel(mask_conditioning_channels=0)` config-gated 加寬 `patchify_proj`；預設不變。
+- `LatentState`/`Modality` 加 `cond_channels` 欄位，跟著 clone/clear/append 流動；投影前在 `_apply_patchify_proj` concat（None 補零）。
+- 新輸入欄位 **zero-init**：`widen_patchify_proj_for_mask_channels` 轉換舊 checkpoint，行為不變、待微調才生效。
+- `VideoConditionByMaskChannels`：(K+1) 語意 mask → 空間下採樣 + 時間 8× 堆疊 → 寫入尾端 driving token（target 保持零，符合論文）。
+- 驗證通過（`verify_mask_channels.py`）：向後相容、zero-init 加寬 forward == baseline、mask pipeline 不 crash。**僅驗證 plumbing，畫質需 Phase 3 微調。**
+
+### Phase 3 — 訓練整合（ltx-trainer）✅ 已完成（程式碼路徑；實訓需 GPU）
+- **決策**：完整整合到 `FlexibleStrategy`；訓練 = LoRA + 解凍 `patchify_proj`（新 mask 欄位無法純 LoRA 訓練）。
+- 新 `DrivingConditionConfig` + `MaskChannelsConditionConfig`（`flexible.py`）；`_apply_driving_condition`(cond-first concat + ΔW) + `_build_mask_channels`(重用 `encode_mask_channels`) → `Modality.cond_channels`。
+- `load_transformer(mask_conditioning_channels=)` 用 `widen_module_patchify_proj_for_mask_channels` 加寬；`ModelConfig.mask_conditioning_channels`；trainer 在 LoRA 模式解凍 patchify_proj。
+- `configs/scail_animation_lora.yaml` + docs。CPU 單元驗證通過（`verify_phase3_trainer.py`）。
+- **本機無 GPU/Linux/checkpoint → 未跑實機訓練**；dataset 前處理（driving latents + 語意 mask）與 validation runner 接線未做。
+
+### Phase 4 — Pipeline + CLI 包裝 ✅ 已完成（程式碼路徑；實跑需 GPU）
+- `LTXModelConfigurator` 讀 `mask_conditioning_channels` → SCAIL checkpoint 自描述、載入自動加寬（不需 runtime wrapper）。
+- `scail_animation.py`：`ScailAnimationPipeline`（仿 `distilled.py` 兩階段）+ 可測純函式 `build_scail_conditionings`（driving + mask 條件組裝）+ `load_masks` + `main()`。
+- `utils/args.py` `scail_animation_arg_parser`（`--driving-video/--mask-path/--mode/--driving-strength`）。
+- CPU 驗證通過（`verify_phase4_pipeline.py`）；`ltx-pipelines/CLAUDE.md` 表格 row。
+- **實機端到端需 GPU + SCAIL-trained checkpoint**（config 含 `mask_conditioning_channels` + 訓練好的 patchify_proj + LoRA）。
+
+## Phase 1 簡化取捨（記錄，Phase 2 需回頭處理）
+- (a) driving 時間座標直接複製 target 的（token-wise），故 driving 需與 target 同 F/H/W。
+- (b) 單一 frozen driving group 的 `attention_mask` 維持 None（= 全連接，target 完全看得到 driving），與 reference cond 一致。
+- (c) **ANIMATION 與 REPLACEMENT 在 Phase 1 產生相同 driving 座標** — mode 差異（z_ref 的 ΔH_ref 高度位移、target 時間原點、mask channel）屬 Phase 2，enum 先保留佔位。
+
+## 參考
+- 論文：arXiv:2606.10804 — *SCAIL-2: Unifying Controlled Character Animation with End-to-end In-Context Conditioning*
+- 官方實作：`zai-org/SCAIL-2`（GitHub / HuggingFace），建構於 Wan 2.1
+- 藍本檔案：`packages/ltx-core/src/ltx_core/conditioning/types/reference_video_cond.py`

docs/tasks.md

@@ -0,0 +1,59 @@
+# SCAIL-2 → LTX-2 移植任務追蹤
+
+狀態圖例：✅ 完成 ｜ 🔄 進行中 ｜ ⬜ 未開始 ｜ ⏸️ 暫緩
+
+相關計畫見 [`plan.md`](./plan.md)。
+
+## Phase 1 — Driving 串接 item（推論期 PoC）
+
+| # | 任務 | 狀態 | 產出 / 備註 |
+|---|---|---|---|
+| 1.1 | 新增 `VideoConditionByDrivingLatent` + `DrivingMode` | ✅ | `packages/ltx-core/src/ltx_core/conditioning/types/driving_video_cond.py`，以 `reference_video_cond.py` 為藍本，實作 ΔW width 偏移、時間對齊 target、`max_pos` 防呆、token 數防呆 |
+| 1.2 | 匯出新 conditioning 類別 | ✅ | `conditioning/types/__init__.py`、`conditioning/__init__.py` |
+| 1.3 | 免-GPU 資料流驗證腳本 | ✅ | 序列長度、ΔW 不重疊、時間對齊、frozen denoise_mask、attention_mask=None、`clear_conditioning` 剝除、超界/token 數防呆 — 全通過；ruff clean |
+| 1.4 | （選配）端到端 smoke run | ✅ | 本機 CPU-only、無 checkpoint，改用小型真實 `LTXModel`（隨機權重）跑 pipeline 實走路徑：`create_noised_state`（含 driving cond）→ `modality_from_latent_state` → **真 transformer forward**（seq 160=target 80+driving 80）→ `clear_conditioning`（→80）。ANIMATION/REPLACEMENT 皆通過：不 crash、輸出 finite、driving frozen 且正確剝除。**僅驗證整合，不評估畫質** |
+
+## Phase 2 — In-context mask channel（checkpoint 手術）
+
+> **決策**：mask channel 時間編碼採**忠實堆疊**——每 latent 幀對應 8 個 pixel 子幀沿 channel 堆疊，`8×(K+1)=56` channel（LTX 時間因子 8，K=6）。非 Wan 的 4×(K+1)=28。
+
+| # | 任務 | 狀態 | 備註 |
+|---|---|---|---|
+| 2.1 | `patchify_proj` 加寬（config-gated） | ✅ | `LTXModel(mask_conditioning_channels=0)`；>0 時 `patchify_proj=Linear(in+mask, inner)`。預設不變 |
+| 2.2 | `LatentState`/`Modality` 帶 `cond_channels` 欄位 | ✅ | `types.py`、`modality.py` 加 `cond_channels: Tensor\|None=None`（patchified [B,T,C]）；`tools.clear_conditioning` 裁切；`helpers.modality_from_latent_state` 帶入 |
+| 2.3 | 投影前 concat/zero-pad cond_channels | ✅ | `transformer_args.py` `_apply_patchify_proj`：寬度不足時用 cond_channels 補齊，None 則補零 |
+| 2.4 | checkpoint zero-init 加寬轉換 | ✅ | `mask_channels_checkpoint.py` `widen_patchify_proj_for_mask_channels`：尾端補零欄，載入舊權重行為不變 |
+| 2.5 | `VideoConditionByMaskChannels`（56ch） | ✅ | `mask_channels_cond.py`：1 環境開關 + K=6 綁定槽 → 空間下採樣 + 時間 8× 堆疊(causal 首幀複製)=56ch，寫入尾端 driving token，target 保持零 |
+| 2.6 | token-append conditioning 延伸 cond_channels | ✅ | `cond_channels.py` `extend_cond_channels`；reference_video/reference_audio/driving/keyframe 皆接上 |
+| 2.7 | Phase 2 驗證（免訓練） | ✅ | `verify_mask_channels.py`：mask=0 向後相容；zero-init 加寬 forward == baseline（任意 cond_channels）；driving+mask pipeline forward 不 crash、cond_channels 形狀/placement 正確、`clear_conditioning` 剝除 |
+| 2.8 | Replacement Mode z_ref 高度位移 ΔH_ref | ⬜ | 仍暫緩（Phase 1 divergence，需獨立 reference token group） |
+
+## Phase 3 — 訓練整合（ltx-trainer）
+
+> **決策**：完整整合到 `FlexibleStrategy`；訓練方式 = **LoRA + 解凍 patchify_proj**（新 mask 欄位無法純 LoRA 訓練）。**本機無 GPU/Linux/checkpoint，只做到 CPU 單元驗證**，實訓需在 GPU 機器跑。
+
+| # | 任務 | 狀態 | 備註 |
+|---|---|---|---|
+| 3.1 | 新增 Driving/MaskChannels ConditionConfig | ✅ | `flexible.py`：`DrivingConditionConfig`(latents_dir/mode/width_offset) + `MaskChannelsConditionConfig`(mask_dir/num_slots)，加入 union + `get_data_sources` |
+| 3.2 | strategy 接線 driving concat + cond_channels | ✅ | `_apply_driving_condition`(cond-first concat + ΔW width 偏移) + `_build_mask_channels`(重用 `encode_mask_channels`，寫前 N driving token) → `Modality.cond_channels` |
+| 3.3 | model_loader + ModelConfig 支援加寬 | ✅ | `widen_module_patchify_proj_for_mask_channels`(live module zero-init) + `load_transformer(mask_conditioning_channels=)` + `ModelConfig.mask_conditioning_channels` |
+| 3.4 | LoRA 模式解凍 patchify_proj | ✅ | `trainer._unfreeze_patchify_proj`：mask_channels>0 時把 video patchify_proj 設 trainable，讓新欄位隨 LoRA 一起訓 |
+| 3.5 | 範例 config + docs | ✅ | `configs/scail_animation_lora.yaml`；`configs/README.md` 與 `docs/training-modes.md` 表格 row |
+| 3.6 | CPU 單元驗證 | ✅ | `verify_phase3_trainer.py`：config round-trip、prepare_training_inputs 建 cond_channels[B,T,56]、driving 前置/mask placement、widened model forward + compute_loss finite、widen helper zero-init 等價 |
+| 3.7 | dataset 前處理（產 driving latents + 語意 mask） | 🟡 | `char_masks/` 前處理已做：`scripts/process_char_masks.py`（label-map 影片/圖 → `[K+1,F_pix,H,W]`，對齊 target latent，nearest 保留整數 label，ch0 環境開關）。`driving_latents/` 沿用既有 `process_videos.py`（driving 影片走 video latent 路徑，同 target 形狀）。**仍缺**：從原始影片產生 label-map 的分割/追蹤步驟（SAM 等，資料集特定，未含） |
+| 3.8 | validation runner 接 driving/mask | ⬜ | 驗證期取樣尚未接 SCAIL 條件（config 內 validation 先停用），與 Phase 4 一起 |
+| 3.9 | 實機訓練跑通 | ⬜ | 需 Linux + GPU + checkpoint，本機無法 |
+
+## Phase 4 — Pipeline + CLI
+
+| # | 任務 | 狀態 | 備註 |
+|---|---|---|---|
+| 4.1 | configurator 讀 `mask_conditioning_channels` | ✅ | `LTXModelConfigurator` 兩分支 `config.get("mask_conditioning_channels",0)` → SCAIL checkpoint 自描述、載入自動加寬（不需 runtime widen wrapper） |
+| 4.2 | `ScailAnimationPipeline` + `build_scail_conditionings` | ✅ | `scail_animation.py`：仿 `distilled.py` 兩階段，注入 SCAIL driving+mask conditioning；`build_scail_conditionings` 為可測純函式；`load_masks` helper |
+| 4.3 | `scail_animation_arg_parser` + `main()` | ✅ | `utils/args.py`：在 `default_2_stage_distilled_arg_parser` 上加 `--driving-video/--mask-path/--mode/--driving-strength` |
+| 4.4 | CPU 驗證 + docs | ✅ | `verify_phase4_pipeline.py`：import、CLI round-trip、conditioning assembly(driving append + mask 在尾端)、configurator 自描述加寬；`ltx-pipelines/CLAUDE.md` pipeline 表格 row |
+| 4.5 | 實機端到端跑通 | ⬜ | 需 GPU + SCAIL-trained checkpoint（含 `mask_conditioning_channels` config + 訓練好的 patchify_proj + LoRA） |
+
+## 決議紀錄
+- **範圍**：先只做 Phase 1（推論期 PoC）。Phase 2+ 待 Phase 1 驗證後再討論。
+- **架構前提**：SCAIL-2 建構於 Wan 2.1，本移植為跨架構移植。

modal/README.md

@@ -0,0 +1,111 @@
+# SCAIL-2 on Modal — free testing environment
+
+Test the SCAIL-2 LTX-2 training/inference **code path** on Modal without a local
+GPU. The free path uses tiny random-init models + synthetic data (no checkpoint,
+no dataset), so it validates that the SCAIL wiring runs on real Linux/GPU — not
+model quality.
+
+> ⚠️ The real 19B LTX-2 model is **not** free to train/run. The `verify`/`smoke`
+> targets are near-free; `train` on a real checkpoint uses a paid GPU.
+
+## 1. One-time setup
+
+1. Sign up at [modal.com](https://modal.com) (the free Starter plan includes a
+   monthly credit allowance — enough for many `verify`/`smoke` runs).
+2. Install and authenticate:
+   ```bash
+   pip install modal
+   modal setup          # opens a browser to link your account / token
+   ```
+
+## 2. Free / near-free checks
+
+From the repo root:
+
+```bash
+# CPU: SCAIL-2 Phase 1-4 plumbing (driving concat, mask channels, zero-init widen,
+# a FlexibleStrategy training step + loss, inference conditioning assembly).
+modal run modal/app.py::verify
+
+# Same checks on a T4 GPU (validates the CUDA path). ~cents.
+modal run modal/app.py::smoke
+```
+
+The first run builds the image (installs the `ltx-core`/`ltx-pipelines`/`ltx-trainer`
+workspace via `uv sync`). `ltx-kernels` (CUDA-compiled) is intentionally skipped —
+attention falls back to PyTorch SDPA, so no CUDA toolchain is required.
+
+Expected tail:
+```
+[OK] Phase 1 driving concat: seq 48 -> 96
+[OK] Phase 2 zero-init widen: in_features=72, output preserved
+[OK] Phase 3 training step: cond_channels (1, 96, 56), loss ...
+[OK] Phase 4 build_scail_conditionings: [driving, mask_channels]
+All SCAIL-2 checks passed on cuda   # (or cpu)
+```
+
+## 3. Scaling up to real training (paid)
+
+`train` runs `packages/ltx-trainer/scripts/train.py` against a real checkpoint.
+You must supply the weights + data via the persistent `scail-data` Volume.
+
+1. Create/populate the Volume (checkpoint, Gemma encoder, preprocessed latents):
+   ```bash
+   modal volume create scail-data          # if not auto-created
+   modal volume put scail-data /local/ltx-2-model.safetensors /model/ltx-2.safetensors
+   modal volume put scail-data /local/gemma /model/gemma
+   modal volume put scail-data /local/preprocessed /data/preprocessed
+   ```
+2. Copy `configs/scail_animation_lora.yaml`, and point its paths at the mounted
+   Volume (everything lands under `/data` in the container):
+   ```yaml
+   model:
+     model_path: "/data/model/ltx-2.safetensors"
+     text_encoder_path: "/data/model/gemma"
+     mask_conditioning_channels: 56       # widens patchify_proj at load
+   data:
+     preprocessed_data_root: "/data/preprocessed"
+   ```
+   (Put your edited config on the Volume too, or bake it into the repo.)
+3. Launch (pick a GPU big enough for the model — the 19B needs an A100):
+   ```bash
+   # edit gpu="A10G" -> "A100" in modal/app.py::train for the full model
+   modal run modal/app.py::train --config-rel /data/scail_animation_lora.yaml
+   ```
+
+### Dataset preprocessing (not yet automated for SCAIL)
+
+The SCAIL training config expects, under `preprocessed_data_root/`:
+
+```
+latents/           # target video latents
+conditions/        # text embeddings
+driving_latents/   # driving video latents (same F/H/W as target)
+char_masks/        # per-sample "mask" = [K+1, F_pix, H, W]  (ch0 env switch, 1..K binding slots)
+```
+
+`latents/`, `conditions/`, and `driving_latents/` come from the existing
+`packages/ltx-trainer/scripts/process_dataset.py` (run it once per video set).
+`char_masks/` is produced by `packages/ltx-trainer/scripts/process_char_masks.py`
+from per-sample **label-map** videos/images (integer pixel labels: `0` =
+environment, `1..K` = characters → binding slots):
+
+```bash
+python packages/ltx-trainer/scripts/process_char_masks.py dataset.csv \
+    --mask-column char_labels --latents-dir ./latents \
+    --output-dir ./char_masks --num-slots 6 --main-media-column media_path
+```
+
+You still need a segmentation/tracking model (e.g. SAM) to *produce* those label
+maps from raw video — that upstream step is dataset-specific and not included.
+
+## Cost notes
+
+| Target  | GPU   | Rough cost | Use |
+|---------|-------|-----------|-----|
+| `verify`| none  | ~free      | validate code path on CPU |
+| `smoke` | T4    | cents      | validate CUDA path |
+| `train` | A10G/A100 | paid   | real fine-tuning (needs checkpoint + data) |
+
+Free credits are best spent on `verify`/`smoke` to catch integration issues before
+committing a paid GPU to a real run. Watch usage in the Modal dashboard.

modal/app.py

@@ -0,0 +1,84 @@
+"""Modal app for testing the SCAIL-2 LTX-2 training/inference path.
+
+Free-tier friendly: `verify` runs on CPU and `smoke` on a cheap T4, both using
+tiny random-init models + synthetic data (no checkpoint, no dataset) so a full run
+costs pennies. `train` is the scale-up entrypoint that runs the real trainer once
+you upload a checkpoint + preprocessed data to the `scail-data` Volume -- that one
+uses a real GPU and is NOT free.
+
+Setup (once):
+    pip install modal && modal setup
+
+Run:
+    modal run modal/app.py::verify     # CPU plumbing check (~free)
+    modal run modal/app.py::smoke      # same checks on a T4 GPU (cents)
+    modal run modal/app.py::train --config-rel configs/scail_animation_lora.yaml
+"""
+
+import subprocess
+from pathlib import Path
+
+import modal
+
+REPO = Path(__file__).parent.parent
+
+# Build the workspace once into the image. ltx-kernels (CUDA-compiled) is excluded
+# from the workspace and not needed -- attention falls back to SDPA.
+image = (
+    modal.Image.debian_slim(python_version="3.11")
+    .apt_install("git")
+    .pip_install("uv")
+    .env({"UV_LINK_MODE": "copy", "UV_PROJECT_ENVIRONMENT": "/root/LTX-2/.venv"})
+    .add_local_dir(
+        str(REPO),
+        "/root/LTX-2",
+        copy=True,
+        ignore=[".git", ".venv", "**/__pycache__", "**/*.pyc", "outputs", "wandb", "**/.pytest_cache"],
+    )
+    .run_commands("cd /root/LTX-2 && uv sync --package ltx-trainer")
+)
+
+app = modal.App("scail-ltx2", image=image)
+
+# Persistent volume for the (large) checkpoint + preprocessed dataset used by `train`.
+data_volume = modal.Volume.from_name("scail-data", create_if_missing=True)
+
+_UV_PY = ["uv", "run", "--package", "ltx-trainer", "python"]
+
+
+def _run(args: list[str]) -> None:
+    subprocess.run([*_UV_PY, *args], cwd="/root/LTX-2", check=True)
+
+
+@app.function(timeout=1800)
+def verify() -> None:
+    """Run the SCAIL-2 plumbing checks on CPU (near-free)."""
+    _run(["modal/checks.py"])
+
+
+@app.function(gpu="T4", timeout=1800)
+def smoke() -> None:
+    """Run the same checks on a T4 GPU to validate the CUDA path (cents)."""
+    _run(["modal/checks.py"])
+
+
+@app.function(gpu="A10G", timeout=60 * 60 * 6, volumes={"/data": data_volume})
+def train(config_rel: str = "configs/scail_animation_lora.yaml") -> None:
+    """Run the real SCAIL trainer. NOT free -- needs a real checkpoint + data.
+
+    Upload your base checkpoint, Gemma text encoder, and preprocessed data to the
+    ``scail-data`` Volume (mounted at /data), and point the config's paths at /data.
+    For the full 19B model use a bigger GPU (e.g. gpu="A100") and expect real cost.
+    """
+    _run(["packages/ltx-trainer/scripts/train.py", config_rel])
+    data_volume.commit()
+
+
+@app.local_entrypoint()
+def main(target: str = "verify", config_rel: str = "configs/scail_animation_lora.yaml") -> None:
+    if target == "smoke":
+        smoke.remote()
+    elif target == "train":
+        train.remote(config_rel)
+    else:
+        verify.remote()

modal/checks.py

@@ -0,0 +1,159 @@
+# ruff: noqa: T201
+"""Device-aware SCAIL-2 smoke checks (CPU or GPU), runnable anywhere.
+
+Consolidates the Phase 1-4 plumbing checks into one script that auto-selects CUDA
+when available. It uses tiny, randomly-initialised models and synthetic data, so it
+needs **no checkpoint and no dataset** -- ideal for a near-free Modal test that
+validates the SCAIL training/inference code path end to end in a real Linux/GPU
+environment before spending credits on the full 19B model.
+
+Run locally:   uv run --package ltx-trainer python modal/checks.py
+On Modal:      modal run modal/app.py::verify        (CPU)
+               modal run modal/app.py::smoke         (T4 GPU)
+"""
+
+from __future__ import annotations
+
+from dataclasses import replace
+
+import torch
+
+from ltx_core.components.noisers import GaussianNoiser
+from ltx_core.components.patchifiers import VideoLatentPatchifier
+from ltx_core.conditioning import DrivingMode, VideoConditionByDrivingLatent
+from ltx_core.model.transformer.mask_channels_checkpoint import (
+    widen_module_patchify_proj_for_mask_channels,
+)
+from ltx_core.model.transformer.model import LTXModel, LTXModelType
+from ltx_core.tools import VideoLatentTools
+from ltx_core.types import SpatioTemporalScaleFactors, VideoLatentShape
+from ltx_pipelines.scail_animation import build_scail_conditionings
+from ltx_pipelines.utils.helpers import create_noised_state, modality_from_latent_state
+from ltx_trainer.timestep_samplers import UniformTimestepSampler
+from ltx_trainer.training_strategies.flexible import (
+    DrivingConditionConfig,
+    FlexibleStrategy,
+    FlexibleStrategyConfig,
+    MaskChannelsConditionConfig,
+    ModalityConfig,
+)
+
+# --- config ---------------------------------------------------------------------
+B, C, F, H, W = 1, 16, 3, 4, 4
+HEADS, HEAD_DIM, LAYERS, INNER = 4, 16, 2, 64
+CTX_LEN, K = 8, 6
+SCALE = SpatioTemporalScaleFactors.default()
+MASK_CH = SCALE.time * (K + 1)  # 56
+F_PIX = (F - 1) * SCALE.time + 1
+N = F * H * W
+
+
+def _tiny_model(mask_channels: int, device: torch.device, dtype: torch.dtype) -> LTXModel:
+    model = LTXModel(
+        model_type=LTXModelType.VideoOnly,
+        num_attention_heads=HEADS,
+        attention_head_dim=HEAD_DIM,
+        in_channels=C,
+        out_channels=C,
+        num_layers=LAYERS,
+        cross_attention_dim=INNER,
+        mask_conditioning_channels=mask_channels,
+    )
+    for p in model.parameters():
+        torch.nn.init.normal_(p, std=0.02)
+    return model.to(device=device, dtype=dtype).eval()
+
+
+def main() -> None:
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    dtype = torch.float32
+    print(f"== SCAIL-2 checks on device={device} "
+          f"({torch.cuda.get_device_name(0) if device.type == 'cuda' else 'cpu'}) ==")
+    torch.manual_seed(0)
+
+    tools = VideoLatentTools(
+        patchifier=VideoLatentPatchifier(patch_size=1),
+        target_shape=VideoLatentShape(batch=B, channels=C, frames=F, height=H, width=W),
+        fps=24.0,
+        scale_factors=SCALE,
+    )
+
+    # Phase 1: driving-latent concat + ΔW RoPE (inference conditioning path).
+    driving = torch.randn(B, C, F, H, W, device=device, dtype=dtype)
+    noiser = GaussianNoiser(generator=torch.Generator(device=device).manual_seed(1))
+    state = create_noised_state(
+        tools, [VideoConditionByDrivingLatent(driving, mode=DrivingMode.ANIMATION)], noiser, dtype, device, 1.0
+    )
+    assert state.latent.shape[1] == 2 * N
+    print(f"[OK] Phase 1 driving concat: seq {N} -> {state.latent.shape[1]}")
+
+    # Phase 2: widen patchify_proj (zero-init) is output-preserving.
+    base = _tiny_model(0, device, dtype)
+    ctx = torch.randn(B, CTX_LEN, INNER, device=device, dtype=dtype)
+    sigma = torch.ones(B, device=device, dtype=dtype)
+    mod = modality_from_latent_state(state, context=ctx, sigma=sigma)
+    with torch.inference_mode():
+        out0, _ = base(video=mod, audio=None, perturbations=None)
+    widen_module_patchify_proj_for_mask_channels(base, MASK_CH)
+    mod_cc = replace(mod, cond_channels=torch.randn(B, 2 * N, MASK_CH, device=device, dtype=dtype))
+    with torch.inference_mode():
+        out1, _ = base(video=mod_cc, audio=None, perturbations=None)
+    assert torch.allclose(out0, out1, atol=1e-4), "zero-init widening changed output"
+    print(f"[OK] Phase 2 zero-init widen: in_features={base.patchify_proj.in_features}, output preserved")
+
+    # Phase 3: FlexibleStrategy driving + mask-channels training step.
+    cfg = FlexibleStrategyConfig(
+        name="flexible",
+        video=ModalityConfig(
+            is_generated=True,
+            latents_dir="video_latents",
+            conditions=[
+                DrivingConditionConfig(type="driving", latents_dir="driving_latents"),
+                MaskChannelsConditionConfig(type="mask_channels", mask_dir="char_masks", num_slots=K),
+            ],
+        ),
+    )
+    strategy = FlexibleStrategy(cfg)
+
+    def latents() -> dict:
+        return {
+            "latents": torch.randn(B, C, F, H, W, device=device),
+            "num_frames": torch.tensor([F]),
+            "height": torch.tensor([H]),
+            "width": torch.tensor([W]),
+            "fps": torch.tensor([24.0]),
+        }
+
+    mask_pix = torch.rand(B, K + 1, F_PIX, H * SCALE.height, W * SCALE.width, device=device)
+    batch = {
+        "video_latents": latents(),
+        "driving_latents": latents(),
+        "char_masks": {"mask": (mask_pix > 0.5).float()},
+        "conditions": {
+            "video_prompt_embeds": torch.randn(B, CTX_LEN, INNER, device=device),
+            "audio_prompt_embeds": torch.randn(B, CTX_LEN, INNER, device=device),
+            "prompt_attention_mask": torch.ones(B, CTX_LEN, device=device),
+        },
+    }
+    inputs = strategy.prepare_training_inputs(batch, UniformTimestepSampler(0.0, 1.0))
+    assert inputs.video.cond_channels.shape == (B, 2 * N, MASK_CH)
+    model = _tiny_model(MASK_CH, device, dtype)
+    with torch.inference_mode():
+        vpred, _ = model(video=inputs.video, audio=None, perturbations=None)
+    loss = strategy.compute_loss(vpred, None, inputs)
+    assert loss.shape == (B,)
+    assert torch.isfinite(loss).all()
+    cc_shape = tuple(inputs.video.cond_channels.shape)
+    print(f"[OK] Phase 3 training step: cond_channels {cc_shape}, loss {loss.item():.4f}")
+
+    # Phase 4: inference conditioning assembly.
+    masks = batch["char_masks"]["mask"]
+    conds = build_scail_conditionings(driving, masks, mode=DrivingMode.ANIMATION)
+    assert len(conds) == 2
+    print("[OK] Phase 4 build_scail_conditionings: [driving, mask_channels]")
+
+    print("\nAll SCAIL-2 checks passed on", device)
+
+
+if __name__ == "__main__":
+    main()

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).
+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.
-
-**Requirements**: `uv sync --frozen --extra fp8-trtllm`
 
 **Usage with QuantizationPolicy:**
 

packages/ltx-core/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "ltx-core"
-version = "v1.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)
+    chunks = []
-    return [BatchedPerturbationConfig([next(it) for _ in range(s)]) for s in sizes]
+    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.
+            cpu_slots_count: Number of pinned CPU buffer slots. ``None`` falls
-                ``None`` = RAM streaming (all blocks pre-loaded with LoRA fusion).
+                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)
+        sync = create_stream_sync(device)
-        gpu_pool = BufferPool(
+        gpu_pool = BufferPool(source.slot_nbytes, gpu_slots_count, device, reuse_barrier=sync.reuse_barrier)
-            source.slot_nbytes,
-            gpu_slots_count,
-            device,
-            reuse_barrier=lambda event: copy_stream.wait_event(event),
-        )
         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":
+        # Move A/B to a GPU-class target (CUDA/MPS) so the B@A aggregation runs on
-            a = a.to(device=device, non_blocking=True)
+        # the device; on a CPU target they stay put. non_blocking only helps CUDA.
-            b = b.to(device=device, non_blocking=True)
+        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.
+        sync: Coordinates copy-vs-compute ordering for the backend
-        target_device: GPU device for compute.
+            (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:
+    ) -> StreamEvent | None:
-        """Enqueue H2D copy + LoRA fusion on the copy stream and wait on compute.
+        """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 --
+        in *gpu_weights*.
-        and instrumentation regions wrapping it -- observe the full transfer time.
+        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):
+        with self._sync.copy_scope():
-            raw[:nbytes].copy_(cpu_buffer[:nbytes], non_blocking=True)
+            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 self._sync.commit_copy()
-        return h2d_event
 
-    def release(self, idx: int, event: torch.cuda.Event) -> None:
+    def release(self, idx: int, event: StreamEvent | None) -> None:
-        """Attach a compute-done event -- waited before this buffer is recycled."""
+        """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."""
+        """Drain outstanding copy/compute work and release all resources."""
-        self._copy_stream.synchronize()
+        self._sync.synchronize()
-        torch.cuda.current_stream(self._target_device).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``."""
+        """Fuse LoRA deltas directly into GPU block weights via ``fuse_rule``.
-        agg_dtype = self._fuse_rule.aggregation_dtype
+        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,
+    registration fails. Pinning is fundamentally a CUDA driver operation; when
-    since pinning is fundamentally a CUDA driver operation.
+    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."""
+        """Release the block weights once its forward pass has been enqueued.
-        compute_done = torch.cuda.Event()
+        The provider guards the buffer against reuse until this block's compute
-        compute_done.record(torch.cuda.current_stream(self._target_device))
+        completes."""
-        self._provider.release(block_idx, event=compute_done)
+        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/conditioning/__init__.py

@@ -5,10 +5,14 @@ from ltx_core.conditioning.item import ConditioningItem
 from ltx_core.conditioning.types import (
     AudioConditionByReferenceLatent,
     ConditioningItemAttentionStrengthWrapper,
+    DrivingMode,
+    VideoConditionByDrivingLatent,
     VideoConditionByKeyframeIndex,
     VideoConditionByLatentIndex,
     VideoConditionByMask,
+    VideoConditionByMaskChannels,
     VideoConditionByReferenceLatent,
+    encode_mask_channels,
 )
 
 __all__ = [
@@ -16,8 +20,12 @@ __all__ = [
     "ConditioningError",
     "ConditioningItem",
     "ConditioningItemAttentionStrengthWrapper",
+    "DrivingMode",
+    "VideoConditionByDrivingLatent",
     "VideoConditionByKeyframeIndex",
     "VideoConditionByLatentIndex",
     "VideoConditionByMask",
+    "VideoConditionByMaskChannels",
     "VideoConditionByReferenceLatent",
+    "encode_mask_channels",
 ]

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

@@ -0,0 +1,28 @@
+"""Helpers for in-context conditioning channels (``LatentState.cond_channels``).
+
+Conditioning channels are extra per-token input features (e.g. SCAIL-2 mask
+channels) that ride alongside the latent in patchified token space ``(B, T, C)``
+and are concatenated onto the latent before the model's first projection. When
+present they must stay length-aligned with the token sequence, so any
+conditioning item that *appends* tokens must also extend the channels.
+"""
+
+from __future__ import annotations
+
+import torch
+
+from ltx_core.types import LatentState
+
+
+def extend_cond_channels(latent_state: LatentState, num_new_tokens: int) -> torch.Tensor | None:
+    """Return ``cond_channels`` extended with ``num_new_tokens`` zero rows, or ``None``.
+
+    Appended tokens (reference / driving / keyframe) carry no in-context signal by default, so they
+    are padded with zeros to preserve the ``T``-alignment invariant. Returns ``None`` unchanged when
+    the state has no conditioning channels (the standard case), keeping non-SCAIL pipelines untouched.
+    """
+    cond_channels = latent_state.cond_channels
+    if cond_channels is None:
+        return None
+    zeros = cond_channels.new_zeros(cond_channels.shape[0], num_new_tokens, cond_channels.shape[2])
+    return torch.cat([cond_channels, zeros], dim=1)

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

@@ -1,8 +1,10 @@
 """Conditioning type implementations."""
 
 from ltx_core.conditioning.types.attention_strength_wrapper import ConditioningItemAttentionStrengthWrapper
+from ltx_core.conditioning.types.driving_video_cond import DrivingMode, VideoConditionByDrivingLatent
 from ltx_core.conditioning.types.keyframe_cond import VideoConditionByKeyframeIndex
 from ltx_core.conditioning.types.latent_cond import VideoConditionByLatentIndex
+from ltx_core.conditioning.types.mask_channels_cond import VideoConditionByMaskChannels, encode_mask_channels
 from ltx_core.conditioning.types.mask_cond import VideoConditionByMask
 from ltx_core.conditioning.types.reference_audio_cond import AudioConditionByReferenceLatent
 from ltx_core.conditioning.types.reference_video_cond import VideoConditionByReferenceLatent
@@ -10,8 +12,12 @@ from ltx_core.conditioning.types.reference_video_cond import VideoConditionByRef
 __all__ = [
     "AudioConditionByReferenceLatent",
     "ConditioningItemAttentionStrengthWrapper",
+    "DrivingMode",
+    "VideoConditionByDrivingLatent",
     "VideoConditionByKeyframeIndex",
     "VideoConditionByLatentIndex",
     "VideoConditionByMask",
+    "VideoConditionByMaskChannels",
     "VideoConditionByReferenceLatent",
+    "encode_mask_channels",
 ]

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

@@ -0,0 +1,162 @@
+"""Driving-video conditioning for SCAIL-2-style end-to-end character animation.
+
+Ports mechanism 1 + 3 of SCAIL-2 (arXiv:2606.10804) to LTX-2: the *driving*
+video latent is concatenated directly into the DiT token sequence (no skeleton /
+pose intermediate), and carries a fixed spatial offset ``width_offset`` (the
+paper's ΔW) on the RoPE width axis so its coordinates stay detached from the
+main video tokens. This is the inference-only PoC path -- it reuses the existing
+frozen-reference-token machinery and does not touch the transformer or RoPE.
+
+This mirrors :class:`ltx_core.conditioning.types.reference_video_cond.VideoConditionByReferenceLatent`
+(same patchify -> positions -> append -> attention-mask flow); the only new
+behaviour is the mode-aware coordinate assignment.
+
+Scope note (Phase 1): SCAIL-2's in-context mask channels (mechanism 2) and the
+reference-latent height shift ΔH_ref of Replacement Mode require model-weight
+surgery / a separate reference token group and are intentionally out of scope
+here. Because LTX handles the reference image as a frame-0 in-place replacement
+(not a separate token group), the driving-token placement is identical for both
+:class:`DrivingMode` values in Phase 1 -- ``mode`` is stored for forward
+compatibility and to document intent, but does not yet alter the driving
+coordinates. See the plan for the deferred Phase 2 work.
+"""
+
+from __future__ import annotations
+
+from enum import Enum
+
+import torch
+
+from ltx_core.components.patchifiers import get_pixel_coords
+from ltx_core.conditioning.cond_channels import extend_cond_channels
+from ltx_core.conditioning.item import ConditioningItem
+from ltx_core.conditioning.mask_utils import update_attention_mask
+from ltx_core.tools import VideoLatentTools
+from ltx_core.types import LatentState, VideoLatentShape
+
+# Default normalization ceiling for the RoPE width axis. Must stay in sync with
+# the model's ``positional_embedding_max_pos[2]`` (see rope.py:precompute_freqs_cis
+# default ``max_pos=[20, 2048, 2048]`` and model.py `_init_` default). Driving
+# width coordinates that reach or exceed this value would wrap under RoPE.
+DEFAULT_MAX_WIDTH_POSITION = 2048
+
+
+class DrivingMode(Enum):
+    """SCAIL-2 conditioning mode. See module docstring for the Phase 1 caveat."""
+
+    ANIMATION = "animation"
+    REPLACEMENT = "replacement"
+
+
+class VideoConditionByDrivingLatent(ConditioningItem):
+    """Append driving-video tokens with a width-axis RoPE offset (SCAIL-2 ΔW).
+
+    The driving tokens are appended after the target sequence as clean latents
+    (placeholder zeros in the noisy latent), kept frozen (``denoise_mask =
+    1 - strength``), temporally aligned to the target, and shifted along the
+    width axis by ``width_offset`` so they occupy ``[ΔW, ΔW + Wv)`` while the
+    target stays at ``[0, Wv)``.
+
+    Args:
+        latent: Driving video latents ``[B, C, F, H, W]``. Must match the target
+            shape (same F/H/W) so tokens align frame-for-frame with the target.
+        mode: SCAIL-2 mode (reserved for Phase 2; see module docstring).
+        width_offset: ΔW in RoPE pixel-space width units. ``None`` (default) uses
+            the target's pixel width (``target_shape.width * scale_factors.width``),
+            placing the driving tokens immediately to the right of the target.
+        strength: 1.0 keeps the driving latent fully clean (frozen); 0.0 would
+            denoise it. Default 1.0.
+        max_width_position: RoPE width normalization ceiling; validation raises if
+            the shifted driving coordinates would reach it. Keep in sync with the
+            model's ``positional_embedding_max_pos[2]``.
+    """
+
+    def __init__(
+        self,
+        latent: torch.Tensor,
+        mode: DrivingMode = DrivingMode.ANIMATION,
+        width_offset: float | None = None,
+        strength: float = 1.0,
+        max_width_position: int = DEFAULT_MAX_WIDTH_POSITION,
+    ):
+        self.latent = latent
+        self.mode = mode
+        self.width_offset = width_offset
+        self.strength = strength
+        self.max_width_position = max_width_position
+
+    def apply_to(
+        self,
+        latent_state: LatentState,
+        latent_tools: VideoLatentTools,
+    ) -> LatentState:
+        """Append driving tokens with target-aligned time and a ΔW width shift."""
+        tokens = latent_tools.patchifier.patchify(self.latent)
+
+        num_target_tokens = latent_tools.patchifier.get_token_count(latent_tools.target_shape)
+        if tokens.shape[1] != num_target_tokens:
+            raise ValueError(
+                "VideoConditionByDrivingLatent expects the driving latent to match the target shape "
+                f"(same F/H/W): got {tokens.shape[1]} driving tokens vs {num_target_tokens} target tokens. "
+                "Resize/resample the driving video to the target resolution and frame count."
+            )
+
+        # Compute the driving tokens' own pixel-space coordinates (same flow as
+        # the base reference conditioning and create_initial_state).
+        latent_coords = latent_tools.patchifier.get_patch_grid_bounds(
+            output_shape=VideoLatentShape.from_torch_shape(self.latent.shape),
+            device=self.latent.device,
+        )
+        positions = get_pixel_coords(
+            latent_coords=latent_coords,
+            scale_factors=latent_tools.scale_factors,
+            causal_fix=latent_tools.causal_fix,
+        ).to(dtype=torch.float32)
+
+        # Temporal alignment: copy the target's time coordinates so the driving
+        # tokens sit on exactly the same time grid as z_t (robust to causal_fix /
+        # fps nuances). Token order is a flattened (f h w) grid identical to the
+        # target's, so a token-wise copy is frame-aligned.
+        positions[:, 0:1, :] = latent_state.positions[:, 0:1, :num_target_tokens].to(dtype=torch.float32)
+
+        # ΔW: shift the driving tokens along the width axis so they stay spatially
+        # detached from the target tokens. Default offset = target pixel width,
+        # giving target=[0, Wv), driving=[Wv, 2*Wv).
+        width_offset = self.width_offset
+        if width_offset is None:
+            width_offset = float(latent_tools.target_shape.width * latent_tools.scale_factors.width)
+        positions[:, 2, ...] = positions[:, 2, ...] + width_offset
+
+        max_width = positions[:, 2, ...].max().item()
+        if max_width >= self.max_width_position:
+            raise ValueError(
+                f"Driving width coordinate {max_width:.1f} reaches the RoPE ceiling "
+                f"{self.max_width_position} and would wrap. Reduce width_offset "
+                f"(currently {width_offset:.1f}) or lower the output width."
+            )
+
+        denoise_mask = torch.full(
+            size=(*tokens.shape[:2], 1),
+            fill_value=1.0 - self.strength,
+            device=self.latent.device,
+            dtype=self.latent.dtype,
+        )
+
+        new_attention_mask = update_attention_mask(
+            latent_state=latent_state,
+            attention_mask=None,
+            num_noisy_tokens=num_target_tokens,
+            num_new_tokens=tokens.shape[1],
+            batch_size=tokens.shape[0],
+            device=self.latent.device,
+            dtype=self.latent.dtype,
+        )
+
+        return LatentState(
+            latent=torch.cat([latent_state.latent, torch.zeros_like(tokens)], dim=1),
+            denoise_mask=torch.cat([latent_state.denoise_mask, denoise_mask], dim=1),
+            positions=torch.cat([latent_state.positions, positions], dim=2),
+            clean_latent=torch.cat([latent_state.clean_latent, tokens], dim=1),
+            attention_mask=new_attention_mask,
+            cond_channels=extend_cond_channels(latent_state, tokens.shape[1]),
+        )

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

@@ -1,6 +1,7 @@
 import torch
 
 from ltx_core.components.patchifiers import get_pixel_coords
+from ltx_core.conditioning.cond_channels import extend_cond_channels
 from ltx_core.conditioning.item import ConditioningItem
 from ltx_core.conditioning.mask_utils import update_attention_mask
 from ltx_core.tools import VideoLatentTools
@@ -81,4 +82,5 @@ class VideoConditionByKeyframeIndex(ConditioningItem):
             positions=torch.cat([latent_state.positions, positions], dim=2),
             clean_latent=torch.cat([latent_state.clean_latent, tokens], dim=1),
             attention_mask=new_attention_mask,
+            cond_channels=extend_cond_channels(latent_state, tokens.shape[1]),
         )

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

@@ -0,0 +1,122 @@
+"""SCAIL-2 in-context mask conditioning (mechanism 2) for LTX-2.
+
+Encodes ``K+1`` semantic pixel-space masks (1 environment switch + ``K`` character
+binding slots) into per-token conditioning channels and writes them onto the
+sequence via :attr:`LatentState.cond_channels`. Following the paper, mask signals
+are carried by the *driving* (and reference) tokens while the noisy target keeps
+an all-zero mask -- so by default the channels are written onto the last ``N``
+tokens of the sequence (the appended driving group), which means **this item must
+be applied after the driving conditioning**.
+
+Channel expansion (LTX-2 faithful port of the paper's ``4(K+1)``): each semantic
+mask is spatially downsampled to the latent grid and temporally stacked along the
+channel dimension by the VAE temporal factor ``t`` (8 for LTX-2, vs 4 for the
+paper's Wan-2.1 backbone), giving ``t*(K+1)`` channels -- ``8*7 = 56`` for the
+default ``K=6``. The model's ``patchify_proj`` must be built with a matching
+``mask_conditioning_channels`` (see :class:`ltx_core.model.transformer.model.LTXModel`).
+"""
+
+from __future__ import annotations
+
+from dataclasses import replace
+
+import torch
+import torch.nn.functional as F
+from einops import rearrange
+
+from ltx_core.conditioning.item import ConditioningItem
+from ltx_core.tools import VideoLatentTools
+from ltx_core.types import LatentState
+
+
+def encode_mask_channels(
+    masks: torch.Tensor,
+    temporal_factor: int,
+    height_lat: int,
+    width_lat: int,
+    frames_lat: int,
+) -> torch.Tensor:
+    """Encode ``[B, K+1, F_pix, H_pix, W_pix]`` masks into ``[B, t*(K+1), F_lat, H_lat, W_lat]``.
+
+    Each semantic mask is area-downsampled to the latent spatial grid, then the pixel frames that
+    map to one latent frame are stacked along the channel dimension (temporal factor ``t``). The
+    causal first latent frame corresponds to a single pixel frame, which is replicated across its
+    ``t`` stacked channels so every latent frame yields a uniform ``t`` channels per semantic class.
+    """
+    b, s, f_pix, _h_pix, _w_pix = masks.shape
+    t = temporal_factor
+    expected_f_pix = (frames_lat - 1) * t + 1
+    if f_pix != expected_f_pix:
+        raise ValueError(
+            f"mask pixel frames ({f_pix}) incompatible with latent frames ({frames_lat}) at temporal "
+            f"factor {t}: expected (F_lat - 1) * {t} + 1 = {expected_f_pix}."
+        )
+
+    # Spatial downsample every (semantic, frame) mask to the latent grid.
+    flat = rearrange(masks.to(dtype=torch.float32), "b s f h w -> (b s f) 1 h w")
+    down = F.interpolate(flat, size=(height_lat, width_lat), mode="area")
+    down = rearrange(down, "(b s f) 1 h w -> b s f h w", b=b, s=s)
+
+    # Temporal stacking. Latent frame 0 = pixel frame 0 (causal), replicated across t channels;
+    # latent frames 1.. group t consecutive pixel frames.
+    first = down[:, :, :1].repeat(1, 1, t, 1, 1).unsqueeze(2)  # [B, S, 1, t, H, W]
+    rest = rearrange(down[:, :, 1:], "b s (fl t) h w -> b s fl t h w", t=t)  # [B, S, F_lat-1, t, H, W]
+    stacked = torch.cat([first, rest], dim=2)  # [B, S, F_lat, t, H, W]
+
+    # Fold (semantic, temporal) into a single channel axis, semantic-major: channel = s * t + ti.
+    return rearrange(stacked, "b s fl t h w -> b (s t) fl h w")  # [B, t*(K+1), F_lat, H_lat, W_lat]
+
+
+class VideoConditionByMaskChannels(ConditioningItem):
+    """Write SCAIL-2 in-context mask channels onto the driving/reference tokens.
+
+    Args:
+        masks: Pixel-space semantic masks ``[B, K+1, F_pix, H_pix, W_pix]``. Channel 0 is the
+            environment switch (whether the environment comes from the reference vs the driving
+            video); channels ``1..K`` are the character binding slots (regions sharing a slot share
+            motion). ``F_pix`` must equal ``(F_lat - 1) * temporal_factor + 1``.
+        applies_to_last_n: Number of trailing tokens to write the mask onto. ``None`` (default) uses
+            the target token count, i.e. the appended driving group when this item runs right after
+            the driving conditioning. The noisy target tokens keep an all-zero mask (paper-faithful).
+    """
+
+    def __init__(self, masks: torch.Tensor, applies_to_last_n: int | None = None):
+        self.masks = masks
+        self.applies_to_last_n = applies_to_last_n
+
+    def apply_to(self, latent_state: LatentState, latent_tools: VideoLatentTools) -> LatentState:
+        shape = latent_tools.target_shape
+        cond = encode_mask_channels(
+            masks=self.masks,
+            temporal_factor=latent_tools.scale_factors.time,
+            height_lat=shape.height,
+            width_lat=shape.width,
+            frames_lat=shape.frames,
+        )
+        tokens = latent_tools.patchifier.patchify(cond)  # [B, T_region, C_mask]
+        b, n_region, c_mask = tokens.shape
+
+        total_tokens = latent_state.latent.shape[1]
+        n = self.applies_to_last_n if self.applies_to_last_n is not None else n_region
+        if n != n_region:
+            raise ValueError(
+                f"applies_to_last_n ({n}) must equal the encoded mask token count ({n_region})."
+            )
+        if n > total_tokens:
+            raise ValueError(
+                f"cannot write {n} mask tokens onto a sequence of only {total_tokens} tokens."
+            )
+
+        cond_channels = latent_state.cond_channels
+        if cond_channels is None:
+            cond_channels = tokens.new_zeros(b, total_tokens, c_mask)
+        else:
+            if cond_channels.shape[2] != c_mask:
+                raise ValueError(
+                    f"existing cond_channels width ({cond_channels.shape[2]}) != mask channels ({c_mask})."
+                )
+            cond_channels = cond_channels.clone()
+
+        start = total_tokens - n
+        cond_channels[:, start : start + n] = tokens.to(dtype=cond_channels.dtype)
+        return replace(latent_state, cond_channels=cond_channels)

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

@@ -4,6 +4,7 @@ from __future__ import annotations
 
 import torch
 
+from ltx_core.conditioning.cond_channels import extend_cond_channels
 from ltx_core.conditioning.mask_utils import update_attention_mask
 from ltx_core.tools import LatentTools
 from ltx_core.types import LatentState
@@ -56,4 +57,5 @@ class AudioConditionByReferenceLatent:
             positions=torch.cat([latent_state.positions, self.positions], dim=2),
             clean_latent=torch.cat([latent_state.clean_latent, tokens], dim=1),
             attention_mask=new_attention_mask,
+            cond_channels=extend_cond_channels(latent_state, tokens.shape[1]),
         )

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

@@ -3,6 +3,7 @@
 import torch
 
 from ltx_core.components.patchifiers import get_pixel_coords
+from ltx_core.conditioning.cond_channels import extend_cond_channels
 from ltx_core.conditioning.item import ConditioningItem
 from ltx_core.conditioning.mask_utils import update_attention_mask
 from ltx_core.tools import VideoLatentTools
@@ -99,4 +100,5 @@ class VideoConditionByReferenceLatent(ConditioningItem):
             positions=torch.cat([latent_state.positions, positions], dim=2),
             clean_latent=torch.cat([latent_state.clean_latent, tokens], dim=1),
             attention_mask=new_attention_mask,
+            cond_channels=extend_cond_channels(latent_state, tokens.shape[1]),
         )

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):
+class PerturbationType(IntEnum):
-    """Types of attention perturbations for STG (Spatio-Temporal Guidance)."""
+    """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_VIDEO_SELF_ATTN = 0
-    SKIP_V2A_CROSS_ATTN = "skip_v2a_cross_attn"
+    SKIP_AUDIO_SELF_ATTN = 1
-    SKIP_VIDEO_SELF_ATTN = "skip_video_self_attn"
+    SKIP_A2V_CROSS_ATTN = 2
-    SKIP_AUDIO_SELF_ATTN = "skip_audio_self_attn"
+    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(
+    def __init__(
-        self, perturbation_type: PerturbationType, block: int, device: DeviceLikeType, dtype: torch.dtype
+        self,
-    ) -> torch.Tensor:
+        perturbations: list[PerturbationConfig],
-        mask = torch.ones((len(self.perturbations),), device=device, dtype=dtype)
+        num_blocks: int,
-        for batch_idx, perturbation in enumerate(self.perturbations):
+        device: DeviceLikeType | None = None,
-            if perturbation.is_perturbed(perturbation_type, block):
+        dtype: torch.dtype | None = None,
-                mask[batch_idx] = 0
+    ) -> 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:
+    def batch_slice(self, start: int, end: int) -> "BatchedPerturbationConfig":
-        mask = self.mask(perturbation_type, block, values.device, values.dtype)
+        """A view over samples ``[start:end]`` of the batch, by slicing the mask tensors.
-        return mask.view(mask.numel(), *([1] * len(values.shape[1:])))
+        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":
+    def empty(
-        return BatchedPerturbationConfig([PerturbationConfig.empty() for _ in range(batch_size)])
+        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)
+        products = _products_for_sd_key(lora_sd_and_strengths, key, rule.aggregation_dtype, fusion_device)
-        deltas = aggregate_lora_products(products, fuse_rule.aggregation_dtype)
+        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
+        # On CUDA/CPU we use autocast(dtype=float32) rather than self.float()
-        # upcasts bf16 weights per-op at kernel level, avoiding the temporary
+        # because it upcasts bf16 weights per-op at kernel level, avoiding the
-        # memory spike of self.float() / self.to(original_dtype).
+        # 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,16 +28,9 @@ 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
 except ImportError:
     flash_attn_interface = None
@@ -44,11 +38,16 @@ 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
+        """Human-readable identifier. Encodes the SDPA priority list so a
-        list so a single-backend pin reads differently from the full-priority
+        single-backend pin reads differently from the full-priority dispatcher walk."""
-        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):
+class MPSSdpaAttention(AttentionCallable):
-    label = "xFormers"
+    """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,
+        self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, heads: int, mask: torch.Tensor | None = None
-        q: torch.Tensor,
-        k: torch.Tensor,
-        v: torch.Tensor,
-        heads: int,
-        mask: torch.Tensor | None = None,
     ) -> torch.Tensor:
-        if memory_efficient_attention is None:
+        if _mps_sdpa_opt is None:
-            raise RuntimeError("XFormersAttention was selected but `xformers` is not installed.")
+            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
-
+        q, k, v = (t.view(b, -1, heads, dim_head).transpose(1, 2) for t in (q, k, v))
-        # xformers expects [B, M, H, K]
-        q, k, v = (t.view(b, -1, heads, dim_head) 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
+        out = _mps_sdpa_opt(q, k, v, attn_mask=mask)
-            # doesn't this remove the padding again??
+        out = out.transpose(1, 2).reshape(b, -1, heads * dim_head)
-            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)
         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
+    """Pick a mask-aware attention. On macOS, Apple's fused MPSGraph kernel via
-    the full priority list (the dispatcher rejects FLASH automatically when a
+    ``mps-sdpa`` (a hard dependency on Apple Silicon, else torch's SDPA on
-    mask is present and walks past it)."""
+    Intel/CPU macs); else SDPA with the full priority list (the dispatcher
-    if memory_efficient_attention is not None:
+    rejects FLASH automatically when a mask is present and walks past it --
-        return XFormersAttention()
+    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:
+class CompiledBlockPerturbationsProcessor(BlockPerturbationsProcessor):
-    """Marks the per-block seq dim dynamic, then delegates to an inner processor.
+    """Per-block input prep for compiled blocks: mark the seq dim dynamic, then attach perturbation
-    Installed by ``compile_transformer`` so the per-block compile artifact stays
+    as config-independent runtime masks so the block traces ONCE.
-    shape-polymorphic. Wraps whatever ``block_input_processor`` was already on
+    The ``mark_dynamic`` calls keep the per-block compile artifact shape-polymorphic; they run in
-    the model -- callers that customised the processor keep their customisation;
+    eager mode (this processor lives outside the compiled region) on the tensors about to cross into
-    only the seq-dim marking is layered on top. Lives outside the compiled
+    the trace. Both keep-masks are then attached UNCONDITIONALLY and the skip flags pinned False, so
-    region, so ``mark_dynamic`` runs in eager mode on the tensors that are
+    the trace is identical for every pass (cond / uncond / STG): the block never sees a flipped
-    about to cross into the trace.
+    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/mask_channels_checkpoint.py

@@ -0,0 +1,92 @@
+"""Checkpoint surgery to enable SCAIL-2 in-context mask channels on a trained LTX-2 model.
+
+Widening ``LTXModel.mask_conditioning_channels`` from 0 to ``C`` grows the video
+``patchify_proj`` weight from ``[inner, in]`` to ``[inner, in + C]``. This helper
+appends ``C`` **zero** input columns to that weight so a converted checkpoint is
+numerically identical to the original until the new columns are finetuned -- the
+extra channels contribute nothing at load time.
+
+Usage::
+
+    sd = load_state_dict(path)
+    widen_patchify_proj_for_mask_channels(sd, mask_channels=56)
+    model = LTXModel(..., mask_conditioning_channels=56)
+    model.load_state_dict(sd)
+"""
+
+from __future__ import annotations
+
+import torch
+
+_VIDEO_PROJ_SUFFIX = "patchify_proj.weight"
+_AUDIO_PROJ_SUFFIX = "audio_patchify_proj.weight"
+
+
+def widen_module_patchify_proj_for_mask_channels(model: torch.nn.Module, mask_channels: int) -> torch.nn.Module:
+    """In place, widen a live ``LTXModel``'s video ``patchify_proj`` to accept ``mask_channels`` extra inputs.
+
+    Replaces ``model.patchify_proj`` with a wider ``Linear`` whose original input columns are copied and
+    whose ``mask_channels`` new columns are zero (so the model reproduces its pre-widening output until
+    those columns are trained). Also sets ``model.mask_conditioning_channels`` for metadata/consistency.
+    Idempotent guard: raises if the model is already widened to a different value.
+    """
+    if mask_channels < 0:
+        raise ValueError(f"mask_channels must be non-negative, got {mask_channels}")
+
+    proj = getattr(model, "patchify_proj", None)
+    if not isinstance(proj, torch.nn.Linear):
+        raise AttributeError("model has no linear 'patchify_proj' to widen")
+
+    already = int(getattr(model, "mask_conditioning_channels", 0))
+    if mask_channels in (0, already):
+        model.mask_conditioning_channels = mask_channels
+        return model
+    if already != 0:
+        raise ValueError(f"patchify_proj already widened for {already} mask channels, refusing to re-widen")
+
+    new_in = proj.in_features + mask_channels
+    wider = torch.nn.Linear(new_in, proj.out_features, bias=proj.bias is not None)
+    wider = wider.to(device=proj.weight.device, dtype=proj.weight.dtype)
+    with torch.no_grad():
+        wider.weight.zero_()
+        wider.weight[:, : proj.in_features].copy_(proj.weight)
+        if proj.bias is not None:
+            wider.bias.copy_(proj.bias)
+    model.patchify_proj = wider
+    model.mask_conditioning_channels = mask_channels
+    return model
+
+
+def widen_patchify_proj_for_mask_channels(
+    state_dict: dict[str, torch.Tensor],
+    mask_channels: int,
+) -> dict[str, torch.Tensor]:
+    """In place, append ``mask_channels`` zero input columns to the video ``patchify_proj`` weight.
+
+    Only the video projection is touched (audio ``audio_patchify_proj`` is left alone). The bias,
+    if present, is unchanged. Returns the same dict for convenience. Idempotency is the caller's
+    responsibility -- calling twice widens twice.
+    """
+    if mask_channels < 0:
+        raise ValueError(f"mask_channels must be non-negative, got {mask_channels}")
+    if mask_channels == 0:
+        return state_dict
+
+    target_keys = [
+        key
+        for key in state_dict
+        if key.endswith(_VIDEO_PROJ_SUFFIX) and not key.endswith(_AUDIO_PROJ_SUFFIX)
+    ]
+    if not target_keys:
+        raise KeyError(
+            f"No '{_VIDEO_PROJ_SUFFIX}' weight found in the state dict; cannot widen for mask channels."
+        )
+
+    for key in target_keys:
+        weight = state_dict[key]  # [inner_dim, in_features]
+        if weight.ndim != 2:
+            raise ValueError(f"Expected 2-D weight for '{key}', got shape {tuple(weight.shape)}")
+        pad = weight.new_zeros(weight.shape[0], mask_channels)
+        state_dict[key] = torch.cat([weight, pad], dim=1)
+
+    return state_dict

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

@@ -38,6 +38,9 @@ class Modality:
             attention. ``None`` means unrestricted (full) attention between
             all tokens. Built incrementally by conditioning items; see
             :class:`~ltx_core.conditioning.types.attention_strength_wrapper.ConditioningItemAttentionStrengthWrapper`.
+        cond_channels: Optional in-context conditioning channels, shape ``(B, T, C_cond)``, aligned token-for-token
+            with ``latent``. Concatenated onto ``latent`` before the first projection (see
+            :meth:`TransformerArgsPreprocessor.prepare`); never noised. ``None`` for standard models.
     """
 
     latent: (
@@ -53,6 +56,7 @@ class Modality:
     enabled: bool = True
     context_mask: torch.Tensor | None = None
     attention_mask: torch.Tensor | None = None
+    cond_channels: torch.Tensor | None = None
 
     def split(self, sizes: list[int]) -> list[Modality]:
         """Split along the batch dimension into chunks of the given sizes."""

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

@@ -73,11 +73,12 @@ class LTXModel(torch.nn.Module):
         caption_projection: torch.nn.Module | None = None,
         audio_caption_projection: torch.nn.Module | None = None,
         cross_attention_adaln: bool = False,
+        mask_conditioning_channels: int = 0,
     ):
         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",
@@ -86,6 +87,10 @@ class LTXModel(torch.nn.Module):
         )
         self._enable_gradient_checkpointing = False
         self.cross_attention_adaln = cross_attention_adaln
+        # Extra per-token input channels (e.g. SCAIL-2 in-context mask channels) concatenated onto the
+        # video latent before ``patchify_proj``. 0 keeps the standard input width; >0 widens the first
+        # projection. Fed via ``Modality.cond_channels`` (see TransformerArgsPreprocessor.prepare).
+        self.mask_conditioning_channels = mask_conditioning_channels
         self.use_middle_indices_grid = use_middle_indices_grid
         self.rope_type = rope_type
         self.double_precision_rope = double_precision_rope
@@ -154,8 +159,11 @@ class LTXModel(torch.nn.Module):
         caption_projection: torch.nn.Module | None = None,
     ) -> None:
         """Initialize video-specific components."""
-        # Video input components
+        # Video input components. When ``mask_conditioning_channels > 0`` the first projection is
-        self.patchify_proj = torch.nn.Linear(in_channels, self.inner_dim, bias=True)
+        # widened to also accept the in-context conditioning channels concatenated onto the latent.
+        self.patchify_proj = torch.nn.Linear(
+            in_channels + self.mask_conditioning_channels, self.inner_dim, bias=True
+        )
         if caption_projection is not None:
             self.caption_projection = caption_projection
 
@@ -358,21 +366,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.
+        """Process transformer blocks for LTX."""
-        Per-block perturbation masks are precomputed here and attached to each
-        modality's ``TransformerArgs`` so the block forward has no per-block
-        identity to specialise on — all blocks share a single Dynamo cache slot.
-        """
-        if perturbations is None:
-            batch_size = (video or audio).x.shape[0]
-            perturbations = BatchedPerturbationConfig.empty(batch_size)
-
         for block_idx, block in enumerate(self.transformer_blocks):
             if video is not None:
                 video = self.block_input_processor(
@@ -424,7 +429,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 +443,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 +487,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 +521,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/model_configurator.py

@@ -69,6 +69,7 @@ class LTXModelConfigurator(ModelConfigurator[LTXModel]):
             caption_projection=caption_projection,
             audio_caption_projection=audio_caption_projection,
             cross_attention_adaln=config.get("cross_attention_adaln", False),
+            mask_conditioning_channels=config.get("mask_conditioning_channels", 0),
         )
 
 
@@ -120,6 +121,7 @@ class LTXVideoOnlyModelConfigurator(ModelConfigurator[LTXModel]):
             apply_gated_attention=config.get("apply_gated_attention", False),
             caption_projection=caption_projection,
             cross_attention_adaln=config.get("cross_attention_adaln", False),
+            mask_conditioning_channels=config.get("mask_conditioning_channels", 0),
         )
 
 

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,
@@ -203,12 +201,46 @@ class TransformerArgsPreprocessor:
         )
         return pe
 
+    def _apply_patchify_proj(self, modality: Modality) -> torch.Tensor:
+        """Project patchified latents, concatenating any in-context conditioning channels first.
+
+        When ``patchify_proj`` expects more input features than the latent provides (a model built
+        with ``mask_conditioning_channels > 0``), the extra width is filled by ``modality.cond_channels``
+        (per-token, never noised). If those channels are absent they default to zeros, so a widened
+        model still runs and, with zero-initialized new projection columns, reproduces the base model's
+        output exactly.
+        """
+        latent = modality.latent
+        expected = self.patchify_proj.in_features
+        actual = latent.shape[-1]
+        if expected != actual:
+            missing = expected - actual
+            if missing < 0:
+                raise ValueError(
+                    f"patchify_proj expects {expected} input features but the latent already has {actual}."
+                )
+            cond_channels = modality.cond_channels
+            if cond_channels is None:
+                cond_channels = latent.new_zeros(latent.shape[0], latent.shape[1], missing)
+            elif cond_channels.shape[-1] != missing:
+                raise ValueError(
+                    f"cond_channels has {cond_channels.shape[-1]} channels but patchify_proj needs {missing} "
+                    f"extra input features (latent {actual} + cond {cond_channels.shape[-1]} != {expected})."
+                )
+            elif cond_channels.shape[1] != latent.shape[1]:
+                raise ValueError(
+                    f"cond_channels token length {cond_channels.shape[1]} must match latent token length "
+                    f"{latent.shape[1]}."
+                )
+            latent = torch.cat([latent, cond_channels.to(dtype=latent.dtype)], dim=-1)
+        return self.patchify_proj(latent)
+
     def prepare(
         self,
         modality: Modality,
         cross_modality: Modality | None = None,  # noqa: ARG002
     ) -> TransformerArgs:
-        x = self.patchify_proj(modality.latent)
+        x = self._apply_patchify_proj(modality)
         batch_size = x.shape[0]
         timestep, embedded_timestep = self._prepare_timestep(
             modality.timesteps, self.adaln, batch_size, modality.latent.dtype

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,19 +49,6 @@ 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

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-core/src/ltx_core/tools.py

@@ -75,6 +75,9 @@ class LatentTools(Protocol):
         clean_latent = latent_state.clean_latent[:, :num_tokens]
         denoise_mask = torch.ones_like(latent_state.denoise_mask)[:, :num_tokens]
         positions = latent_state.positions[:, :, :num_tokens]
+        cond_channels = (
+            latent_state.cond_channels[:, :num_tokens] if latent_state.cond_channels is not None else None
+        )
 
         return LatentState(
             latent=latent,
@@ -82,6 +85,7 @@ class LatentTools(Protocol):
             positions=positions,
             clean_latent=clean_latent,
             attention_mask=None,
+            cond_channels=cond_channels,
         )
 
 

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

@@ -193,6 +193,11 @@ class LatentState:
         clean_latent: Initial state of the latent before denoising, may include conditioning latents.
         attention_mask: Optional 2D self-attention mask of shape (B, T, T). Values in [0, 1] where 1 = full attention,
             0 = no attention. None means full attention everywhere. Built incrementally by conditioning items.
+        cond_channels: Optional in-context conditioning channels in patchified token space, shape (B, T, C_cond),
+            aligned token-for-token with ``latent``. These extra per-token feature channels (e.g. SCAIL-2 mask
+            channels) are concatenated onto the latent right before the model's first projection; they are never
+            noised or denoised. ``None`` means no conditioning channels (standard models). When present, the token
+            length T must always match ``latent``; token-appending conditioning items extend it with zeros.
     """
 
     latent: torch.Tensor
@@ -200,6 +205,7 @@ class LatentState:
     positions: torch.Tensor
     clean_latent: torch.Tensor
     attention_mask: torch.Tensor | None = None
+    cond_channels: torch.Tensor | None = None
 
     def clone(self) -> "LatentState":
         return LatentState(
@@ -208,4 +214,5 @@ class LatentState:
             positions=self.positions.clone(),
             clean_latent=self.clean_latent.clone(),
             attention_mask=self.attention_mask.clone() if self.attention_mask is not None else None,
+            cond_channels=self.cond_channels.clone() if self.cond_channels is not None else None,
         )

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];
+            }    
+        }
+    }
+}