diff --git a/examples/diffusers/Dockerfile.dreamverse b/examples/diffusers/Dockerfile.dreamverse index bec0f90760d6..55495245a219 100644 --- a/examples/diffusers/Dockerfile.dreamverse +++ b/examples/diffusers/Dockerfile.dreamverse @@ -30,6 +30,24 @@ 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, @@ -37,6 +55,7 @@ RUN . /opt/venv/bin/activate \ 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"). @@ -44,9 +63,11 @@ 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 diff --git a/examples/diffusers/ltx23/bake_bench.py b/examples/diffusers/ltx23/bake_bench.py index b63810b9d175..3d3b1e4b2ba4 100644 --- a/examples/diffusers/ltx23/bake_bench.py +++ b/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 diff --git a/examples/diffusers/ltx23/config.py b/examples/diffusers/ltx23/config.py index 89638a576fe6..a2026ce32703 100755 --- a/examples/diffusers/ltx23/config.py +++ b/examples/diffusers/ltx23/config.py @@ -28,12 +28,18 @@ /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: diff --git a/examples/diffusers/patches/flash-attn-cute-fa4-tuple-fix.patch b/examples/diffusers/patches/flash-attn-cute-fa4-tuple-fix.patch new file mode 100644 index 000000000000..eb0b9ff88d08 --- /dev/null +++ b/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 + +