vram_footprint.h

Go to the documentation of this file.

// SPDX-License-Identifier: Apache-2.0
#pragma once
 
#include <algorithm>
#include <cstdint>
#include <string>
#include <utility>
 
namespace entropic {
 
constexpr uint64_t kBaseKvPerTokenF16 = 16ull * 1024ull;
 
constexpr int kAllLayersSentinel = 99;
 
enum class Offload {
    none,             
    full,             
    partial_unknown,  
};
 
struct FootprintInputs {
    uint64_t weights_bytes = 0;   
    uint64_t mmproj_bytes = 0;    
    int gpu_layers = -1;          
    int context_length = 0;       
    std::string cache_type_k = "f16";  
    std::string cache_type_v = "f16";  
    int vram_reserve_mb = 0;      
 
    int max_sessions = 1;
 
    uint64_t draft_bytes = 0;
 
    int block_count = 0;           
    uint64_t block_bytes = 0;      
    uint64_t non_block_bytes = 0;  
 
    uint64_t host_expert_bytes = 0;
};
 
struct FootprintEstimate {
    bool known = false;      
    uint64_t bytes = 0;      
    const char* reason = ""; 
};
 
inline double kv_scale_for_cache_type(const std::string& type) {
    static const std::pair<const char*, double> kScales[] = {
        {"f32",  2.0},      // 4 B/elem
        {"f16",  1.0},      // 2 B/elem  — the baseline
        {"bf16", 1.0},      // 2 B/elem
        {"q8_0", 0.53125},  // 34 B / 32 elem = 1.0625 B/elem
        {"q5_1", 0.375},    // 24 B / 32 elem = 0.75   B/elem
        {"q5_0", 0.34375},  // 22 B / 32 elem = 0.6875 B/elem
        {"q4_1", 0.3125},   // 20 B / 32 elem = 0.625  B/elem
        {"q4_0", 0.28125},  // 18 B / 32 elem = 0.5625 B/elem
    };
    double scale = 1.0;
    for (const auto& entry : kScales) {
        if (type == entry.first) {
            scale = entry.second;
            break;
        }
    }
    return scale;
}
 
inline Offload classify_offload(int gpu_layers) {
    Offload placement = Offload::partial_unknown;
    if (gpu_layers < 0 || gpu_layers >= kAllLayersSentinel) {
        placement = Offload::full;
    } else if (gpu_layers == 0) {
        placement = Offload::none;
    }
    return placement;
}
 
inline double kv_bytes_per_token(const FootprintInputs& in) {
    const double avg = (kv_scale_for_cache_type(in.cache_type_k)
                        + kv_scale_for_cache_type(in.cache_type_v)) / 2.0;
    return static_cast<double>(kBaseKvPerTokenF16) * avg;
}
 
inline FootprintEstimate estimate_vram_footprint(const FootprintInputs& in) {
    FootprintEstimate est;
    const Offload placement = classify_offload(in.gpu_layers);
    // v2.13.1: a partial offload IS priceable once the GGUF's shape has been
    // read. Without it we still decline, because guessing here refuses
    // configurations that demonstrably work (gh#142) — but the shape is now
    // available from a metadata-only read, so declining is the fallback
    // rather than the rule.
    const bool have_shape = in.block_count > 0 && in.block_bytes > 0;
    if (placement == Offload::partial_unknown && !have_shape) {
        est.reason = "partial offload: layer count is not knowable without "
                     "GGUF metadata, so the budget gate stays open";
        return est;
    }
    uint64_t resident = 0;
    if (placement == Offload::full) {
        // The projector follows the weights: offload none and it lives in
        // host RAM.
        resident = in.weights_bytes + in.mmproj_bytes;
    } else if (placement == Offload::partial_unknown) {
        // Per-block cost from the block tensors ALONE. Dividing the whole
        // file by the layer count folds in embeddings and the output head,
        // which do not move with `gpu_layers`, and so overstates every
        // layer.
        const auto per_block =
            in.block_bytes / static_cast<uint64_t>(in.block_count);
        const auto offloaded = static_cast<uint64_t>(in.gpu_layers);
        resident = per_block * offloaded;
    }
    // Experts held host-side by `cpu_moe_layers` are not on the card.
    resident -= std::min(resident, in.host_expert_bytes);
    // The draft head is resident whenever speculative decode is configured,
    // at any placement.
    resident += in.draft_bytes;
    const int ctx = in.context_length > 0 ? in.context_length : 0;
    const int sessions = in.max_sessions > 0 ? in.max_sessions : 1;
    // gh#144 (v2.12.0): context_length is per session; the pool allocates
    // ctx * sessions cells. Deliberately an over-count for an iSWA model
    // like Gemma-4: kBaseKvPerTokenF16 is a flat rate and does not model
    // sliding-window layers (5:1 SWA at window 512 under swa_full=false),
    // so the real figure is well below this. That is the safe direction
    // per this header's own rule — never silently under-count.
    const uint64_t kv = static_cast<uint64_t>(
        static_cast<double>(ctx) * static_cast<double>(sessions)
        * kv_bytes_per_token(in));
    est.known = true;
    est.bytes = resident + kv
        + static_cast<uint64_t>(in.vram_reserve_mb) * 1024ull * 1024ull;
    return est;
}
 
inline int recommend_context_length(const FootprintInputs& in,
                                    uint64_t available_bytes) {
    FootprintInputs probe = in;
    probe.context_length = 0;
    const FootprintEstimate fixed = estimate_vram_footprint(probe);
    if (!fixed.known || fixed.bytes >= available_bytes) {
        return 0;
    }
    // gh#144 (v2.12.0): divide the budget across the pool, or the
    // recommendation is a PER-SESSION window N times too large — handing an
    // operator a config that aborts, from the very function whose job is to
    // offer one that works.
    const int sessions = in.max_sessions > 0 ? in.max_sessions : 1;
    const double per_token =
        kv_bytes_per_token(in) * static_cast<double>(sessions);
    const uint64_t kv_budget = available_bytes - fixed.bytes;
    const uint64_t fits = static_cast<uint64_t>(
        static_cast<double>(kv_budget) / per_token);
    const uint64_t requested = in.context_length > 0
        ? static_cast<uint64_t>(in.context_length) : 0ull;
    uint64_t chosen = fits < requested ? fits : requested;
    chosen -= chosen % 512ull;
    return static_cast<int>(chosen);
}
 
}  // namespace entropic