Skip to content

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

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

Pfannkuchensack wants to merge 78 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.

Runs each installed model family through a launched scratch server and compares PSNR, time and driver-level VRAM against accepted references.
Graphs come from the frontend builders via an offline graph export; the generate submission plan is shared with the workbench.
VAE estimators return a WorkingMemoryEstimate built from fitted and analytic terms; model_on_device accepts it as working_memory=.
The AutoencoderKL coefficients and their measured points move to backend/working_memory/data, checked by a data test; reservations are unchanged.
A probe records each estimated model_on_device window's device-level demand inside a calibration scope; results go to a per-device-class store under databases/working_memory.
Admin endpoints start, inspect, reset and export runs; the reserved queue origin is refused for ordinary enqueues.
VAE estimators price a measured slope and base per conv backend, allocator and dtype instead of scaling one bf16 constant by element size; CogView4 gets its own lines.
SD/SDXL nodes and their estimate share one run-dtype rule; on ROCm Windows an SDXL fp32 1024² decode reserves 6.2 instead of 14.6 GB and the UNet stays resident.
A lower second reading at a measured size no longer tilts a fitted line, and points that include cast weight growth no longer fit one; two shipped lines move by under 1 %.
…n fact

Environment ids now include windows, linux or other, so calibrations from different operating systems are kept apart; the OS build string moves to soft fact os_build.
… runs

start --calibrate submits each case through the calibration endpoint and saves the store export per server lifetime; scripts/working_memory/build_shipped.py turns exports into probe points and refits.
Probe points bound estimates from below only, and the builder refuses data that the shared bound rule rejects.
# Conflicts:
#	docs/src/content/docs/development/Architecture/model-manager.mdx
…anager

Adds a working-memory section to the model detail: plan, calibrate, stop and reset per GPU, with the measured, built-in and calibrated estimate per size.
User fits from finished calibration attempts price VAE decodes within their measured range (+10 % + 256 MiB on the fitted part) and never below built-in values above it.
New plan and summary routes; summary reports lines shared with other models and returns 404 for unknown models.
…calibration

Judges allocator fragmentation per calibration attempt with one rule shared with the shipped-data builder, and moves other models off the GPU before a measured calibration lock.
User fits use the shipped 5 % margin; the plan sizes cases from the capacity the cache will grant, including memory the driver keeps after earlier runs.
The table marks uncounted sizes with their reason and explains each value column.
# Conflicts:
#	invokeai/backend/util/vae_working_memory.py
#	tests/app/invocations/test_flux2_working_memory.py
# Conflicts:
#	invokeai/backend/util/vae_working_memory.py
Main added useNsfwChecker and useWatermark to the generation project settings; exported calibration and E2E graphs run without both, as the defaults do.
…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 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.
@lstein lstein added the 7.0.0 label Oct 19, 2026
@lstein lstein moved this to 7.0 Theme: Modular Design in Invoke - Community Roadmap Oct 19, 2026
…orking memory

The node prices its denoise through the estimator framework now, so the test patches that call instead of the removed _estimate_working_memory method.
…stworthy

The session processor starts after every other service; windows of failed or canceled nodes, failed counter resets and locks racing a window are stored invalid.
Keyed requests are identified by model, graph and batch data, are created once per run and admit all of their items or none; the reserved-origin 400 is declared.
…ser-calibration

# Conflicts:
#	docs/src/content/docs/development/Architecture/model-manager.mdx
#	invokeai/app/services/working_memory_calibration/working_memory_calibration_default.py
#	invokeai/backend/model_manager/load/load_base.py
#	invokeai/frontend/api/openapi.json
#	tests/app/services/session_processor/conftest.py
#	tests/app/services/working_memory_calibration/test_working_memory_calibration_service.py
…-memory-denoise

# Conflicts:
#	docs/src/content/docs/development/Architecture/model-manager.mdx
#	invokeai/app/invocations/z_image/z_image_denoise.py
#	invokeai/backend/working_memory/probe.py
…mory-tiled-vae

# Conflicts:
#	docs/src/content/docs/development/Architecture/model-manager.mdx
…river loop

A measuring process that raises, including a failed synchronization or peak reset, is recorded as a failed point and kept out of the point file while the run goes on.
@Pfannkuchensack
Pfannkuchensack marked this pull request as ready for review October 21, 2026 00:45

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 services PRs that change app services

Projects

Status: 7.0 Theme: Modular Design

Development

Successfully merging this pull request may close these issues.

3 participants