feat(working-memory): reserve FLUX.2 denoise working memory from measured lines

[Pfannkuchensack]

Stacked on #9919 (and through it #9910, #9899, #9898, #9897, #9894). Until those merge, the diff also shows their commits. This PR's own change is 4ba43cda56.

Summary

The FLUX.2 denoise reserved a hand-set 0.42 MiB per attended token, scaled linearly by transformer width, plus a fixed 1 GiB, on every platform. Fresh-process measurements on an RTX 4090 and an RX 9060 XT show this reserves too much. The denoise now runs through the working-memory estimator framework with one measured line per width.

FLUX.2 Klein 9B (width 4096), bf16	1024²	2048²	2048² + one 2024 px reference
Before (every platform)	2959 MiB	8120 MiB	14894 MiB
RTX 4090, native allocator	1788	6556	12814
RTX 4090, cudaMallocAsync	1186	4320	8434
RX 9060 XT, expandable segments (Windows)	1333	4169	7891
RX 9060 XT, native (Windows)	2005	6747	12970
Unmeasured: XPU, MPS, CPU, Linux ROCm, other allocators	at least the previous values

Estimator

• denoise.flux2 prices a fitted line per width: Klein 4B 3072, Klein 9B 4096, [dev] 6144. The size is the attended tokens (image + reference + caption) × batch. Batch and reference images lie on the line.
• Regional bias, the SDPA score matrix, LoRA patches, the dequantization copy and the fp8 cast transient are analytic terms.
• The node passes model_on_device(working_memory=…). User calibration stays VAE-only.

Measurement

• New scripts/calibrate_flux2_denoise_working_memory.py. Its old denoise section is removed, and so is the old script's VAE section.
• The fresh-process plumbing shared with the Z-Image script moves to scripts/working_memory/fresh_points.py: the point records, the server runtime, the child protocol and the driver. The cold-and-warm rule moves to invokeai/backend/working_memory/measurement.py.
• The fit-eligibility rule is one helper (fit.off_line_reasons).

Related Issues / Discussions

• Builds on feat(working-memory): reserve Z-Image denoise working memory from measured lines #9919 (Z-Image denoise) and the stack below it.
• Supersedes the per-width scaling of the [bug]: OOM using Hildegard nodes. #9500 fix (fix(flux2): estimate working memory for denoise and both VAE directions #9519).

QA Instructions

Measured points

• 144 fresh-process points ship in working_memory/data/denoise.json:
  • RTX 4090 (torch 2.13+cu130), native and cudaMallocAsync;
  • RX 9060 XT (torch 2.13+rocm10.0, Windows, AOTriton fused), expandable and native.
• Grid: random-init transformer, 2 blocks per stream, at each width, bf16:
  • 512–2048 px with a 512-token caption, plus 77-token captions;
  • 1 and 3 references, batch 2;
  • a large-token grid up to 33025 tokens (2048 px + one 2024 px reference).
• Each point sits at 1.05–1.26× its estimate. Two 77-token caption points sit at 1.42× and 1.40× on the native allocators and are listed exceptions: caption tokens cost less than image tokens, and the line has to cover the 512-caption points.
• 4 vs 2 blocks: equal on the RTX 4090; +5–7 % on the RX 9060 XT.
• Forced math kernel: 8.74 bytes per score element against the 17 charged.

In-server check with real models, through the calibrate route. Real-model demand divided by the 2-block point at the same tokens:

Card	Model	Ratio
RTX 4090, native	Klein 4B fp8	0.96–0.99
RTX 4090, native	Klein 9B fp8	0.96–0.97
RX 9060 XT, expandable	Klein 4B bf16	0.81–0.95
RX 9060 XT, expandable	Klein 9B fp8, after its fp8 cast term	0.71–0.77

A/B, family E2E tool, interleaved before → after → before, 1 warm-up + 3 measured, Klein 4B and Klein 9B at 1024²–2048²:

• RTX 4090, Klein 9B at 2048²: the transformer goes from 84 % to 94 % resident, and the time from 38.5–39.1 s to 37.5–38.0 s. Other cases are unchanged within run-to-run spread.
• RX 9060 XT: unchanged within the spread, transformer 100 % resident in both. A constant 288 MiB shared-memory reading appears equally before and after.

Tests

• uv run --no-sync pytest -n 4 over the working-memory, calibration-service, FLUX.2/Z-Image/VAE invocation, builder, calibration-script, fresh_points and run_app environment suites: 3805 passed.
• Ruff is clean. VAE and Z-Image data, fits and reservations are byte-identical to the base.
• Mutation checks, each caught: estimator terms against the frozen legacy formula, the legacy flag, the moved cold-and-warm rule.

Review

Resolved findings

• The calibration scripts duplicated about 120 lines of plumbing; it is now shared.
• The fit-eligibility rule was copied per estimator; it is now one helper.
• The old FLUX.2 script kept a stale VAE section and stale text; both are removed.
• The docs called 33025 tokens the node's maximum. The node caps neither side, so they now say "measured to 33025 tokens, extrapolated linearly beyond".
• An old 28068-token check row is retired: its shape is unreachable (several references are resized to 1024 px each), it was measured on a superseded torch build, and that range is now measured directly.

Remaining limits

• Linux ROCm native stays on the previous values (unmeasured).
• The RX 9060 XT [dev] lines are measured only to 1440 px and without three references.
• Above 33025 tokens (4096 px, or more references) the lines extrapolate linearly.

Compatibility / Rollout

• No persisted state. No API change.
• scripts/calibrate_flux2_working_memory.py keeps only its dispatch and SDPA checks.

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