From f715d6dc157125f3c6f8a90326605b88eeb22baf Mon Sep 17 00:00:00 2001 From: Elnur Abdullaev Date: Wed, 8 Apr 2026 11:31:30 +0200 Subject: [PATCH 1/3] implement key structs and logic for expert cache --- ggml/include/ggml-backend.h | 14 + ggml/src/ggml-backend.cpp | 562 ++++++++++++++++++++++++++++++++++-- tests/CMakeLists.txt | 1 + tests/test-expert-cache.cpp | 241 ++++++++++++++++ 4 files changed, 800 insertions(+), 18 deletions(-) create mode 100644 tests/test-expert-cache.cpp diff --git a/ggml/include/ggml-backend.h b/ggml/include/ggml-backend.h index 9fd3f7f32a0..aefca0b3af6 100644 --- a/ggml/include/ggml-backend.h +++ b/ggml/include/ggml-backend.h @@ -340,6 +340,20 @@ extern "C" { // Set a callback to be called for each resulting node during graph compute GGML_API void ggml_backend_sched_set_eval_callback(ggml_backend_sched_t sched, ggml_backend_sched_eval_callback callback, void * user_data); + // Expert cache for MoE weight offloading: enable/disable and query stats + GGML_API void ggml_backend_sched_set_expert_cache( + ggml_backend_sched_t sched, int32_t n_slots); + + GGML_API void ggml_backend_sched_get_expert_cache_stats( + ggml_backend_sched_t sched, + int64_t * n_hits, + int64_t * n_misses, + int64_t * n_fate_hits, + int64_t * bytes_saved, + int64_t * bytes_copied); + + GGML_API void ggml_backend_sched_reset_expert_cache_stats(ggml_backend_sched_t sched); + // // Utils // diff --git a/ggml/src/ggml-backend.cpp b/ggml/src/ggml-backend.cpp index 22c656996cc..b418956e6bd 100644 --- a/ggml/src/ggml-backend.cpp +++ b/ggml/src/ggml-backend.cpp @@ -682,6 +682,8 @@ struct ggml_backend_sched_split { struct ggml_cgraph graph; }; +struct ggml_backend_sched_expert_cache; + struct ggml_backend_sched { bool is_reset; // true if the scheduler has been reset since the last graph split bool is_alloc; @@ -736,8 +738,147 @@ struct ggml_backend_sched { int debug_realloc; int debug_graph_size; int debug_prev_graph_size; + + struct ggml_backend_sched_expert_cache * expert_cache = nullptr; +}; + +// Expert cache for MoE weight offloading. +// Tracks which expert rows are already present in the GPU input_cpy tensor +// and skips redundant CPU->GPU transfers during graph reuse. +struct ggml_expert_cache_entry { + struct ggml_tensor * cpu_tensor; // key: original CPU weight tensor pointer + void * gpu_data; // last-seen input_cpy->data (staleness check) + ggml_bitset_t * populated; // [n_expert] bitset: which rows are current (n_slots_alloc==0 only) + ggml_bitset_t * fate_mask; // [n_expert] bitset: rows loaded by FATE this call (n_slots_alloc==0 only) + int32_t * fate_ids; // [n_expert] array: last token's routing + int32_t n_fate_ids; // number of valid fate IDs + int64_t n_expert; // total number of experts + size_t expert_size; // bytes per expert row + // N-slot LFRU fields (only used when n_slots_alloc > 0): + // LFRU eviction score: freq / age where age = lru_counter - lru_clock[slot] + 1 + // Evict slot with minimum score (least frequently, most stale). Empty slots score 0. + int32_t n_slots_alloc; // number of GPU slots (0 = full-copy dedup mode) + int32_t * slot_to_expert; // [n_slots_alloc] slot_id → expert_id (-1 = empty slot) + int32_t * expert_to_slot; // [n_expert] expert_id → slot_id (-1 = not cached) + int64_t * lru_clock; // [n_slots_alloc] last-access timestamp per slot + int32_t * slot_freq; // [n_slots_alloc] access frequency per slot (LFRU) + int64_t lru_counter; // monotonic counter, incremented on each cache access + int64_t last_decay_counter; // lru_counter at last frequency decay + // Write-back skip optimization: avoid redundant GPU ids_tensor write when mapping is unchanged. + bool slot_ids_dirty; // true if any slot→expert mapping changed since last write + struct ggml_tensor * last_ids_tensor; // ids_tensor pointer we last wrote slot_ids into + int32_t * prefill_freq; // [n_expert] frequency count from prefill routing (for warmup) }; +struct ggml_backend_sched_expert_cache { + struct ggml_expert_cache_entry * entries; + int n_entries; + int max_entries; + bool enabled; + int32_t n_slots; // 0 = dedup-only mode, >0 = LRU N-slot mode (VRAM savings) + int64_t n_hits; + int64_t n_misses; + int64_t n_fate_hits; + int64_t bytes_saved; + int64_t bytes_copied; +}; + +static struct ggml_expert_cache_entry * expert_cache_find_or_create( + struct ggml_backend_sched_expert_cache * cache, + struct ggml_tensor * cpu_tensor, + int64_t n_expert, + size_t expert_size) { + for (int i = 0; i < cache->n_entries; i++) { + if (cache->entries[i].cpu_tensor == cpu_tensor) { + return &cache->entries[i]; + } + } + if (cache->n_entries >= cache->max_entries) { + cache->max_entries = cache->max_entries > 0 ? cache->max_entries * 2 : 64; + cache->entries = (struct ggml_expert_cache_entry *)realloc( + cache->entries, cache->max_entries * sizeof(struct ggml_expert_cache_entry)); + if (!cache->entries) { + GGML_ABORT("expert cache: allocation failed"); + } + } + struct ggml_expert_cache_entry * entry = &cache->entries[cache->n_entries]; + memset(entry, 0, sizeof(*entry)); + + ggml_bitset_t * populated = (ggml_bitset_t *)calloc(ggml_bitset_size(n_expert), sizeof(ggml_bitset_t)); + ggml_bitset_t * fate_mask = (ggml_bitset_t *)calloc(ggml_bitset_size(n_expert), sizeof(ggml_bitset_t)); + int32_t * fate_ids = (int32_t *)calloc(n_expert, sizeof(int32_t)); + int32_t * prefill_freq = (int32_t *)calloc(n_expert, sizeof(int32_t)); + + if (!populated || !fate_mask || !fate_ids || !prefill_freq) { + free(populated); + free(fate_mask); + free(fate_ids); + free(prefill_freq); + GGML_ABORT("expert cache: allocation failed"); + } + + entry->cpu_tensor = cpu_tensor; + entry->gpu_data = nullptr; + entry->n_expert = n_expert; + entry->expert_size = expert_size; + entry->populated = populated; + entry->fate_mask = fate_mask; + entry->fate_ids = fate_ids; + entry->n_fate_ids = 0; + entry->prefill_freq = prefill_freq; + + // N-slot LRU initialization + const int32_t n_slots_alloc = cache->n_slots; + entry->n_slots_alloc = n_slots_alloc; + entry->lru_counter = 0; + entry->last_decay_counter = 0; + entry->slot_ids_dirty = true; // first token must always write + entry->last_ids_tensor = nullptr; + if (n_slots_alloc > 0) { + entry->slot_to_expert = (int32_t *)malloc((size_t)n_slots_alloc * sizeof(int32_t)); + entry->expert_to_slot = (int32_t *)malloc((size_t)n_expert * sizeof(int32_t)); + entry->lru_clock = (int64_t *)calloc((size_t)n_slots_alloc, sizeof(int64_t)); + entry->slot_freq = (int32_t *)calloc((size_t)n_slots_alloc, sizeof(int32_t)); + if (!entry->slot_to_expert || !entry->expert_to_slot || !entry->lru_clock || !entry->slot_freq) { + free(entry->slot_to_expert); + free(entry->expert_to_slot); + free(entry->lru_clock); + free(entry->slot_freq); + GGML_ABORT("expert cache: slot allocation failed"); + } + for (int32_t i = 0; i < n_slots_alloc; i++) entry->slot_to_expert[i] = -1; + for (int64_t i = 0; i < n_expert; i++) entry->expert_to_slot[i] = -1; + } else { + entry->slot_to_expert = nullptr; + entry->expert_to_slot = nullptr; + entry->lru_clock = nullptr; + entry->slot_freq = nullptr; + } + + cache->n_entries++; + return entry; +} + +// LFRU eviction: find slot with minimum freq/age score (cross-multiply to avoid float). +// Empty slots (freq=0) are preferred. Ties broken by first-encountered. +static int32_t expert_cache_find_evict_slot( + const struct ggml_expert_cache_entry * entry) { + int32_t evict_slot = 0; + int64_t best_freq = entry->slot_freq[0]; + int64_t best_age = entry->lru_counter - entry->lru_clock[0] + 1; + for (int32_t s = 1; s < entry->n_slots_alloc; s++) { + const int64_t freq = entry->slot_freq[s]; + const int64_t age = entry->lru_counter - entry->lru_clock[s] + 1; + // evict s if freq/age < best_freq/best_age (cross-multiply) + if (freq * best_age < best_freq * age) { + evict_slot = s; + best_freq = freq; + best_age = age; + } + } + return evict_slot; +} + #define hash_id(tensor) ggml_hash_find_or_insert(&sched->hash_set, tensor) #define tensor_backend_id(tensor) sched->hv_tensor_backend_ids[hash_id(tensor)] #define tensor_id_copy(id, backend_id, copy_id) sched->hv_tensor_copies[(id) * sched->n_backends * sched->n_copies + (backend_id) * sched->n_copies + (copy_id)] @@ -1263,7 +1404,24 @@ void ggml_backend_sched_split_graph(ggml_backend_sched_t sched, struct ggml_cgra if (tensor_id_copy(src_id, cur_backend_id, 0) == NULL) { ggml_backend_t backend = sched->backends[cur_backend_id]; for (int c = 0; c < sched->n_copies; c++) { - struct ggml_tensor * tensor_copy = ggml_dup_tensor_layout(sched->ctx, src); + // For MoE expert weights in N-slot mode, allocate only N GPU slots + // instead of the full expert dimension — this is where VRAM is saved. + const int32_t ec_n_slots = (sched->expert_cache && + sched->expert_cache->n_slots > 0) ? sched->expert_cache->n_slots : 0; + // N-slot tensors are only safe for unquantized types. + // Quantized types (MXFP4, Q4, etc.) use CUDA MMQ kernels that read + // a few bytes past the last expert row; N-slot tensors have no room + // for this and trigger illegal memory access. Fall back to full-size + // tensor + dedup-only path for quantized types. + const bool use_slot_tensor = (ec_n_slots > 0 && + node->op == GGML_OP_MUL_MAT_ID && j == 0 && + ggml_n_dims(src) == 3 && + ggml_backend_buffer_is_host(src->buffer) && + ggml_backend_buffer_get_usage(src->buffer) == GGML_BACKEND_BUFFER_USAGE_WEIGHTS && + !ggml_is_quantized(src->type)); + struct ggml_tensor * tensor_copy = use_slot_tensor + ? ggml_new_tensor_3d(sched->ctx, src->type, src->ne[0], src->ne[1], ec_n_slots) + : ggml_dup_tensor_layout(sched->ctx, src); ggml_format_name(tensor_copy, "%s#%s#%d", ggml_backend_name(backend), src->name, c); if (sched->n_copies > 1) { ggml_set_input(tensor_copy); @@ -1447,7 +1605,8 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s struct ggml_backend_sched_split * splits = sched->splits; ggml_tensor * prev_ids_tensor = nullptr; - std::vector ids; + std::vector ids; // original expert IDs (never modified after download) + std::vector slot_ids; // remapped slot IDs for N-slot LFRU write-back std::vector used_ids; for (int split_id = 0; split_id < sched->n_splits; split_id++) { @@ -1539,29 +1698,314 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s expert_size_copy + padding_end); }; - int id = 0; - while (!ggml_bitset_get(used_ids.data(), id)) { - id++; - } - int32_t first_id = id; - int32_t last_id = first_id; + if (sched->expert_cache && sched->expert_cache->enabled) { + struct ggml_expert_cache_entry * entry = expert_cache_find_or_create( + sched->expert_cache, input, n_expert, expert_size); + + // Staleness check: if the GPU copy buffer changed, all slots are invalid. + if (entry->gpu_data != input_cpy->data) { + memset(entry->populated, 0, ggml_bitset_size(n_expert) * sizeof(ggml_bitset_t)); + entry->gpu_data = input_cpy->data; + entry->slot_ids_dirty = true; + entry->last_ids_tensor = nullptr; + if (entry->n_slots_alloc > 0) { + for (int32_t i = 0; i < entry->n_slots_alloc; i++) entry->slot_to_expert[i] = -1; + for (int64_t i = 0; i < n_expert; i++) entry->expert_to_slot[i] = -1; + memset(entry->lru_clock, 0, (size_t)entry->n_slots_alloc * sizeof(int64_t)); + memset(entry->slot_freq, 0, (size_t)entry->n_slots_alloc * sizeof(int32_t)); + entry->lru_counter = 0; + entry->last_decay_counter = 0; + } + } - for (++id; id < n_expert; ++id) { - if (!ggml_bitset_get(used_ids.data(), id)) { - continue; + // Route to LFRU when GPU tensor is N-slot sized (both prefill and decode). + // Overflow (n_unique > n_slots) is handled gracefully with forced eviction. + const bool gpu_is_slot_tensor = (entry->n_slots_alloc > 0 && + input_cpy->ne[2] == (int64_t)entry->n_slots_alloc); + + if (gpu_is_slot_tensor) { + // ========================================================= + // N-slot LFRU mode: fixed GPU slots, LFRU eviction, id remap. + // LFRU eviction score: freq / age (higher = more valuable). + // Implemented via cross-multiply to avoid floating point. + // ========================================================= + + // One-shot warmup: on first decode token, pre-populate slots from prefill frequencies. + if (entry->lru_counter == 0 && entry->prefill_freq != nullptr) { + // Find top-N experts by prefill frequency. + // Simple O(n_expert * n_slots) selection — runs once. + std::vector> freq_sorted; // (freq, expert_id) + for (int32_t eid = 0; eid < (int32_t)n_expert; eid++) { + if (entry->prefill_freq[eid] > 0) { + freq_sorted.push_back({entry->prefill_freq[eid], eid}); + } + } + std::sort(freq_sorted.begin(), freq_sorted.end(), std::greater<>()); + + const int32_t n_warmup = std::min((int32_t)freq_sorted.size(), entry->n_slots_alloc); + for (int32_t i = 0; i < n_warmup; i++) { + const int32_t eid = freq_sorted[i].second; + const int32_t slot = i; // slots 0..n_warmup-1 + { + ggml_backend_tensor_set_async(split_backend, input_cpy, + (const uint8_t *)input->data + (size_t)eid * expert_size, + (size_t)slot * expert_size, expert_size); + } + entry->slot_to_expert[slot] = eid; + entry->expert_to_slot[eid] = slot; + entry->lru_clock[slot] = 1; // initial timestamp + entry->slot_freq[slot] = entry->prefill_freq[eid]; // carry over frequency + sched->expert_cache->bytes_copied += (int64_t)expert_size; + } + entry->lru_counter = 1; + entry->slot_ids_dirty = true; + + // Clear prefill_freq to prevent re-warmup on cache invalidation. + memset(entry->prefill_freq, 0, (size_t)n_expert * sizeof(int32_t)); + } + + // P1: Frequency decay to prevent old high-freq experts from dominating. + // Decay every n_slots_alloc accesses (~once per token in steady state). + if (entry->lru_counter - entry->last_decay_counter >= entry->n_slots_alloc) { + for (int32_t s = 0; s < entry->n_slots_alloc; s++) { + entry->slot_freq[s] = (entry->slot_freq[s] + 1) >> 1; // halve, floor 1 + } + entry->last_decay_counter = entry->lru_counter; + } + + // Count unique experts; fall back gracefully on overflow. + int32_t n_unique = 0; + for (int32_t eid = 0; eid < (int32_t)n_expert; eid++) { + if (ggml_bitset_get(used_ids.data(), eid)) n_unique++; + } + + if (n_unique > entry->n_slots_alloc) { + // Prefill overflow: n_unique experts > n_slots GPU slots. + // Run full LFRU eviction with slot_ids remap. + // Silent-corruption fix: all experts get a slot assignment; no slot-0 fallback. + GGML_LOG_DEBUG("expert-cache: %d unique > %d slots (overflow, using forced eviction)\n", + n_unique, entry->n_slots_alloc); + for (int32_t eid = 0; eid < (int32_t)n_expert; eid++) { + if (!ggml_bitset_get(used_ids.data(), eid)) continue; + const int32_t evict_slot = expert_cache_find_evict_slot(entry); + const int32_t old_eid = entry->slot_to_expert[evict_slot]; + if (old_eid >= 0) entry->expert_to_slot[old_eid] = -1; + { + ggml_backend_tensor_set_async(split_backend, input_cpy, + (const uint8_t *)input->data + (size_t)eid * expert_size, + (size_t)evict_slot * expert_size, expert_size); + } + entry->slot_to_expert[evict_slot] = eid; + entry->expert_to_slot[eid] = evict_slot; + entry->lru_clock[evict_slot] = ++entry->lru_counter; + entry->slot_freq[evict_slot] = 1; + sched->expert_cache->n_misses++; + sched->expert_cache->bytes_copied += (int64_t)expert_size; + } + // Remap ids → slot_ids and write to GPU. + // Overflow: n_unique > n_slots, so some early-loaded experts may have + // been evicted by later ones. Map evicted experts to slot 0 (known quality + // degradation; only affects overflow prefill batches, not decode). + entry->slot_ids_dirty = true; + slot_ids.resize(ids.size()); + for (int64_t i1 = 0; i1 < ids_tensor->ne[1]; i1++) { + for (int64_t i0 = 0; i0 < ids_tensor->ne[0]; i0++) { + const size_t pos = i1 * (ids_tensor->nb[1] / sizeof(int32_t)) + + i0 * (ids_tensor->nb[0] / sizeof(int32_t)); + const int32_t sl = entry->expert_to_slot[ids[pos]]; + slot_ids[pos] = (sl >= 0) ? sl : 0; + } + } + ggml_backend_tensor_set_async(ids_backend, ids_tensor, + slot_ids.data(), 0, ggml_nbytes(ids_tensor)); + entry->last_ids_tensor = ids_tensor; + entry->slot_ids_dirty = false; + } else { + // FATE pre-load: load predicted experts before processing hits/misses + for (int32_t i = 0; i < entry->n_fate_ids; i++) { + int32_t fid = entry->fate_ids[i]; + if (fid >= 0 && fid < (int32_t)n_expert && + ggml_bitset_get(used_ids.data(), fid) && + entry->expert_to_slot[fid] < 0) { + // FATE miss: evict and load + const int32_t evict_slot = expert_cache_find_evict_slot(entry); + const int32_t old_eid = entry->slot_to_expert[evict_slot]; + if (old_eid >= 0) { + entry->expert_to_slot[old_eid] = -1; + } + { + ggml_backend_tensor_set_async(split_backend, input_cpy, + (const uint8_t *)input->data + (size_t)fid * expert_size, + (size_t)evict_slot * expert_size, expert_size); + } + entry->slot_to_expert[evict_slot] = fid; + entry->expert_to_slot[fid] = evict_slot; + entry->lru_clock[evict_slot] = ++entry->lru_counter; + entry->slot_freq[evict_slot] = 1; + entry->slot_ids_dirty = true; + sched->expert_cache->n_fate_hits++; + sched->expert_cache->bytes_copied += (int64_t)expert_size; + } + } + // Normal LFRU path: per-expert hit/miss with eviction. + for (int32_t eid = 0; eid < (int32_t)n_expert; eid++) { + if (!ggml_bitset_get(used_ids.data(), eid)) continue; + + const int32_t slot = entry->expert_to_slot[eid]; + if (slot >= 0) { + // Cache hit: update LFRU frequency and recency. + entry->lru_clock[slot] = ++entry->lru_counter; + entry->slot_freq[slot]++; + sched->expert_cache->n_hits++; + sched->expert_cache->bytes_saved += (int64_t)expert_size; + } else { + // Cache miss: LFRU eviction. + const int32_t evict_slot = expert_cache_find_evict_slot(entry); + const int32_t old_eid = entry->slot_to_expert[evict_slot]; + if (old_eid >= 0) { + entry->expert_to_slot[old_eid] = -1; + } + + // Copy new expert into eviction slot. + { + ggml_backend_tensor_set_async(split_backend, input_cpy, + (const uint8_t *)input->data + (size_t)eid * expert_size, + (size_t)evict_slot * expert_size, expert_size); + } + + entry->slot_to_expert[evict_slot] = eid; + entry->expert_to_slot[eid] = evict_slot; + entry->lru_clock[evict_slot] = ++entry->lru_counter; + entry->slot_freq[evict_slot] = 1; + + // Mapping changed — must write slot_ids back to GPU. + entry->slot_ids_dirty = true; + + sched->expert_cache->n_misses++; + sched->expert_cache->bytes_copied += (int64_t)expert_size; + } + } + + // P0: Skip slot_ids write-back if no mapping changed AND this + // entry last wrote to the same ids_tensor (handles w13/w2 sharing). + // When w13 writes to ids_tensor, w2 still needs to write its own + // slot mapping even if its own mapping is unchanged, because w13 + // overwrote the shared tensor. last_ids_tensor tracks this. + if (entry->slot_ids_dirty || ids_tensor != entry->last_ids_tensor) { + slot_ids.resize(ids.size()); + for (int64_t i1 = 0; i1 < ids_tensor->ne[1]; i1++) { + for (int64_t i0 = 0; i0 < ids_tensor->ne[0]; i0++) { + const size_t pos = i1 * (ids_tensor->nb[1] / sizeof(int32_t)) + + i0 * (ids_tensor->nb[0] / sizeof(int32_t)); + const int32_t sl = entry->expert_to_slot[ids[pos]]; + GGML_ASSERT(sl >= 0 && "LFRU: expert not loaded"); + slot_ids[pos] = sl; + } + } + ggml_backend_tensor_set_async(ids_backend, ids_tensor, + slot_ids.data(), 0, ggml_nbytes(ids_tensor)); + entry->last_ids_tensor = ids_tensor; + entry->slot_ids_dirty = false; + } + + // P5: FATE recording using used_ids bitset (dedup, no duplicates). + entry->n_fate_ids = 0; + for (int32_t eid = 0; eid < (int32_t)n_expert; eid++) { + if (ggml_bitset_get(used_ids.data(), eid)) { + entry->fate_ids[entry->n_fate_ids++] = eid; + } + } + } + + } else { + // ========================================================= + // Dedup-only mode (n_slots_alloc == 0): no VRAM saving. + // Keeps all experts in GPU but skips re-copying within a batch. + // FATE pre-populates predicted experts before the next token. + // ========================================================= + + // Phase 1: FATE pre-populate — copy predicted experts that are + // needed this token but not yet on GPU (from previous token's routing). + memset(entry->fate_mask, 0, ggml_bitset_size(n_expert) * sizeof(ggml_bitset_t)); + for (int32_t i = 0; i < entry->n_fate_ids; i++) { + int32_t fid = entry->fate_ids[i]; + if (fid >= 0 && fid < (int32_t)n_expert && + ggml_bitset_get(used_ids.data(), fid) && + !ggml_bitset_get(entry->populated, fid)) { + copy_experts(fid, fid); + ggml_bitset_set(entry->populated, fid); + ggml_bitset_set(entry->fate_mask, fid); + sched->expert_cache->n_fate_hits++; + sched->expert_cache->bytes_copied += (int64_t)expert_size; + } + } + + // Phase 2: copy remaining needed experts; skip already-populated. + int32_t first_id = -1, last_id = -1; + for (int32_t eid = 0; eid < (int32_t)n_expert; eid++) { + if (!ggml_bitset_get(used_ids.data(), eid)) { + continue; + } + if (ggml_bitset_get(entry->populated, eid)) { + if (first_id >= 0) { copy_experts(first_id, last_id); first_id = -1; } + if (!ggml_bitset_get(entry->fate_mask, eid)) { + sched->expert_cache->n_hits++; + sched->expert_cache->bytes_saved += (int64_t)expert_size; + } + continue; + } + sched->expert_cache->n_misses++; + sched->expert_cache->bytes_copied += (int64_t)expert_size; + ggml_bitset_set(entry->populated, eid); + if (first_id < 0) { first_id = last_id = eid; } + else if (eid == last_id + 1) { last_id = eid; } + else { copy_experts(first_id, last_id); first_id = last_id = eid; } + } + if (first_id >= 0) { copy_experts(first_id, last_id); } + + // Accumulate prefill routing frequencies for LFRU warmup. + if (entry->n_slots_alloc > 0) { + for (int32_t eid = 0; eid < (int32_t)n_expert; eid++) { + if (ggml_bitset_get(used_ids.data(), eid)) { + entry->prefill_freq[eid]++; + } + } + } + + // Phase 3: record routing for FATE on next token (deduped via used_ids bitset). + entry->n_fate_ids = 0; + for (int32_t eid = 0; eid < (int32_t)n_expert; eid++) { + if (ggml_bitset_get(used_ids.data(), eid)) { + entry->fate_ids[entry->n_fate_ids++] = eid; + } + } + } + } else { + // cache disabled: copy all needed expert rows in one pass + int id = 0; + while (!ggml_bitset_get(used_ids.data(), id)) { + id++; } + int32_t first_id = id; + int32_t last_id = first_id; + + for (++id; id < n_expert; ++id) { + if (!ggml_bitset_get(used_ids.data(), id)) { + continue; + } - if (id == last_id + 1) { + if (id == last_id + 1) { + last_id = id; + continue; + } + + copy_experts(first_id, last_id); + + first_id = id; last_id = id; - continue; } - copy_experts(first_id, last_id); - - first_id = id; - last_id = id; } - copy_experts(first_id, last_id); } else { // try async copy, but if not possible, we can still use a sync copy without synchronizing the dst backend, since we handle the synchronization here with multiple copies and events // TODO: add public function to facilitate this, since applications do not have direct access to the backend interface @@ -1719,6 +2163,20 @@ void ggml_backend_sched_free(ggml_backend_sched_t sched) { free(sched->context_buffer); free(sched->graph.nodes); free(sched->graph.leafs); + if (sched->expert_cache) { + for (int i = 0; i < sched->expert_cache->n_entries; i++) { + free(sched->expert_cache->entries[i].populated); + free(sched->expert_cache->entries[i].fate_mask); + free(sched->expert_cache->entries[i].fate_ids); + free(sched->expert_cache->entries[i].slot_to_expert); + free(sched->expert_cache->entries[i].expert_to_slot); + free(sched->expert_cache->entries[i].lru_clock); + free(sched->expert_cache->entries[i].slot_freq); + free(sched->expert_cache->entries[i].prefill_freq); + } + free(sched->expert_cache->entries); + free(sched->expert_cache); + } free(sched); } @@ -1732,6 +2190,27 @@ void ggml_backend_sched_reset(ggml_backend_sched_t sched) { sched->is_reset = true; } sched->is_alloc = false; + // Invalidate expert cache populated state: GPU buffer data may be stale after reset. + // Entries (cpu_tensor keys) are preserved to avoid re-discovering tensors next graph. + if (sched->expert_cache) { + for (int i = 0; i < sched->expert_cache->n_entries; i++) { + struct ggml_expert_cache_entry * e = &sched->expert_cache->entries[i]; + memset(e->populated, 0, ggml_bitset_size(e->n_expert) * sizeof(ggml_bitset_t)); + memset(e->fate_mask, 0, ggml_bitset_size(e->n_expert) * sizeof(ggml_bitset_t)); + e->gpu_data = nullptr; + if (e->n_slots_alloc > 0) { + for (int32_t j = 0; j < e->n_slots_alloc; j++) e->slot_to_expert[j] = -1; + for (int64_t j = 0; j < e->n_expert; j++) e->expert_to_slot[j] = -1; + memset(e->lru_clock, 0, (size_t)e->n_slots_alloc * sizeof(int64_t)); + memset(e->slot_freq, 0, (size_t)e->n_slots_alloc * sizeof(int32_t)); + e->lru_counter = 0; + e->last_decay_counter = 0; + } + if (e->prefill_freq) { + memset(e->prefill_freq, 0, (size_t)e->n_expert * sizeof(int32_t)); + } + } + } } void ggml_backend_sched_reserve_size(ggml_backend_sched_t sched, struct ggml_cgraph * measure_graph, size_t * sizes) { @@ -1824,6 +2303,53 @@ void ggml_backend_sched_set_eval_callback(ggml_backend_sched_t sched, ggml_backe sched->callback_eval_user_data = user_data; } +void ggml_backend_sched_set_expert_cache(ggml_backend_sched_t sched, int32_t n_slots) { + GGML_ASSERT(sched); + const bool enabled = (n_slots != 0); + if (enabled && !sched->expert_cache) { + sched->expert_cache = (struct ggml_backend_sched_expert_cache *)calloc(1, sizeof(struct ggml_backend_sched_expert_cache)); + GGML_ASSERT(sched->expert_cache && "expert cache: allocation failed"); + } + if (sched->expert_cache) { + sched->expert_cache->enabled = enabled; + sched->expert_cache->n_slots = (n_slots > 0) ? n_slots : 0; + } +} + +void ggml_backend_sched_get_expert_cache_stats( + ggml_backend_sched_t sched, + int64_t * n_hits, + int64_t * n_misses, + int64_t * n_fate_hits, + int64_t * bytes_saved, + int64_t * bytes_copied) { + GGML_ASSERT(sched); + int64_t h = 0, m = 0, f = 0, s = 0, c = 0; + if (sched->expert_cache) { + h = sched->expert_cache->n_hits; + m = sched->expert_cache->n_misses; + f = sched->expert_cache->n_fate_hits; + s = sched->expert_cache->bytes_saved; + c = sched->expert_cache->bytes_copied; + } + if (n_hits) *n_hits = h; + if (n_misses) *n_misses = m; + if (n_fate_hits) *n_fate_hits = f; + if (bytes_saved) *bytes_saved = s; + if (bytes_copied) *bytes_copied = c; +} + +void ggml_backend_sched_reset_expert_cache_stats(ggml_backend_sched_t sched) { + GGML_ASSERT(sched); + if (sched->expert_cache) { + sched->expert_cache->n_hits = 0; + sched->expert_cache->n_misses = 0; + sched->expert_cache->n_fate_hits = 0; + sched->expert_cache->bytes_saved = 0; + sched->expert_cache->bytes_copied = 0; + } +} + int ggml_backend_sched_get_n_splits(ggml_backend_sched_t sched) { GGML_ASSERT(sched); return sched->n_splits; diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 5e87c8b34e1..da63b9f86b1 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -209,6 +209,7 @@ llama_build_and_test( peg-parser/tests.h ) llama_build_and_test(test-regex-partial.cpp) +llama_build_and_test(test-expert-cache.cpp) if (NOT ${CMAKE_SYSTEM_PROCESSOR} MATCHES "s390x") set(MODEL_NAME "tinyllamas/stories15M-q4_0.gguf") diff --git a/tests/test-expert-cache.cpp b/tests/test-expert-cache.cpp new file mode 100644 index 00000000000..ac111fa5046 --- /dev/null +++ b/tests/test-expert-cache.cpp @@ -0,0 +1,241 @@ +// SPDX-License-Identifier: MIT +#include "ggml.h" +#include "ggml-alloc.h" +#include "ggml-backend.h" +#include +#include +#include + +static void test_expert_cache_enable_disable() { + ggml_backend_t cpu = ggml_backend_init_by_type(GGML_BACKEND_DEVICE_TYPE_CPU, nullptr); + GGML_ASSERT(cpu); + + ggml_backend_t backends[] = { cpu }; + ggml_backend_sched_t sched = ggml_backend_sched_new( + backends, nullptr, 1, 512, false, false); + GGML_ASSERT(sched); + + // Stats should be zero before enabling + int64_t hits = -1, misses = -1, fate = -1, saved = -1, copied = -1; + ggml_backend_sched_get_expert_cache_stats(sched, &hits, &misses, &fate, &saved, &copied); + GGML_ASSERT(hits == 0); + GGML_ASSERT(misses == 0); + GGML_ASSERT(fate == 0); + GGML_ASSERT(saved == 0); + GGML_ASSERT(copied == 0); + + ggml_backend_sched_set_expert_cache(sched, true); + ggml_backend_sched_set_expert_cache(sched, false); + + ggml_backend_sched_get_expert_cache_stats(sched, &hits, &misses, &fate, &saved, &copied); + GGML_ASSERT(hits == 0); + GGML_ASSERT(misses == 0); + + ggml_backend_sched_free(sched); + ggml_backend_free(cpu); + + printf(" PASS: test_expert_cache_enable_disable\n"); +} + +static void test_expert_cache_reset_survives() { + ggml_backend_t cpu = ggml_backend_init_by_type(GGML_BACKEND_DEVICE_TYPE_CPU, nullptr); + ggml_backend_t backends[] = { cpu }; + ggml_backend_sched_t sched = ggml_backend_sched_new( + backends, nullptr, 1, 512, false, false); + + ggml_backend_sched_set_expert_cache(sched, true); + + for (int i = 0; i < 10; i++) { + ggml_backend_sched_reset(sched); + } + + int64_t h, m, f, s, c; + ggml_backend_sched_get_expert_cache_stats(sched, &h, &m, &f, &s, &c); + GGML_ASSERT(h == 0 && m == 0 && f == 0 && s == 0 && c == 0); + + ggml_backend_sched_free(sched); + ggml_backend_free(cpu); + + printf(" PASS: test_expert_cache_reset_survives\n"); +} + +static void test_expert_cache_toggle() { + ggml_backend_t cpu = ggml_backend_init_by_type(GGML_BACKEND_DEVICE_TYPE_CPU, nullptr); + ggml_backend_t backends[] = { cpu }; + ggml_backend_sched_t sched = ggml_backend_sched_new( + backends, nullptr, 1, 512, false, false); + + ggml_backend_sched_set_expert_cache(sched, false); + ggml_backend_sched_set_expert_cache(sched, true); + ggml_backend_sched_set_expert_cache(sched, false); + + int64_t h, m, f, s, c; + ggml_backend_sched_get_expert_cache_stats(sched, &h, &m, &f, &s, &c); + GGML_ASSERT(h == 0); + + ggml_backend_sched_free(sched); + ggml_backend_free(cpu); + + printf(" PASS: test_expert_cache_toggle\n"); +} + +static void test_expert_cache_null_stats() { + ggml_backend_t cpu = ggml_backend_init_by_type(GGML_BACKEND_DEVICE_TYPE_CPU, nullptr); + ggml_backend_t backends[] = { cpu }; + ggml_backend_sched_t sched = ggml_backend_sched_new( + backends, nullptr, 1, 512, false, false); + + ggml_backend_sched_set_expert_cache(sched, true); + ggml_backend_sched_get_expert_cache_stats(sched, nullptr, nullptr, nullptr, nullptr, nullptr); + + int64_t hits; + ggml_backend_sched_get_expert_cache_stats(sched, &hits, nullptr, nullptr, nullptr, nullptr); + GGML_ASSERT(hits == 0); + + ggml_backend_sched_free(sched); + ggml_backend_free(cpu); + + printf(" PASS: test_expert_cache_null_stats\n"); +} + +static void test_expert_cache_disabled_no_crash() { + ggml_backend_t cpu = ggml_backend_init_by_type(GGML_BACKEND_DEVICE_TYPE_CPU, nullptr); + ggml_backend_t backends[] = { cpu }; + ggml_backend_sched_t sched = ggml_backend_sched_new( + backends, nullptr, 1, 512, false, false); + + // Never enable -- sched must work normally + ggml_backend_sched_reset(sched); + + int64_t h, m, f, s, c; + ggml_backend_sched_get_expert_cache_stats(sched, &h, &m, &f, &s, &c); + GGML_ASSERT(h == 0 && m == 0); + + ggml_backend_sched_free(sched); + ggml_backend_free(cpu); + + printf(" PASS: test_expert_cache_disabled_no_crash\n"); +} + +// Integration smoke test: exercises the actual cache copy logic via a minimal +// MUL_MAT_ID compute graph with CPU expert weights and a GPU copy target. +// On the second token with identical expert IDs the cache must record hits and +// bytes_saved > 0. The test is skipped when no GPU backend is present; full +// behavioural correctness is validated in Task 5 (perplexity-parity run). +static void test_expert_cache_hit_on_second_token() { + ggml_backend_t gpu = ggml_backend_init_by_type(GGML_BACKEND_DEVICE_TYPE_GPU, nullptr); + if (!gpu) { + printf(" SKIP: test_expert_cache_hit_on_second_token (no GPU backend)\n"); + return; + } + + const int n_expert = 8; + const int n_ff = 16; + const int n_embd = 8; + const int n_expert_used = 2; + // n_tokens must be >= GGML_OP_OFFLOAD_MIN_BATCH (default 32) for MUL_MAT_ID + // to be offloaded to GPU, which is required to trigger the cross-backend copy path + const int n_tokens = 32; + + ggml_backend_t cpu = ggml_backend_init_by_type(GGML_BACKEND_DEVICE_TYPE_CPU, nullptr); + GGML_ASSERT(cpu); + + ggml_backend_t backends[2] = { gpu, cpu }; + ggml_backend_buffer_type_t buft_gpu = ggml_backend_get_default_buffer_type(gpu); + ggml_backend_buffer_type_t buft_cpu = ggml_backend_get_default_buffer_type(cpu); + ggml_backend_buffer_type_t bufts[2] = { buft_gpu, buft_cpu }; + + // op_offload=true is required for weights on the last (CPU) backend to be + // offloaded to GPU, which triggers the cross-backend expert copy path + ggml_backend_sched_t sched = ggml_backend_sched_new(backends, bufts, 2, 8192, false, true); + GGML_ASSERT(sched); + ggml_backend_sched_set_expert_cache(sched, true); + + struct ggml_init_params iparams = { 1024*1024, nullptr, true }; + struct ggml_context * ctx = ggml_init(iparams); + GGML_ASSERT(ctx); + + // Expert weight tensor: [n_embd, n_ff, n_expert] on CPU (host weights) + struct ggml_tensor * experts = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_embd, n_ff, n_expert); + ggml_set_name(experts, "ffn_up_exps"); + + // Allocate CPU buffer sized correctly for this tensor (accounts for alignment) + ggml_backend_buffer_t cpu_buf = ggml_backend_buft_alloc_buffer(buft_cpu, + ggml_backend_buft_get_alloc_size(buft_cpu, experts)); + GGML_ASSERT(cpu_buf); + ggml_backend_buffer_set_usage(cpu_buf, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); + + struct ggml_tallocr tallocr = ggml_tallocr_new(cpu_buf); + ggml_tallocr_alloc(&tallocr, experts); + + // Zero-fill the weight data + memset(experts->data, 0, ggml_nbytes(experts)); + + // Input: [n_embd, n_expert_used, n_tokens] + struct ggml_tensor * input = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_embd, n_expert_used, n_tokens); + + // IDs: [n_expert_used, n_tokens] i32 — selects experts 0 and 1 + struct ggml_tensor * ids = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, n_expert_used, n_tokens); + + struct ggml_tensor * out = ggml_mul_mat_id(ctx, experts, input, ids); + + struct ggml_cgraph * graph = ggml_new_graph(ctx); + ggml_build_forward_expand(graph, out); + + ggml_backend_sched_alloc_graph(sched, graph); + + // IDs: experts 0 and 1 interleaved across all tokens (same on every run) + const int n_id_vals = n_expert_used * n_tokens; + int32_t * id_vals = (int32_t *)malloc(n_id_vals * sizeof(int32_t)); + GGML_ASSERT(id_vals); + for (int i = 0; i < n_id_vals; i++) { id_vals[i] = i % n_expert_used; } + ggml_backend_tensor_set(ids, id_vals, 0, n_id_vals * sizeof(int32_t)); + free(id_vals); + + // Zero-fill input activations + const size_t input_bytes = (size_t)n_embd * n_expert_used * n_tokens * sizeof(float); + float * indata = (float *)calloc(n_embd * n_expert_used * n_tokens, sizeof(float)); + GGML_ASSERT(indata); + ggml_backend_tensor_set(input, indata, 0, input_bytes); + free(indata); + + // Token 1: cold cache — expect misses >= n_expert_used + ggml_backend_sched_graph_compute(sched, graph); + + int64_t h1, m1, f1, s1, c1; + ggml_backend_sched_get_expert_cache_stats(sched, &h1, &m1, &f1, &s1, &c1); + printf(" After token 1: hits=%lld misses=%lld fate=%lld saved=%lld\n", + (long long)h1, (long long)m1, (long long)f1, (long long)s1); + GGML_ASSERT(m1 >= n_expert_used); + + // Token 2: same IDs — cache should hit and save bytes (ids already set above) + ggml_backend_sched_graph_compute(sched, graph); + + int64_t h2, m2, f2, s2, c2; + ggml_backend_sched_get_expert_cache_stats(sched, &h2, &m2, &f2, &s2, &c2); + printf(" After token 2: hits=%lld misses=%lld fate=%lld saved=%lld\n", + (long long)h2, (long long)m2, (long long)f2, (long long)s2); + // Either direct hits or FATE pre-population should fire on the second token + GGML_ASSERT(h2 > h1 || f2 > f1); + GGML_ASSERT(s2 > 0); + + ggml_free(ctx); + ggml_backend_buffer_free(cpu_buf); + ggml_backend_sched_free(sched); + ggml_backend_free(gpu); + ggml_backend_free(cpu); + + printf(" PASS: test_expert_cache_hit_on_second_token\n"); +} + +int main() { + printf("test-expert-cache:\n"); + test_expert_cache_enable_disable(); + test_expert_cache_reset_survives(); + test_expert_cache_toggle(); + test_expert_cache_null_stats(); + test_expert_cache_disabled_no_crash(); + test_expert_cache_hit_on_second_token(); + printf("ALL TESTS PASSED\n"); + return 0; +} From ff16f5a3837bf8c3c766769a2284eba7c54c86f4 Mon Sep 17 00:00:00 2001 From: Elnur Abdullaev Date: Wed, 8 Apr 2026 11:31:42 +0200 Subject: [PATCH 2/3] add auto sizing and resolve --n-cpu-moe conflicts --- common/arg.cpp | 12 +++++ common/common.cpp | 1 + common/common.h | 1 + include/llama.h | 13 +++++ src/llama-context.cpp | 113 +++++++++++++++++++++++++++++++++++++++++- src/llama-cparams.h | 1 + 6 files changed, 139 insertions(+), 2 deletions(-) diff --git a/common/arg.cpp b/common/arg.cpp index 2e0f46db519..f303992af54 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -2328,6 +2328,18 @@ common_params_context common_params_parser_init(common_params & params, llama_ex } } ).set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_N_CPU_MOE_DRAFT")); + add_opt(common_arg( + {"--expert-cache-slots"}, "N", + "N-slot LRU expert cache: N GPU slots per MoE layer (saves VRAM vs full GPU copy); use with --cpu-moe. " + "N=0 disables (default), N=-1 dedup-only, N=-2 or 'auto' fills available VRAM", + [](common_params & params, const std::string & value) { + if (value == "auto") { + params.expert_cache_n_slots = -2; + } else { + params.expert_cache_n_slots = std::stoi(value); + } + } + ).set_env("LLAMA_ARG_EXPERT_CACHE_SLOTS")); GGML_ASSERT(params.n_gpu_layers < 0); // string_format would need to be extended for a default >= 0 add_opt(common_arg( {"-ngl", "--gpu-layers", "--n-gpu-layers"}, "N", diff --git a/common/common.cpp b/common/common.cpp index 16f78debd02..72270b0a1b3 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1474,6 +1474,7 @@ struct llama_context_params common_context_params_to_llama(const common_params & cparams.offload_kqv = !params.no_kv_offload; cparams.no_perf = params.no_perf; cparams.op_offload = !params.no_op_offload; + cparams.expert_cache_n_slots = params.expert_cache_n_slots; cparams.swa_full = params.swa_full; cparams.kv_unified = params.kv_unified; diff --git a/common/common.h b/common/common.h index 020b6a721ff..1d2513ccdbb 100644 --- a/common/common.h +++ b/common/common.h @@ -543,6 +543,7 @@ struct common_params { bool no_op_offload = false; // globally disable offload host tensor operations to device bool no_extra_bufts = false; // disable extra buffer types (used for weight repacking) bool no_host = false; // bypass host buffer allowing extra buffers to be used + int32_t expert_cache_n_slots = 0; // N-slot LRU expert cache: 0=off, -1=dedup-only, N>0=N GPU slots per layer (--expert-cache-slots) bool single_turn = false; // single turn chat conversation diff --git a/include/llama.h b/include/llama.h index bf2bff8dac6..f7edcde0afc 100644 --- a/include/llama.h +++ b/include/llama.h @@ -367,6 +367,7 @@ extern "C" { bool offload_kqv; // offload the KQV ops (including the KV cache) to GPU bool no_perf; // measure performance timings bool op_offload; // offload host tensor operations to device + int32_t expert_cache_n_slots; // N-slot LRU expert cache: 0=off, -1=dedup-only, -2=auto (fill available VRAM), N>0=N GPU slots per layer bool swa_full; // use full-size SWA cache (https://github.com/ggml-org/llama.cpp/pull/13194#issuecomment-2868343055) // NOTE: setting to false when n_seq_max > 1 can cause bad performance in some cases // ref: https://github.com/ggml-org/llama.cpp/pull/13845#issuecomment-2924800573 @@ -991,6 +992,18 @@ extern "C" { // and is not necessary to call it explicitly in most cases LLAMA_API void llama_synchronize(struct llama_context * ctx); + // Expert cache statistics for MoE weight offloading (--expert-cache-slots / --n-cpu-moe). + // All output pointers are optional (may be NULL). Reset counters with llama_expert_cache_stats_reset(). + LLAMA_API void llama_expert_cache_stats( + const struct llama_context * ctx, + int64_t * n_hits, + int64_t * n_misses, + int64_t * n_fate_hits, + int64_t * bytes_saved, + int64_t * bytes_copied); + + LLAMA_API void llama_expert_cache_stats_reset(struct llama_context * ctx); + // Token logits obtained from the last call to llama_decode() // The logits for which llama_batch.logits[i] != 0 are stored contiguously // in the order they have appeared in the batch. diff --git a/src/llama-context.cpp b/src/llama-context.cpp index a808e3e4542..4a2500850be 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -160,8 +160,9 @@ llama_context::llama_context( cparams.n_ubatch = std::min(cparams.n_batch, params.n_ubatch == 0 ? params.n_batch : params.n_ubatch); - cparams.op_offload = params.op_offload; - cparams.kv_unified = params.kv_unified; + cparams.op_offload = params.op_offload; + cparams.expert_cache_n_slots = params.expert_cache_n_slots; + cparams.kv_unified = params.kv_unified; // initialized later cparams.pipeline_parallel = false; @@ -365,6 +366,15 @@ llama_context::llama_context( } llama_context::~llama_context() { + if (cparams.expert_cache_n_slots != 0 && sched) { + int64_t hits, misses, fate, saved, copied; + ggml_backend_sched_get_expert_cache_stats(sched.get(), &hits, &misses, &fate, &saved, &copied); + if (hits + misses > 0) { + LLAMA_LOG_INFO("%s: expert-cache: hits=%" PRId64 " (%.1f%%) misses=%" PRId64 " fate=%" PRId64 " saved=%.1fMB copied=%.1fMB\n", + __func__, hits, 100.0 * hits / (hits + misses), misses, fate, + saved / 1048576.0, copied / 1048576.0); + } + } if (!model.hparams.no_alloc) { for (size_t i = 0; i < backend_ptrs.size(); ++i) { ggml_backend_t backend = backend_ptrs[i]; @@ -409,6 +419,65 @@ void llama_context::sched_reserve() { sched.reset(ggml_backend_sched_new(backend_ptrs.data(), backend_buft.data(), backend_ptrs.size(), max_nodes, cparams.pipeline_parallel, cparams.op_offload)); + // Auto-size expert cache: resolve -2 to a concrete slot count from available GPU VRAM. + if (cparams.expert_cache_n_slots == -2) { + const auto & hparams = model.hparams; + if (hparams.n_expert > 0) { + // Find first GPU backend to query free memory. + bool resolved = false; + for (const auto & backend : backends) { + auto * dev = ggml_backend_get_device(backend.get()); + if (!dev || ggml_backend_dev_type(dev) == GGML_BACKEND_DEVICE_TYPE_CPU) { + continue; + } + size_t free_mem = 0, total_mem = 0; + ggml_backend_dev_memory(dev, &free_mem, &total_mem); + + // Estimate per-expert size: gate_up fused = n_embd * 2 * n_ff_exp * type_size. + // Use bf16 (2 bytes) as conservative fallback. + const uint32_t n_ff_exp = hparams.n_ff_exp > 0 + ? hparams.n_ff_exp + : (hparams.n_expert > 0 ? hparams.n_ff(0) / hparams.n_expert : 0); + if (n_ff_exp == 0) { break; } + + const size_t expert_size_est = (size_t)hparams.n_embd * 2 * n_ff_exp * 2; // 2 bytes (bf16) + + // Count MoE layers (layers with experts). + uint32_t n_moe_layers = 0; + for (uint32_t il = 0; il < hparams.n_layer; il++) { + if (hparams.n_expert > 0) { n_moe_layers++; } + } + if (n_moe_layers == 0) { break; } + + // Each MoE layer has ~2 expert weight tensors (gate_up + down). + const size_t per_slot_cost = (size_t)n_moe_layers * 2 * expert_size_est; + + // Use 80% of free VRAM, bounded by [4, n_expert]. + const size_t budget = (size_t)(free_mem * 0.8); + int32_t auto_slots = per_slot_cost > 0 ? (int32_t)(budget / per_slot_cost) : 0; + auto_slots = std::max(auto_slots, (int32_t)4); + auto_slots = std::min(auto_slots, (int32_t)hparams.n_expert); + + cparams.expert_cache_n_slots = auto_slots; + LLAMA_LOG_INFO("%s: expert-cache auto-sized to %d slots (%.0f MiB free, %.1f MiB/slot)\n", + __func__, auto_slots, free_mem / 1048576.0, per_slot_cost / 1048576.0); + resolved = true; + break; + } + if (!resolved) { + LLAMA_LOG_WARN("%s: expert-cache auto: no GPU found or invalid model params, disabling\n", __func__); + cparams.expert_cache_n_slots = 0; + } + } else { + LLAMA_LOG_WARN("%s: expert-cache auto: model has no experts, disabling\n", __func__); + cparams.expert_cache_n_slots = 0; + } + } + + if (cparams.expert_cache_n_slots != 0) { + ggml_backend_sched_set_expert_cache(sched.get(), cparams.expert_cache_n_slots); + } + llama_memory_context_ptr mctx; if (memory) { LLAMA_LOG_DEBUG("%s: reserving full memory module\n", __func__); @@ -562,6 +631,9 @@ void llama_context::sched_reserve() { LLAMA_LOG_WARN("%s: compute buffer allocation failed, retrying without pipeline parallelism\n", __func__); cparams.pipeline_parallel = false; sched.reset(ggml_backend_sched_new(backend_ptrs.data(), backend_buft.data(), backend_ptrs.size(), max_nodes, false, cparams.op_offload)); + if (cparams.expert_cache_n_slots != 0) { + ggml_backend_sched_set_expert_cache(sched.get(), cparams.expert_cache_n_slots); + } gf = graph_reserve(n_tokens, n_seqs, n_tokens, mctx.get()); } if (!gf) { @@ -2217,6 +2289,27 @@ llm_graph_cb llama_context::graph_get_cb() const { } } } + + // Expert cache: when expert weights are on CPU (via --n-cpu-moe / tensor_buft_overrides) + // the scheduler normally assigns MUL_MAT_ID to CPU, preventing cross-backend copies + // and leaving the expert weight cache inactive. Force MUL_MAT_ID to the layer's GPU + // backend so the scheduler creates a CPU→GPU copy that the cache can intercept. + if (cparams.expert_cache_n_slots != 0 && + cur->op == GGML_OP_MUL_MAT_ID && + il >= 0 && + cur->src[0] != nullptr && + cur->src[0]->buffer != nullptr && + ggml_backend_buffer_is_host(cur->src[0]->buffer) && + ggml_backend_buffer_get_usage(cur->src[0]->buffer) == GGML_BACKEND_BUFFER_USAGE_WEIGHTS) { + const auto & dev_layer = model.dev_layer(il); + for (const auto & backend : backends) { + if (ggml_backend_get_device(backend.get()) == dev_layer && + ggml_backend_supports_op(backend.get(), cur)) { + ggml_backend_sched_set_tensor_backend(sched.get(), cur, backend.get()); + break; + } + } + } }; } @@ -2910,6 +3003,7 @@ llama_context_params llama_context_default_params() { /*.offload_kqv =*/ true, /*.no_perf =*/ true, /*.op_offload =*/ true, + /*.expert_cache_n_slots =*/ 0, /*.swa_full =*/ true, /*.kv_unified =*/ false, /*.sampler =*/ nullptr, @@ -3068,6 +3162,21 @@ void llama_synchronize(llama_context * ctx) { ctx->synchronize(); } +void llama_expert_cache_stats( + const llama_context * ctx, + int64_t * n_hits, + int64_t * n_misses, + int64_t * n_fate_hits, + int64_t * bytes_saved, + int64_t * bytes_copied) { + ggml_backend_sched_get_expert_cache_stats( + ctx->get_sched(), n_hits, n_misses, n_fate_hits, bytes_saved, bytes_copied); +} + +void llama_expert_cache_stats_reset(llama_context * ctx) { + ggml_backend_sched_reset_expert_cache_stats(ctx->get_sched()); +} + float * llama_get_logits(llama_context * ctx) { ctx->synchronize(); diff --git a/src/llama-cparams.h b/src/llama-cparams.h index 9d359474132..c5aa84a9f6d 100644 --- a/src/llama-cparams.h +++ b/src/llama-cparams.h @@ -37,6 +37,7 @@ struct llama_cparams { bool no_perf; bool warmup; bool op_offload; + int32_t expert_cache_n_slots; bool kv_unified; bool pipeline_parallel; From 0e5d061c8675ea46b42d5f02a5f51c290cb9500b Mon Sep 17 00:00:00 2001 From: Elnur Abdullaev Date: Wed, 8 Apr 2026 11:31:54 +0200 Subject: [PATCH 3/3] add support for llama-bench --- tools/llama-bench/llama-bench.cpp | 60 ++++++++++++++++++++++++++++--- 1 file changed, 56 insertions(+), 4 deletions(-) diff --git a/tools/llama-bench/llama-bench.cpp b/tools/llama-bench/llama-bench.cpp index 0b395b460e8..0374dd99ec7 100644 --- a/tools/llama-bench/llama-bench.cpp +++ b/tools/llama-bench/llama-bench.cpp @@ -341,6 +341,7 @@ struct cmd_params { std::vector use_direct_io; std::vector embeddings; std::vector no_op_offload; + std::vector expert_cache; std::vector no_host; std::vector fit_params_target; std::vector fit_params_min_ctx; @@ -385,6 +386,7 @@ static const cmd_params cmd_params_defaults = { /* use_direct_io */ { false }, /* embeddings */ { false }, /* no_op_offload */ { false }, + /* expert_cache */ { 0 }, /* no_host */ { false }, /* fit_params_target */ { 0 }, /* fit_params_min_ctx */ { 0 }, @@ -456,6 +458,7 @@ static void print_usage(int /* argc */, char ** argv) { printf(" -ot --override-tensor =;...\n"); printf(" (default: disabled)\n"); printf(" -nopo, --no-op-offload <0|1> (default: 0)\n"); + printf(" -ec, --expert-cache-slots (default: %s)\n", join(cmd_params_defaults.expert_cache, ",").c_str()); printf(" --no-host <0|1> (default: %s)\n", join(cmd_params_defaults.no_host, ",").c_str()); printf("\n"); printf( @@ -817,6 +820,18 @@ static cmd_params parse_cmd_params(int argc, char ** argv) { } auto p = string_split(argv[i], split_delim); params.no_op_offload.insert(params.no_op_offload.end(), p.begin(), p.end()); + } else if (arg == "-ec" || arg == "--expert-cache-slots") { + if (++i >= argc) { + invalid_param = true; + break; + } + std::string val = argv[i]; + if (val == "auto") { + params.expert_cache.push_back(-2); + } else { + auto p = string_split(val, split_delim); + params.expert_cache.insert(params.expert_cache.end(), p.begin(), p.end()); + } } else if (arg == "--no-host") { if (++i >= argc) { invalid_param = true; @@ -1087,6 +1102,9 @@ static cmd_params parse_cmd_params(int argc, char ** argv) { if (params.no_op_offload.empty()) { params.no_op_offload = cmd_params_defaults.no_op_offload; } + if (params.expert_cache.empty()) { + params.expert_cache = cmd_params_defaults.expert_cache; + } if (params.no_host.empty()) { params.no_host = cmd_params_defaults.no_host; } @@ -1138,6 +1156,7 @@ struct cmd_params_instance { bool use_direct_io; bool embeddings; bool no_op_offload; + int32_t expert_cache; bool no_host; size_t fit_target; uint32_t fit_min_ctx; @@ -1216,8 +1235,9 @@ struct cmd_params_instance { cparams.offload_kqv = !no_kv_offload; cparams.flash_attn_type = flash_attn ? LLAMA_FLASH_ATTN_TYPE_ENABLED : LLAMA_FLASH_ATTN_TYPE_DISABLED; cparams.embeddings = embeddings; - cparams.op_offload = !no_op_offload; - cparams.swa_full = false; + cparams.op_offload = !no_op_offload; + cparams.expert_cache_n_slots = expert_cache; + cparams.swa_full = false; return cparams; } @@ -1243,6 +1263,7 @@ static std::vector get_cmd_params_instances(const cmd_param for (const auto & noh : params.no_host) for (const auto & embd : params.embeddings) for (const auto & nopo : params.no_op_offload) + for (const auto & ec : params.expert_cache) for (const auto & nb : params.n_batch) for (const auto & nub : params.n_ubatch) for (const auto & tk : params.type_k) @@ -1284,6 +1305,7 @@ static std::vector get_cmd_params_instances(const cmd_param /* .use_direct_io= */ dio, /* .embeddings = */ embd, /* .no_op_offload= */ nopo, + /* .expert_cache = */ ec, /* .no_host = */ noh, /* .fit_target = */ fpt, /* .fit_min_ctx = */ fpc, @@ -1321,6 +1343,7 @@ static std::vector get_cmd_params_instances(const cmd_param /* .use_direct_io= */ dio, /* .embeddings = */ embd, /* .no_op_offload= */ nopo, + /* .expert_cache = */ ec, /* .no_host = */ noh, /* .fit_target = */ fpt, /* .fit_min_ctx = */ fpc, @@ -1358,6 +1381,7 @@ static std::vector get_cmd_params_instances(const cmd_param /* .use_direct_io= */ dio, /* .embeddings = */ embd, /* .no_op_offload= */ nopo, + /* .expert_cache = */ ec, /* .no_host = */ noh, /* .fit_target = */ fpt, /* .fit_min_ctx = */ fpc, @@ -1400,6 +1424,7 @@ struct test { bool use_direct_io; bool embeddings; bool no_op_offload; + int32_t expert_cache; bool no_host; size_t fit_target; uint32_t fit_min_ctx; @@ -1440,6 +1465,7 @@ struct test { use_direct_io = inst.use_direct_io; embeddings = inst.embeddings; no_op_offload = inst.no_op_offload; + expert_cache = inst.expert_cache; no_host = inst.no_host; fit_target = inst.fit_target; fit_min_ctx = inst.fit_min_ctx; @@ -1500,7 +1526,7 @@ struct test { "type_k", "type_v", "n_gpu_layers", "n_cpu_moe", "split_mode", "main_gpu", "no_kv_offload", "flash_attn", "devices", "tensor_split", "tensor_buft_overrides", "use_mmap", "use_direct_io", "embeddings", - "no_op_offload", "no_host", "fit_target", "fit_min_ctx", + "no_op_offload", "expert_cache", "no_host", "fit_target", "fit_min_ctx", "n_prompt", "n_gen", "n_depth", "test_time", "avg_ns", "stddev_ns", "avg_ts", "stddev_ts" }; @@ -1513,7 +1539,7 @@ struct test { if (field == "build_number" || field == "n_batch" || field == "n_ubatch" || field == "n_threads" || field == "poll" || field == "model_size" || field == "model_n_params" || field == "n_gpu_layers" || field == "main_gpu" || field == "n_prompt" || field == "n_gen" || field == "n_depth" || field == "avg_ns" || - field == "stddev_ns" || field == "no_op_offload" || field == "n_cpu_moe" || + field == "stddev_ns" || field == "no_op_offload" || field == "expert_cache" || field == "n_cpu_moe" || field == "fit_target" || field == "fit_min_ctx") { return INT; } @@ -1594,6 +1620,7 @@ struct test { std::to_string(use_direct_io), std::to_string(embeddings), std::to_string(no_op_offload), + std::to_string(expert_cache), std::to_string(no_host), std::to_string(fit_target), std::to_string(fit_min_ctx), @@ -1788,6 +1815,9 @@ struct markdown_printer : public printer { if (field == "no_op_offload") { return 4; } + if (field == "expert_cache") { + return 2; + } if (field == "no_host") { return 4; } @@ -1828,6 +1858,9 @@ struct markdown_printer : public printer { if (field == "no_op_offload") { return "nopo"; } + if (field == "expert_cache") { + return "ec"; + } if (field == "no_host") { return "noh"; } @@ -1921,6 +1954,9 @@ struct markdown_printer : public printer { if (params.no_op_offload.size() > 1 || params.no_op_offload != cmd_params_defaults.no_op_offload) { fields.emplace_back("no_op_offload"); } + if (params.expert_cache.size() > 1 || params.expert_cache != cmd_params_defaults.expert_cache) { + fields.emplace_back("expert_cache"); + } if (params.no_host.size() > 1 || params.no_host != cmd_params_defaults.no_host) { fields.emplace_back("no_host"); } @@ -2317,6 +2353,9 @@ int main(int argc, char ** argv) { } } + // Reset expert cache counters before the timed runs so stats reflect this test only. + llama_expert_cache_stats_reset(ctx); + for (int i = 0; i < params.reps; i++) { llama_memory_clear(llama_get_memory(ctx), false); @@ -2395,6 +2434,19 @@ int main(int argc, char ** argv) { fflush(p->fout); } + // Print expert cache stats if the cache is active. + if (t.expert_cache != 0) { + int64_t ec_hits = 0, ec_misses = 0, ec_fate = 0, ec_saved = 0, ec_copied = 0; + llama_expert_cache_stats(ctx, &ec_hits, &ec_misses, &ec_fate, &ec_saved, &ec_copied); + const double hit_pct = (ec_hits + ec_misses > 0) + ? 100.0 * ec_hits / (ec_hits + ec_misses) : 0.0; + fprintf(stderr, "expert-cache: hits=%" PRId64 " (%.1f%%) misses=%" PRId64 + " fate=%" PRId64 " saved=%.1fMiB copied=%.1fMiB\n", + ec_hits, hit_pct, ec_misses, ec_fate, + ec_saved / 1048576.0, ec_copied / 1048576.0); + fflush(stderr); + } + if (p_err) { p_err->print_test(t); fflush(p_err->fout);