Switch to HDR IC-LoRA app with on-the-fly PromptEncoder + bundled ltx-2-internal

Replaces distilled app with HDRICLoraPipeline. Prompt encoded through Gemma each run. Bundles ltx-2-internal for the HDR code path (not on public Lightricks/LTX-2 main).

README.md

@@ -1,8 +1,8 @@
 ---
-title: LTX 2.3 Distilled
-emoji: 📚
+title: LTX 2.3 HDR
+emoji: 🎞️
 colorFrom: indigo
-colorTo: green
+colorTo: purple
 sdk: gradio
 sdk_version: 6.8.0
 python_version: '3.12'
@@ -10,4 +10,4 @@ app_file: app.py
 pinned: false
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+LTX-2.3 HDR IC-LoRA Gradio app (ZeroGPU). Requires `HF_TOKEN` secret with access to `Lightricks/LTX-2.3` and `diffusers-internal-dev/LTX-HDR-LoRA`.

app.py

@@ -1,43 +1,77 @@
 import os
 import subprocess
 import sys
+from pathlib import Path
 
 # Disable torch.compile / dynamo before any torch import
 os.environ["TORCH_COMPILE_DISABLE"] = "1"
 os.environ["TORCHDYNAMO_DISABLE"] = "1"
+# OpenImageIO / OpenEXR readers need this to be set before cv2 import
+os.environ["OPENCV_IO_ENABLE_OPENEXR"] = "1"
 
 # Install xformers for memory-efficient attention
-subprocess.run([sys.executable, "-m", "pip", "install", "xformers==0.0.32.post2", "--no-build-isolation"], check=False)
-
-# Clone LTX-2 repo and install packages
-LTX_REPO_URL = "https://github.com/Lightricks/LTX-2.git"
-LTX_REPO_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "LTX-2")
-LTX_COMMIT_SHA = "ae855f8538843825f9015a419cf4ba5edaf5eec2"
+subprocess.run(
+    [sys.executable, "-m", "pip", "install", "xformers==0.0.32.post2", "--no-build-isolation"],
+    check=False,
+)
 
-if not os.path.exists(LTX_REPO_DIR):
-    print(f"Cloning {LTX_REPO_URL}...")
-    os.makedirs(LTX_REPO_DIR)
-    subprocess.run(["git", "init", LTX_REPO_DIR], check=True)
-    subprocess.run(["git", "remote", "add", "origin", LTX_REPO_URL], cwd=LTX_REPO_DIR, check=True)
-    subprocess.run(["git", "fetch", "--depth", "1", "origin", LTX_COMMIT_SHA], cwd=LTX_REPO_DIR, check=True)
-    subprocess.run(["git", "checkout", LTX_COMMIT_SHA], cwd=LTX_REPO_DIR, check=True)
+# Video / HDR dependencies
+subprocess.run(
+    [sys.executable, "-m", "pip", "install",
+     "imageio[ffmpeg]", "scikit-image", "opencv-python-headless",
+     "decord", "num2words", "OpenImageIO"],
+    check=False,
+)
+subprocess.run([sys.executable, "-m", "pip", "install", "num2words"], check=True)
 
+# Reinstall torchaudio to match the torch CUDA version on this space.
+_tv = subprocess.run(
+    [sys.executable, "-c", "import torch; print(torch.__version__)"],
+    capture_output=True, text=True,
+)
+if _tv.returncode == 0:
+    _full_ver = _tv.stdout.strip()
+    _cuda_suffix = _full_ver.split("+")[-1] if "+" in _full_ver else "cu124"
+    _base_ver = _full_ver.split("+")[0]
+    print(f"Detected torch {_full_ver}, reinstalling matching torchaudio...")
+    subprocess.run(
+        [sys.executable, "-m", "pip", "install", "--force-reinstall", "--no-deps",
+         f"torchaudio=={_base_ver}",
+         "--index-url", f"https://download.pytorch.org/whl/{_cuda_suffix}"],
+        check=False,
+    )
 
-print("Installing ltx-core and ltx-pipelines from cloned repo...")
+# ─────────────────────────────────────────────────────────────────────────────
+# ltx-core / ltx-pipelines source
+#
+# The HDRICLoraPipeline and its supporting modules (ltx_core.hdr,
+# ltx_pipelines.utils.blocks, load_video_conditioning_hdr, apply_hdr_decode_postprocess,
+# save_exr_tensor, encode_exr_sequence_to_mp4) are NOT on the public main
+# branch at the pinned commit used by the outpaint app. We install from the
+# local ltx-2-internal checkout so the HDR code path actually exists.
+# ─────────────────────────────────────────────────────────────────────────────
+LTX_INTERNAL = Path(os.environ.get(
+    "LTX_INTERNAL_PATH",
+    str(Path(__file__).resolve().parent / "ltx-2-internal"),
+))
+LTX_CORE_PKG = LTX_INTERNAL / "packages" / "ltx-core"
+LTX_PIPELINES_PKG = LTX_INTERNAL / "packages" / "ltx-pipelines"
+
+print(f"Installing ltx-core + ltx-pipelines from {LTX_INTERNAL}...")
 subprocess.run(
-    [sys.executable, "-m", "pip", "install", "--force-reinstall", "--no-deps", "-e",
-     os.path.join(LTX_REPO_DIR, "packages", "ltx-core"),
-     "-e", os.path.join(LTX_REPO_DIR, "packages", "ltx-pipelines")],
+    [sys.executable, "-m", "pip", "install",
+     "--force-reinstall", "--no-deps",
+     "-e", str(LTX_CORE_PKG),
+     "-e", str(LTX_PIPELINES_PKG)],
     check=True,
 )
-
-sys.path.insert(0, os.path.join(LTX_REPO_DIR, "packages", "ltx-pipelines", "src"))
-sys.path.insert(0, os.path.join(LTX_REPO_DIR, "packages", "ltx-core", "src"))
+sys.path.insert(0, str(LTX_PIPELINES_PKG / "src"))
+sys.path.insert(0, str(LTX_CORE_PKG / "src"))
 
 import logging
 import random
 import tempfile
-from pathlib import Path
+import zipfile
 
 import torch
 torch._dynamo.config.suppress_errors = True
@@ -48,13 +82,18 @@ import gradio as gr
 import numpy as np
 from huggingface_hub import hf_hub_download, snapshot_download
 
-from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
+from ltx_core.model.video_vae import TilingConfig
 from ltx_core.quantization import QuantizationPolicy
-from ltx_pipelines.distilled import DistilledPipeline
-from ltx_pipelines.utils.args import ImageConditioningInput
-from ltx_pipelines.utils.media_io import encode_video
+from ltx_pipelines.hdr_ic_lora import HDRICLoraPipeline, _make_tiling_config
+from ltx_pipelines.utils.blocks import PromptEncoder
+from ltx_pipelines.utils.media_io import (
+    encode_exr_sequence_to_mp4,
+    get_videostream_metadata,
+    save_exr_tensor,
+)
+from ltx_pipelines.utils.types import OffloadMode
 
-# Force-patch xformers attention into the LTX attention module.
+# xformers attention patch (same as the outpaint app).
 from ltx_core.model.transformer import attention as _attn_mod
 print(f"[ATTN] Before patch: memory_efficient_attention={_attn_mod.memory_efficient_attention}")
 try:
@@ -66,253 +105,302 @@ except Exception as e:
 
 logging.getLogger().setLevel(logging.INFO)
 
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Constants + model downloads
+# ─────────────────────────────────────────────────────────────────────────────
 MAX_SEED = np.iinfo(np.int32).max
-DEFAULT_PROMPT = (
-    "An astronaut hatches from a fragile egg on the surface of the Moon, "
-    "the shell cracking and peeling apart in gentle low-gravity motion. "
-    "Fine lunar dust lifts and drifts outward with each movement, floating "
-    "in slow arcs before settling back onto the ground."
-)
-DEFAULT_FRAME_RATE = 24.0
 
-# Resolution presets: (width, height)
+# Frames must satisfy (n-1) % 8 == 0. Aspect-ratio canvas sizes (divisible by 32).
 RESOLUTIONS = {
-    "high": {"16:9": (1536, 1024), "9:16": (1024, 1536), "1:1": (1024, 1024)},
-    "low": {"16:9": (768, 512), "9:16": (512, 768), "1:1": (768, 768)},
+    "low":  {"16:9": (768, 512), "9:16": (512, 768), "1:1": (768, 768),
+             "4:3": (768, 576), "3:4": (576, 768), "21:9": (768, 384)},
+    "high": {"16:9": (1536, 1024), "9:16": (1024, 1536), "1:1": (1024, 1024),
+             "4:3": (1536, 1152), "3:4": (1152, 1536), "21:9": (1536, 768)},
 }
 
-# Model repos
 LTX_MODEL_REPO = "Lightricks/LTX-2.3"
+DISTILLED_CHECKPOINT = "ltx-2.3-22b-distilled-1.1.safetensors"
+SPATIAL_UPSCALER = "ltx-2.3-spatial-upscaler-x2-1.1.safetensors"
 GEMMA_REPO = "google/gemma-3-12b-it-qat-q4_0-unquantized"
 
-# Download model checkpoints
+HDR_LORA_REPO = "diffusers-internal-dev/LTX-HDR-LoRA"
+HDR_LORA_FILENAME = "comfyui_models_loras_ltxv_ltx2_ltx-2.3-22b-ic-lora-hdr-0.9 (4).safetensors"
+
 print("=" * 80)
-print("Downloading LTX-2.3 distilled model + Gemma...")
+print("Downloading LTX-2.3 distilled + spatial upsampler + Gemma + HDR IC-LoRA...")
 print("=" * 80)
 
-checkpoint_path = hf_hub_download(repo_id=LTX_MODEL_REPO, filename="ltx-2.3-22b-distilled-1.1.safetensors")
-spatial_upsampler_path = hf_hub_download(repo_id=LTX_MODEL_REPO, filename="ltx-2.3-spatial-upscaler-x2-1.1.safetensors")
+checkpoint_path = hf_hub_download(repo_id=LTX_MODEL_REPO, filename=DISTILLED_CHECKPOINT)
+spatial_upsampler_path = hf_hub_download(repo_id=LTX_MODEL_REPO, filename=SPATIAL_UPSCALER)
+hdr_lora_path = hf_hub_download(repo_id=HDR_LORA_REPO, filename=HDR_LORA_FILENAME)
 gemma_root = snapshot_download(repo_id=GEMMA_REPO)
 
 print(f"Checkpoint: {checkpoint_path}")
 print(f"Spatial upsampler: {spatial_upsampler_path}")
+print(f"HDR IC-LoRA: {hdr_lora_path}")
 print(f"Gemma root: {gemma_root}")
 
-# Initialize pipeline WITH text encoder
-pipeline = DistilledPipeline(
-    distilled_checkpoint_path=checkpoint_path,
-    spatial_upsampler_path=spatial_upsampler_path,
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Text encoding: on-the-fly Gemma -> (video_context, audio_context) for each
+# prompt. HDRICLoraPipeline expects a `.pt` path at __init__, so we bootstrap
+# one from an empty prompt, then overwrite `pipeline.text_embeddings` in
+# memory each generate call.
+# ─────────────────────────────────────────────────────────────────────────────
+_DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+_DTYPE = torch.bfloat16
+
+prompt_encoder = PromptEncoder(
+    checkpoint_path=checkpoint_path,
     gemma_root=gemma_root,
-    loras=[],
-    quantization=QuantizationPolicy.fp8_cast(),
+    dtype=_DTYPE,
+    device=_DEVICE,
 )
 
-# Preload all models for ZeroGPU tensor packing.
-print("Preloading all models (including Gemma)...")
-ledger = pipeline.model_ledger
-_transformer = ledger.transformer()
-_video_encoder = ledger.video_encoder()
-_video_decoder = ledger.video_decoder()
-_audio_decoder = ledger.audio_decoder()
-_vocoder = ledger.vocoder()
-_spatial_upsampler = ledger.spatial_upsampler()
-_text_encoder = ledger.text_encoder()
-_embeddings_processor = ledger.gemma_embeddings_processor()
-
-ledger.transformer = lambda: _transformer
-ledger.video_encoder = lambda: _video_encoder
-ledger.video_decoder = lambda: _video_decoder
-ledger.audio_decoder = lambda: _audio_decoder
-ledger.vocoder = lambda: _vocoder
-ledger.spatial_upsampler = lambda: _spatial_upsampler
-ledger.text_encoder = lambda: _text_encoder
-ledger.gemma_embeddings_processor = lambda: _embeddings_processor
-print("All models preloaded (including Gemma text encoder)!")
 
-print("=" * 80)
-print("Pipeline ready!")
-print("=" * 80)
+def encode_prompt_to_contexts(prompt: str) -> tuple[torch.Tensor, torch.Tensor]:
+    """Run Gemma + embeddings processor to produce (video_context, audio_context).
 
+    HDRICLoraPipeline only consumes video_context; audio_context is stored for
+    shape-compat with the `.pt` interface but ignored during HDR generation.
+    """
+    (out,) = prompt_encoder([prompt])
+    v = out.video_encoding
+    a = out.audio_encoding if out.audio_encoding is not None else torch.zeros(0, device=v.device, dtype=v.dtype)
+    return v, a
 
-def log_memory(tag: str):
-    if torch.cuda.is_available():
-        allocated = torch.cuda.memory_allocated() / 1024**3
-        peak = torch.cuda.max_memory_allocated() / 1024**3
-        free, total = torch.cuda.mem_get_info()
-        print(f"[VRAM {tag}] allocated={allocated:.2f}GB peak={peak:.2f}GB free={free / 1024**3:.2f}GB total={total / 1024**3:.2f}GB")
+
+# Bootstrap a dummy embeddings file so HDRICLoraPipeline can construct.
+print("Pre-computing bootstrap text embeddings...")
+_bootstrap_v, _bootstrap_a = encode_prompt_to_contexts("")
+_bootstrap_emb_path = Path(tempfile.gettempdir()) / "ltx_hdr_bootstrap_emb.pt"
+torch.save({"video_context": _bootstrap_v, "audio_context": _bootstrap_a}, _bootstrap_emb_path)
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Initialize pipeline
+# ─────────────────────────────────────────────────────────────────────────────
+# HDRICLoraPipeline is video-only (no audio path). HDR transform (LogC3) and
+# reference_downscale_factor are auto-detected from the LoRA metadata.
+pipeline = HDRICLoraPipeline(
+    distilled_checkpoint_path=checkpoint_path,
+    spatial_upsampler_path=spatial_upsampler_path,
+    hdr_lora=hdr_lora_path,
+    text_embeddings_path=str(_bootstrap_emb_path),
+    quantization=QuantizationPolicy.fp8_cast(),
+    offload_mode=OffloadMode.NONE,
+)
+print(f"HDRICLoraPipeline ready. HDR transform: {pipeline.hdr_transform}, "
+      f"ref_downscale={pipeline.reference_downscale_factor}")
+print("=" * 80)
 
 
-def detect_aspect_ratio(image) -> str:
-    """Detect the closest aspect ratio (16:9, 9:16, or 1:1) from an image."""
-    if image is None:
+# ─────────────────────────────────────────────────────────────────────────────
+# UI helpers
+# ─────────────────────────────────────────────────────────────────────────────
+def detect_aspect_ratio(video_path) -> str:
+    if video_path is None:
         return "16:9"
-    if hasattr(image, "size"):
-        w, h = image.size
-    elif hasattr(image, "shape"):
-        h, w = image.shape[:2]
-    else:
+    try:
+        meta = get_videostream_metadata(str(video_path))
+        ratio = meta.width / meta.height
+    except Exception:
         return "16:9"
-    ratio = w / h
-    candidates = {"16:9": 16 / 9, "9:16": 9 / 16, "1:1": 1.0}
+    candidates = {
+        "16:9": 16 / 9, "9:16": 9 / 16, "1:1": 1.0,
+        "4:3": 4 / 3, "3:4": 3 / 4, "21:9": 21 / 9,
+    }
     return min(candidates, key=lambda k: abs(ratio - candidates[k]))
 
 
-def on_image_upload(image, high_res):
-    """Auto-set resolution when image is uploaded."""
-    aspect = detect_aspect_ratio(image)
+def get_target_resolution(target_aspect: str, high_res: bool) -> tuple[int, int]:
     tier = "high" if high_res else "low"
-    w, h = RESOLUTIONS[tier][aspect]
-    return gr.update(value=w), gr.update(value=h)
+    return RESOLUTIONS[tier].get(target_aspect, RESOLUTIONS[tier]["16:9"])
 
 
-def on_highres_toggle(image, high_res):
-    """Update resolution when high-res toggle changes."""
-    aspect = detect_aspect_ratio(image)
-    tier = "high" if high_res else "low"
-    w, h = RESOLUTIONS[tier][aspect]
-    return gr.update(value=w), gr.update(value=h)
-
-
-@spaces.GPU(duration=75)
+def on_video_upload(video):
+    if video is None:
+        return gr.update(), gr.update(), gr.update()
+    try:
+        meta = get_videostream_metadata(str(video))
+        dur = round(min(meta.frames / meta.fps, 6.7), 1)
+        fps = float(meta.fps)
+    except Exception:
+        dur, fps = 3.0, 24.0
+    return gr.update(value=dur), gr.update(value=detect_aspect_ratio(video)), gr.update(value=fps)
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Generation
+# ─────────────────────────────────────────────────────────────────────────────
+@spaces.GPU(duration=300)
 @torch.inference_mode()
 def generate_video(
-    input_image,
+    input_video,
     prompt: str,
     duration: float,
-    enhance_prompt: bool = True,
-    seed: int = 42,
-    randomize_seed: bool = True,
-    height: int = 1024,
-    width: int = 1536,
+    frame_rate: float,
+    target_aspect: str,
+    high_res: bool,
+    seed: int,
+    randomize_seed: bool,
+    high_quality_hdr: bool,
+    export_exr: bool,
     progress=gr.Progress(track_tqdm=True),
 ):
+    current_seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed)
     try:
         torch.cuda.reset_peak_memory_stats()
-        log_memory("start")
 
-        current_seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed)
+        if input_video is None:
+            raise ValueError("Please upload a source video.")
 
-        frame_rate = DEFAULT_FRAME_RATE
+        video_path = str(input_video)
+
+        # Frame count must satisfy (n-1) % 8 == 0
         num_frames = int(duration * frame_rate) + 1
         num_frames = ((num_frames - 1 + 7) // 8) * 8 + 1
 
-        print(f"Generating: {height}x{width}, {num_frames} frames ({duration}s), seed={current_seed}")
+        target_w, target_h = get_target_resolution(target_aspect, high_res)
 
-        images = []
-        if input_image is not None:
-            output_dir = Path("outputs")
-            output_dir.mkdir(exist_ok=True)
-            temp_image_path = output_dir / f"temp_input_{current_seed}.jpg"
-            if hasattr(input_image, "save"):
-                input_image.save(temp_image_path)
-            else:
-                temp_image_path = Path(input_image)
-            images = [ImageConditioningInput(path=str(temp_image_path), frame_idx=0, strength=1.0)]
+        print(f"[HDR] {target_h}x{target_w}, frames={num_frames}, fps={frame_rate}, "
+              f"seed={current_seed}, aspect={target_aspect}, hq_hdr={high_quality_hdr}")
 
-        tiling_config = TilingConfig.default()
-        video_chunks_number = get_video_chunks_number(num_frames, tiling_config)
+        # Encode prompt -> (video_context, audio_context) and swap into the
+        # pipeline. Gemma is loaded, used, and freed inside prompt_encoder.
+        print(f"[HDR] Encoding prompt: {prompt!r}")
+        video_context, audio_context = encode_prompt_to_contexts(prompt or "")
+        pipeline.text_embeddings = (video_context, audio_context)
 
-        log_memory("before pipeline call")
+        # Tiling config: smaller spatial tile on lower-VRAM targets
+        tiling_config = _make_tiling_config(spatial_tile=768 if not high_res else 1280)
 
-        video, audio = pipeline(
-            prompt=prompt,
+        hdr_video = pipeline(
             seed=current_seed,
-            height=int(height),
-            width=int(width),
+            height=int(target_h),
+            width=int(target_w),
             num_frames=num_frames,
-            frame_rate=frame_rate,
-            images=images,
+            frame_rate=float(frame_rate),
+            video_conditioning=[(video_path, 1.0)],
             tiling_config=tiling_config,
-            enhance_prompt=enhance_prompt,
+            high_quality_hdr=high_quality_hdr,
         )
 
-        log_memory("after pipeline call")
+        # hdr_video is [f, h, w, c] linear HDR float. Write EXR frames, then
+        # tonemap to a libx264 mp4 for in-browser preview.
+        work_dir = Path(tempfile.mkdtemp(prefix="ltx-hdr-"))
+        exr_dir = work_dir / "exr"
+        exr_dir.mkdir(parents=True, exist_ok=True)
+
+        for j in range(hdr_video.shape[0]):
+            save_exr_tensor(
+                hdr_video[j].detach().cpu(),
+                str(exr_dir / f"frame_{j:05d}.exr"),
+                half=False,
+            )
+        del hdr_video
+
+        preview_mp4 = work_dir / "preview.mp4"
+        encode_exr_sequence_to_mp4(exr_dir, preview_mp4, frame_rate=float(frame_rate))
 
-        output_path = tempfile.mktemp(suffix=".mp4")
-        encode_video(
-            video=video,
-            fps=frame_rate,
-            audio=audio,
-            output_path=output_path,
-            video_chunks_number=video_chunks_number,
-        )
+        exr_zip_path = None
+        if export_exr:
+            exr_zip_path = work_dir / "exr.zip"
+            with zipfile.ZipFile(exr_zip_path, "w", compression=zipfile.ZIP_STORED) as zf:
+                for exr in sorted(exr_dir.glob("frame_*.exr")):
+                    zf.write(exr, arcname=exr.name)
 
-        log_memory("after encode_video")
-        return str(output_path), current_seed
+        return str(preview_mp4), (str(exr_zip_path) if exr_zip_path else None), current_seed
 
     except Exception as e:
         import traceback
-        log_memory("on error")
-        print(f"Error: {str(e)}\n{traceback.format_exc()}")
-        return None, current_seed
+        print(f"Error: {e}\n{traceback.format_exc()}")
+        return None, None, current_seed
 
 
-with gr.Blocks(title="LTX-2.3 Distilled") as demo:
-    gr.Markdown("# LTX-2.3 Distilled (22B): Fast Audio-Video Generation")
-    gr.Markdown(
-        "Fast and high quality video + audio generation "
-        "[[model]](https://huggingface.co/Lightricks/LTX-2.3) "
-        "[[code]](https://github.com/Lightricks/LTX-2)"
-    )
+# ─────────────────────────────────────────────────────────────────────────────
+# Gradio UI — LTX 2.3 HDR
+# ─────────────────────────────────────────────────────────────────────────────
+css = """
+.generate-btn { min-height: 52px !important; font-size: 1.1em !important; }
+footer { display: none !important; }
+video { object-fit: contain !important; }
+"""
+
+theme = gr.themes.Citrus(
+    primary_hue=gr.themes.colors.purple,
+    secondary_hue=gr.themes.colors.purple,
+    neutral_hue=gr.themes.colors.gray,
+)
+
+with gr.Blocks(title="LTX 2.3 HDR", css=css, theme=theme) as demo:
+    gr.Markdown("""
+# LTX 2.3 HDR ✨
+Video-to-video HDR via LTX-2.3 + [HDR IC-LoRA](https://huggingface.co/diffusers-internal-dev/LTX-HDR-LoRA).
+Output is linear HDR (LogC3 inverse decoded — auto-detected from LoRA metadata). The preview mp4 is a fixed-EV sRGB tonemap; the EXR zip contains the full linear float frames for grading.
+    """)
 
     with gr.Row():
-        with gr.Column():
-            input_image = gr.Image(label="Input Image (Optional)", type="pil")
+        with gr.Column(scale=1):
+            input_video = gr.Video(label="Source Video")
+
             prompt = gr.Textbox(
                 label="Prompt",
-                info="for best results - make it as elaborate as possible",
-                value="Make this image come alive with cinematic motion, smooth animation",
-                lines=3,
-                placeholder="Describe the motion and animation you want...",
+                info="Describes the scene being regenerated in HDR. Encoded through Gemma on each run.",
+                lines=2,
+                placeholder="a cinematic sunset over mountains, high dynamic range, bright sky, deep shadows",
             )
-            
-            with gr.Row():
-                duration = gr.Slider(label="Duration (seconds)", minimum=1.0, maximum=10.0, value=3.0, step=0.1)
-                with gr.Column():
-                    enhance_prompt = gr.Checkbox(label="Enhance Prompt", value=False)
-                    high_res = gr.Checkbox(label="High Resolution", value=True)
 
-            generate_btn = gr.Button("Generate Video", variant="primary", size="lg")
+            with gr.Row():
+                target_aspect = gr.Dropdown(
+                    label="Aspect Ratio",
+                    choices=["16:9", "9:16", "1:1", "4:3", "3:4", "21:9"],
+                    value="16:9",
+                )
+                duration = gr.Slider(
+                    label="Duration (s)", minimum=1.0, maximum=6.7, value=3.0, step=0.1,
+                )
+                frame_rate = gr.Number(label="FPS", value=24.0, precision=2)
+
+            generate_btn = gr.Button(
+                "Generate HDR", variant="primary", size="lg", elem_classes=["generate-btn"],
+            )
 
             with gr.Accordion("Advanced Settings", open=False):
+                high_res = gr.Checkbox(label="High Resolution (2×)", value=False)
+                high_quality_hdr = gr.Checkbox(
+                    label="High-Quality HDR",
+                    value=False,
+                    info="Generates at 2× internal frame count, keeps every other frame. ~2× slower, smoother temporally.",
+                )
+                export_exr = gr.Checkbox(
+                    label="Export EXR frames (zip)",
+                    value=True,
+                    info="Per-frame linear OpenEXR (float32). Uncheck if you only want the mp4 preview.",
+                )
                 seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, value=10, step=1)
                 randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)
-                with gr.Row():
-                    width = gr.Number(label="Width", value=1536, precision=0)
-                    height = gr.Number(label="Height", value=1024, precision=0)
-
-        with gr.Column():
-            output_video = gr.Video(label="Generated Video", autoplay=True)
-
-    # Auto-detect aspect ratio from uploaded image and set resolution
-    input_image.change(
-        fn=on_image_upload,
-        inputs=[input_image, high_res],
-        outputs=[width, height],
-    )
 
-    # Update resolution when high-res toggle changes
-    high_res.change(
-        fn=on_highres_toggle,
-        inputs=[input_image, high_res],
-        outputs=[width, height],
+        with gr.Column(scale=1):
+            output_video = gr.Video(label="Preview (sRGB tonemap)", autoplay=True, height=480)
+            output_exr = gr.File(label="EXR frames (zip)")
+
+    input_video.change(
+        fn=on_video_upload,
+        inputs=[input_video],
+        outputs=[duration, target_aspect, frame_rate],
     )
 
     generate_btn.click(
         fn=generate_video,
         inputs=[
-            input_image, prompt, duration, enhance_prompt,
-            seed, randomize_seed, height, width,
+            input_video, prompt, duration, frame_rate, target_aspect, high_res,
+            seed, randomize_seed, high_quality_hdr, export_exr,
         ],
-        outputs=[output_video, seed],
+        outputs=[output_video, output_exr, seed],
     )
 
 
-css = """
-.fillable{max-width: 1200px !important}
-.progress-text {color: white}
-"""
-
 if __name__ == "__main__":
-    demo.launch(theme=gr.themes.Citrus(), css=css)
-
+    demo.launch()

ltx-2-internal/LICENSE

@@ -0,0 +1,381 @@
+                         LTX-2 Community License Agreement
+                           License date: January 5, 2026
+
+
+By using or distributing any portion or element of LTX-2, you agree
+to be bound by this Agreement.
+
+   1. Definitions.
+
+      "Agreement" means the terms and conditions for the license, use,
+      reproduction, and distribution of LTX-2 and the Complementary
+      Materials, as specified in this document.
+
+      "Control" means the direct or indirect ownership of more than
+      fifty percent (50%) of the voting securities or other ownership
+      interests, or the power to direct the management and policies of
+      such Entity through voting rights, contract, or otherwise.
+
+      "Data" means a collection of information and/or content extracted
+      from the dataset used with LTX-2, including to train, pretrain,
+      or otherwise evaluate LTX-2. The Data is not licensed under this
+      Agreement.
+
+      "Derivatives of LTX-2" means all modifications to LTX-2, works
+      based on LTX-2, or any other model which is created or initialized
+      by transfer of patterns of the weights, parameters, activations or
+      output of LTX-2, to the other model, in order to cause the other
+      model to perform similarly to LTX-2, including – but not limited
+      to - distillation methods entailing the use of intermediate data
+      representations or methods based on the generation of synthetic
+      data by LTX-2 for training the other model. For clarity, Derivatives
+      of LTX-2 include: (i) any fine-tuned or adapted weights, parameters,
+      or checkpoints derived from LTX-2; (ii) derivative model architectures
+      that incorporate or are based upon LTX-2's architecture; and
+      (iii) any modified or extended versions of the Complementary
+      Materials. All intellectual property rights in Derivatives of LTX-2
+      shall be subject to the terms of this Agreement, and you may not
+      claim exclusive ownership rights in any Derivatives of LTX-2 that
+      would restrict the rights granted herein.
+
+      "Entity" means any individual, corporation, partnership, limited
+      liability company, or other legal entity. For purposes of this
+      Agreement, an Entity shall be deemed to include, on an aggregative
+      basis, all subsidiaries, affiliates, and other companies under
+      common Control with such Entity. When determining whether an Entity
+      meets any threshold under this Agreement (including revenue
+      thresholds), all subsidiaries, affiliates, and companies under
+      common Control shall be considered collectively.
+
+      "Harm" includes but is not limited to physical, mental,
+      psychological, financial and reputational damage, pain, or loss.
+
+      "Licensor" or "Lightricks" means the owner that is granting the
+      license under this Agreement. For the purposes of this Agreement,
+      the Licensor is Lightricks Ltd.
+
+      "LTX-2" means the large language models, text/image/video/audio/3D
+      generation models, and multimodal large language models and their
+      software and algorithms, including trained model weights, parameters
+      (including optimizer states), machine-learning model code,
+      inference-enabling code, training-enabling code, fine-tuning
+      enabling code, accompanying source code, scripts, documentation,
+      tutorials, examples, and all other elements of the foregoing
+      distributed and made publicly available by Lightricks (including,
+      for example, at https://github.com/Lightricks/LTX-2) for the LTX-2
+      model released on January 5, 2026. This license is applicable to
+      all LTX-2 versions released since January 5, 2026, and all future
+      releases of LTX-2 under this license.
+
+      "Output" means the results of operating LTX-2 as embodied in
+      informational content resulting therefrom.
+
+      "you" (or "your") means an individual or legal Entity licensing
+      LTX-2 in accordance with this Agreement and/or making use of LTX-2
+      for whichever purpose and in any field of use, including usage of
+      LTX-2 in an end-use application - e.g. chatbot, translator, image
+      generator.
+
+   2. Grant of License. Subject to the terms and conditions of this
+      Agreement, you are granted a non-exclusive, worldwide,
+      non-transferable and royalty-free limited license under Licensor's
+      intellectual property or other rights owned by Licensor embodied
+      in LTX-2 to use, reproduce, prepare, distribute, publicly display,
+      publicly perform, sublicense, copy, create derivative works of,
+      and make modifications to LTX-2, for any purpose, subject to the
+      restrictions set forth in Attachment A; provided however, that
+      Entities with annual revenues of at least $10,000,000 (the
+      "Commercial Entities") are required to obtain a paid commercial
+      use license in order to use LTX-2 and Derivatives of LTX-2,
+      subject to the terms and provisions of a different license (the
+      "Commercial Use Agreement"), as will be provided by the Licensor.
+      Commercial Entities interested in such a commercial license are
+      required to [contact Licensor](https://ltx.io/model/licensing).
+      Any commercial use of LTX-2 or Derivatives of LTX-2 by the
+      Commercial Entities not in accordance with this Agreement and/or
+      the Commercial Use Agreement is strictly prohibited and shall be
+      deemed a material breach of this Agreement. Such material breach
+      will be subject, in addition to any license fees owed to Licensor
+      for the period such Commercial Entity used LTX-2 (as will be
+      determined by Licensor), to liquidated damages, which will be paid
+      to Licensor immediately upon demand, in an amount equal to double
+      the amount that would otherwise have been paid by you for the
+      relevant period of time. Such amount reflects a reasonable estimation
+      of the losses and administrative costs incurred due to such breach.
+      You agree and understand that this remedy does not limit the Licensor's
+      right to pursue other remedies available at law or equity.
+
+   3. Distribution and Redistribution. You may host for third parties
+      remote access purposes (e.g. software-as-a-service), reproduce
+      and distribute copies of LTX-2 or Derivatives of LTX-2 thereof in
+      any medium, with or without modifications, provided that you meet
+      the following conditions:
+
+      (a) Use-based restrictions as referenced in paragraph 4 and all
+          provisions of Attachment A MUST be included as an enforceable
+          provision by you in any type of legal agreement (e.g. a
+          license) governing the use and/or distribution of LTX-2 or
+          Derivatives of LTX-2, and you shall give notice to subsequent
+          users you distribute to, that LTX-2 or Derivatives of LTX-2
+          are subject to paragraph 4 and Attachment A in their entirety,
+          including all use restrictions and acceptable use policies;
+
+      (b) You must provide any third party recipients of LTX-2 or
+          Derivatives of LTX-2 a copy of this Agreement, including all
+          attachments and use policies. Any Derivative of LTX-2 (as
+          defined in Section 1, including but not limited to fine-tuned
+          weights, modified training code, models trained on Outputs, or
+          any other derivative) must be distributed exclusively under
+          the terms of this Agreement with a complete copy of this
+          license included;
+
+      (c) You must cause any modified files to carry prominent notices
+          stating that you changed the files;
+
+      (d) You must retain all copyright, patent, trademark, and
+          attribution notices excluding those notices that do not
+          pertain to any part of LTX-2, Derivatives of LTX-2.
+
+      You may add your own copyright statement to your modifications and
+      may provide additional or different license terms and conditions -
+      respecting paragraph 3(a) - for use, reproduction, or distribution
+      of your modifications, or for any such Derivatives of LTX-2 as a
+      whole, provided your use, reproduction, and distribution of LTX-2
+      otherwise complies with the conditions stated in this Agreement,
+      and you provide a complete copy of this Agreement with any such
+      use, reproduction and distribution of LTX-2 and any Derivatives
+      thereof.
+
+   4. Use-based restrictions. The restrictions set forth in Attachment A
+      are considered Use-based restrictions. Therefore, you cannot use
+      LTX-2 and the Derivatives of LTX-2 in violation of the specified
+      restricted uses. You may use LTX-2 subject to this Agreement,
+      including only for lawful purposes and in accordance with the
+      Agreement. "Use" may include creating any content with, fine-tuning,
+      updating, running, training, evaluating and/or re-parametrizing
+      LTX-2. You shall require all of your users who use LTX-2 or a
+      Derivative of LTX-2 to comply with the terms of this paragraph 4.
+
+   5. The Output You Generate. Except as set forth herein, Licensor
+      claims no rights in the Output you generate using LTX-2. You are
+      accountable for input you insert into LTX-2, the Output you
+      generate and its subsequent uses. No use of the Output can
+      contravene any provision as stated in the Agreement.
+
+   6. Updates and Runtime Restrictions. To the maximum extent permitted
+      by law, Licensor reserves the right to restrict (remotely or
+      otherwise) usage of LTX-2 in violation of this Agreement, update
+      LTX-2 through electronic means, or modify the Output of LTX-2
+      based on updates. You shall undertake reasonable efforts to use
+      the latest version of LTX-2. Any use of the non-current version
+      of LTX-2 is done solely at your risk.
+
+   7. Export Controls and Sanctions Compliance. You acknowledge that
+      LTX-2, Derivatives of LTX-2 may be subject to export control laws
+      and regulations, including but not limited to the U.S. Export
+      Administration Regulations and sanctions programs administered by
+      the Office of Foreign Assets Control (OFAC). You represent and
+      warrant that you and any users of LTX-2 are not (i) located in,
+      organized under the laws of, or ordinarily resident in any country
+      or territory subject to comprehensive sanctions; (ii) identified
+      on any U.S. government restricted party list, including the
+      Specially Designated Nationals and Blocked Persons List; or
+      (iii) otherwise prohibited from receiving LTX-2 under applicable
+      law. You shall not export, re-export, or transfer LTX-2, directly
+      or indirectly, in violation of any applicable export control or
+      sanctions laws or regulations. You agree to comply with all
+      applicable trade control laws and shall indemnify and hold
+      Licensor harmless from any claims arising from your failure to
+      comply with such laws.
+
+   8. Trademarks and related. Nothing in this Agreement permits you to
+      make use of Licensor's trademarks, trade names, logos or to
+      otherwise suggest endorsement or misrepresent the relationship
+      between the parties; and any rights not expressly granted herein
+      are reserved by the Licensor.
+
+   9. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides LTX-2 on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or
+      conditions of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS
+      FOR A PARTICULAR PURPOSE. You are solely responsible for
+      determining the appropriateness of using or redistributing LTX-2
+      and Derivatives of LTX-2 and assume any risks associated with
+      your exercise of permissions under this Agreement.
+
+   10. Limitation of Liability. In no event and under no legal theory,
+       whether in tort (including negligence), contract, or otherwise,
+       unless required by applicable law (such as deliberate and grossly
+       negligent acts) or agreed to in writing, shall Licensor be liable
+       to you for damages, including any direct, indirect, special,
+       incidental, or consequential damages of any character arising as
+       a result of this Agreement or out of the use or inability to use
+       LTX-2 (including but not limited to damages for loss of goodwill,
+       work stoppage, computer failure or malfunction, or any and all
+       other commercial damages or losses), even if Licensor has been
+       advised of the possibility of such damages.
+
+   11. Accepting Warranty or Additional Liability. While redistributing
+       LTX-2 and Derivatives of LTX-2, you may, provided you do not
+       violate the terms of this Agreement, choose to offer and charge
+       a fee for, acceptance of support, warranty, indemnity, or other
+       liability obligations. However, in accepting such obligations,
+       you may act only on your own behalf and on your sole
+       responsibility, not on behalf of Licensor, and only if you agree
+       to indemnify, defend, and hold Licensor harmless for any liability
+       incurred by, or claims asserted against Licensor, by reason of
+       your accepting any such warranty or additional liability.
+
+   12. Governing Law. This Agreement and all relations, disputes, claims
+       and other matters arising hereunder (including non-contractual
+       disputes or claims) will be governed exclusively by, and construed
+       exclusively in accordance with, the laws of the State of New York.
+       To the extent permitted by law, choice of laws rules and the
+       United Nations Convention on Contracts for the International Sale
+       of Goods will not apply. For the purposes of adjudicating any
+       action or proceeding to enforce the terms of this Agreement, you
+       hereby irrevocably consent to the exclusive jurisdiction of, and
+       venue in, the federal and state courts located in the County of
+       New York within the State of New York. The prevailing party in
+       any claim or dispute between the parties under this Agreement
+       will be entitled to reimbursement of its reasonable attorneys'
+       fees and costs. You hereby waive the right to a trial by jury,
+       to participate in a class or representative action (including in
+       arbitration), or to combine individual proceedings in court or
+       in arbitration without the consent of all parties.
+
+   13. Term and Termination. This Agreement is effective upon your
+       acceptance and continues until terminated. Licensor may terminate
+       this Agreement immediately upon written notice to you if you
+       breach any provision of this Agreement, including but not limited
+       to violations of the use restrictions in Attachment A or
+       unauthorized commercial use. Upon termination: (a) all rights
+       granted to you under this Agreement will immediately cease;
+       (b) you must immediately cease all use of LTX-2 and Derivatives
+       of LTX-2; (c) you must delete or destroy all copies of LTX-2
+       and Derivatives of LTX-2 in your possession or control; and
+       (d) you must notify any third parties to whom you distributed
+       LTX-2 or Derivatives of LTX-2 of the termination. Sections 8-13,
+       and Section 15 shall survive termination of this Agreement.
+       Termination does not relieve you of any obligations incurred
+       prior to termination, including payment obligations under
+       Section 2. In addition, if You commence a lawsuit or other
+       proceedings (including a cross-claim or counterclaim in a lawsuit)
+       against Licensor or any person or entity alleging that LTX-2 or
+       any Output, or any portion of any of the foregoing, infringe any
+       intellectual property or other right owned or licensable by you,
+       then all licenses granted to you under this Agreement shall
+       terminate as of the date such lawsuit or other proceeding is filed.
+
+   14. Disputes and Arbitration. All disputes arising in connection with
+       this Agreement shall be finally settled by arbitration under the
+       Rules of Arbitration of the International Chamber of Commerce
+       ("ICC Rules"), by one (1) arbitrator appointed in accordance with
+       the ICC Rules. The seat of arbitration shall be New York, NY, USA,
+       and the proceedings shall be conducted in English. The arbitrator
+       shall be empowered to grant any relief that a court could grant.
+       Judgment on the arbitration award may be entered by any court
+       having jurisdiction thereof. Each party waives its right to a
+       trial by jury and to participate in any class or representative
+       action.
+
+   15. If any provision of this Agreement is held to be
+       invalid, illegal
+       or unenforceable, the remaining provisions shall be unaffected
+       thereby and remain valid as if such provision had not been set
+       forth herein.
+
+   END OF TERMS AND CONDITIONS
+
+   ATTACHMENT A: Use Restrictions
+
+      When using the Outputs, LTX-2 and any Derivatives thereof, you
+      will comply with the Acceptable Use Policy. In addition, you
+      agree not to use the Outputs, LTX-2 or its Derivatives in any
+      of the following ways:
+
+      1. In any way that violates any applicable national, federal,
+         state, local or international law or regulation;
+
+      2. For the purpose of exploiting, Harming or attempting to
+         exploit or Harm minors in any way;
+
+      3. To generate or disseminate false information and/or content
+         with the purpose of Harming others;
+
+      4. To generate or disseminate personal identifiable information
+         that can be used to Harm an individual;
+
+      5. To generate or disseminate information and/or content (e.g.
+         images, code, posts, articles), and place the information
+         and/or content in any context (e.g. bot generating tweets)
+         without expressly and intelligibly disclaiming that the
+         information and/or content is machine generated;
+
+      6. To defame, disparage or otherwise harass others;
+
+      7. To impersonate or attempt to impersonate (e.g. deepfakes)
+         others without their consent;
+
+      8. For fully automated decision making that adversely impacts an
+         individual's legal rights or otherwise creates or modifies a
+         binding, enforceable obligation;
+
+      9. For any use intended to or which has the effect of
+         discriminating against or Harming individuals or groups based
+         on online or offline social behavior or known or predicted
+         personal or personality characteristics;
+
+      10. To exploit any of the vulnerabilities of a specific group of
+          persons based on their age, social, physical or mental
+          characteristics, in order to materially distort the behavior
+          of a person pertaining to that group in a manner that causes
+          or is likely to cause that person or another person physical
+          or psychological Harm;
+
+      11. For any use intended to or which has the effect of
+          discriminating against individuals or groups based on legally
+          protected characteristics or categories;
+
+      12. To provide medical advice and medical results interpretation;
+
+      13. To generate or disseminate information for the purpose to be
+          used for administration of justice, law enforcement,
+          immigration or asylum processes, such as predicting an
+          individual will commit fraud/crime commitment (e.g. by text
+          profiling, drawing causal relationships between assertions
+          made in documents, indiscriminate and arbitrarily-targeted use);
+
+      14. To generate and/or disseminate malware (including – but not
+          limited to – ransomware) or any other content to be used for
+          the purpose of harming electronic systems;
+
+      15. To engage in, promote, incite, or facilitate discrimination
+          or other unlawful or harmful conduct in the provision of
+          employment, employment benefits, credit, housing, or other
+          essential goods and services;
+
+      16. To engage in, promote, incite, or facilitate the harassment,
+          abuse, threatening, or bullying of individuals or groups of
+          individuals;
+
+      17. For military, warfare, nuclear industries or applications,
+          weapons development, or any use in connection with activities
+          that may cause death, personal injury, or severe physical or
+          environmental damage;
+
+      18. For commercial use only: To train, improve, or fine-tune any
+          other machine learning model, artificial intelligence system,
+          or competing model, except for Derivatives of LTX-2 as
+          expressly permitted under this Agreement;
+
+      19. To circumvent, disable, or interfere with any technical
+          limitations, safety features, content filters, or use
+          restrictions implemented in LTX-2 by Licensor;
+
+      20. To use LTX-2 or Derivatives of LTX-2 in any product, service,
+          or application that directly competes with Licensor's
+          commercial products or services, or is designed to replace or
+          substitute Licensor's offerings in the market, without
+          obtaining a separate commercial license from Licensor.

ltx-2-internal/README.md

@@ -0,0 +1,66 @@
+# LTX-2
+
+
+https://github.com/user-attachments/assets/1922e7f0-cc5d-430c-9e1a-93c589e85fb4
+
+
+## Setup
+
+To set up ltx-2 repository for development:
+* `uv sync` creates and installs the virtual environment. To activate the environment, run `source .venv/bin/activate`.
+* Set up pre-commit:
+```bash
+$ PIP_NO_INPUT=1 pre-commit install-hooks --config .pre-commit-config.yaml
+$ pre-commit install
+```
+* `pytest` runs all tests in the repository. `pytest --pyargs packages/ltx-core` runs only tests specific to `ltx-core`.
+* `uv build --package ltx-core` builds only the ltx-core package.
+* Dependencies declared in the main `pyproject.toml` file apply to the entire repository. Always prefer specifying dependencies for specific packages.
+
+## Folder structure
+
+```
+├── README.md
+├── packages
+│   └── ltx-core
+│       ├── README.md
+│       ├── pyproject.toml
+│       ├── src
+│       │   └── ltx_core
+│       │       ├── __init__.py
+│       │       ├── # package-specific modules go here
+│       └── tests
+│           └── ltx_core
+│               ├── # test folder structure mirrors src of the package.
+│               ├── conftest.py
+│               └── unit
+│                   └── test_placeholder.py
+├── pyproject.toml
+├── pytest.toml
+└── uv.lock
+```
+
+
+## Comments mentioning ComfyUI or internal company code
+
+Comments relating to ComfyUI, its LTX repo or any other internal
+Lightricks resource should be wrapped in LT_INTERNAL tags. Tagged
+sections should be placed inside comments or docstrings. The
+following examples demonstrate proper usage:
+```python
+# LT_INTERNAL: check PR to understand better
+
+# conditioning is applied before the step LT_INTERNAL: http://...
+
+# condition is applied after the step to integrate LT_INTERNAL_BEGIN http://... LT_INTERNAL_END
+
+"""
+Consult with the paper, https://arxiv.org
+LT_INTERNAL_BEGIN
+Most important PRs to check
+http://...
+http://...
+LT_INTERNAL_END
+"""
+
+```

ltx-2-internal/README.public.md

@@ -0,0 +1,121 @@
+# LTX-2
+
+[![Website](https://img.shields.io/badge/Website-LTX-181717?logo=google-chrome)](https://ltx.io)
+[![Model](https://img.shields.io/badge/HuggingFace-Model-orange?logo=huggingface)](https://huggingface.co/Lightricks/LTX-2.3)
+[![Demo](https://img.shields.io/badge/Demo-Try%20Now-brightgreen?logo=data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAABQAAAAUCAYAAACNiR0NAAAAAXNSR0IArs4c6QAAAERlWElmTU0AKgAAAAgAAYdpAAQAAAABAAAAGgAAAAAAA6ABAAMAAAABAAEAAKACAAQAAAABAAAAFKADAAQAAAABAAAAFAAAAACy3fD9AAACmElEQVQ4Ea1VP2haYRA/fRo0mESRIIqb2IwxuNUl0CGFQBC6OAWcikMottCpqYtDQIgdQsBFhAjZqiQhbhmySJBOgmNU0EGCg9r61Bivd0ffoykE0iQH37/77n7f3e/uqQFIPB7P/N3d3QeDwfAFEedZ91ghnyH5JM1m87dWq6UavF6vdTKZfDcajW/p4rE49+wIFMj33Gq1vlNo+kxg758KpiETqP/29vaXweVyqaS0aBfPXEfGFwTjWCwM+KBQoWA4HAJx/KDNvxcmTTGbzYAH8SljOp2C2+2GjY0NqNfrcHFxAXNzc2LDfCuKIq78KBdFOwsgGzidTnA4HHBzcwO9Xg8sFgtsbm7C3t4eVCoVaDQa0O12YXl5GUwmk5z5cZ/PB6PRCNrttgADFQUXFhbw8PAQVVXF3d1dJAeMx+P0zn0Jh8OYz+eRADCRSGAqlcLxeIz7+/u4tLSEjKUDZrNZ8U4mk0jR4fr6Op6enoru+voa0+k0rq2tYTAYxE6ng9QiSLRgrVZDv9+PFLkA6kUhT+GEC8C8XF5ewtHRkejICShiaDabwPvj42NJm3k7ODiQdDl9Fr0ocqJpdXUVIpEIdz7Y7XZRr6ysQDQahXK5LORvbW1p5rC9vQ2UifAooBqHuVxO0vt72tnZwWq1qqtisRgWCgU5ZzIZPDk50fdUUEmZvxTmAgKBgAxunT/fJpRKJWmhUCgEVDi4uroSG46kWCzC4uKitNVgMICzszOhSgA5fiJZhp4Lbbh1KARpbF65D/lx3vMdP05Vlkf5zKIDyukFJi7N6AVwNAhVsdlsM+LsjaZ56sq8kyQUqs4P6rsAKV49B4x4Padf7Y9Kv9+fEmiBQH8S4Gsa5v8EHpL9VwL7xH8BvwEcd4ccVf02KQAAAABJRU5ErkJggg==)](https://console.ltx.video/playground)
+[![Paper](https://img.shields.io/badge/Paper-PDF-EC1C24?logo=adobeacrobatreader&logoColor=white)](https://arxiv.org/abs/2601.03233)
+[![Discord](https://img.shields.io/badge/Join-Discord-5865F2?logo=discord)](https://discord.gg/ltxplatform)
+
+**LTX-2** is the first DiT-based audio-video foundation model that contains all core capabilities of modern video generation in one model: synchronized audio and video, high fidelity, multiple performance modes, production-ready outputs, API access, and open access.
+
+<div align="center">
+  <video src="https://github.com/user-attachments/assets/4414adc0-086c-43de-b367-9362eeb20228" width="70%" poster=""> </video>
+</div>
+
+## 🚀 Quick Start
+
+```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 following 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)
+  * [`ltx-2.3-22b-distilled-1.1.safetensors`](https://huggingface.co/Lightricks/LTX-2.3/blob/main/ltx-2.3-22b-distilled-1.1.safetensors) - [Download](https://huggingface.co/Lightricks/LTX-2.3/resolve/main/ltx-2.3-22b-distilled-1.1.safetensors)
+
+**Spatial Upscaler** - Required for current two-stage pipeline implementations in this repository
+  * [`ltx-2.3-spatial-upscaler-x2-1.1.safetensors`](https://huggingface.co/Lightricks/LTX-2.3/blob/main/ltx-2.3-spatial-upscaler-x2-1.1.safetensors) - [Download](https://huggingface.co/Lightricks/LTX-2.3/resolve/main/ltx-2.3-spatial-upscaler-x2-1.1.safetensors)
+  * [`ltx-2.3-spatial-upscaler-x1.5-1.0.safetensors`](https://huggingface.co/Lightricks/LTX-2.3/blob/main/ltx-2.3-spatial-upscaler-x1.5-1.0.safetensors) - [Download](https://huggingface.co/Lightricks/LTX-2.3/resolve/main/ltx-2.3-spatial-upscaler-x1.5-1.0.safetensors)
+
+**Temporal Upscaler** - Supported by the model and will be required for future pipeline implementations
+  * [`ltx-2.3-temporal-upscaler-x2-1.0.safetensors`](https://huggingface.co/Lightricks/LTX-2.3/blob/main/ltx-2.3-temporal-upscaler-x2-1.0.safetensors) - [Download](https://huggingface.co/Lightricks/LTX-2.3/resolve/main/ltx-2.3-temporal-upscaler-x2-1.0.safetensors)
+
+**Distilled LoRA** - Required for current two-stage pipeline implementations in this repository (except DistilledPipeline and ICLoraPipeline)
+  * [`ltx-2.3-22b-distilled-lora-384-1.1.safetensors`](https://huggingface.co/Lightricks/LTX-2.3/blob/main/ltx-2.3-22b-distilled-lora-384-1.1.safetensors) - [Download](https://huggingface.co/Lightricks/LTX-2.3/resolve/main/ltx-2.3-22b-distilled-lora-384-1.1.safetensors)
+
+**Gemma Text Encoder** (download all assets from the repository)
+  * [`Gemma 3`](https://huggingface.co/google/gemma-3-12b-it-qat-q4_0-unquantized/tree/main)
+
+**LoRAs**
+  * [`LTX-2.3-22b-IC-LoRA-Union-Control`](https://huggingface.co/Lightricks/LTX-2.3-22b-IC-LoRA-Union-Control) - [Download](https://huggingface.co/Lightricks/LTX-2.3-22b-IC-LoRA-Union-Control/resolve/main/ltx-2.3-22b-ic-lora-union-control-ref0.5.safetensors)
+  * [`LTX-2.3-22b-IC-LoRA-Motion-Track-Control`](https://huggingface.co/Lightricks/LTX-2.3-22b-IC-LoRA-Motion-Track-Control) - [Download](https://huggingface.co/Lightricks/LTX-2.3-22b-IC-LoRA-Motion-Track-Control/resolve/main/ltx-2.3-22b-ic-lora-motion-track-control-ref0.5.safetensors)
+  * [`LTX-2-19b-IC-LoRA-Detailer`](https://huggingface.co/Lightricks/LTX-2-19b-IC-LoRA-Detailer) - [Download](https://huggingface.co/Lightricks/LTX-2-19b-IC-LoRA-Detailer/resolve/main/ltx-2-19b-ic-lora-detailer.safetensors)
+  * [`LTX-2-19b-IC-LoRA-Pose-Control`](https://huggingface.co/Lightricks/LTX-2-19b-IC-LoRA-Pose-Control) - [Download](https://huggingface.co/Lightricks/LTX-2-19b-IC-LoRA-Pose-Control/resolve/main/ltx-2-19b-ic-lora-pose-control.safetensors)
+  * [`LTX-2-19b-LoRA-Camera-Control-Dolly-In`](https://huggingface.co/Lightricks/LTX-2-19b-LoRA-Camera-Control-Dolly-In) - [Download](https://huggingface.co/Lightricks/LTX-2-19b-LoRA-Camera-Control-Dolly-In/resolve/main/ltx-2-19b-lora-camera-control-dolly-in.safetensors)
+  * [`LTX-2-19b-LoRA-Camera-Control-Dolly-Left`](https://huggingface.co/Lightricks/LTX-2-19b-LoRA-Camera-Control-Dolly-Left) - [Download](https://huggingface.co/Lightricks/LTX-2-19b-LoRA-Camera-Control-Dolly-Left/resolve/main/ltx-2-19b-lora-camera-control-dolly-left.safetensors)
+  * [`LTX-2-19b-LoRA-Camera-Control-Dolly-Out`](https://huggingface.co/Lightricks/LTX-2-19b-LoRA-Camera-Control-Dolly-Out) - [Download](https://huggingface.co/Lightricks/LTX-2-19b-LoRA-Camera-Control-Dolly-Out/resolve/main/ltx-2-19b-lora-camera-control-dolly-out.safetensors)
+  * [`LTX-2-19b-LoRA-Camera-Control-Dolly-Right`](https://huggingface.co/Lightricks/LTX-2-19b-LoRA-Camera-Control-Dolly-Right) - [Download](https://huggingface.co/Lightricks/LTX-2-19b-LoRA-Camera-Control-Dolly-Right/resolve/main/ltx-2-19b-lora-camera-control-dolly-right.safetensors)
+  * [`LTX-2-19b-LoRA-Camera-Control-Jib-Down`](https://huggingface.co/Lightricks/LTX-2-19b-LoRA-Camera-Control-Jib-Down) - [Download](https://huggingface.co/Lightricks/LTX-2-19b-LoRA-Camera-Control-Jib-Down/resolve/main/ltx-2-19b-lora-camera-control-jib-down.safetensors)
+  * [`LTX-2-19b-LoRA-Camera-Control-Jib-Up`](https://huggingface.co/Lightricks/LTX-2-19b-LoRA-Camera-Control-Jib-Up) - [Download](https://huggingface.co/Lightricks/LTX-2-19b-LoRA-Camera-Control-Jib-Up/resolve/main/ltx-2-19b-lora-camera-control-jib-up.safetensors)
+  * [`LTX-2-19b-LoRA-Camera-Control-Static`](https://huggingface.co/Lightricks/LTX-2-19b-LoRA-Camera-Control-Static) - [Download](https://huggingface.co/Lightricks/LTX-2-19b-LoRA-Camera-Control-Static/resolve/main/ltx-2-19b-lora-camera-control-static.safetensors)
+  * [`LTX-2.3-22b-IC-LoRA-HDR`](https://huggingface.co/Lightricks/LTX-2.3-22b-IC-LoRA-HDR) - HDR IC-LoRA and pre-computed text embeddings for `HDRICLoraPipeline`
+
+### Available Pipelines
+
+* **[TI2VidTwoStagesPipeline](packages/ltx-pipelines/src/ltx_pipelines/ti2vid_two_stages.py)** - Production-quality text/image-to-video with 2x upsampling (recommended)
+* **[TI2VidTwoStagesHQPipeline](packages/ltx-pipelines/src/ltx_pipelines/ti2vid_two_stages_hq.py)** - Same two-stage flow as above but uses the res_2s second-order sampler (fewer steps, better quality)
+* **[TI2VidOneStagePipeline](packages/ltx-pipelines/src/ltx_pipelines/ti2vid_one_stage.py)** - Single-stage generation for quick prototyping
+* **[DistilledPipeline](packages/ltx-pipelines/src/ltx_pipelines/distilled.py)** - Fastest inference with 8 predefined sigmas
+* **[ICLoraPipeline](packages/ltx-pipelines/src/ltx_pipelines/ic_lora.py)** - Video-to-video and image-to-video transformations (uses distilled model.)
+* **[KeyframeInterpolationPipeline](packages/ltx-pipelines/src/ltx_pipelines/keyframe_interpolation.py)** - Interpolate between keyframe images
+* **[A2VidPipelineTwoStage](packages/ltx-pipelines/src/ltx_pipelines/a2vid_two_stage.py)** - Audio-to-video generation conditioned on an input audio file
+* **[RetakePipeline](packages/ltx-pipelines/src/ltx_pipelines/retake.py)** - Regenerate a specific time region of an existing video
+* **[HDRICLoraPipeline](packages/ltx-pipelines/src/ltx_pipelines/hdr_ic_lora.py)** - Video-to-video with HDR output (linear float frames via LogC3 inverse decode, suitable for EXR export and tonemapping)
+
+### ⚡ Optimization Tips
+
+* **Use DistilledPipeline** - Fastest inference with only 8 predefined sigmas (8 steps stage 1, 4 steps stage 2)
+* **Enable FP8 quantization** - Enables lower memory footprint: `--quantization fp8-cast` (CLI) or `quantization=QuantizationPolicy.fp8_cast()` (Python). Fp8-cast should be used with bf16 checkpoints, it shall downcast them on the fly. For Hopper GPUs with TensorRT-LLM, use `--quantization fp8-scaled-mm` for FP8 scaled matrix multiplication. Fp8-scaled-mm should be used with fp8 checkpoints.
+* **Install attention optimizations** - Use xFormers (`uv sync --extra xformers`) or [Flash Attention 3](https://github.com/Dao-AILab/flash-attention) for Hopper GPUs
+* **Use gradient estimation** - Reduce inference steps from 40 to 20-30 while maintaining quality (see [pipeline documentation](packages/ltx-pipelines/README.md#denoising-loop-optimization))
+* **Skip memory cleanup** - If you have sufficient VRAM, disable automatic memory cleanup between stages for faster processing
+* **Choose single-stage pipeline** - Use `TI2VidOneStagePipeline` for faster generation when high resolution isn't required
+
+## ✍️ Prompting for LTX-2
+
+When writing prompts, focus on detailed, chronological descriptions of actions and scenes. Include specific movements, appearances, camera angles, and environmental details - all in a single flowing paragraph. Start directly with the action, and keep descriptions literal and precise. Think like a cinematographer describing a shot list. Keep within 200 words. For best results, build your prompts using this structure:
+
+- Start with main action in a single sentence
+- Add specific details about movements and gestures
+- Describe character/object appearances precisely
+- Include background and environment details
+- Specify camera angles and movements
+- 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>
+
+### Automatic Prompt Enhancement
+
+LTX-2 pipelines support automatic prompt enhancement via an `enhance_prompt` parameter.
+
+## 🔌 ComfyUI Integration
+
+To use our model with ComfyUI, please follow the instructions at <https://github.com/Lightricks/ComfyUI-LTXVideo/>.
+
+## 📦 Packages
+
+This repository is organized as a monorepo with three main packages:
+
+* **[ltx-core](packages/ltx-core/)** - Core model implementation, inference stack, and utilities
+* **[ltx-pipelines](packages/ltx-pipelines/)** - High-level pipeline implementations for text-to-video, image-to-video, and other generation modes
+* **[ltx-trainer](packages/ltx-trainer/)** - Training and fine-tuning tools for LoRA, full fine-tuning, and IC-LoRA
+
+Each package has its own README and documentation. See the [Documentation](#-documentation) section below.
+
+## 📚 Documentation
+
+Each package includes comprehensive documentation:
+
+* **[LTX-Core README](packages/ltx-core/README.md)** - Core model implementation, inference stack, and utilities
+* **[LTX-Pipelines README](packages/ltx-pipelines/README.md)** - High-level pipeline implementations and usage guides
+* **[LTX-Trainer README](packages/ltx-trainer/README.md)** - Training and fine-tuning documentation with detailed guides

ltx-2-internal/conftest.py

@@ -0,0 +1,217 @@
+from __future__ import annotations
+
+import gc
+import os
+from pathlib import Path
+from typing import TYPE_CHECKING, Callable
+
+import av
+import pytest
+import torch
+import torch.nn.functional as F
+from torch._prims_common import DeviceLikeType
+
+if TYPE_CHECKING:
+    from ltx_core.guidance.perturbations import BatchedPerturbationConfig
+    from ltx_core.model.transformer import Modality
+
+torch.use_deterministic_algorithms(True)
+os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
+
+MODELS_PATH = Path(os.getenv("MODELS_PATH", "/models"))
+CHECKPOINTS_DIR = MODELS_PATH / "comfyui_models" / "checkpoints"
+LORAS_DIR = MODELS_PATH / "comfyui_models" / "loras"
+
+GEMMA_ROOT = MODELS_PATH / "comfyui_models" / "text_encoders" / "gemma-3-12b-it-qat-q4_0-unquantized_readout_proj"
+GEMMA_FLATTENED_ROOT = MODELS_PATH / "comfyui_models" / "text_encoders" / "gemma-3-12b-it-qat-q4_0-unquantized"
+
+LTX_2_0_CHECKPOINT_PATH = CHECKPOINTS_DIR / "ltx-2-19b-dev.safetensors"
+LTX_2_0_CHECKPOINT_FP8_PATH = CHECKPOINTS_DIR / "ltx-2-19b-dev-fp8.safetensors"
+LTX_2_3_CHECKPOINT_PATH = CHECKPOINTS_DIR / "ltx-2.3-22b-dev.safetensors"
+LTX_2_0_DISTILLED_CHECKPOINT_PATH = CHECKPOINTS_DIR / "ltx-2-19b-distilled.safetensors"
+LTX_2_3_DISTILLED_CHECKPOINT_PATH = CHECKPOINTS_DIR / "ltx-2.3-22b-distilled.safetensors"
+
+LTX_2_0_SPATIAL_UPSAMPLER_PATH = CHECKPOINTS_DIR / "ltx2-spatial-upscaler-x2-1.0.bf16.safetensors"
+LTX_2_3_SPATIAL_UPSAMPLER_PATH = (
+    MODELS_PATH / "comfyui_models" / "latent_upscale_models" / "ltx-2.3-spatial-upscaler-x2-1.0.safetensors"
+)
+
+LTX_2_0_DISTILLED_LORA_PATH = LORAS_DIR / "ltxv" / "ltx2" / "ltx-av-distilled-from-42500-lora-384_comfy.safetensors"
+LTX_2_3_DISTILLED_LORA_PATH = LORAS_DIR / "ltxv" / "ltx2" / "ltx-2.3-22b-distilled-lora-384.safetensors"
+LTX_2_3_IC_LORA_PATH = LORAS_DIR / "ltxv" / "ltx2" / "ltx-2.3-22b-ic-lora-union-control-ref0.5.safetensors"
+LTX_2_3_HDR_IC_LORA_PATH = LORAS_DIR / "ltxv" / "ltx2" / "ltx-2.3-22b-ic-lora-hdr.safetensors"
+LTX_2_3_HDR_TEXT_EMBEDDINGS_PATH = LORAS_DIR / "ltxv" / "ltx2" / "ltx-2.3-22b-ic-lora-hdr-scene-emb.pt"
+
+
+def _psnr(pred: torch.Tensor, target: torch.Tensor, max_val: float = 1.0, eps: float = 1e-8) -> torch.Tensor:
+    """
+    Compute Peak Signal-to-Noise Ratio (PSNR) between two images (or batches of images).
+
+    Args:
+        pred:   Predicted image tensor, shape (..., H, W) or (..., C, H, W)
+        target: Ground truth image tensor, same shape as `pred`
+        max_val: Maximum possible pixel value of the images.
+                 For images in [0, 1] use 1.0, for [0, 255] use 255.0, etc.
+        eps:    Small value to avoid log of zero.
+
+    Returns:
+        psnr: PSNR value (in dB).
+    """
+    if pred.shape != target.shape:
+        raise ValueError(f"Shape mismatch: pred {pred.shape}, target {target.shape}")
+
+    pred = pred.float()
+    target = target.float()
+
+    if pred.dim() > 1:
+        mse = F.mse_loss(pred, target, reduction="none")
+        dims = list(range(mse.dim()))
+        mse = mse.mean(dim=dims)
+    else:
+        mse = F.mse_loss(pred, target, reduction="mean")
+
+    psnr_val = 10.0 * torch.log10((max_val**2) / (mse + eps))
+    return psnr_val
+
+
+def _psnr_per_frame(pred: torch.Tensor, target: torch.Tensor, max_val: float = 1.0, eps: float = 1e-8) -> torch.Tensor:
+    """
+    Compute per-frame PSNR between two video tensors.
+
+    Args:
+        pred:   Predicted video tensor, shape (T, C, H, W) or higher
+        target: Ground truth video tensor, same shape as `pred`
+        max_val: Maximum possible pixel value.
+        eps:    Small value to avoid log of zero.
+
+    Returns:
+        psnr: Per-frame PSNR values (in dB).
+    """
+    if pred.shape != target.shape:
+        raise ValueError(f"Shape mismatch: pred {pred.shape}, target {target.shape}")
+
+    pred = pred.float()
+    target = target.float()
+
+    if pred.dim() < 4:
+        raise ValueError("Expected at least 4D tensor (T, C, H, W) for per-frame PSNR.")
+
+    mse = F.mse_loss(pred, target, reduction="none")
+    mse = mse.mean(dim=(-3, -2, -1))
+    psnr_val = 10.0 * torch.log10((max_val**2) / (mse + eps))
+    return psnr_val
+
+
+def _decode_video_from_file(path: str, device: DeviceLikeType) -> tuple[torch.Tensor, torch.Tensor | None]:
+    container = av.open(path)
+    try:
+        video_stream = next(s for s in container.streams if s.type == "video")
+        audio_stream = next((s for s in container.streams if s.type == "audio"), None)
+
+        frames = []
+        audio = [] if audio_stream else None
+
+        streams_to_decode = [video_stream]
+        if audio_stream:
+            streams_to_decode.append(audio_stream)
+
+        for frame in container.decode(*streams_to_decode):
+            if isinstance(frame, av.VideoFrame):
+                tensor = torch.tensor(frame.to_rgb().to_ndarray(), dtype=torch.uint8, device=device).unsqueeze(0)
+                frames.append(tensor)
+            elif isinstance(frame, av.AudioFrame):
+                audio.append(torch.tensor(frame.to_ndarray(), dtype=torch.float32, device=device).unsqueeze(0))
+
+        if audio:
+            audio = torch.cat(audio)
+    finally:
+        container.close()
+
+    return torch.cat(frames), audio
+
+
+class MockTransformer:
+    """Mock transformer that tracks calls and returns the input latents.
+
+    Records ``call_count`` and ``batch_sizes`` for assertions. Output is
+    simply ``video.latent`` / ``audio.latent`` (identity), which is
+    deterministic and split-safe for adapter equivalence tests.
+    """
+
+    def __init__(self) -> None:
+        self.call_count = 0
+        self.batch_sizes: list[int] = []
+
+    @staticmethod
+    def _validate_modality(modality: Modality) -> None:
+        """Check that all tensor fields share the same batch dimension.
+
+        The real transformer's ``_prepare_timestep`` does
+        ``view(batch_size, ...)`` on sigma embeddings, which crashes if
+        sigma's batch dim doesn't match latent's.  This check catches
+        such mismatches without needing real model weights.
+        """
+        batch_size = modality.latent.shape[0]
+        for name, tensor in [
+            ("sigma", modality.sigma),
+            ("timesteps", modality.timesteps),
+            ("positions", modality.positions),
+            ("context", modality.context),
+        ]:
+            if tensor.shape[0] != batch_size:
+                msg = f"Modality.{name} batch dim {tensor.shape[0]} != latent batch dim {batch_size}"
+                raise ValueError(msg)
+
+    def __call__(
+        self,
+        video: Modality | None = None,
+        audio: Modality | None = None,
+        perturbations: BatchedPerturbationConfig | None = None,  # noqa: ARG002
+    ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
+        self.call_count += 1
+        if video is not None:
+            self._validate_modality(video)
+        if audio is not None:
+            self._validate_modality(audio)
+        b = video.latent.shape[0] if video is not None else audio.latent.shape[0]
+        self.batch_sizes.append(b)
+        return (
+            video.latent if video is not None else None,
+            audio.latent if audio is not None else None,
+        )
+
+
+@pytest.fixture
+def mock_transformer() -> MockTransformer:
+    """Return a fresh MockTransformer instance."""
+    return MockTransformer()
+
+
+@pytest.fixture(autouse=True)
+def _cleanup_cuda_memory() -> None:
+    """Free GPU memory before and after each test to prevent OOM across test modules."""
+    gc.collect()
+    torch.cuda.empty_cache()
+
+    yield
+
+    gc.collect()
+    torch.cuda.empty_cache()
+
+
+@pytest.fixture
+def psnr() -> Callable[[torch.Tensor, torch.Tensor, float, float], torch.Tensor]:
+    """Fixture that returns the PSNR function."""
+    return _psnr
+
+
+@pytest.fixture
+def psnr_per_frame() -> Callable[[torch.Tensor, torch.Tensor, float, float], torch.Tensor]:
+    """Fixture that returns the per-frame PSNR function."""
+    return _psnr_per_frame
+
+
+@pytest.fixture
+def decode_video_from_file() -> Callable[[str, DeviceLikeType], tuple[torch.Tensor, torch.Tensor | None]]:
+    """Fixture that returns the function to decode a video from a file."""
+    return _decode_video_from_file

ltx-2-internal/internal/ltx-kernels/MANIFEST.in

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

ltx-2-internal/internal/ltx-kernels/README.md

@@ -0,0 +1,43 @@
+# ltx-kernels
+
+Custom CUDA/C++ All2All communication kernels for multi-GPU tensor parallelism. Used by the sequence parallel inference path in `ltx-core`.
+
+## Requirements
+
+- CUDA toolkit (nvcc) matching your GPU architecture
+- PyTorch with CUDA support
+- Linux
+
+## Building
+
+From the repository root:
+
+```bash
+uv pip install -e internal/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 internal/ltx-kernels --no-build-isolation
+
+# Multiple architectures
+TORCH_CUDA_ARCH_LIST="9.0 9.0a 10.0 12.0" uv pip install -e internal/ltx-kernels --no-build-isolation
+```
+
+## Testing
+
+Tests require a CUDA GPU:
+
+```bash
+uv run pytest internal/ltx-kernels/tests/ -v
+```
+
+## Operations
+
+- **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.

ltx-2-internal/internal/ltx-kernels/csrc/all2all/all2all.cpp

@@ -0,0 +1,408 @@
+/**
+ * @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)
+    : 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);
+  // 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);
+
+  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);
+
+  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);
+
+  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)")
+      .def(pybind11::init<int, int, int, int, int, at::ScalarType>())
+      .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.");
+}

ltx-2-internal/internal/ltx-kernels/csrc/all2all/all2all.hpp

@@ -0,0 +1,231 @@
+/**
+ * @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 <pybind11/pybind11.h>
+#include <pybind11/pytypes.h>
+#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
+
+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)
+   */
+  All2All(int rank, int world_size, int num_tokens, int hidden_dim, int num_sms, at::ScalarType tensor_dtype);
+
+  /**
+   * @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);
+};
+
+} // namespace all2all
+} // namespace ltx_kernels

ltx-2-internal/internal/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) {
+  // 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,
+                                    NUM_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) {
+  // 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,
+                                    NUM_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) {
+  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);
+    } 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);
+    }
+
+    // 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) {
+  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);
+    } 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);
+    }
+
+    // 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

ltx-2-internal/internal/ltx-kernels/csrc/all2all/cuda/allgather.cu

@@ -0,0 +1,197 @@
+/**
+ * @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) {
+
+  // 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,
+                                    NUM_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) {
+  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);
+    } 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);
+    } 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

ltx-2-internal/internal/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);
+
+/**
+ * @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);
+
+/**
+ * @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);
+
+} // namespace all2all_cuda
+} // namespace all2all
+} // namespace ltx_kernels

ltx-2-internal/internal/ltx-kernels/csrc/include/cuda/configs.cuh

@@ -0,0 +1,90 @@
+/**
+ * @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 Timeout for barrier synchronization in GPU clock cycles.
+ *
+ * If a barrier wait exceeds this timeout, the kernel will trap to indicate
+ * a deadlock or communication failure. The value ~25 billion cycles corresponds
+ * to approximately 10-100 seconds depending on GPU clock speed.
+ *
+ * Example: At 2.5 GHz boost clock, 25B cycles ≈ 10 seconds.
+ */
+constexpr uint64_t NUM_TIMEOUT_CYCLES = 25000000000ull;
+
+// =============================================================================
+// 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

ltx-2-internal/internal/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

ltx-2-internal/internal/ltx-kernels/csrc/include/cuda/utils.cuh

@@ -0,0 +1,356 @@
+/**
+ * @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 (use NUM_TIMEOUT_CYCLES)
+ */
+__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) {
+  // Ensure all threads in the block have completed their work
+  __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) {
+      bool is_recv_full = ld_volatile_global(barrier_signal_ptrs[rank] + thread_id) == expected_count;
+      if (is_recv_full) {
+        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, ld_volatile_global(barrier_signal_ptrs[rank] + thread_id));
+        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) {
+  // Ensure all threads in the block have completed their work
+  __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) {
+      bool is_recv_full = ld_volatile_global(barrier_signal_ptrs[rank] + thread_id) == expected_from_each_source;
+      if (is_recv_full) {
+        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, ld_volatile_global(barrier_signal_ptrs[rank] + thread_id));
+        trap();
+      }
+    }
+    // Reset barrier for next use
+    atomicSub_system(barrier_signal_ptrs[rank] + thread_id, expected_from_each_source);
+  }
+}
+
+} // namespace ltx_kernels

ltx-2-internal/internal/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

ltx-2-internal/internal/ltx-kernels/pyproject.toml

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

ltx-2-internal/internal/ltx-kernels/setup.py

@@ -0,0 +1,39 @@
+from pathlib import Path
+
+import setuptools
+from torch.utils.cpp_extension import BuildExtension, CUDAExtension
+
+
+def _nvidia_include_dirs() -> list[str]:
+    """Include dirs from pip-installed nvidia packages (e.g. cusparse headers)."""
+    try:
+        import nvidia  # noqa: PLC0415
+
+        return [str(p) for pkg in Path(nvidia.__path__[0]).iterdir() if (p := pkg / "include").is_dir()]
+    except ImportError:
+        return []
+
+
+if __name__ == "__main__":
+    cxx_flags = ["-O3", "-Wall", "-Wextra", "-Werror", "-Wno-unused-parameter", "-Wno-attributes"]
+    nvcc_flags = ["-O3"]
+    extra_compile_args = {
+        "cxx": cxx_flags,
+        "nvcc": nvcc_flags,
+    }
+    ROOT = Path(__file__).resolve().parent
+    setuptools.setup(
+        ext_modules=[
+            CUDAExtension(
+                name="all2all_cpp",
+                include_dirs=[str(ROOT / "csrc/all2all"), str(ROOT / "csrc/include"), *_nvidia_include_dirs()],
+                sources=[
+                    "csrc/all2all/all2all.cpp",
+                    "csrc/all2all/cuda/all2all_heads.cu",
+                    "csrc/all2all/cuda/allgather.cu",
+                ],
+                extra_compile_args=extra_compile_args,
+            )
+        ],
+        cmdclass={"build_ext": BuildExtension},
+    )

ltx-2-internal/internal/ltx-kernels/src/ltx_kernels/__init__.py

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

ltx-2-internal/internal/ltx-kernels/src/ltx_kernels/all_to_all.py

@@ -0,0 +1,90 @@
+"""All2All communication primitives for distributed attention."""
+
+from typing import Any
+
+import torch
+import torch.distributed as dist
+from all2all_cpp import All2All as All2AllCpp
+
+
+class All2All:
+    """IPC-based All2All communication for distributed head-parallel attention.
+
+    This class manages GPU memory buffers and IPC handles to enable efficient
+    cross-GPU communication for attention head redistribution.
+
+    Args:
+        rank: Local rank of this process.
+        world_size: Total number of processes in the distributed group.
+        seqlen: Maximum sequence length to allocate buffers for.
+        hidden_dim: Hidden dimension size (num_heads * head_dim).
+        num_sms: Number of SMs to use for kernel execution.
+        tensor_dtype: Data type for tensors (e.g., torch.bfloat16).
+        group: PyTorch distributed process group.
+    """
+
+    def __init__(
+        self,
+        rank: int,
+        world_size: int,
+        seqlen: int,
+        hidden_dim: int,
+        num_sms: int,
+        tensor_dtype: torch.dtype,
+        group: torch.distributed.ProcessGroup | None = None,
+    ) -> None:
+        self.rank = rank
+        self.world_size = world_size
+        self.num_sms = num_sms
+        self.tensor_dtype = tensor_dtype
+        self.buffer_size = int(seqlen * hidden_dim * tensor_dtype.itemsize)
+
+        # Initialize the C++ runtime
+        self.runtime = All2AllCpp(rank, world_size, seqlen, hidden_dim, num_sms, tensor_dtype)
+
+        # Exchange IPC handles across all ranks
+        ipc_handles: list[Any] = [None] * world_size
+        local_ipc_handle = self.runtime.get_local_ipc_handle()
+        dist.all_gather_object(ipc_handles, local_ipc_handle, group)
+
+        self.runtime.sync(ipc_handles)
+
+    def send_recv_heads(self, x: torch.Tensor, *, copy_out: bool = False) -> torch.Tensor:
+        """Exchange attention heads across ranks (All2All pattern).
+
+        Args:
+            x: Input tensor of shape [batch, tokens, heads, head_dim].
+            copy_out: If True, copy result to a new tensor instead of using buffer.
+
+        Returns:
+            Output tensor with redistributed heads.
+        """
+        return self.runtime.send_recv_heads(x, copy_out)
+
+    def gather_heads(self, x: torch.Tensor, *, copy_out: bool = False) -> torch.Tensor:
+        """Gather heads back to original distribution (reverse All2All).
+
+        Args:
+            x: Input tensor with distributed heads.
+            copy_out: If True, copy result to a new tensor instead of using buffer.
+
+        Returns:
+            Output tensor with gathered heads.
+        """
+        return self.runtime.gather_heads(x, copy_out)
+
+    def allgather(self, x: torch.Tensor, *, copy_out: bool = False) -> torch.Tensor:
+        """Allgather operation across all ranks.
+
+        Args:
+            x: Input tensor to gather.
+            copy_out: If True, copy result to a new tensor instead of using buffer.
+
+        Returns:
+            Gathered tensor from all ranks.
+        """
+        return self.runtime.allgather(x, copy_out)
+
+    def destroy(self) -> None:
+        """Release IPC handles and GPU memory buffers."""
+        self.runtime.destroy()

ltx-2-internal/internal/test-txt2img-parity/conftest.py

@@ -0,0 +1,85 @@
+"""Session-level import stubs for txt2img JAX modules.
+
+Why this exists
+---------------
+``txt2img.jax.models.transformers.attention_block`` (which defines
+``SelfAttention`` and ``CrossAttention``) has two top-level imports that fail
+in CPU/GPU environments:
+
+1. ``from ltx_video.utils.skip_layer_strategy import SkipLayerStrategy``
+   ``ltx_video`` is a Lightricks-internal package that is not on PyPI and is
+   not installed in this environment.  Its ``SkipLayerStrategy`` enum is only
+   used in ``SelfAttention.__call__`` for STG skip-layer strategies, which are
+   never exercised by the parity tests.
+
+2. ``txt2img.jax.models.transformers.attention.__init__`` unconditionally
+   imports ``TokamaxSplashAttention``, which depends on ``tokamax`` — a
+   Google-internal TPU kernel library unavailable on CPU/GPU.  The splash
+   attention path is only reachable when ``attention_op="tokamax_splash"``;
+   parity tests always use ``attention_op="explicit"``.
+
+Python executes every line of a module file before exposing any of its
+contents, so even though ``SelfAttention`` itself never touches these paths,
+the file fails to load and nothing inside it is importable.
+
+The fix: inject minimal ``sys.modules`` entries *before* any test module is
+imported.  pytest loads ``conftest.py`` first, so the stubs are in place when
+the test files run their top-level imports.  The stubs are never called at
+runtime — they exist solely to satisfy the module loader.
+"""
+
+import sys
+import types
+from enum import Enum, auto
+
+
+# ---------------------------------------------------------------------------
+# 1. ltx_video.utils.skip_layer_strategy
+# ---------------------------------------------------------------------------
+class SkipLayerStrategy(Enum):
+    AttentionSkip = auto()
+    AttentionValues = auto()
+    Residual = auto()
+    TransformerBlock = auto()
+
+
+_ltx_video = types.ModuleType("ltx_video")
+_ltx_video_utils = types.ModuleType("ltx_video.utils")
+_ltx_video_sls = types.ModuleType("ltx_video.utils.skip_layer_strategy")
+_ltx_video_sls.SkipLayerStrategy = SkipLayerStrategy  # type: ignore[attr-defined]
+
+sys.modules.setdefault("ltx_video", _ltx_video)
+sys.modules.setdefault("ltx_video.utils", _ltx_video_utils)
+sys.modules.setdefault("ltx_video.utils.skip_layer_strategy", _ltx_video_sls)
+
+
+# ---------------------------------------------------------------------------
+# 2. tokamax — TPU-only splash/ring attention kernels
+# ---------------------------------------------------------------------------
+class _SplashConfig:
+    pass
+
+
+class _SplashKernel:
+    SplashConfig = _SplashConfig
+
+
+class _SplashMask:
+    Mask = None
+
+
+_tpu_splash = types.ModuleType("tokamax._src.ops.experimental.tpu.splash_attention")
+_tpu_splash.splash_attention_kernel = _SplashKernel()  # type: ignore[attr-defined]
+_tpu_splash.ring_attention_kernel = None  # type: ignore[attr-defined]
+_tpu_splash.splash_attention_mask = _SplashMask()  # type: ignore[attr-defined]
+
+for _name in [
+    "tokamax",
+    "tokamax._src",
+    "tokamax._src.ops",
+    "tokamax._src.ops.experimental",
+    "tokamax._src.ops.experimental.tpu",
+]:
+    sys.modules.setdefault(_name, types.ModuleType(_name))
+
+sys.modules.setdefault("tokamax._src.ops.experimental.tpu.splash_attention", _tpu_splash)

ltx-2-internal/internal/test-txt2img-parity/pyproject.toml

@@ -0,0 +1,38 @@
+[project]
+name = "test-txt2img-parity"
+version = "0.1.0"
+description = "Parity tests verifying PyTorch LTX-2 matches the txt2img JAX reference implementation"
+requires-python = ">=3.10"
+dependencies = [
+    "ltx-core",
+    "jax",
+    "flax",
+    "txt2img",
+    "pydantic",
+    "google-cloud-storage",
+    "diffusers",
+    "pytest",
+]
+
+[build-system]
+requires = ["uv_build>=0.9.8,<0.10.0"]
+build-backend = "uv_build"
+
+[tool.pytest.ini_options]
+addopts = ["--import-mode=importlib"]
+pythonpath = ["src"]
+
+[tool.uv.sources]
+ltx-core = { path = "../../packages/ltx-core", editable = true }
+txt2img = { git = "https://github.com/LightricksResearch/txt2img", branch = "main" }
+
+# txt2img pins requires-python = "==3.12.12" but works on >=3.10 — relaxed here.
+# txt2img base deps include local-only editable packages (ltx-video, dataset_metadata, lttensor)
+# that don't exist on PyPI. The actual ML deps (jax, flax, etc.) are in dependency groups,
+# not base deps, so they're never installed automatically. We strip everything and declare
+# exactly what we need explicitly.
+[[tool.uv.dependency-metadata]]
+name = "txt2img"
+version = "0.1.0"
+requires-python = ">=3.10"
+requires-dist = []

ltx-2-internal/internal/test-txt2img-parity/src/test_txt2img_parity/testing_facilities.py

@@ -0,0 +1,282 @@
+import logging
+import time
+from collections.abc import Callable, Generator
+from contextlib import ExitStack, contextmanager
+from copy import deepcopy
+from functools import wraps
+from typing import Any
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+import numpy.typing as npt
+import torch
+import torch.utils._pytree as pytree
+from jax import grad
+from txt2img.jax.utils.torch_compat import jax_statedict_to_torch
+
+Shape = tuple[int, ...]
+ShapeDict = dict[str, Shape]
+NumpyTensorDict = dict[str, npt.NDArray]
+JaxTensorTree = jax.Array | dict[str, jax.Array]
+TorchTensorTree = torch.Tensor | dict[str, torch.Tensor]
+TorchFunction = Callable[[TorchTensorTree], torch.Tensor]
+JaxFunction = Callable[[JaxTensorTree], jax.Array]
+
+
+@contextmanager
+def torch_default_dtype(enable_x64: bool) -> Generator[None, None, None]:
+    previous_dtype = torch.get_default_dtype()
+    new_dtype = torch.float64 if enable_x64 else torch.float32
+    torch.set_default_dtype(new_dtype)
+    try:
+        yield
+    finally:
+        torch.set_default_dtype(previous_dtype)
+
+
+@contextmanager
+def jax_64_precision() -> Generator[None, None, None]:
+    previous_state = jax.config.read("jax_enable_x64")  # Save current state
+    jax.config.update("jax_enable_x64", True)  # Enable x64
+    try:
+        yield
+    finally:
+        jax.config.update("jax_enable_x64", previous_state)  # Restore original state
+
+
+def measure_time(func: Callable[..., Any]) -> Callable[..., tuple[Any, float]]:
+    """Decorator that measures the execution time of a function."""
+
+    @wraps(func)
+    def wrapper(*args: object, **kwargs: object) -> tuple[object, float]:
+        start_time = time.perf_counter()
+        result = func(*args, **kwargs)
+        end_time = time.perf_counter()
+        execution_time = end_time - start_time
+        return result, execution_time
+
+    return wrapper
+
+
+def ensure_tolerance(desc: str, data_torch: npt.NDArray, data_jax: npt.NDArray, atol: float, rtol: float) -> None:
+    """
+    Ensures that the outputs of the torch and jax functions are within the given tolerance.
+
+    Args:
+        desc (str): The tensors description
+        data_torch (npt.NDArray): torch predictions
+        data_jax (npt.NDArray): jax predictions
+        atol (float): absolute tolerance
+        rtol (float): relative tolerance
+    """
+    torch_shape = data_torch.shape
+    jax_shape = data_jax.shape
+    assert torch_shape == jax_shape, f"{desc} Shapes do not match: {torch_shape} != {jax_shape}"
+
+    if data_torch.dtype in (bool, np.bool_):
+        abs_err = np.bitwise_xor(data_jax, data_torch).sum() / np.prod(data_jax.shape)
+        rtol = 0
+        atol = 0
+    else:
+        abs_err = abs(data_torch - data_jax).max()
+
+    passed = np.isclose(data_torch, data_jax, atol=atol, rtol=rtol).all()
+    logger = logging.getLogger(__name__)
+    logger.info(
+        "%s: max_abs_err=%.6e, atol=%.6e, rtol=%.6e, passed=%s",
+        desc,
+        float(abs_err),
+        atol,
+        rtol,
+        passed,
+    )
+    assert passed, f"{desc} Outputs do not match, got absolute error {abs_err}"
+
+
+def _compare_module(
+    data: npt.NDArray | NumpyTensorDict,  # NOTE - some layers may define several inputs, so we use a dictionary
+    torch_function: TorchFunction,
+    jax_function: JaxFunction,
+    atol: float,
+    rtol: float,
+    forward_only: bool,
+) -> None:
+    """
+    Compare a torch module against its jax equivalent.
+
+    Args:
+        data (npt.NDArray | NumpyTensorDict): Dictionary containing the input tensors
+        torch_function (TorchFunction): Torch function to test
+        jax_function (JaxFunction): Jax function to test
+        atol (float): Absolute tolerance.
+        rtol (float): Relative tolerance.
+        forward_only (bool): if True, only the forward pass will be tested
+    """
+    logger = logging.getLogger(__name__)
+
+    # run torch function
+    inputs_torch = pytree.tree_map(torch.from_numpy, deepcopy(data))
+    output_torch, torch_time = measure_time(torch_function)(inputs_torch)
+    logger.info("torch_time: %s seconds", torch_time)
+
+    # Run jax function
+    inputs_jax = pytree.tree_map(jax.numpy.array, deepcopy(data))
+    output_jax, jax_time = measure_time(jax_function)(inputs_jax)
+    logger.info("jax_time: %s seconds", jax_time)
+
+    # Handle few outputs
+    if isinstance(output_torch, tuple):
+        assert isinstance(output_jax, tuple), "torch returned tuple but jax did not"
+        assert len(output_jax) == len(output_torch), (
+            f"torch and jax output counts do not match: {len(output_torch)} != {len(output_jax)}"
+        )
+    else:
+        output_torch, output_jax = (output_torch,), (output_jax,)
+
+    # Convert to numpy
+    for i, (single_output_torch, single_output_jax) in enumerate(
+        zip(output_torch, output_jax, strict=True),
+    ):
+        data_torch = single_output_torch.cpu().detach().numpy()
+
+        cpu_device = jax.devices("cpu")[0]
+        data_jax = np.array(jax.device_put(single_output_jax, cpu_device))
+
+        ensure_tolerance(f"Output {i}", data_torch, data_jax, atol, rtol)
+
+    if not forward_only:
+        cpu_device = jax.devices("cpu")[0]
+
+        def _to_numpy(x: jax.Array) -> npt.NDArray:
+            return np.array(jax.device_put(x, cpu_device))
+
+        # Backward torch
+        num_outputs = len(output_torch)
+        for i in range(num_outputs):
+            inputs_torch = pytree.tree_map(torch.from_numpy, deepcopy(data))
+            pytree.tree_map(lambda x: (x.requires_grad_(), x.retain_grad()), inputs_torch)
+            output_torch, torch_time = measure_time(torch_function)(inputs_torch)
+            if num_outputs > 1:
+                output_torch = output_torch[i]
+            output_torch.sum().backward()
+            grad_torch = pytree.tree_map(lambda x: x.grad.numpy() if x.grad is not None else None, inputs_torch)
+            logger.warning("torch_time (backward) for output %d: %s seconds", i, torch_time)
+
+            inputs_jax = pytree.tree_map(jax.numpy.array, deepcopy(data))
+
+            def _make_loss_func(idx: int, n_outputs: int, jax_fn: JaxFunction) -> Callable[[JaxTensorTree], jax.Array]:
+                def _loss(x: JaxTensorTree) -> jax.Array:
+                    return jnp.sum(jax_fn(x)[idx]) if n_outputs > 1 else jnp.sum(jax_fn(x))
+
+                return _loss
+
+            loss_func = _make_loss_func(i, num_outputs, jax_function)
+
+            grad_jax, jax_time = jax.grad(measure_time(loss_func), has_aux=True)(inputs_jax)
+            logger.warning("jax_time (backward) for output %d: %s seconds", i, jax_time)
+
+            if isinstance(grad_torch, dict):
+                for key, torch_grad in grad_torch.items():
+                    grad_jax_np = _to_numpy(grad_jax[key])
+                    if torch_grad is not None:
+                        ensure_tolerance(f"Grad {key} for output {i}", torch_grad, grad_jax_np, atol, rtol)
+                    else:
+                        assert bool((grad_jax_np == 0).all())
+            else:
+                ensure_tolerance(f"Grad for output {i}", grad_torch, _to_numpy(grad_jax), atol, rtol)
+
+
+def torch_jax_comparison_test(
+    shapes: Shape | ShapeDict,
+    torch_function: TorchFunction,
+    jax_function: JaxFunction,
+    atol: float = 1e-5,
+    rtol: float = 1e-5,
+    seed: int = 42,
+    forward_only: bool = False,
+    use_x64_precision: bool = True,  # Default to True to preserve current behavior
+) -> None:
+    """
+    Facilitates the comparison test between torch and jax functions.
+
+    Args:
+        shapes (Shape | ShapeDict): a mapping between input name and its shape that will be generated
+        torch_function (TorchFunction): The torch function equivalent to test
+        jax_function (JaxFunction): The jax function equivalent to test
+        atol (float, optional): Absolute tolerance of error. Defaults to 1e-5.
+        rtol (float, optional): Relative tolerance of error. Defaults to 1e-5.
+        seed (int): the seed to generate random numpy arrays with
+        forward_only (bool): if True, only the forward pass will be tested
+        use_x64_precision (bool): Flag to enable jax x64 precision mode. Defaults to True.
+    """
+    rng = np.random.default_rng(seed)
+    torch.manual_seed(seed)
+    dtype = np.float64 if use_x64_precision else np.float32
+    if isinstance(shapes, dict):
+        data = {name: rng.random(shape).astype(dtype) for name, shape in shapes.items()}  # Multi input
+    else:
+        data = rng.random(shapes).astype(dtype)  # Single input
+
+    with ExitStack() as stack:
+        if use_x64_precision:
+            stack.enter_context(jax_64_precision())
+            stack.enter_context(torch_default_dtype(use_x64_precision))
+        with jax.default_matmul_precision("highest"):
+            _compare_module(data, torch_function, jax_function, atol, rtol, forward_only)
+
+
+def compare_gradients_by_parameters(
+    torch_module: torch.nn.Module,
+    jax_module: object,
+    shapes: ShapeDict,
+    loss_fn_torch: Callable[[Any], torch.Tensor],
+    loss_fn_jax: Callable[[Any], jax.Array],
+    use_x64_precision: bool,
+    params_jax: dict,
+    atol: float,
+    rtol: float,
+) -> None:
+    def _convert_masks_to_boolean(data: dict[str, Any]) -> dict[str, Any]:
+        if "attention_mask" in data:
+            data["attention_mask"] = data["attention_mask"] >= 0.5
+        return data
+
+    rng = np.random.default_rng(42)
+    torch_module.train()
+    torch_module.zero_grad()
+    dtype = np.float64 if use_x64_precision else np.float32
+    data = {name: rng.random(shape).astype(dtype) for name, shape in shapes.items()}
+    torch_input = pytree.tree_map(torch.from_numpy, deepcopy(data))
+    jax_input = pytree.tree_map(jax.numpy.array, deepcopy(data))
+
+    torch_input = _convert_masks_to_boolean(torch_input)
+    torch_input["encoder_attention_mask"] = torch.ones_like(torch_input.pop("encoder_attention_segment_ids")).to(
+        torch.bool
+    )
+
+    torch_output = torch_module(**torch_input)
+    torch_loss = loss_fn_torch(torch_output)
+    torch_loss.backward()
+
+    def jax_loss(params: dict, inputs: dict) -> jax.Array:
+        outputs = jax_module.apply(params, **inputs)
+        ret = loss_fn_jax(outputs)
+        return ret
+
+    jax_input = _convert_masks_to_boolean(jax_input)
+    jax_input["encoder_attention_segment_ids"] = jnp.ones_like(data["encoder_attention_segment_ids"], dtype=jnp.uint8)
+
+    jax_grads = grad(jax_loss)(params_jax, jax_input)
+    jax_grads_to_torch_statedict = jax_statedict_to_torch(jax_grads)
+
+    for name, param in torch_module.named_parameters():
+        torch_grad = param.grad.detach().numpy() if param.grad is not None else None
+        jax_grad = np.array(jax_grads_to_torch_statedict[name])
+
+        if torch_grad is not None and jax_grad is not None:
+            ensure_tolerance(f"PARAMETER Grad for output {name}", torch_grad, jax_grad, atol, rtol)
+        elif torch_grad is None and not (jax_grad == 0).all():
+            raise ValueError("Mismatch in gradient presence: torch gradient is None while the jax is non-zero.")
+        elif torch_grad is not None and (jax_grad == 0).all():
+            raise ValueError("Mismatch in gradient presence: torch gradient is not None while the jax is zero.")

ltx-2-internal/packages/ltx-bench/README.md

@@ -0,0 +1,65 @@
+# ltx-bench
+
+Profiling and benchmarking tools for LTX-2.
+
+## Running benchmarks
+
+```bash
+# All benchmarks (requires H100 GPU + model checkpoints at /models/)
+uv run pytest -m bench --bench
+
+# With Nsight Systems profiling
+uv run pytest -m bench --bench --nsys
+
+# Custom output directory and iterations
+uv run pytest -m bench --bench \
+  --bench-output=./my_results \
+  --bench-warmup=3 \
+  --bench-repeats=10
+```
+
+## Benchmark types
+
+### Components
+
+Individual model components benchmarked in isolation:
+
+| Benchmark | What it measures |
+|-----------|-----------------|
+| `transformer_forward_video_only` | X0Model forward pass (video only) |
+| `transformer_forward_av` | X0Model forward pass (video + audio) |
+| `video_vae_decode` | Video VAE decoder |
+| `video_vae_encode` | Video VAE encoder |
+| `audio_decode` | Audio VAE decoder + vocoder |
+| `audio_encode` | Audio VAE encoder |
+| `text_encoder_forward` | Gemma text encoder (positive + negative prompt) |
+| `spatial_upsampler` | Latent spatial upsampler 2x |
+
+### Pipelines
+
+End-to-end pipeline benchmarks with per-stage instrumentation:
+
+| Benchmark | Pipeline |
+|-----------|----------|
+| `distilled` | 8-step distilled inference |
+| `ti2vid_one_stage` | Text-to-video one-stage |
+| `ti2vid_two_stages` | Text-to-video two-stage with upsampling |
+| `keyframe_interpolation` | Keyframe interpolation two-stage |
+| `ic_lora` | IC-LoRA with depth conditioning |
+| `a2vid_two_stage` | Audio-to-video two-stage |
+| `retake` | Video retake with regeneration |
+
+## Trend analysis
+
+Generate trend charts from GCS historical data:
+
+```bash
+python -m ltx_bench.trends --branch main --last 20 --output ./trend_report
+```
+
+## Report output
+
+Each run produces a `REPORT.md` with:
+- Per-component GPU time, peak VRAM, and peak RAM
+- Per-pipeline total time with stage breakdown (text encoding, denoising, VAE decode, etc.)
+- Denoising throughput in s/it

ltx-2-internal/packages/ltx-bench/benchmarks/components/conftest.py

@@ -0,0 +1,140 @@
+"""Model builders, re-exports, and helpers for component benchmarks.
+
+Component benchmarks build models explicitly via ``SingleGPUModelBuilder``
+rather than going through ``ModelLedger`` (which is a pipeline-level abstraction).
+Shared constants (paths, shape helpers) are re-exported from the parent conftest.
+"""
+
+from __future__ import annotations
+
+import torch
+
+# Re-export from parent conftest for test files
+from benchmarks.conftest import (  # noqa: F401
+    CHECKPOINT_PATH,
+    DISTILLED_LORA_PATH,
+    GEMMA_ROOT_PATH,
+    IC_LORA_PATH,
+    UPSAMPLER_PATH,
+    make_audio_latent_state,
+    make_noised_states,
+    make_video_latent_state,
+)
+from ltx_core.loader.single_gpu_model_builder import SingleGPUModelBuilder as Builder
+from ltx_core.model.audio_vae import (
+    AUDIO_VAE_DECODER_COMFY_KEYS_FILTER,
+    AUDIO_VAE_ENCODER_COMFY_KEYS_FILTER,
+    VOCODER_COMFY_KEYS_FILTER,
+    AudioDecoder,
+    AudioDecoderConfigurator,
+    AudioEncoder,
+    AudioEncoderConfigurator,
+    Vocoder,
+    VocoderConfigurator,
+)
+from ltx_core.model.model_protocol import ModelType
+from ltx_core.model.transformer import LTXV_MODEL_COMFY_RENAMING_MAP, LTXModelConfigurator, X0Model
+from ltx_core.model.upsampler import LatentUpsampler, LatentUpsamplerConfigurator
+from ltx_core.model.video_vae import (
+    VAE_DECODER_COMFY_KEYS_FILTER,
+    VAE_ENCODER_COMFY_KEYS_FILTER,
+    VideoDecoder,
+    VideoDecoderConfigurator,
+    VideoEncoder,
+    VideoEncoderConfigurator,
+)
+from ltx_core.text_encoders.gemma import (
+    GEMMA_LLM_KEY_OPS,
+    GemmaTextEncoderConfigurator,
+    module_ops_from_gemma_root,
+)
+from ltx_core.text_encoders.gemma.encoders.encoder_configurator import GEMMA_MODEL_OPS
+from ltx_core.utils import find_matching_file
+
+_DTYPE = torch.bfloat16
+_DEVICE = "cuda"
+
+
+def _build(builder: Builder[ModelType]) -> ModelType:
+    return builder.build(device=torch.device(_DEVICE), dtype=_DTYPE).to(_DEVICE).eval()
+
+
+def build_transformer(checkpoint_path: str) -> X0Model:
+    builder = Builder(
+        model_path=checkpoint_path,
+        model_class_configurator=LTXModelConfigurator,
+        model_sd_ops=LTXV_MODEL_COMFY_RENAMING_MAP,
+    )
+    return X0Model(_build(builder)).eval()
+
+
+def build_video_decoder(checkpoint_path: str) -> VideoDecoder:
+    builder = Builder(
+        model_path=checkpoint_path,
+        model_class_configurator=VideoDecoderConfigurator,
+        model_sd_ops=VAE_DECODER_COMFY_KEYS_FILTER,
+    )
+    return _build(builder)
+
+
+def build_video_encoder(checkpoint_path: str) -> VideoEncoder:
+    builder = Builder(
+        model_path=checkpoint_path,
+        model_class_configurator=VideoEncoderConfigurator,
+        model_sd_ops=VAE_ENCODER_COMFY_KEYS_FILTER,
+    )
+    return _build(builder)
+
+
+def build_audio_decoder(checkpoint_path: str) -> AudioDecoder:
+    builder = Builder(
+        model_path=checkpoint_path,
+        model_class_configurator=AudioDecoderConfigurator,
+        model_sd_ops=AUDIO_VAE_DECODER_COMFY_KEYS_FILTER,
+    )
+    return _build(builder)
+
+
+def build_audio_encoder(checkpoint_path: str) -> AudioEncoder:
+    builder = Builder(
+        model_path=checkpoint_path,
+        model_class_configurator=AudioEncoderConfigurator,
+        model_sd_ops=AUDIO_VAE_ENCODER_COMFY_KEYS_FILTER,
+    )
+    return _build(builder)
+
+
+def build_vocoder(checkpoint_path: str) -> Vocoder:
+    builder = Builder(
+        model_path=checkpoint_path,
+        model_class_configurator=VocoderConfigurator,
+        model_sd_ops=VOCODER_COMFY_KEYS_FILTER,
+    )
+    return _build(builder)
+
+
+def build_text_encoder(checkpoint_path: str, gemma_root_path: str) -> torch.nn.Module:
+    if not gemma_root_path:
+        msg = "gemma_root_path is required to build text encoder"
+        raise ValueError(msg)
+    module_ops = module_ops_from_gemma_root(gemma_root_path)
+    model_folder = find_matching_file(gemma_root_path, "model*.safetensors").parent
+    weight_paths = [str(p) for p in model_folder.rglob("*.safetensors")]
+    builder = Builder(
+        model_path=(checkpoint_path, *weight_paths),
+        model_class_configurator=GemmaTextEncoderConfigurator,
+        model_sd_ops=GEMMA_LLM_KEY_OPS,
+        module_ops=(GEMMA_MODEL_OPS, *module_ops),
+    )
+    return _build(builder)
+
+
+def build_upsampler(upsampler_path: str) -> LatentUpsampler:
+    if not upsampler_path:
+        msg = "upsampler_path is required to build upsampler"
+        raise ValueError(msg)
+    builder = Builder(
+        model_path=upsampler_path,
+        model_class_configurator=LatentUpsamplerConfigurator,
+    )
+    return _build(builder)

ltx-2-internal/packages/ltx-bench/benchmarks/components/test_audio_vae.py

@@ -0,0 +1,119 @@
+"""Benchmark: audio VAE encoder, decoder, and vocoder.
+
+Measures GPU time for the audio latent <-> spectrogram path (encoder/decoder)
+and spectrogram -> waveform (vocoder). Audio shapes depend only on video
+duration (num_frames / fps), not spatial resolution. Resolution column shows
+the actual input tensor shape (excluding batch).
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from pathlib import Path
+
+import pytest
+import torch
+
+from benchmarks.conftest import COMPONENT_SIZES, InputSize, audio_latent_torch_shape
+from ltx_bench.core import BenchmarkResult
+from ltx_bench.nvtx import nvtx_range
+from ltx_core.model.audio_vae import AudioDecoder, AudioEncoder, Vocoder
+
+from .conftest import (
+    CHECKPOINT_PATH,
+    build_audio_decoder,
+    build_audio_encoder,
+    build_vocoder,
+)
+
+_MODEL_NAME = Path(CHECKPOINT_PATH).stem
+
+
+def _shape_label(tensor: torch.Tensor) -> str:
+    """Format tensor shape (excluding batch dim) as a resolution label."""
+    return "x".join(str(d) for d in tensor.shape[1:])
+
+
+@pytest.fixture(scope="module")
+def audio_decoder() -> AudioDecoder:
+    return build_audio_decoder(CHECKPOINT_PATH)
+
+
+@pytest.fixture(scope="module")
+def audio_encoder() -> AudioEncoder:
+    return build_audio_encoder(CHECKPOINT_PATH)
+
+
+@pytest.fixture(scope="module")
+def vocoder() -> Vocoder:
+    return build_vocoder(CHECKPOINT_PATH)
+
+
+def _make_spectrogram(
+    audio_decoder: AudioDecoder,
+    size: InputSize,
+) -> torch.Tensor:
+    """Decode a random audio latent to produce a realistic spectrogram."""
+    device = next(audio_decoder.parameters()).device
+    dtype = next(audio_decoder.parameters()).dtype
+    latents = torch.randn(*audio_latent_torch_shape(size), device=device, dtype=dtype)
+    with torch.inference_mode():
+        return audio_decoder(latents)
+
+
+@pytest.mark.bench
+@pytest.mark.parametrize("size", COMPONENT_SIZES, ids=[s.label for s in COMPONENT_SIZES])
+@torch.inference_mode()
+def test_audio_decoder_forward(
+    audio_decoder: AudioDecoder,
+    bench_run: Callable[..., BenchmarkResult],
+    size: InputSize,
+) -> None:
+    """Benchmark AudioDecoder forward pass (latent → spectrogram)."""
+    device = next(audio_decoder.parameters()).device
+    dtype = next(audio_decoder.parameters()).dtype
+    latents = torch.randn(*audio_latent_torch_shape(size), device=device, dtype=dtype)
+
+    def _decode() -> None:
+        with nvtx_range("audio_decoder_forward"):
+            audio_decoder(latents)
+
+    bench_run(_decode, bench_type="audio_decoder_forward", model_name=_MODEL_NAME, resolution=_shape_label(latents))
+
+
+@pytest.mark.bench
+@pytest.mark.parametrize("size", COMPONENT_SIZES, ids=[s.label for s in COMPONENT_SIZES])
+@torch.inference_mode()
+def test_audio_encoder_forward(
+    audio_encoder: AudioEncoder,
+    audio_decoder: AudioDecoder,
+    bench_run: Callable[..., BenchmarkResult],
+    size: InputSize,
+) -> None:
+    """Benchmark AudioEncoder forward pass (spectrogram → latent)."""
+    spectrogram = _make_spectrogram(audio_decoder, size)
+
+    def _encode() -> None:
+        with nvtx_range("audio_encoder_forward"):
+            audio_encoder(spectrogram)
+
+    bench_run(_encode, bench_type="audio_encoder_forward", model_name=_MODEL_NAME, resolution=_shape_label(spectrogram))
+
+
+@pytest.mark.bench
+@pytest.mark.parametrize("size", COMPONENT_SIZES, ids=[s.label for s in COMPONENT_SIZES])
+@torch.inference_mode()
+def test_vocoder_forward(
+    vocoder: Vocoder,
+    audio_decoder: AudioDecoder,
+    bench_run: Callable[..., BenchmarkResult],
+    size: InputSize,
+) -> None:
+    """Benchmark vocoder (HiFi-GAN) forward pass (spectrogram → waveform)."""
+    spectrogram = _make_spectrogram(audio_decoder, size)
+
+    def _forward() -> None:
+        with nvtx_range("vocoder_forward"):
+            vocoder(spectrogram)
+
+    bench_run(_forward, bench_type="vocoder_forward", model_name=_MODEL_NAME, resolution=_shape_label(spectrogram))

ltx-2-internal/packages/ltx-bench/benchmarks/components/test_component_loading.py

@@ -0,0 +1,68 @@
+"""Benchmark: model loading time for each component (safetensors → GPU)."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from pathlib import Path
+
+import pytest
+
+from benchmarks.conftest import warmup_io
+from ltx_bench.core import BenchmarkResult
+
+from .conftest import (
+    CHECKPOINT_PATH,
+    GEMMA_ROOT_PATH,
+    UPSAMPLER_PATH,
+    build_audio_decoder,
+    build_audio_encoder,
+    build_text_encoder,
+    build_transformer,
+    build_upsampler,
+    build_video_decoder,
+    build_video_encoder,
+    build_vocoder,
+)
+
+_CP = Path(CHECKPOINT_PATH).stem
+_UP = Path(UPSAMPLER_PATH).stem
+_GP = Path(GEMMA_ROOT_PATH).name
+
+_LOAD_CASES: list[tuple[str, str, Callable[[], object]]] = [
+    ("text_encoder_load", _GP, lambda: build_text_encoder(CHECKPOINT_PATH, GEMMA_ROOT_PATH)),
+    ("transformer_load", _CP, lambda: build_transformer(CHECKPOINT_PATH)),
+    ("video_encoder_load", _CP, lambda: build_video_encoder(CHECKPOINT_PATH)),
+    ("video_decoder_load", _CP, lambda: build_video_decoder(CHECKPOINT_PATH)),
+    ("audio_encoder_load", _CP, lambda: build_audio_encoder(CHECKPOINT_PATH)),
+    ("audio_decoder_load", _CP, lambda: build_audio_decoder(CHECKPOINT_PATH)),
+    ("vocoder_load", _CP, lambda: build_vocoder(CHECKPOINT_PATH)),
+    ("upsampler_load", _UP, lambda: build_upsampler(UPSAMPLER_PATH)),
+]
+
+
+@pytest.fixture(scope="module", autouse=True)
+def _warmup_page_cache() -> None:
+    """Read checkpoint files into OS page cache before loading benchmarks."""
+    warmup_io()
+
+
+@pytest.mark.bench
+@pytest.mark.parametrize(
+    ("bench_type", "model_name", "builder"),
+    _LOAD_CASES,
+    ids=[c[0] for c in _LOAD_CASES],
+)
+def test_component_load(
+    bench_type: str,
+    model_name: str,
+    builder: Callable[[], object],
+    bench_run: Callable[..., BenchmarkResult],
+) -> None:
+    """Benchmark loading a single component from safetensors to GPU."""
+    bench_run(
+        builder,
+        bench_type=bench_type,
+        model_name=model_name,
+        warmup=0,
+        repeats=1,
+    )

ltx-2-internal/packages/ltx-bench/benchmarks/components/test_lora_fusion.py

@@ -0,0 +1,56 @@
+"""Benchmark: static LoRA fusion via ``apply_loras()`` with different LoRA sizes."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from pathlib import Path
+
+import pytest
+import torch
+
+from ltx_bench.core import BenchmarkResult
+from ltx_core.loader import LTXV_LORA_COMFY_RENAMING_MAP
+from ltx_core.loader.fuse_loras import apply_loras
+from ltx_core.loader.primitives import LoraStateDictWithStrength, StateDict
+from ltx_core.loader.sft_loader import SafetensorsModelStateDictLoader
+from ltx_core.model.transformer import LTXV_MODEL_COMFY_RENAMING_MAP
+
+from .conftest import CHECKPOINT_PATH, DISTILLED_LORA_PATH, IC_LORA_PATH
+
+_LOADER = SafetensorsModelStateDictLoader()
+
+_LORA_CASES = [
+    ("lora_fusion_distilled", DISTILLED_LORA_PATH),
+    ("lora_fusion_ic_lora", IC_LORA_PATH),
+]
+
+
+@pytest.fixture(scope="module")
+def model_sd() -> StateDict:
+    """Load the transformer state dict on GPU, matching production (SingleGPUModelBuilder)."""
+    return _LOADER.load([CHECKPOINT_PATH], sd_ops=LTXV_MODEL_COMFY_RENAMING_MAP, device=torch.device("cuda"))
+
+
+@pytest.mark.bench
+@pytest.mark.parametrize(
+    ("bench_type", "lora_path"),
+    _LORA_CASES,
+    ids=[c[0] for c in _LORA_CASES],
+)
+def test_lora_fusion(
+    bench_type: str,
+    lora_path: str,
+    model_sd: StateDict,
+    bench_run: Callable[..., BenchmarkResult],
+) -> None:
+    """Benchmark static LoRA fusion (apply_loras) with a single LoRA checkpoint."""
+    lora_sd = _LOADER.load([lora_path], sd_ops=LTXV_LORA_COMFY_RENAMING_MAP)
+    lora_pairs = [LoraStateDictWithStrength(lora_sd, 1.0)]
+
+    bench_run(
+        lambda: apply_loras(model_sd, lora_pairs, dtype=torch.bfloat16, destination_sd=model_sd),
+        bench_type=bench_type,
+        model_name=Path(CHECKPOINT_PATH).stem,
+        warmup=1,
+        repeats=3,
+    )

ltx-2-internal/packages/ltx-bench/benchmarks/components/test_text_encoder.py

@@ -0,0 +1,45 @@
+"""Benchmark: Gemma 3 text encoder forward pass (positive + negative prompt)."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from pathlib import Path
+
+import pytest
+import torch
+
+from ltx_bench.core import BenchmarkResult
+from ltx_bench.nvtx import nvtx_range
+from ltx_core.text_encoders.gemma.encoders.base_encoder import GemmaTextEncoder
+from ltx_pipelines.utils.constants import DEFAULT_NEGATIVE_PROMPT
+
+from .conftest import CHECKPOINT_PATH, GEMMA_ROOT_PATH, build_text_encoder
+
+PROMPT = (
+    "A medium close-up shot of a person speaking directly to the camera with natural lighting "
+    "and a softly blurred background."
+)
+
+
+@pytest.fixture(scope="module")
+def text_encoder() -> GemmaTextEncoder:
+    return build_text_encoder(CHECKPOINT_PATH, GEMMA_ROOT_PATH)
+
+
+@pytest.mark.bench
+@torch.inference_mode()
+def test_text_encoder_forward(
+    text_encoder: GemmaTextEncoder,
+    bench_run: Callable[..., BenchmarkResult],
+) -> None:
+    """Benchmark Gemma text encoder: encode positive + negative prompt.
+
+    Uses ``encode_text()`` which is the same function the pipeline calls,
+    encoding both prompts sequentially.
+    """
+
+    def _forward() -> None:
+        with nvtx_range("text_encoder_forward"):
+            [text_encoder.encode(p) for p in [PROMPT, DEFAULT_NEGATIVE_PROMPT]]
+
+    bench_run(_forward, bench_type="text_encoder_forward", model_name=Path(GEMMA_ROOT_PATH).name)

ltx-2-internal/packages/ltx-bench/benchmarks/components/test_transformer.py

@@ -0,0 +1,282 @@
+"""Benchmark: transformer (X0Model) forward pass using production Modality construction.
+
+Measures GPU time for a single denoising step through the 48-layer
+dual-stream transformer. Tests video-only and audio+video variants.
+Input latent states are sized to match stage-1 half-resolution.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable, Generator
+from pathlib import Path
+
+import pytest
+import torch
+
+from benchmarks.conftest import COMPONENT_SIZES, InputSize
+from ltx_bench.core import BenchmarkResult
+from ltx_bench.nvtx import nvtx_range
+from ltx_core.block_streaming import DISK_CPU_SLOTS, BlockStreamingWrapper, StreamingModelBuilder
+from ltx_core.components.patchifiers import AudioPatchifier, VideoLatentPatchifier, VideoLatentShape
+from ltx_core.guidance.perturbations import BatchedPerturbationConfig, PerturbationConfig
+from ltx_core.model.transformer import LTXV_MODEL_COMFY_RENAMING_MAP, LTXModelConfigurator
+from ltx_core.model.transformer.model import X0Model
+from ltx_core.tools import AudioLatentTools, VideoLatentTools
+from ltx_core.types import AudioLatentShape, VideoPixelShape
+from ltx_pipelines.utils.constants import VIDEO_LATENT_CHANNELS, VIDEO_SCALE_FACTORS
+from ltx_pipelines.utils.denoisers import _repeat_state
+from ltx_pipelines.utils.helpers import modality_from_latent_state
+
+from .conftest import (
+    CHECKPOINT_PATH,
+    build_transformer,
+    make_video_latent_state,
+)
+
+_MODEL_NAME = Path(CHECKPOINT_PATH).stem
+_FPS = 24.0
+
+# Dummy context token count for benchmarking (roughly matches a typical short prompt).
+_CONTEXT_TOKENS = 64
+
+
+@pytest.fixture(scope="module")
+def transformer() -> X0Model:
+    return build_transformer(CHECKPOINT_PATH)
+
+
+def _make_av_modalities(
+    transformer: torch.nn.Module,
+    height: int,
+    width: int,
+    num_frames: int,
+    batch: int,
+) -> tuple:
+    """Build batched video + audio modalities and return (video, audio, perturbations, total_tokens)."""
+    model = transformer.velocity_model
+    device = next(model.parameters()).device
+    dtype = next(model.parameters()).dtype
+
+    pixel_shape = VideoPixelShape(batch=1, frames=num_frames, height=height, width=width, fps=_FPS)
+    v_shape = VideoLatentShape.from_pixel_shape(pixel_shape, VIDEO_LATENT_CHANNELS, VIDEO_SCALE_FACTORS)
+    a_shape = AudioLatentShape.from_video_pixel_shape(pixel_shape)
+
+    v_tools = VideoLatentTools(VideoLatentPatchifier(patch_size=1), v_shape, _FPS)
+    a_tools = AudioLatentTools(AudioPatchifier(patch_size=1), a_shape)
+
+    video_state = v_tools.create_initial_state(device=device, dtype=dtype)
+    audio_state = a_tools.create_initial_state(device=device, dtype=dtype)
+
+    v_context_dim = model.transformer_blocks[0].attn2.to_k.in_features
+    a_context_dim = model.transformer_blocks[0].audio_attn2.to_k.in_features
+
+    sigma = torch.full((batch,), 0.5, device=device, dtype=torch.float32)
+    v_context = torch.randn(batch, _CONTEXT_TOKENS, v_context_dim, device=device, dtype=dtype)
+    a_context = torch.randn(batch, _CONTEXT_TOKENS, a_context_dim, device=device, dtype=dtype)
+
+    video = modality_from_latent_state(_repeat_state(video_state, batch), v_context, sigma)
+    audio = modality_from_latent_state(_repeat_state(audio_state, batch), a_context, sigma)
+    perturbations = BatchedPerturbationConfig([PerturbationConfig.empty()] * batch)
+
+    total_tokens = video.latent.shape[1] + audio.latent.shape[1]  # per batch item
+    return video, audio, perturbations, total_tokens
+
+
+@pytest.mark.bench
+@pytest.mark.parametrize("size", COMPONENT_SIZES, ids=[s.label for s in COMPONENT_SIZES])
+@torch.inference_mode()
+def test_transformer_forward_video_only(
+    transformer: X0Model,
+    bench_run: Callable[..., BenchmarkResult],
+    size: InputSize,
+) -> None:
+    """Benchmark X0Model forward pass with video-only input (no audio).
+
+    Uses ``modality_from_latent_state()`` to construct the Modality, matching
+    the production code path exactly.
+    """
+    model = transformer.velocity_model
+    device = next(model.parameters()).device
+    dtype = next(model.parameters()).dtype
+
+    video_state = make_video_latent_state(size, device=device, dtype=dtype)
+    context_dim = model.transformer_blocks[0].attn2.to_k.in_features
+    context = torch.randn(1, _CONTEXT_TOKENS, context_dim, device=device, dtype=dtype)
+    sigma = torch.tensor([0.5], device=device, dtype=torch.float32)
+    perturbations = BatchedPerturbationConfig.empty(1)
+
+    video = modality_from_latent_state(video_state, context, sigma)
+
+    def _forward() -> None:
+        with nvtx_range("transformer_forward_video_only"):
+            transformer(video, None, perturbations)
+
+    bench_run(
+        _forward,
+        bench_type="transformer_forward_video_only",
+        model_name=_MODEL_NAME,
+        resolution=size.label,
+    )
+
+
+# ---------------------------------------------------------------------------
+# Batched forward-pass benchmarks — data for throughput vs token-count graphs.
+#
+# Measures a single transformer forward call at varying batch sizes and
+# resolutions (token counts), with and without layer streaming.
+# To compare sequential (4 x B=1) vs batched (1 x B=4) step time,
+# multiply the B=1 result by the number of guidance passes.
+# ---------------------------------------------------------------------------
+
+_BATCH_SIZES = [
+    pytest.param(1, id="B1"),
+    pytest.param(2, id="B2"),
+    pytest.param(4, id="B4"),
+]
+
+
+@pytest.mark.bench
+@pytest.mark.parametrize("size", COMPONENT_SIZES, ids=[s.label for s in COMPONENT_SIZES])
+@pytest.mark.parametrize("batch", _BATCH_SIZES)
+@torch.inference_mode()
+def test_transformer_av_batch(
+    transformer: X0Model,
+    bench_run: Callable[..., BenchmarkResult],
+    size: InputSize,
+    batch: int,
+) -> None:
+    """Benchmark a single AV forward call at a given batch size (no streaming)."""
+    video, audio, perturbations, total_tokens = _make_av_modalities(
+        transformer, size.height, size.width, size.num_frames, batch=batch
+    )
+
+    def _forward() -> None:
+        with nvtx_range(f"av_b{batch}"):
+            transformer(video, audio, perturbations)
+
+    result = bench_run(
+        _forward,
+        warmup=1,
+        repeats=3,
+        bench_type=f"transformer_av_b{batch}",
+        model_name=_MODEL_NAME,
+        resolution=size.label,
+    )
+    result.extra["total_tokens"] = total_tokens
+    result.extra["batch"] = batch
+    result.extra["streaming"] = False
+
+
+def _build_streaming_wrapper(
+    cpu_slots_count: int | None = None,
+    gpu_slots_count: int | None = None,
+) -> BlockStreamingWrapper:
+    """Build a streaming wrapper for the transformer checkpoint."""
+    builder = StreamingModelBuilder(
+        model_class_configurator=LTXModelConfigurator,
+        model_path=CHECKPOINT_PATH,
+        model_sd_ops=LTXV_MODEL_COMFY_RENAMING_MAP,
+        blocks_attr="velocity_model.transformer_blocks",
+        blocks_prefix="transformer_blocks",
+        state_dict_prefix="velocity_model.",
+        model_wrapper=lambda m: X0Model(m).eval(),
+    )
+    return builder.build(
+        target_device=torch.device("cuda"),
+        dtype=torch.bfloat16,
+        cpu_slots_count=cpu_slots_count,
+        gpu_slots_count=gpu_slots_count,
+    )
+
+
+@pytest.fixture(scope="module")
+def streaming_transformer() -> Generator[BlockStreamingWrapper, None, None]:
+    """Module-scoped RAM-streaming wrapper."""
+    wrapped = _build_streaming_wrapper()
+    yield wrapped
+    wrapped.teardown()
+
+
+@pytest.fixture(scope="module")
+def disk_streaming_transformer() -> Generator[BlockStreamingWrapper, None, None]:
+    """Module-scoped disk-streaming wrapper."""
+    wrapped = _build_streaming_wrapper(cpu_slots_count=DISK_CPU_SLOTS)
+    yield wrapped
+    wrapped.teardown()
+
+
+# ---------------------------------------------------------------------------
+# AV batch — RAM streaming
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.bench
+@pytest.mark.parametrize("size", COMPONENT_SIZES, ids=[s.label for s in COMPONENT_SIZES])
+@pytest.mark.parametrize("batch", _BATCH_SIZES)
+@torch.inference_mode()
+def test_transformer_av_batch_streaming(
+    streaming_transformer: BlockStreamingWrapper,
+    bench_run: Callable[..., BenchmarkResult],
+    size: InputSize,
+    batch: int,
+) -> None:
+    """Benchmark a single AV forward call at a given batch size (with RAM streaming)."""
+
+    video, audio, perturbations, total_tokens = _make_av_modalities(
+        streaming_transformer, size.height, size.width, size.num_frames, batch=batch
+    )
+
+    def _forward() -> None:
+        with nvtx_range(f"av_b{batch}_streaming"):
+            streaming_transformer(video, audio, perturbations)
+
+    result = bench_run(
+        _forward,
+        warmup=1,
+        repeats=3,
+        bench_type=f"transformer_av_b{batch}_streaming",
+        model_name=_MODEL_NAME,
+        resolution=size.label,
+    )
+    result.extra["total_tokens"] = total_tokens
+    result.extra["batch"] = batch
+    result.extra["streaming"] = True
+
+
+# ---------------------------------------------------------------------------
+# AV batch — disk streaming
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.bench
+@pytest.mark.parametrize("size", COMPONENT_SIZES, ids=[s.label for s in COMPONENT_SIZES])
+@pytest.mark.parametrize("batch", _BATCH_SIZES)
+@torch.inference_mode()
+def test_transformer_av_batch_disk_streaming(
+    disk_streaming_transformer: BlockStreamingWrapper,
+    bench_run: Callable[..., BenchmarkResult],
+    size: InputSize,
+    batch: int,
+) -> None:
+    """Benchmark a single AV forward call at a given batch size (with disk streaming)."""
+
+    video, audio, perturbations, total_tokens = _make_av_modalities(
+        disk_streaming_transformer, size.height, size.width, size.num_frames, batch=batch
+    )
+
+    def _forward() -> None:
+        with nvtx_range(f"av_b{batch}_disk_streaming"):
+            disk_streaming_transformer(video, audio, perturbations)
+
+    result = bench_run(
+        _forward,
+        warmup=1,
+        repeats=3,
+        bench_type=f"transformer_av_b{batch}_disk_streaming",
+        model_name=_MODEL_NAME,
+        resolution=size.label,
+    )
+    result.extra["total_tokens"] = total_tokens
+    result.extra["batch"] = batch
+    result.extra["streaming"] = True
+    result.extra["disk"] = True

ltx-2-internal/packages/ltx-bench/benchmarks/components/test_upsampler.py

@@ -0,0 +1,67 @@
+"""Benchmark: spatial latent upsampler 2x using the production ``upsample_video`` wrapper.
+
+Measures GPU time for the learned 2x spatial upsampling of video latents,
+including the per-channel un-normalize / normalize steps. Input latents are
+sized to match stage-1 half-resolution (output is full pipeline resolution).
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from pathlib import Path
+
+import pytest
+import torch
+
+from benchmarks.conftest import COMPONENT_SIZES, InputSize, video_latent_torch_shape
+from ltx_bench.core import BenchmarkResult
+from ltx_bench.nvtx import nvtx_range
+from ltx_core.model.upsampler import LatentUpsampler, upsample_video
+from ltx_core.model.video_vae import VideoEncoder
+
+from .conftest import (
+    CHECKPOINT_PATH,
+    UPSAMPLER_PATH,
+    build_upsampler,
+    build_video_encoder,
+)
+
+
+@pytest.fixture(scope="module")
+def upsampler() -> LatentUpsampler:
+    return build_upsampler(UPSAMPLER_PATH)
+
+
+@pytest.fixture(scope="module")
+def video_encoder() -> VideoEncoder:
+    return build_video_encoder(CHECKPOINT_PATH)
+
+
+@pytest.mark.bench
+@pytest.mark.parametrize("size", COMPONENT_SIZES, ids=[s.label for s in COMPONENT_SIZES])
+@torch.inference_mode()
+def test_spatial_upsampler(
+    upsampler: LatentUpsampler,
+    video_encoder: VideoEncoder,
+    bench_run: Callable[..., BenchmarkResult],
+    size: InputSize,
+) -> None:
+    """Benchmark spatial upsampler 2x using the production ``upsample_video`` wrapper.
+
+    Includes the per-channel un-normalize / normalize steps that the
+    real pipeline applies around the upsampler.
+    """
+    device = next(upsampler.parameters()).device
+    dtype = next(upsampler.parameters()).dtype
+    latents = torch.randn(*video_latent_torch_shape(size), device=device, dtype=dtype)
+
+    def _upsample() -> None:
+        with nvtx_range("upsampler_forward"):
+            upsample_video(latents, video_encoder, upsampler)
+
+    bench_run(
+        _upsample,
+        bench_type="upsampler_forward",
+        model_name=Path(UPSAMPLER_PATH).stem,
+        resolution=size.label,
+    )

ltx-2-internal/packages/ltx-bench/benchmarks/components/test_video_vae.py

@@ -0,0 +1,127 @@
+"""Benchmark: video VAE decoder and encoder (tiled and non-tiled).
+
+Measures GPU time for video latent → pixel (decode) and pixel → latent (encode)
+paths. Input tensors are sized to match stage-1 half-resolution latents.
+Non-tiled variants run at the smallest size only (larger sizes OOM).
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from pathlib import Path
+
+import pytest
+import torch
+
+from benchmarks.conftest import COMPONENT_SIZES, InputSize, pixel_torch_shape, video_latent_torch_shape
+from ltx_bench.core import BenchmarkResult
+from ltx_bench.nvtx import nvtx_range
+from ltx_core.model.video_vae import TilingConfig, VideoDecoder, VideoEncoder
+from ltx_pipelines.utils.constants import LTX_2_3_PARAMS
+
+from .conftest import (
+    CHECKPOINT_PATH,
+    build_video_decoder,
+    build_video_encoder,
+)
+
+_MODEL_NAME = Path(CHECKPOINT_PATH).stem
+HALF_1080P_121 = InputSize(544, 960, 121)
+
+
+@pytest.fixture(scope="module")
+def video_decoder() -> VideoDecoder:
+    return build_video_decoder(CHECKPOINT_PATH)
+
+
+@pytest.fixture(scope="module")
+def video_encoder() -> VideoEncoder:
+    return build_video_encoder(CHECKPOINT_PATH)
+
+
+# --- Decode ---
+
+
+@pytest.mark.bench
+@torch.inference_mode()
+def test_video_vae_decode_non_tiled(
+    video_decoder: VideoDecoder,
+    bench_run: Callable[..., BenchmarkResult],
+) -> None:
+    """Benchmark non-tiled video VAE decoding (HALF_1080P_121 only — larger sizes OOM)."""
+    device = next(video_decoder.parameters()).device
+    dtype = next(video_decoder.parameters()).dtype
+    latents = torch.randn(*video_latent_torch_shape(HALF_1080P_121), device=device, dtype=dtype)
+    generator = torch.Generator(device=device).manual_seed(LTX_2_3_PARAMS.seed)
+
+    def _decode() -> None:
+        with nvtx_range("video_decoder_non_tiled"):
+            for _ in video_decoder.decode_video(latents, None, generator=generator):
+                pass
+
+    bench_run(_decode, bench_type="video_decoder_non_tiled", model_name=_MODEL_NAME, resolution=HALF_1080P_121.label)
+
+
+@pytest.mark.bench
+@pytest.mark.parametrize("size", COMPONENT_SIZES, ids=[s.label for s in COMPONENT_SIZES])
+@torch.inference_mode()
+def test_video_vae_decode_tiled(
+    video_decoder: VideoDecoder,
+    bench_run: Callable[..., BenchmarkResult],
+    size: InputSize,
+) -> None:
+    """Benchmark tiled video VAE decoding across all sizes."""
+    device = next(video_decoder.parameters()).device
+    dtype = next(video_decoder.parameters()).dtype
+    latents = torch.randn(*video_latent_torch_shape(size), device=device, dtype=dtype)
+    generator = torch.Generator(device=device).manual_seed(LTX_2_3_PARAMS.seed)
+    tiling_config = TilingConfig.default()
+
+    def _decode() -> None:
+        with nvtx_range("video_decoder_tiled"):
+            for _ in video_decoder.decode_video(latents, tiling_config, generator=generator):
+                pass
+
+    bench_run(_decode, bench_type="video_decoder_tiled", model_name=_MODEL_NAME, resolution=size.label)
+
+
+# --- Encode ---
+
+
+@pytest.mark.bench
+@torch.inference_mode()
+def test_video_vae_encode_non_tiled(
+    video_encoder: VideoEncoder,
+    bench_run: Callable[..., BenchmarkResult],
+) -> None:
+    """Benchmark non-tiled video VAE encoding (HALF_1080P_121 only — larger sizes OOM)."""
+    device = next(video_encoder.parameters()).device
+    dtype = next(video_encoder.parameters()).dtype
+    pixels = torch.randn(*pixel_torch_shape(HALF_1080P_121), device=device, dtype=dtype)
+
+    def _encode() -> None:
+        with nvtx_range("video_encoder_non_tiled"):
+            video_encoder(pixels)
+
+    bench_run(_encode, bench_type="video_encoder_non_tiled", model_name=_MODEL_NAME, resolution=HALF_1080P_121.label)
+
+
+@pytest.mark.bench
+@pytest.mark.parametrize("size", COMPONENT_SIZES, ids=[s.label for s in COMPONENT_SIZES])
+@torch.inference_mode()
+def test_video_vae_encode_tiled(
+    video_encoder: VideoEncoder,
+    bench_run: Callable[..., BenchmarkResult],
+    size: InputSize,
+) -> None:
+    """Benchmark tiled video VAE encoding across all sizes."""
+    device = next(video_encoder.parameters()).device
+    dtype = next(video_encoder.parameters()).dtype
+    pixels = torch.randn(*pixel_torch_shape(size), device=device, dtype=dtype)
+    tiling_config = TilingConfig.default()
+
+    def _encode() -> None:
+        with nvtx_range("video_encoder_tiled"):
+            video_encoder.tiled_encode(pixels, tiling_config)
+
+    bench_run(_encode, bench_type="video_encoder_tiled", model_name=_MODEL_NAME, resolution=size.label)

ltx-2-internal/packages/ltx-bench/benchmarks/conftest.py

@@ -0,0 +1,190 @@
+"""Model configuration, fixtures, and helpers for benchmark tests."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+
+import torch
+
+from conftest import (
+    GEMMA_ROOT,
+    LTX_2_3_CHECKPOINT_PATH,
+    LTX_2_3_DISTILLED_CHECKPOINT_PATH,
+    LTX_2_3_DISTILLED_LORA_PATH,
+    LTX_2_3_IC_LORA_PATH,
+    LTX_2_3_SPATIAL_UPSAMPLER_PATH,
+)
+from ltx_core.components.noisers import GaussianNoiser
+from ltx_core.components.patchifiers import AudioPatchifier, VideoLatentPatchifier, VideoLatentShape
+from ltx_core.tools import AudioLatentTools, VideoLatentTools
+from ltx_core.types import AudioLatentShape, LatentState, VideoPixelShape
+from ltx_pipelines.utils.constants import (
+    VIDEO_LATENT_CHANNELS,
+    VIDEO_SCALE_FACTORS,
+)
+
+# ---------------------------------------------------------------------------
+# Input sizes
+# ---------------------------------------------------------------------------
+
+
+@dataclass(frozen=True)
+class InputSize:
+    """Video generation target dimensions."""
+
+    height: int
+    width: int
+    num_frames: int
+    frame_rate: float = 24.0
+
+    @property
+    def label(self) -> str:
+        return f"{self.height}x{self.width}x{self.num_frames}"
+
+
+# Pipeline sizes specify video output (height, width, frames) for end-to-end benchmarks.
+PIPELINE_SIZES = [
+    InputSize(1088, 1920, 121),  # 1080p 121 frames ~ 5s @ 24 fps
+    InputSize(1088, 1920, 241),  # 1080p 241 frames ~ 10s @ 24 fps
+    InputSize(1472, 2560, 121),  # 1440p 121 frames ~ 5s @ 24 fps
+    InputSize(1472, 2560, 241),  # 1440p 241 frames ~ 10s @ 24 fps
+]
+
+# Component sizes measure the model's latent "base tile" (2x downsampled from pipeline output).
+COMPONENT_SIZES = [InputSize(s.height // 2, s.width // 2, s.num_frames) for s in PIPELINE_SIZES]
+
+# Model paths resolved at import time.
+CHECKPOINT_PATH = LTX_2_3_CHECKPOINT_PATH.resolve().as_posix()
+DISTILLED_CHECKPOINT = LTX_2_3_DISTILLED_CHECKPOINT_PATH.resolve().as_posix()
+GEMMA_ROOT_PATH = str(GEMMA_ROOT) if GEMMA_ROOT.exists() else ""
+UPSAMPLER_PATH = str(LTX_2_3_SPATIAL_UPSAMPLER_PATH) if LTX_2_3_SPATIAL_UPSAMPLER_PATH.exists() else ""
+DISTILLED_LORA_PATH = LTX_2_3_DISTILLED_LORA_PATH.resolve().as_posix()
+IC_LORA_PATH = LTX_2_3_IC_LORA_PATH.resolve().as_posix()
+
+PIPELINE_ASSETS_DIR = Path(__file__).resolve().parent.parent.parent / "ltx-pipelines" / "tests" / "assets"
+if not PIPELINE_ASSETS_DIR.exists():
+    import warnings
+
+    warnings.warn(f"Pipeline assets not found at {PIPELINE_ASSETS_DIR}", stacklevel=1)
+
+
+# ---------------------------------------------------------------------------
+# IO warmup — read checkpoint files into OS page cache so first benchmark
+# iteration is not penalised by network/disk IO.
+# ---------------------------------------------------------------------------
+def _read_chunked(path: Path, chunk_size: int = 64 * 1024 * 1024) -> None:
+    """Read a file in chunks to warm OS page cache without allocating the full file in Python memory."""
+    with path.open("rb") as f:
+        while f.read(chunk_size):
+            pass
+
+
+def warmup_io() -> None:
+    """Read checkpoint and text-encoder files into OS page cache."""
+    for p in [LTX_2_3_CHECKPOINT_PATH, LTX_2_3_DISTILLED_CHECKPOINT_PATH, LTX_2_3_SPATIAL_UPSAMPLER_PATH]:
+        if p.exists():
+            _read_chunked(p)
+    if GEMMA_ROOT.exists():
+        for p in GEMMA_ROOT.rglob("*.safetensors"):
+            _read_chunked(p)
+
+
+# ---------------------------------------------------------------------------
+# Shape computation from InputSize
+# ---------------------------------------------------------------------------
+
+
+def pixel_shape_for(size: InputSize) -> VideoPixelShape:
+    """Build a VideoPixelShape from an InputSize."""
+    return VideoPixelShape(
+        batch=1,
+        frames=size.num_frames,
+        width=size.width,
+        height=size.height,
+        fps=size.frame_rate,
+    )
+
+
+def video_latent_shape_for(size: InputSize) -> VideoLatentShape:
+    """Compute the video latent shape for a given InputSize."""
+    return VideoLatentShape.from_pixel_shape(
+        pixel_shape_for(size),
+        latent_channels=VIDEO_LATENT_CHANNELS,
+        scale_factors=VIDEO_SCALE_FACTORS,
+    )
+
+
+def audio_latent_shape_for(size: InputSize) -> AudioLatentShape:
+    """Compute the audio latent shape for a given InputSize."""
+    return AudioLatentShape.from_video_pixel_shape(pixel_shape_for(size))
+
+
+def video_latent_torch_shape(size: InputSize) -> tuple[int, ...]:
+    """Return the (B, C, T, H, W) tuple for torch.randn."""
+    return video_latent_shape_for(size).to_torch_shape()
+
+
+def pixel_torch_shape(size: InputSize) -> tuple[int, ...]:
+    """Return the (B, C, T, H, W) pixel tensor shape."""
+    return (1, 3, size.num_frames, size.height, size.width)
+
+
+def audio_latent_torch_shape(size: InputSize) -> tuple[int, ...]:
+    """Return the (B, C, T, M) audio latent tensor shape."""
+    return audio_latent_shape_for(size).to_torch_shape()
+
+
+DEFAULT_PATCH_SIZE = 1
+
+
+# ---------------------------------------------------------------------------
+# Latent state helpers
+# ---------------------------------------------------------------------------
+
+
+def make_video_latent_state(
+    size: InputSize,
+    *,
+    device: torch.device,
+    dtype: torch.dtype,
+    patch_size: int = DEFAULT_PATCH_SIZE,
+) -> LatentState:
+    """Create a patchified video LatentState for a given InputSize."""
+    shape = video_latent_shape_for(size)
+    tools = VideoLatentTools(
+        patchifier=VideoLatentPatchifier(patch_size=patch_size),
+        target_shape=shape,
+        fps=size.frame_rate,
+    )
+    return tools.create_initial_state(device=device, dtype=dtype)
+
+
+def make_audio_latent_state(
+    size: InputSize,
+    *,
+    device: torch.device,
+    dtype: torch.dtype,
+    patch_size: int = DEFAULT_PATCH_SIZE,
+) -> LatentState:
+    """Create a patchified audio LatentState for a given InputSize."""
+    audio_shape = audio_latent_shape_for(size)
+    tools = AudioLatentTools(
+        AudioPatchifier(patch_size=patch_size),
+        audio_shape,
+    )
+    return tools.create_initial_state(device=device, dtype=dtype)
+
+
+def make_noised_states(
+    size: InputSize,
+    *,
+    device: torch.device,
+    dtype: torch.dtype,
+    seed: int = 42,
+) -> tuple[LatentState, LatentState]:
+    """Create noised video and audio LatentStates for a given InputSize."""
+    noiser = GaussianNoiser(generator=torch.Generator(device=device).manual_seed(seed))
+    video_state = noiser(make_video_latent_state(size, device=device, dtype=dtype))
+    audio_state = noiser(make_audio_latent_state(size, device=device, dtype=dtype))
+    return video_state, audio_state

ltx-2-internal/packages/ltx-bench/benchmarks/pipelines/conftest.py

@@ -0,0 +1,10 @@
+"""Re-export benchmark helpers for pipeline tests."""
+
+from benchmarks.conftest import (  # noqa: F401
+    CHECKPOINT_PATH,
+    DISTILLED_CHECKPOINT,
+    DISTILLED_LORA_PATH,
+    GEMMA_ROOT_PATH,
+    PIPELINE_ASSETS_DIR,
+    UPSAMPLER_PATH,
+)

ltx-2-internal/packages/ltx-bench/benchmarks/pipelines/test_pipeline_distilled.py

@@ -0,0 +1,88 @@
+"""Benchmark: full DistilledPipeline end-to-end (8-step fast inference)."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from pathlib import Path
+
+import pytest
+import torch
+
+from benchmarks.conftest import PIPELINE_SIZES, InputSize
+from ltx_bench import regions
+from ltx_bench.core import BenchmarkResult
+from ltx_bench.instrument import instrument
+from ltx_bench.nvtx import nvtx_range
+from ltx_core.model.audio_vae.audio_vae import decode_audio
+from ltx_core.model.upsampler.model import upsample_video
+from ltx_core.model.video_vae import TilingConfig
+from ltx_core.model.video_vae.video_vae import VideoDecoder
+from ltx_core.text_encoders.gemma.encoders.base_encoder import GemmaTextEncoder
+from ltx_pipelines.distilled import DistilledPipeline
+from ltx_pipelines.utils.blocks import DiffusionStage
+from ltx_pipelines.utils.constants import LTX_2_3_PARAMS
+from ltx_pipelines.utils.samplers import euler_denoising_loop
+
+from .conftest import DISTILLED_CHECKPOINT, GEMMA_ROOT_PATH, UPSAMPLER_PATH
+
+
+@pytest.fixture(scope="module")
+def pipeline() -> DistilledPipeline:
+    return DistilledPipeline(
+        distilled_checkpoint_path=DISTILLED_CHECKPOINT,
+        spatial_upsampler_path=UPSAMPLER_PATH,
+        gemma_root=GEMMA_ROOT_PATH,
+        loras=[],
+    )
+
+
+@pytest.mark.bench
+@pytest.mark.parametrize("size", PIPELINE_SIZES, ids=[s.label for s in PIPELINE_SIZES])
+@torch.inference_mode()
+def test_pipeline_distilled_e2e(
+    pipeline: DistilledPipeline,
+    bench_run: Callable[..., BenchmarkResult],
+    size: InputSize,
+) -> None:
+    """Benchmark distilled pipeline end-to-end (8-step schedule).
+
+    Two-stage generation with distilled sigma schedule.
+    """
+    tiling_config = TilingConfig.default()
+
+    def _run() -> None:
+        with nvtx_range("pipeline_distilled_e2e"):
+            video_frames, _audio = pipeline(
+                prompt="A cat sitting on a windowsill watching birds outside.",
+                seed=LTX_2_3_PARAMS.seed,
+                height=size.height,
+                width=size.width,
+                num_frames=size.num_frames,
+                frame_rate=size.frame_rate,
+                images=[],
+                tiling_config=tiling_config,
+            )
+            for _ in video_frames:
+                pass
+
+    instrument.activate()
+    instrument.patch(DiffusionStage._build_transformer, regions.MODEL_LOAD)
+    instrument.patch(GemmaTextEncoder.encode, regions.TEXT_ENCODING)
+    instrument.patch(euler_denoising_loop, regions.DENOISING)
+    instrument.patch(VideoDecoder.forward, regions.VAE_VIDEO_DECODE, aggregate=True)
+    instrument.patch(decode_audio, regions.VAE_AUDIO_DECODE)
+    instrument.patch(upsample_video, regions.SPATIAL_UPSAMPLE)
+    instrument.patch_tqdm()
+    try:
+        result = bench_run(
+            _run,
+            warmup=0,
+            repeats=1,
+            bench_type="distilled",
+            model_name=Path(DISTILLED_CHECKPOINT).stem,
+            resolution=size.label,
+        )
+    finally:
+        instrument.deactivate()
+
+    instrument.collect_stages(result)

ltx-2-internal/packages/ltx-bench/benchmarks/pipelines/test_pipeline_two_stage.py

@@ -0,0 +1,178 @@
+"""Benchmark: TI2VidTwoStagesPipeline — baseline, RAM streaming, disk streaming."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from pathlib import Path
+
+import pytest
+import torch
+
+from ltx_bench import regions
+from ltx_bench.core import BenchmarkResult
+from ltx_bench.instrument import instrument
+from ltx_bench.nvtx import nvtx_range
+from ltx_core.block_streaming.builder import StreamingModelBuilder
+from ltx_core.block_streaming.disk import DiskBlockReader, DiskTensorReader, LoraSource
+from ltx_core.block_streaming.pool import WeightPool
+from ltx_core.block_streaming.provider import WeightsProvider
+from ltx_core.block_streaming.source import DiskWeightSource
+from ltx_core.loader import LTXV_LORA_COMFY_RENAMING_MAP, LoraPathStrengthAndSDOps
+from ltx_core.model.audio_vae.audio_vae import decode_audio
+from ltx_core.model.upsampler.model import upsample_video
+from ltx_core.model.video_vae import TilingConfig
+from ltx_core.model.video_vae.video_vae import VideoDecoder
+from ltx_core.text_encoders.gemma.encoders.base_encoder import GemmaTextEncoder
+from ltx_pipelines.ti2vid_two_stages import TI2VidTwoStagesPipeline
+from ltx_pipelines.utils.blocks import DiffusionStage, PromptEncoder
+from ltx_pipelines.utils.constants import DEFAULT_NEGATIVE_PROMPT, LTX_2_3_PARAMS
+from ltx_pipelines.utils.samplers import euler_denoising_loop
+from ltx_pipelines.utils.types import OffloadMode
+
+from .conftest import CHECKPOINT_PATH, DISTILLED_LORA_PATH, GEMMA_ROOT_PATH, UPSAMPLER_PATH
+
+_PROMPT = "A cat sitting on a windowsill watching birds outside."
+_RESOLUTION = f"{LTX_2_3_PARAMS.stage_2_height}x{LTX_2_3_PARAMS.stage_2_width}x{LTX_2_3_PARAMS.num_frames}"
+
+_DISTILLED_LORA = [
+    LoraPathStrengthAndSDOps(DISTILLED_LORA_PATH, 0.6, LTXV_LORA_COMFY_RENAMING_MAP),
+]
+
+
+def _create_pipeline(offload_mode: OffloadMode = OffloadMode.NONE) -> TI2VidTwoStagesPipeline:
+    return TI2VidTwoStagesPipeline(
+        checkpoint_path=CHECKPOINT_PATH,
+        distilled_lora=_DISTILLED_LORA,
+        spatial_upsampler_path=UPSAMPLER_PATH,
+        gemma_root=GEMMA_ROOT_PATH,
+        loras=[],
+        offload_mode=offload_mode,
+    )
+
+
+@pytest.fixture(scope="module")
+def pipeline() -> TI2VidTwoStagesPipeline:
+    return _create_pipeline()
+
+
+def _run_pipeline(pipeline: TI2VidTwoStagesPipeline) -> None:
+    tiling_config = TilingConfig.default()
+    with nvtx_range("pipeline_two_stage_e2e"):
+        video_frames, _audio = pipeline(
+            prompt=_PROMPT,
+            negative_prompt=DEFAULT_NEGATIVE_PROMPT,
+            seed=LTX_2_3_PARAMS.seed,
+            height=LTX_2_3_PARAMS.stage_2_height,
+            width=LTX_2_3_PARAMS.stage_2_width,
+            num_frames=LTX_2_3_PARAMS.num_frames,
+            frame_rate=LTX_2_3_PARAMS.frame_rate,
+            num_inference_steps=LTX_2_3_PARAMS.num_inference_steps,
+            video_guider_params=LTX_2_3_PARAMS.video_guider_params,
+            audio_guider_params=LTX_2_3_PARAMS.audio_guider_params,
+            images=[],
+            tiling_config=tiling_config,
+        )
+        for _ in video_frames:
+            pass
+
+
+def _patch_pipeline_stages() -> None:
+    """Register instrument patches for the common pipeline stages."""
+    instrument.patch(PromptEncoder.__call__, regions.PROMPT_ENCODER)
+    instrument.patch(DiffusionStage.__call__, regions.DIFFUSION_STAGE)
+    instrument.patch(GemmaTextEncoder.encode, regions.TEXT_ENCODING)
+    instrument.patch(euler_denoising_loop, regions.DENOISING)
+    instrument.patch(VideoDecoder.forward, regions.VAE_VIDEO_DECODE, aggregate=True)
+    instrument.patch(decode_audio, regions.VAE_AUDIO_DECODE)
+    instrument.patch(upsample_video, regions.SPATIAL_UPSAMPLE)
+    instrument.patch_tqdm()
+
+
+def _patch_streaming_stages() -> None:
+    """Register instrument patches for block-streaming internals."""
+    instrument.patch(WeightsProvider.get, regions.PROVIDER_GET, aggregate=True)
+    instrument.patch(WeightsProvider._copy_to_gpu, regions.H2D_COPY, aggregate=True)
+    instrument.patch(WeightsProvider._fuse_block_loras, regions.LORA_FUSION, aggregate=True)
+    instrument.patch(WeightPool.__init__, regions.WEIGHT_POOL_INIT, aggregate=True)
+
+
+def _patch_disk_streaming_stages() -> None:
+    """Register instrument patches specific to disk streaming."""
+    instrument.patch(DiskBlockReader.read_into, regions.DISK_READ_INTO, aggregate=True)
+    instrument.patch(DiskTensorReader.__init__, regions.DISK_READER_INIT, aggregate=True)
+    instrument.patch(StreamingModelBuilder._build_disk_source, regions.BUILD_DISK_SOURCE)
+    instrument.patch(DiskWeightSource.get, regions.DISK_SOURCE_GET, aggregate=True)
+    instrument.patch(StreamingModelBuilder._load_non_block_weights, regions.LOAD_NON_BLOCK_WEIGHTS)
+    instrument.patch(LoraSource.__init__, regions.LORA_SOURCE_INIT, aggregate=True)
+    instrument.patch(LoraSource.get_delta, regions.LORA_GET_DELTA, aggregate=True)
+
+
+def _patch_ram_streaming_stages() -> None:
+    """Register instrument patches specific to RAM streaming."""
+    instrument.patch(StreamingModelBuilder._build_pinned_source, regions.BUILD_PINNED_SOURCE)
+
+
+def _instrument_and_run(
+    pipeline: TI2VidTwoStagesPipeline,
+    bench_run: Callable[..., BenchmarkResult],
+    bench_type: str,
+    extra_patches: Callable[[], None] | None = None,
+) -> BenchmarkResult:
+    instrument.activate()
+    _patch_pipeline_stages()
+    if extra_patches is not None:
+        extra_patches()
+    try:
+        result = bench_run(
+            lambda: _run_pipeline(pipeline),
+            warmup=0,
+            repeats=1,
+            bench_type=bench_type,
+            model_name=Path(CHECKPOINT_PATH).stem,
+            resolution=_RESOLUTION,
+        )
+    finally:
+        instrument.deactivate()
+
+    instrument.collect_stages(result)
+    return result
+
+
+@pytest.mark.bench
+@torch.inference_mode()
+def test_pipeline_two_stage_e2e(
+    pipeline: TI2VidTwoStagesPipeline,
+    bench_run: Callable[..., BenchmarkResult],
+) -> None:
+    """Baseline: full model on GPU, no streaming."""
+    _instrument_and_run(pipeline, bench_run, "ti2vid_two_stages")
+
+
+@pytest.mark.bench
+@torch.inference_mode()
+def test_pipeline_two_stage_ram_streaming(
+    bench_run: Callable[..., BenchmarkResult],
+) -> None:
+    """RAM streaming: weights pinned in CPU RAM, streamed to GPU per-layer."""
+    pipeline = _create_pipeline(OffloadMode.CPU)
+
+    def _patches() -> None:
+        _patch_streaming_stages()
+        _patch_ram_streaming_stages()
+
+    _instrument_and_run(pipeline, bench_run, "ti2vid_two_stages_cpu_offload", extra_patches=_patches)
+
+
+@pytest.mark.bench
+@torch.inference_mode()
+def test_pipeline_two_stage_disk_streaming(
+    bench_run: Callable[..., BenchmarkResult],
+) -> None:
+    """Disk streaming: weights read from safetensors on demand."""
+    pipeline = _create_pipeline(OffloadMode.DISK)
+
+    def _patches() -> None:
+        _patch_streaming_stages()
+        _patch_disk_streaming_stages()
+
+    _instrument_and_run(pipeline, bench_run, "ti2vid_two_stages_disk_offload", extra_patches=_patches)

ltx-2-internal/packages/ltx-bench/benchmarks/pipelines/test_pipeline_two_stage_hq.py

@@ -0,0 +1,101 @@
+"""Benchmark: full TI2VidTwoStagesHQPipeline end-to-end (res2s sampler)."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from pathlib import Path
+
+import pytest
+import torch
+
+from benchmarks.conftest import PIPELINE_SIZES, InputSize
+from ltx_bench import regions
+from ltx_bench.core import BenchmarkResult
+from ltx_bench.instrument import instrument
+from ltx_bench.nvtx import nvtx_range
+from ltx_core.loader import LTXV_LORA_COMFY_RENAMING_MAP, LoraPathStrengthAndSDOps
+from ltx_core.model.audio_vae.audio_vae import decode_audio
+from ltx_core.model.upsampler.model import upsample_video
+from ltx_core.model.video_vae import TilingConfig
+from ltx_core.model.video_vae.video_vae import VideoDecoder
+from ltx_core.text_encoders.gemma.encoders.base_encoder import GemmaTextEncoder
+from ltx_pipelines.ti2vid_two_stages_hq import TI2VidTwoStagesHQPipeline
+from ltx_pipelines.utils.blocks import DiffusionStage
+from ltx_pipelines.utils.constants import DEFAULT_NEGATIVE_PROMPT, LTX_2_3_HQ_PARAMS
+from ltx_pipelines.utils.samplers import res2s_audio_video_denoising_loop
+
+from .conftest import CHECKPOINT_PATH, DISTILLED_LORA_PATH, GEMMA_ROOT_PATH, UPSAMPLER_PATH
+
+
+@pytest.fixture(scope="module")
+def pipeline() -> TI2VidTwoStagesHQPipeline:
+    distilled_lora = [
+        LoraPathStrengthAndSDOps(DISTILLED_LORA_PATH, 1.0, LTXV_LORA_COMFY_RENAMING_MAP),
+    ]
+    return TI2VidTwoStagesHQPipeline(
+        checkpoint_path=CHECKPOINT_PATH,
+        distilled_lora=distilled_lora,
+        distilled_lora_strength_stage_1=0.25,
+        distilled_lora_strength_stage_2=0.5,
+        spatial_upsampler_path=UPSAMPLER_PATH,
+        gemma_root=GEMMA_ROOT_PATH,
+        loras=(),
+    )
+
+
+@pytest.mark.bench
+@pytest.mark.parametrize("size", PIPELINE_SIZES, ids=[s.label for s in PIPELINE_SIZES])
+@torch.inference_mode()
+def test_pipeline_two_stage_hq_e2e(
+    pipeline: TI2VidTwoStagesHQPipeline,
+    bench_run: Callable[..., BenchmarkResult],
+    size: InputSize,
+) -> None:
+    """Benchmark HQ two-stage text-to-video pipeline end-to-end.
+
+    Uses the res2s second-order sampler.
+    Stage 1: lower-resolution generation with CFG guidance.
+    Stage 2: target-resolution refinement with distilled LoRA.
+    """
+    tiling_config = TilingConfig.default()
+
+    def _run() -> None:
+        with nvtx_range("pipeline_two_stage_hq_e2e"):
+            video_frames, _audio = pipeline(
+                prompt="A cat sitting on a windowsill watching birds outside.",
+                negative_prompt=DEFAULT_NEGATIVE_PROMPT,
+                seed=LTX_2_3_HQ_PARAMS.seed,
+                height=size.height,
+                width=size.width,
+                num_frames=size.num_frames,
+                frame_rate=size.frame_rate,
+                num_inference_steps=LTX_2_3_HQ_PARAMS.num_inference_steps,
+                video_guider_params=LTX_2_3_HQ_PARAMS.video_guider_params,
+                audio_guider_params=LTX_2_3_HQ_PARAMS.audio_guider_params,
+                images=[],
+                tiling_config=tiling_config,
+            )
+            for _ in video_frames:
+                pass
+
+    instrument.activate()
+    instrument.patch(DiffusionStage._build_transformer, regions.MODEL_LOAD)
+    instrument.patch(GemmaTextEncoder.encode, regions.TEXT_ENCODING)
+    instrument.patch(res2s_audio_video_denoising_loop, regions.DENOISING)
+    instrument.patch(VideoDecoder.forward, regions.VAE_VIDEO_DECODE, aggregate=True)
+    instrument.patch(decode_audio, regions.VAE_AUDIO_DECODE)
+    instrument.patch(upsample_video, regions.SPATIAL_UPSAMPLE)
+    instrument.patch_tqdm()
+    try:
+        result = bench_run(
+            _run,
+            warmup=0,
+            repeats=1,
+            bench_type="ti2vid_two_stages_hq",
+            model_name=Path(CHECKPOINT_PATH).stem,
+            resolution=size.label,
+        )
+    finally:
+        instrument.deactivate()
+
+    instrument.collect_stages(result)

ltx-2-internal/packages/ltx-bench/pyproject.toml

@@ -0,0 +1,14 @@
+[project]
+name = "ltx-bench"
+version = "0.1.0"
+description = "Profiling and benchmarking tools for LTX-2"
+readme = "README.md"
+requires-python = ">=3.10"
+dependencies = ["ltx-core", "ltx-pipelines", "psutil", "torch~=2.7"]
+
+[project.entry-points.pytest11]
+ltx_bench = "ltx_bench.conftest"
+
+[build-system]
+requires = ["uv_build>=0.9.8,<0.10.0"]
+build-backend = "uv_build"

ltx-2-internal/packages/ltx-bench/src/ltx_bench/__init__.py

@@ -0,0 +1,31 @@
+from ltx_bench.core import BenchmarkResult, MemoryTimelineEntry, MemoryTimelineTracker, benchmark, cuda_timer
+from ltx_bench.instrument import instrument
+from ltx_bench.memory_graph import generate_memory_graph
+from ltx_bench.nvtx import nvtx_range
+from ltx_bench.storage import (
+    TrendRun,
+    download_trend,
+    generate_report_md,
+    load_results,
+    merge_reports,
+    save_run,
+    upload_to_gcs,
+)
+
+__all__ = [
+    "BenchmarkResult",
+    "MemoryTimelineEntry",
+    "MemoryTimelineTracker",
+    "TrendRun",
+    "benchmark",
+    "cuda_timer",
+    "download_trend",
+    "generate_memory_graph",
+    "generate_report_md",
+    "instrument",
+    "load_results",
+    "merge_reports",
+    "nvtx_range",
+    "save_run",
+    "upload_to_gcs",
+]

ltx-2-internal/packages/ltx-bench/src/ltx_bench/conftest.py

@@ -0,0 +1,215 @@
+"""pytest plugin for ltx-bench: --bench flag, fixtures, and result collection."""
+
+from __future__ import annotations
+
+import gc
+import logging
+import os
+from collections.abc import Callable, Generator
+from pathlib import Path
+from typing import Any
+
+import pytest
+import torch
+
+from ltx_bench.core import BenchmarkResult, benchmark
+from ltx_bench.instrument import instrument
+from ltx_bench.memory_graph import generate_memory_graph
+from ltx_bench.storage import save_run, upload_to_gcs
+
+logger = logging.getLogger("ltx-bench")
+
+
+def pytest_addoption(parser: pytest.Parser) -> None:
+    group = parser.getgroup("bench", "LTX benchmarking options")
+    group.addoption("--bench", action="store_true", default=False, help="Run benchmark tests")
+    group.addoption(
+        "--nsys",
+        action="store_true",
+        default=False,
+        help="Enable NVTX markers for Nsight Systems profiling",
+    )
+    group.addoption(
+        "--bench-mode",
+        choices=["timing", "memory"],
+        default="timing",
+        help="Measurement mode: 'timing' for CUDA event timing, 'memory' for memory timeline profiling",
+    )
+    group.addoption(
+        "--bench-output",
+        default="packages/ltx-bench/bench_results",
+        help="Directory for benchmark result files",
+    )
+    group.addoption("--bench-warmup", type=int, default=2, help="Number of warmup iterations")
+    group.addoption("--bench-repeats", type=int, default=5, help="Number of measured iterations")
+    group.addoption("--bench-gcs-upload", action="store_true", default=False, help="Upload results to GCS")
+    group.addoption("--bench-gcs-bucket", default=None, help="Override default GCS bucket name")
+
+
+def pytest_configure(config: pytest.Config) -> None:
+    config.addinivalue_line("markers", "bench: mark test as a benchmark (requires --bench to run)")
+    config._bench_results: list[BenchmarkResult] = []  # type: ignore[attr-defined]
+    if config.getoption("--nsys"):
+        os.environ["LTX_BENCH_NVTX"] = "1"
+    else:
+        os.environ.pop("LTX_BENCH_NVTX", None)
+    instrument.set_bench_mode(config.getoption("--bench-mode"))
+
+
+def pytest_collection_modifyitems(config: pytest.Config, items: list[pytest.Item]) -> None:
+    if not config.getoption("--bench"):
+        skip = pytest.mark.skip(reason="Benchmarks require --bench flag")
+        for item in items:
+            if "bench" in item.keywords:
+                item.add_marker(skip)
+
+
+@pytest.fixture(scope="session")
+def bench_config(request: pytest.FixtureRequest) -> dict[str, int | str]:
+    """Expose benchmark CLI options as a dict fixture."""
+    return {
+        "warmup": request.config.getoption("--bench-warmup"),
+        "repeats": request.config.getoption("--bench-repeats"),
+        "output_dir": request.config.getoption("--bench-output"),
+        "bench_mode": request.config.getoption("--bench-mode"),
+    }
+
+
+@pytest.fixture
+def bench_collect(request: pytest.FixtureRequest) -> list[BenchmarkResult]:
+    """Callable list -- append BenchmarkResults to register them for saving."""
+    return request.config._bench_results  # type: ignore[attr-defined]
+
+
+_BenchRunFn = Callable[..., BenchmarkResult]
+
+
+@pytest.fixture
+def bench_run(
+    bench_config: dict[str, int | str],
+    bench_collect: list[BenchmarkResult],
+) -> _BenchRunFn:
+    """Convenience fixture: calls ``benchmark()`` with session config and auto-collects."""
+    mode = str(bench_config["bench_mode"])
+
+    def _run(
+        fn: Callable[[], Any],
+        *,
+        bench_type: str = "",
+        model_name: str = "",
+        warmup: int | None = None,
+        repeats: int | None = None,
+        **extra: Any,  # noqa: ANN401
+    ) -> BenchmarkResult:
+        extra.setdefault("dtype", "bfloat16")
+        effective_warmup = warmup if warmup is not None else int(bench_config["warmup"])
+        effective_repeats = repeats if repeats is not None else int(bench_config["repeats"])
+        # Memory mode defaults to 1 repeat for a clean timeline (unless overridden)
+        if mode == "memory" and repeats is None:
+            effective_repeats = 1
+        result = benchmark(
+            fn,
+            warmup=effective_warmup,
+            repeats=effective_repeats,
+            bench_type=bench_type,
+            model_name=model_name,
+            bench_mode=mode,
+            **extra,
+        )
+        bench_collect.append(result)
+        return result
+
+    return _run
+
+
+@pytest.fixture(autouse=True)
+def _bench_cleanup(request: pytest.FixtureRequest) -> Generator[None, None, None]:
+    """Free GPU memory before and after each benchmark test."""
+    if "bench" not in request.keywords:
+        yield
+        return
+
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+
+    yield
+
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+
+
+def _format_summary(results: list[BenchmarkResult]) -> str:
+    """Render a mode-appropriate summary table."""
+    if not results:
+        return ""
+
+    mode = results[0].bench_mode
+
+    if mode == "memory":
+        lines = [
+            "",
+            "## Benchmark Summary (Memory Mode)",
+            "",
+            "| Benchmark | Peak VRAM (MB) | Peak USS (MB) | Peak RSS (MB) | Non-evictable (MB) | Samples |",
+            "|-----------|----------------|---------------|---------------|----------------|---------|",
+        ]
+        for r in results:
+            lines.append(
+                f"| {r.name} | {r.peak_vram_mb:.0f}"
+                f" | {r.peak_uss_mb:.0f} | {r.peak_rss_mb:.0f}"
+                f" | {r.peak_non_evictable_mb:.0f} | {len(r.memory_timeline)} |"
+            )
+    else:
+        lines = [
+            "",
+            "## Benchmark Summary (Timing Mode)",
+            "",
+            "| Benchmark | GPU (ms) | Wall (ms) |",
+            "|-----------|----------|-----------|",
+        ]
+        for r in results:
+            lines.append(f"| {r.name} | {r.gpu_time_ms:.1f} | {r.wall_time_ms:.1f} |")
+
+    lines.append("")
+    return "\n".join(lines)
+
+
+def pytest_sessionfinish(session: pytest.Session, exitstatus: int) -> None:  # noqa: ARG001
+    results: list[BenchmarkResult] = session.config._bench_results  # type: ignore[attr-defined]
+    if results:
+        output_dir = session.config.getoption("--bench-output")
+
+        run_dir = save_run(results, output_dir)
+        logger.info("Saved %d results to %s", len(results), run_dir)
+        logger.info(_format_summary(results))
+
+        # Generate memory graphs for pipeline results only (component benchmarks
+        # are too short to produce useful timelines).
+        if results[0].bench_mode == "memory":
+            for r in results:
+                if r.memory_timeline and r.extra.get("stage_wall_regions"):
+                    safe_name = r.name.replace(" / ", "_").replace(" ", "_")
+                    graph_path = Path(run_dir) / f"{safe_name}_memory.png"
+                    generate_memory_graph(r, graph_path)
+                    r.memory_graph_path = str(graph_path)
+                    logger.info("Memory graph: %s", graph_path)
+
+        if session.config.getoption("--bench-gcs-upload"):
+            branch = results[0].git_branch if results else ""
+            commit = results[0].git_commit if results else ""
+            bucket_kwarg: dict[str, str] = {}
+            custom_bucket = session.config.getoption("--bench-gcs-bucket")
+            if custom_bucket:
+                bucket_kwarg["bucket_name"] = custom_bucket
+            uri = upload_to_gcs(
+                run_dir / "results.jsonl",
+                branch=branch,
+                commit=commit,
+                **bucket_kwarg,
+            )
+            if uri:
+                logger.info("Uploaded to %s", uri)
+            else:
+                logger.warning("GCS upload skipped or failed")

ltx-2-internal/packages/ltx-bench/src/ltx_bench/core.py

@@ -0,0 +1,499 @@
+from __future__ import annotations
+
+import functools
+import os
+import subprocess
+import threading
+import time
+from collections.abc import Callable, Generator
+from contextlib import contextmanager
+from dataclasses import dataclass, field, replace
+from datetime import datetime, timezone
+from statistics import median
+from typing import Any
+
+import torch
+
+# ---------------------------------------------------------------------------
+# Data classes
+# ---------------------------------------------------------------------------
+
+_MB = 1024 * 1024
+
+
+@dataclass(frozen=True)
+class MemoryTimelineEntry:
+    """Single sample in a memory timeline, combining smaps_rollup and VRAM."""
+
+    elapsed_s: float
+
+    # /proc/pid/smaps_rollup fields (MB)
+    rss_mb: float
+    pss_mb: float
+    private_clean_mb: float
+    private_dirty_mb: float
+    shared_clean_mb: float
+    shared_dirty_mb: float
+    anonymous_mb: float
+    swap_mb: float
+    uss_mb: float  # Private_Clean + Private_Dirty + Private_Hugetlb
+
+    vram_allocated_mb: float
+    vram_reserved_mb: float
+
+
+@dataclass
+class BenchmarkResult:
+    """Single benchmark measurement with timing, memory, and metadata."""
+
+    bench_type: str
+    model_name: str
+    bench_mode: str = "timing"
+
+    # Timing (populated in timing mode only)
+    all_wall_times_ms: list[float] = field(default_factory=list)
+    all_gpu_times_ms: list[float] = field(default_factory=list)
+    num_steps: int = 0
+    step_times_gpu_ms: list[float] = field(default_factory=list)
+
+    # Memory peaks (populated in memory mode only)
+    peak_vram_mb: float = 0.0
+    peak_uss_mb: float = 0.0
+    peak_rss_mb: float = 0.0
+    peak_non_evictable_mb: float = 0.0
+
+    # Memory timeline (populated in memory mode only)
+    memory_timeline: list[MemoryTimelineEntry] = field(default_factory=list)
+    memory_graph_path: str = ""
+
+    # Metadata
+    git_commit: str = ""
+    git_branch: str = ""
+    gpu_name: str = ""
+    dtype: str = ""
+    resolution: str = ""
+    timestamp: str = ""
+    extra: dict[str, Any] = field(default_factory=dict)
+
+    @property
+    def gpu_time_ms(self) -> float:
+        """Median GPU time across all measured iterations."""
+        return median(self.all_gpu_times_ms) if self.all_gpu_times_ms else 0.0
+
+    @property
+    def wall_time_ms(self) -> float:
+        """Median wall-clock time across all measured iterations."""
+        return median(self.all_wall_times_ms) if self.all_wall_times_ms else 0.0
+
+    @property
+    def name(self) -> str:
+        """Full benchmark name: bench_type / model_name [/ resolution]."""
+        parts = [self.bench_type, self.model_name]
+        if self.resolution:
+            parts.append(self.resolution)
+        return " / ".join(parts)
+
+
+# ---------------------------------------------------------------------------
+# smaps_rollup parsing
+# ---------------------------------------------------------------------------
+
+
+@dataclass(frozen=True)
+class SmapsSnapshot:
+    """All useful fields from /proc/pid/smaps_rollup, in bytes."""
+
+    rss: int = 0
+    pss: int = 0
+    private_clean: int = 0
+    private_dirty: int = 0
+    private_hugetlb: int = 0
+    shared_clean: int = 0
+    shared_dirty: int = 0
+    anonymous: int = 0
+    swap: int = 0
+
+    @property
+    def uss(self) -> int:
+        return self.private_clean + self.private_dirty + self.private_hugetlb
+
+    @property
+    def shared(self) -> int:
+        return self.rss - self.uss
+
+    @staticmethod
+    def read(pid: int) -> SmapsSnapshot:
+        """Parse /proc/pid/smaps_rollup in a single pass."""
+        rss = pss = private_clean = private_dirty = private_hugetlb = 0
+        shared_clean = shared_dirty = anonymous = swap = 0
+        with open(f"/proc/{pid}/smaps_rollup", "rb") as f:
+            for line in f:
+                if line.startswith(b"Rss:"):
+                    rss = int(line.split()[1]) * 1024
+                elif line.startswith(b"Pss:"):
+                    # Pss: is the aggregate; skip Pss_Dirty/Pss_Anon/etc.
+                    if line[3:4] == b":":
+                        pss = int(line.split()[1]) * 1024
+                elif line.startswith(b"Private_Clean:"):
+                    private_clean = int(line.split()[1]) * 1024
+                elif line.startswith(b"Private_Dirty:"):
+                    private_dirty = int(line.split()[1]) * 1024
+                elif line.startswith(b"Private_Hugetlb:"):
+                    private_hugetlb = int(line.split()[1]) * 1024
+                elif line.startswith(b"Shared_Clean:"):
+                    shared_clean = int(line.split()[1]) * 1024
+                elif line.startswith(b"Shared_Dirty:"):
+                    shared_dirty = int(line.split()[1]) * 1024
+                elif line.startswith(b"Anonymous:"):
+                    anonymous = int(line.split()[1]) * 1024
+                elif line.startswith(b"Swap:") and line[4:5] == b":":
+                    swap = int(line.split()[1]) * 1024
+        return SmapsSnapshot(
+            rss=rss,
+            pss=pss,
+            private_clean=private_clean,
+            private_dirty=private_dirty,
+            private_hugetlb=private_hugetlb,
+            shared_clean=shared_clean,
+            shared_dirty=shared_dirty,
+            anonymous=anonymous,
+            swap=swap,
+        )
+
+
+# ---------------------------------------------------------------------------
+# Memory timeline tracker (memory mode)
+# ---------------------------------------------------------------------------
+
+
+class MemoryTimelineTracker:
+    """Polls smaps_rollup + VRAM from a background thread, recording a timeline.
+
+    Default 15 ms polling interval provides high resolution without
+    meaningful overhead (each /proc read takes ~4 ms on a large process).
+    """
+
+    def __init__(self, interval_s: float = 0.015) -> None:
+        self._pid = os.getpid()
+        self._interval = interval_s
+        self._timeline: list[MemoryTimelineEntry] = []
+        self._stop = threading.Event()
+        self._thread: threading.Thread | None = None
+        self._start_time: float = 0.0
+        self._has_cuda = torch.cuda.is_available()
+
+    def _sample(self) -> None:
+        snap = SmapsSnapshot.read(self._pid)
+        elapsed = time.perf_counter() - self._start_time
+        # memory_allocated / memory_reserved are thread-safe stat queries
+        # (atomic counter reads in the caching allocator), safe from any thread.
+        if self._has_cuda:
+            vram_alloc = torch.cuda.memory_allocated() / _MB
+            vram_resv = torch.cuda.memory_reserved() / _MB
+        else:
+            vram_alloc = 0.0
+            vram_resv = 0.0
+        self._timeline.append(
+            MemoryTimelineEntry(
+                elapsed_s=elapsed,
+                rss_mb=snap.rss / _MB,
+                pss_mb=snap.pss / _MB,
+                private_clean_mb=snap.private_clean / _MB,
+                private_dirty_mb=snap.private_dirty / _MB,
+                shared_clean_mb=snap.shared_clean / _MB,
+                shared_dirty_mb=snap.shared_dirty / _MB,
+                anonymous_mb=snap.anonymous / _MB,
+                swap_mb=snap.swap / _MB,
+                uss_mb=snap.uss / _MB,
+                vram_allocated_mb=vram_alloc,
+                vram_reserved_mb=vram_resv,
+            )
+        )
+
+    def sample_vram(self) -> None:
+        """Sample VRAM from the main thread and update the last timeline entry.
+
+        CUDA memory queries must run on the main thread to avoid interfering
+        with the allocator.  Call this after ``stop()`` or between iterations.
+        """
+        if not self._timeline or not self._has_cuda:
+            return
+        self._timeline[-1] = replace(
+            self._timeline[-1],
+            vram_allocated_mb=torch.cuda.memory_allocated() / _MB,
+            vram_reserved_mb=torch.cuda.memory_reserved() / _MB,
+        )
+
+    def start(self) -> None:
+        self._timeline.clear()
+        self._start_time = time.perf_counter()
+        self._sample()
+        self._stop.clear()
+        self._thread = threading.Thread(target=self._poll, daemon=True)
+        self._thread.start()
+
+    def _poll(self) -> None:
+        while not self._stop.is_set():
+            self._sample()
+            self._stop.wait(self._interval)
+
+    def stop(self) -> list[MemoryTimelineEntry]:
+        """Stop polling and return the full timeline."""
+        self._stop.set()
+        if self._thread:
+            self._thread.join()
+        self._sample()
+        return self._timeline
+
+    @property
+    def start_time(self) -> float:
+        """The ``perf_counter`` origin used by the timeline."""
+        return self._start_time
+
+    @property
+    def peak_rss_mb(self) -> float:
+        return max((e.rss_mb for e in self._timeline), default=0.0)
+
+    @property
+    def peak_uss_mb(self) -> float:
+        return max((e.uss_mb for e in self._timeline), default=0.0)
+
+    @property
+    def peak_uss_entry(self) -> MemoryTimelineEntry | None:
+        """Return the timeline sample where USS was highest."""
+        return max(self._timeline, key=lambda e: e.uss_mb) if self._timeline else None
+
+
+# ---------------------------------------------------------------------------
+# CUDA event timer (timing mode)
+# ---------------------------------------------------------------------------
+
+
+@contextmanager
+def cuda_timer(nvtx_label: str = "") -> Generator[dict[str, float], None, None]:
+    """Time a CUDA region using CUDA events (primary) and wall clock (secondary).
+
+    Uses CUDA events for GPU timing — no synchronize() between start and end events.
+    A synchronize() before start drains any prior queued work. end.synchronize()
+    waits only for the end event.
+
+    If *nvtx_label* is provided, the region is also marked with an NVTX range
+    for nsys profiling.
+    """
+    result: dict[str, float] = {}
+    torch.cuda.synchronize()
+
+    if nvtx_label:
+        torch.cuda.nvtx.range_push(nvtx_label)
+
+    start_event = torch.cuda.Event(enable_timing=True)
+    end_event = torch.cuda.Event(enable_timing=True)
+
+    wall_start = time.perf_counter()
+    start_event.record()
+    try:
+        yield result
+    finally:
+        end_event.record()
+        wall_end = time.perf_counter()
+
+        if nvtx_label:
+            torch.cuda.nvtx.range_pop()
+
+        end_event.synchronize()
+        result["gpu_ms"] = start_event.elapsed_time(end_event)
+        result["wall_ms"] = (wall_end - wall_start) * 1000.0
+
+
+# ---------------------------------------------------------------------------
+# Git / GPU metadata
+# ---------------------------------------------------------------------------
+
+
+@functools.lru_cache(maxsize=1)
+def _git_info() -> tuple[str, str]:
+    """Return (commit_hash, branch_name) from git, or empty strings on failure."""
+    try:
+        commit = subprocess.check_output(
+            ["git", "rev-parse", "--short", "HEAD"],
+            stderr=subprocess.DEVNULL,
+            text=True,
+        ).strip()
+        branch = subprocess.check_output(
+            ["git", "rev-parse", "--abbrev-ref", "HEAD"],
+            stderr=subprocess.DEVNULL,
+            text=True,
+        ).strip()
+    except (subprocess.CalledProcessError, FileNotFoundError):
+        return "", ""
+    return commit, branch
+
+
+# ---------------------------------------------------------------------------
+# Internal measurement paths
+# ---------------------------------------------------------------------------
+
+
+def _run_timing(
+    fn: Callable[[], Any],
+    *,
+    warmup: int,
+    repeats: int,
+    bench_type: str,
+    on_loop_start: Callable[[], None] | None,
+    on_iter_end: Callable[[], None] | None,
+) -> tuple[list[float], list[float]]:
+    """Timing-only path: CUDA events + wall clock. No memory measurement."""
+    for _ in range(warmup):
+        fn()
+
+    if on_loop_start:
+        on_loop_start()
+
+    all_wall: list[float] = []
+    all_gpu: list[float] = []
+
+    for _ in range(repeats):
+        with cuda_timer(nvtx_label=bench_type) as t:
+            fn()
+        all_gpu.append(t["gpu_ms"])
+        all_wall.append(t["wall_ms"])
+        if on_iter_end:
+            on_iter_end()
+
+    return all_wall, all_gpu
+
+
+def _run_memory(
+    fn: Callable[[], Any],
+    *,
+    warmup: int,
+    repeats: int,
+    on_loop_start: Callable[[], None] | None,
+    on_iter_end: Callable[[], None] | None,
+) -> tuple[list[MemoryTimelineEntry], float, float]:
+    """Memory-only path: smaps_rollup timeline + VRAM. No timing wrappers.
+
+    Returns (timeline, peak_vram_mb, start_time) where start_time is
+    the ``perf_counter`` origin for aligning wall-clock region annotations.
+    """
+    for _ in range(warmup):
+        fn()
+
+    torch.cuda.reset_peak_memory_stats()
+
+    if on_loop_start:
+        on_loop_start()
+
+    tracker = MemoryTimelineTracker()
+    tracker.start()
+
+    try:
+        for _ in range(repeats):
+            fn()
+            tracker.sample_vram()  # VRAM queries on main thread only
+            if on_iter_end:
+                on_iter_end()
+    finally:
+        timeline = tracker.stop()
+        tracker.sample_vram()  # final VRAM snapshot
+    peak_vram_mb = torch.cuda.max_memory_allocated() / _MB if torch.cuda.is_available() else 0.0
+    return timeline, peak_vram_mb, tracker.start_time
+
+
+# ---------------------------------------------------------------------------
+# Public API
+# ---------------------------------------------------------------------------
+
+
+def benchmark(
+    fn: Callable[[], Any],
+    *,
+    warmup: int = 2,
+    repeats: int = 5,
+    bench_type: str = "",
+    model_name: str = "",
+    bench_mode: str = "timing",
+    on_loop_start: Callable[[], None] | None = None,
+    on_iter_end: Callable[[], None] | None = None,
+    **extra: Any,  # noqa: ANN401
+) -> BenchmarkResult:
+    """Run ``fn`` with warmup iterations, then measure ``repeats`` iterations.
+
+    *bench_mode* selects the measurement strategy:
+    - ``"timing"``: CUDA events + wall clock. No memory tracking.
+    - ``"memory"``: smaps_rollup timeline + VRAM. No timing wrappers.
+
+    Returns a :class:`BenchmarkResult` populated for the chosen mode.
+    """
+    if bench_mode not in ("timing", "memory"):
+        msg = f"bench_mode must be 'timing' or 'memory', got {bench_mode!r}"
+        raise ValueError(msg)
+    if repeats < 1:
+        msg = f"repeats must be >= 1, got {repeats}"
+        raise ValueError(msg)
+    if warmup < 0:
+        msg = f"warmup must be >= 0, got {warmup}"
+        raise ValueError(msg)
+
+    git_commit, git_branch = _git_info()
+    gpu_name = torch.cuda.get_device_name() if torch.cuda.is_available() else ""
+
+    # Extract known extra fields
+    num_steps = int(extra.pop("num_steps", 0))
+    step_times_gpu_ms = list(extra.pop("step_times_gpu_ms", ()))
+    dtype = extra.pop("dtype", "")
+    resolution = extra.pop("resolution", "")
+
+    common = {
+        "bench_type": bench_type,
+        "model_name": model_name,
+        "bench_mode": bench_mode,
+        "num_steps": num_steps,
+        "step_times_gpu_ms": step_times_gpu_ms,
+        "git_commit": git_commit,
+        "git_branch": git_branch,
+        "gpu_name": gpu_name,
+        "dtype": dtype,
+        "resolution": resolution,
+        "timestamp": datetime.now(tz=timezone.utc).isoformat(),
+        "extra": extra,
+    }
+
+    if bench_mode == "timing":
+        all_wall, all_gpu = _run_timing(
+            fn,
+            warmup=warmup,
+            repeats=repeats,
+            bench_type=bench_type,
+            on_loop_start=on_loop_start,
+            on_iter_end=on_iter_end,
+        )
+        return BenchmarkResult(
+            all_wall_times_ms=all_wall,
+            all_gpu_times_ms=all_gpu,
+            **common,
+        )
+
+    timeline, peak_vram_mb, mem_start_time = _run_memory(
+        fn,
+        warmup=warmup,
+        repeats=repeats,
+        on_loop_start=on_loop_start,
+        on_iter_end=on_iter_end,
+    )
+
+    # Derive peak values from the snapshot where USS was highest.
+    peak_entry = max(timeline, key=lambda e: e.uss_mb) if timeline else None
+
+    # Store the timeline origin so instrument wall-clock regions can be aligned.
+    extra["memory_start_time"] = mem_start_time
+
+    return BenchmarkResult(
+        peak_vram_mb=peak_vram_mb,
+        peak_uss_mb=peak_entry.uss_mb if peak_entry else 0.0,
+        peak_rss_mb=peak_entry.rss_mb if peak_entry else 0.0,
+        peak_non_evictable_mb=peak_entry.private_dirty_mb if peak_entry else 0.0,
+        memory_timeline=timeline,
+        **common,
+    )

ltx-2-internal/packages/ltx-bench/src/ltx_bench/instrument.py

@@ -0,0 +1,375 @@
+"""Per-stage pipeline instrumentation using CUDA events.
+
+Instruments production functions by monkey-patching them at test time.
+Zero footprint in production code — all patching happens in benchmark tests
+and is fully restored on ``deactivate()``.
+"""
+
+from __future__ import annotations
+
+import functools
+import re
+import sys
+import time
+from collections.abc import Callable
+from contextlib import contextmanager
+from typing import TYPE_CHECKING, Any, Generator
+
+if TYPE_CHECKING:
+    from ltx_bench.core import BenchmarkResult
+
+import torch
+
+from ltx_bench.nvtx import nvtx_range
+
+
+class BenchInstrument:
+    """Singleton that collects per-stage CUDA event timings when active.
+
+    **Lifecycle:**
+    1. ``activate()``
+    2. ``patch()`` / ``patch_tqdm()`` — explicit per-target
+    3. Run pipeline
+    4. ``deactivate()`` — restores all patches
+    5. ``collect()`` or ``collect_stages()``
+
+    **Threading:** Not thread-safe.  All instrumented code must run on a single
+    thread and (for correct CUDA event timing) on the default CUDA stream.
+
+    **Warmup/repeats:** When instrumentation is active, the benchmarked function
+    must be called exactly once (``warmup=0, repeats=1``).
+
+    Usage::
+
+        instrument.activate()
+        instrument.patch(some_func, "region_name")
+        instrument.patch_tqdm()
+        try:
+            result = bench_run(fn, warmup=0, repeats=1, ...)
+        finally:
+            instrument.deactivate()
+
+        instrument.collect_stages(result)
+    """
+
+    def __init__(self) -> None:
+        self._active = False
+        self._bench_mode = "timing"
+        self._regions: list[tuple[str, str | None, torch.cuda.Event, torch.cuda.Event]] = []
+        self._wall_regions: list[tuple[str, float, float]] = []
+        self._name_counts: dict[str, int] = {}
+        self._current_region: str | None = None
+        self._patches: list[tuple[object, str, Any]] = []
+        self._aggregate_names: set[str] = set()
+        self.tqdm_stats: dict[str, dict[str, Any]] = {}
+
+    @property
+    def active(self) -> bool:
+        return self._active
+
+    @property
+    def current_region(self) -> str | None:
+        return self._current_region
+
+    def set_bench_mode(self, mode: str) -> None:
+        """Set the bench mode so ``activate()`` uses the correct strategy."""
+        self._bench_mode = mode
+
+    def activate(self) -> None:
+        """Enable instrumentation.  Call ``patch()`` after this.
+
+        In timing mode, records per-region CUDA event timings.
+        In memory mode, records wall-clock region boundaries (no CUDA events).
+        Call ``set_bench_mode()`` before ``activate()`` to select the mode.
+        """
+        self._active = True
+        self._regions.clear()
+        self._wall_regions.clear()
+        self._name_counts.clear()
+        self._current_region = None
+        self._patches.clear()
+        self._aggregate_names.clear()
+        self.tqdm_stats = {}
+
+    def deactivate(self) -> None:
+        """Disable instrumentation and restore all patches."""
+        self._active = False
+        self._current_region = None
+        for obj, attr, original in reversed(self._patches):
+            setattr(obj, attr, original)
+        self._patches.clear()
+
+    # ------------------------------------------------------------------
+    # Patching
+    # ------------------------------------------------------------------
+
+    def patch(self, target: Callable[..., Any], region: str, *, aggregate: bool = False) -> None:
+        """Monkey-patch a function to record a CUDA event region.
+
+        ``target`` is the actual callable — a module-level function or an unbound
+        class method (e.g. ``VideoDecoder.forward``).
+
+        When ``aggregate=True``, repeated invocations under the same parent
+        region (auto-tagged ``region[1]``, ``region[2]``, ...) are collapsed into
+        a single entry by :meth:`collect_stages`. Invocations under different
+        parents are kept separate as ``region[0]``, ``region[1]``, etc.
+
+        The function is patched at its definition site AND at every module that
+        imported it (via ``from X import func``).  All patches are restored on
+        ``deactivate()``.
+        """
+        owner, attr, original = self._resolve_callable(target)
+        if aggregate:
+            self._aggregate_names.add(region)
+
+        @functools.wraps(original)
+        def wrapper(*args: Any, **kwargs: Any) -> Any:  # noqa: ANN401
+            with self.region(region):
+                return original(*args, **kwargs)
+
+        # Preserve staticmethod descriptor when patching class attrs.
+        raw_descriptor = owner.__dict__.get(attr) if isinstance(owner, type) else None
+        to_set: Any = staticmethod(wrapper) if isinstance(raw_descriptor, staticmethod) else wrapper
+
+        # Patch at definition site.
+        setattr(owner, attr, to_set)
+        self._patches.append((owner, attr, raw_descriptor if raw_descriptor is not None else original))
+
+        # Patch at all import locations — including aliases like
+        # ``from pkg import func as alias``.
+        for mod in list(sys.modules.values()):
+            if mod is None or mod is owner:
+                continue
+            try:
+                for name, val in list(vars(mod).items()):
+                    if val is original:
+                        setattr(mod, name, wrapper)
+                        self._patches.append((mod, name, original))
+            except Exception:
+                pass
+
+    def patch_tqdm(self) -> None:
+        """Replace ``tqdm.tqdm`` with a capturing subclass at all import locations.
+
+        Captures denoising step count and it/s into ``tqdm_stats``, keyed by the
+        active ``region()`` name at the time the progress bar closes.
+        """
+        import tqdm as tqdm_module  # noqa: PLC0415
+
+        original_tqdm = tqdm_module.tqdm
+        inst = self
+
+        class _CaptureTqdm(original_tqdm):  # type: ignore[misc]
+            def close(self) -> None:
+                if inst.active and not self.disable and inst.current_region and self.n > 0:
+                    inst.tqdm_stats[inst.current_region] = {
+                        "rate": self.format_dict.get("rate"),
+                        "elapsed": self.format_dict.get("elapsed"),
+                        "total": self.format_dict.get("total"),
+                    }
+                super().close()
+
+        # Patch the tqdm module itself.
+        tqdm_module.tqdm = _CaptureTqdm
+        self._patches.append((tqdm_module, "tqdm", original_tqdm))
+
+        # Patch modules that did ``from tqdm import tqdm`` (including aliases).
+        for mod in list(sys.modules.values()):
+            if mod is None or mod is tqdm_module:
+                continue
+            try:
+                for name, val in list(vars(mod).items()):
+                    if val is original_tqdm:
+                        setattr(mod, name, _CaptureTqdm)
+                        self._patches.append((mod, name, original_tqdm))
+            except Exception:
+                pass
+
+    # ------------------------------------------------------------------
+    # Region recording
+    # ------------------------------------------------------------------
+
+    @contextmanager
+    def region(self, name: str) -> Generator[None, None, None]:
+        """Context manager that records a named region.
+
+        Timing mode: CUDA events for GPU timing.
+        Memory mode: wall-clock timestamps for graph annotation.
+
+        When ``LTX_BENCH_NVTX`` is enabled, each region is also wrapped in an
+        NVTX range so it appears in Nsight Systems timelines.
+        """
+        tag = self._tag(name)
+        prev_region = self._current_region
+        self._current_region = tag
+
+        if self._bench_mode == "memory":
+            wall_start = time.perf_counter()
+            try:
+                with nvtx_range(tag):
+                    yield
+            finally:
+                self._wall_regions.append((tag, wall_start, time.perf_counter()))
+                self._current_region = prev_region
+        else:
+            start = torch.cuda.Event(enable_timing=True)
+            end = torch.cuda.Event(enable_timing=True)
+            start.record()
+            # Append in start order (before entering the region) so that a
+            # region's parent always precedes its children in ``_regions``.
+            self._regions.append((tag, prev_region, start, end))
+            try:
+                with nvtx_range(tag):
+                    yield
+            finally:
+                end.record()
+                self._current_region = prev_region
+
+    # ------------------------------------------------------------------
+    # Collection
+    # ------------------------------------------------------------------
+
+    def collect(self) -> list[tuple[str, str | None, float]]:
+        """Synchronize the GPU and return all region timings as (tag, parent, ms).
+
+        Must be called after ``deactivate()``.
+        """
+        if not self._regions:
+            return []
+
+        torch.cuda.synchronize()
+        return [(tag, parent, start.elapsed_time(end)) for tag, parent, start, end in self._regions]
+
+    def collect_wall_regions(self) -> list[tuple[str, float, float]]:
+        """Return wall-clock region boundaries as (name, start_s, end_s) tuples.
+
+        Only populated in memory mode.  Times are absolute ``perf_counter`` values;
+        subtract the benchmark start time to get elapsed seconds.
+        """
+        return list(self._wall_regions)
+
+    def collect_stages(self, result: BenchmarkResult) -> None:
+        """Collect region timings and tqdm stats into ``result.extra``.
+
+        Aggregates tiled regions (VAE decode, streaming internals) whose names
+        share a prefix into single entries.
+        Must be called after ``deactivate()``.
+
+        In memory mode, also stores wall-clock region boundaries for graph
+        annotation.
+        """
+        if self._bench_mode == "memory":
+            result.extra["stage_wall_regions"] = [
+                {"name": name, "start_s": start, "end_s": end} for name, start, end in self._wall_regions
+            ]
+        else:
+            stages, counts, parents = self._aggregate_regions(self.collect())
+            result.extra["stage_regions"] = stages
+            result.extra["stage_region_counts"] = counts
+            result.extra["stage_region_parents"] = parents
+
+        result.extra["tqdm_stats"] = dict(self.tqdm_stats)
+
+    def _aggregate_regions(
+        self, raw: list[tuple[str, str | None, float]]
+    ) -> tuple[dict[str, float], dict[str, int], dict[str, str | None]]:
+        """Aggregate tagged variants of regions registered with ``aggregate=True``.
+
+        Invocations with the same name but different parent aggregates are
+        kept separate (disambiguated with ``[N]`` suffixes). Unaggregated
+        regions pass through with their original tag.
+        """
+        stages: dict[str, float] = {}
+        counts: dict[str, int] = {}
+        parents: dict[str, str | None] = {}
+
+        agg_bucket: dict[tuple[str, str | None], str] = {}
+        name_counter: dict[str, int] = {}
+        tag_to_display: dict[str, str] = {}
+
+        for tag, parent_tag, ms in raw:
+            parent_display = tag_to_display.get(parent_tag) if parent_tag is not None else None
+            base = re.sub(r"\[\d+\]$", "", tag)
+            prefix = base if base in self._aggregate_names else None
+
+            if prefix is not None:
+                bucket_key = (prefix, parent_display)
+                if bucket_key not in agg_bucket:
+                    n = name_counter.get(prefix, 0)
+                    display = prefix if n == 0 else f"{prefix}[{n}]"
+                    name_counter[prefix] = n + 1
+                    agg_bucket[bucket_key] = display
+                    stages[display] = 0.0
+                    counts[display] = 0
+                    parents[display] = parent_display
+                display = agg_bucket[bucket_key]
+                stages[display] += ms
+                counts[display] += 1
+            else:
+                display = tag
+                stages[display] = ms
+                parents[display] = parent_display
+
+            tag_to_display[tag] = display
+
+        # Drop zero-time aggregate buckets.
+        stages = {k: v for k, v in stages.items() if v > 0.0 or k not in counts}
+        counts = {k: v for k, v in counts.items() if stages.get(k, 0.0) > 0.0}
+        parents = {k: v for k, v in parents.items() if k in stages}
+        return stages, counts, parents
+
+    # ------------------------------------------------------------------
+    # Internal
+    # ------------------------------------------------------------------
+
+    def _tag(self, name: str) -> str:
+        """Return a unique tag, appending [N] for repeated names."""
+        count = self._name_counts.get(name, 0)
+        self._name_counts[name] = count + 1
+        if count == 0:
+            return name
+        return f"{name}[{count}]"
+
+    @staticmethod
+    def _resolve_callable(target: Callable[..., Any]) -> tuple[object, str, Any]:
+        """Resolve a callable to ``(owner, attr_name, original_value)``.
+
+        For module-level functions the owner is the defining module.
+        For class methods (e.g. ``VideoDecoder.forward``) the owner is the class.
+        """
+        module_name = getattr(target, "__module__", None)
+        qualname = getattr(target, "__qualname__", None)
+        if module_name is None or qualname is None:
+            msg = f"target must have __module__ and __qualname__, got {target!r}"
+            raise ValueError(msg)
+        if module_name not in sys.modules:
+            msg = f"Module {module_name!r} for {target!r} not found in sys.modules"
+            raise ValueError(msg)
+
+        mod = sys.modules[module_name]
+        parts = qualname.split(".")
+
+        # Traverse to the owner (module for top-level funcs, class for methods).
+        owner: Any = mod
+        try:
+            for part in parts[:-1]:
+                owner = getattr(owner, part)
+            attr = parts[-1]
+            resolved = getattr(owner, attr)
+        except AttributeError as exc:
+            msg = f"Cannot traverse {qualname!r} in module {module_name!r}: {exc}"
+            raise ValueError(msg) from exc
+
+        if resolved is not target:
+            msg = (
+                f"Resolved {module_name}.{qualname} points to a different object than "
+                f"the target passed to patch(). The callable may have been renamed or "
+                f"replaced since import."
+            )
+            raise ValueError(msg)
+
+        return owner, attr, resolved
+
+
+# Module-level singleton.
+instrument = BenchInstrument()

ltx-2-internal/packages/ltx-bench/src/ltx_bench/memory_graph.py

@@ -0,0 +1,146 @@
+"""Generate memory timeline graphs from benchmark results."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from ltx_bench.core import BenchmarkResult
+
+_STAGE_COLORS = [
+    "#4e79a7",
+    "#f28e2b",
+    "#e15759",
+    "#76b7b2",
+    "#59a14f",
+    "#edc948",
+    "#b07aa1",
+]
+
+
+def _aggregate_regions(
+    raw_regions: list[dict[str, Any]],
+    start_time: float,
+) -> list[tuple[str, float, float]]:
+    """Merge consecutive regions with the same base name into a single span.
+
+    ``vae_video_decode``, ``vae_video_decode[1]``, ... become one
+    ``vae_video_decode`` span covering the full range.
+    """
+    groups: list[tuple[str, float, float]] = []
+    for region in raw_regions:
+        name = region["name"].split("[")[0]  # strip [N] suffix
+        t0 = region["start_s"] - start_time
+        t1 = region["end_s"] - start_time
+        if groups and groups[-1][0] == name:
+            # Extend existing group
+            groups[-1] = (name, groups[-1][1], t1)
+        else:
+            groups.append((name, t0, t1))
+    return groups
+
+
+def _draw_stage_regions(
+    axes: list[Any],
+    stage_wall_regions: list[dict[str, Any]],
+    start_time: float,
+) -> None:
+    """Draw shaded vertical spans and labels for aggregated stage regions."""
+    groups = _aggregate_regions(stage_wall_regions, start_time)
+
+    for i, (_name, t0, t1) in enumerate(groups):
+        color = _STAGE_COLORS[i % len(_STAGE_COLORS)]
+        for ax in axes:
+            ax.axvspan(t0, t1, alpha=0.15, color=color, zorder=0)
+
+    # Labels on the top axis, horizontal, staggered to avoid overlap.
+    ax_top = axes[0]
+    y_levels = [1.02, 1.10]
+    prev_right_data = -float("inf")
+    level_idx = 0
+    for i, (name, t0, t1) in enumerate(groups):
+        color = _STAGE_COLORS[i % len(_STAGE_COLORS)]
+        mid = (t0 + t1) / 2
+        # Estimate label width in data coords: ~0.8 data-units per char at fontsize 7
+        # on a typical 16-inch-wide figure.  This is a rough heuristic.
+        x_range = axes[0].get_xlim()[1] - axes[0].get_xlim()[0] if axes[0].get_xlim()[1] > 0 else 1.0
+        char_width = x_range / 120  # approx chars that fit in the x range
+        half_label = len(name) * char_width / 2
+        label_left = mid - half_label
+        level_idx = (level_idx + 1) % len(y_levels) if label_left < prev_right_data else 0
+        prev_right_data = mid + half_label
+        ax_top.annotate(
+            name,
+            xy=(mid, y_levels[level_idx]),
+            xycoords=("data", "axes fraction"),
+            ha="center",
+            va="bottom",
+            fontsize=7,
+            fontweight="bold",
+            color=color,
+            rotation=0,
+            annotation_clip=False,
+        )
+
+
+def generate_memory_graph(result: BenchmarkResult, output_path: Path) -> Path:
+    """Create a two-panel PNG showing host memory and VRAM over time.
+
+    Top panel: RSS, USS, Anonymous, Private_Clean, Private_Dirty (host).
+    Bottom panel: VRAM allocated + reserved (GPU).
+    Stage regions (if present) are drawn as shaded vertical spans.
+
+    Returns *output_path* for chaining.
+    """
+    import matplotlib  # noqa: PLC0415
+
+    matplotlib.use("Agg")
+    import matplotlib.pyplot as plt  # noqa: PLC0415
+
+    timeline = result.memory_timeline
+    if not timeline:
+        msg = "No memory timeline data to plot"
+        raise ValueError(msg)
+
+    t = [e.elapsed_s for e in timeline]
+
+    fig, (ax_host, ax_vram) = plt.subplots(
+        2,
+        1,
+        figsize=(16, 9),
+        sharex=True,
+        gridspec_kw={"height_ratios": [2, 1]},
+    )
+    fig.suptitle(result.name, fontsize=12, fontweight="bold", y=0.98)
+
+    # --- Host memory panel ---
+    ax_host.plot(t, [e.rss_mb for e in timeline], label="RSS", linewidth=1.5)
+    ax_host.plot(t, [e.uss_mb for e in timeline], label="USS", linewidth=1.5)
+    ax_host.plot(t, [e.anonymous_mb for e in timeline], label="Anonymous", linewidth=1.2, linestyle="--")
+    ax_host.plot(t, [e.private_clean_mb for e in timeline], label="Private_Clean", linewidth=1.0, linestyle=":")
+    ax_host.plot(t, [e.private_dirty_mb for e in timeline], label="Private_Dirty", linewidth=1.0, linestyle="-.")
+    ax_host.set_ylabel("Host Memory (MB)")
+    ax_host.legend(loc="upper right", fontsize=7, framealpha=0.8)
+    ax_host.grid(alpha=0.3)
+
+    # --- VRAM panel ---
+    ax_vram.plot(t, [e.vram_allocated_mb for e in timeline], label="Allocated", linewidth=1.5)
+    ax_vram.plot(t, [e.vram_reserved_mb for e in timeline], label="Reserved", linewidth=1.2, linestyle="--")
+    ax_vram.set_ylabel("VRAM (MB)")
+    ax_vram.set_xlabel("Elapsed (s)")
+    ax_vram.legend(loc="upper right", fontsize=7, framealpha=0.8)
+    ax_vram.grid(alpha=0.3)
+
+    # --- Stage region annotations ---
+    stage_regions = result.extra.get("stage_wall_regions")
+    start_time = result.extra.get("memory_start_time", 0.0)
+    if stage_regions:
+        _draw_stage_regions([ax_host, ax_vram], stage_regions, start_time)
+
+    fig.subplots_adjust(top=0.84)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    fig.savefig(output_path, dpi=150)
+    plt.close(fig)
+
+    return output_path

ltx-2-internal/packages/ltx-bench/src/ltx_bench/nvtx.py

@@ -0,0 +1,26 @@
+from __future__ import annotations
+
+import os
+from collections.abc import Generator
+from contextlib import contextmanager
+
+import torch
+
+
+def _env_flag_enabled(name: str) -> bool:
+    value = os.getenv(name, "")
+    return value.lower() in {"1", "true", "yes", "on"}
+
+
+@contextmanager
+def nvtx_range(name: str) -> Generator[None, None, None]:
+    """NVTX range marker for nsys profiling; no-op unless LTX_BENCH_NVTX is enabled."""
+    if not _env_flag_enabled("LTX_BENCH_NVTX"):
+        yield
+        return
+
+    torch.cuda.nvtx.range_push(name)
+    try:
+        yield
+    finally:
+        torch.cuda.nvtx.range_pop()