feat: prefetch streamed layers during compute

docs/performance.md

@@ -55,21 +55,28 @@ See [backend selection](./backend.md) for full syntax.
 
 ## Run models that don't fit in VRAM (CPU streaming).
 
-`--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:
+`--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. The following 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.
+- `--resident-layers <N|auto>` sets the maximum number of leading parameter-bearing graph-cut segments kept resident (default: `-1`; `auto` is an alias for `-1`). `-1` uses as many as the VRAM budget permits, `0` keeps none, and a positive `N` keeps up to `N`.
+- `--layer-prefetch-depth <N>` sets the maximum number of future parameter-bearing graph-cut segments copied through a separate transfer backend or queue when supported while the active segment computes (default: `0`). `0` disables asynchronous prefetching, `1` overlaps the next segment, and larger values provide deeper lookahead when the VRAM budget permits.
 
-The three flags stack. The recommended shape for "biggest model my card can host":
+Both controls require `--stream-layers` and a non-zero `--max-vram`. Prefetch is budgeted before residency; requested limits are reduced as needed, a positive VRAM cap is never exceeded to force either one, and the active segment is never evicted. Residents persist only across repeated sampling steps, and stale graph state is released automatically. An explicit non-negative residency limit uses an unmerged plan, which may add dispatch overhead.
+
+These flags stack. The recommended shape for "biggest model my card can host":
 
 ```shell
 sd-cli --diffusion-model flux1-dev.safetensors ... \
-       --offload-to-cpu --max-vram -1 --stream-layers
+       --offload-to-cpu --max-vram -1 --stream-layers \
+       --resident-layers auto --layer-prefetch-depth 1
 ```
 
 - `--offload-to-cpu`: params in RAM, staged as needed.
 - `--max-vram -1`: use most of the free VRAM as the compute budget, spare 1 GiB headroom, let the graph-cut segmenter split each forward pass to fit.
 - `--stream-layers`: on top of the segmenter, stream individual transformer blocks so their weights don't all need to be resident at once.
+- `--resident-layers auto`: use the remaining budget for a leading resident prefix.
+- `--layer-prefetch-depth 1`: prepare the next parameter-bearing segment concurrently with the current segment's computation.
 
 Ordered from fastest to smallest-VRAM: no flags → `--offload-to-cpu` → `--offload-to-cpu --max-vram <N>` → `--offload-to-cpu --max-vram <N> --stream-layers`. Each step down costs a few percent of throughput to buy more room; combined they can run models roughly 3-4x larger than the raw VRAM would allow.
 

examples/cli/main.cpp

@@ -894,7 +894,8 @@ int main(int argc, const char* argv[]) {
         }
     }
 
-    sd_ctx_params_t sd_ctx_params = ctx_params.to_sd_ctx_params_t(cli_params.taesd_preview);
+    sd_ctx_params_t sd_ctx_params                = ctx_params.to_sd_ctx_params_t(cli_params.taesd_preview);
+    sd_layer_stream_params_t layer_stream_params = ctx_params.to_sd_layer_stream_params_t();
 
     SDImageVec results;
     int num_results             = 0;
@@ -904,7 +905,7 @@ int main(int argc, const char* argv[]) {
         num_results = 1;
         results.push_back(gen_params.init_image.release());
     } else {
-        SDCtxPtr sd_ctx(new_sd_ctx(&sd_ctx_params));
+        SDCtxPtr sd_ctx(new_sd_ctx_with_layer_stream(&sd_ctx_params, &layer_stream_params));
 
         if (sd_ctx == nullptr) {
             LOG_INFO("new_sd_ctx_t failed");

examples/common/common.cpp

@@ -505,6 +505,11 @@ ArgOptions SDContextParams::get_options() {
          "maximum VRAM budget in GiB for graph-cut segmented execution. Accepts a single value or assignments by backend/device, e.g. 6 or cuda0=6,vulkan0=4. 0 disables graph splitting; a negative value auto-detects free VRAM, sparing the specified value",
          0,
          &max_vram},
+        {"",
+         "--resident-layers",
+         "maximum leading parameter-bearing graph-cut segments kept resident with --stream-layers: -1 selects automatically from the VRAM budget, auto is an alias for -1, 0 keeps none, N keeps up to N (default: -1)",
+         0,
+         &resident_layers_spec},
     };
 
     options.int_options = {
@@ -513,6 +518,10 @@ ArgOptions SDContextParams::get_options() {
          "number of threads to use during computation (default: -1). "
          "If threads <= 0, then threads will be set to the number of CPU physical cores",
          &n_threads},
+        {"",
+         "--layer-prefetch-depth",
+         "number of future parameter-bearing graph-cut segments to prefetch with --stream-layers (default: 0; 0 disables prefetching, 1 overlaps the next segment, 2+ enables deeper lookahead when VRAM permits)",
+         &layer_prefetch_depth},
     };
 
     options.bool_options = {
@@ -721,6 +730,25 @@ bool SDContextParams::resolve(SDMode mode) {
         n_threads = sd_get_num_physical_cores();
     }
 
+    std::string resident_spec = resident_layers_spec;
+    std::transform(resident_spec.begin(), resident_spec.end(), resident_spec.begin(), [](unsigned char c) {
+        return static_cast<char>(std::tolower(c));
+    });
+    if (resident_spec == "auto") {
+        resident_layers = -1;
+    } else {
+        try {
+            size_t parsed   = 0;
+            resident_layers = std::stoi(resident_layers_spec, &parsed);
+            if (parsed != resident_layers_spec.size()) {
+                LOG_ERROR("error: --resident-layers must be auto, -1, 0, or a positive integer");
+                return false;
+            }
+        } catch (const std::exception&) {
+            LOG_ERROR("error: --resident-layers must be auto, -1, 0, or a positive integer");
+            return false;
+        }
+    }
     build_embedding_map();
 
     return true;
@@ -754,6 +782,14 @@ bool SDContextParams::validate(SDMode mode) {
         LOG_ERROR("error: vae_format must be 'auto', 'flux', 'sd3', 'flux2', or 'wan'");
         return false;
     }
+    if (resident_layers < -1) {
+        LOG_ERROR("error: --resident-layers must be auto, -1, 0, or a positive integer");
+        return false;
+    }
+    if (layer_prefetch_depth < 0) {
+        LOG_ERROR("error: --layer-prefetch-depth must be >= 0");
+        return false;
+    }
 
     return true;
 }
@@ -832,6 +868,8 @@ std::string SDContextParams::to_string() const {
         << "  offload_params_to_cpu: " << (offload_params_to_cpu ? "true" : "false") << ",\n"
         << "  max_vram: \"" << max_vram << "\",\n"
         << "  stream_layers: " << (stream_layers ? "true" : "false") << ",\n"
+        << "  resident_layers: " << resident_layers << ",\n"
+        << "  layer_prefetch_depth: " << layer_prefetch_depth << ",\n"
         << "  eager_load: " << (eager_load ? "true" : "false") << ",\n"
         << "  backend: \"" << backend << "\",\n"
         << "  params_backend: \"" << params_backend << "\",\n"
@@ -914,6 +952,14 @@ sd_ctx_params_t SDContextParams::to_sd_ctx_params_t(bool taesd_preview) {
     return sd_ctx_params;
 }
 
+sd_layer_stream_params_t SDContextParams::to_sd_layer_stream_params_t() const {
+    sd_layer_stream_params_t params;
+    sd_layer_stream_params_init(&params);
+    params.resident_layers      = resident_layers;
+    params.layer_prefetch_depth = layer_prefetch_depth;
+    return params;
+}
+
 SDGenerationParams::SDGenerationParams() {
     sd_sample_params_init(&sample_params);
     sd_sample_params_init(&high_noise_sample_params);

examples/common/common.h

@@ -146,12 +146,15 @@ struct SDContextParams {
     std::map<std::string, std::string> embedding_map;
     std::vector<sd_embedding_t> embedding_vec;
 
-    rng_type_t rng_type         = CUDA_RNG;
-    rng_type_t sampler_rng_type = RNG_TYPE_COUNT;
-    bool offload_params_to_cpu  = false;
-    std::string max_vram        = "0";
-    bool stream_layers          = false;
-    bool eager_load             = false;
+    rng_type_t rng_type              = CUDA_RNG;
+    rng_type_t sampler_rng_type      = RNG_TYPE_COUNT;
+    bool offload_params_to_cpu       = false;
+    std::string max_vram             = "0";
+    bool stream_layers               = false;
+    std::string resident_layers_spec = "-1";
+    int resident_layers              = -1;
+    int layer_prefetch_depth         = 0;
+    bool eager_load                  = false;
     std::string backend;
     std::string params_backend;
     std::string split_mode;
@@ -183,6 +186,7 @@ struct SDContextParams {
     bool resolve_and_validate(SDMode mode);
     std::string to_string() const;
     sd_ctx_params_t to_sd_ctx_params_t(bool taesd_preview);
+    sd_layer_stream_params_t to_sd_layer_stream_params_t() const;
 };
 
 struct SDGenerationParams {

examples/server/main.cpp

@@ -85,8 +85,9 @@ int main(int argc, const char** argv) {
     LOG_DEBUG("%s", ctx_params.to_string().c_str());
     LOG_DEBUG("%s", default_gen_params.to_string().c_str());
 
-    sd_ctx_params_t sd_ctx_params = ctx_params.to_sd_ctx_params_t(false);
-    SDCtxPtr sd_ctx(new_sd_ctx(&sd_ctx_params));
+    sd_ctx_params_t sd_ctx_params                = ctx_params.to_sd_ctx_params_t(false);
+    sd_layer_stream_params_t layer_stream_params = ctx_params.to_sd_layer_stream_params_t();
+    SDCtxPtr sd_ctx(new_sd_ctx_with_layer_stream(&sd_ctx_params, &layer_stream_params));
 
     if (sd_ctx == nullptr) {
         LOG_ERROR("new_sd_ctx_t failed");

include/stable-diffusion.h

@@ -237,6 +237,12 @@ typedef struct {
     const char* model_args;
 } sd_ctx_params_t;
 
+typedef struct {
+    uint32_t struct_size;      // Set by sd_layer_stream_params_init; permits future extension
+    int resident_layers;       // With stream_layers: maximum leading graph-cut segments kept resident (-1 = automatic, 0 = none)
+    int layer_prefetch_depth;  // With stream_layers: future parameter-bearing graph-cut segments prefetched during compute (0 = disabled)
+} sd_layer_stream_params_t;
+
 typedef struct {
     uint32_t sample_rate;
     uint32_t channels;
@@ -477,8 +483,11 @@ SD_API void sd_hires_params_init(sd_hires_params_t* hires_params);
 
 SD_API void sd_ctx_params_init(sd_ctx_params_t* sd_ctx_params);
 SD_API char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params);
+SD_API void sd_layer_stream_params_init(sd_layer_stream_params_t* params);
 
 SD_API sd_ctx_t* new_sd_ctx(const sd_ctx_params_t* sd_ctx_params);
+SD_API sd_ctx_t* new_sd_ctx_with_layer_stream(const sd_ctx_params_t* sd_ctx_params,
+                                              const sd_layer_stream_params_t* layer_stream_params);
 SD_API void free_sd_ctx(sd_ctx_t* sd_ctx);
 SD_API void free_sd_audio(sd_audio_t* audio);