feat: prefetch streamed layers during compute (#1905)

Co-authored-by: leejet <leejet714@gmail.com>
This commit is contained in:
assouanandleejet authored and GitHub committed 2026-09-06 16:35:30 +08:00
1 parent 6b3edaaf32
commit 6c57cc3b38
15 files changed
+562 -11

No files matched your search

+2 -1
View File
@@ -58,7 +58,8 @@ See [backend selection](./backend.md) for full syntax.
`--offload-to-cpu` alone keeps every parameter in system RAM and stages it to the runtime backend on first use, then leaves it resident there. If the diffusion model is larger than the runtime backend's free memory (e.g. Flux dev at bf16 on an 8 GiB GPU), that residency stops fitting during the sampling loop and generation fails. Two additional flags make it fit by trading a small amount of speed for room:
- `--max-vram <GiB>` sets a VRAM budget the graph-cut segmenter respects. It cuts each forward pass into segments sized to fit the budget, running them in sequence and freeing intermediate activations between them. Negative values auto-detect free VRAM and spare the given amount (`--max-vram -1` uses most of the free VRAM and keeps ~1 GiB headroom), a positive value caps the budget, `0` disables segmentation.
- `--stream-layers` streams the diffusion model's transformer blocks one at a time. Each block's parameters are copied from the CPU to the runtime backend just before it runs and evicted when the residency budget is reached. Prefetching hides most of the copy latency behind compute. This flag only takes effect when the diffusion params backend is CPU, so it must be combined with `--offload-to-cpu` (or an explicit `--params-backend diffusion=cpu`); a warning is logged and the flag is ignored otherwise.
- `--stream-layers` streams the diffusion model's transformer blocks one at a time. While one block computes, the next block's parameters are prefetched automatically from the CPU on a separate transfer queue; parameters are evicted when the residency budget is reached. This flag only takes effect when the diffusion params backend is CPU, so it must be combined with `--offload-to-cpu` (or an explicit `--params-backend diffusion=cpu`); a warning is logged and the flag is ignored otherwise.
- `--disable-prefetch` disables the asynchronous next-block prefetch while retaining synchronous `--stream-layers` execution. This is mainly useful for debugging or backends where transfer and compute do not overlap effectively.
The three flags stack. The recommended shape for "biggest model my card can host":