Run FLUX.1-dev three times faster

TL;DR

Modal walks through how they got FLUX.1-dev inference under three seconds per 1024² image — a 3× speedup over the diffusers baseline — to make a self-hosted endpoint competitive with proprietary APIs on both latency and price. The recipe is two layers: ~1.5× from “standard” Torch-compile + fused QKV + channels-last memory layout, then another ~2× from First Block Caching (a TEACache-style approximate activation cache that skips diffusion steps whose first transformer block doesn’t change much). Cold-start regression from max-autotune is then clawed back ~30× via Torch’s FX/Inductor/Triton caches plus Modal’s memory snapshots. The exact same playbook should port to Flux.2 and other large DiTs.

Key claims

Baseline FLUX.1-dev in diffusers + bf16 generates 1024² in ~6.75 s; the target to match proprietary APIs is “under three seconds” [§Implement the baseline].
Standard optimizations alone (torch.compile with max-autotune-no-cudagraphs, fuse_qkv_projections, memory_format=torch.channels_last) deliver a ~1.5× speedup, “driven mostly by the Torch compiler” [§Apply standard optimizations].
Fusing Q/K/V into one matmul QKV = X @ W_qkv exposes more parallelism and removes the compiler’s need to verify that three reads of X alias the same tensor [§Expose more parallelism to the GPU with fused QKV].
First Block Caching (from ParaAttention, based on the TEACache paper) gives an additional ~2× by skipping diffusion steps whose post-first-block residual change is below a threshold; Modal uses residual_diff_threshold=0.12 (default 0.08), trading “slight changes in results” for speed [§Apply approximate caching].
The approximate-cache threshold gives a “smoother tradeoff between performance improvement and quality degradation than other techniques that do the same, like quantization” [§Apply approximate caching].
max-autotune blows up boot time from seconds to tens of minutes; persisting TORCHINDUCTOR_FX_GRAPH_CACHE, CUDA_CACHE_PATH, TORCHINDUCTOR_CACHE_DIR, TRITON_CACHE_DIR on a shared volume + Modal memory snapshots cuts cold-start ~30× [§Cut cold start latency by 30x].
Torch’s “megacache” (whole-graph) was kept in the pipeline but did not contribute a large speedup at the time of writing [§Cut cold start latency by 30x].

Method

The pipeline is the standard diffusers.FluxPipeline in bf16. Modal applies four layers of optimization:

1. Torch compile. Both pipe.transformer and pipe.vae.decode are wrapped with torch.compile(mode="max-autotune-no-cudagraphs", dynamic=True). Inductor config tweaks: conv_1x1_as_mm, coordinate_descent_tuning, coordinate_descent_check_all_directions, shape_padding, epilogue_fusion=False. A dummy 1024² forward triggers compile on container enter.

2. Fused QKV. pipe.transformer.fuse_qkv_projections() and same on the VAE — concatenates W_q, W_k, W_v so attention becomes one big matmul against X instead of three.

3. Channels-last layout. .to(memory_format=torch.channels_last) on transformer and VAE. Reorders feature-map tensors so all channels for a given (h, w) are contiguous, improving locality for spatially-local-but-channel-global ops like convolutions.

4. First Block Caching. para_attn.first_block_cache.diffusers_adapters.apply_cache_on_pipe(pipe, residual_diff_threshold=0.12). At each diffusion timestep, the model runs the first transformer block; if the residual delta is below threshold, the rest of the forward pass is skipped and the previous step’s output is reused. This is approximate — output diverges slightly from a no-cache run, but visually the trade is favourable up to 0.12.

5. Cold-start mitigation. max-autotune profiles many kernel variants at compile time, costing tens of minutes per cold start. Modal points Inductor’s FX graph cache, Triton’s cache, CUDA’s JIT cache, and the megacache to a persistent modal.Volume (/cache), so the second container’s compile reuses artifacts. Memory Snapshots are then taken with @modal.enter(snap=True) after from_pretrained on CPU; @modal.enter(snap=False) later moves to CUDA. Snapshot turns the import torch + weight-load IO into a single mmap.

Results

Baseline diffusers FLUX.1-dev (bf16, 1024², averaged over varied prompts): ~6.75 s per image.
After standard optimizations (compile + fused QKV + channels-last): ~1.5× faster (~4.5 s implied).
After adding First Block Caching at threshold 0.12: another ~2×, putting the system under the 3-second target — i.e. ~3× end-to-end vs. baseline.
Cold start: ~30× reduction via Inductor/Triton caches + memory snapshots, bringing max-autotune from “tens of minutes” back to something compatible with autoscaling.
The post claims this is enough to “match or beat” proprietary FLUX serving providers on price, given Modal’s autoscaling.

No FID / image-quality numbers are reported; quality is evaluated qualitatively (animations comparing denoising-step evolution at matched wall-clock rates). The 0.12 vs 0.08 caching threshold choice is also a qualitative judgement (“images looked better”).

Why it’s interesting

The recipe is directly transferable: any large DiT served from diffusers (Flux.2, Wan-family, SD3.5, etc.) gets the same ~1.5× from compile + fused-QKV + channels-last basically for free, and the First Block Caching idea generalises to any timestep-iterative model where adjacent steps are highly correlated — which is most diffusion / flow-matching models in practice. For Luma’s video DiTs specifically, the per-step correlation is even higher than for images (motion is smooth), so the FBC speedup ceiling is plausibly larger. The cold-start playbook (persist Inductor/Triton caches on a shared volume, then snapshot post-from_pretrained CPU state) is also a reasonable template for any team running max-autotune in production behind an autoscaler. The piece is a blog post, not a paper, so there’s no ablation against quantization, distillation, or step-reduction — those remain orthogonal axes worth combining.