perf: LTX-2.3 SPEED on FA4 attention (10.4s to 7.3s)

examples/diffusers/Dockerfile.dreamverse

@@ -30,23 +30,44 @@ RUN apt-get update \
 RUN . /opt/venv/bin/activate \
  && uv pip install --prerelease=allow ai-dynamo==1.1.1
 
+# FA4 (flash-attn-cute) source swap. dreamverse's baked-in install is the
+# XOR-op/flash-attention fork FastVideo's June-26 pyproject pinned (commit
+# f4c68c7f) -- stale against cutlass-dsl 4.5+, which moved/renamed APIs the
+# fork's kernels call (cute.core.ThrMma -> cute.ThrMma/cute.atom.ThrMma;
+# cute.make_fragment -> make_rmem_tensor). FastVideo hit this exact break and
+# fixed it 2026-07-06 (hao-ai-lab/FastVideo#1564): they dropped the fork and
+# now pin flash-attn-4 straight from the real upstream, Dao-AILab/flash-attention,
+# at a CuTe-DSL-4.6-compatible commit (= upstream PR #2648). Reproduce that
+# exact pin -- do not bump to a newer/moving commit without re-validating.
+ARG FA4_CUTE_REF=82d6441eec5d4dfec120153db2c0145ae855a083
+RUN . /opt/venv/bin/activate \
+ && pip uninstall -y flash-attn-cute \
+ && git clone https://github.com/Dao-AILab/flash-attention.git /tmp/fa-upstream \
+ && cd /tmp/fa-upstream && git checkout ${FA4_CUTE_REF} \
+ && cd / \
+ && uv pip install --no-build-isolation /tmp/fa-upstream/flash_attn/cute \
+ && rm -rf /tmp/fa-upstream
+
 # deepinfra patches, applied to the INSTALLED (non-editable) fastvideo package.
 # The patches are git-format (a/fastvideo/... b/fastvideo/...); -p1 from
 # site-packages strips a/ -> fastvideo/... . dreamverse's fastvideo is @4f3ad3f6,
 # the exact commit these patches were authored against, so they apply clean.
 COPY patches/ltx23_gpu_worker_megacache.patch /tmp/megacache.patch
 COPY patches/x264-threads-cap.patch /tmp/x264-threads-cap.patch
 COPY patches/pyav-scenecut.patch /tmp/pyav-scenecut.patch
+COPY patches/flash-attn-cute-fa4-tuple-fix.patch /tmp/flash-attn-cute-fa4-tuple-fix.patch
 # NOTE: verify by grepping the patched source, NOT by importing fastvideo --
 # importing pulls in triton, which fails to init its driver in a GPU-less build
 # ("0 active drivers").
 RUN SP=/opt/venv/lib/python3.12/site-packages \
  && patch -p1 -d "$SP" --forward --fuzz=0 < /tmp/megacache.patch \
  && patch -p1 -d "$SP" --forward --fuzz=0 < /tmp/x264-threads-cap.patch \
  && patch -p1 -d "$SP" --forward --fuzz=0 < /tmp/pyav-scenecut.patch \
+ && patch -p1 -d "$SP" --forward --fuzz=0 < /tmp/flash-attn-cute-fa4-tuple-fix.patch \
  && grep -q 'torch.compiler.load_cache_artifacts' "$SP/fastvideo/worker/gpu_worker.py" \
  && grep -q 'FASTVIDEO_X264_THREADS' "$SP/fastvideo/entrypoints/video_generator.py" \
  && grep -q 'FASTVIDEO_X264_SCENECUT' "$SP/fastvideo/entrypoints/video_generator.py" \
+ && grep -q ')\[:2\]' "$SP/fastvideo/attention/utils/flash_attn_cute.py" \
  && echo "patches applied + verified"
 
 ENV FASTVIDEO_VIDEO_CODEC=libx264

examples/diffusers/ltx23/bake_bench.py

@@ -12,12 +12,20 @@
 logging.basicConfig(level=logging.INFO, format="%(message)s")
 sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
 from lib.pool import SubprocessPool
+from ltx23.config import profile_attention_backend
 
 OUT = os.environ.get("BENCH_OUT", "/tmp/bench-out")
 os.makedirs(OUT, exist_ok=True)
 W, H, NF, FPS, STEPS, GS = 1920, 1088, 121, 24, 5, 1.0
 SHAPE = "%dx%d@%df" % (W, H, NF)
-MAX_WARM_S = float(os.environ.get("LTX23_MAX_WARM_S", "14"))
+# Derive from the profile (same source as ltx23/worker.py + warmup.py) instead
+# of hardcoding a backend -- a hardcoded value silently drifts from whatever
+# the served profile actually uses, baking the WRONG kernels into Mega-Cache.
+PROFILE = os.environ.get("LTX23_PROFILE", "speed")
+ATTENTION_BACKEND = os.environ.get(
+    "FASTVIDEO_ATTENTION_BACKEND", profile_attention_backend(PROFILE)
+)
+MAX_WARM_S = float(os.environ.get("LTX23_MAX_WARM_S", "10"))
 I2V_IMAGE = os.environ.get("LTX23_I2V_IMAGE", "/work/joy_nordic_woman_mid.png")
 T2V_PROMPTS = [
     ("musician", "A street musician in her thirties sings and strums an acoustic guitar on a sunny city sidewalk, natural human face, photorealistic, sharp focus."),
@@ -39,7 +47,7 @@ def req(tag: str, prompt: str, **extra) -> dict:
 async def main(mode: str) -> int:
     pool = SubprocessPool(
         model_path=os.environ.get("LTX23_MODEL_PATH", "/models/ltx-2.3-distilled-diffusers"),
-        num_gpus=1, enable_optimizations=True, attention_backend="TORCH_SDPA",
+        num_gpus=1, enable_optimizations=True, attention_backend=ATTENTION_BACKEND,
         model_factory_dotted="ltx23.factory:load_model", model_label="ltx23-bake",
     )
     warm_times: list[tuple[str, float]] = []
@@ -62,11 +70,17 @@ async def main(mode: str) -> int:
         print("BAKE_BENCH_DONE mode=%s" % mode, flush=True)
     finally:
         await pool.shutdown()
-    slow = [(tag, dt) for tag, dt in warm_times[1:] or warm_times if dt > MAX_WARM_S]
-    if slow:
-        print("LATENCY GATE FAILED (> %.0fs warm): %s -- do NOT ship this bake; "
-              "check the preset boot log and recent env/package changes." % (MAX_WARM_S, slow), flush=True)
-        return 1
+    # Latency gate only applies to the t2v pass: there the first gen is the cold
+    # compile and the rest (warm_times[1:]) are genuinely warm cache hits. The
+    # i2v pass has a single gen that is ALWAYS a cold compile (it exists to
+    # extend the Mega-Cache blob with i2v graphs), so there is no warm sample to
+    # gate -- skip it there to avoid a guaranteed false-positive failure.
+    if mode == "t2v":
+        slow = [(tag, dt) for tag, dt in warm_times[1:] if dt > MAX_WARM_S]
+        if slow:
+            print("LATENCY GATE FAILED (> %.0fs warm): %s -- do NOT ship this bake; "
+                  "check the preset boot log and recent env/package changes." % (MAX_WARM_S, slow), flush=True)
+            return 1
     return 0
 
 

examples/diffusers/ltx23/config.py

@@ -28,12 +28,18 @@
 <model>/spatial_upscaler, vae_tiling False, the Blackwell Inductor
 knobs (factory.py: shape_padding=False etc.), LD_LIBRARY_PATH unset, cu128 env.
 
-Attention backend is per-profile (see profile_attention_backend): QUALITY uses
-FLASH_ATTN (FastVideo's SM100/SM103 fast kernels, what prod runs); SPEED uses
-TORCH_SDPA -- the config we validated end-to-end on B200 (~10s, quality OK'd by
-Johan). FastVideo's own streaming_demo speed path uses FLASH_ATTN (faster,
-~4.5s), but SPEED+FLASH_ATTN is NOT something we have validated, so we ship the
-attention we tested. Revisit with worker.py --attention-backend.
+Attention backend is per-profile (see profile_attention_backend): both QUALITY
+and SPEED use FLASH_ATTN -- FastVideo's SM100/SM103 FA4 kernels, matching their
+actual streaming_demo speed recipe. SPEED previously shipped TORCH_SDPA as a
+stand-in (~10.4s, quality OK'd by Johan) because FA4 (flash-attn-cute) didn't
+build: the pinned XOR-op fork was stale against cutlass-dsl 4.5+ API moves
+(cute.core.ThrMma, cute.make_fragment renamed/moved). FastVideo hit the same
+break and fixed it 2026-07-06 (hao-ai-lab/FastVideo#1564): they dropped the fork
+and now pin flash-attn-4 straight from Dao-AILab/flash-attention upstream at a
+CuTe-DSL-4.6-compatible commit. Reproducing that exact fix (+ FastVideo's own
+matching `[:2]` tuple-unpack fix in their flash_attn_cute.py wrapper) gets FA4
+working end-to-end: ~7.3s warm, quality OK'd by Johan as on par or better than
+TORCH_SDPA. Override with worker.py --attention-backend to revisit TORCH_SDPA.
 
 Changing any of these requires a coordinated re-bake (the compile cache is keyed
 on the kwargs + shape). Do NOT hand-edit values away from the FastVideo source.
@@ -130,15 +136,18 @@ def profile_uses_nvfp4(profile: str) -> bool:
 def profile_attention_backend(profile: str) -> str:
     """Attention backend per profile.
 
-    QUALITY -> FLASH_ATTN (FastVideo's SM100/SM103 fast kernels; what prod runs).
-    SPEED   -> TORCH_SDPA: the config validated end-to-end on B200 (~10s,
-    quality OK'd). FastVideo's streaming_demo speed path actually uses FLASH_ATTN
-    (faster, ~4.5s), but we have NOT validated SPEED+FLASH_ATTN ourselves, so we
-    ship the attention we tested. Bake (warmup) and serve (worker) both derive
-    the default from here so the compile cache always matches; override with
-    worker.py --attention-backend to revisit FLASH_ATTN.
+    Both QUALITY and SPEED -> FLASH_ATTN (FastVideo's SM100/SM103 FA4 kernels),
+    matching FastVideo's actual recipes. Validated for SPEED 2026-07-08: ~7.3s
+    warm (down from TORCH_SDPA's ~10.4s), quality OK'd by Johan as on par or
+    better. Requires the flash-attn-4 install pinned to a CuTe-DSL-4.6-compatible
+    commit (see Dockerfile.dreamverse) + the matching fastvideo wrapper patch
+    (patches/flash-attn-cute-fa4-tuple-fix.patch) -- without both, FA4 fails to
+    import/run and this setting would silently break the worker. Bake (warmup)
+    and serve (worker) both derive the default from here so the compile cache
+    always matches; override with worker.py --attention-backend to revisit
+    TORCH_SDPA.
     """
-    return "TORCH_SDPA" if (profile or "quality").strip().lower() == "speed" else "FLASH_ATTN"
+    return "FLASH_ATTN"
 
 
 def profile_default_optimizations(profile: str) -> bool:

examples/diffusers/patches/flash-attn-cute-fa4-tuple-fix.patch

@@ -0,0 +1,22 @@
+diff --git a/fastvideo/attention/utils/flash_attn_cute.py b/fastvideo/attention/utils/flash_attn_cute.py
+index 99cb5c7..38d2d88 100644
+--- a/fastvideo/attention/utils/flash_attn_cute.py
++++ b/fastvideo/attention/utils/flash_attn_cute.py
+@@ -37,7 +37,7 @@ def _flash_attn_cute_forward(
+         softcap=0.0,
+         num_splits=1,
+         pack_gqa=None,
+-    )
++    )[:2]
+     return out, lse
+
+
+@@ -129,7 +129,7 @@ def _flash_attn_cute_varlen_forward(
+         softcap=0.0,
+         num_splits=1,
+         pack_gqa=None,
+-    )
++    )[:2]
+     return out, lse
+
+