diff --git a/common/common.cpp b/common/common.cpp index f81bf76..a4ee9a2 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1761,18 +1761,10 @@ bool common_prompt_should_continue_cached( int32_t n_cached, int32_t last_user, int32_t n_new) { - if (n_keep < 0 || n_cached <= 0 || n_new <= 0) { - return false; - } - if (last_user <= n_keep || last_user >= n_new) { - return false; - } - // Only when retokenize stopped near the original prompt, not after a - // near-complete text match of the generation. - if (n_keep > n_cached / 2 || n_cached <= n_keep + 64) { - return false; - } - return true; + // Stateless API requests may belong to unrelated conversations. Only use + // cached history when the entire prefix matches and no new messages would + // be skipped before the final user turn. + return n_cached > 0 && n_keep == n_cached && last_user == n_cached && n_new > n_cached; } int32_t common_cmoe_first_prefill_at(int32_t decode_ubatch, int32_t prefill_ubatch) { diff --git a/tests/test-hybrid-lcp.cpp b/tests/test-hybrid-lcp.cpp index 8ae6cdb..b775eb3 100644 --- a/tests/test-hybrid-lcp.cpp +++ b/tests/test-hybrid-lcp.cpp @@ -351,8 +351,14 @@ int main() { "empty cache"); } - expect(common_prompt_should_continue_cached(25, 29850, 34000, 34498), - "long generate plus a new last-user turn continues from KV"); + expect(!common_prompt_should_continue_cached(25, 29850, 34000, 34498), + "unrelated history must never be replaced by cached generation"); + expect(!common_prompt_should_continue_cached(7, 16383, 56, 600), + "a fresh short request must not inherit an almost full context"); + expect(!common_prompt_should_continue_cached(512, 512, 640, 1000), + "new messages before the final user turn must not be skipped"); + expect(common_prompt_should_continue_cached(512, 512, 512, 600), + "an exact complete prefix may continue without changing the prompt"); expect(!common_prompt_should_continue_cached(25, 50, 40, 80), "short cached tail is not a continuation"); expect(!common_prompt_should_continue_cached(25, 29850, 10, 34498), diff --git a/tools/server/server.cpp b/tools/server/server.cpp index a3b2a8b..f7dd453 100644 --- a/tools/server/server.cpp +++ b/tools/server/server.cpp @@ -10,11 +10,15 @@ #include "common.h" #include "fit.h" #include "llama.h" +#include "src/llama-ext.h" #include "log.h" #include #include +#include #include +#include +#include #include #include // for std::thread::hardware_concurrency @@ -85,6 +89,100 @@ static server_http_context::handler_t ex_wrapper(server_http_context::handler_t }; } +static std::optional aurago_prepare_startup_manifest( + const common_params & params, + const llama_context * ctx) { + const char * manifest_path = std::getenv("AURAGO_STARTUP_MANIFEST"); + if (manifest_path == nullptr || manifest_path[0] == '\0') { + return std::nullopt; + } + + std::ifstream input(manifest_path); + if (!input.good()) { + throw std::runtime_error("AuraGo startup manifest is unavailable"); + } + + json manifest; + input >> manifest; + if (!manifest.is_object()) { + throw std::runtime_error("AuraGo startup manifest is invalid"); + } + + const std::string requested_device = json_value(manifest, "requested_device", std::string()); + const std::string alias = json_value(manifest, "alias", std::string()); + const uint32_t requested_ctx = json_value(manifest, "context_size", uint32_t(0)); + const uint32_t requested_batch = json_value(manifest, "batch_size", uint32_t(0)); + const uint32_t requested_ubatch = json_value(manifest, "ubatch_size", uint32_t(0)); + const bool gpu_requested = !requested_device.empty() && requested_device != "cpu"; + + std::map device_memory; + const llama_memory_breakdown memory_breakdown = llama_get_memory_breakdown(ctx); + for (const auto & [buft, memory] : memory_breakdown) { + if (ggml_backend_buft_is_host(buft)) { + continue; + } + ggml_backend_dev_t device = ggml_backend_buft_get_device(buft); + if (device == nullptr) { + continue; + } + auto & total = device_memory[ggml_backend_dev_name(device)]; + total.model += memory.model; + total.context += memory.context; + total.compute += memory.compute; + } + + std::string actual_device = "cpu"; + llama_memory_breakdown_data actual_memory; + if (gpu_requested) { + auto selected = device_memory.find(requested_device); + if (selected != device_memory.end()) { + actual_device = selected->first; + actual_memory = selected->second; + } + } + + const bool gpu_offload = gpu_requested + && actual_device == requested_device + && params.n_gpu_layers <= -2 + && actual_memory.model > 0; + const bool kv_offload = gpu_requested + && actual_device == requested_device + && !params.no_kv_offload + && actual_memory.context > 0; + const bool alias_verified = !alias.empty() && params.model_alias.count(alias) == 1; + const bool ling = alias == "aurago-ling"; + const bool phase = ling && json_value(manifest, "active_backend", std::string()) == "cuda"; + const uint32_t reserve = phase ? 2048 : requested_batch; + const bool phase_verified = phase + ? params.n_parallel == 1 && requested_batch == 64 && requested_ubatch == 64 + && params.cmoe_n_batch_prefill == 2048 && params.cmoe_n_ubatch_prefill == 2048 + && params.cmoe_n_batch_decode == 64 && params.cmoe_n_ubatch_decode == 64 + : params.cmoe_n_batch_prefill == 0 && params.cmoe_n_ubatch_prefill == 0 + && params.cmoe_n_batch_decode == 0 && params.cmoe_n_ubatch_decode == 0; + const char * kvflash = std::getenv("LLAMA_KVFLASH"); + const bool full_context = !ling || (kvflash && std::string(kvflash) == "0"); + const bool memory_profile_verified = + llama_n_ctx(ctx) == requested_ctx + && llama_n_batch(ctx) == reserve + && llama_n_ubatch(ctx) == (phase ? 2048 : requested_ubatch) + && phase_verified && full_context + && !params.fit_params + && params.enable_reasoning == 0 + && alias_verified + && (gpu_requested + ? gpu_offload && kv_offload + : params.n_gpu_layers == 0 && actual_device == "cpu"); + + manifest["prefill_batch_size"] = phase ? 2048 : 0; + manifest["prefill_ubatch_size"] = phase ? 2048 : 0; + manifest["kvflash_tokens"] = 0; + manifest["actual_device"] = actual_device; + manifest["gpu_offload"] = gpu_offload; + manifest["kv_offload"] = kv_offload; + manifest["memory_profile_verified"] = memory_profile_verified; + return manifest; +} + int llama_server(int argc, char ** argv) { std::setlocale(LC_NUMERIC, "C"); @@ -182,6 +280,7 @@ int llama_server(common_params & params, int argc, char ** argv) { server_child child; // only used in non-router mode server_routes routes(params, ctx_server); server_tools tools; + std::optional aurago_startup_manifest; std::optional models_routes{}; if (is_router_server) { @@ -233,6 +332,21 @@ int llama_server(common_params & params, int argc, char ** argv) { ctx_http.get ("/health", ex_wrapper(routes.get_health)); // public endpoint (no API key check) ctx_http.get ("/v1/health", ex_wrapper(routes.get_health)); // public endpoint (no API key check) ctx_http.get ("/metrics", ex_wrapper(routes.get_metrics)); + ctx_http.get ("/startup-manifest", ex_wrapper([&aurago_startup_manifest](const server_http_req &) { + auto res = std::make_unique(); + if (!aurago_startup_manifest.has_value()) { + res->status = 404; + res->data = safe_json_to_str({ + {"error", { + {"message", "startup manifest is not configured"}, + {"type", "not_found_error"}, + }} + }); + return res; + } + res->data = safe_json_to_str(*aurago_startup_manifest); + return res; + })); ctx_http.get ("/props", ex_wrapper(routes.get_props)); ctx_http.post("/props", ex_wrapper(routes.post_props)); ctx_http.get ("/models", ex_wrapper(routes.get_models)); // public endpoint (no API key check) @@ -456,6 +570,16 @@ int llama_server(common_params & params, int argc, char ** argv) { } routes.update_meta(ctx_server); + try { + aurago_startup_manifest = aurago_prepare_startup_manifest(params, ctx_server.get_llama_context()); + } catch (const std::exception &) { + clean_up(); + if (ctx_http.thread.joinable()) { + ctx_http.thread.join(); + } + SRV_ERR("%s", "failed to prepare AuraGo startup manifest\n"); + return 1; + } ctx_http.is_ready.store(true); SRV_INF("%s", "model loaded\n");