feat: --stream-layers for streaming weights from CPU during generation - #1576
Conversation
Surface area for the unified-streaming design. No behaviour change
when --stream-layers is unset — the dispatch in compute<T>() (added in
the follow-up commit on src/ggml_extend.hpp) short-circuits to the
upstream walker.
- `sd_ctx_params_t::stream_layers` field (include/stable-diffusion.h).
- `--stream-layers` boolean CLI flag (examples/common/common.{h,cpp}).
- `sd::ggml_graph_cut::SegmentResidency` enum + Segment::residency
field + annotate_residency declaration (src/ggml_graph_cut.h).
- `sd::layer_registry::Registry` with register/move primitives using
the proven dup-copy-swap idiom on tensor->buffer/data/extra
(new src/layer_registry.{h,cpp}).
- Conditioner subclasses gain a virtual set_stream_layers_enabled
(src/conditioner.hpp). UpscalerGGML forwards to its inner runner
(src/upscaler.{h,cpp}).
Builds on the foundation commit (--stream-layers + planner annotation + executor scaffolding) with: ## chunk-K residency A parallel `resident_*` offload track on `GGMLRunner` keeps a fraction of the diffusion model's params on GPU permanently across sampling steps, amortising H2D over many invocations. - Members: `resident_offload_ctx`, `resident_offload_pairs`, `resident_runtime_params_buffer`, `resident_param_set`, `resident_state_token` (parallel to the existing `partial_offload_*` per-segment track). - `offload_resident_params(tensors)` / `restore_resident_params()` use the same dup-copy-swap idiom as `offload_partial_params` but write to the resident slot and persist across `compute()` calls. - `offload_partial_params` filters tensors already in `resident_param_set` so per-segment offload skips them. `restore_resident_params` is hooked into `~GGMLRunner()` and `free_params_buffer()` to keep swap pointers valid through teardown. - `compute_streaming_segments<T>` reads `graph_cut_plan_cache_.graph_cut_plan` (the unmerged base plan), annotates it, gathers the union of RESIDENT segments' param tensors, and offloads them once. Compute itself proceeds on the merged plan for fused-graph efficiency. A commutative pointer-hash state-token detects when a different plan is in play and rebuilds the resident set. `annotate_residency` updates: - "Any param-bearing segment exists" sanity replaces the `segments[0].input_param_bytes == 0` early-return (wrong for diffusion models whose first segment is a small prelude). - Greedy cumulative-bytes loop handles heterogeneous segment sizes (small prelude + large transformer layers). - Resets `seg.residency = STREAMED` at entry so cached plans don't carry forward stale RESIDENT marks from a previous larger-budget call. - Don't reserve a `prefetch_segments * largest_segment` window; async prefetch is no longer used (see below). ## Multi-runner safety - Per-runner free-VRAM clamp at compute time in `resolve_graph_cut_plan`. Each runner queries `ggml_backend_dev_memory(runtime_backend)` and clamps `effective_budget = min(max_vram, free - 512 MB)` per call. Without this, after the LLM committed ~7 GB chunk-K resident the diffusion runner still believed it had the whole budget and OOM'd. - `--stream-layers` is restricted to diffusion runners only (diffusion_model + high_noise_diffusion_model) — matches PR leejet#1477's scope and avoids one-shot runners (LLM, VAE, clip_vision, upscaler) claiming permanent chunk-K state that starves the diffusion model. - `GGMLRunner::release_streaming_residency()` (public trampoline to `restore_resident_params`) is called from `decode_first_stage()` on diffusion_model + high_noise_diffusion_model right before VAE decode. Without it the 6.5 GB chunk-K residency from sampling would starve VAE's compute buffer (~4.5 GB at full image resolution) and OOM at decode. ## `LORA_APPLY_AUTO` picks runtime when streaming or CPU-offload is on Immediate mode bakes LoRA into weights at load time by running a forward pass over every weight tensor — allocates a full-model-size (~11 GB on Z-Image bf16) compute buffer on the runtime backend in one shot and OOMs on any VRAM-constrained setup, which is the whole reason `--stream-layers` / `--offload-to-cpu` exist. AUTO previously only picked runtime for quantized models; now `stream_layers || offload_params_to_cpu` is also a trigger. ## Conservative streaming v4 (current shipping configuration) After observing edge-case failures with prefetch + multi-LoRA / + non-default --guidance/--flow-shift / + batch_count > 1, the final configuration is: - chunk-K residency skips itself when `weight_adapter != nullptr`. The state-token hashes tensor *pointers*, not data, so it can't detect MultiLoraAdapter modifications across batch images / steps; the symptom was colored static noise on batch image 2+. - Async prefetch is hard-disabled (`prefetch_enabled = false`). The `compute_streaming_segments_prefetch<T>` implementation stays in the file for future reference but is currently dead. Two correctness problems forced the disable: * Multi-LoRA workloads: graph_compute_async + per-segment pending offload races MultiLoraAdapter's per-layer patch_weight reads. * batch_count > 1 + non-default --guidance + --flow-shift: the smaller merged segments required to fit two prefetched buffers in --max-vram accumulate FP error across the extra boundary- cache roundtrips → collapses to pure white frames. - `resolve_graph_cut_plan` always passes the full `effective_budget` to the planner (no `/4` shrinking). Produces the upstream walker's large merged segments — the validated configuration. - The chunk-K hook reserves room for the LARGEST merged segment's params: `chunk_k_budget = max_graph_vram_bytes - largest_merged_segment`. Without this reservation chunk-K could grow large enough that the active merged-segment offload OOMs. ## Verification - Z-Image Q8 with `--stream-layers off`: byte-identical to upstream walker. - Z-Image bf16 1024x688 + 2 LoRAs + `--cfg-scale 1.0 --guidance 3.5 --flow-shift 3.0` + batch_count=2 (the production REST API failure mode that motivated the correctness work): both batch images clean and distinct. - Smoke-matrix-verified across Z-Image Q8/bf16, HiDream, Qwen, Flux schnell, SD3.5, SDXL, Anima, WAN.
2575d97 to
5cd49ca
Compare
|
Loved the previous PR, madly in love with this new one. It's going for a spin on my server asap 😍 |
I'm happy for it, thank you. Sorry, but now the preformance is slower than the prev. PR, and missing some optimizations yet which was implemented earlier. See Future work in the PR description. |
leejet
left a comment
There was a problem hiding this comment.
My main suggestion is to keep this PR focused on the chunk-K residency behavior, make the active path easier to reason about, and move the disabled async prefetch implementation to the follow-up PR where it will actually be fixed/enabled.
…treaming # Conflicts: # src/stable-diffusion.cpp
--stream-layers for streaming weights from CPU during generation--stream-layers for streaming weights from CPU during generation
|
Thank you for your contribution. |
Brings the upstream src-layout reorg (leejet#1615 model/, core/, conditioning/, runtime/, extensions/), the new offload path (leejet#1601 pinned host buffer, leejet#1576 --stream-layers), vram-limit propagation (leejet#1583), APG/PiD/ideogram4, and the photomaker->generation-extension move (leejet#1618). Conflict resolution (4 files): - model.h / stable-diffusion.cpp: union the fork's LONGCAT_AVATAR version with upstream's PiD/Ideogram4; keep the avatar deferred-DiT-load + per-frame timestep zeroing, adopt upstream's alloc error-checks + generation-extensions alloc loop (pmid is now an extension); keep whisper-encoder alloc. - conditioner.hpp: keep both set_keep_params_resident + set_stream_layers_enabled. - ggml_extend.hpp: keep the fork's coherent offload system (lap-32 pinned alloc, lap-32.2 H2D pipelining, partial/all-param restore, umT5 free-then-reload null fix, lap-28 F16-KV/mask attention) and fold upstream's persistent_externals snapshot + observed_max_effective_budget reset alongside; flash_skip_kv_pad opt-out coexists with upstream leejet#1453's unconditional kv-pad removal. - Repointed fork-only headers (longcat_avatar/audio, nava, nava example) at the new nested include paths.
04b55bba Merge pull request leejet#8 from V-Sekai-fire/sync/upstream-refresh 667804ef Sync ggml/ to upstream refresh. 52c7f4a8 Merge pull request leejet#7 from V-Sekai-fire/repoint-urls-at-v-sekai-fire 4eebc4a9 Point URLs at V-Sekai-fire d4d15152 Merge pull request leejet#6 from v-sekai-fabric/magi/citation 47fe255e Add CITATION.cff for the canonical ggml tree d646939a Merge pull request leejet#5 from v-sekai-fabric/gate-volk-behind-ifdef bcccfc16 ggml-vulkan: gate volk include behind #ifdef GGML_VULKAN_VOLK 06d5dab5 Merge pull request #4 from v-sekai-fabric/drop-empty-gitmodules 1302b98d Drop empty .gitmodules (CLAUDE.md submodule blocklist doctrine) eee2d470 Merge pull request #3 from v-sekai-fabric/revert-broken-dl-port 0d9079c4 Revert "Port turboquant-godot: guard ggml-backend-dl.h behind GGML_BACKEND_DL" 273a4fcf Merge pull request #2 from v-sekai-fabric/turboquant-godot-port 32567060 Port turboquant-godot: route ggml-vulkan through volk loader e3e3ec26 Port turboquant-godot: guard ggml-backend-dl.h behind GGML_BACKEND_DL 89bf35b3 Merge pull request #1 from weftspun/sam3-flash-attn-dk16-dk56 0827e29a metal : flash_attn_ext head_dim=16 and head_dim=56 3404c951 vulkan : local perf-log naming for ADD, GROUP_NORM and CONT c21b1a10 metal : implement diag_mask_inf REVERT: e20c3a1 ggml-cuda : add native FP8 matmul with cuBLASLt REVERT: 032b699 ggml : support FP8 casts across compute backends REVERT: 8e800ce ggml : remove standalone regular Hadamard op REVERT: 373c7f1 ggml-vulkan : add native INT8 convrot support REVERT: 4b053d2 ggml : fix HIP build for INT8 convrot REVERT: 86803f3 ggml : add native INT8 convrot support REVERT: 8846b79 cmake : add config version support (leejet#1582) REVERT: 30bf868 ggml : bump version to 0.19.0 (leejet#1581) REVERT: 77be358 sync : llama.cpp REVERT: 0f8a392 ggml : add aarch64 HWCAP fallbacks and fix fp16 variant detection (llama/25554) REVERT: 8bb3846 sycl: fix UE4M3 parsing (llama/25608) REVERT: f6515fc sycl: *glu flat path (llama/26354) REVERT: ff7002b sycl : Support DSv4 OPs: LIGHTNING_INDEXER,DSV4_HC_COMB,DSV4_HC_POST,DSV4_HC_PRE (llama/26568) REVERT: e271907 sycl : fix error Error OP FLASH_ATTN_EXT on arc770 (llama/26441) REVERT: 3b4aa5b sycl : enhance OP set_rows to support all missed data types (llama/26515) REVERT: c0011fe cuda: fix warnings for unused variable/function (llama/26688) REVERT: 02e0bca metal : avoid `threadgroup` matrix array instantiation in kernel_lightning_indexer (llama/26646) REVERT: c035d79 ci : onboard AMD ROCm CI with gfx1151 fixes (llama/26544) REVERT: 0b0a78b vulkan: fix submission batching size, add debug tools for diagnosing causes of DeviceLost drivers errors (llama/26371) REVERT: f10de3a mtmd/ggml: add ggml_build_forward_order (llama/26649) REVERT: d0e4951 vulkan backend ops: implemented GATED_LINEAR_ATTN (llama/25601) REVERT: 90951f9 ggml : bump version to 0.18.1 (leejet#1578) REVERT: 46c86bd sync : llama.cpp REVERT: e75fcb4 sycl: parallelize the non-contiguous concat kernel (llama/25852) REVERT: dbe27be Extended SYCL oneDNN SDPA to non-FP16 KV caches (Q4_0–Q8_0 and FP32) (#25874) REVERT: e8e7ea4 ggml: use dynamic allocation for split graph inputs (llama/22789) REVERT: f80f881 opencl: route large q6_K lm_head to the flat GEMV (llama/26427) REVERT: 3e12f44 CUDA: Fix data-races when reusing SMEM in block_reduce (llama/26385) REVERT: 6521490 metal: implement DSv4 Lightning Indexer (llama/25893) REVERT: 5899365 metal : add SILU_BACK (llama/25982) REVERT: 6be1ee1 metal : add F16 support for bin ops (llama/26465) REVERT: 09ebda8 opencl: limit local workgroup size for GLU operation (llama/26383) REVERT: 91778e3 metal: implement DeepSeek V4 hyper-connections (llama/26459) REVERT: af8c565 opencl: bugfix increment ref_count in ggml_backend_opencl_init() (llama/26162) REVERT: 12fbc7a sycl: fix classification of iGPUs (llama/26105) REVERT: b5c81eb ggml-webgpu: add support for f16 repeat (llama/26307) REVERT: a8f4315 vulkan: extend topk_moe fusion to support sqrt(softplus) (llama/26124) REVERT: c7e4593 vulkan: add POOL_1D op (llama/25431) REVERT: f3022a1 vulkan: Introduce driver version check for Windows Intel GPU to mitigate crashing (llama/25192) REVERT: 369bf0a cuda: extract Q2_0 elements via __byte_perm (llama/25603) REVERT: ec1eb4f SYCL: add oneMKL GEMM flash attention for XMX-accelerated prompt proc… (#25025) REVERT: 0635a0f support the missed types in cpy (llama/26005) REVERT: e2a5c53 ggml-zendnn : group matmul direct API for mul_mat_id (llama/25918) REVERT: 97d7638 sycl : support dev2dev memcpy by DEV2DEV_MEMCPY_FORWARD (llama/26234) REVERT: 776124b Support q2 mul_mat (llama/26231) REVERT: 2413b70 sycl: fuse RMS_NORM + MUL (llama/26015) REVERT: 39f1fd3 ggml-webgpu: improve flash_attn_vec for quantized KV at long contexts (llama/25956) REVERT: f04850c vulkan: Support quantized concat (llama/25684) REVERT: 3d68e2b Test support for alternative conv layout (llama/25617) REVERT: 49ed848 ggml-cuda: Allow transpose-free gemmv computation (llama/26171) REVERT: 78de606 sync : whisper.cpp REVERT: 06ca976 ggml : bump version to 0.18.0 (leejet#1576) REVERT: dfeb865 typo: init_model does not create tensors (leejet#1572) REVERT: 4429529 sync : llama.cpp REVERT: b5cd818 CUDA: add Q2_0 support (llama/25707) REVERT: 560511d metal: fix memory unwire if model is freed without any GPU operations (llama/26082) REVERT: 99ac6e0 ggml : Fix issue with kleidiai ci and stringop overflow warning (llama/26277) REVERT: 111adea enhance UT to show all real unsupported backends (llama/25234) REVERT: dbc6bc2 ggml-cuda : disable MMQ on devices with less than 48 KiB shared memory (llama/26141) REVERT: 44fd996 sycl: contiguous fast path + 32-bit index math for unary elementwise ops (llama/25946) REVERT: 0e1747a RPC: add tensor_memset (llama/25912) REVERT: dd628ca add rdna3.5, and 3 to mmq configs so they can be tuned independently. (llama/26199) REVERT: be64c58 ggml-webgpu: Fix some binding alias issues to support all archs, fix recurrent-state-rollback test (llama/25931) REVERT: 98d5e77 opencl: skip the Adreno KQ/KQV image kernels for multi-stream batches (llama/26189) REVERT: d727fbb ggml : set output of view src (llama/25729) REVERT: 4a17c1b vulkan: add iq4_nl support back to FA (llama/24585) REVERT: 260279d ggml-cuda: add chunked SSD matmul for Mamba-2 prefill acceleration (llama/22675) REVERT: 23b34d5 sycl: fix use-after-return of the SDPA scale in the oneDNN flash-attention path (llama/25880) REVERT: 5265d8f ggml-metal: FWHT kernel for metal backend (llama/25924) REVERT: 3fc55d4 Disable -ffast-math on HIP (llama/25495) REVERT: 0929483 sycl(build): parallelize ocloc invocations (llama/25903) REVERT: 27bbeaf ggml : adjust logic for offloading ops to weight's backend (llama/25832) REVERT: bc40151 ggml-cpu: Enable BF16 tiled gemm optimization on PowerPC (llama/26068) REVERT: 1f64a93 opencl: fix fused RMS norm mul view offset (llama/26085) REVERT: 99a582e hexagon: partial im2col support (llama/26007) REVERT: 3ca7e8c Update ggml/src/gguf.cpp : Defined virtual keyword for destructor of gguf_writer_base (llama/25867) REVERT: 0b28b84 HIP: remove rocWMMA FlashAttention (llama/26046) REVERT: 4c0b27d opencl: cache compiled cl_program binaries on disk (llama/26050) REVERT: 6a92a50 opencl: do not treat NULL-mask flash attention as causal (llama/25771) REVERT: c2a30bb hexagon: further improved pipeline of the core bits (L2, DMA, MM, FA) (llama/26049) REVERT: 5c4119e hexagon: fix Windows crash when op_poll is enabled (llama/26029) REVERT: 2fbff1e CUDA: fix external compilation of q1_0 MMQ (llama/25778) REVERT: 17a2a7c metal : add f16 type support to leaky relu (llama/25981) REVERT: e7059df CUDA: Improve NVFP4 W4A4 activation quantization (llama/25730) REVERT: 874aa98 hexagon: activation ops update (llama/25974) REVERT: 596deb3 ggml: enable PowerPC backend variants on AIX (llama/25983) REVERT: a72a524 webgpu : add CONV_2D_DW (depthwise conv2d) kernel (llama/25847) REVERT: 08130cf cuda: GET_ROWS quants (llama/25962) REVERT: e4330b9 hexagon: check tensor type when reusing descriptors (llama/25968) REVERT: 37a89db cuda: add sqrt_softplus in topk-moe for dsv4 (llama/25896) REVERT: c4fa8af kleidiai : warn once when a weight type has no KleidiAI kernel (llama/25701) REVERT: 3ca9985 vulkan: Refactor vk_queue to use per-instance mutexes and unique handles (llama/23570) REVERT: b186c92 ggml-openvino: Add GGML_BACKEND_DL_IMPL invocation for OpenVINO backend (llama/25795) REVERT: 924cd99 CUDA: vectorize same-type get_rows with int4 copy (llama/25929) REVERT: 09e1fe8 hexagon: add CLAMP op (llama/25934) REVERT: 0fd2904 opencl: Support broadcast for Adreno MUL_MAT and honor `view_offs` for Adreno Q8_0 MUL_MAT for llama-server multi-stream (llama/25910) REVERT: a0c4f2a opencl: load and use `kernel_gemm_moe_q6_k_f32_ns` from bin kernel lib (llama/25797) REVERT: 454ea6f opencl: read/write MoE dp4a activation tiles to local memory as 128-bit (vectorized LD/ST perf opt) for Adreno GPUs (llama/25810) REVERT: f68fee5 opencl: transpose q4_K noshuffle scales for coalesced reads (llama/25805) REVERT: 65a9776 tests : initialize all tensors in test_dsv4_hc to avoid NaNs in sentinel tensors (llama/25822) git-subtree-dir: ggml git-subtree-split: 04b55bba4877e0d269e2db68911608485c08eb3a
Summary
This is the successor of #1477. That earlier PR did the same thing (stream model weights from CPU to GPU so larger models fit), but it ran as a parallel system alongside the existing graph-cut planner (#1476) and exposed many user-facing flags. Both points came up in the review.
This PR rebuilds the feature on top of the graph-cut planner instead of running alongside it. There is one new boolean flag,
--stream-layers. When it is off, behavior is byte-identical to upstream master.The change is split into two commits:
--stream-layersflag, aSegmentResidencyannotation pass on the existing planner, and a small layer registry used by the runner. No behavior change when the flag is off.weight_adapteris attached, prefer the runtime LoRA mode when streaming is on because immediate mode OOMs).Usage
Streaming kicks in when both
--max-vram(or its auto sentinel-1) and--stream-layersare set. Weights need a place to stream from, so--offload-to-cpuis implicitly enabled if you forget it (with a log line).--max-vram -1auto-detects free VRAM and reserves 1 GiB headroom. Pass a positive value (e.g.--max-vram 9) to set the budget explicitly.Tested
--max-vram 4 --offload-to-cpu --stream-layers: Z-Image Q8 and bf16, HiDream, Flux schnell, SD3.5 large, SDXL, WAN 2.2 5B. All generated valid output, nocudaMallocfailures.--stream-layers off: byte-identical PNG to upstream walker.--cfg-scale 1.0 --guidance 3.5 --flow-shift 3.0. Both batch images come out correct.Performance
The async-prefetch path that was in earlier iterations of this branch is intentionally disabled in this PR because of correctness regressions it caused with runtime LoRA and with
batch_count > 1combined with non-default--guidance/--flow-shift. The implementation stays in the source for the follow-up to build on, but the engagement gate isfalse.chunk-K residency is still active and saves H2D for the resident segments across all sampling steps. The wallclock benefit varies by model and
--max-vrambudget.Future work
The headline perf win, keeping the GPU near 100 percent utilization during streaming, needs PR #1477's
chunk_graph.hpphelper ported on top of this foundation. That caches a fused cgraph for K base layers so the host does not need to issue per-layer kernel launches between H2D copies. It is the planned next PR, written specifically against this branch.Other items I have queued for follow-ups:
--stream-layers diffusion,llm,vae, once cross-runner residency eviction is in place. The current PR scopes streaming to the diffusion runner only.SDCPP_STREAM_PROFILEenv-var-gated per-stage timing breakdown, useful for tuning.Checklist