Compare commits
6
Commits
e5eb5edb58
...
deee33503b
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
deee33503b | ||
|
|
f2d79d89f5 | ||
|
|
bf68c5446e | ||
|
|
7398c40eda | ||
|
|
be7f3b3172 | ||
|
|
2b94398ed7 |
+7
-5
@@ -2575,15 +2575,17 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
|
||||
).set_env("LLAMA_ARG_CPU_MOE"));
|
||||
add_opt(common_arg(
|
||||
{"-ncmoe", "--n-cpu-moe"}, "N",
|
||||
"keep the Mixture of Experts (MoE) weights of the first N layers in the CPU",
|
||||
[](common_params & params, int value) {
|
||||
if (value < 0) {
|
||||
"keep the Mixture of Experts (MoE) weights of the first N layers in the CPU; "
|
||||
"fractional N offloads part of the boundary layer at tensor granularity",
|
||||
[](common_params & params, const std::string & value) {
|
||||
const double n = std::stod(value);
|
||||
if (n < 0) {
|
||||
throw std::invalid_argument("invalid value");
|
||||
}
|
||||
for (int i = 0; i < value; ++i) {
|
||||
for (const std::string & re : llm_ffn_exps_cpu_block_regexes(n)) {
|
||||
// keep strings alive and avoid leaking memory by storing them in a static vector
|
||||
static std::list<std::string> buft_overrides;
|
||||
buft_overrides.push_back(llm_ffn_exps_block_regex(i));
|
||||
buft_overrides.push_back(re);
|
||||
params.tensor_buft_overrides.push_back({buft_overrides.back().c_str(), ggml_backend_cpu_buffer_type()});
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1080,6 +1080,28 @@ inline llama_model_tensor_buft_override llm_ffn_exps_cpu_override() {
|
||||
return { LLM_FFN_EXPS_REGEX, ggml_backend_cpu_buffer_type() };
|
||||
}
|
||||
|
||||
// ATSInfer-style tensor-granularity static placement of MoE expert weights.
|
||||
// Offloads the expert weights of the first floor(n) layers to the CPU, plus a
|
||||
// subset of the boundary layer's three expert tensors for the fractional part.
|
||||
// Expert tensors are dropped from the GPU in ascending performance-density
|
||||
// order (gate, then up), keeping the higher-value down_proj resident longest.
|
||||
inline std::vector<std::string> llm_ffn_exps_cpu_block_regexes(double n_cpu_moe) {
|
||||
std::vector<std::string> regexes;
|
||||
const int n_full = n_cpu_moe > 0 ? (int) n_cpu_moe : 0;
|
||||
for (int i = 0; i < n_full; ++i) {
|
||||
regexes.push_back(llm_ffn_exps_block_regex(i));
|
||||
}
|
||||
const int k = (int) ((n_cpu_moe - n_full) * 3.0 + 0.5);
|
||||
if (k >= 3) {
|
||||
regexes.push_back(llm_ffn_exps_block_regex(n_full));
|
||||
} else if (k == 2) {
|
||||
regexes.push_back(string_format("blk\\.%d\\.ffn_(gate|up)_(ch|)exps", n_full));
|
||||
} else if (k == 1) {
|
||||
regexes.push_back(string_format("blk\\.%d\\.ffn_gate_(ch|)exps", n_full));
|
||||
}
|
||||
return regexes;
|
||||
}
|
||||
|
||||
//
|
||||
// training utils
|
||||
//
|
||||
|
||||
@@ -557,6 +557,16 @@ extern "C" {
|
||||
|
||||
LLAMA_API const struct llama_model * llama_get_model (const struct llama_context * ctx);
|
||||
LLAMA_API llama_memory_t llama_get_memory (const struct llama_context * ctx);
|
||||
|
||||
// On-demand device (VRAM) residency: free the model's GPU weight buffers (keeping a host
|
||||
// shadow so no reload is needed) to hand VRAM to another model, and rebuild them on demand.
|
||||
// Any subsequent llama_decode auto-restores. If evict_kv is true, the KV cache device buffers
|
||||
// are also evicted to a host shadow (for when the KV would not fit alongside the other model),
|
||||
// at the cost of a D2H/H2D copy of the live KV each cycle; otherwise the KV stays resident.
|
||||
// Intended for time-sharing a single GPU between several always-loaded models.
|
||||
LLAMA_API void llama_context_release_device(struct llama_context * ctx, bool evict_kv);
|
||||
LLAMA_API void llama_context_restore_device(struct llama_context * ctx);
|
||||
LLAMA_API bool llama_context_weights_resident(const struct llama_context * ctx);
|
||||
LLAMA_API enum llama_pooling_type llama_pooling_type(const struct llama_context * ctx); // TODO: rename to llama_get_pooling_type
|
||||
|
||||
LLAMA_API const struct llama_vocab * llama_model_get_vocab(const struct llama_model * model);
|
||||
|
||||
@@ -728,6 +728,34 @@ void llama_context::synchronize() {
|
||||
t_compute_start_us = 0;
|
||||
}
|
||||
|
||||
void llama_context::release_device(bool evict_kv) {
|
||||
if (!model.weights_resident() && !(evict_kv && !kv_device_evicted)) {
|
||||
return;
|
||||
}
|
||||
// ensure no compute is in flight before freeing the device buffers
|
||||
synchronize();
|
||||
model.release_device_weights();
|
||||
if (evict_kv && memory && !kv_device_evicted) {
|
||||
memory->release_device_buffers();
|
||||
kv_device_evicted = true;
|
||||
}
|
||||
}
|
||||
|
||||
void llama_context::restore_device() {
|
||||
if (model.weights_resident() && !kv_device_evicted) {
|
||||
return;
|
||||
}
|
||||
model.restore_device_weights();
|
||||
if (kv_device_evicted && memory) {
|
||||
memory->restore_device_buffers();
|
||||
kv_device_evicted = false;
|
||||
}
|
||||
// make sure all weight/KV uploads have completed before any compute reads them
|
||||
if (sched) {
|
||||
ggml_backend_sched_synchronize(sched.get());
|
||||
}
|
||||
}
|
||||
|
||||
const llama_model & llama_context::get_model() const {
|
||||
return model;
|
||||
}
|
||||
@@ -1705,6 +1733,12 @@ int llama_context::decode(const llama_batch & batch_inp) {
|
||||
return -1;
|
||||
}
|
||||
|
||||
// on-demand VRAM: if the weights were released while this model was idle, bring them back
|
||||
// to the device before building the compute graph
|
||||
if (!model.weights_resident()) {
|
||||
restore_device();
|
||||
}
|
||||
|
||||
const auto & vocab = model.vocab;
|
||||
const auto & hparams = model.hparams;
|
||||
|
||||
@@ -3674,6 +3708,18 @@ void llama_synchronize(llama_context * ctx) {
|
||||
ctx->synchronize();
|
||||
}
|
||||
|
||||
void llama_context_release_device(llama_context * ctx, bool evict_kv) {
|
||||
ctx->release_device(evict_kv);
|
||||
}
|
||||
|
||||
void llama_context_restore_device(llama_context * ctx) {
|
||||
ctx->restore_device();
|
||||
}
|
||||
|
||||
bool llama_context_weights_resident(const llama_context * ctx) {
|
||||
return ctx->get_model().weights_resident();
|
||||
}
|
||||
|
||||
float * llama_get_logits(llama_context * ctx) {
|
||||
ctx->synchronize();
|
||||
|
||||
|
||||
@@ -57,6 +57,14 @@ struct llama_context {
|
||||
|
||||
void synchronize();
|
||||
|
||||
// on-demand device (VRAM) residency: free / rebuild the model's GPU weight buffers so that
|
||||
// VRAM can be time-shared with another model. release_device() synchronizes first; decode()
|
||||
// auto-restores when it finds the weights have been released. If evict_kv is set, the KV cache
|
||||
// device buffers are also evicted to a host shadow (for when the KV would not fit alongside the
|
||||
// other model) - this adds a D2H/H2D copy of the live KV on each cycle.
|
||||
void release_device(bool evict_kv = false);
|
||||
void restore_device();
|
||||
|
||||
const llama_model & get_model() const;
|
||||
const llama_cparams & get_cparams() const;
|
||||
|
||||
@@ -345,6 +353,9 @@ private:
|
||||
|
||||
bool sched_need_reserve = true;
|
||||
|
||||
// on-demand device residency: true while the KV cache device buffers have been evicted to host
|
||||
bool kv_device_evicted = false;
|
||||
|
||||
ggml_backend_t backend_cpu = nullptr;
|
||||
std::vector<ggml_backend_ptr> backends;
|
||||
|
||||
|
||||
@@ -272,6 +272,17 @@ void llama_kv_cache_iswa::state_read(llama_io_read_i & io, llama_seq_id seq_id,
|
||||
kv_swa->state_read(io, seq_id, flags);
|
||||
}
|
||||
|
||||
void llama_kv_cache_iswa::release_device_buffers() {
|
||||
kv_base->release_device_buffers();
|
||||
kv_swa->release_device_buffers();
|
||||
}
|
||||
|
||||
bool llama_kv_cache_iswa::restore_device_buffers() {
|
||||
bool ok = kv_base->restore_device_buffers();
|
||||
ok = kv_swa->restore_device_buffers() && ok;
|
||||
return ok;
|
||||
}
|
||||
|
||||
llama_kv_cache * llama_kv_cache_iswa::get_base() const {
|
||||
return kv_base.get();
|
||||
}
|
||||
|
||||
@@ -83,6 +83,9 @@ public:
|
||||
void state_write(llama_io_write_i & io, llama_seq_id seq_id = -1, llama_state_seq_flags flags = 0) const override;
|
||||
void state_read (llama_io_read_i & io, llama_seq_id seq_id = -1, llama_state_seq_flags flags = 0) override;
|
||||
|
||||
void release_device_buffers() override;
|
||||
bool restore_device_buffers() override;
|
||||
|
||||
//
|
||||
// llama_kv_cache_iswa specific API
|
||||
//
|
||||
|
||||
+83
-2
@@ -371,7 +371,9 @@ void llama_kv_cache::clear(bool data) {
|
||||
|
||||
if (data) {
|
||||
for (auto & [_, buf] : ctxs_bufs) {
|
||||
ggml_backend_buffer_clear(buf.get(), 0);
|
||||
if (buf) { // may be null if evicted for on-demand VRAM sharing
|
||||
ggml_backend_buffer_clear(buf.get(), 0);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -681,6 +683,9 @@ llama_pos llama_kv_cache::seq_pos_max(llama_seq_id seq_id) const {
|
||||
std::map<ggml_backend_buffer_type_t, size_t> llama_kv_cache::memory_breakdown() const {
|
||||
std::map<ggml_backend_buffer_type_t, size_t> ret;
|
||||
for (const auto & [ctx, buf] : ctxs_bufs) {
|
||||
if (!buf) { // may be null if evicted for on-demand VRAM sharing
|
||||
continue;
|
||||
}
|
||||
ggml_backend_buffer_type_t buft = ggml_backend_buffer_get_type(buf.get());
|
||||
|
||||
if (hparams.no_alloc) {
|
||||
@@ -1801,12 +1806,88 @@ size_t llama_kv_cache::total_size() const {
|
||||
size_t size = 0;
|
||||
|
||||
for (const auto & [_, buf] : ctxs_bufs) {
|
||||
size += ggml_backend_buffer_get_size(buf.get());
|
||||
if (buf) { // may be null if evicted for on-demand VRAM sharing
|
||||
size += ggml_backend_buffer_get_size(buf.get());
|
||||
}
|
||||
}
|
||||
|
||||
return size;
|
||||
}
|
||||
|
||||
void llama_kv_cache::release_device_buffers() {
|
||||
// NOTE: the caller must have synchronized the backend so nothing references these buffers.
|
||||
// The KV cache is read-write, so its host shadow is (re)captured fresh on every release.
|
||||
if (dev_released) {
|
||||
return;
|
||||
}
|
||||
dev_shadows.assign(ctxs_bufs.size(), device_buffer_shadow{});
|
||||
size_t freed = 0;
|
||||
for (size_t i = 0; i < ctxs_bufs.size(); ++i) {
|
||||
ggml_context * ctx = ctxs_bufs[i].first.get();
|
||||
ggml_backend_buffer_t buf = ctxs_bufs[i].second.get();
|
||||
if (buf == nullptr || ggml_backend_buffer_is_host(buf) || ggml_backend_buffer_get_size(buf) == 0) {
|
||||
continue; // only real device (VRAM) buffers are evictable
|
||||
}
|
||||
auto & sh = dev_shadows[i];
|
||||
sh.releasable = true;
|
||||
sh.buft = ggml_backend_buffer_get_type(buf);
|
||||
|
||||
// capture live contents compactly in stable iteration order (skip views, which alias a base)
|
||||
size_t total = 0;
|
||||
for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) {
|
||||
if (t->view_src == nullptr) { total += ggml_nbytes(t); }
|
||||
}
|
||||
sh.data.resize(total);
|
||||
size_t off = 0;
|
||||
for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) {
|
||||
if (t->view_src != nullptr) { continue; }
|
||||
const size_t n = ggml_nbytes(t);
|
||||
ggml_backend_tensor_get(t, sh.data.data() + off, 0, n);
|
||||
off += n;
|
||||
}
|
||||
|
||||
freed += ggml_backend_buffer_get_size(buf);
|
||||
ctxs_bufs[i].second.reset(); // free the device buffer
|
||||
for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) {
|
||||
t->buffer = nullptr;
|
||||
t->data = nullptr;
|
||||
}
|
||||
}
|
||||
dev_released = true;
|
||||
if (freed > 0) {
|
||||
LLAMA_LOG_INFO("%s: released %.2f MiB of KV cache from device\n", __func__, freed / 1024.0 / 1024.0);
|
||||
}
|
||||
}
|
||||
|
||||
bool llama_kv_cache::restore_device_buffers() {
|
||||
if (!dev_released) {
|
||||
return true;
|
||||
}
|
||||
for (size_t i = 0; i < ctxs_bufs.size(); ++i) {
|
||||
auto & sh = dev_shadows[i];
|
||||
if (!sh.releasable) {
|
||||
continue;
|
||||
}
|
||||
ggml_context * ctx = ctxs_bufs[i].first.get();
|
||||
ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors_from_buft(ctx, sh.buft);
|
||||
if (buf == nullptr) {
|
||||
LLAMA_LOG_ERROR("%s: failed to reallocate KV cache device buffer (out of VRAM?)\n", __func__);
|
||||
return false;
|
||||
}
|
||||
size_t off = 0;
|
||||
for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) {
|
||||
if (t->view_src != nullptr) { continue; }
|
||||
const size_t n = ggml_nbytes(t);
|
||||
ggml_backend_tensor_set(t, sh.data.data() + off, 0, n);
|
||||
off += n;
|
||||
}
|
||||
ctxs_bufs[i].second.reset(buf);
|
||||
}
|
||||
dev_released = false;
|
||||
dev_shadows.clear(); // recaptured on next release
|
||||
return true;
|
||||
}
|
||||
|
||||
size_t llama_kv_cache::size_k_bytes() const {
|
||||
size_t size_k_bytes = 0;
|
||||
|
||||
|
||||
@@ -149,6 +149,10 @@ public:
|
||||
void state_write(llama_io_write_i & io, llama_seq_id seq_id = -1, llama_state_seq_flags flags = 0) const override;
|
||||
void state_read (llama_io_read_i & io, llama_seq_id seq_id = -1, llama_state_seq_flags flags = 0) override;
|
||||
|
||||
// on-demand device (VRAM) residency (see llama_memory_i)
|
||||
void release_device_buffers() override;
|
||||
bool restore_device_buffers() override;
|
||||
|
||||
//
|
||||
// llama_kv_cache specific API
|
||||
//
|
||||
@@ -265,6 +269,16 @@ private:
|
||||
// ggml contexts for the KV cache along with the allocated backend buffers:
|
||||
std::vector<std::pair<ggml_context_ptr, ggml_backend_buffer_ptr>> ctxs_bufs;
|
||||
|
||||
// on-demand device eviction: host shadow of each device buffer's live contents (re-captured on
|
||||
// every release since the KV is read-write), plus the buffer type needed to reallocate it.
|
||||
struct device_buffer_shadow {
|
||||
ggml_backend_buffer_type_t buft = nullptr;
|
||||
bool releasable = false; // device (VRAM) buffer that was freed
|
||||
std::vector<uint8_t> data; // captured tensor bytes (compact, iteration order)
|
||||
};
|
||||
std::vector<device_buffer_shadow> dev_shadows; // parallel to ctxs_bufs
|
||||
bool dev_released = false;
|
||||
|
||||
// the current index from where we start searching for a free slot in the ring buffer of KV cells (see find_slot())
|
||||
// note: this is not part of the KV state and it's only used to speed-up the find_slot() method
|
||||
std::vector<uint32_t> v_heads;
|
||||
|
||||
@@ -201,6 +201,18 @@ void llama_memory_hybrid::state_read(llama_io_read_i & io, llama_seq_id seq_id,
|
||||
mem_recr->state_read(io, seq_id, flags);
|
||||
}
|
||||
|
||||
void llama_memory_hybrid::release_device_buffers() {
|
||||
// evict the attention KV (grows with context); the recurrent state uses the no-op default
|
||||
mem_attn->release_device_buffers();
|
||||
mem_recr->release_device_buffers();
|
||||
}
|
||||
|
||||
bool llama_memory_hybrid::restore_device_buffers() {
|
||||
bool ok = mem_attn->restore_device_buffers();
|
||||
ok = mem_recr->restore_device_buffers() && ok;
|
||||
return ok;
|
||||
}
|
||||
|
||||
llama_kv_cache * llama_memory_hybrid::get_mem_attn() const {
|
||||
return mem_attn.get();
|
||||
}
|
||||
|
||||
@@ -76,6 +76,9 @@ public:
|
||||
void state_write(llama_io_write_i & io, llama_seq_id seq_id = -1, llama_state_seq_flags flags = 0) const override;
|
||||
void state_read (llama_io_read_i & io, llama_seq_id seq_id = -1, llama_state_seq_flags flags = 0) override;
|
||||
|
||||
void release_device_buffers() override;
|
||||
bool restore_device_buffers() override;
|
||||
|
||||
//
|
||||
// llama_memory_hybrid specific API
|
||||
//
|
||||
|
||||
@@ -124,6 +124,16 @@ struct llama_memory_i {
|
||||
|
||||
virtual void state_write(llama_io_write_i & io, llama_seq_id seq_id = -1, llama_state_seq_flags flags = 0) const = 0;
|
||||
virtual void state_read (llama_io_read_i & io, llama_seq_id seq_id = -1, llama_state_seq_flags flags = 0) = 0;
|
||||
|
||||
//
|
||||
// on-demand device (VRAM) residency
|
||||
//
|
||||
|
||||
// Free this memory's device (VRAM) buffers to a host shadow and rebuild them on demand, to
|
||||
// hand VRAM to another model while keeping the cached data (no re-prefill). The caller must
|
||||
// have synchronized the backend first. Default: no-op (the memory stays resident).
|
||||
virtual void release_device_buffers() {}
|
||||
virtual bool restore_device_buffers() { return true; }
|
||||
};
|
||||
|
||||
using llama_memory_ptr = std::unique_ptr<llama_memory_i>;
|
||||
|
||||
@@ -1015,6 +1015,16 @@ struct llama_model::impl {
|
||||
// contexts where the model tensors metadata is stored as well as the corresponding buffers:
|
||||
std::vector<std::pair<ggml_context_ptr, std::vector<ggml_backend_buffer_ptr>>> ctxs_bufs;
|
||||
|
||||
// on-demand device (VRAM) weight residency: per-ctxs_bufs slot describing whether the
|
||||
// device weight buffer can be freed and rebuilt from a host shadow (see release/restore_device_weights)
|
||||
struct weight_release_slot {
|
||||
ggml_backend_buffer_type_t buft = nullptr; // buffer type to reallocate on restore
|
||||
bool releasable = false; // single-buffer alloc-path device (VRAM) buffer
|
||||
bool resident = true; // currently allocated on device
|
||||
std::vector<uint8_t> shadow; // host copy, captured lazily on first release
|
||||
};
|
||||
std::vector<weight_release_slot> weight_release; // parallel to ctxs_bufs
|
||||
|
||||
buft_list_t cpu_buft_list;
|
||||
std::map<ggml_backend_dev_t, buft_list_t> gpu_buft_list;
|
||||
|
||||
@@ -1608,6 +1618,23 @@ bool llama_model_base::load_tensors(llama_model_loader & ml) {
|
||||
|
||||
pimpl->ctxs_bufs.emplace_back(std::move(ctx_ptr), std::move(bufs));
|
||||
|
||||
// record on-demand release metadata (parallel to ctxs_bufs). Only the single-buffer
|
||||
// alloc path on a device (non-host) buffer can be freed and rebuilt from a host shadow;
|
||||
// the mmap buffer_from_host_ptr path and host (CPU) buffers are left resident.
|
||||
{
|
||||
llama_model::impl::weight_release_slot slot;
|
||||
const bool mmap_host_ptr_path = ml.use_mmap && use_mmap_buffer && buffer_from_host_ptr_supported && is_default_buft;
|
||||
auto & last_bufs = pimpl->ctxs_bufs.back().second;
|
||||
if (!mmap_host_ptr_path && last_bufs.size() == 1) {
|
||||
ggml_backend_buffer_t b = last_bufs[0].get();
|
||||
if (b != nullptr && !ggml_backend_buffer_is_host(b)) {
|
||||
slot.releasable = true;
|
||||
slot.buft = buft;
|
||||
}
|
||||
}
|
||||
pimpl->weight_release.push_back(std::move(slot));
|
||||
}
|
||||
|
||||
ctx_buf_maps.emplace_back(ctx, buf_map);
|
||||
}
|
||||
|
||||
@@ -1662,6 +1689,97 @@ ggml_tensor * llama_model_base::create_tensor(llama_model_loader & ml, const LLM
|
||||
tn, ne, flags);
|
||||
}
|
||||
|
||||
void llama_model::release_device_weights() const {
|
||||
// NOTE: the caller is responsible for synchronizing the backend scheduler first, so that no
|
||||
// compute is in flight referencing these buffers when they are freed.
|
||||
size_t freed = 0;
|
||||
for (size_t i = 0; i < pimpl->ctxs_bufs.size(); ++i) {
|
||||
auto & slot = pimpl->weight_release[i];
|
||||
if (!slot.releasable || !slot.resident) {
|
||||
continue;
|
||||
}
|
||||
ggml_context * ctx = pimpl->ctxs_bufs[i].first.get();
|
||||
ggml_backend_buffer_t buf = pimpl->ctxs_bufs[i].second[0].get();
|
||||
const size_t sz = ggml_backend_buffer_get_size(buf);
|
||||
|
||||
// lazily capture a host shadow of the weights (read-only, so captured only once).
|
||||
// Store tensor bytes compactly in stable iteration order (independent of buffer layout /
|
||||
// alignment padding) so restore does not depend on the re-allocated offsets matching.
|
||||
// View tensors alias their base, so they are skipped (restored implicitly via their base).
|
||||
if (slot.shadow.empty()) {
|
||||
size_t total = 0;
|
||||
for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) {
|
||||
if (t->view_src != nullptr) continue;
|
||||
total += ggml_nbytes(t);
|
||||
}
|
||||
slot.shadow.resize(total);
|
||||
size_t off = 0;
|
||||
for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) {
|
||||
if (t->view_src != nullptr) continue;
|
||||
const size_t n = ggml_nbytes(t);
|
||||
ggml_backend_tensor_get(t, slot.shadow.data() + off, 0, n);
|
||||
off += n;
|
||||
}
|
||||
}
|
||||
|
||||
// free the device (VRAM) buffer and clear the now-dangling tensor pointers so that
|
||||
// ggml_backend_alloc_ctx_tensors_from_buft reallocates them cleanly on restore
|
||||
pimpl->ctxs_bufs[i].second[0].reset();
|
||||
for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) {
|
||||
t->buffer = nullptr;
|
||||
t->data = nullptr;
|
||||
}
|
||||
slot.resident = false;
|
||||
freed += sz;
|
||||
}
|
||||
if (freed > 0) {
|
||||
LLAMA_LOG_INFO("%s: released %.2f MiB of device weights\n", __func__, freed / 1024.0 / 1024.0);
|
||||
}
|
||||
}
|
||||
|
||||
bool llama_model::restore_device_weights() const {
|
||||
size_t restored = 0;
|
||||
for (size_t i = 0; i < pimpl->ctxs_bufs.size(); ++i) {
|
||||
auto & slot = pimpl->weight_release[i];
|
||||
if (!slot.releasable || slot.resident) {
|
||||
continue;
|
||||
}
|
||||
ggml_context * ctx = pimpl->ctxs_bufs[i].first.get();
|
||||
ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors_from_buft(ctx, slot.buft);
|
||||
if (buf == nullptr) {
|
||||
LLAMA_LOG_ERROR("%s: failed to reallocate device weight buffer (out of VRAM?)\n", __func__);
|
||||
return false;
|
||||
}
|
||||
ggml_backend_buffer_set_usage(buf, GGML_BACKEND_BUFFER_USAGE_WEIGHTS);
|
||||
|
||||
// re-upload weights from the compact host shadow, using the same stable iteration order
|
||||
// and view-skipping as the capture above
|
||||
size_t off = 0;
|
||||
for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) {
|
||||
if (t->view_src != nullptr) continue;
|
||||
const size_t n = ggml_nbytes(t);
|
||||
ggml_backend_tensor_set(t, slot.shadow.data() + off, 0, n);
|
||||
off += n;
|
||||
}
|
||||
pimpl->ctxs_bufs[i].second[0].reset(buf);
|
||||
slot.resident = true;
|
||||
restored += ggml_backend_buffer_get_size(buf);
|
||||
}
|
||||
if (restored > 0) {
|
||||
LLAMA_LOG_INFO("%s: restored %.2f MiB of device weights\n", __func__, restored / 1024.0 / 1024.0);
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
bool llama_model::weights_resident() const {
|
||||
for (const auto & slot : pimpl->weight_release) {
|
||||
if (slot.releasable && !slot.resident) {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
std::string llama_model::arch_name() const {
|
||||
return llm_arch_name(arch);
|
||||
}
|
||||
@@ -1714,6 +1832,11 @@ std::map<ggml_backend_buffer_type_t, size_t> llama_model::memory_breakdown() con
|
||||
ret[buft] += ggml_backend_alloc_ctx_tensors_from_buft_size(ctx.get(), buft);
|
||||
} else {
|
||||
for (const auto & buf : bufs) {
|
||||
// buf may be null if the device weights were released for on-demand VRAM sharing
|
||||
// (see release_device_weights); skip it in the breakdown
|
||||
if (!buf) {
|
||||
continue;
|
||||
}
|
||||
// GGML_ASSERT(ggml_backend_buffer_get_base(buf.get()) != nullptr); // multi_buffer does not have a defined base
|
||||
ret[ggml_backend_buffer_get_type(buf.get())] += ggml_backend_buffer_get_size(buf.get());
|
||||
}
|
||||
|
||||
@@ -663,6 +663,14 @@ struct llama_model {
|
||||
|
||||
const struct ggml_tensor * get_tensor(const char * name) const;
|
||||
|
||||
// on-demand device (VRAM) weight residency: free the GPU weight buffers (keeping a host
|
||||
// shadow) and rebuild them on demand. Used to time-share VRAM between models without
|
||||
// reloading or losing the KV cache. release_device_weights() requires the caller to have
|
||||
// synchronized the backend scheduler first.
|
||||
void release_device_weights() const;
|
||||
bool restore_device_weights() const;
|
||||
bool weights_resident() const;
|
||||
|
||||
float get_rope_freq_base (const llama_cparams & cparams, int il) const;
|
||||
float get_rope_freq_scale(const llama_cparams & cparams, int il) const;
|
||||
|
||||
|
||||
@@ -25,6 +25,8 @@
|
||||
#include <filesystem>
|
||||
#include <utility>
|
||||
#include <fstream>
|
||||
#include <thread>
|
||||
#include <atomic>
|
||||
|
||||
// fix problem with std::min and std::max
|
||||
#if defined(_WIN32)
|
||||
@@ -35,6 +37,16 @@
|
||||
#include <windows.h>
|
||||
#endif
|
||||
|
||||
// POSIX file locking + inotify doorbell for the cross-process VRAM arbiter (see vram_share_* below)
|
||||
#if !defined(_WIN32)
|
||||
#include <sys/file.h>
|
||||
#include <sys/stat.h>
|
||||
#include <sys/inotify.h>
|
||||
#include <fcntl.h>
|
||||
#include <unistd.h>
|
||||
#include <poll.h>
|
||||
#endif
|
||||
|
||||
using json = nlohmann::ordered_json;
|
||||
|
||||
constexpr int HTTP_POLLING_SECONDS = 1;
|
||||
@@ -871,6 +883,8 @@ public:
|
||||
}
|
||||
|
||||
~server_context_impl() {
|
||||
// stop the VRAM warden thread (and release the token) before tearing anything down
|
||||
vram_share_shutdown();
|
||||
if (!sleeping) {
|
||||
// destroy() is already called when entering sleeping state
|
||||
// we don't call it again here to avoid double free
|
||||
@@ -894,6 +908,188 @@ private:
|
||||
llama_model * model_dft = nullptr;
|
||||
llama_context * ctx_dft = nullptr;
|
||||
|
||||
// Cross-process VRAM arbiter: when LLAMA_SLEEP_VRAM_ONLY is set, several always-loaded
|
||||
// llama-server processes time-share one GPU. A single flock() on <arena>/token.lock is the
|
||||
// baton: "resident (weights in VRAM) iff I hold the token". A model warms up (locks + restores
|
||||
// weights) right before it decodes, and goes cold (releases weights + unlocks) when it idles,
|
||||
// so only one model holds VRAM at a time and the KV cache is never evicted (no re-prefill).
|
||||
bool vram_only = false; // release-mode sleep active (LLAMA_SLEEP_VRAM_ONLY)
|
||||
bool vram_evict_kv = false; // also evict the KV cache to host on release (LLAMA_SLEEP_EVICT_KV)
|
||||
bool vram_flock = false; // cross-process flock coordination available
|
||||
std::atomic<bool> vram_cold{true}; // weights currently released (read by warden thread)
|
||||
int vram_lock_fd = -1; // fd for <arena>/token.lock
|
||||
|
||||
// inotify "doorbell": a waiter that wants the VRAM token touches <arena>/doorbell/<pid>, which
|
||||
// wakes the current holder's warden thread so it releases immediately instead of only on idle.
|
||||
std::string vram_doorbell_dir;
|
||||
std::string vram_pid_str;
|
||||
int vram_inotify_fd = -1;
|
||||
std::thread vram_warden;
|
||||
std::atomic<bool> vram_warden_run{false};
|
||||
|
||||
// open the shared arena (flock token + doorbell dir). Idempotent; sets vram_only/vram_flock.
|
||||
void vram_arena_open() {
|
||||
if (getenv("LLAMA_SLEEP_VRAM_ONLY") == nullptr) {
|
||||
return;
|
||||
}
|
||||
vram_only = true;
|
||||
vram_evict_kv = getenv("LLAMA_SLEEP_EVICT_KV") != nullptr;
|
||||
#if !defined(_WIN32)
|
||||
if (vram_lock_fd >= 0) {
|
||||
return; // already open
|
||||
}
|
||||
const char * arena_env = getenv("LLAMA_VRAM_ARENA");
|
||||
const std::string arena = arena_env ? arena_env : "/dev/shm/llama-vram";
|
||||
mkdir(arena.c_str(), 0777);
|
||||
const std::string lock_path = arena + "/token.lock";
|
||||
vram_lock_fd = open(lock_path.c_str(), O_RDWR | O_CREAT, 0666);
|
||||
if (vram_lock_fd >= 0) {
|
||||
vram_flock = true;
|
||||
vram_pid_str = std::to_string(getpid());
|
||||
vram_doorbell_dir = arena + "/doorbell";
|
||||
mkdir(vram_doorbell_dir.c_str(), 0777);
|
||||
SRV_INF("VRAM arbiter: cross-process GPU sharing via %s (doorbell %s)\n",
|
||||
lock_path.c_str(), vram_doorbell_dir.c_str());
|
||||
} else {
|
||||
SRV_WRN("VRAM arbiter: cannot open %s, running VRAM-only sleep without cross-process lock\n", lock_path.c_str());
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
// Acquire the VRAM token BEFORE uploading this model's weights, so any model currently resident
|
||||
// on the GPU releases first and the load uploads into free VRAM instead of racing it (which
|
||||
// could OOM, e.g. the 4B task model holding VRAM while a large model loads). Blocks until free.
|
||||
// Called from load_model() right before common_init_from_params().
|
||||
void vram_acquire_for_load() {
|
||||
vram_arena_open();
|
||||
if (!vram_only) {
|
||||
return;
|
||||
}
|
||||
#if !defined(_WIN32)
|
||||
if (vram_flock) {
|
||||
vram_ring_doorbell(); // nudge whoever is resident to release
|
||||
flock(vram_lock_fd, LOCK_EX); // block until the GPU is free
|
||||
}
|
||||
#endif
|
||||
vram_cold = false; // we hold the token; weights will be resident after the upload
|
||||
}
|
||||
|
||||
// Start the doorbell warden. Called from init() after the (coordinated) load. The model stays
|
||||
// warm (holding the token acquired in vram_acquire_for_load) and serves its first request
|
||||
// without a re-warm; the warden releases it when another model rings the doorbell.
|
||||
void vram_share_init() {
|
||||
vram_arena_open();
|
||||
if (!vram_only) {
|
||||
return;
|
||||
}
|
||||
vram_cold = false; // resident and holding the token after the coordinated load
|
||||
#if !defined(_WIN32)
|
||||
if (vram_flock && vram_inotify_fd < 0) {
|
||||
vram_inotify_fd = inotify_init1(IN_NONBLOCK);
|
||||
if (vram_inotify_fd >= 0) {
|
||||
inotify_add_watch(vram_inotify_fd, vram_doorbell_dir.c_str(), IN_CLOSE_WRITE);
|
||||
vram_warden_run = true;
|
||||
vram_warden = std::thread([this]{ vram_warden_loop(); });
|
||||
}
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
// warden thread: block on the doorbell; when another process rings (wants the token) and we
|
||||
// currently hold it, ask the loop to yield (release) at its next idle point. Only touches the
|
||||
// thread-safe queue - never the GPU - so it cannot race a decode.
|
||||
void vram_warden_loop() {
|
||||
#if !defined(_WIN32)
|
||||
char buf[4096];
|
||||
while (vram_warden_run.load()) {
|
||||
struct pollfd pfd { vram_inotify_fd, POLLIN, 0 };
|
||||
int pr = poll(&pfd, 1, 500); // 500ms so we periodically re-check the run flag
|
||||
if (pr <= 0) {
|
||||
continue;
|
||||
}
|
||||
ssize_t n = read(vram_inotify_fd, buf, sizeof(buf));
|
||||
if (n <= 0) {
|
||||
continue;
|
||||
}
|
||||
bool foreign_ring = false;
|
||||
for (char * p = buf; p < buf + n; ) {
|
||||
struct inotify_event * ev = (struct inotify_event *) p;
|
||||
if (ev->len > 0 && vram_pid_str != ev->name) {
|
||||
foreign_ring = true; // someone else wants the GPU
|
||||
}
|
||||
p += sizeof(struct inotify_event) + ev->len;
|
||||
}
|
||||
if (foreign_ring && !vram_cold) {
|
||||
queue_tasks.request_yield();
|
||||
}
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
// ring the doorbell so the current token holder releases promptly
|
||||
void vram_ring_doorbell() {
|
||||
#if !defined(_WIN32)
|
||||
if (!vram_flock || vram_doorbell_dir.empty()) {
|
||||
return;
|
||||
}
|
||||
const std::string f = vram_doorbell_dir + "/" + vram_pid_str;
|
||||
int fd = open(f.c_str(), O_WRONLY | O_CREAT | O_TRUNC, 0666);
|
||||
if (fd >= 0) {
|
||||
ssize_t w = write(fd, "1", 1);
|
||||
(void) w;
|
||||
close(fd);
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
// acquire the VRAM token (blocking) and bring weights back to the device. Called on the loop
|
||||
// thread right before a decode, so it never races compute.
|
||||
void vram_ensure_warm() {
|
||||
if (!vram_only || !vram_cold) {
|
||||
return;
|
||||
}
|
||||
#if !defined(_WIN32)
|
||||
if (vram_flock) {
|
||||
vram_ring_doorbell(); // nudge the current holder to release
|
||||
flock(vram_lock_fd, LOCK_EX); // blocks until the current holder goes cold
|
||||
}
|
||||
#endif
|
||||
llama_context_restore_device(ctx_tgt);
|
||||
if (ctx_dft != nullptr) {
|
||||
llama_context_restore_device(ctx_dft);
|
||||
}
|
||||
vram_cold = false;
|
||||
}
|
||||
|
||||
void vram_share_shutdown() {
|
||||
#if !defined(_WIN32)
|
||||
vram_warden_run = false;
|
||||
if (vram_warden.joinable()) {
|
||||
vram_warden.join();
|
||||
}
|
||||
if (vram_inotify_fd >= 0) { close(vram_inotify_fd); vram_inotify_fd = -1; }
|
||||
if (vram_lock_fd >= 0) { flock(vram_lock_fd, LOCK_UN); close(vram_lock_fd); vram_lock_fd = -1; }
|
||||
#endif
|
||||
}
|
||||
|
||||
// release the device weights (keeping KV + host shadow) and drop the VRAM token so another
|
||||
// model can warm up. Called on the loop thread when the server goes idle.
|
||||
void vram_go_cold() {
|
||||
if (!vram_only || vram_cold) {
|
||||
return;
|
||||
}
|
||||
llama_context_release_device(ctx_tgt, vram_evict_kv);
|
||||
if (ctx_dft != nullptr) {
|
||||
llama_context_release_device(ctx_dft, vram_evict_kv);
|
||||
}
|
||||
#if !defined(_WIN32)
|
||||
if (vram_flock) {
|
||||
flock(vram_lock_fd, LOCK_UN);
|
||||
}
|
||||
#endif
|
||||
vram_cold = true;
|
||||
}
|
||||
|
||||
common_speculative_init_result_ptr spec_init;
|
||||
|
||||
common_context_seq_rm_type ctx_tgt_seq_rm_type = COMMON_CONTEXT_SEQ_RM_TYPE_NO;
|
||||
@@ -951,6 +1147,25 @@ private:
|
||||
|
||||
void handle_sleeping_state(bool new_state) {
|
||||
GGML_ASSERT(sleeping != new_state);
|
||||
|
||||
// Lightweight VRAM-only sleep: instead of a full unload/reload (which frees RAM and the
|
||||
// KV cache and pays a full reload on wake), just release the model's device (VRAM) weight
|
||||
// buffers to a host shadow (keeping the context, KV cache and host weights) and drop the
|
||||
// shared VRAM token. Waking is handled lazily by vram_ensure_warm() right before the next
|
||||
// decode (which re-acquires the token first), so a single GPU is time-shared between models
|
||||
// with no re-prefill.
|
||||
if (vram_only && ctx_tgt != nullptr) {
|
||||
if (new_state) {
|
||||
SRV_INF("%s", "server entering sleeping state (VRAM-only: releasing device weights, keeping KV cache)\n");
|
||||
vram_go_cold();
|
||||
} else {
|
||||
SRV_INF("%s", "server exiting sleeping state (VRAM-only: weights restored on next decode)\n");
|
||||
// token is acquired and weights restored by vram_ensure_warm() before decode
|
||||
}
|
||||
sleeping = new_state;
|
||||
return;
|
||||
}
|
||||
|
||||
if (new_state) {
|
||||
SRV_INF("%s", "server is entering sleeping state\n");
|
||||
destroy();
|
||||
@@ -1142,6 +1357,11 @@ private:
|
||||
params_base.load_progress_callback_user_data = &load_progress_text;
|
||||
}
|
||||
|
||||
// VRAM arbiter: acquire the GPU token before uploading weights, so any resident model (e.g.
|
||||
// the warm 4B task model) releases first and this load uploads into free VRAM instead of
|
||||
// racing it and OOM-ing. No-op unless LLAMA_SLEEP_VRAM_ONLY is set.
|
||||
vram_acquire_for_load();
|
||||
|
||||
llama_init = common_init_from_params(params_base);
|
||||
|
||||
model_tgt = llama_init->model();
|
||||
@@ -1409,6 +1629,11 @@ private:
|
||||
handle_sleeping_state(sleeping);
|
||||
});
|
||||
|
||||
// VRAM arbiter: start the doorbell warden. The load was coordinated (vram_acquire_for_load
|
||||
// grabbed the token before uploading), so we stay warm holding the token and serve the first
|
||||
// request without a re-warm; the warden releases us when another model rings the doorbell.
|
||||
vram_share_init();
|
||||
|
||||
metrics.init();
|
||||
|
||||
if (params_base.cache_idle_slots) {
|
||||
@@ -3589,6 +3814,10 @@ private:
|
||||
n_empty_consecutive = 0;
|
||||
}
|
||||
|
||||
// VRAM arbiter: make sure our weights are resident (acquiring the shared GPU token first)
|
||||
// before we decode. Runs on the loop thread, so it never races an in-flight decode.
|
||||
vram_ensure_warm();
|
||||
|
||||
const int ret = llama_decode(ctx_tgt, batch_view);
|
||||
|
||||
metrics.on_decoded(slots);
|
||||
@@ -4010,8 +4239,12 @@ struct server_res_generator : server_res_spipe {
|
||||
server_response_reader rd;
|
||||
server_res_generator(server_queue & queue_tasks, server_response & queue_results, int sleep_idle_seconds, bool bypass_sleep = false)
|
||||
: rd(queue_tasks, queue_results, HTTP_POLLING_SECONDS) {
|
||||
// fast path in case sleeping is disabled
|
||||
bypass_sleep |= sleep_idle_seconds < 0;
|
||||
// fast path in case sleeping is disabled. Note: the VRAM arbiter (LLAMA_SLEEP_VRAM_ONLY)
|
||||
// can put the server to sleep via the cross-process doorbell even when idle-sleep is
|
||||
// disabled (sleep_idle_seconds < 0), so in that case requests must still wake it - otherwise
|
||||
// a doorbell-slept server would hang, never returning from a request.
|
||||
static const bool vram_arbiter = getenv("LLAMA_SLEEP_VRAM_ONLY") != nullptr;
|
||||
bypass_sleep |= (sleep_idle_seconds < 0 && !vram_arbiter);
|
||||
if (!bypass_sleep) {
|
||||
queue_tasks.wait_until_no_sleep();
|
||||
}
|
||||
|
||||
@@ -116,6 +116,12 @@ void server_queue::wait_until_no_sleep() {
|
||||
}
|
||||
}
|
||||
|
||||
void server_queue::request_yield() {
|
||||
std::unique_lock<std::mutex> lock(mutex_tasks);
|
||||
yield_requested = true;
|
||||
condition_tasks.notify_all();
|
||||
}
|
||||
|
||||
void server_queue::terminate() {
|
||||
std::unique_lock<std::mutex> lock(mutex_tasks);
|
||||
running = false;
|
||||
@@ -129,6 +135,9 @@ void server_queue::start_loop(int64_t idle_sleep_ms) {
|
||||
constexpr auto max_wait_time = std::chrono::seconds(1);
|
||||
auto should_sleep = [&]() -> bool {
|
||||
// caller must hold mutex_tasks
|
||||
if (yield_requested) {
|
||||
return true; // another process rang the VRAM doorbell - release now
|
||||
}
|
||||
if (idle_sleep_ms < 0) {
|
||||
return false;
|
||||
}
|
||||
@@ -178,6 +187,7 @@ void server_queue::start_loop(int64_t idle_sleep_ms) {
|
||||
if (should_sleep()) {
|
||||
QUE_INF("%s", "entering sleeping state\n");
|
||||
sleeping = true;
|
||||
yield_requested = false; // consumed
|
||||
callback_sleeping_state(true);
|
||||
req_stop_sleeping = false;
|
||||
// wait until we are requested to exit sleeping state
|
||||
@@ -195,13 +205,17 @@ void server_queue::start_loop(int64_t idle_sleep_ms) {
|
||||
condition_tasks.notify_all(); // notify wait_until_no_sleep()
|
||||
break; // process new tasks
|
||||
} else {
|
||||
// wait for new tasks or timeout for checking sleeping condition
|
||||
// wait for new tasks, a VRAM yield request, or timeout for checking sleeping condition
|
||||
bool res = condition_tasks.wait_for(lock, max_wait_time, [&]{
|
||||
return (!queue_tasks.empty() || !running);
|
||||
return (!queue_tasks.empty() || !running || yield_requested);
|
||||
});
|
||||
if (res) {
|
||||
if (res && !queue_tasks.empty()) {
|
||||
break; // new task arrived or terminate
|
||||
}
|
||||
if (!running) {
|
||||
break;
|
||||
}
|
||||
// otherwise (timeout or yield request), loop again to re-check should_sleep
|
||||
// otherwise, loop again to check sleeping condition
|
||||
}
|
||||
}
|
||||
|
||||
@@ -16,6 +16,7 @@ private:
|
||||
bool running = false;
|
||||
bool sleeping = false;
|
||||
bool req_stop_sleeping = false;
|
||||
bool yield_requested = false; // set by request_yield() when another process wants the VRAM token
|
||||
int64_t time_last_task = 0;
|
||||
|
||||
// queues
|
||||
@@ -51,6 +52,11 @@ public:
|
||||
// returns immediately if not sleeping
|
||||
void wait_until_no_sleep();
|
||||
|
||||
// request that the loop go to sleep (release VRAM) as soon as it is idle - called from the
|
||||
// VRAM-arbiter warden thread when another process rings the doorbell for the GPU token.
|
||||
// Thread-safe; wakes the loop so it releases promptly instead of at the next idle poll.
|
||||
void request_yield();
|
||||
|
||||
bool is_sleeping() {
|
||||
std::unique_lock<std::mutex> lock(mutex_tasks);
|
||||
return sleeping;
|
||||
|
||||
Reference in New Issue
Block a user