add sd-fit-params: fit placement to free device memory using measured dry runs

[gianni-cor]

Adds sd-fit-params, a placement helper that measures the requested generation workload in a metadata-only dry run and emits the CLI arguments needed to fit it into available device memory.

What changed

• Adds sd_fit_params(ctx_params, workload, result) plus sd_fit_workload_init / sd_fit_result_free.
• Adds the sd-fit-params example tool, which prints reusable --backend, --params-backend, --vae-tiling, --max-vram, and --stream-layers arguments.
• Adds metadata-only measurement mode for GGMLRunner::compute(): graphs are built and compute buffers are measured without reading tensor data or allocating the real buffers.
• Measures the real image/video generation path, including image-conditioned graphs and a second VAE-tiling pass when applicable.
• Adds planner tests covering default fit, resident multi-device placement, time-share CPU fallback, VAE tiling, and diffusion CPU params plus layer streaming.
• Documents the planner fallback order in examples/fit-params/README.md.
• Documents library API usage, including scalar workload fields versus full sd_img_gen_params_t / sd_vid_gen_params_t requests, result ownership, status handling, and applying stream_layers with the original max_vram.

Planner order

The planner tries faster, more resident placements first, then falls back to lower-VRAM options:

1. Keep the default first-GPU placement if all params plus the largest compute buffer fit.
2. Try resident multi-device placement, sorted by module parameter size.
3. Retry resident placement with measured VAE tiled compute when full-resolution VAE compute does not fit.
4. Try time-share GPU placement with module params loaded per phase from disk.
5. Retry time-share placement with measured VAE tiled compute.
6. Split splittable modules across multiple GPUs.
7. For splittable diffusion modules, keep diffusion params in CPU RAM and emit --stream-layers when a nonzero --max-vram graph budget is available.
8. Fall back to CPU runtime for modules that still do not fit on GPU.

Related fixes

• Removes the obsolete Wan video spatial-tiling bypass so video VAE tiling can be measured and used.
• Silences the expected ggml GGUF-reader probe errors when the GGUFReader fallback succeeds.
• Adds SD_FIT_DEBUG_DEVICES to exercise planner paths against simulated GPU capacities.

Validation

• Local build: cmake --build build --target sd-fit-params test-fit-params -j 8.
• Local tests: ctest --test-dir build --output-on-failure.
• CUDA 13.3 build and full test run on 2x RTX 5090.
• Multi-GPU validation on 2x RTX 5090 for resident-spread and layer-split plans.
• LTX-2.3 22B + Gemma 3 12B test on 2x RTX 5090: sd-fit-params --max-vram 8 recommends --params-backend diffusion=cpu --stream-layers; the emitted args completed a one-step sd-cli video generation.
• Existing --auto-fit heuristic behavior is unchanged.

[iancris]

Measure mode destroys the graph context before returning, so callers that read tensors out of it afterwards dereference freed memory.

The measure branch calls free_compute_ctx() here; the ordinary path returns through execute_graph() and leaves compute_ctx alive until the next reset_compute_ctx(). Two callers depend on that lifetime, and both pass no_return = true precisely because they collect their own outputs from the graph:

• src/model/diffusion/control.hpp:426 then iterates control_outputs_ggml at :433-437, calling sd::make_sd_tensor_from_ggml<float>(control) on tensors build_graph allocated in compute_ctx.
• src/model/adapter/lora.hpp:1180 then calls ggml_backend_tensor_copy(final_tensor, original_tensor) at :1182-1186, where final_tensor came from build_lora_graph.

Both are reachable from a fit run: sd_fit_params drives the real generation pipeline, and LoRAs and a full sd_img_gen_params_t are advertised workload inputs.

Impact: a use-after-free inside a dry run documented as metadata-only — a crash, a GGML_ABORT on the type check, or a memcpy of arbitrary length.

Suggested fix: remove the free_compute_ctx() call so the measure branch matches the non-measure path and the context survives until the normal lifecycle frees it. That closes both dangling-read sites at once.

Worth deciding separately: even with the context alive, these tensors are never allocated in measure mode, so make_sd_tensor_from_ggml would read a null data. no_return = true means "the caller harvests the graph itself" — something measure mode cannot currently satisfy.

[iancris]

ggml_graph_node(gf, -1) is not the graph's result whenever cached or debug tensors were expanded after it.

get_compute_graph (src/core/ggml_extend.hpp:2043-2060) names the result first, then appends more nodes:

auto result = ggml_graph_node(gf, -1);
ggml_set_name(result, final_result_name.c_str());
for (const auto& entry : debug_tensors)    { ggml_build_forward_expand(gf, entry.first); }
for (const auto& entry : cache_tensor_map) { ggml_build_forward_expand(gf, entry.second); }

The real path accounts for this and looks the result up by name — ggml_get_tensor(compute_ctx, final_result_name.c_str()) at :3051. This branch indexes the last node instead, so once cache_tensor_map is non-empty it returns a cache entry: GGMLRunner::cache() wraps a view in a fresh ggml_cont (:3279-3283), which is a node and lands after the result.

Both cache() call sites are on paths this PR targets — src/model/vae/wan_vae.hpp:1372 (feat_idx: temporal state) and src/model/diffusion/control.hpp:397 (guided_hint).

Impact: for WanVAE build_graph_partial, measure mode returns a feature-cache entry instead of the VAE output. Tiled, that trips the shape assert at :925; untiled, it silently feeds a wrong-shaped latent into the next measured graph, so the reported sizes describe graphs that never get built.

Suggested fix: resolve the output by name, matching execute_graph, so nodes appended after the naming cannot be mistaken for the result:

ggml_tensor* out = ggml_get_tensor(compute_ctx, final_result_name.c_str());
if (out == nullptr) { free_compute_ctx(); return std::nullopt; }

[iancris]

Measure mode hardcodes four dimensions, so every rank-aware wrapper downstream sees a different shape than it would in a real run.

sd::zeros<T>({out->ne[0], out->ne[1], out->ne[2], out->ne[3]}) always produces rank 4. The real path builds its tensor through sd::shape_from_ggml (src/core/tensor_ggml.hpp:39-46), which pushes only ggml_n_dims(tensor) entries. So a [768, 77, 1, 1] text-encoder output is rank 2 in a real run and rank 4 here.

That difference is load-bearing, because the wrappers branch on rank: restore_trailing_singleton_dims(..., x.dim()) at src/model/diffusion/unet.hpp:838 and its equivalents in wan.hpp and flux.hpp, MiniMaxH3VAE::ensure_video_shape, and the sd::ops::* shape assertions.

Impact: a wrong-rank tensor either asserts outright or propagates a wrong shape into the next measured graph — so the compute sizes this run reports belong to graphs that would never be built for real.

Suggested fix: reuse the same shape helper the real path uses, so measure mode and read_graph_tensor agree on rank by construction rather than by coincidence:

result = sd::zeros<T>(sd::shape_from_ggml(out));

[iancris]

The planner assumes only one module's compute buffer is live at a time, which is false for any workload carrying a ControlNet.

Both the check-first tier here and the resident tier (std::max<int64_t>(max_compute[di], compute) at :231 and :261) budget one compute buffer per device. examples/fit-params/README.md:66-68 states that assumption outright, and it holds for TE → diffusion → VAE, where each runner releases its buffer before the next runs.

It does not hold for ControlNet. ControlNet::compute passes free_compute_buffer = false (src/model/diffusion/control.hpp:426), and so does UNetModel (src/model/diffusion/unet.hpp:838) — the fourth positional argument of GGMLRunner::compute (src/core/ggml_extend.hpp:3293-3298). ControlNet is invoked from inside the denoise callback (src/stable-diffusion.cpp:2584) while the diffusion buffer is still resident, and neither is released until the sampling loop returns (:3019-3024).

Impact: for a ControlNet workload the device is under-budgeted by the entire ControlNet compute buffer, so the emitted plan is projected to fit and then OOMs at real generation time — the silent-under-budget outcome this tool exists to prevent.

Suggested fix: partition modules by whether their runners are invoked with free_compute_buffer == false inside one phase, and sum the compute buffers within a partition rather than taking the maximum across it. Treating CONTROL_NET and DIFFUSION as concurrent is the minimum that closes this case.

[iancris]

The final CPU fallback performs no capacity check at all, so a model far too large for the machine still reports success.

This last tier sets placed and on_cpu unconditionally. Host RAM is never enumerated and never compared against params_bytes. plan_placement then sets plan->valid = true on every path (:458) and returns true, so the guard at src/stable-diffusion.cpp:4218 — if (!planned || !plan.valid) return SD_FIT_FAILURE; — can only fire when canonicalize_backend_keys fails. The only other SD_FIT_FAILURE is the user_set_placement branch at :4225-4227.

So the status documented as "could not find a placement projected to fit" is unreachable for the case it names.

Impact: a 70 GB diffusion model on a 4 GB machine returns SD_FIT_SUCCESS with --backend diffusion=cpu, and the real run then exhausts host memory. examples/fit-params/README.md:161-163 and docs/backend.md both promise otherwise, so a caller branching on the status gets a false green — worse than no status at all, since the point of the API is to be trusted before committing to a run.

Suggested fix: enumerate available host RAM and compare it against the params still unplaced when this tier is reached, returning SD_FIT_FAILURE when they do not fit — that makes the documented status reachable and the success signal meaningful. If measuring host RAM is unwanted, the alternative is to surface a cpu_fallback_unverified flag on the result and weaken the documented contract to match what the code actually checks.

[iancris]

fit_module_ defaults to SDBackendModule::TE, so any unregistered runner's measurement is silently booked against the text encoder.

set_fit_module is called from exactly two places: src/stable-diffusion.cpp:349, inside register_runner_params, and src/upscaler.cpp:99. Anything that is a GGMLRunner but does not pass through register_runner_params therefore keeps this default and reports as TE.

LoraModel is exactly that case — it derives from GGMLRunner (src/model/adapter/lora.hpp:24) and is never registered — so its measurement lands under TE at src/core/ggml_extend.hpp:2341. That measurement is a graph over the full merged weight set, potentially the largest compute buffer in the run.

The aggregation at src/stable-diffusion.cpp:4173 then takes std::max(m.compute_bytes, record.compute_bytes), so the text encoder's compute absorbs an unrelated graph's peak.

Impact: the planner sizes and places te to accommodate memory belonging to a module it does not model, and the LoRA graph's real requirement is attributed to a module that may end up on a different device entirely. Placements are then derived from numbers that describe no real module.

Suggested fix: give the enum an explicit unset value and default to it, so an unattributed measurement is distinguishable from a genuine text-encoder one instead of silently impersonating it. Then either call set_fit_module on every runner, or drop records whose module is unset during aggregation.

[iancris]

The video-capability check is consulted only when no complete request is set, which is the opposite of the path the CLI and the README both recommend.

sd_version_supports_video_generation() sits behind workload->image_gen_params == nullptr, so it never runs once a full request is supplied.

Meanwhile examples/fit-params/main.cpp:231 picks the mode purely from gen_params.video_frames > 1, video_frames defaults to 1 (examples/common/common.h:238), and sd-fit-params exposes no -M/--mode flag of its own the way sd-cli does. So sd-fit-params -m wan2.2.gguf -W 832 -H 480 sets workload.image_gen_params (main.cpp:257-258), and sd_fit_params runs generate_image against a video-only model.

Impact: fitting a Wan / LTX / HunyuanVideo / MiniMax-H3 model without remembering --video-frames N either hard-fails with SD_FIT_ERROR and "measurement dry run failed", or measures a pipeline that will never run. The README's own advice — use the full request form for the most accurate plan — steers users straight into it.

Suggested fix: apply the version check regardless of which request pointer is set, so model capability decides the pipeline rather than which input form the caller happened to use. Either reject an image_gen_params workload on a video-only model with SD_FIT_ERROR and a clear message, or promote it to the video path. Deriving the mode in main.cpp from the model's declared capability rather than from video_frames would close the CLI half of it too.

[iancris]

The headroom reserve exists only on the zero branch, so supplying an explicit budget silently removes it.

gib == 0 reserves 512 MiB (:54), but the gib > 0 branch clamps to free_bytes and the gib < 0 branch subtracts only what the user asked for. So --max-vram 16 on a card reporting 15.5 GiB free yields a budget of exactly free_bytes, and the placement tiers pack right up to it — nothing left for driver context, allocator fragmentation, or the cache buffers this measurement does not model.

The PR's own test_resident_spread (tests/test-fit-params.cpp:66-76) lands on exactly 8 GiB of an 8 GiB budget, so the zero-headroom case is currently encoded as the expected result.

Impact: any positive --max-vram at or above free memory produces a plan with no margin, which OOMs on the real run for reasons the projection cannot see.

Suggested fix: apply the reserve on every branch rather than only the default one, so the headroom is a property of the budget rather than of which branch computed it:

const int64_t reserve = 512 * MiB;
if (gib > 0.f) {
    d.budget_bytes = std::min<int64_t>((int64_t)(gib * 1024.0 * 1024.0 * 1024.0), d.free_bytes - reserve);

Separate judgement call, flagged rather than asserted: at runtime --max-vram is a graph-splitting budget (src/core/ggml_graph_cut.cpp:230-241), not a whole-device cap, so reusing the same number as a device budget conflates two different quantities.

[iancris]

Measure mode returns before the cache buffer is ever allocated, so stateful models are measured in a form they never actually run in.

The real path allocates a separate backend buffer for cached tensors — cache_buffer = ggml_backend_alloc_ctx_tensors(cache_ctx, runtime_backend) at :2439, driven from copy_cache_tensors_to_cache_buffer at :3048. This branch returns before execute_graph, with two consequences:

1. cache_buffer bytes never reach any graph_memory_measurement.
2. cache_ctx stays null, so get_cache_tensor_by_name returns null and every graph is built in its cold form.

Combined with gen.sample_params.sample_steps = 1 (src/stable-diffusion.cpp:4103, :4126), only the first and smallest iteration is ever observed.

Impact: Wan VAE feature caches and the ControlNet guided-hint cache are unmeasured VRAM, so for any cached model the per-module numbers sit strictly below the steady-state requirement. The plan is projected to fit and then does not.

Suggested fix: add the projected cache-tensor bytes to graph_memory_measurement — they are enumerable from cache_tensor_map at measure time, without allocating anything, which keeps the dry run metadata-only. Measuring a second warm-cache iteration would additionally capture the steady-state graph shape rather than only the cold one.

[aegioscy]

Could we structure this loader like SdGenHandlers? Extract the shared image/audio loaders, define one named handler per input, and let load_generation_inputs iterate them. This should be an ordered container rather than an unordered_map, because image dimensions are resolved by earlier loads and are needed by later inputs. It would keep this function as a small orchestration loop while making each input path easier to extend independently.

[aegioscy]

For example, the shared loaders can move out of load_generation_inputs:

static bool load_image_file(SDGenerationParams& params,
                            const std::string& path,
                            SDImageOwner& image,
                            bool resize = true,
                            int channels = 3) {
    if (path.empty()) return true;

    const bool use_size = resize && params.width_and_height_are_set();
    if (!load_sd_image_from_file(image.put(), path.c_str(),
                                 use_size ? params.width : 0,
                                 use_size ? params.height : 0,
                                 channels)) {
        fprintf(stderr, "failed to load image from '%s'\\n", path.c_str());
        return false;
    }
    params.set_width_and_height_if_unset(image.get().width, image.get().height);
    return true;
}

static bool load_audio_file(const std::string& path, SDAudioOwner& audio) {
    std::vector<float> samples;
    uint32_t sample_rate = 0;
    uint32_t channels = 0;
    if (!load_wav_from_file(path, samples, sample_rate, channels)) {
        fprintf(stderr, "failed to load WAV audio from '%s'\\n", path.c_str());
        return false;
    }
    audio.reset(std::move(samples), sample_rate, channels);
    return true;
}

Then the input sequence can be expressed like the addon handlers, while preserving the required ordering:

using LoadInput = bool (*)(SDGenerationParams&, SDMode, bool);
struct LoadInputHandler {
    const char* name;
    LoadInput load;
};

static const LoadInputHandler LOAD_INPUT_HANDLERS[] = {
    {"init_image", [](auto& p, auto, auto) {
        return load_image_file(p, p.init_image_path, p.init_image);
    }},
    {"end_image", [](auto& p, auto, auto) {
        return load_image_file(p, p.end_image_path, p.end_image);
    }},
    {"ref_images", [](auto& p, auto, auto) {
        p.ref_images.clear();
        for (const auto& path : p.ref_image_paths) {
            SDImageOwner image({0, 0, 3, nullptr});
            if (!load_image_file(p, path, image, false)) return false;
            p.ref_images.push_back(std::move(image));
        }
        return true;
    }},
    {"validate", [](auto& p, auto mode, auto) {
        return p.validate(mode);
    }},
    {"mask_image", [](auto& p, auto, auto) {
        return load_image_file(p, p.mask_image_path, p.mask_image, true, 1);
    }},
    // ... ref_videos, audio, control_image, control_video, etc.
};

static bool load_generation_inputs(SDGenerationParams& params,
                                   SDMode mode,
                                   bool verbose) {
    for (const auto& handler : LOAD_INPUT_HANDLERS) {
        if (!handler.load(params, mode, verbose)) return false;
    }
    return true;
}

If generic lambdas are undesirable here, the handlers can use explicit parameter types. The important part is keeping this as an ordered array/vector rather than a map, since image dimensions and validation create dependencies between handlers.