partial_offload.h

Go to the documentation of this file.

// SPDX-License-Identifier: Apache-2.0
#pragma once
 
#include <algorithm>
#include <cstdint>
#include <cstdlib>
 
namespace entropic {
 
constexpr uint64_t kPartialOffloadReserveBytes = 2048ull * 1024 * 1024;
 
inline int partial_gpu_layers_for(uint64_t file_bytes,
                                  uint64_t free_vram_bytes,
                                  int n_layers = 30) {
    if (file_bytes == 0 || n_layers <= 0) { return 0; }
    if (free_vram_bytes <= kPartialOffloadReserveBytes) { return 0; }
    const uint64_t budget = free_vram_bytes - kPartialOffloadReserveBytes;
    // Per-layer cost approximated from the file: exact enough to choose a
    // split, and it needs no GGUF metadata this header refuses to read.
    const uint64_t per_layer =
        file_bytes / static_cast<uint64_t>(n_layers);
    if (per_layer == 0) { return n_layers; }
    const auto fits = static_cast<int>(budget / per_layer);
    return std::clamp(fits, 0, n_layers);
}
 
inline bool host_can_hold_warm_load(
    uint64_t file_bytes, uint64_t available_bytes,
    uint64_t headroom_bytes = 1024ull * 1024 * 1024) {
    if (available_bytes == 0) { return true; }
    // WARM and ACTIVE residencies can overlap across the reload boundary.
    const uint64_t peak = file_bytes + (file_bytes / 2);
    return available_bytes >= peak + headroom_bytes;
}
 
constexpr uint64_t kLargeModelThresholdBytes = 10ull * 1024 * 1024 * 1024;
 
inline bool large_model_tests_waived(uint64_t file_bytes) {
    if (file_bytes <= kLargeModelThresholdBytes) { return false; }
    const char* v = std::getenv("ENTROPIC_SKIP_LARGE_MODEL_TESTS");
    return v != nullptr && v[0] == '1';
}
 
constexpr uint64_t kMemlockUnlimited = UINT64_MAX;
 
inline bool mlock_refused(bool use_mlock, uint64_t file_bytes,
                          int gpu_layers, uint64_t memlock_limit_bytes,
                          uint64_t host_expert_bytes = 0) {
    // v2.13.1: `gpu_layers < 0` exempts because the host side is then
    // TRANSIENT — nothing stays in host RAM to be pinned. An expert split
    // breaks that premise: `cpu_moe_layers` deliberately leaves expert
    // weights host-RESIDENT while every layer reports as offloaded, so the
    // sentinel alone would exempt exactly the configuration the rule exists
    // to catch. Reachable since v2.13.1, when `gpu_layers: auto` began
    // deriving expert splits and emitting -1 alongside them.
    const bool host_side_is_transient =
        gpu_layers < 0 && host_expert_bytes == 0;
    const bool exempt = !use_mlock
        || host_side_is_transient
        || file_bytes <= kLargeModelThresholdBytes      // never refuse small
        || memlock_limit_bytes == kMemlockUnlimited;
    if (exempt) { return false; }
    // Under an expert split it is the HOST-resident expert bytes that get
    // pinned, not the whole file; pricing the file would refuse splits that
    // comfortably fit the limit.
    const uint64_t pinned =
        host_expert_bytes > 0 ? host_expert_bytes : file_bytes;
    return pinned > memlock_limit_bytes;
}
 
inline bool gpu_offload_refused(uint64_t file_bytes, int gpu_layers,
                                uint64_t free_vram_bytes) {
    const bool exempt = gpu_layers == 0
        || free_vram_bytes == 0
        || file_bytes <= kLargeModelThresholdBytes;
    return !exempt && free_vram_bytes <= kPartialOffloadReserveBytes;
}
 
}  // namespace entropic