fix(flux2): estimate working memory for denoise and both VAE directions - #9519
fix(flux2): estimate working memory for denoise and both VAE directions#9519Pfannkuchensack wants to merge 12 commits into
Conversation
Sorry, something went wrong.
The FLUX.2 path called model_on_device() with no working_mem_bytes anywhere, so the model cache reserved only the default device_working_mem_gb and filled the rest of the card with the model. Reference images make that fatal rather than merely tight: their latents are concatenated onto the image stream, so three 1024x1024 references quadruple the attended sequence of a 1024x1024 generation -- 6.5GB of activations against a 3GB reservation. Measured on CUDA in bf16 as peak reserved memory: transformer activations scale linearly at ~0.39 MB/token (no O(seq^2) term, SDPA) and are independent of block count; the FLUX.2 VAE costs ~2170 (decode) / ~1070 (encode) bytes per pixel per element byte, so a 1024x1024 decode peaks at ~4.3GB. Add Flux2DenoiseInvocation._estimate_working_memory() and estimate_vae_working_memory_flux2(), and pass them at every load site so the cache evicts enough to make room instead of hitting the shortfall as an OOM. Closes invoke-ai#9500
JPPhoto
left a comment
There was a problem hiding this comment.
Please fix:
-
invokeai/backend/util/vae_working_memory.py:110: CUDA-linear estimate underestimates FLUX.2 VAE attention on ROCm, where code notes materialized attention. High-resolution encode/decode can still OOM. Effect: fix fails on supported ROCm. Likelihood: Normal ROCm high-resolution use. Recovery: lower resolution; nodes expose no tiling control. Test: Run ROCm Torch 2.10 at 1024/1536px; compare peak reserved memory with estimate.
-
invokeai/app/invocations/flux2_denoise.py:468: Regional prompting passes full float S x S mask via invokeai/backend/flux2/extensions/regional_prompting_extension.py:41; if SDPA selects math fallback, score workspace is quadratic, but estimate adds only mask storage. PyTorch documents backend-dependent SDPA dispatch and math intermediates here. Effect: high-resolution regional prompts can still OOM. Likelihood: Plausible backend/resolution edge. Recovery: disable regional prompting or lower resolution. Test: Measure torch.cuda.max_memory_reserved() for masked 1024/2048px FLUX.2 forwards on Torch 2.7/2.10.
Suggestions:
- Consider backend-specific peak-memory calibration for ROCm and regional-mask attention.
Sorry, something went wrong.
The FLUX.2 working-memory estimates were linear in the sequence length, which holds only while SDPA picks a fused kernel. That is a property of the torch build, not of FLUX.2: ROCm's fused kernels cap the head dim at 128 and reject arbitrary additive masks, so both the VAE's 512-wide mid-block head and the dense S x S bias regional prompting attaches fall through to the math fallback and materialize the score matrix -- ~17GB for a 1536px decode, and heads x S^2 for a masked forward. Rather than assume either way, ask torch: sdpa_score_matrix_bytes() queries can_use_flash/efficient/cudnn_attention for the real head dim, dtype and mask, and adds 13 bytes per score element only when no fused kernel is eligible. Measured on CUDA with SDPBackend.MATH forced: 12.9 bytes/element at 4k tokens, 10.3 at 8k, 9.7 at 16k, identical for bf16, fp16 and fp32 because the fallback's softmax intermediates are always fp32. On CUDA every shape reports fused, so the term is zero and the existing calibration is untouched. Non-CUDA devices keep the fused assumption -- torch exposes no equivalent query there, and guessing would reserve double-digit GB on no evidence.
|
Cant really test the ROCm stuff maybe @lstein can take a look |
Sorry, something went wrong.
JPPhoto
left a comment
There was a problem hiding this comment.
Merge blockers:
- invokeai/backend/util/attention.py:71-74 treats MPS as fused, but Torch 2.7.1 routes MPS SDPA through math, materializing Q @ K^T (dispatch, math path). vae_working_memory.py:158-164 therefore omits about 3.5 GB at 1024px. Effect: MPS Flux2 VAE decode can OOM after cache admission. Likelihood: Normal MPS 1024px+ decode. Recovery: Lower resolution or CPU fallback. Test: Run Torch 2.7.1 MPS 1024px decode; assert nonzero score reservation.
Other findings/issues:
- invokeai/backend/util/attention.py:93-94 treats every probe exception as fused. Probe OOM or backend failure then returns zero score bytes. Effect: Under-reservation and forward OOM. Likelihood: Low free VRAM or backend mismatch. Recovery: Clear cache or lower resolution. Test: Inject torch.cuda.OutOfMemoryError from torch.empty; assert conservative budgeting.
- invokeai/backend/util/attention.py:76-92 probes native Torch eligibility, but Flux2 can use another Diffusers backend through dispatch_attention_fn (Diffusers dispatch). Forced native-math bypasses the probe. Effect: CUDA math attention can materialize O(S^2) while estimate adds zero. Likelihood: Custom attention backend users. Recovery: Restore fused backend or lower resolution. Test: Force _native_math; assert score bytes are included.
Suggestions:
- Instead of assuming non-CUDA is fused, probe the active backend or conservatively charge math on MPS.
- Instead of returning fused on exceptions, budget unknown probe results as math.
Sorry, something went wrong.
The score-matrix term probed torch's CUDA eligibility helpers and read everything else as fused. That was wrong twice over: MPS has no fused SDPA kernel at all and runs the MPSGraph math transcription, so a 1024px VAE decode was admitted ~3.5GB short; and a failed probe returned "fused" too, turning "we don't know" into the one answer that can OOM. Ask `_fused_sdp_choice` instead -- the same dispatch query `scaled_dot_product_attention` runs to pick its kernel. Torch registers it for CPU, CUDA/ROCm and XPU only, so the call raises on exactly the devices that fall through to `math`, and every other failure lands on the conservative side by the same branch. Diffusers models do not reach torch's SDPA directly, so also consult `dispatch_attention_fn`'s active backend: a user on `_native_math` materializes the score matrix on hardware whose probe reports fused. Only the transformer needs this -- the FLUX.2 VAE's mid-block attention still calls SDPA itself through `AttnProcessor2_0` -- and a test pins that asymmetry. On CUDA with the stock backend every one of these terms remains zero.
JPPhoto
left a comment
There was a problem hiding this comment.
Fix:
- invokeai/backend/util/attention.py:78-114: _diffusers_attention_dispatch() caches mutable backend state permanently. After one native estimate, switching Diffusers to _native_math still returns torch, omitting the SxS score allocation and risking OOM. Diffusers supports changing this state via attention_backend() and model.set_attention_backend(). Effect: underestimation during later high-resolution denoise. Likelihood: plausible in long-lived/custom-backend processes. Recovery: clear the private cache or restart. Test: estimate once with native, switch to _native_math without clearing the cache, and verify the score term appears.
Suggestions:
-
Instead of permanently caching the dispatcher result, read the active backend at estimate time or key the cache by backend state.
-
Consider passing the transformer's effective backend into the estimator so per-model overrides cannot disagree with the memory calculation.
Sorry, something went wrong.
`_diffusers_attention_dispatch()` was `lru_cache`d, so the first estimate in a process pinned the answer forever. A switch to `_native_math` after that kept reserving zero for the S x S score matrix -- the exact case the lookup was added to catch. Read it live; it is a dict lookup against an already-imported module, priced once per invocation. The torch probe had the same defect one level down: its answer depends on the global SDPA kernel toggles, which `sdpa_kernel()` and `enable_flash_sdp()` flip at runtime. That probe allocates and dispatches, so it stays cached -- but keyed on the toggles, so a switch invalidates it. Per-model overrides need no plumbing: `set_attention_backend()` stamps its choice onto the process-wide registry as well as onto the model's processors, deliberately, so the estimate sees it without holding the model it is priced ahead of. A test pins that propagation.
|
@JPPhoto _diffusers_attention_dispatch() caching mutable state — the lru_cache is gone; the active backend is read live on every estimate. It's a dict lookup against an already-imported module, priced once per invocation. Your test is in as test_a_backend_switch_is_not_masked_by_an_earlier_estimate: estimate with native, switch to _native_math, estimate again with nothing cleared, assert the score term appears — then assert it disappears again on the way back. Mutation-checked: restoring the decorator turns 5 tests red. While fixing it I found the same defect one level down. The torch probe was cached permanently too, and its answer depends on the global SDPA toggles that sdpa_kernel() and enable_flash_sdp() flip at runtime. That probe allocates and dispatches, so it stays cached — but the toggles are now part of the key. Measured on a 4090: 0 → 12.3 GB inside sdpa_kernel([MATH]) → 0 again, no cache cleared anywhere. Passing the transformer's effective backend into the estimator — I looked at this and deliberately didn't build it, because it would add nothing. ModelMixin.set_attention_backend() stamps the choice onto the model's attention processors and calls _AttentionBackendRegistry.set_active_backend(), with the comment "Important to set the active backend so that it propagates gracefully throughout". So the process-wide lookup already covers per-model overrides — verified against a real ModelMixin, and pinned by test_a_model_level_override_reaches_the_registry so it surfaces if diffusers ever changes that. reset_attention_backend() clears only the processors and leaves the registry pinned, which errs toward over-reserving. That's why the estimate doesn't need the model in hand — it's priced before the transformer is loaded. Standing caveat: ROCm and MPS are still covered only by simulating their dispatch through the real code path. I have neither to measure on. On CUDA every one of these terms remains zero, so the change is a no-op for the hardware the constants were calibrated on. |
Sorry, something went wrong.
JPPhoto
left a comment
There was a problem hiding this comment.
Issues:
-
invokeai/backend/util/attention.py:133-145,170-195: Cache key omits SDPA priority order. sdpa_kernel(..., set_priority=True) can make MATH first while all four booleans stay unchanged; PyTorch selects the first eligible backend in that order. Context manager, selector. A prior fused result can therefore suppress a later score-matrix reservation. Effect: silent under-reservation and possible OOM. Likelihood: plausible advanced runtime configuration. Recovery: restart or clear private probe cache. Test: cache a fused answer, enter sdpa_kernel([MATH, FLASH_ATTENTION, EFFICIENT_ATTENTION, CUDNN_ATTENTION], set_priority=True), force the probe to return MATH, and assert it re-probes.
-
invokeai/backend/util/vae_working_memory.py:144-167, invokeai/app/invocations/flux2_vae_decode.py:56-69: VAE estimate uses only spatial dimensions; decode accepts unrestricted batch tensors and passes the full batch to vae.decode. A (2, 32, H, W) input receives the same reservation as batch 1, while activations and materialized score memory scale with batch. Effect: possible OOM despite the cache reservation. Likelihood: plausible LatentsField edge. Recovery: split into batch-1 decodes or reject non-1 batches. Test: decode a 2-sample latent tensor and verify working_mem_bytes scales or batch input is rejected.
Suggestions:
-
Instead of caching only enable flags, include SDPA priority and deterministic state in the cache key, or remove this cache.
-
Consider enforcing batch 1 for FLUX.2 VAE decode; otherwise pass effective batch size through the estimator.
Sorry, something went wrong.
… batch The probe's cache key held the four per-backend enable flags, but torch takes the *first eligible* backend in a priority order that `sdpa_kernel(..., set_priority=True)` reorders while leaving every flag untouched -- measured: same flags, EFFICIENT outside and MATH inside. A fused answer cached before the switch would suppress the score-matrix reservation after it. Rather than adding the priority order to the key -- the next thing to forget is always one more -- drop the cache. The probe costs ~6us against a multi-second forward, so there is nothing to protect. `vae.decode` is also handed whatever batch the latents carry, and a LatentsField is not pinned to one, so an estimate built from H and W alone gave a two-sample decode a single sample's reservation. Measured at 1024px: 4.23GB at batch 1, 7.96GB at 2, 11.89GB at 3 -- linear, slightly sub-linear per sample, so the scaled single-sample estimate stays an upper bound. The score matrix is (batch, heads, S, S) and scales with it.
JPPhoto
left a comment
There was a problem hiding this comment.
Issues:
- invokeai/app/invocations/flux2_denoise.py:263-279,479-493,529-570: Batched initial latents can reach denoise, but the new reservation ignores batch size. B=2 runs allocate roughly twice the activations, and reference tensors are repeated per batch. Effect: under-reservation and possible OOM. Likelihood: low in stock UI, reachable via API/custom graphs. Recovery: reject B > 1 or scale the estimate by batch. Test: add a B=2 denoise reservation test.
Suggestions:
- Consider passing the actual batch size into _estimate_working_memory and scaling sequence activations and score-matrix cost.
Sorry, something went wrong.
The node had `b` in hand from preparing the latents and never passed it, so a two-sample run reserved one sample's activations and the cache admitted it to a card that could not run it. Batched latents do not come from the stock UI, but the API and custom graphs reach this node. Batch multiplies the token count and nothing else. Measured on the Klein geometry with a reduced block count: 4608 tokens at B=1 peaks at 2570MB, the same 4608 at B=2 at 5126MB, and 9728 tokens at B=1 at 5584MB -- per total token that is 0.554-0.578MB across every combination, so batch and sequence are interchangeable. Reference latents are repeated per sample by `ensure_batch_size`, so they scale too, and the score matrix is (batch, heads, S, S). The fixed base does not scale -- it covers weight casts and allocator slack -- and neither does the regional bias, built as (1, 1, S, S) and broadcast.
Reviewers
Requested changes must be addressed to merge this pull request.
Summary
Fix. The FLUX.2 path called model_on_device() without a working_mem_bytes estimate in every one of its load sites — the denoise node, the VAE decode and encode nodes, and the reference-image encode inside Flux2RefImageExtension. Every other base (SD1/SDXL, FLUX.1, SD3, CogView4, Qwen-Image, Wan, Anima, Krea-2) passes one. Without it the model cache reserves only the default device_working_mem_gb (3 GB) and fills the remainder of the card with the model, so it has no idea that the operation about to run needs far more than that.
Reference images are what turns this from tight into fatal. FLUX.2 concatenates reference latents onto the image stream, so attaching three 1024×1024 references to a 1024×1024 generation takes the attended sequence from 4 608 to 16 896 tokens — and the activation footprint from ~1.7 GB to ~6.5 GB, against a 3 GB reservation. Tile-based refiner workflows do exactly this, once per tile.
How. Two estimators, both calibrated against measured peak reserved memory (the conservative quantity, including allocator overhead), passed at every load site:
Measurements (RTX 4090, bf16, peak reserved, each point in a fresh subprocess so allocator history cannot contaminate it):
The last row is worth noting: 2048px with no references and 1024px with three references produce the same sequence length and the same peak, which is what makes a single per-token constant the right model.
Related Issues / Discussions
Closes #9500
Note on the report: the reporter states the same workflow ran fine on 6.13. I could not confirm a 6.13 → 6.14 regression. _get_vram_available, _load_locked_model and _offload_unlocked_models are logically identical between v6.13.8 and v6.14.0-rc1 for single-GPU (the large model_cache.py diff is multi-GPU plumbing), and the only FLUX.2 memory change in 6.14 — tiling the reference-image VAE encode — reduces peak usage. The defect described here is present in both versions; given how narrow the margin is (see QA below), 6.13 most likely just got lucky more often. That also matches the reporter's own "sometimes it goes through, other times it OOMs at the Nth tile".
QA Instructions
Unit tests: pytest tests/app/invocations/test_flux2_working_memory.py (35 tests). The measured peaks above are pinned as both lower and upper bounds on the estimate, so a future constant change that would reintroduce the OOM — or one that over-reserves so hard the cache pushes the transformer to RAM — fails the suite. The wiring tests were mutation-checked: removing working_mem_bytes= from the denoise call site fails two of them.
End-to-end on a 24 GB card (RTX 4090), device_working_mem_gb: 3, enable_partial_loading: true, Klein 9B fp8 (17.35 GB resident as bf16), three 1024×1024 reference images, 4 steps, Qwen3 encoder on CPU.
Roomy card, 1024×1024 output — both pass, but look at the residency:
That ~0.5 GB of margin in the "before" row is the whole bug. It is not a comfortable pass; it is a coin flip that lands differently depending on what else the cache happens to be holding — precisely the "sometimes it goes through, other times it OOMs at the Nth tile" the issue describes. With the estimate the cache deliberately holds 2.3 GB of the transformer back in RAM.
Tight card, 1328×1328 output — the difference stops being theoretical. A second process pinned 4.5 GB (scripts/allocate_vram.py) to emulate a card with a desktop and other apps on it, and PYTORCH_CUDA_ALLOC_CONF left at the stock native allocator:
Near-identical residency, wildly different outcomes: without the reservation the forward has to claw its working set out of a card the cache believes is fine, and the allocator spends its time synchronizing and releasing cached blocks instead of computing.
Reproducing the reporter's hard OOM. I was not able to make the baseline OOM outright on this hardware — with partial loading enabled it degrades into the PCIe-thrashing case above instead of raising. A card that cannot fall back that way (partial loading disabled, or a model that must be fully resident) is where the same shortfall surfaces as torch.cuda.OutOfMemoryError.
On the choice of constant. Both estimators target peak reserved memory, not peak allocated, consistent with every other estimator in vae_working_memory.py. At 19 689 tokens the denoise allocates ~3.3 GB but reserves ~7.4 GB; targeting the allocated figure would be enough on backend:cudaMallocAsync and would reintroduce the bug for everyone on the default allocator. The cost is that cudaMallocAsync users reserve more than they strictly need.
Merge Plan
Nothing special — backend only, no schema, DB or redux changes. Node versions are unbumped, matching the precedent set by #9305 (the equivalent Qwen-Image working-memory fix), since no node interface changed.
Checklist