auto_placement.h

Go to the documentation of this file.

// SPDX-License-Identifier: LGPL-3.0-or-later
#pragma once
 
#include "gguf_metadata.h"
#include "vram_footprint.h"
 
#include <cstdint>
 
namespace entropic {
 
struct AutoPlacement {
    bool known = false;       
 
    int gpu_layers = 0;
    int cpu_moe_layers = 0;   
    uint64_t bytes = 0;       
    bool fully_resident = false;  
 
    int n_ubatch = 0;
 
    uint64_t margin_bytes = 0;
    const char* reason = "";  
};
 
constexpr uint64_t kComputeBytesAtUbatch512 = 1222ull * 1024ull * 1024ull;
 
constexpr int kComputeReferenceUbatch = 512;
 
constexpr int kDefaultUbatch = 512;
 
constexpr int kUbatchLadder[] = {512, 256, 128};
 
inline uint64_t compute_allowance_bytes(const FootprintInputs& base,
                                        int n_ubatch) {
    const int ubatch = n_ubatch > 0 ? n_ubatch : kDefaultUbatch;
    const uint64_t contexts = base.draft_bytes > 0 ? 2ull : 1ull;
    const uint64_t scaled =
        kComputeBytesAtUbatch512 * static_cast<uint64_t>(ubatch)
        / static_cast<uint64_t>(kComputeReferenceUbatch) * contexts;
    // Only the EXCESS over `vram_reserve_mb`. That reserve is already inside
    // the estimate and vram_footprint.h names covering compute buffers as
    // its whole purpose, so holding back the full figure on top of it would
    // charge the same bytes twice — the exact defect this change exists to
    // remove, reintroduced one layer up. At the default 512 MB reserve and a
    // 128-token ubatch the two nearly cancel, which is why the 26B-A4B at
    // IQ2 still comes out fully resident, as measured.
    const uint64_t reserve =
        static_cast<uint64_t>(base.vram_reserve_mb) * 1024ull * 1024ull;
    return scaled > reserve ? scaled - reserve : 0ull;
}
 
inline bool placement_fits(const FootprintInputs& base,
                           const GgufShape& shape,
                           int gpu_layers,
                           int cpu_moe_layers,
                           uint64_t budget_bytes,
                           uint64_t* used) {
    FootprintInputs in = base;
    // A candidate naming every layer IS full residency; hand the estimator
    // the sentinel so it takes the full branch and counts the embeddings
    // and output head, which the partial branch deliberately excludes.
    in.gpu_layers =
        (gpu_layers < 0 || gpu_layers >= shape.block_count) ? -1 : gpu_layers;
    in.block_count = shape.block_count;
    in.block_bytes = shape.block_bytes;
    in.non_block_bytes = shape.non_block_bytes;
    in.host_expert_bytes =
        shape.expert_bytes_per_block() * static_cast<uint64_t>(cpu_moe_layers);
    const FootprintEstimate est = estimate_vram_footprint(in);
    if (used != nullptr) { *used = est.bytes; }
    return est.known && est.bytes <= budget_bytes;
}
 
inline bool try_fully_resident(const GgufShape& shape,
                               const FootprintInputs& base,
                               uint64_t budget, AutoPlacement* out) {
    uint64_t used = 0;
    if (!placement_fits(base, shape, shape.block_count, 0, budget, &used)) {
        return false;
    }
    out->known = true;
    out->gpu_layers = -1;  // the sentinel, not a count — see AutoPlacement.
    out->fully_resident = true;
    out->bytes = used;
    out->reason = "fully resident";
    return true;
}
 
inline bool try_expert_offload(const GgufShape& shape,
                               const FootprintInputs& base,
                               uint64_t budget, AutoPlacement* out) {
    if (shape.expert_count <= 0 || shape.expert_bytes == 0) { return false; }
    uint64_t used = 0;
    for (int moe = 1; moe <= shape.block_count; ++moe) {
        if (!placement_fits(base, shape, shape.block_count, moe,
                            budget, &used)) {
            continue;
        }
        out->known = true;
        out->gpu_layers = -1;  // all layers; experts are what moved.
        out->fully_resident = true;
        out->cpu_moe_layers = moe;
        out->bytes = used;
        out->reason = "all layers resident, experts host-side";
        return true;
    }
    return false;
}
 
inline bool try_fewer_layers(const GgufShape& shape,
                             const FootprintInputs& base,
                             uint64_t budget, AutoPlacement* out) {
    uint64_t used = 0;
    for (int layers = shape.block_count - 1; layers >= 0; --layers) {
        if (!placement_fits(base, shape, layers, 0, budget, &used)) {
            continue;
        }
        out->known = true;
        out->gpu_layers = layers;
        out->bytes = used;
        out->reason = "partial offload";
        return true;
    }
    return false;
}
 
inline bool try_smaller_ubatch(const GgufShape& shape,
                               const FootprintInputs& base,
                               uint64_t free_vram_bytes, int configured,
                               AutoPlacement* out) {
    const int current = configured > 0 ? configured : kDefaultUbatch;
    for (const int ub : kUbatchLadder) {
        if (ub >= current) { continue; }
        const uint64_t allowance = compute_allowance_bytes(base, ub);
        const uint64_t budget =
            free_vram_bytes > allowance ? free_vram_bytes - allowance : 0;
        if (!try_fully_resident(shape, base, budget, out)) { continue; }
        out->n_ubatch = ub;
        out->reason = "fully resident at a reduced n_ubatch";
        return true;
    }
    return false;
}
 
inline AutoPlacement derive_auto_placement(const GgufShape& shape,
                                           const FootprintInputs& base,
                                           uint64_t free_vram_bytes,
                                           int n_ubatch = 0) {
    AutoPlacement out;
    if (!shape.known || free_vram_bytes == 0) {
        out.reason = "GGUF shape or free VRAM unknown";
        return out;
    }
    // The headroom the operator asked for is already inside the estimate
    // (vram_reserve_mb), so it is not subtracted twice — that double-count
    // is what v2.13.0 shipped.
    //
    // What IS subtracted is a compute-buffer allowance, because the
    // estimator explicitly does not count graph/activation scratch and says
    // so: llama.cpp sizes it by ubatch and model internals, and
    // vram_footprint.h measured 1222 MiB for one gemma-4 E4B MTP context at
    // a 512-token ubatch — more than twice the 512 MB default reserve, and
    // paid TWICE by a speculative configuration holding two contexts.
    //
    // Auto is choosing on the operator's behalf, so it must leave room for
    // the thing it cannot price rather than hand back a placement that fits
    // the estimate and fails the load. This is an allowance, not a
    // derivation; it is deliberately generous, and the cost of being wrong
    // is a layer left on the host rather than a failed load.
    const uint64_t allowance = compute_allowance_bytes(base, n_ubatch);
    const uint64_t budget =
        free_vram_bytes > allowance ? free_vram_bytes - allowance : 0;
    // Order is the measurement: ubatch costs ONE axis, experts and layers
    // cost two. Spend the cheap currency first.
    if (!try_fully_resident(shape, base, budget, &out)
        && !try_smaller_ubatch(shape, base, free_vram_bytes, n_ubatch, &out)
        && !try_expert_offload(shape, base, budget, &out)
        && !try_fewer_layers(shape, base, budget, &out)) {
        out.reason = "nothing fits, even with no layers offloaded";
        return out;
    }
    // The budget already has the compute allowance netted out, so this is
    // the real physical remainder, not an accounting one. Recomputed rather
    // than threaded back out of four call sites: `try_smaller_ubatch` moves
    // the budget, so the rung it chose is the one to price against.
    const uint64_t chosen = compute_allowance_bytes(
        base, out.n_ubatch > 0 ? out.n_ubatch : n_ubatch);
    const uint64_t effective =
        free_vram_bytes > chosen ? free_vram_bytes - chosen : 0;
    out.margin_bytes = effective > out.bytes ? effective - out.bytes : 0;
    return out;
}
 
}  // namespace entropic