feat: prefetch streamed layers during compute - #1905
Open
assouan wants to merge 2 commits into
Open
Conversation
Adds async layer prefetching through `ModelManager`, allowing upcoming segments to be loaded ahead of execution. Layer streaming is now configurable end-to-end with resident-layer and prefetch-depth limits, exposed through the public API and the new `--resident-layers` and `--layer-prefetch-depth` CLI/server options. The runner derives the streaming policy from the graph cut, accounts for streaming allocations when enforcing VRAM budgets, and dynamically evicts or falls back when needed to stay within budget. Existing `new_sd_ctx` callers remain fully compatible.
Reworks stream-layer prefetch bookkeeping to use graph-cut segment parameter allocations directly, with a new shared segment-parameter map used by prefetching, residency retention, and eviction. Prefetch scheduling now explicitly targets parameter-bearing segments and keeps async transfers limited to non-shared cross-segment params. The docs, CLI option text, and public header comments were updated to match this behavior.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
This PR adds asynchronous lookahead to
--stream-layers. While the currentparameter-bearing graph-cut segment is computing, parameters for future segments
can be copied from CPU memory through a separate transfer backend or queue when
supported. This changes layer streaming from a serial load-then-compute path into
a rolling pipeline that can hide parameter-transfer latency behind useful GPU
work.
--layer-prefetch-depth <N>controls that pipeline. Its default0preservesthe existing synchronous behavior,
1overlaps the next segment with the activesegment, and larger values request deeper lookahead. The effective depth is
reduced automatically when it does not fit: prefetch never forces a positive
--max-vrambudget to be exceeded, and the active segment is never evicted.When asynchronous transfer is unavailable or a runtime prefetch allocation
fails, execution safely continues with synchronous streaming.
The implementation also adds
--resident-layers <N|auto>as a maximum, ratherthan a required count. This control is important for the prefetch pipeline:
--resident-layers 0 --layer-prefetch-depth 1disables optional residency and,when the prefetch fits the VRAM budget, maintains the minimal rolling window of
one active segment and one segment loading ahead.
auto(the default, representedby
-1) retains as many leading segments as the remaining budget permits, whilea positive value caps that resident prefix.
The streaming path additionally:
shared by resident, active, and prefetched segments;
same allocation if that segment is resident;
pressure would otherwise prevent prefetching;
not retain one-shot preparation graphs across the sampling loop;
versioned
sd_layer_stream_params_textension andnew_sd_ctx_with_layer_stream();their VRAM-budget semantics.
Related Issue / Discussion
RunnerWeightManager.Additional Information
Effective behavior
--layer-prefetch-depth 0--layer-prefetch-depth 1--layer-prefetch-depth NNfuture segments, reduced to fit the VRAM budget--resident-layers 0 --layer-prefetch-depth 1--resident-layers autoor-1--resident-layers NNleading parameter-bearing segments resident--max-vram 0Both controls only apply when
--stream-layersis enabled and the diffusionparameters use the CPU backend.
Verification
220/220compile and link steps),including the CLI and server.
(
13/13incremental steps); the CUDA server target also built successfully.sampling steps, and seed 42:
--max-vram cuda1=2 --resident-layers 0 --layer-prefetch-depth 1selectedzero residents and depth one; all 102 queued prefetches were activated.
--max-vram cuda1=8 --resident-layers 50 --layer-prefetch-depth 10selected22 residents and depth ten; all 80 queued prefetches were activated.
failures, and all 22 output SHA-256 hashes matched their reference runs.
git diff --checkpasses, and targetedclang-formatverification reports noremaining replacements in the modified C/C++ lines.
CUDA hardware was available for integration testing. Vulkan, ROCm, and Metal
were not exercised on physical hardware in this environment.
Checklist