kv-cache : resolve the partial KV residency set once for the model

common/arg.cpp

@@ -2459,9 +2459,9 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
     ).set_env("LLAMA_ARG_LIVE_CONTEXT_WORKSPACE"));
     add_opt(common_arg(
         {"--kv-gpu-layers"}, "N",
-        string_format("with --no-kv-offload, keep the first N independently owned attention KV layers device-resident "
-                      "for standard and direct hybrid caches. Unsupported specialized caches ignore this option "
-                      "(default: %d)", params.kv_gpu_layers),
+        string_format("with --no-kv-offload, keep N independently owned attention KV layers device-resident. "
+                      "Layers are selected in turn from each owning device. Unsupported cache layouts keep "
+                      "fewer, or none (default: %d)", params.kv_gpu_layers),
         [](common_params & params, int value) {
             if (value < 0) {
                 throw std::invalid_argument("--kv-gpu-layers must not be negative");

include/llama.h

@@ -364,7 +364,7 @@ extern "C" {
         uint32_t n_rs_seq;              // number of recurrent-state snapshots per seq for rollback (0 = no rollback) [EXPERIMENTAL]
         uint32_t n_outputs_max;         // max outputs in a ubatch (0 = n_batch)
         uint32_t n_outputs_max_per_seq; // max outputs per sequence (0 = n_outputs_max)
-        uint32_t kv_gpu_layers;         // with offload_kqv=false, keep this many standard or direct hybrid attention KV layers on their assigned devices
+        uint32_t kv_gpu_layers;         // with offload_kqv=false, request up to this many independently owned attention KV layers on their assigned devices
         int32_t  n_threads;             // number of threads to use for generation
         int32_t  n_threads_batch;       // number of threads to use for batch processing
 

src/llama-context.cpp

@@ -435,12 +435,9 @@ llama_context::llama_context(
             cparams.live_context_workspace = false;
         }
 
+        // create_memory lowered kv_gpu_layers to the layers that became device-resident
         if (!cparams.offload_kqv && cparams.kv_gpu_layers > 0) {
-            if (memory && memory->get_supports_partial_kv()) {
-                cparams.offload_attn_compute = cparams.offload_attn_compute || cparams.op_offload;
-            } else {
-                LLAMA_LOG_WARN("%s: partial GPU KV residency is not supported for this memory layout; ignoring kv_gpu_layers\n", __func__);
-            }
+            cparams.offload_attn_compute = cparams.offload_attn_compute || cparams.op_offload;
         }
     }
 

src/llama-kv-cache.cpp

@@ -287,7 +287,10 @@ llama_kv_cache::llama_kv_cache(
 
         const char * dev_name = "CPU";
 
-        const bool layer_offload = offload || n_gpu_resident < placement.gpu_resident_layers;
+        // Sub-caches claim each selected layer once.
+        const bool layer_picked = placement.gpu_resident_ils.count(il) > 0 &&
+            (!placement.gpu_resident_done || placement.gpu_resident_done->count(il) == 0);
+        const bool layer_offload = offload || layer_picked;
         ggml_backend_buffer_type_t buft = llama_kv_cache_get_host_buft(model, il, placement.cpu_pinned);
 
         if (layer_offload) {
@@ -299,6 +302,9 @@ llama_kv_cache::llama_kv_cache(
 
         if (!offload && layer_offload) {
             ++n_gpu_resident;
+            if (placement.gpu_resident_done) {
+                placement.gpu_resident_done->insert(il);
+            }
         }
 
         LLAMA_LOG_DEBUG("%s: layer %3d: dev = %s\n", __func__, il, dev_name);
@@ -339,9 +345,8 @@ llama_kv_cache::llama_kv_cache(
         layers.push_back({ il, k, v, k_store_quantize, v_store_quantize, k_stream, v_stream, });
     }
 
-    if (!offload && placement.gpu_resident_layers > 0) {
-        LLAMA_LOG_INFO("%s: partial GPU KV residency: %u of %u requested owned attention layers device-resident\n",
-                __func__, n_gpu_resident, placement.gpu_resident_layers);
+    if (!offload && n_gpu_resident > 0) {
+        LLAMA_LOG_DEBUG("%s: %u attention layers device-resident\n", __func__, n_gpu_resident);
     }
 
     if (reuse) {

src/llama-kv-cache.h

@@ -138,10 +138,6 @@ class llama_kv_cache : public llama_memory_i {
 
     bool get_can_shift() const override;
 
-    bool get_supports_partial_kv() const override {
-        return true;
-    }
-
     void clear(bool data) override;
 
     bool seq_rm  (llama_seq_id seq_id,                              llama_pos p0, llama_pos p1) override;

src/llama-memory-hybrid.h

@@ -66,10 +66,6 @@ class llama_memory_hybrid : public llama_memory_i {
     bool recurrent_sparse_snapshots_supported() const override;
     bool recurrent_set_sparse_snapshot_mode(bool enabled, int32_t selected_token) override;
 
-    bool get_supports_partial_kv() const override {
-        return true;
-    }
-
     void clear(bool data) override;
 
     bool seq_rm  (llama_seq_id seq_id,                              llama_pos p0, llama_pos p1) override;

src/llama-memory.h

@@ -5,6 +5,7 @@
 
 #include <map>
 #include <memory>
+#include <set>
 #include <functional>
 
 struct llama_ubatch;
@@ -16,8 +17,13 @@ class llama_io_read_i;
 
 struct llama_memory_placement_options {
     bool cpu_pinned = false;
-    uint32_t gpu_resident_layers = 0;
     bool recurrent_offload = false;
+
+    // One residency set is shared by all sub-caches of this context.
+    std::set<uint32_t> gpu_resident_ils;
+
+    // The first cache that owns a selected layer claims it; auxiliary copies remain on the host.
+    std::shared_ptr<std::set<uint32_t>> gpu_resident_done;
 };
 
 struct llama_memory_params {
@@ -121,10 +127,6 @@ struct llama_memory_i {
     virtual bool recurrent_sparse_snapshots_supported() const { return false; }
     virtual bool recurrent_set_sparse_snapshot_mode(bool, int32_t) { return false; }
 
-    virtual bool get_supports_partial_kv() const {
-        return false;
-    }
-
     //
     // ops
     //