feat: pool VRAM for streamed layer - #1906
Open
assouan wants to merge 3 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.
Reuse fixed-size VRAM slots for active, prefetched, and resident layers, avoiding repeated buffer allocations while respecting streaming limits.
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 the opt-in
--stream-layer-poolmode for streamed diffusion layers. Instead of repeatedly allocating and freeing backend buffers as segment parameters move through the streaming pipeline, it keeps one contiguous backend allocation alive for the sampling loop and reuses it in place.The allocation is divided into fixed-capacity slots. Every slot is sized for the largest backend allocation footprint among the private parameters of streamed graph-cut segments, so a segment can be loaded by overwriting an available slot without recreating its backend buffer. The main target is the minimal
--resident-layers 0 --layer-prefetch-depth 1pipeline: one slot is active while the next slot is filled, then the two slots alternate throughout sampling.The implementation also:
--max-vrambudget, while the budget continues to limit optional slots, prefetch, and residency;The flag is disabled by default and is enabled by its presence.
Related Issue / Discussion
feat: prefetch streamed layers during compute.This is intentionally a stacked change. Its branch was created from the head branch of #1905 because the pool directly reuses the graph-cut segment metadata, streaming policy, residency ownership, prefetch queue, and
ModelManagerstaging lifecycle introduced there. Starting frommasterwould require duplicating that foundation or maintaining a temporary incompatible implementation, creating avoidable conflicts in the same hot paths.This PR should therefore be reviewed and merged after #1905. Once #1905 lands, rebasing or retargeting this branch onto
masterwill leave only the pool-specific commit and diff.Additional Information
Effective behavior
--stream-layer-pool--resident-layers 0 --layer-prefetch-depth 1 --stream-layer-poolVerification
--resident-layers 0 --layer-prefetch-depth 1 --stream-layer-poolallocated one contiguous two-slot pool (2 x 678.81 MiB) and used no per-transfer allocation fallback;0.1 GiBbudget selected one mandatory678.81 MiBslot, reduced prefetch depth to zero, and completed successfully;0(0 / 4,325,376differing channels).git diff --checkpasses.Checklist