feat(laguna): use KVFlash pool for MoE placement (#448)

Author: davide221

server/src/common/kvflash_placement.h

@@ -3,7 +3,7 @@
 // Any MoE / weight-offload backend that places experts against a VRAM budget
 // must decide how much KV to reserve.  Reserving for `max_ctx` forces experts
 // cold at high max_ctx even when KVFlash bounds the *resident* KV to a fixed
-// pool.  This helper centralises the rule so every backend (qwen35moe today,
+// pool.  This helper centralises the rule so every backend (qwen35moe, laguna,
 // DeepSeek-V4 / future MoE next) inherits the "pool bounds the expert-placement
 // cliff" win without re-deriving the byte math.
 #pragma once

server/src/laguna/laguna_backend.cpp

@@ -19,6 +19,7 @@
 #include "../common/moe_hybrid_types.h"
 #include "../common/moe_hybrid_types_impl.h"
 #include "../common/moe_hybrid_placement.h"
+#include "../common/kvflash_placement.h"
 #include "../common/moe_hybrid_ffn_eval.h"
 #include "../common/moe_hybrid_storage.h"
 #include "../common/moe_hybrid_routing_stats.h"
@@ -121,27 +122,42 @@ KvFlashConfig LagunaBackend::kvflash_config() const {
     return pc;
 }
 
-void LagunaBackend::kvflash_read_config() {
+void LagunaBackend::kvflash_resolve_drafter() {
     if (std::getenv("DFLASH_KVFLASH")) {
         kvflash_drafter_path_ = kvflash_find_drafter(args_.target_path.c_str());
     }
+}
+
+KvFlashAutoBudget LagunaBackend::make_kvflash_budget(int64_t gpu_free) const {
+    KvFlashAutoBudget b;
+    b.free_bytes      = gpu_free;
+    b.bytes_per_token = (int64_t)w_.n_layer * w_.n_head_kv * 2 *
+        (int64_t)ggml_row_size(args_.kv_type, w_.head_dim);
+    b.reserve_bytes   = (int64_t)(1.5 * 1073741824.0) +
+        (kvflash_drafter_path_.empty() ? 0 : (int64_t)(1.7 * 1073741824.0));
+    return b;
+}
+
+void LagunaBackend::kvflash_read_config() {
+    kvflash_resolve_drafter();
     // "auto" sizes from the GPU (weights resident, cache not yet allocated):
     // laguna pools ALL n_layer layers at the configured KV quant.
-    KvFlashAutoBudget kvf_budget;
-    {
-        size_t gpu_free = 0, gpu_total = 0;
-        if (ggml_backend_dev_t dev = ggml_backend_get_device(backend_)) {
-            ggml_backend_dev_memory(dev, &gpu_free, &gpu_total);
-        }
-        kvf_budget.free_bytes      = (int64_t)gpu_free;
-        kvf_budget.bytes_per_token = (int64_t)w_.n_layer * w_.n_head_kv * 2 *
-            (int64_t)ggml_row_size(args_.kv_type, w_.head_dim);
-        kvf_budget.reserve_bytes   = (int64_t)(1.5 * 1073741824.0) +
-            (kvflash_drafter_path_.empty() ? 0 : (int64_t)(1.7 * 1073741824.0));
+    size_t gpu_free = 0, gpu_total = 0;
+    if (ggml_backend_dev_t dev = ggml_backend_get_device(backend_)) {
+        ggml_backend_dev_memory(dev, &gpu_free, &gpu_total);
     }
+    KvFlashAutoBudget kvf_budget = make_kvflash_budget((int64_t)gpu_free);
     kvflash_tokens_ = kvflash_pool_from_env(args_.max_ctx, kvflash_config(),
-                                            !kvflash_drafter_path_.empty(),
+                                            kvflash_scorer_expected(),
                                             kvf_budget);
+    if (kvflash_tokens_ > 0 && placement_all_hot_full_kv_) {
+        std::printf("[laguna][kvflash] disabled: placement all-hot at max_ctx %d, "
+                    "pool not needed\n", args_.max_ctx);
+        std::fflush(stdout);
+        kvflash_tokens_ = 0;
+        kvflash_tau_ = 64;
+        kvflash_drafter_path_.clear();
+    }
     if (kvflash_tokens_ > 0) {
         const char * tau = std::getenv("DFLASH_KVFLASH_TAU");
         kvflash_tau_ = std::max(1, tau ? std::atoi(tau) : 64);
@@ -266,6 +282,7 @@ bool LagunaBackend::unpark(const std::string & what) {
         }
         cache_.kv_k_type = args_.kv_type;
         cache_.kv_v_type = args_.kv_type;
+        kvflash_read_config();
         if (!create_laguna_target_cache(w_, args_.max_ctx, backend_, cache_,
                                         kvflash_tokens_)) {
             std::fprintf(stderr, "[unpark] cache: %s\n", dflash27b_last_error());
@@ -1382,6 +1399,7 @@ static inline uint64_t elapsed_us(HybridClock::time_point t0, HybridClock::time_
 
 bool LagunaBackend::init_hybrid_mode() {
     const char * hotness_path = std::getenv("DFLASH_LAGUNA_HOTNESS");
+    placement_all_hot_full_kv_ = false;
 
     // Step 1: Load model WITHOUT expert data to GPU (partial load)
     TargetLoadPlan _hybrid_plan;
@@ -1454,13 +1472,29 @@ bool LagunaBackend::init_hybrid_mode() {
 
     const uint64_t kv_bytes_per_tok = (uint64_t)w_.n_layer * 2 *
         (uint64_t)w_.n_head_kv * (uint64_t)w_.head_dim * 2;
-    const uint64_t kv_total = kv_bytes_per_tok * (uint64_t)max_context;
 
     const uint64_t warm_cache_bytes = 200ULL * 1024 * 1024;
     const uint64_t safety_bytes = 512ULL * 1024 * 1024;
     const uint64_t core_bytes =
         moe_hybrid_core_bytes_from_memory("laguna", gpu_free, gpu_total);
 
+    kvflash_resolve_drafter();
+    const int kvf_pool = kvflash_pool_from_env(
+        max_context, kvflash_config(), kvflash_scorer_expected(),
+        make_kvflash_budget((int64_t)gpu_free));
+    const auto kvf_dec = dflash::common::kvflash_placement_decision(
+        kv_bytes_per_tok, max_context, kvf_pool,
+        gpu_total, core_bytes, total_expert_bytes,
+        warm_cache_bytes, safety_bytes, /*draft_bytes=*/0);
+    const uint64_t kv_total = kvf_dec.kv_total;
+    const int kv_ctx_log = kvf_dec.kv_ctx;
+    placement_all_hot_full_kv_ = kvf_dec.all_hot_full_kv;
+    if (kvf_dec.pool_reduced) {
+        std::printf("[laguna][kvflash] placement reserves pool KV (%d tokens, "
+                    "not max_ctx %d) -> experts stay hot\n", kvf_pool, max_context);
+        std::fflush(stdout);
+    }
+
     uint64_t expert_budget = 0;
     if (gpu_total > core_bytes + kv_total + warm_cache_bytes + safety_bytes) {
         expert_budget = gpu_total - core_bytes - kv_total - warm_cache_bytes - safety_bytes;
@@ -1511,7 +1545,7 @@ bool LagunaBackend::init_hybrid_mode() {
                 gpu_total / 1024.0 / 1024.0 / 1024.0,
                 core_bytes / 1024.0 / 1024.0 / 1024.0,
                 kv_total / 1024.0 / 1024.0 / 1024.0,
-                max_context,
+                kv_ctx_log,
                 warm_cache_bytes / 1024.0 / 1024.0,
                 safety_bytes / 1024.0 / 1024.0,
                 expert_budget / 1024.0 / 1024.0 / 1024.0,

server/src/laguna/laguna_backend.h

@@ -142,11 +142,19 @@ class LagunaBackend : public ModelBackend {
     int          kvflash_tau_    = 64;
     bool         kvflash_drafter_failed_ = false;
     bool kvflash_active() const { return kvflash_tokens_ > 0; }
+    // True iff hybrid placement showed all experts fit hot with the FULL
+    // max_ctx KV reservation. In that case KVFlash is redundant and can be
+    // disabled before cache creation.
+    bool placement_all_hot_full_kv_ = false;
     // Drafter rescore + repage every effective-tau generated tokens
     // (lazy-loads the drafter + cross-tokenizer scorer on first need).
     void kvflash_maybe_reselect(const std::vector<int32_t> & history, int generated);
     // Pager protections (SWA tail) shared by the floor and attach.
     KvFlashConfig kvflash_config() const;
+    // Shared pool sizing inputs for placement and runtime allocation.
+    void kvflash_resolve_drafter();
+    bool kvflash_scorer_expected() const { return !kvflash_drafter_path_.empty(); }
+    KvFlashAutoBudget make_kvflash_budget(int64_t gpu_free) const;
     // Read DFLASH_KVFLASH and round/clamp; call before cache creation.
     void kvflash_read_config();
     // Attach the pager to the freshly created cache (init / unpark).

server/test/test_kvflash_placement.cpp

@@ -80,6 +80,18 @@ int main() {
         expect(d.kv_ctx == 32768,  "C5: keeps full ctx");
     }
 
-    std::printf("PASS: kvflash placement decision (5 cases)\n");
+    // Case 6 — laguna/poolside-style no-draft placement: same pool rule, no
+    // drafter reserve. This keeps the shared helper covered for both current
+    // backends that use it.
+    {
+        auto d = kvflash_placement_decision(kv_per_tok, 131072, /*pool=*/16384,
+                                            gpu, core, experts, warm, safety,
+                                            /*draft_bytes=*/0);
+        expect(!d.all_hot_full_kv, "C6: full KV still forces experts cold");
+        expect(d.kv_ctx == 16384,  "C6: laguna reserves POOL ctx");
+        expect(d.pool_reduced,     "C6: laguna pool reduction engaged");
+    }
+
+    std::printf("PASS: kvflash placement decision (6 cases)\n");
     return 0;
 }