fix(working-memory): reserve VAE working memory per dtype and allocator

[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.}