Skip to content

feat: pool VRAM for streamed layer - #1906

Open
assouan wants to merge 3 commits into
leejet:masterfrom
assouan:feature/implement-stream-layers-pool
Open

feat: pool VRAM for streamed layer#1906
assouan wants to merge 3 commits into
leejet:masterfrom
assouan:feature/implement-stream-layers-pool

Conversation

@assouan

@assouan assouan commented Aug 23, 2026

Copy link
Copy Markdown

Summary

This PR adds the opt-in --stream-layer-pool mode 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 1 pipeline: one slot is active while the next slot is filled, then the two slots alternate throughout sampling.

The implementation also:

  • allocates only the peak slot count required by the effective current, prefetch, and residency policy rather than every slot the budget could hold;
  • keeps prefetch ahead of optional residents when slots are constrained and never evicts the active segment;
  • attempts progressively smaller monolithic pools when the requested allocation fails, without changing the slot size or falling back to per-transfer backend allocations;
  • treats the mandatory one-slot pool as execution memory, so it is still attempted when its size exceeds the effective --max-vram budget, while the budget continues to limit optional slots, prefetch, and residency;
  • keeps the pool alive when asynchronous lookahead is unavailable or fails, allowing current segments to continue through synchronous transfers into the same reusable slots;
  • reuses a sufficiently large pool for the whole sampling lifecycle and remembers reduced runtime caps instead of repeating failed allocation or prefetch attempts;
  • exposes the mode through an additive layer-stream option and constructor, preserving the existing public structure layouts and entry points.

The flag is disabled by default and is enabled by its presence.

Related Issue / Discussion

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 ModelManager staging lifecycle introduced there. Starting from master would 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 master will leave only the pool-specific commit and diff.

Additional Information

Effective behavior

Configuration Pool behavior
No --stream-layer-pool Existing streamed-layer allocation behavior is unchanged
--resident-layers 0 --layer-prefetch-depth 1 --stream-layer-pool Two reusable slots when the budget permits: one current and one future segment
Budget below one slot The mandatory one-slot pool is attempted with a warning; optional prefetch and residency are reduced
Requested pool allocation fails The slot count is reduced down to one while preserving one contiguous allocation and the original slot size
One-slot allocation fails Context creation fails clearly; there is no hidden fallback to per-transfer allocations
Asynchronous prefetch fails The pool remains active and subsequent current segments continue through synchronous transfers into pool slots

Verification

  • The standard Windows MinGW Release build completed for the shared library, CLI, and server.
  • CUDA builds completed for architectures 89 and 120a.
  • Vulkan compilation completed against Vulkan 1.3.275; physical Vulkan inference was not available in this environment.
  • CUDA integration with MiniMax-H3 confirmed:
    • --resident-layers 0 --layer-prefetch-depth 1 --stream-layer-pool allocated one contiguous two-slot pool (2 x 678.81 MiB) and used no per-transfer allocation fallback;
    • an effective 0.1 GiB budget selected one mandatory 678.81 MiB slot, reduced prefetch depth to zero, and completed successfully;
    • pool and non-pool reference runs produced 22 decoded frames with a maximum absolute channel difference of 0 (0 / 4,325,376 differing channels).
  • Focused local validation covered pool planning, allocation degradation, slot reuse, resident eviction order, prefetch failure isolation, lifecycle cleanup, and C API compatibility.
  • git diff --check passes.

Checklist

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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant