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

[lstein]

Review of this PR's own diff (475dd91a95..1f55613faa, on top of #9919). I found no crash and no wrong-input wiring, and the main claims hold up (see the end). One finding needs a measurement or a term before merge; the rest are small.

Needs a fix or a justification

1. Dropping the 1 GiB base uncovers regimes that no point measures.
The CUDA native entry has b_denoise_bf16_h* = 0 at every width (denoise.json:116-121). The old formula added a flat 1 GiB. All 144 points are a fully resident 2-block model without a mask, so two real cases have less headroom than before, and nothing measures them:

• Partially loaded [dev] bf16. Weights streamed per forward under partial loading are not priced; probe.py:39-40 and the model-manager docs say so. [dev]'s largest single-block linear (to_qkv_mlp_proj, 6144 × 55296) is about 680 MB in bf16, and it is alive during its own matmul alongside the activations. At 1024² with one 1024² reference (8704 tokens) on CUDA native, the reservation drops from 6.35 GiB to 4.92 GiB. The line's margin over the envelope at that size is only about 250–600 MB. The real-model checks cover resident fp8 Klein 4B/9B; there is no [dev] window, and [dev] is where the gap is largest.
• Regional prompting. No point has has_attention_mask=true, so a masked forward prices as the fused line plus the bias, with no base. On CUDA the mask goes to the memory-efficient kernel, whose footprint with a bias is unmeasured. It may also pad the bias to an 8-aligned stride, making a second bias-sized copy; I have not verified this on this torch build. Example: Klein 9B at 1440² with two 512-token regional captions, S = 9124. That reserves 3.71 GB plus the bias, against 5.09 GB plus the bias before.

The cache's device_working_mem_gb floor (model_cache.py:2242) covers small sizes, but not these.

Ask: measure a partially loaded [dev] on CUDA (and ideally one masked forward), or price the largest non-resident layer as a term.

Smaller items

2. Fun-ControlNet lock under calibration (flux2_denoise.py:~553). The branch locks with plain reservation_bytes, while the transformer's calibration lock keeps calibration_reservation(R) free. Loading the branch can therefore use up the calibration headroom, and the window then spills or OOMs instead of measuring. No calibration graph includes a control branch today, so this cannot happen yet. Still, the comment "Locked with the same working memory" is wrong inside a calibration scope. Passing probe.calibration_reservation(...) when a scope is open would fix it.
3. Linux CUDA uses the Windows-measured lines without the unmeasured floor: MEASURED has ("cudnn", …, None). If that is intentional, say so in the module docstring.
4. The dry-run text in scripts/calibrate_flux2_denoise_working_memory.py is wrong (:37, :500). It says it prices "with the shipped values of cuDNN under the native allocator". A CPU is an unmeasured platform, so unmeasured_floor raises it to the legacy constants: --dry-run --widths 4096 shows 2.89 GiB at 4608 tokens, not the fit's ~1.75 GiB.
5. measured_bytes_and_method can record a point low (measurement.py:650). If the cold window has a driver reading and the warm one does not, the function falls back to torch counters only and drops the cold window's driver figure. On HIP that figure is ~245–300 MiB above reserved for a first forward. Scripts only.

Test gaps (a plausible regression passes)

• Script model geometry. No test ties the script's _transformer (calibrate_flux2_denoise_working_memory.py:283-292) to the loader's construction (model_loaders/flux.py). Mutating num_attention_heads=hidden_size // HEAD_DIM to 24 measures every width at Klein 4B's inner width and fits it under 4096/6144, and no test fails. Comparing inner_dim == hidden_size, or comparing against the loader's config, would catch this.
• Node wiring. No node test passes a regional mask or a negative caption. Both of these mutations survive:
  • regional_attn_mask=... if ref_image_seq_len == 0 else None changed to None;
  • text_seq_len=max(txt.shape[1], neg_txt...) changed to txt.shape[1].
• fresh_points.measure_in_this_process. The test does not pin the order of apply_server_runtime before measure, which is the contract it describes.
• Duplicated tests. TestFlux2DenoiseWorkingMemoryEstimate and TestFlux2DenoiseBatchIsBudgeted in test_flux2_working_memory.py largely duplicate the new file's coverage. Their docstrings still cite the retired measurements.
• test_each_widths_slope_covers_the_rx_9070_xts_closely evaluates on the legacy entry. TestLegacyReservation already freezes those constants.

Checked and holding

• The node passes the same features as before:
  • tokens = image + reference + max of the CFG pair's captions (the passes run sequentially);
  • batch from x.shape[0], with references repeated per sample;
  • the mask dropped when references are present.
• The control, LoRA and dequant terms are unchanged; the fp8 cast term only adds.
• The new ValueError is unreachable: the node always runs in bf16, and every variant maps to 3072, 4096 or 6144.
• The legacy constants are 330300, 440401 and 660601 bytes per token plus a 1 GiB base, matching the old arithmetic. Every unmeasured platform resolves to these values or above.
• The non-legacy FLUX.2 entries' coefficients are reproduced exactly from their points. Every point is covered at ≥1.05×; above 1.3× there are only the two listed exceptions (1.417 and 1.391).
• Refitting on points ≤9216 tokens keeps every long point below its line, which supports the extrapolation.
• Z-Image and VAE data and reservations are identical to the base. The off_line_reasons refactor keeps Z-Image's labels and fit membership.
• user_calibrated=False is honoured in user_fit, fit_user_lines and the calibration fits service.
• The Z-Image script refactor onto fresh_points preserves the child protocol, the provenance string, the cold-OOM rule and the CLI guards.

Not checked: GPU behaviour. Finding 1 rests on layer sizes and the shipped data, not on an observed OOM.

[joshistoast]

lgtm

[Pfannkuchensack]

Thanks — finding 1 was right. Measured on a 4090 (CUDA native), FLUX.2 [dev] bf16 partially loaded (15 of 64 GB resident) needed 5509 / 5459 / 3305 MB at 8704 / 8612 / 4608 tokens against 5285 / 5229 / 2798 MB reserved (0.96× / 0.96× / 0.85×). The reservation now adds the largest layer's stored bytes (679 MB for [dev] bf16; fp8/GGUF are charged their stored size), which covers all three at 1.05–1.08×. A masked forward at S=9124 cost one extra bias-sized buffer (+158 MB); the padded copy is now priced. Spec bumped to v2. Items 2–5 and the test gaps are fixed; your mutations now fail the tests.

[lstein]

Re-reviewed 07f31d88ec (head fdc0aaba79). All findings from the previous review are resolved:

• Finding 1: the largest stored layer is now reserved for partial loading, and the padded bias copy for masked attention. Your [dev] measurements sit at 1.05–1.08× with the new layer term.
• Shipped data: the only change in denoise.json is the FLUX.2 version bump to 2; the points, coefficients and Z-Image data are identical.
• Items 2–5 and the test gaps: fixed. I ran 1106 focused tests, all passing, and re-ran the mutations: the geometry, regional mask, negative caption, runtime order, streamed-layer and padded-bias terms, control-branch calibration reservation and mixed driver windows all fail the tests now.

Nit (non-blocking): changing the calibration script's weights_resident=True (scripts/calibrate_flux2_denoise_working_memory.py:~274) to False passes every test. If that happened, each newly measured point would carry streamed_layer_bytes, _line would put it off the line, and a calibration run would silently fit nothing. A test asserting that a script point's features have streamed_layer_bytes == 0 (or that it stays in_fit) would pin it.