feat: pool VRAM for streamed layer

docs/performance.md

@@ -55,21 +55,31 @@ 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.
+- `--stream-layer-pool` keeps one contiguous backend allocation alive for the sampling loop instead of repeatedly allocating buffers for segment-private streamed parameters (disabled by default; its presence enables it). The allocation is divided into fixed-capacity slots, each sized for the largest private parameter set of any streamed segment. This trades additional reserved VRAM for fewer backend allocations and frees during inference.
 
-The three flags stack. The recommended shape for "biggest model my card can host":
+These 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 an optional allocation, 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. With `--resident-layers 0 --layer-prefetch-depth 1 --stream-layer-pool`, the active and next segments alternate between two reusable slots when the budget permits.
+
+The pool always attempts to allocate the one slot required to execute the largest segment, even when that slot exceeds the effective `--max-vram` budget; the budget limits additional slots, prefetch, and residency. If the requested contiguous allocation fails, the slot count is reduced without changing the slot size. Failure to allocate even one slot stops inference, while failure of asynchronous lookahead only disables prefetch: current segments continue to load synchronously through the same pool, which remains allocated until sampling ends.
+
+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,11 @@ 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();
+    const uint32_t layer_stream_options          = ctx_params.stream_layer_pool
+                                                       ? SD_LAYER_STREAM_OPTION_POOL
+                                                       : SD_LAYER_STREAM_OPTION_NONE;
 
     SDImageVec results;
     int num_results             = 0;
@@ -904,7 +908,9 @@ 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_options(&sd_ctx_params,
+                                                             &layer_stream_params,
+                                                             layer_stream_options));
 
         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,13 +518,21 @@ 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 = {
         {"",
          "--stream-layers",
          "enable residency+prefetch streaming on top of --max-vram (no effect without --max-vram; defaults to false)",
          true, &stream_layers},
+        {"",
+         "--stream-layer-pool",
+         "reuse one contiguous VRAM pool divided into slots sized to the largest streamed segment (defaults to false)",
+         true, &stream_layer_pool},
         {"",
          "--eager-load",
          "load all params into the params backend at model-load time instead of lazily on first use (defaults to false)",
@@ -721,6 +734,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 +786,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 +872,9 @@ 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"
+        << "  stream_layer_pool: " << (stream_layer_pool ? "true" : "false") << ",\n"
         << "  eager_load: " << (eager_load ? "true" : "false") << ",\n"
         << "  backend: \"" << backend << "\",\n"
         << "  params_backend: \"" << params_backend << "\",\n"
@@ -914,6 +957,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,16 @@ 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 stream_layer_pool           = false;
+    bool eager_load                  = false;
     std::string backend;
     std::string params_backend;
     std::string split_mode;
@@ -183,6 +187,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,14 @@ 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();
+    const uint32_t layer_stream_options          = ctx_params.stream_layer_pool
+                                                       ? SD_LAYER_STREAM_OPTION_POOL
+                                                       : SD_LAYER_STREAM_OPTION_NONE;
+    SDCtxPtr sd_ctx(new_sd_ctx_with_layer_stream_options(&sd_ctx_params,
+                                                         &layer_stream_params,
+                                                         layer_stream_options));
 
     if (sd_ctx == nullptr) {
         LOG_ERROR("new_sd_ctx_t failed");

include/stable-diffusion.h

@@ -237,6 +237,17 @@ 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;
+
+enum sd_layer_stream_option_t {
+    SD_LAYER_STREAM_OPTION_NONE = 0,
+    SD_LAYER_STREAM_OPTION_POOL = 1u << 0,  // Reuse a fixed-slot VRAM pool for streamed parameters
+};
+
 typedef struct {
     uint32_t sample_rate;
     uint32_t channels;
@@ -477,8 +488,15 @@ 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 sd_ctx_t* new_sd_ctx_with_layer_stream_options(
+    const sd_ctx_params_t* sd_ctx_params,
+    const sd_layer_stream_params_t* layer_stream_params,
+    uint32_t layer_stream_options);
 SD_API void free_sd_ctx(sd_ctx_t* sd_ctx);
 SD_API void free_sd_audio(sd_audio_t* audio);