[perf] LTX2 FP4-linear on DGX Spark (GB10/sm_121): ~24% faster denoise at 1080p

examples/inference/basic/ltx2_fp4_ab.py

@@ -0,0 +1,354 @@
+# SPDX-License-Identifier: Apache-2.0
+import json
+import os
+import time
+from collections import OrderedDict
+from pathlib import Path
+
+import torch
+import torch._inductor.config
+from fastvideo import VideoGenerator
+from fastvideo.configs.pipelines.base import PipelineConfig
+from fastvideo.layers.quantization.nvfp4_config import NVFP4Config
+from fastvideo.utils import maybe_download_model
+
+VALIDATION_JSON = (
+    Path(__file__).resolve().parents[2] / "training" / "finetune" / "ltx2" / "validation.json"
+)
+
+# Override with a local snapshot or converted directory when needed, e.g.
+#   export LTX2_MODEL_PATH=/raid/$USER/hf/FastVideo/LTX2-Distilled-Diffusers
+MODEL_ID = os.path.expandvars(
+    os.path.expanduser(os.getenv("LTX2_MODEL_PATH", "FastVideo/LTX2-Distilled-Diffusers"))
+)
+# A/B arm: FP4 linear ON (default) vs bf16 baseline, selected by FP4_LINEAR env.
+_ARM = "fp4" if os.getenv("FP4_LINEAR", "1") == "1" else "bf16"
+OUTPUT_DIR = Path(f"outputs_video/ltx2_fp4_ab/{_ARM}")
+
+# FLASH_ATTN is dead on sm_121; LTX2 ran with SDPA on Spark (catalogue). Override
+# with LTX2_BACKEND if needed.
+os.environ["FASTVIDEO_ATTENTION_BACKEND"] = os.getenv("LTX2_BACKEND", "TORCH_SDPA")
+os.environ["FASTVIDEO_STAGE_LOGGING"] = "1"
+
+# Tune Inductor flags
+config = torch._inductor.config
+config.conv_1x1_as_mm = True  # treat 1x1 convolutions as matrix muls
+config.coordinate_descent_tuning = True
+config.coordinate_descent_check_all_directions = True
+config.epilogue_fusion = False  # do not fuse pointwise ops into matmuls
+
+
+def load_validation_entries(path: Path) -> list[dict]:
+    with path.open(encoding="utf-8") as f:
+        data = json.load(f)
+
+    if isinstance(data, dict) and isinstance(data.get("data"), list):
+        return [entry for entry in data["data"] if isinstance(entry, dict)]
+
+    raise ValueError(f"Unsupported validation format in {path}. Expected {{'data': [...]}}.")
+
+
+def print_stage_breakdown(
+    result: dict,
+    run_idx: int,
+    num_runs: int,
+) -> float | None:
+    logging_info = result.get("logging_info")
+    if logging_info is None:
+        print(f"[{run_idx}/{num_runs}] Stage breakdown unavailable: no logging_info")
+        return None
+
+    stages = getattr(logging_info, "stages", None)
+    if not stages:
+        print(f"[{run_idx}/{num_runs}] Stage breakdown unavailable: no stage timings")
+        return None
+
+    print(f"[{run_idx}/{num_runs}] Stage breakdown:")
+    total = 0.0
+    for stage_name, stage_metrics in stages.items():
+        exec_time = float(stage_metrics.get("execution_time", 0.0))
+        total += exec_time
+        print(f"  - {stage_name}: {exec_time:.3f}s")
+    print(f"  - total(stage sum): {total:.3f}s")
+    return total
+
+
+def extract_sr_forward_latency(
+    result: dict,
+) -> tuple[float | None, list[tuple[str, float]], list[str]]:
+    logging_info = result.get("logging_info")
+    if logging_info is None:
+        return None, [], []
+
+    stages = getattr(logging_info, "stages", None)
+    if not stages:
+        return None, [], []
+
+    stage_names = list(stages.keys())
+    sr_match_substr = os.getenv("FASTVIDEO_SR_LATENCY_STAGE_SUBSTR", "").strip().lower()
+
+    sr_stage_entries: list[tuple[str, float]] = []
+    for stage_name, stage_metrics in stages.items():
+        stage_name_l = stage_name.lower()
+        if sr_match_substr:
+            is_sr_stage = sr_match_substr in stage_name_l
+        else:
+            is_sr_stage = (
+                "srdenoisingstage" in stage_name_l
+                or "sr_denoising" in stage_name_l
+                or "upsample" in stage_name_l
+                or ("refine" in stage_name_l and "denois" in stage_name_l)
+            )
+        if not is_sr_stage:
+            continue
+        exec_time = float(stage_metrics.get("execution_time", 0.0))
+        sr_stage_entries.append((stage_name, exec_time))
+
+    if not sr_stage_entries:
+        return None, [], stage_names
+    return sum(x[1] for x in sr_stage_entries), sr_stage_entries, stage_names
+
+
+def collect_stage_times(
+    result: dict,
+    stage_times: dict[str, list[float]],
+    stage_order: OrderedDict[str, None],
+) -> None:
+    logging_info = result.get("logging_info")
+    if logging_info is None:
+        return
+    stages = getattr(logging_info, "stages", None)
+    if not stages:
+        return
+    for stage_name, stage_metrics in stages.items():
+        stage_order.setdefault(stage_name, None)
+        exec_time = float(stage_metrics.get("execution_time", 0.0))
+        stage_times.setdefault(stage_name, []).append(exec_time)
+
+
+def print_stage_averages(
+    stage_times: dict[str, list[float]],
+    stage_order: OrderedDict[str, None],
+    measured_runs: int,
+) -> None:
+    if measured_runs <= 0:
+        return
+    if not stage_times:
+        print("No stage timings collected for measured runs.")
+        return
+
+    print(f"Average stage times over {measured_runs} measured runs:")
+    total_avg = 0.0
+    for stage_name in stage_order.keys():
+        times = stage_times.get(stage_name, [])
+        if not times:
+            continue
+        avg = sum(times) / len(times)
+        total_avg += avg
+        print(f"  - {stage_name}: {avg:.3f}s")
+    print(f"  - total(stage sum avg): {total_avg:.3f}s")
+
+
+def resolve_refine_upsampler_path(model_root: str) -> Path:
+    root = Path(model_root)
+    candidates = [
+        root / "spatial_upscaler",
+        root / "spatial_upsampler",
+    ]
+
+    env_path = os.getenv("LTX2_REFINE_UPSAMPLER_PATH")
+    if env_path:
+        candidates.insert(0, Path(os.path.expandvars(os.path.expanduser(env_path))))
+
+    for candidate in candidates:
+        if (candidate / "config.json").is_file():
+            return candidate
+
+    checked = "\n".join(f"  - {candidate}" for candidate in candidates)
+    raise FileNotFoundError(
+        "Could not find an LTX2 refine upsampler directory.\n"
+        "Checked:\n"
+        f"{checked}"
+    )
+
+
+def main() -> None:
+    if not VALIDATION_JSON.exists():
+        raise FileNotFoundError(f"Validation file not found: {VALIDATION_JSON}")
+
+    validation_entries = load_validation_entries(VALIDATION_JSON)
+    if not validation_entries:
+        raise ValueError(f"No validation entries found in {VALIDATION_JSON}")
+
+    benchmark_entry = validation_entries[0]
+    prompt = benchmark_entry.get("caption")
+    if not isinstance(prompt, str) or not prompt.strip():
+        raise ValueError("First validation entry is missing a usable caption")
+
+    # Default 2 warmup + 5 measured (mean), matching FastVideo's perf-benchmark
+    # convention (.buildkite/performance-benchmarks/tests/wan-t2v-1.3b.json).
+    num_runs = int(os.getenv("LTX2_NUM_RUNS", "7"))
+    warmup_runs = int(os.getenv("LTX2_WARMUP_RUNS", "2"))
+    avg_window = num_runs - warmup_runs
+    measured_start_idx = max(warmup_runs, num_runs - avg_window)
+
+    model_root = maybe_download_model(MODEL_ID)
+    refine_upsampler_path = resolve_refine_upsampler_path(model_root)
+    print(f"Using refine upsampler: {refine_upsampler_path}")
+
+    pipeline_config = PipelineConfig.from_pretrained(model_root)
+    if os.getenv("FP4_LINEAR", "1") == "1":
+        pipeline_config.dit_config.quant_config = NVFP4Config()
+        print(">>> ARM = FP4 linear ON (NVFP4Config, refine-profile)")
+    else:
+        print(">>> ARM = bf16 baseline (no quant_config)")
+    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
+    # torch.compile OFF by default for the A/B (de-risk inductor on sm_121 + isolate
+    # the FP4-linear effect). Set LTX2_COMPILE=1 for the tuned perf number.
+    _COMPILE = os.getenv("LTX2_COMPILE", "0") == "1"
+    print(f">>> torch.compile = {_COMPILE}")
+    torch_compile_kwargs = {
+        "backend": "inductor",
+        "fullgraph": True,
+        # Uncomment for final best-performance profiling. It is disabled
+        # for faster development iteration because autotuning is slow.
+        # "mode": "max-autotune-no-cudagraphs",
+        "dynamic": False,
+    }
+
+    generator = VideoGenerator.from_pretrained(
+        model_root,
+        num_gpus=1,
+        ltx2_refine_enabled=True,
+        ltx2_refine_upsampler_path=str(refine_upsampler_path),
+        refine_lora_path="",  # keep refine LoRA disabled in this repo's typed adapter
+        ltx2_refine_lora_path="",  # keep refine LoRA disabled for distilled model
+        ltx2_refine_num_inference_steps=2,
+        ltx2_refine_guidance_scale=1.0,
+        ltx2_refine_add_noise=True,
+        pipeline_config=pipeline_config,
+        enable_torch_compile=_COMPILE,
+        enable_torch_compile_text_encoder=_COMPILE,
+        enable_torch_compile_vae=_COMPILE,
+        torch_compile_kwargs=torch_compile_kwargs,
+        torch_compile_kwargs_vae=torch_compile_kwargs,
+        dit_cpu_offload=False,
+        text_encoder_cpu_offload=False,
+        vae_cpu_offload=False,
+        # Tiling ON by default: an untiled high-res decode OOMs the GB10's unified
+        # memory and locks the box. Set LTX2_VAE_TILING=0 only if you know it fits.
+        ltx2_vae_tiling=(os.getenv("LTX2_VAE_TILING", "1") == "1"),
+    )
+
+    run_times: list[float] = []
+    e2e_times: list[float] = []
+    sr_forward_times: list[float] = []
+    non_stage_overhead_times: list[float] = []
+    stage_times: dict[str, list[float]] = {}
+    stage_order: OrderedDict[str, None] = OrderedDict()
+
+    try:
+        for i in range(num_runs):
+            output_path = OUTPUT_DIR / f"output_ltx2_basic_t2v_run_{i + 1}.mp4"
+            if output_path.exists():
+                output_path.unlink()
+                print(f"[{i + 1}/{num_runs}] Removed existing file: {output_path}")
+
+            print(f"[{i + 1}/{num_runs}] Generating: {output_path}")
+            if os.environ.get("FASTVIDEO_STAGE_LOGGING") == "0" and torch.cuda.is_available():
+                torch.cuda.synchronize()
+
+            start = time.perf_counter()
+            result = generator.generate_video(
+                prompt=prompt,
+                output_path=str(output_path),
+                fps=24,
+                seed=10,
+                save_video=True,
+                guidance_scale=1.0,
+                # Start small (env-overridable) so the first A/B can't OOM-lock the
+                # box; scale up to 1088x1920x121 once stable.
+                height=int(os.getenv("LTX2_HEIGHT", benchmark_entry.get("height", 1088))),
+                width=int(os.getenv("LTX2_WIDTH", benchmark_entry.get("width", 1920))),
+                num_frames=int(os.getenv("LTX2_FRAMES", "121")),
+                num_inference_steps=int(os.getenv("LTX2_STEPS", "5")),
+                # image_path="examples/inference/basic/prompt1.png",
+                # ltx2_image_crf=0.0
+            )
+            if os.environ.get("FASTVIDEO_STAGE_LOGGING") == "0":
+                torch.cuda.synchronize()
+
+            elapsed = result.get("generation_time") if isinstance(result, dict) else None
+            e2e_elapsed = result.get("e2e_latency") if isinstance(result, dict) else None
+            if elapsed is None:
+                elapsed = time.perf_counter() - start
+            if e2e_elapsed is None:
+                e2e_elapsed = time.perf_counter() - start
+
+            run_times.append(elapsed)
+            e2e_times.append(e2e_elapsed)
+            print(f"[{i + 1}/{num_runs}] Generation time: {elapsed:.2f}s")
+            print(f"[{i + 1}/{num_runs}] End-to-end latency: {e2e_elapsed:.2f}s")
+
+            if isinstance(result, dict):
+                stage_sum = print_stage_breakdown(result, i + 1, num_runs)
+                if stage_sum is not None:
+                    non_stage_overhead = e2e_elapsed - stage_sum
+                    print(f"[{i + 1}/{num_runs}] Non-stage overhead (e2e - stage sum): {non_stage_overhead:.3f}s")
+                    if i >= measured_start_idx:
+                        non_stage_overhead_times.append(non_stage_overhead)
+
+                sr_forward_total, sr_stage_entries, stage_names = extract_sr_forward_latency(result)
+                if sr_forward_total is None:
+                    print(f"[{i + 1}/{num_runs}] SR forward latency unavailable")
+                    if stage_names:
+                        print(f"    Available stage keys: {', '.join(stage_names)}")
+                        print("    Tip: set FASTVIDEO_SR_LATENCY_STAGE_SUBSTR=<substring> to match your SR stage key.")
+                else:
+                    print(f"[{i + 1}/{num_runs}] SR forward latency: {sr_forward_total:.3f}s")
+                    for sr_stage_name, sr_exec_time in sr_stage_entries:
+                        print(f"    - {sr_stage_name}: {sr_exec_time:.3f}s")
+                    if i >= measured_start_idx:
+                        sr_forward_times.append(sr_forward_total)
+
+                if i >= measured_start_idx:
+                    collect_stage_times(result, stage_times, stage_order)
+
+        measured_times = run_times[measured_start_idx:]
+        avg_time = sum(measured_times) / len(measured_times)
+        print(
+            f"Average video generation time over {len(measured_times)} runs "
+            f"(runs {measured_start_idx + 1}-{len(run_times)}, skipping first {warmup_runs} warmup runs): "
+            f"{avg_time:.2f}s"
+        )
+
+        measured_e2e_times = e2e_times[measured_start_idx:]
+        avg_e2e_time = sum(measured_e2e_times) / len(measured_e2e_times)
+        print(
+            f"Average end-to-end latency over {len(measured_e2e_times)} runs "
+            f"(runs {measured_start_idx + 1}-{len(e2e_times)}, skipping first {warmup_runs} warmup runs): "
+            f"{avg_e2e_time:.2f}s"
+        )
+
+        if sr_forward_times:
+            avg_sr_forward = sum(sr_forward_times) / len(sr_forward_times)
+            print(f"Average SR forward latency over {len(sr_forward_times)} runs: {avg_sr_forward:.3f}s")
+        else:
+            print("Average SR forward latency unavailable (no SR stages matched).")
+
+        print_stage_averages(stage_times, stage_order, len(measured_times))
+
+        if non_stage_overhead_times:
+            avg_non_stage_overhead = sum(non_stage_overhead_times) / len(non_stage_overhead_times)
+            print(
+                "Average non-stage overhead over "
+                f"{len(non_stage_overhead_times)} measured runs: {avg_non_stage_overhead:.3f}s"
+            )
+        else:
+            print("Average non-stage overhead unavailable (no stage timings).")
+    finally:
+        generator.shutdown()
+
+
+if __name__ == "__main__":
+    main()