Merge pull request #8 from kmbandy/feature-kv-cache-improvements

Tiered kv cache implementation along with vattention (mad-lab-cache)

Author: kmbandy

.gitignore

@@ -144,7 +144,7 @@ poetry.toml
 /*.code-workspace
 /.windsurf/
 # emscripten
-a.out.*
+a.out.*x
 .gitnexus
 .claude/
 .agent.md
@@ -175,4 +175,5 @@ ggml/ggml-version.cmake
 tests/libgguf-model-data.a
 *.a
 tests/cmake_install.cmake
+Makefile
 *.cmake

common/arg.cpp

@@ -31,6 +31,7 @@
 #include <list>
 #include <regex>
 #include <set>
+#include <sstream>
 #include <string>
 #include <thread> // for hardware_concurrency
 #include <vector>
@@ -1310,6 +1311,81 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
             params.cache_ram_mib = value;
         }
     ).set_env("LLAMA_ARG_CACHE_RAM").set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}));
+    add_opt(common_arg(
+        {"--kv-tiered"}, "VRAM%,RAM%,SSD%",
+        "enable tiered KV cache with percentages for VRAM, RAM, SSD (e.g., 25,25,50 for 25% VRAM, 25% RAM, 50% SSD)",
+        [](common_params & params, const std::string & value) {
+            params.kv_tiered_enabled = true;
+            // Parse the percentage string
+            std::vector<std::string> parts;
+            std::stringstream ss(value);
+            std::string part;
+            while (std::getline(ss, part, ',')) {
+                parts.push_back(part);
+            }
+            if (parts.size() == 3) {
+                params.kv_tier_hot_pct = std::stof(parts[0]);
+                params.kv_tier_warm_pct = std::stof(parts[1]);
+                params.kv_tier_cold_pct = std::stof(parts[2]);
+            }
+        }
+    ).set_env("LLAMA_ARG_KT_TIERED").set_examples({LLAMA_EXAMPLE_SERVER}));
+    add_opt(common_arg(
+        {"--tier-ssd-path"}, "PATH",
+        "directory for cold tier (SSD) storage",
+        [](common_params & params, const std::string & value) {
+            params.kv_tier_ssd_path = value;
+        }
+    ).set_env("LLAMA_ARG_TIER_SSD_PATH").set_examples({LLAMA_EXAMPLE_SERVER}));
+    add_opt(common_arg(
+        {"--tier-eviction-policy"}, "POLICY",
+        "eviction policy: 0=LRU, 1=LFU, 2=attention, 3=hybrid (default)",
+        [](common_params & params, int value) {
+            params.kv_tier_eviction_policy = value;
+        }
+    ).set_env("LLAMA_ARG_TIER_EVICTION_POLICY").set_examples({LLAMA_EXAMPLE_SERVER}));
+    add_opt(common_arg(
+        {"--tier-compression"}, "TYPE",
+        "compression type: 0=none, 1=int4, 2=int8, 3=lz4, 4=quantized",
+        [](common_params & params, int value) {
+            params.kv_tier_compression = value;
+        }
+    ).set_env("LLAMA_ARG_TIER_COMPRESSION").set_examples({LLAMA_EXAMPLE_SERVER}));
+    add_opt(common_arg(
+        {"--tier-attention-threshold"}, "THRESH",
+        "attention threshold for eviction (0.0-1.0, default: 0.1)",
+        [](common_params & params, int value) {
+            params.kv_tier_attention_threshold = float(value) / 100.0f;
+        }
+    ).set_env("LLAMA_ARG_TIER_ATTENTION_THRESHOLD").set_examples({LLAMA_EXAMPLE_SERVER}));
+    add_opt(common_arg(
+        {"--kv-warm-device"}, "DEVICE",
+        "HIP device index to use as warm KV cache tier (e.g. 1 for 6900XT eGPU); requires --kv-tiered",
+        [](common_params & params, int value) {
+            params.kv_warm_device = value;
+        }
+    ).set_env("LLAMA_ARG_KV_WARM_DEVICE").set_examples({LLAMA_EXAMPLE_SERVER}));
+    add_opt(common_arg(
+        {"--kv-semantic-index"}, "PATH",
+        "path to embedding model (GGUF) for semantic KV cache indexing (e.g. bge-small); empty = disabled",
+        [](common_params & params, const std::string & value) {
+            params.kv_semantic_index = value;
+        }
+    ).set_env("LLAMA_ARG_KV_SEMANTIC_INDEX").set_examples({LLAMA_EXAMPLE_SERVER}));
+    add_opt(common_arg(
+        {"--kv-semantic-threshold"}, "THRESHOLD",
+        "minimum cosine similarity threshold for semantic prefetch hints (default: 0.65)",
+        [](common_params & params, const std::string & value) {
+            params.kv_semantic_threshold = std::stof(value);
+        }
+    ).set_env("LLAMA_ARG_KV_SEMANTIC_THRESHOLD").set_examples({LLAMA_EXAMPLE_SERVER}));
+    add_opt(common_arg(
+        {"--kv-semantic-topk"}, "K",
+        "number of prefetch hints to return (default: 5)",
+        [](common_params & params, int value) {
+            params.kv_semantic_top_k = value;
+        }
+    ).set_env("LLAMA_ARG_KV_SEMANTIC_TOPK").set_examples({LLAMA_EXAMPLE_SERVER}));
     add_opt(common_arg(
         {"-kvu", "--kv-unified"},
         {"-no-kvu", "--no-kv-unified"},
@@ -2354,6 +2430,27 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
             }
         }
     ).set_env("LLAMA_ARG_N_GPU_LAYERS"));
+    add_opt(common_arg(
+        {"--weight-paging"},
+        "enable NVMe→VRAM demand paging for model weights (allows models larger than VRAM)",
+        [](common_params & params) {
+            params.weight_paging_enabled = true;
+        }
+    ).set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_WEIGHT_PAGING"));
+    add_opt(common_arg(
+        {"--weight-paging-slots"}, "N",
+        "number of VRAM slots for weight paging (-1 = auto, default: -1)",
+        [](common_params & params, const std::string & value) {
+            params.weight_paging_slots = std::stoi(value);
+        }
+    ).set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_WEIGHT_PAGING_SLOTS"));
+    add_opt(common_arg(
+        {"--weight-paging-prefetch"},
+        "enable async prefetch of next layer (default: enabled)",
+        [](common_params & params) {
+            params.weight_paging_prefetch = true;
+        }
+    ).set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_WEIGHT_PAGING_PREFETCH"));
     add_opt(common_arg(
         {"-sm", "--split-mode"}, "{none,layer,row,tensor}",
         "how to split the model across multiple GPUs, one of:\n"

common/common.cpp

@@ -1343,7 +1343,7 @@ common_init_result_ptr common_init_from_params(common_params & params) {
         common_set_adapter_lora(lctx, params.lora_adapters);
     }
 
-    if (params.warmup) {
+    if (params.warmup && !params.weight_paging_enabled) {
         LOG_WRN("%s: warming up the model with an empty run - please wait ... (--no-warmup to disable)\n", __func__);
 
         llama_set_warmup(lctx, true);
@@ -1488,6 +1488,11 @@ struct llama_model_params common_model_params_to_llama(common_params & params) {
     mparams.progress_callback_user_data = params.load_progress_callback_user_data;
     mparams.no_alloc                    = params.no_alloc;
 
+    // Weight paging parameters
+    mparams.weight_paging_enabled  = params.weight_paging_enabled;
+    mparams.weight_paging_slots    = params.weight_paging_slots;
+    mparams.weight_paging_prefetch = params.weight_paging_prefetch;
+
     return mparams;
 }
 

common/common.h

@@ -533,6 +533,11 @@ struct common_params {
     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
 
+    // weight paging parameters
+    bool    weight_paging_enabled = false;  // enable NVMe→VRAM demand paging for model weights
+    int32_t weight_paging_slots   = -1;     // number of VRAM slots for weight paging (-1 = auto)
+    bool    weight_paging_prefetch = false;  // enable async prefetch of next layer
+
     bool single_turn       = false; // single turn chat conversation
 
     ggml_type cache_type_k = GGML_TYPE_F16; // KV cache data type for the K
@@ -573,6 +578,21 @@ struct common_params {
     int32_t checkpoint_every_nt = 8192;  // make a checkpoint every n tokens during prefill
     int32_t cache_ram_mib       = 8192;  // -1 = no limit, 0 - disable, 1 = 1 MiB, etc.
 
+    // tiered KV cache parameters
+    bool kv_tiered_enabled    = false;   // enable tiered KV cache
+    float kv_tier_hot_pct     = 25.0f;   // hot tier percentage (VRAM)
+    float kv_tier_warm_pct    = 25.0f;   // warm tier percentage (RAM)
+    float kv_tier_cold_pct    = 50.0f;   // cold tier percentage (SSD)
+    std::string kv_tier_ssd_path = "";   // SSD path for cold tier storage
+    int kv_tier_eviction_policy = 3;     // 0=LRU, 1=LFU, 2=attention, 3=hybrid (default)
+    int kv_tier_compression   = 1;       // 0=none, 1=int4, 2=int8, 3=lz4, 4=quantized
+    float kv_tier_attention_threshold = 0.1f;  // attention threshold for eviction
+    int   kv_warm_device        = -1;    // HIP device index for warm KV tier (-1 = disabled)
+    int   kv_tier_total_ctx     = 0;     // full ctx budget across all tiers (set at load time)
+    std::string kv_semantic_index = "";   // path to embedding model for semantic KV index (empty = disabled)
+    float kv_semantic_threshold = 0.65f;  // minimum cosine similarity threshold for prefetch hints
+    int   kv_semantic_top_k     = 5;      // number of prefetch hints to return
+
     std::string hostname      = "127.0.0.1";
     std::string public_path   = "";                                                                         // NOLINT
     std::string api_prefix    = "";                                                                         // NOLINT

common/fit.cpp

@@ -53,9 +53,10 @@ static std::vector<llama_device_memory_data> common_get_device_memory_data(
     }, &ud);
 
     llama_model_params mparams_copy = *mparams;
-    mparams_copy.no_alloc  = true;
-    mparams_copy.use_mmap  = false;
-    mparams_copy.use_mlock = false;
+    mparams_copy.no_alloc       = true;
+    mparams_copy.use_mmap       = false;
+    mparams_copy.use_mlock      = false;
+    mparams_copy.use_direct_io  = false;
 
     llama_model * model = llama_model_load_from_file(path_model, mparams_copy);
     if (model == nullptr) {

ggml/src/ggml-cpu/ops.cpp

@@ -12,7 +12,7 @@
 #include <cfloat>
 #include <cmath>
 
-extern "C" GGML_API int turbo3_cpu_wht_group_size;
+extern "C" int turbo3_cpu_wht_group_size;
 
 // ggml_compute_forward_dup
 

ggml/src/ggml-cuda/common.cuh

@@ -1193,7 +1193,7 @@ struct ggml_cuda_graph {
     std::vector<node_properties> node_props;
 
     bool is_enabled() const {
-        static const bool disable_cuda_graphs_due_to_env = (getenv("GGML_CUDA_DISABLE_GRAPHS") != nullptr);
+        const bool disable_cuda_graphs_due_to_env = (getenv("GGML_CUDA_DISABLE_GRAPHS") != nullptr);
         return !(disable_due_to_gpu_arch || disable_cuda_graphs_due_to_env);
     }
 #endif

ggml/src/ggml-cuda/ggml-cuda.cu

@@ -1419,6 +1419,54 @@ static cudaError_t ggml_cuda_cpy_tensor_2d(
     const int64_t i1_diff = i1_high - i1_low;
 
     const char * x = src_ptr + i1_low*nb1 + i2*nb2 + i3*nb3;
+
+#if defined(GGML_USE_HIP)
+    // hipMemcpy2DAsync with hipMemcpyDeviceToDevice requires P2P access between devices.
+    // For mixed-architecture multi-GPU setups (e.g. gfx1201 + gfx1030) without P2P,
+    // it fails asynchronously with hipErrorNoBinaryForGpu when the internal copy kernel
+    // isn't compiled for the destination device. Detect cross-device copies via pointer
+    // attributes and fall back to row-wise hipMemcpyPeerAsync, which stages through the
+    // host when P2P is unavailable. Mirrors the fix in ggml_cuda_Memcpy2DPeerAsync.
+    hipPointerAttribute_t src_attr = {};
+    hipPointerAttribute_t dst_attr = {};
+    const bool src_ok = hipPointerGetAttributes(&src_attr, x)       == hipSuccess;
+    const bool dst_ok = hipPointerGetAttributes(&dst_attr, dst_ptr) == hipSuccess;
+    if (src_ok && dst_ok && src_attr.device != dst_attr.device) {
+        const int src_dev = src_attr.device;
+        const int dst_dev = dst_attr.device;
+        if (nb0 == ts && nb1 == ts*ne0/bs) {
+            return hipMemcpyPeerAsync(dst_ptr, dst_dev, x, src_dev, i1_diff*nb1, stream);
+        } else if (nb0 == ts) {
+            const size_t row_bytes = ts*ne0/bs;
+            for (int64_t i1 = 0; i1 < i1_diff; i1++) {
+                hipError_t err = hipMemcpyPeerAsync(
+                    dst_ptr + i1*row_bytes, dst_dev,
+                    x + i1*nb1,             src_dev,
+                    row_bytes, stream);
+                if (err != hipSuccess) {
+                    return err;
+                }
+            }
+            return hipSuccess;
+        } else {
+            for (int64_t i1 = 0; i1 < i1_diff; i1++) {
+                const char * rx = x + i1*nb1;
+                char       * rd = dst_ptr + i1*ts*ne0/bs;
+                for (int64_t e = 0; e < ne0; e++) {
+                    hipError_t err = hipMemcpyPeerAsync(
+                        rd + e*(ts/bs), dst_dev,
+                        rx + e*nb0,     src_dev,
+                        ts/bs, stream);
+                    if (err != hipSuccess) {
+                        return err;
+                    }
+                }
+            }
+            return hipSuccess;
+        }
+    }
+#endif
+
     if (nb0 == ts && nb1 == ts*ne0/bs) {
         return cudaMemcpyAsync(dst_ptr, x, i1_diff*nb1, cudaMemcpyDeviceToDevice, stream);
     } else if (nb0 == ts) {

include/llama.h

@@ -321,6 +321,11 @@ extern "C" {
         bool use_extra_bufts; // use extra buffer types (used for weight repacking)
         bool no_host;         // bypass host buffer allowing extra buffers to be used
         bool no_alloc;        // only load metadata and simulate memory allocations
+
+        // Weight paging parameters (NVMe→VRAM demand paging)
+        bool weight_paging_enabled;    // enable NVMe→VRAM demand paging for model weights
+        int32_t weight_paging_slots;   // number of VRAM slots for weight paging (-1 = auto)
+        bool weight_paging_prefetch;   // enable async prefetch of next layer
     };
 
     struct llama_sampler_seq_config {

plans/errors/command.md

@@ -0,0 +1,28 @@
+Here is what I use to start the server:
+
+llama-server-feature \
+          --model ~/models/MiniMax-M2.7-UD-IQ3_S-00001-of-00003.gguf \
+          -ngl 999 \
+          --device ROCm0 \
+          --ctx-size 196608 \
+          --cache-type-k turbo4 \
+          --cache-type-v turbo4 \
+          --flash-attn on \
+          --metrics \
+          --port 8080 \
+          --host 0.0.0.0 \
+          --no-mmap \
+          --kv-tiered 12.5,50,37.5 \
+          --tier-ssd-path ~/kv-cold \
+          --tier-eviction-policy 3 \
+          --parallel 1 \
+          --kv-semantic-index ~/models/bge-small-en-v1.5-q8_0.gguf \
+          --kv-semantic-threshold 0.65 \
+          --kv-semantic-topk 5 \
+          --alias default \
+          --direct-io \
+          --kv-warm-device 1 \
+          --weight-paging \
+          --fit off \
+          --weight-paging-slots 8 \
+          --jinja