diff --git a/src/arbiterAI/modelRuntime.cpp b/src/arbiterAI/modelRuntime.cpp index da76092..6894a5f 100644 --- a/src/arbiterAI/modelRuntime.cpp +++ b/src/arbiterAI/modelRuntime.cpp @@ -2197,6 +2197,18 @@ llama_context *ModelRuntime::getLlamaContext(const std::string &model) const return nullptr; } +std::vector *ModelRuntime::kvCacheTokens(const std::string &model) +{ + std::lock_guard lock(m_mutex); + + auto it=m_models.find(model); + if(it!=m_models.end()&&it->second.state==ModelState::Loaded) + { + return &it->second.kvCacheTokens; + } + return nullptr; +} + std::optional ModelRuntime::getLoadedModelInfo(const std::string &model) const { std::lock_guard lock(m_mutex); diff --git a/src/arbiterAI/modelRuntime.h b/src/arbiterAI/modelRuntime.h index 10a1b83..7a4a4c4 100644 --- a/src/arbiterAI/modelRuntime.h +++ b/src/arbiterAI/modelRuntime.h @@ -87,6 +87,7 @@ struct LoadedModel { llama_model *llamaModel=nullptr; llama_context *llamaCtx=nullptr; RuntimeOptions activeOptions; // llama.cpp options active for this loaded model + std::vector kvCacheTokens; // tokens currently in llamaCtx's KV cache (seq 0), for cache_prompt prefix reuse }; class ModelRuntime { @@ -209,6 +210,12 @@ class ModelRuntime { /// Returns nullptr if not loaded or not a local model. llama_context *getLlamaContext(const std::string &model) const; + /// Access the KV-cache token record for a loaded model (nullptr if not + /// loaded). Mirrors the tokens decoded into the context's KV cache on + /// sequence 0 so cache_prompt requests can reuse the common prefix. + /// Callers must hold the inference mutex while reading or mutating it. + std::vector *kvCacheTokens(const std::string &model); + /// Get the ModelInfo for a loaded model. std::optional getLoadedModelInfo(const std::string &model) const; diff --git a/src/arbiterAI/providers/llama.cpp b/src/arbiterAI/providers/llama.cpp index 54ecdba..f3240d1 100644 --- a/src/arbiterAI/providers/llama.cpp +++ b/src/arbiterAI/providers/llama.cpp @@ -241,6 +241,13 @@ ErrorCode Llama::getEmbeddings(const EmbeddingRequest &request, std::lock_guard inferenceLock(runtime.getInferenceMutex()); + // Embedding batches overwrite the context's KV cache — the completion + // prefix record is no longer valid + if(std::vector *kvRecord=runtime.kvCacheTokens(request.model)) + { + kvRecord->clear(); + } + // Combine input text std::string inputText; std::visit([&inputText](auto &&arg) @@ -616,6 +623,64 @@ std::string Llama::formatHarmonyPrompt(const CompletionRequest &request, return prompt; } +int kvPrefixReuseLength(const std::vector &cachedTokens, + const std::vector &promptTokens) +{ + // Keep at least one prompt token to decode so the final position has + // fresh logits for sampling + int maxReuse=std::min(static_cast(cachedTokens.size()), + static_cast(promptTokens.size())-1); + int reused=0; + while(reused &promptTokens, std::vector *kvRecord) +{ + llama_memory_t mem=llama_get_memory(ctx); + const int nTokens=static_cast(promptTokens.size()); + int reused=0; + + if(request.cache_prompt.value_or(false)&&kvRecord&&!kvRecord->empty()) + { + reused=kvPrefixReuseLength(*kvRecord, promptTokens); + } + + if(reused>0&&llama_memory_seq_rm(mem, 0, reused, -1)) + { + spdlog::info("[llama] cache_prompt: reusing {} of {} prompt tokens from KV cache", + reused, nTokens); + } + else + { + if(reused>0) + { + spdlog::warn("[llama] cache_prompt: partial KV erase unsupported — full prefill"); + } + spdlog::debug("[llama] clearing KV cache, prompt tokens={}", nTokens); + llama_memory_clear(mem, true); + reused=0; + } + + if(kvRecord) + { + kvRecord->clear(); + } + return reused; +} + ErrorCode Llama::runInference(llama_model *model, llama_context *ctx, const CompletionRequest &request, const ModelInfo &modelInfo, std::string &result, int &promptTokens, int &completionTokens, @@ -655,17 +720,17 @@ ErrorCode Llama::runInference(llama_model *model, llama_context *ctx, tokensList.resize(nTokens); promptTokens=nTokens; - // Clear KV cache for fresh inference - spdlog::debug("[llama] clearing KV cache, prompt tokens={}", nTokens); - llama_memory_clear(llama_get_memory(ctx), true); + std::vector *kvRecord=ModelRuntime::instance().kvCacheTokens(request.model); + int reusedTokens=prepareKvCache(ctx, request, tokensList, kvRecord); int nBatch=static_cast(llama_n_batch(ctx)); llama_batch batch=llama_batch_init(std::max(nBatch, 512), 0, 1); - // Process prompt (timed) — chunk into n_batch-sized pieces + // Process prompt (timed) — chunk into n_batch-sized pieces, skipping any + // prefix already in the KV cache std::chrono::steady_clock::time_point promptStart=std::chrono::steady_clock::now(); - for(int start=0; start=nTokens); @@ -712,6 +777,7 @@ ErrorCode Llama::runInference(llama_model *model, llama_context *ctx, int nCur=nTokens; completionTokens=0; + std::vector generatedInKv; // Set up sampler chain llama_sampler_chain_params samplerParams=llama_sampler_chain_default_params(); @@ -848,6 +914,7 @@ ErrorCode Llama::runInference(llama_model *model, llama_context *ctx, llama_batch_free(batch); return ErrorCode::GenerationError; } + generatedInKv.push_back(nextToken); } std::chrono::steady_clock::time_point genEnd=std::chrono::steady_clock::now(); @@ -856,6 +923,14 @@ ErrorCode Llama::runInference(llama_model *model, llama_context *ctx, llama_sampler_free(samplerChain); llama_batch_free(batch); + // Record what now sits in the KV cache so the next cache_prompt request + // can reuse the common prefix + if(kvRecord) + { + *kvRecord=tokensList; + kvRecord->insert(kvRecord->end(), generatedInKv.begin(), generatedInKv.end()); + } + return ErrorCode::Success; } @@ -907,15 +982,15 @@ ErrorCode Llama::runInferenceWithTokens(llama_model *model, llama_context *ctx, int nTokens=static_cast(promptTokens.size()); promptTokenCount=nTokens; - spdlog::debug("[llama] clearing KV cache, prompt tokens={}", nTokens); - llama_memory_clear(llama_get_memory(ctx), true); + std::vector *kvRecord=ModelRuntime::instance().kvCacheTokens(request.model); + int reusedTokens=prepareKvCache(ctx, request, promptTokens, kvRecord); int nBatch=static_cast(llama_n_batch(ctx)); llama_batch batch=llama_batch_init(std::max(nBatch, 512), 0, 1); std::chrono::steady_clock::time_point promptStart=std::chrono::steady_clock::now(); - for(int start=0; start generatedInKv; llama_sampler_chain_params samplerParams=llama_sampler_chain_default_params(); llama_sampler *samplerChain=llama_sampler_chain_init(samplerParams); @@ -1098,6 +1174,7 @@ ErrorCode Llama::runInferenceWithTokens(llama_model *model, llama_context *ctx, llama_batch_free(batch); return ErrorCode::GenerationError; } + generatedInKv.push_back(nextToken); } std::chrono::steady_clock::time_point genEnd=std::chrono::steady_clock::now(); @@ -1106,6 +1183,14 @@ ErrorCode Llama::runInferenceWithTokens(llama_model *model, llama_context *ctx, llama_sampler_free(samplerChain); llama_batch_free(batch); + // Record what now sits in the KV cache so the next cache_prompt request + // can reuse the common prefix + if(kvRecord) + { + *kvRecord=promptTokens; + kvRecord->insert(kvRecord->end(), generatedInKv.begin(), generatedInKv.end()); + } + return ErrorCode::Success; } diff --git a/src/arbiterAI/providers/llama.h b/src/arbiterAI/providers/llama.h index cd76433..88be52b 100644 --- a/src/arbiterAI/providers/llama.h +++ b/src/arbiterAI/providers/llama.h @@ -15,6 +15,12 @@ struct llama_context; namespace arbiterAI { +/// How many leading tokens of promptTokens are already present in the KV +/// cache (per cachedTokens) and can skip prefill. Always leaves at least +/// one prompt token to decode so the final position has fresh logits. +int kvPrefixReuseLength(const std::vector &cachedTokens, + const std::vector &promptTokens); + class Llama : public BaseProvider { public: Llama(); diff --git a/src/server/routes.cpp b/src/server/routes.cpp index b77d018..0c8fa6f 100644 --- a/src/server/routes.cpp +++ b/src/server/routes.cpp @@ -1648,6 +1648,10 @@ void handleChatCompletions(const httplib::Request &req, httplib::Response &res) arbiterRequest.stop=requestJson.at("stop").get>(); } + // llama.cpp extension: reuse the KV cache for the common prompt prefix + if(requestJson.contains("cache_prompt")) + arbiterRequest.cache_prompt=requestJson.at("cache_prompt").get(); + // Parse tools in OpenAI format: [{type: "function", function: {...}}] if(requestJson.contains("tools")) { diff --git a/tests/llamaProviderTests.cpp b/tests/llamaProviderTests.cpp index b535c68..f02852e 100644 --- a/tests/llamaProviderTests.cpp +++ b/tests/llamaProviderTests.cpp @@ -3,6 +3,7 @@ #include "arbiterAI/arbiterAI.h" #include "arbiterAI/chatClient.h" #include "arbiterAI/modelRuntime.h" +#include "arbiterAI/providers/llama.h" #include "arbiterAI/telemetryCollector.h" #include "arbiterAI/modelManager.h" @@ -431,4 +432,49 @@ TEST_F(LlamaConfigInjectionTest, InjectWithoutVariantsFails) EXPECT_EQ(loadResult, ErrorCode::InvalidRequest); } +// ── cache_prompt prefix reuse ──────────────────────────────────────────── + +TEST(KvPrefixReuse, EmptyCacheReusesNothing) +{ + EXPECT_EQ(kvPrefixReuseLength({}, {1, 2, 3}), 0); +} + +TEST(KvPrefixReuse, GrowingPromptReusesWholeCache) +{ + // Typical agent turn: previous prompt + generated reply + new tool result + EXPECT_EQ(kvPrefixReuseLength({1, 2, 3, 4}, {1, 2, 3, 4, 5, 6, 7}), 4); +} + +TEST(KvPrefixReuse, IdenticalPromptLeavesOneTokenToDecode) +{ + // The last position must be re-decoded so sampling has fresh logits + EXPECT_EQ(kvPrefixReuseLength({1, 2, 3, 4}, {1, 2, 3, 4}), 3); +} + +TEST(KvPrefixReuse, DivergenceTruncatesAtMismatch) +{ + EXPECT_EQ(kvPrefixReuseLength({1, 2, 3, 4}, {1, 2, 9, 4, 5}), 2); +} + +TEST(KvPrefixReuse, ShorterPromptCapsBelowPromptLength) +{ + EXPECT_EQ(kvPrefixReuseLength({1, 2, 3, 4, 5, 6}, {1, 2, 3}), 2); +} + +TEST(KvPrefixReuse, SingleTokenPromptNeverReuses) +{ + EXPECT_EQ(kvPrefixReuseLength({1, 2, 3}, {1}), 0); +} + +TEST(KvPrefixReuse, CompletelyDifferentPromptReusesNothing) +{ + EXPECT_EQ(kvPrefixReuseLength({7, 8, 9}, {1, 2, 3}), 0); +} + +TEST(KvPrefixReuse, KvCacheTokensRequiresLoadedModel) +{ + ModelRuntime::reset(); + EXPECT_EQ(ModelRuntime::instance().kvCacheTokens("not-loaded"), nullptr); +} + } // namespace arbiterAI