fix(flux2): estimate working memory for denoise and both VAE directions by Pfannkuchensack · Pull Request #9519 · invoke-ai/InvokeAI

Pfannkuchensack · 2026-08-19T20:23:18Z

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:

Flux2DenoiseInvocation._estimate_working_memory() — FLUX.2 attention runs through SDPA and never materializes the O(seq²) score matrix, so activations scale linearly with the total attended sequence. Measured slope on the Klein 9B geometry in bf16: ~0.39 MB per token, flat from 1.5k to 28k tokens. It is also independent of block count (a no-grad forward frees each block's intermediates as it goes), so the same constant covers Klein 4B and 9B. Image, reference and text tokens all count; LoRA sidecar patches and the regional-prompting additive bias get their own terms.
estimate_vae_working_memory_flux2() — the FLUX.2 VAE scales linearly in pixel area at ~2170 (decode) / ~1070 (encode) bytes per pixel per element byte, which rounds to the same 2200/1100 constants the FLUX.1 VAE already uses. A 1024×1024 decode peaks at ~4.3 GB, a 1536×1536 decode at ~9.6 GB. The tiled branch bounds the estimate by one tile, matching the 512px tiling the reference-image encode already forces.

Measurements (RTX 4090, bf16, peak reserved, each point in a fresh subprocess so allocator history cannot contaminate it):

image tokens	reference tokens	total seq	measured peak
4 096 (1024px)	0	4 608	1.66 GB
4 096	4 096 (1 ref)	8 704	3.25 GB
4 096	12 288 (3 refs)	16 896	6.38 GB
6 889 (1328px)	12 288	19 689	7.35 GB
6 889	20 667 (3×1328px refs)	28 068	10.70 GB
16 384 (2048px)	0	16 896	6.38 GB

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:

	transformer residency	free VRAM for the forward	result
before	17 350 MB (100 %)	~7 GB against a 6.5 GB need	completes
after	15 046 MB (86.7 %)	~9.5 GB	completes

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:

	transformer residency	result
before	14 182 MB (81.7 %)	still not finished after 12 minutes; cancelled
after	13 894 MB (80.1 %)	completes in 58 s (denoise 35 s)

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

The PR has a short but descriptive title, suitable for a changelog
Tests added / updated (if applicable)
❗Changes to a redux slice have a corresponding migration
Documentation added / updated (if applicable)
Updated What's New copy (if doing a release after this PR)

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

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.

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.

Pfannkuchensack · 2026-08-24T21:02:24Z

Cant really test the ROCm stuff maybe @lstein can take a look

JPPhoto

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.

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

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.

`_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.

Pfannkuchensack · 2026-08-26T22:37:27Z

@JPPhoto
Both findings from the last review are addressed in 12e1c9a.

_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.

JPPhoto

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.

… 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

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.

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.

JPPhoto

invokeai/app/invocations/flux2_denoise.py:268-277,338-353,479-496: With batched init_latents and add_noise=True, b comes from built-in batch-1 noise. Broadcasting creates batch-2 x, but reservation remains batch 1. Effect: possible OOM. Likelihood: custom/API batched img2img. Recovery: reject or split batches. Test: use (2,32,H,W) init latents with default noise and assert reservation uses batch 2.

Suggestions:

Consider deriving batch size from x.shape[0] after blending, or expanding noise to the init-latent batch.

`b` was read from the noise tensor, which this node builds at batch 1 from width/height/seed whenever `add_noise` is set. Batched init latents then broadcast against it in the img2img preblend, producing a two-sample `x` against a one-sample reservation. Read `x.shape[0]` instead. It is already in scope at the estimate -- the blend, the pack and the BN normalize all run above it, and none of them change the batch -- and it is the only thing that knows how many samples reach the transformer. Expanding the noise to match would have worked too, but that changes the noise, and with it the output. Note that the same `b` still feeds `generate_img_ids_flux2`, which is a correctness question rather than a memory one and is left alone here.

JPPhoto

Approved, pending feedback from users of other architectures.

lstein

Adversarial review of 58ef4ae. The core of this is right — the load sites are complete, the batch plumbing is correct, and the SDPA dispatch probe is a much better answer than assuming. One blocker below, plus two ROCm test failures I can reproduce.

All numbers below are measured on ROCm 7.1 / torch 2.10 / gfx1100, bf16, peak reserved, each point in a fresh process.

Blocker

1. The per-token constant is width-calibrated on Klein 9B but applied to FLUX.2 [dev], under-reserving ~27%

invokeai/app/invocations/flux2_denoise.py:684

per_token_bytes = int(0.4 * MB)

Per-token activation cost is linear in the transformer's hidden width, not only in the token count. Measured slope, taken between seq 4608 and 9216 on a reduced block count (2 dual + 2 single), everything else identical:

variant geometry	hidden	Δ @ 4608 tok	Δ @ 9216 tok	slope
Klein 9B (32 heads × 128)	4096	1662 MB	3332 MB	0.362 MB/tok
[dev] (48 heads × 128)	6144	2482 MB	5014 MB	0.550 MB/tok

0.550 / 0.362 = 1.516 ≈ 6144 / 4096. Exactly linear in width. As a check that this transfers to your calibration: my Klein-9B Δ at 4608 tokens is 1662 MB against your MEASURED_DENOISE row of 1702 MB, and the block-count independence you document reproduces (4 blocks, same number).

FLUX.2 [dev] reaches this node as a first-class path — Flux2VariantType.Dev, invokeai/app/invocations/flux2_dev_model_loader.py emitting a TransformerField, and shipped starter models under # region FLUX.2 [dev] in starter_models.py. configs/flux2_variant.py records its width directly:

_HIDDEN_SIZE: dict[Flux2VariantType, int] = {
    Flux2VariantType.Klein4B: 3072,
    Flux2VariantType.Klein9B: 4096,
    Flux2VariantType.Dev: 6144,  # 48 heads × 128 head_dim
}

Triggering sequence. FLUX.2 [dev] via Flux2DevModelLoader → Flux2Denoise at 1024×1024 with three 1024×1024 reference images = 16 896 attended tokens.

reserved: 16896 × int(0.4 MB) + 1 GB = 7.60 GiB
required: scaling your own measured 6538 MB row by 1.516 ≈ 9.9 GB

~2 GB short — the same shortfall #9500 describes, on the 32B model where partial loading makes this estimate the thing that decides residency. Calibrating on Klein 9B is conservative for Klein 4B (3072, narrower) and wrong for [dev] (6144, wider).

This reads as an oversight rather than a decision, because the score-matrix term already accounts for [dev]'s width:

FLUX2_MAX_ATTENTION_HEADS = 48   # "(Klein 4B: 24, Klein 9B: 32, FLUX.2 dev: 48)"

Only the dominant linear term doesn't. transformer_config is already in scope at flux2_denoise.py:321, well above the estimate, so _HIDDEN_SIZE[variant] / 4096 is available without new plumbing. It is a no-op for Klein 9B, so the pinned MEASURED_DENOISE fixtures stay green; using the real head count in the same pass also removes the 48-vs-32 over-count in the score-matrix term.

Medium

2. Two of the new tests fail on ROCm

tests/app/invocations/test_flux2_working_memory.py, 82 tests, 2 failed, 80 passed:

TestSdpaBackendProbe::test_this_build_reports_its_own_dispatch — asserts masked_bytes > 0 on HIP, got 0:

>           assert masked_bytes > 0
E           assert 0 > 0

Head dim 128 with an additive mask reports EFFICIENT_ATTENTION on this build, not MATH. Direct probe of torch.ops.aten._fused_sdp_choice, gfx1100:

device	dtype	head_dim	mask	choice
cuda (hip)	bf16	128	no	FLASH_ATTENTION
cuda (hip)	bf16	128	yes	EFFICIENT_ATTENTION
cuda (hip)	bf16	512	no	MATH
cuda (hip)	bf16	512	yes	MATH

The code is fine — it asks torch and gets the right answer. The test pins a hardware claim that is false here, and so do three docstrings and a commit message ("ROCm's fused kernels cap the head dim at 128 and reject arbitrary additive masks", attention.py:47, flux2_denoise.py _estimate_working_memory). The head-dim half is correct and load-bearing; the mask half is not. Worth softening both the assertion and the prose to "whatever torch says".

TestSdpaBackendProbe::test_constant_upper_bounds_a_forced_math_forward[1-4096-512]:

>       assert estimate >= measured
E       assert 218103808 >= 226492416

Measured bytes per score element on ROCm, forced SDPBackend.MATH:

heads	seq	head_dim	bytes/element
1	4096	512	13.62
1	8192	512	10.78
1	16384	512	9.76
4	4096	128	10.16
48	4608	128	9.59

So SDPA_MATH_BYTES_PER_SCORE_ELEMENT = 13 is not an upper bound at the smallest parametrized point, which contradicts the docstring's "13 is an upper bound on every measured point from 4k tokens up" — 4096 is 4k tokens, and that is where it fails. The practical shortfall is 4.7 MB, i.e. exactly the allocator-rounding noise the same docstring waves off, so this is a test-assertion bug rather than a memory bug. But it is red for any ROCm CI or developer running the suite, so it needs either a small margin in the constant or a lower bound on the parametrization.

3. MPS gains a multi-GB term on an unvalidated constant, on every FLUX.2 estimate

torch.ops.aten._fused_sdp_choice is registered for CPU, CUDA/ROCm and XPU only, so on MPS the probe always raises and always takes the materializing branch — as documented and intended. The consequence at ordinary settings is large: a plain 1024×1024 Klein generation, no references, no mask, gains 48 × 4608² × 13 = 13.2 GB. MPS's _get_vram_available budgets against psutil.virtual_memory().available, so that reservation will push the transformer wholly to RAM on essentially every Mac.

Whether MPS really materializes at that scale is untested, as the PR says. That is fine as a conservative default, but this is the one place where conservative is not cheap. Two things would take most of the sting out without giving up the safety:

Use the model's real head count instead of FLUX2_MAX_ATTENTION_HEADS. Klein 9B is 32, not 48 — a third off immediately, and it falls out of the fix for #1 anyway.
Log at info level when the score-matrix term actually fires, so a Mac user who suddenly sees 0% residency can tell why. Right now the only log is the _warn_unknown_diffusers_dispatch path, which is a different cause.

Non-blocking

warnings.catch_warnings() is not thread-safe. _torch_sdpa_materializes_score_matrix (attention.py) mutates the process-global filter list. With multi-GPU session workers, and with the probe cache deliberately removed so this now runs on every estimate, two concurrent estimates can interleave enter/exit and leave a stale filter set behind. torch._C._set_print_stack_traces-style suppression or just accepting the warnings would avoid the global.
Device inconsistency between the two VAE nodes. flux2_vae_decode.py passes device=vae_info.compute_device; flux2_vae_encode.py passes device=TorchDevice.choose_torch_device(). For a cpu_only VAE these differ and the probe then answers for the wrong device. (The encode node's own compute path has the same pre-existing issue, which is presumably why it was written this way — worth aligning both, or at least noting it.)
The VAE model runs short below ~1024px on a materializing backend. Stock-config AutoencoderKLFlux2 decode, measured vs. estimate_vae_working_memory_flux2:

px	measured	estimate (linear + score)
512	1.72 GB	1.28 GB
768	3.78 GB	3.45 GB
1024	6.70 GB	7.55 GB

The 3 GB device_working_mem_gb floor covers 512px; 768px is ~0.33 GB short. Small, and only on backends where the score matrix is real, but it means the combined model is not an upper bound everywhere the way the 1024px+ points suggest.

invokeai/app/invocations/ernie_image_vae_decode.py:48 decodes through the same AutoencoderKLFlux2 with a bare model_on_device(). Out of scope for this PR, but estimate_vae_working_memory_flux2 applies there too (with care for the patchified latent shape) — worth a follow-up issue.

Attacks that found nothing

Recording these so the clean areas are visible:

Load-site coverage is complete. The two remaining bare model_on_device() calls on a FLUX.2 VAE — flux2_denoise.py:171 and flux2_pid_decode.py:155 — only read BN buffers and vae.config respectively, and never run the VAE.
x.shape[0] is the right batch and is stable across the img2img preblend, pack_flux2 and _bn_normalize above it; nothing between the estimate and denoise() changes it, and ensure_batch_size repeats the reference latents to match.
The regional-bias skip matches the runtime skip. The estimate gates on ref_image_seq_len == 0, the forward on img_cond_seq is None. They can only disagree for a non-None extension with zero-length latents, and a falsy kontext_conditioning leaves the extension None anyway.
regional_attention_bias_bytes matches the real allocation — get_joint_attention_kwargs builds one (S, S) tensor of dtype, unsqueezed to (1, 1, S, S), so numel × element_size is right and the broadcast correctly does not scale with batch.
CFG runs pos and neg as sequential forwards, not a concatenated batch, so max(txt.shape[1], neg_txt.shape[1]) is the correct text term.
_AttentionBackendRegistry.get_active_backend() does return a 2-tuple in the pinned diffusers 0.39.0, and AttentionBackendName.NATIVE.value == "native", so the backend, _ = ... unpack and the name == "native" branch both hold — the except path is not silently swallowing every estimate into MATH.
The transformer/VAE dispatch asymmetry is real. diffusers.models.transformers.transformer_flux2 routes through dispatch_attention_fn; autoencoder_kl_flux2 does not, reaching SDPA via AttnProcessor2_0. The via_diffusers_dispatch split and the test pinning it are correct.
The VAE geometry constants are correct. block_out_channels[-1] == 512 and 2 ** (len(block_out_channels) - 1) == 8, so the single 512-wide head and _FLUX2_VAE_SPATIAL_COMPRESSION = 8 both check out. patch_size=(2, 2) only sizes the BatchNorm, it is not an extra compression stage, so LATENT_SCALE_FACTOR and the //8 mid-block sequence length are right for both directions.
The tiled reference-encode estimate stays above the untiled cost for images below the 512px tiling threshold, so forcing the tiled branch unconditionally does not under-reserve.
Reading .model and .compute_device outside the lock is the established pattern for these estimators and does not race the load.

fishd72 · 2026-08-29T15:26:58Z

I set up a test environment for my RX9070XT 16GB using ROCm7.1 following the instructions from @Pfannkuchensack on Discord (you're a legend, thanks).

(FWIW, this is Fedora 44 KDE, using kernel 7.1.10-200.fc44.x86_64 and the default amdgpu driver stack).

I set up a fresh DB in a test root, did an in-place import from my live install root model folder and included the following env vars:

MIOPEN_FIND_MODE=2
MIGRAPHX_MLIR_USE_SPECIFIC_OPS=attention
FLASH_ATTENTION_TRITON_AMD_ENABLE=TRUE
FLASH_ATTENTION_TRITON_AMD_AUTOTUNE=FALSE
TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1
TORCH_BLAS_PREFER_HIPBLASLT=1
PYTORCH_HIP_ALLOC_CONF=garbage_collection_threshold:0.7,max_split_size_mb:256

and my invokeai.yml contains the following to mirror my current setup:

precision: bfloat16
attention_type: torch-sdp

I was able to test SDXL, Anima Base-v1.0 and KREA2 (krea2_turbo-Q4_K_M.gguf) models successfully at 1024x1024 with a simple prompt.

I then imported the FLUX2 models (FLUX.2 Klein 4B (GGUF Q4)) using the starter model link in the model manager, and have run successive single generations and batches of up to ten images at at time (1024x1024 resolution) successfully without issue. (In the current live version, I seem to get about 4 generations in, then a get a hard crash of the graphics subsystem. The generations complete successfully, but the GPU driver is quite unhappy.)

In case of interest:

[2026-08-29 16:47:47,277]::[InvokeAI]::INFO --> Graph stats: e943a1f2-8956-44db-9601-6a482b6a383f
                          Node   Calls   Seconds VRAM Change
                       integer       1    0.001s     +0.000G
      flux2_klein_model_loader       1    0.000s     +0.000G
                        string       1    0.000s     +0.000G
                 core_metadata       1    0.001s     +0.000G
      flux2_klein_text_encoder       1    0.001s     +0.000G
                       collect       1    0.001s     +0.000G
                 flux2_denoise       1    2.687s     +0.000G
              flux2_vae_decode       1    2.922s     +0.000G
TOTAL GRAPH EXECUTION TIME:   5.615s
TOTAL GRAPH WALL TIME:   5.622s
RAM used by InvokeAI process: 12.36G (+0.000G)
RAM used to load models: 2.95G
VRAM in use: 8.430G
RAM cache statistics:
   Model cache hits: 3
   Model cache misses: 0
   Models cached: 5
   Models cleared from cache: 0
   Cache high water mark: 10.48/12.92G

It should be noted, at one point in my testing I neglected to include the environment variables listed above, and during those tests I would receive OOM errors during FLUX2 generation and it would cause parts of KWIN (KDE Desktop Environment) to crash giving the error "Desktop effects were restarted due to a graphics reset". So the env variables are clearly important.

I'll try to replicate this testing now on macOS as well.

…ount Per-token activation cost is linear in the transformer's hidden width, and the constant was calibrated on Klein 9B (4096) but applied to every variant. FLUX.2 [dev] is 6144 and reaches this node as a first-class path, so 1024x1024 with three references reserved 7.6GB against ~10GB needed -- the same shortfall invoke-ai#9500 describes, on the model where partial loading makes the estimate decide residency. Measured slope between 4608 and 9216 tokens, everything else held fixed: 0.291 MB/tok at 3072, 0.386 at 4096, 0.555 at 6144 -- 0.755 / 1.00 / 1.438 against width ratios of 0.75 / 1.00 / 1.50. Scale by width, and take the head count from the same number instead of always charging the widest. Also: raise SDPA_MATH_BYTES_PER_SCORE_ELEMENT to 14, which ROCm's 13.62 at the smallest measured shape needs; drop the claim that ROCm rejects additive masks, which gfx1100 disproves; log at info when the score-matrix term fires, since it decides residency and nothing else said so; guard the warning-filter swap with a lock now that the probe runs on every estimate; and give the VAE encode node the same compute device as the decode node.

…chensack/InvokeAI into fix/flux2_working_memory

Pfannkuchensack · 2026-08-29T17:41:04Z

@fishd72 @lstein can you run the python scripts/calibrate_flux2_working_memory.py --csv flux2_rocm.csv script to get the right numbers for the vae stuff.

fishd72 · 2026-08-29T18:17:22Z

Done:

torch 2.10.0+rocm7.1 | device AMD Radeon RX 9070 XT | hip=7.1.25424 | dtype=bfloat16

=== 1. SDPA dispatch: which kernel would this build pick? ===
The score-matrix term is charged only where the answer is MATH (or the query raises).

 head_dim   mask                 kernel
----------------------------------------
      128  False        FLASH_ATTENTION
      128   True    EFFICIENT_ATTENTION
      512  False        FLASH_ATTENTION
      512   True    EFFICIENT_ATTENTION

  head_dim 512 is the FLUX.2 VAE mid-block; 128 is the transformer, and the masked row
  is regional prompting. MATH on the 512 row means the VAE estimate needs the score term.

=== 2. SDPA_MATH_BYTES_PER_SCORE_ELEMENT ===
Peak reserved per score element with MATH forced. Shipped constant: 14.

 heads     seq  head_dim  reserved(GiB)  bytes/elem
----------------------------------------------------
     1    4096       512          0.256       16.38
     1    8192       512          0.842       13.47
     1   16384       512          2.998       11.99
     4    4096       128          0.803       12.84
    48    4608       128         11.098       11.69

  max = 16.38; shipped constant is 14 -> SHORT

=== 3. VAE linear constants (2200 decode / 1100 encode) ===
`implied_k` backs the score-matrix term out, so it is comparable to those literals; fit the
constant on the rows whose `math` column matches what this build really does (section 1).
`covered` is the question that matters: is the shipped estimate an upper bound here?
Caveat: forcing math on a build that HAS a fused kernel is not equivalent to a build that
has none -- on CUDA the forced-math decode measures *below* the fused one, because the
memory-efficient kernel's workspace is the larger term there. Only a real run on the
materializing build calibrates it.

op         px  math  measured(GiB)  estimate(GiB)  implied_k  covered
----------------------------------------------------------------------
decode    512 False          2.154          1.074       4412       NO
decode    512  True          2.154          1.293       3964       NO
decode    768 False          4.721          2.417       4297       NO
decode    768  True          4.721          3.524       3289       NO
decode   1024 False          8.391          4.297       4296       NO
decode   1024  True          8.391          7.797       2504       NO
decode   1280 False          6.691          6.714       2193      yes
decode   1280  True          8.307         15.259        n/a      yes
decode   1536 False          9.355          9.668       2129      yes
decode   1536  True         15.191         27.387        n/a      yes
encode    512 False          1.080          0.537       2212       NO
encode    512  True          1.080          0.756       1764       NO
encode    768 False          2.428          1.208       2210       NO
encode    768  True          2.428          2.316       1202       NO
encode   1024 False          4.377          2.148       2241       NO
encode   1024  True          4.377          5.648        449      yes
encode   1280 False          6.838          3.357       2241       NO
encode   1280  True          7.230         11.902        n/a      yes
encode   1536 False          9.846          4.834       2240       NO
encode   1536  True         14.910         22.553        n/a      yes

  decode (fused): implied_k max = 4412, shipped = 2200 -> SHORT
  decode (math): implied_k max = 3964, shipped = 2200 -> SHORT
  encode (fused): implied_k max = 2241, shipped = 1100 -> SHORT
  encode (math): implied_k max = 1764, shipped = 1100 -> SHORT

  Points the shipped estimate does NOT cover:
    decode 512px force_math=False: short by 1.08 GiB
    decode 512px force_math=True: short by 0.86 GiB
    decode 768px force_math=False: short by 2.30 GiB
    decode 768px force_math=True: short by 1.20 GiB
    decode 1024px force_math=False: short by 4.09 GiB
    decode 1024px force_math=True: short by 0.59 GiB
    encode 512px force_math=False: short by 0.54 GiB
    encode 512px force_math=True: short by 0.32 GiB
    encode 768px force_math=False: short by 1.22 GiB
    encode 768px force_math=True: short by 0.11 GiB
    encode 1024px force_math=False: short by 2.23 GiB
    encode 1280px force_math=False: short by 3.48 GiB
    encode 1536px force_math=False: short by 5.01 GiB

=== 4. Denoise per-token constant and its width scaling ===
Shipped: 0.4 MB/token at hidden=4096, scaled linearly by width.

variant     hidden  blocks  short(GiB)  long(GiB)  MB/token   ratio
--------------------------------------------------------------------
klein_4b      3072       2       1.195      2.611    0.3147   1.000
klein_9b      4096       2       1.604      3.434    0.4067   1.000
dev           6144       2       2.551      5.201    0.5890   1.448

  hidden=4096: measured 0.4067 MB/token, shipped 0.4 -> SHORT
  hidden=3072: slope ratio 0.774 against width ratio 0.750
  hidden=4096: slope ratio 1.000 against width ratio 1.000
  hidden=6144: slope ratio 1.448 against width ratio 1.500

  Those two columns should match: that is the claim that width scales the constant.
  klein_4b   1024px + 3 refs: estimate  5.95 GiB, measured slope implies  5.19 GiB
  klein_9b   1024px + 3 refs: estimate  7.60 GiB, measured slope implies  6.71 GiB
  dev        1024px + 3 refs: estimate 10.90 GiB, measured slope implies  9.72 GiB

…ents Two ROCm runs came in through the new calibration script and both shipped constants were under their worst point. SDPA_MATH_BYTES_PER_SCORE_ELEMENT goes 14 -> 17 (gfx1201 costs 16.38 for the shape where CUDA costs 12.88), and the per-token activation constant 0.40 -> 0.42 MB (gfx1201 measures 0.4067 at the reference width). Both are now pinned against all three platforms' measured points rather than only being self-consistent. The runs also disprove the ROCm framing this feature carried: gfx1100 reports MATH for the VAE's 512-wide head, gfx1201 reports FLASH. Two cards, same vendor, same torch, opposite answers -- which is the case for asking torch rather than hard-coding a rule, but it means the docstrings could not keep saying "ROCm caps the head dim at 128". The VAE linear constants are left alone despite measuring short on gfx1201. That run had MIOPEN_FIND_MODE=2 and a HIP allocator garbage_collection threshold set, and its series is non-monotonic above 1024px -- peak reserved falls as resolution rises, which is what a GC threshold does to this measurement. The two AMD cards are also 1.8x apart. The gap is documented where the constants are defined; the script now reports the environment and flags a non-monotonic series so the next run cannot be ambiguous about it.

…chensack/InvokeAI into fix/flux2_working_memory

A clean ROCm run — the earlier one had MIOPEN_FIND_MODE=2, worth a uniform 1.28x, and a HIP allocator GC threshold that clipped the high-resolution points — puts the gfx1201 numbers at 3453/2688 bytes per pixel per element byte against cuDNN's 2185/1072. Flat across 512-1024px on both, so the linear model holds; only the coefficient moves. It is MIOpen's convolution workspaces, not the attention term: identical on the fused path. Shipping the MIOpen numbers everywhere would add ~60% to every cuDNN decode for nothing, so the constant follows the backend, keyed on torch.version.hip rather than the device string (a HIP build reports device.type == "cuda"). The two operations also stop sharing a ratio. "Encoding costs half of decoding" holds on cuDNN (0.49) and not on MIOpen (0.78), so it was a backend property masquerading as an architectural one. MIOPEN_FIND_MODE=2 is deliberately not budgeted for: it is not the default and would tax everyone else. Noted where the constants are defined.

lstein · 2026-08-29T22:36:04Z

torch 2.10.0+rocm7.1 | device AMD Radeon PRO W7900 Dual Slot  | hip=7.1.25424 | dtype=bfloat16
env: MIOPEN_FIND_MODE=, PYTORCH_HIP_ALLOC_CONF=

=== 3. VAE linear constants (this build selects 3500 decode / 2750 encode) ===
`implied_k` backs the score-matrix term out, so it is comparable to those literals; fit the
constant on the rows whose `math` column matches what this build really does (section 1).
`covered` is the question that matters: is the shipped estimate an upper bound here?
Caveat: forcing math on a build that HAS a fused kernel is not equivalent to a build that
has none -- on CUDA the forced-math decode measures *below* the fused one, because the
memory-efficient kernel's workspace is the larger term there. Only a real run on the
materializing build calibrates it.

op         px  math  measured(GiB)  estimate(GiB)  implied_k  covered
----------------------------------------------------------------------
decode    512 False          1.721          1.975       2980      yes
decode    512  True          1.721          1.975       2980      yes
decode    768 False          3.781          5.190       2218      yes
decode    768  True          3.781          5.190       2218      yes
decode   1024 False          6.703         11.086       1256      yes
decode   1024  True          6.703         11.086       1256      yes
encode    512 False          1.312          1.608       2144      yes
encode    512  True          1.312          1.608       2144      yes
encode    768 False          2.953          4.366       1464      yes
encode    768  True          2.953          4.366       1464      yes
encode   1024 False          5.250          9.621        512      yes
encode   1024  True          5.250          9.621        512      yes

  decode (fused): implied_k max = 2980, shipped = 3500 -> OK
  decode (math): implied_k max = 2980, shipped = 3500 -> OK
  encode (fused): implied_k max = 2144, shipped = 2750 -> OK
  encode (math): implied_k max = 2144, shipped = 2750 -> OK

The W7900 run shows forced math and the fused path measuring identically to three decimals at every resolution, with the total flat-linear in area. The score matrix does not add to the convolution peak: the mid-block sits alone at the 8x-downsampled bottleneck, so the full-resolution feature maps are not live while it runs, and peak reserved is a high-water mark rather than a running total. Measured on cuDNN, forcing math stays *below* the fused path until 1536px and then exceeds it by 2.6GB against the 21.5GB the term prices standalone. Summing reserved 11.1GB for a 1024px gfx1100 decode that measures 6.7. Take the max: 7.0GB. A max model is weakest at the crossover, and one measured point sits there -- a 768px encode with cuDNN's constant and a materializing kernel wants 1.80GB against a 1.35GB max, reproducibly. It is not reachable as a shortfall: the cache floors every reservation at device_working_mem_gb and the whole crossover region is below it. Pinned rather than rounded away. The MIOpen decode constant also goes 3500 -> 3600; the W7900 asks for 3525 at 512px where gfx1201 asks for 3453. Its encode column agrees with gfx1201's to the byte, so this is MIOpen rather than a per-card quirk.

Pfannkuchensack requested review from JPPhoto, blessedcoolant, dunkeroni and lstein as code owners August 19, 2026 20:23

github-actions Bot added python PRs that change python files invocations PRs that change invocations backend PRs that change backend files python-tests PRs that change python tests labels Aug 19, 2026

Merge branch 'main' into fix/flux2_working_memory

7b4a437

lstein assigned JPPhoto Aug 24, 2026

lstein added the 6.14.0 label Aug 24, 2026

lstein added this to Invoke - Community Roadmap Aug 24, 2026

lstein moved this to 6.14.x Theme: USER EXPERIENCE in Invoke - Community Roadmap Aug 24, 2026

lstein added 6.14.1 and removed 6.14.0 labels Aug 24, 2026

JPPhoto requested changes Aug 24, 2026

View reviewed changes

Pfannkuchensack requested a review from JPPhoto August 24, 2026 21:02

Merge branch 'main' into fix/flux2_working_memory

bef3270

lstein self-assigned this Aug 24, 2026

Merge branch 'main' into fix/flux2_working_memory

99dc0c6

JPPhoto requested changes Aug 25, 2026

View reviewed changes

JPPhoto and others added 2 commits August 25, 2026 03:54

Merge branch 'main' into fix/flux2_working_memory

6985eb3

Pfannkuchensack requested a review from JPPhoto August 26, 2026 03:12

Merge branch 'main' into fix/flux2_working_memory

6f499da

JPPhoto requested changes Aug 26, 2026

View reviewed changes

Merge branch 'main' into fix/flux2_working_memory

4924adb

JPPhoto self-requested a review August 27, 2026 21:14

JPPhoto requested changes Aug 27, 2026

View reviewed changes

Pfannkuchensack requested a review from JPPhoto August 28, 2026 00:51

JPPhoto requested changes Aug 28, 2026

View reviewed changes

Pfannkuchensack requested a review from JPPhoto August 28, 2026 13:48

JPPhoto requested changes Aug 28, 2026

View reviewed changes

Pfannkuchensack requested a review from JPPhoto August 28, 2026 23:38

JPPhoto approved these changes Aug 29, 2026

View reviewed changes

lstein requested changes Aug 29, 2026

View reviewed changes

Pfannkuchensack added 4 commits August 29, 2026 19:04

Merge branch 'main' into fix/flux2_working_memory

9dd8779

Merge branch 'fix/flux2_working_memory' of https://github.com/Pfannku…

2554823

…chensack/InvokeAI into fix/flux2_working_memory

Merge branch 'main' into fix/flux2_working_memory

591cf21

Pfannkuchensack added 4 commits August 29, 2026 22:31

Merge branch 'fix/flux2_working_memory' of https://github.com/Pfannku…

3658f95

…chensack/InvokeAI into fix/flux2_working_memory

Merge branch 'main' into fix/flux2_working_memory

0d259e5

Pfannkuchensack requested review from JPPhoto and lstein August 29, 2026 22:22

Uh oh!

Conversation

Pfannkuchensack commented Aug 19, 2026

Summary

Related Issues / Discussions

QA Instructions

Merge Plan

Checklist

Uh oh!

JPPhoto left a comment

Choose a reason for hiding this comment

Uh oh!

Pfannkuchensack commented Aug 24, 2026

Uh oh!

JPPhoto left a comment

Choose a reason for hiding this comment

Uh oh!

JPPhoto left a comment

Choose a reason for hiding this comment

Uh oh!

Pfannkuchensack commented Aug 26, 2026

Uh oh!

JPPhoto left a comment

Choose a reason for hiding this comment

Uh oh!

JPPhoto left a comment

Choose a reason for hiding this comment

Uh oh!

JPPhoto left a comment

Choose a reason for hiding this comment

Uh oh!

JPPhoto left a comment

Choose a reason for hiding this comment

Uh oh!

lstein left a comment

Choose a reason for hiding this comment

Blocker

1. The per-token constant is width-calibrated on Klein 9B but applied to FLUX.2 [dev], under-reserving ~27%

Medium

2. Two of the new tests fail on ROCm

3. MPS gains a multi-GB term on an unvalidated constant, on every FLUX.2 estimate

Non-blocking

Attacks that found nothing

Uh oh!

fishd72 commented Aug 29, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Pfannkuchensack commented Aug 29, 2026

Uh oh!

fishd72 commented Aug 29, 2026

Uh oh!

lstein commented Aug 29, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

fishd72 commented Aug 29, 2026 •

edited

Loading