Skip to content

feat: prefetch streamed layers during compute - #1905

Merged
leejet merged 4 commits into
leejet:masterfrom
assouan:feature/implement-stream-layers-prefetch
Sep 6, 2026
Merged

leejet merged 4 commits into
leejet:masterfrom
assouan:feature/implement-stream-layers-prefetch

Conversation

@assouan

@assouan assouan commented Aug 22, 2026

Copy link
Copy Markdown
Contributor

Summary

This PR adds asynchronous lookahead to --stream-layers. While the current
parameter-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 default 0 preserves
the existing synchronous behavior, 1 overlaps the next segment with the active
segment, and larger values request deeper lookahead. The effective depth is
reduced automatically when it does not fit: prefetch never forces a positive
--max-vram budget 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, rather
than a required count. This control is important for the prefetch pipeline:
--resident-layers 0 --layer-prefetch-depth 1 disables 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, represented
by -1) retains as many leading segments as the remaining budget permits, while
a positive value caps that resident prefix.

The streaming path additionally:

  • gives the requested prefetch window priority over optional resident segments;
  • accounts for canonical backend allocation sizes and deduplicates parameters
    shared by resident, active, and prefetched segments;
  • keeps cross-segment shared parameters on the safe synchronous path;
  • reuses prefetched allocations when a segment becomes active and can retain the
    same allocation if that segment is resident;
  • evicts residents from the end of the resident prefix when runtime memory
    pressure would otherwise prevent prefetching;
  • releases stale residents and prefetches when the graph plan changes, and does
    not retain one-shot preparation graphs across the sampling loop;
  • preserves the existing C API layout by exposing the controls through the
    versioned sd_layer_stream_params_t extension and
    new_sd_ctx_with_layer_stream();
  • exposes the controls consistently through the CLI and server and documents
    their VRAM-budget semantics.

Related Issue / Discussion

  • Reimplements the asynchronous-prefetch goal of #1626 on the current runner weight-staging architecture. That PR overlapped the next segment's H2D transfer with the current segment's compute, but its author closed it after #1644 removed the partial-buffer machinery it depended on and required the idea to be ported to RunnerWeightManager.
  • Builds on the resident-layer control discussion in #1830. In this PR, residency is a VRAM-bounded maximum and also provides the explicit zero-resident mode needed for the minimal active-plus-prefetch pipeline.

Additional Information

Effective behavior

Configuration Behavior
--layer-prefetch-depth 0 Existing synchronous layer streaming
--layer-prefetch-depth 1 Prefetch the next segment while the active segment computes, when it fits
--layer-prefetch-depth N Prefetch up to N future segments, reduced to fit the VRAM budget
--resident-layers 0 --layer-prefetch-depth 1 No optional residents; one active plus one prefetched segment when budget permits
--resident-layers auto or -1 Use the remaining budget for an automatically selected resident prefix
--resident-layers N Keep at most N leading parameter-bearing segments resident
--max-vram 0 Graph-cut segmentation remains disabled, so residency and prefetch controls have no effect

Both controls only apply when --stream-layers is enabled and the diffusion
parameters use the CPU backend.

Verification

  • CPU Release build completed successfully (220/220 compile and link steps),
    including the CLI and server.
  • CUDA Release build completed successfully for the shared library and CLI
    (13/13 incremental steps); the CUDA server target also built successfully.
  • Focused tests passed for:
    • VRAM-budget, resident-limit, and prefetch-depth policy boundaries;
    • physical staging-buffer release across partial tensor lifetimes;
    • resident eviction and stale-prefetch cleanup after graph-plan changes;
    • current and undersized versioned C API structures;
    • compatibility of an executable built against the previous public API.
  • CUDA integration was exercised with MiniMax-H3 at 256x256, 22 frames, two
    sampling steps, and seed 42:
    • --max-vram cuda1=2 --resident-layers 0 --layer-prefetch-depth 1 selected
      zero residents and depth one; all 102 queued prefetches were activated.
    • --max-vram cuda1=8 --resident-layers 50 --layer-prefetch-depth 10 selected
      22 residents and depth ten; all 80 queued prefetches were activated.
    • Both configurations generated all 22 frames without warnings or transfer
      failures, and all 22 output SHA-256 hashes matched their reference runs.
  • git diff --check passes, and targeted clang-format verification reports no
    remaining 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

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.
@leejet

leejet commented Sep 6, 2026

Copy link
Copy Markdown
Owner

Thanks for the PR! I kept the core prefetch functionality, but removed a few parts:

  • Prefetch depth: I don't think deeper prefetching helps much. If compute is faster than transfer, transfer is still the bottleneck; if compute is slower, a single prefetch is enough to overlap it.
  • Resident layers: this should be handled automatically by the algorithm rather than exposed as a user-facing option.
  • Eviction/residency and cross-forward state: these are currently too tightly coupled and conflict with the new weight-management model I'm planning.

I also reorganized the code, and the resulting implementation is now much simpler and more straightforward.

@leejet
leejet merged commit 6c57cc3 into leejet:master Sep 6, 2026
9 checks passed
danielhanchen added a commit to unslothai/stable-diffusion.cpp that referenced this pull request Sep 21, 2026
* fix: preserve "token_refiner" token for MiniMax H3 LoRAs (leejet#1864)

* fix: fail with a message when MiniMax-H3 is run in img_gen mode (leejet#1863)

* feat: support INT8 ConvRot safetensors (leejet#1857)

* fix: replace free_compute_buffer with runner_done in vae (leejet#1872)

* sync: update ggml (leejet#1873)

* fix(ci): trigger builds for ggml updates

* feat: add taeh3 support (leejet#1874)

* fix: prevent gallocr hash overflow in tiny graph-cut segments (leejet#1880)

* fix: re-clamp streaming VRAM budget to currently free memory (leejet#1878)

* fix: mark graph cuts with both a prefix and a suffix (leejet#1883)

* fix: make max_order of lms sampler configurable (leejet#1885)

* fix: guard against missing sampler/scheduler names (leejet#1887)

* chore: format code

* fix: use sd_get_preview_interval() (leejet#1907)

* feat: configurable image / video compression (leejet#1909)

* feat: support standard Qwen3-VL weights for MiniMax-H3 (leejet#1910)

* feat: load scaled FP8 weights without upfront conversion (leejet#1913)

* fix: match exact weights in LLM config detection (leejet#1923)

* feat: add LTX-2.5 support (leejet#1893)

Co-authored-by: leejet <leejet714@gmail.com>

* fix: correct MiniMax H3 reference audio encoding (leejet#1886)

* fix: correct MiniMax H3 audio Euler steps (leejet#1908)

* feat: use backend-native FP8 matmul when supported (leejet#1916)

* sync: update ggml

* feat: additional `--preview-interval` values (leejet#1915)

Co-authored-by: leejet <leejet714@gmail.com>

* feat: support numbering for preview images (leejet#1895)

* fix: use carrier sampling for MiniMax H3 audio (leejet#1924)

* feat: generalize temporal tiling across video VAEs (leejet#1926)

* feat: prefetch streamed layers during compute (leejet#1905)

Co-authored-by: leejet <leejet714@gmail.com>

* refactor: unify runner lifecycles and weight residency (leejet#1940)

* feat: add verbose logging and log-level selection (leejet#1941)

* feat: enable single-GPU auto-fit with tiered parameter placement (leejet#1942)

* fix: reuse graph cut plans across CFG passes (leejet#1943)

* refactor: split ggml extensions and move implementations to cpp files (leejet#1945)

* fix: preserve K-quantized embedding weights (leejet#1936)

* fix: correct SDXL embeddings loading (leejet#1939)

* refactor: unify model source and weight lifecycle management (leejet#1956)

* docs: reflect GGML_MAX_NAME value change in rpc docs (and in ggml_extend assert) (leejet#1950)

* refactor: split generation pipeline out of stable-diffusion.cpp (leejet#1957)

* fix: enable VAE decode tiling fallback without auto-fit (leejet#1932)

* fix: preserve BF16 embedding weights for get_rows (leejet#1959)

* fix: handle invalid option numbers (leejet#1961)

* feat: expose the loaded model version name through the public API (leejet#1962)

* feat: add SenseNova U1.5 support (leejet#1935)

* fix: reuse graph plans when scale parameters change (leejet#1963)

* feat: add linear and attention scale overrides (leejet#1964)

* feat: preserve explicit backend assignments during auto-fit (leejet#1967)

* fix: guard GPU memory capacity and propagate encoding failures (leejet#1958)

* fix: bound plain-text runs in parse_prompt_attention regex (leejet#1919)

* feat: add Wan2.2 S2V (audio+img-to-video) support (leejet#1925)

Co-authored-by: leejet <leejet714@gmail.com>

* fix: validate vision projector output dim against LLM hidden size (leejet#1918)

* fix: resolve MSVC narrowing conversion warnings (leejet#1969)

* feat: Add generation parameters into video metadata (leejet#1901)

Co-authored-by: leejet <leejet714@gmail.com>

* feat: support external Hugging Face tokenizer JSON files (leejet#1973)

* refactor: require external Gemma 2 and GPT-OSS tokenizers (leejet#1974)

* feat: support Brownian tree noise in all noise injection samplers (leejet#1899)

* fix: use tokenizer-specific pre-tokenization rules (leejet#1975)

* fix: remove vision_model. from ununsed tensors (leejet#1983)

* perf: eliminate temporary allocations in Philox rounds (leejet#1982)

* fix: honor flash attention flag in LLM text encoder attention (leejet#1987)

* refactor: remove obsolete unused tensor filtering (leejet#1984)

* perf: pad small attention heads to 64 for MMA Flash Attention (leejet#1992)

* perf: update ggml for faster direct convolutions (leejet#1993)

* perf: accelerate VAE direct 3D convolutions (leejet#1996)

* fix: propagate CUDA driver dependency to shared library consumers

* fix: prevent clip_preprocess center crop from exceeding the resized image (leejet#1995)

* perf: reduce CPU overhead in graph execution and sampling (leejet#1997)

* perf: parallelize host tensor elementwise and broadcast ops (leejet#1998)

* feat: support building with upstream ggml (leejet#1999)

* feat: add Qwen Image 2.1 support (leejet#1994)

* feat: restore legacy fp8 handling when building with upstream ggml (leejet#2001)

* fix: avoid passing ggml logs as format strings (leejet#2002)

* feat: add LLaDA-Image support (leejet#1968)

Co-authored-by: leejet <leejet714@gmail.com>

* feat: add native CUDA SageAttention support (leejet#2005)

* fix: avoid narrowing conversion in SigVQ patch embedding and format code

* docs: update CONTRIBUTING.md

---------

Co-authored-by: stduhpf <stephduh@live.fr>
Co-authored-by: leejet <leejet714@gmail.com>
Co-authored-by: LostRuins Concedo <39025047+LostRuins@users.noreply.github.com>
Co-authored-by: fszontagh <51741446+fszontagh@users.noreply.github.com>
Co-authored-by: Wagner Bruna <wbruna@users.noreply.github.com>
Co-authored-by: vmobilis <75476228+vmobilis@users.noreply.github.com>
Co-authored-by: Piotr Wilkin (ilintar) <ilintar@gmail.com>
Co-authored-by: jk212h20 <101200018+jk212h20@users.noreply.github.com>
Co-authored-by: assouan <750048+assouan@users.noreply.github.com>
Co-authored-by: nan <zjn32202153@gmail.com>
Co-authored-by: Hmission <62598659+Hmission@users.noreply.github.com>
Co-authored-by: LED-M <105789115+xledx@users.noreply.github.com>
Co-authored-by: Maphist0 <28743569+Maphist0@users.noreply.github.com>
Co-authored-by: George <35490284+noctrex@users.noreply.github.com>
Co-authored-by: Санька Четвёртый <CAHbKA-IV@mail.ru>
Co-authored-by: Lin Xuhao <linxuhao84@gmail.com>
Co-authored-by: Fabrice Aneche <akhenakh@users.noreply.github.com>
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.

2 participants