Skip to content

fix(working-memory): reserve measured working memory for tiled VAE decodes and encodes - #9946

Open
Pfannkuchensack wants to merge 275 commits into
invoke-ai:mainfrom
Pfannkuchensack:feat/working-memory-tiled-vae
Open

Pfannkuchensack wants to merge 275 commits into
invoke-ai:mainfrom
Pfannkuchensack:feat/working-memory-tiled-vae

Conversation

@Pfannkuchensack

@Pfannkuchensack Pfannkuchensack commented Oct 19, 2026 •

Copy link
Copy Markdown
Member

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

Summary

Tiled VAE operations reserved the untiled line at the tile area × 1.25, and vae.flux added analytic full-image terms on top. Fresh-process measurements on an RTX 4090 and an RX 9060 XT show the formula is wrong in both directions:

  • Decodes up to 4096 px over-reserve up to 1.6×.
  • Above 4096 px, decodes and encodes under-reserve, down to 0.38×. The decoded tiles plus the concatenated image (the assembly) outgrow the tile's own peak.

The tiled estimates now come from measured lines:

  • Tiled encode: the untiled line at the tile, clamped per axis, plus a fitted tiled line in the image's pixels.
  • Tiled decode: the larger of that and a fitted assembly line.
  • Fit: the tiled line is fitted on what tiling adds over the untiled demand measured at the tile, so it carries over to tiles that were not measured. The tiled and assembly lines are fitted together, by a deterministic search over the upper envelopes.
  • SD estimate: it now clamps a covering tile to a single pass and resolves tile_size <= 0 to the VAE's own tile, as the nodes do.
  • User calibration: a tiled operation prices its tile on the larger of the user's and the shipped untiled line. What tiling adds stays built-in.

Default allocators, measured demand in GiB, and the reservation as a multiple of it (main → this PR):

Operation (tile) Card 4096 px 8192 px 12288 px
SDXL fp32 decode (512) RTX 4090 1.63: 1.44 → 1.43 4.94: 0.48 → 1.34 10.10: 1.37
SDXL fp16 decode (1024) RTX 4090 4.31: 1.32 → 1.10 4.32: 1.32 → 1.25 6.17: 1.05
FLUX.1 decode (1024) RTX 4090 4.28: 1.40 → 1.11 4.34: 1.67 → 1.26 6.19: 1.07
SDXL fp32 encode (512) RTX 4090 1.06: 0.97 → 1.09 1.66: 0.62 → 1.12 2.66: 1.14
FLUX.1 decode (1024) RX 9060 XT 6.66: 1.35 → 1.07 6.98: 1.47 → 1.09 7.78: 1.09
SDXL fp32 decode (512) RX 9060 XT 1.87: 1.02 → 1.26 3.07: 0.62 → 1.55 6.58: 1.33
SDXL fp32 encode (512) RX 9060 XT 1.29: 1.01 → 1.08 1.92: 0.68 → 1.09 2.92: 1.12

Related Issues / Discussions

QA Instructions

Measured points

  • 737 tiled points and 118 untiled reference points ship in working_memory/data/vae.json. They come from the new scripts/calibrate_vae_tiling_working_memory.py, which runs each operation the way its node runs it, inside the model_on_device window.
  • Cards: RTX 4090 (torch 2.13+cu130), native and cudaMallocAsync; RX 9060 XT (torch 2.13+rocm10.0, Windows), expandable segments and native.
  • Grid: vae.flux, vae.flux2 and vae.sd (fp16/fp32), decode and encode.
    • Tiles 512 and 1024 at 1024–12288 px.
    • Tiles 256 and 768 at 2048–8192 px.
    • Tile 1536 at 2048–12288 px, RTX 4090 only.
    • Untiled references at every tile side.
  • Tile 1536 was first priced as a holdout, outside the fit: 0.98–1.33×. In the fit it reaches 1.06–1.31×.

Bounds

  • Every point is covered (the lowest estimate/measured is 1.044).
  • Capping points stay ≤ 1.3× or are listed in UPPER_BOUND_EXCEPTIONS with their cause and exact ratio: 16 on the default allocators, 90 under cudaMallocAsync and MIOpen native.
  • Points below 512 px (a tile or image the app never sets) bound from below only.

Tests

  • uv run --no-sync pytest -n 4 over the working-memory, calibration-service, invocation, builder and calibration-script suites: 4785 passed.
  • Ruff is clean.
  • Denoise data and reservations are unchanged.
  • Mutation checks are each caught. They cover:
    • the old 1.25 margin;
    • adding the assembly line instead of taking the max;
    • the offset taken from the shipped line;
    • a user line ignored above shipped;
    • the small-size rules;
    • the tie-breaks;
    • the margin on a negative intercept.

Review

Resolved findings

  • The tiled line did not carry over to unmeasured tiles. The holdout tiles 256 and 768 measured 0.61–0.99×. The tiled line is now fitted over the measured untiled reference, and those tiles are in the fit.
  • Extrapolation past 8192 px under-reserved, down to 0.63× at 12288 px. 12288 px is now measured.
  • User calibration no longer reached tiled operations. Their tile is now priced on max(user, shipped).
  • Smaller fixes:
    • the per-sample batch base of the assembly line;
    • negative SD tile fields;
    • FLUX fp16/fp32 stand-ins that copied SDXL's lines;
    • a builder path that skipped tiled lines silently;
    • duplicate formula tests.

Remaining limits

  • One tiled line serves every tile. Under cudaMallocAsync and MIOpen native it over-reserves by up to about 2.35× and 1.72× at some tile and size combinations. That is extra offloading, never a failure.
  • Measured to 12288 px.
    • At app tiles the estimates stay above the trend of the last measured segment to at least 15,800 px.
    • Four exceptions reach only about 12,800–15,100 px: two under cudaMallocAsync, one MIOpen-native encode, and FLUX.2 decode at a 1024 px tile on cuDNN native.
    • The upscale tab's 16× scale reaches past that range.
  • Unmeasured platforms floor at the per-value maximum of the measured ones. On CPU/XPU/MPS that is up to 1.29×; on Linux ROCm up to 2.5× for encodes, mostly from the untiled stand-in. It is conservative by design until those platforms are measured.
  • Tile 1536 is measured on the RTX 4090 only.

Compatibility / Rollout

  • vae.sd, vae.flux and vae.flux2 move to estimator version 3. Stored user calibrations of those estimators from an earlier version are shown as unusable and need a new calibration. Neither version has been released yet.
  • No API change. vae.json grows to about 1.3 MB (32 KB gzipped).

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)

View with [code]smith Autofix with [code]smith
Need help on this PR? Tag @codesmith-bot with what you need.

…e baseline

Raises otherRawBytes to the measured 215114 B on every route; the 3791 B growth are the Model Manager working-memory strings.
…f the model page

The section starts collapsed and shows the GPU with the calibration state, or the run progress, in its header; the plan loads on first open.
The mock prices sizes above the measured range with the larger of the user line and built-in, as the server does.
Main added useNsfwChecker and useWatermark to the generation project settings; exported calibration and E2E graphs run without both, as the defaults do.
…sured lines

Prices the Z-Image denoise through an EstimatorSpec whose per-token lines are fitted from fresh-process measurements on CUDA and ROCm, with an fp8 cast term and the legacy constants as the floor where nothing was measured.
Fitted lines can now run over tokens; the builder takes measured point files, and calibration scripts share the server runtime and the probe's measurement window.
# Conflicts:
#	docs/src/content/docs/development/Architecture/model-manager.mdx
#	invokeai/backend/model_manager/load/load_base.py
# Conflicts:
#	docs/src/content/docs/development/Architecture/model-manager.mdx
#	invokeai/backend/util/vae_working_memory.py
#	invokeai/backend/working_memory/coefficients.py
#	invokeai/backend/working_memory/data/vae.json
#	invokeai/backend/working_memory/platform.py
#	scripts/calibrate_flux2_working_memory.py
#	tests/app/invocations/test_flux2_working_memory.py
#	tests/app/invocations/test_vae_working_memory_backend.py
#	tests/app/invocations/test_z_image_tiled_decode.py
#	tests/backend/util/test_rocm_sdpa_head_dim_guard.py
#	tests/backend/working_memory/test_shipped_coefficients.py
…ured lines

Prices the FLUX.2 denoise through an EstimatorSpec with one measured line per transformer width, fitted from fresh-process points up to 33025 attended tokens on CUDA and ROCm, with the legacy constants as the floor where nothing was measured.
Calibration scripts share their fresh-process plumbing and the cold-and-warm rule.
# Conflicts:
#	invokeai/app/services/invocation_services.py
…ser-calibration

# Conflicts:
#	invokeai/frontend/webv2/performance/browser-baseline.json
#	invokeai/frontend/webv2/scripts/mock-backend-fixtures.test.mjs
#	invokeai/frontend/webv2/scripts/mock-backend.mjs
…e baseline

Raises otherRawBytes to the measured 224640 B on every route; the 3791 B growth are the Model Manager working-memory strings.
…codes and encodes

Tiled operations price the untiled line at the tile plus a fitted tiled line, and decodes the larger of that and a fitted assembly line, from fresh-process points to 12288 px with tiles 256–1536 on CUDA and ROCm.
User calibration prices the tile; the SD estimate clamps a covering tile and resolves tile fields at or below 0 like the nodes.
@github-actions github-actions Bot added api python PRs that change python files invocations PRs that change invocations backend PRs that change backend files services PRs that change app services frontend-deps PRs that change frontend dependencies frontend PRs that change frontend files labels Oct 19, 2026
…d trusts no reused server pid

Once the runner is dead, stop reads heartbeat.json again and kills and checks the server it names now.
The heartbeat records the serving process with the create time captured at its start.
…n unlock

A failed lock or paced continuation that drops the last pin now returns cast weights to their counted dtypes too, all in one release path.
Records that are detached or about to be evicted as stale are skipped, and the docstrings state how a tied weight split across devices is re-pointed.
…r the fits lock

Outcomes read from the queue are kept per run and item id, so a renumbered database cannot let one run's item decide another's attempt.
The summary, the attempt check and resets judge them under the lock removals hold, so a reset cannot race a judgement.
…lan from accepts_vae

The planner asks the shared rule the model loaders and Generate's picker read, latent channels and variant included, instead of restating it.
…ixel line

Only a spec keyed by no platform trait may carry a hand-set value; a keyed spec's value off a pixel line fails the data test.
A fit-grid forward whose estimator line does not admit its features ships with in_fit false, beside the checks.
@Pfannkuchensack

Copy link
Copy Markdown
Member Author

@JPPhoto, on your review of 2026-10-26:

The planner, calibration service and low-VRAM docs belong to #9910; the fixes below reached this PR through the merge 365899b.

3. JPPhoto (review 2026-10-26), blocker: cached non-canceled verdict outlives a cleared retry

Already fixed: 3022a69 caches outcomes per item (a retry is a new item), and f1080e2 records each retry on its run as the queue commits it. A cleared retry is looked up, not found, counted as canceled, and the attempt becomes aborted_attempt. Test (270dc24): test_a_retry_of_a_judged_failed_step_cleared_before_the_next_refresh_aborts_its_attempt, your exact sequence.

4. JPPhoto: low-vram-mode.mdx discard statement and API runs

Agreed. 116e505 limits the sentence to Model Manager calibrations and says an API run without an attempt_id keeps its sizes if stopped. 47168bf names how to remove them (DELETE /api/v2/models/working_memory/runs/{run_id} or a reset in the Model Manager).

5. JPPhoto: _verdicts grows across reset cycles

_verdicts went away in 3022a69; its successor _item_outcomes and _run_files were still not evicted on reset. Fixed in dcf8ced: the three resets drop the removed runs' entries under the fits' lock. Test: test_each_reset_forgets_what_the_service_kept_of_the_runs_it_removed.

Follow-ups in 7787606:

  • Outcomes are keyed by run and item id (a recreated database can reuse item ids). Tests: test_a_new_attempt_whose_item_id_an_old_run_recorded_counts_when_judged_after_the_old_run, test_an_old_run_whose_item_id_a_new_attempt_took_is_not_judged_by_the_new_item.
  • Every judgement runs under the removal lock. Test: test_every_judgement_of_item_outcomes_holds_the_lock_removals_hold. The write-back race itself is not deterministically testable, since the lock rules out the interleaving.

6. JPPhoto, "probably not fix": PID reuse between communicate() (scripts/family_e2e/graphs.py:119) and kill_group (scripts/family_e2e/server.py:470)

Agreed, no change: the code belongs to the family E2E PRs (#9894/#9899), and reuse needs every group member gone and the PID counter wrapping inside the gap before terminate().

Related, on disk rather than in memory: the removed_run_ids tombstones of keyed runs also grow with every reset; that predates this stack and is filed as #10056.

…ibration

# Conflicts:
#	docs/src/content/docs/development/Architecture/model-manager.mdx
#	invokeai/app/util/torch_cuda_allocator.py
#	invokeai/backend/working_memory/coefficients.py
#	invokeai/backend/working_memory/fit.py
#	invokeai/backend/working_memory/platform.py
#	tests/backend/model_manager/load/model_cache/test_model_cache_unlock_restores_casts.py
#	tests/backend/working_memory/test_fit.py
#	tests/backend/working_memory/test_shipped_coefficients.py
The runner checks the run is still running, spawns the server and records it in the heartbeat in one run-lock hold, so stop either sees the server or no server starts.
A server whose identity cannot be recorded is killed before the lock is released.
Points of estimators keyed by no platform trait name no platform in build_shipped.
…er model identity

Calibrate requests carry graph_models, the hash, base, type, format and variant of every model the graph names as the client compiled it; a difference from the current records, or from a keyed repeat's run, is refused with 409 model_changed.
The model manager and the family E2E tool send the identities of the records they compiled over.
# Conflicts:
#	invokeai/backend/working_memory/fit.py
#	scripts/working_memory/build_shipped.py
# Conflicts:
#	docs/src/content/docs/development/Architecture/model-manager.mdx
#	invokeai/backend/working_memory/bounds.py
#	invokeai/backend/working_memory/fit.py
#	scripts/working_memory/build_shipped.py
#	tests/backend/working_memory/test_fit.py
#	tests/backend/working_memory/test_shipped_coefficients.py
# Conflicts:
#	invokeai/backend/working_memory/bounds.py
# Conflicts:
#	docs/src/content/docs/development/Architecture/model-manager.mdx
#	invokeai/backend/working_memory/bounds.py
#	invokeai/backend/working_memory/coefficients.py
#	invokeai/backend/working_memory/fit.py
#	scripts/working_memory/build_shipped.py
…ibration

# Conflicts:
#	docs/src/content/docs/development/Guides/family-e2e.mdx
#	scripts/family_e2e/__main__.py
#	scripts/family_e2e/client.py
#	scripts/family_e2e/records.py
#	scripts/family_e2e/references.py
#	scripts/family_e2e/report.py
#	scripts/family_e2e/runner.py
#	scripts/family_e2e/server.py
#	tests/family_e2e/fake_invokeai_server.py
#	tests/family_e2e/test_cli.py
#	tests/family_e2e/test_references.py
#	tests/family_e2e/test_report.py
#	tests/family_e2e/test_runner.py

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

7.0.0 api backend PRs that change backend files docs PRs that change docs frontend PRs that change frontend files frontend-deps PRs that change frontend dependencies invocations PRs that change invocations python PRs that change python files python-tests PRs that change python tests Root services PRs that change app services

Projects

Status: 7.0 Theme: Modular Design

Development

Successfully merging this pull request may close these issues.

4 participants