[Pfannkuchensack]

Summary

A model that only partly fits in VRAM copies its offloaded weights to the GPU on every forward. Today each copy is a
blocking t.to(device) from pageable RAM on the compute stream, so the GPU waits for every copy. The new setting
async_weight_streaming changes that. It is opt-in and off by default, like fp8_compute and int8_compute.

With it on, the weights are pinned in place and copied on a per-device side stream, so the copies overlap with
compute. Images stay bit-identical. In the runs below it cuts the time per image of a partially loaded model by a
third (Linux) to more than half (Windows).

How it works:

• Pinning. The RAM copy of the streamed weights is page-locked in place with cudaHostRegister
  (hipHostRegister on ROCm), byte-exact per storage, through the runtime library torch uses. This happens after the
  model lock settles, outside the global load lock and the cache lock.
• Self-test. A start-up self-test checks the driver behavior this relies on. If it fails, copies stay synchronous
  and one log line says why.
• Copies. cast_to_device() sends CPU-to-GPU copies of dense weights to the side stream. At most 384 MiB of copies
  are in flight. The compute stream waits on each copy's event, and small tensors go through a pinned staging buffer.
• Pin lifetime. Pins stay while the model is in the RAM cache, so models that take turns in VRAM pin their weights
  once. Pinned RAM is capped at 40% of physical RAM. At the cap, Invoke unpins weights that no running model streams
  right now, the longest unused first, and only as many as the new registration needs.
• VRAM. The model cache reserves min(384 MiB, streamed bytes) of VRAM for copies in flight. The side stream's
  cached blocks are returned at unlock, so the VAE decode peak does not grow.
• Fallbacks. These keep the synchronous copy: GGUF, SDNQ and other tensor subclasses, LoRA sidecar patches,
  keep_ram_copy_of_weights: false, and devices other than CUDA/ROCm.

Related Issues / Discussions

None.

QA Instructions

Performance

Setup:

• family E2E, 1024², 8 steps;
• VRAM cache capped at 8 GB, so the transformer loads only partly;
• max_cache_ram_gb: 40, so the text encoder and the transformer stay in the RAM cache between images;
• off/on interleaved, 1 warm-up and 3 measured images per case, arm and round, median time per image.

GPU, OS	Model	Transformer in VRAM	off	on
RTX 4090, Windows	ERNIE-Image Turbo bf16	~18%	20.3 s	8.7 s
RTX 4090, Windows	Z-Image Turbo bf16	~24%	15.1 s	6.7 s
RTX PRO 6000 Blackwell, Linux (Verda VM)	ERNIE-Image Turbo bf16	~31%	6.05 s	3.71 s
RTX PRO 6000 Blackwell, Linux (Verda VM)	Z-Image Turbo bf16	~40%	4.56 s	3.09 s
RX 9060 XT (PCIe 4.0 x4), Windows, ROCm	Krea-2 Turbo fp8	~38%	61.9 s	33.7 s

• Every measured image is pixel-identical between off and on.
• The VRAM peak with the setting stays within 50 MB of the peak without it.

When the RAM cache is too small. The RAM cache can be too small to hold the models of one generation. Without
max_cache_ram_gb, the RAM cache is capped at the VRAM, or at max_cache_vram_gb when that is set. In that case
Invoke reloads the text encoder and the transformer for every image, and each reload pays for the first pin again. On
Linux, pinning weights still mapped from their safetensors file costs about 0.3 s/GiB, because the driver makes the
file's pages private.

With the default RAM cache on that VM (8 GB), ERNIE-Image Turbo took 9.4 s per image with the setting and 8.5 s
without it. The docs say this, and the log says it once when a model with pinned weights leaves the RAM cache.

Tests

• tests/backend/model_manager/load/model_cache/cached_model/test_async_weight_streaming.py: 36 passed on an RTX 4090
  (Windows); on Linux CUDA the same tests passed before the last one gained its log-line check. Ten mutations of
  the pin lifecycle each fail the test meant to catch them.
• tests/backend/model_manager/load, tests/test_config.py and tests/app/services/config: 1723 passed.
  test_16_channel_vae_loader fails locally because its fixture is an LFS pointer in that checkout; it is unrelated.
• On Linux CUDA, tests/backend/model_manager/load/model_cache, tests/backend/model_manager/load/test_loaded_model.py
  and tests/test_config.py pass.
• Ruff check and format are clean. OpenAPI, schema.ts, the settings reference and the docs search index are
  regenerated, and the docs build passes.

To try it, set async_weight_streaming: true in invokeai.yaml and run a model larger than the VRAM cache, or lower
max_cache_vram_gb. The log shows async_weight_streaming: weights in RAM stream to cuda:0 asynchronously. once.

Review

Resolved before opening:

• GGUF and SDNQ. Their weights crashed the pin collection; tensor subclasses are now skipped.
• Shared pages. Weights from safetensors.load_file share pages, and a copy that spans two registrations fails.
  Each storage is now registered by its exact range, and the self-test checks this.
• Pending driver error. A refused cudaHostRegister left CUDA error 712 pending until the next unrelated
  cudaMalloc. Registration now goes through the runtime library and clears the error at once.
• Self-test leaks. The self-test never frees a buffer it could not unregister, and it does not pass when its last
  unregistration fails.
• VAE decode peak. The side stream's cached blocks added 846 MiB at VAE decode; they are released at unlock.
• Re-pinning and the cap. Models that took turns in VRAM re-pinned everything on every lock, and at the cap a model
  could not reuse room held by its own pins. Pins now stay, are marked idle at the last unlock, and are reclaimed per
  storage.
• Unload within a lock. Weights that a node moves to RAM within its lock (Wan's second expert, Ideogram 4's second
  branch) are pinned for the rest of that lock.

Remaining limits:

• RAM cache. The gain needs a RAM cache that holds the models of one generation (see QA).
• Locked RAM. Pinned RAM stays page-locked while its model is in the RAM cache, up to 40% of physical RAM.
• Not tested:
  • two GPUs at once;
  • two models whose streamed weights together exceed the cap, on Linux.

Compatibility / Rollout

New setting, default off. With the default, the copy path is unchanged. The OpenAPI schema gains one field.

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)

---

^{Open in [code]smith · Need help on this PR? Tag @codesmith-bot with what you need.}