Repository navigation
fix(working-memory): reserve measured working memory for tiled VAE decodes and encodes - #9946
Pfannkuchensack wants to merge 275 commits into
Conversation
…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.
…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.
…ipelines in the plan
…e default VAE of SD mains
…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.
|
@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 4. JPPhoto: Agreed. 116e505 limits the sentence to Model Manager calibrations and says an API run without an 5. JPPhoto:
Follow-ups in 7787606:
6. JPPhoto, "probably not fix": PID reuse between 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 Related, on disk rather than in memory: the |
…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
…king-memory-user-calibration
…king-memory-denoise
…king-memory-flux2-denoise
…king-memory-tiled-vae
Summary
Tiled VAE operations reserved the untiled line at the tile area × 1.25, and
vae.fluxadded 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:The tiled estimates now come from measured lines:
tile_size <= 0to the VAE's own tile, as the nodes do.Default allocators, measured demand in GiB, and the reservation as a multiple of it (main → this PR):
Related Issues / Discussions
QA Instructions
Measured points
working_memory/data/vae.json. They come from the newscripts/calibrate_vae_tiling_working_memory.py, which runs each operation the way its node runs it, inside themodel_on_devicewindow.cudaMallocAsync; RX 9060 XT (torch 2.13+rocm10.0, Windows), expandable segments and native.vae.flux,vae.flux2andvae.sd(fp16/fp32), decode and encode.Bounds
UPPER_BOUND_EXCEPTIONSwith their cause and exact ratio: 16 on the default allocators, 90 undercudaMallocAsyncand MIOpen native.Tests
uv run --no-sync pytest -n 4over the working-memory, calibration-service, invocation, builder and calibration-script suites: 4785 passed.Review
Resolved findings
Remaining limits
cudaMallocAsyncand 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.cudaMallocAsync, one MIOpen-native encode, and FLUX.2 decode at a 1024 px tile on cuDNN native.Compatibility / Rollout
vae.sd,vae.fluxandvae.flux2move 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.vae.jsongrows to about 1.3 MB (32 KB gzipped).Checklist
What's Newcopy (if doing a release after this PR)Need help on this PR? Tag
@codesmith-botwith what you need.