[Pfannkuchensack]

Stacked on #9896. Until that merges, the diff also shows its commit; this PR's change is ef1bbec220.

Summary

VAE working memory was priced as pixels × element size × k, with one k per convolution backend fitted on bf16 decodes. Measurements show three problems with that:

• The demand does not scale with element size. On MIOpen, fp16 needs about a third of bf16, and on both backends fp32 needs less per value byte than bf16.
• The allocator changes demand by 25–35 %.
• The CogView4 VAE is a larger stack and was under-reserved by almost 2×.

On an RX 9060 XT this evicted the SDXL UNet on every 1024² fp32 image.

The five AutoencoderKL-family estimators now price a measured slope and base per (conv backend, allocator, dtype). The values come from a shipped upper-envelope fit (+5 %) over the measured points, and a test checks they reproduce from the points.

1024², reservation in MiB	before	after
RX 9060 XT (ROCm Windows default), SDXL fp32 decode	14650	6225
RX 9060 XT, SDXL fp16 decode	7200	2499
RTX 4090 cudaMallocAsync, SDXL fp32 decode	9050	6124
RTX 4090 native, FLUX bf16 decode	4400	4604
CogView4 decode, cuDNN native	4400	8759
ROCm Windows with expandable_segments:False, FLUX bf16 decode	7200	11009

• Unmeasured platforms (MPS, XPU, CPU, other allocators, ROCm off Windows except the native allocator) are priced through one declared rule and never below the previous values.
• ROCm native on Linux keeps the earlier Linux W7900 / RX 9070 XT bf16 fit, with the previous values for fp16/fp32.
• SD/SDXL: the node and its estimate share one run-dtype rule. An fp32=False decode runs fp16 whatever the loaded dtype.
• fp32 weight term: the fixed +250 MiB is replaced by the real weight growth of the cast.

Related Issues / Discussions

Builds on #9896. Measurement report and runs: SD/SDXL VAE decode measurements (both cards, fp16/bf16/fp32, native/cudaMallocAsync, isolated and in real generations).

QA Instructions

• Measurements:

  • Random-init VAEs from the real configs, one fresh process per point, two runs each, with demand = max(driver, reserved, allocated) delta.
  • Covered: SDXL fp16/fp32, FLUX.1 and FLUX.2 in bf16/fp32, and CogView4 bf16; decode and encode at 512–1536²; on the RTX 4090 (native and cudaMallocAsync) and the RX 9060 XT (expandable and native).
  • All 279 points live in working_memory/data/vae.json.

• A/B on the RX 9060 XT (SDXL txt2img, 25 steps, E2E tool; before = upstream/main, after = this branch):

  Case	before	after
  fp32 768²	VAE 86–92 % resident	6.37 s, fully resident
  fp32 1024²	VAE 13–33 % resident, UNet reloaded per image, ~15.1 s	10.87 s, fully resident
  fp16 1280²	19.22 s, budget short	17.02 s
  fp16 512–1024²	—	unchanged (3.4 / 6.4 / 10.7 s)
  fp32 1280²	VAE 0 % resident	48–94 % resident; does not fit 16 GB

  The two runs were not interleaved.

• Tests: affected suites 4945 passed, 0 failed. Each changed file also passes alone.

  • Mutation checks: hand-edited fitted values, stand-ins below the previous values, hidden out-of-fit points, a dropped base term, broken batch scaling, a changed margin, and ignored allocator or OS detection are each caught.

Review

Resolved findings

• The CogView4 VAE (block_out_channels [128, 512, 1024, 1024], 3 layers per block) had been borrowing the family's lines. It is now measured and has its own lines.
• ROCm on Windows with the native allocator (the override Invoke's own log suggests) was 30 % under-reserved. It now has its own fitted values, with five listed exceptions to the 1.3× upper bound where its first-run curve is concave.
• MIOpen expandable-segments values, measured on Windows gfx1200 only, no longer apply off Windows.
• The SD run-dtype rule is shared with the nodes.
• A dtype with no coefficient line raises a clear ValueError.
• Tests are isolated from the host's allocator environment.

Remaining limits

• Tiled decode/encode still reserve 1.27–1.49× what they need; a refit follows.
• The SDXL fp32 decode on unmeasured platforms is 90 MiB below the previous value, because the cast term now prices the real growth (160 MiB) instead of 250 MiB.
• Linux ROCm values rest on the older Linux points and the previous constants.

Compatibility / Rollout

• Estimator version 2: coefficient names and data format changed. No persisted state; the reservations change as listed above.

Checklist

• The PR has a short but descriptive title, suitable for a changelog
• Meaningful regression coverage added / updated where needed; obsolete tests/code removed
• Persisted-state and API changes include required migrations / compatibility validation
• Relevant performance/efficiency opportunities considered; material claims have evidence
• Material review findings resolved and relevant checks rerun
• Documentation added / updated (if applicable)
• Updated What's New copy (if doing a release after this PR)

---

^{Need help on this PR? Tag @codesmith-bot with what you need.}

[JPPhoto]

@Pfannkuchensack findings/issues to resolve:

• Should fix 1. invokeai/backend/util/vae_working_memory.py:160-169: What this is: _compute_dtype() picks the first floating-point parameter wider than one byte. Problem: For an SDNQ VAE, that can be an fp32 bias even though packed weights dequantize to bf16. Fix: Prefer the packed tensor’s compute_dtype. Effect: On cuDNN/native, the wrong fp32 line adds about 3.1 GiB for a 1024px decode, evicting other models or forcing earlier tiling. Likelihood: Supported SDNQ VAE with bf16 compute and fp32 side parameters. Recovery: Tile the decode or use a non-SDNQ VAE. Test: Reproduced with an SDNQ weight and fp32 bias; add a regression test asserting the bf16 line. If we skip it: Generation can use more memory and run slower.

• Should fix 2. invokeai/backend/util/vae_working_memory.py:285; invokeai/backend/model_manager/load/model_cache/cached_model/cached_model_only_full_load.py:129,187: What this is: The estimate adds cast growth only while parameters are below fp32 size. Problem: With keep_ram_copy_of_weights=false, an fp32 VAE can be offloaded to CPU without restoring its original dtype, while the cache keeps its original fp16 byte count. The next estimate adds zero for the cast. Fix: Restore the dtype on offload or update cache byte accounting. Effect: A repeated fp32 run can OOM when VRAM is tight; the VAE’s uncounted growth is about 160 MiB. Likelihood: No RAM copy and an offload after an fp32 VAE run. Recovery: Restart to reload the original weights, or tile/reduce the output size. Test: Exercise cast, offload, and reload with no RAM copy; assert accounting includes the expanded size. If we skip it: Repeated runs can fail under memory pressure.

• Should fix 3. scripts/calibrate_flux2_working_memory.py:79-84,612: What this is: The CLI accepts --dtype float16. Problem: FLUX.2 VAE ships only bf16 and fp32 lines, so --only vae --dtype float16 raises KeyError: 'k_decode_fp16'. Fix: Reject or skip unsupported VAE dtype combinations with a clear message. Effect: VAE calibration aborts before measuring. Likelihood: Any operator choosing the accepted float16 option. Recovery: Use bfloat16 or float32. Test: Reproduced on the pinned checkout; add a CLI test for a clear rejection. If we skip it: Operators get a traceback instead of calibration results.

Suggested follow-up work (pre-existing):

• invokeai/app/invocations/vae/flux_vae_encode.py:95 and invokeai/app/invocations/vae/flux_vae_decode.py:97: When the first parameter is an SDNQ packed weight, the input is cast to its packed dtype. These lines are unchanged from the PR base. Suggested issue: use the SDNQ compute dtype for VAE inputs.

[JPPhoto]

@Pfannkuchensack See prior review.

[Pfannkuchensack]

@JPPhoto Thanks. All three points are addressed in e3bb17e, and each has a test that fails without its fix.

1. SDNQ compute dtype. _compute_dtype now takes the packed weights' compute_dtype when the VAE carries one; otherwise it falls back to the first wide floating parameter. A real SDNQTensor (int8 storage, bf16 compute) next to an fp32 bias now prices on the bf16 line. Limit: an SDNQ VAE that kept whole layers unquantized in fp32 would be priced below their real cost. I have not verified whether any shipped SDNQ VAE does that; the docstring says so.

2. Cast weights offloaded without a RAM copy. Instead of making the cache's byte count mutable, the offload restores the dtype the cache counts:

   • CachedModelOnlyFullLoad records each weight's dtype when it keeps no RAM copy. On unload it moves weights a node cast in place back to the CPU in that dtype.
   • CachedModelWithPartialLoad, the default for VAEs on dedicated VRAM, does the same per key on its CPU offload path.
   • move_shared takes the dtype, so tied weights stay one tensor.
   • Packed SDNQ/GGUF tensors are never recast.

   This matches what the RAM-copy path already did. Tests cast, offload and reload with no RAM copy, and assert the dtype, total_bytes, the tie and the VRAM bytes. The partial-load case is CUDA/MPS-only and passed on an RTX 4090. Remaining gap, shared with the RAM-copy path: in a partially loaded model, weights that were already on the CPU when the node cast them stay fp32.

3. --dtype float16. The FLUX.2 calibration script now refuses a VAE dtype with no shipped line whenever the vae section runs. It exits with code 2 and names the supported dtypes, which it reads from the spec. A new CLI test covers this.

Follow-up you suggested. The FLUX VAE nodes cast inputs with next(iter(vae.parameters())).dtype, which is a packed SDNQ weight's storage dtype. This is a separate root cause and unchanged here; I'll file it as its own issue.

[JPPhoto]

Merge blocker

invokeai/backend/model_manager/load/model_cache/cached_model/cached_model_with_partial_load.py:513 can leave weights at the wrong dtype after a partial offload. A VAE upcast to fp32 has both CPU and GPU weights changed. The offload skips weights already on CPU and restores only the GPU weights it selects. Remaining weights stay fp32 while cache accounting still records fp16 sizes. A later load can move fp32 weights into VRAM without accounting for their extra size, causing overcommit or OOM.

Restore all changed weights to their counted dtypes on partial offload, including weights already on CPU, or make cache sizing track their actual dtypes. Add a regression test for partial residency, fp32 cast, and partial offload.

[Pfannkuchensack]

@JPPhoto Thanks. Fixed in 47eac03. The fix does more than restore the dtype on a partial offload: the cache now puts every weight a holder cast back to its counted dtype when the last lock ends.

The fix

• When the last lock on an entry ends, the cache returns every weight a holder cast in place back to the dtype it counts. This covers parameters and persistent buffers, on whichever device they sit, and the weights already on the CPU.
  • On the CPU it re-points them at the RAM copy when there is one, and casts them where they sit otherwise.
  • A tied group gets one tensor under all of its names.
  • Both CachedModelOnlyFullLoad and CachedModelWithPartialLoad do this, and so do cpu_only entries, which never get a load pass.
• Load and offload passes therefore never see cast weights, so the per-key cast on the offload path is gone.
• The SD VAE nodes still cast with a plain vae.to(run_dtype). Every run now starts from the loaded dtype, so _upcast_weight_bytes prices the full upcast each time. Before, a VAE left in fp32 reserved nothing for its next upcast, and with partial residency that also overcommitted.
• The restore re-points Parameter.data and _buffers; it does not use load_state_dict(assign=True). A construction on another thread patches register_parameter process-wide, and this path does not go through that function. So the restore takes no MODEL_LOAD_LOCK, and an unlock never waits behind a construction. In multi-GPU that would have held back a finished image for up to another device's cold load.
• Only the last unlock restores, so an inner lock's unlock keeps the cast its outer holder runs with.
• If a restore fails, the entry is dropped with a warning, so the node keeps its finished result.
• Packed GGUF and SDNQ tensors never differ from their counted dtype, so they are never selected.

Tests

• A new file, test_model_cache_unlock_restores_casts.py, goes through ModelCache lock → cast → unlock, on meta and CUDA, with and without a RAM copy. It covers:
  • partial residency, then a further load and a partial offload: your case;
  • the full-load class, and a cpu_only entry;
  • a module with mixed dtypes and a persistent buffer, and a partly cast model;
  • nesting, a failed restore, and a restore during a live init_empty_weights() on another thread.
• After every step the tests check:
  • the dtype of every key;
  • exact values;
  • per-key device placement;
  • storage shared with the RAM copy;
  • ties;
  • device bytes against cur_vram_bytes() and total bytes against total_bytes().
• The node tests run l2i and i2l through a real ModelCache. Two decodes price the same nonzero upcast, and a failed encode leaves the VAE in its loaded dtype.
• On an RTX 4090 I also ran a real fp16 SDXL VAE about half resident, with and without a RAM copy: fp32 decode, offload and reload, then a second fp32 decode. After each step everything was fp16 and the byte counts held. The two decodes were bit-identical.

[JPPhoto]

The latest commit restores cast weights to their counted dtypes on the last cache unlock, resolving the partial-offload blocker. But there's a merge blocker:

1. invokeai/backend/model_manager/load/model_cache/counted_dtypes.py:35: A tied group of distinct Parameter objects can be cast to fp32 and back before unlock. The helper skips every member once its dtype matches the counted dtype, leaving separate storages. Multi-GPU shared-weight adoption can create this alias shape via load_state_dict(assign=True) in cached_model_only_full_load.py:74-75. The cache counts one storage while two remain resident, risking memory overcommit and OOM. Effect: later cache loads may fail. Likelihood: requires a holder to change a tied model's dtype; the cache supports these aliases. Recovery: evict and reload the entry. Test: reproduced 32 resident bytes versus 16 counted; add a shared-store cast-round-trip test asserting storage identity and byte counts.

Suggested follow-up work (pre-existing):

2. invokeai/app/invocations/vae/flux_vae_encode.py:95 and invokeai/app/invocations/vae/flux_vae_decode.py:97: An SDNQ packed first weight can make input tensors cast to uint8. Suggested issue: cast inputs to the VAE compute dtype.