vram_footprint.h

Go to the documentation of this file.

// SPDX-License-Identifier: Apache-2.0
#pragma once
 
#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;      
};
 
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);
    if (placement == Offload::partial_unknown) {
        est.reason = "partial offload: layer count is not knowable without "
                     "GGUF metadata, so the budget gate stays open";
        return est;
    }
    // The projector follows the weights: offload none and it lives in host RAM.
    const uint64_t resident = (placement == Offload::full)
        ? in.weights_bytes + in.mmproj_bytes
        : 0ull;
    const int ctx = in.context_length > 0 ? in.context_length : 0;
    const uint64_t kv = static_cast<uint64_t>(
        static_cast<double>(ctx) * 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;
    }
    const double per_token = kv_bytes_per_token(in);
    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