diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/paged-attn.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/paged-attn.cu index 6d7e2c6b8..6f44b47dd 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/paged-attn.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/paged-attn.cu @@ -3,6 +3,7 @@ #include "fattn-common.cuh" #include "mma.cuh" +#include #include #include #include @@ -18,6 +19,20 @@ static __host__ __device__ __forceinline__ int32_t paged_attn_ceil_div( return (int32_t) (dividend / divisor + (dividend % divisor != 0)); } +// Most partitions a launch may split the context into: each one keeps a +// head-dim row of halves and a float2 per output row, and the context alone +// asks for one per 1024 tokens. A 2048-row prompt chunk 118K tokens deep ran +// the R9700 out of memory on those partials, beside the weights and the KV +// cache. Its rows give such a chunk its parallelism, so it covers the +// context with fewer, longer partitions instead. +static int32_t paged_attn_partials_cap(int64_t output_rows, int32_t head_dim) { + constexpr int64_t budget = 256ll << 20; + const int64_t per_partition = + output_rows * ((int64_t) head_dim * (int64_t) sizeof(half) + (int64_t) sizeof(float2)); + const int64_t cap = per_partition > 0 ? budget / per_partition : PAGED_ATTN_MAX_PARTITIONS; + return (int32_t) std::max(1, std::min(cap, PAGED_ATTN_MAX_PARTITIONS)); +} + // Partition count for n_blocks of context: enough partitions to cover the // blocks and to reach min_partitions for occupancy, but never more than the // blocks themselves or the cap. Host and device must agree on this so the @@ -978,6 +993,10 @@ static bool try_launch_paged_attn( if (min_partitions > partition_limit) { min_partitions = partition_limit; } + const int32_t partition_cap = paged_attn_partials_cap(output_rows, D); + if (min_partitions > partition_cap) { + min_partitions = partition_cap; + } const int32_t tree_blocks = paged_attn_ceil_div(tree_width, block_size); const int64_t partitionable_blocks = block_table->ne[0] + (parent_ids ? tree_blocks : 0); @@ -993,7 +1012,7 @@ static bool try_launch_paged_attn( const int32_t live_blocks = paged_attn_ceil_div(live_tokens, block_size); int32_t n_partitions = paged_attn_partitions( - live_blocks, min_partitions, PAGED_ATTN_MAX_PARTITIONS); + live_blocks, min_partitions, partition_cap); // Test/debug override used to exercise both the direct and partials paths // independently of device-specific occupancy. Values outside the valid @@ -2162,11 +2181,15 @@ static bool try_launch_paged_attn_wmma(ggml_backend_cuda_context & ctx, ggml_ten if (min_partitions > partition_limit) { min_partitions = partition_limit; } + const int32_t partition_cap = paged_attn_partials_cap((int64_t) n_rows * n_head, 256); + if (min_partitions > partition_cap) { + min_partitions = partition_cap; + } if (min_partitions > partitionable_blocks) { min_partitions = partitionable_blocks; } int32_t n_partitions = - paged_attn_partitions(n_logical_blocks, min_partitions, PAGED_ATTN_MAX_PARTITIONS); + paged_attn_partitions(n_logical_blocks, min_partitions, partition_cap); // Test/debug override, applied after the occupancy floor like the decode // launcher: the floor above would otherwise silently raise any forced // value <= 32 back to 32 and defeat partition bisection experiments. diff --git a/server/src/common/concurrency/seq_engine.h b/server/src/common/concurrency/seq_engine.h index 33c93744d..f58de81d4 100644 --- a/server/src/common/concurrency/seq_engine.h +++ b/server/src/common/concurrency/seq_engine.h @@ -78,6 +78,13 @@ struct StepPlanLimits { int max_prefill_tokens_per_sequence = 512; int max_prefill_tokens_total = 512; int prefill_allocation_quantum = 512; + // While requests decode, the share of wall time a long prompt may take: + // after each of its chunks, the decoding requests step alone for the + // rest. 1: a prompt chunk rides every step. + double prefill_time_share = 1.0; + // A prompt is long when it brings more new tokens than this; 0: more + // than one step's chunk (max_prefill_tokens_per_sequence). + int time_share_prompt_tokens = 0; }; // Select the oldest eligible sequences, then distribute the step's token diff --git a/server/src/common/model_backend.h b/server/src/common/model_backend.h index 296ec68c4..9bcc9df27 100644 --- a/server/src/common/model_backend.h +++ b/server/src/common/model_backend.h @@ -269,9 +269,9 @@ struct ModelBackend { virtual int restore_point_spacing() const { return 0; } // True when the backend ends a prompt forward at every restore point - // (GenerateRequest::restore_points). A tool-result request then keeps its - // own last boundary among them, so a hit that skips the capture and a - // miss that captures the tools head cut prefill alike. + // (GenerateRequest::restore_points). A request then keeps its own last + // boundary among them, where its snapshot goes, so a hit that skips the + // capture and a miss that captures cut prefill alike. virtual bool prefill_cuts_at_restore_points() const { return false; } // System-memory bytes a snapshot of the first `tokens` positions would diff --git a/server/src/qwen35/concurrency/qwen35_seq_engine.h b/server/src/qwen35/concurrency/qwen35_seq_engine.h index fa2cacfff..420f61631 100644 --- a/server/src/qwen35/concurrency/qwen35_seq_engine.h +++ b/server/src/qwen35/concurrency/qwen35_seq_engine.h @@ -88,6 +88,10 @@ class Qwen35SeqEngine final : public SeqEngine { const PrefixStorePlan & plan) override; StepResult step(const StepPlan & plan) override; + // A step that carries a prompt chunk decodes without the drafter, one + // token per request. Beside a prompt of more than 8192 new tokens (a + // document, not an agent's next step), the decoding requests get half + // the time back in steps of their own, with it. StepPlanLimits step_plan_limits(int decode_rows) const override { const bool mixed = decode_rows > 0; const int per_sequence = mixed ? 512 : 2048; @@ -103,6 +107,8 @@ class Qwen35SeqEngine final : public SeqEngine { per_sequence, std::min(max_prefills_ * per_sequence, total_cap), prefill_quantum_, + /*prefill_time_share=*/mixed ? 0.5 : 1.0, + /*time_share_prompt_tokens=*/8192, }; } diff --git a/server/src/qwen4exp/qwen4exp_seq_engine.cpp b/server/src/qwen4exp/qwen4exp_seq_engine.cpp index 5c62b3cff..714cab679 100644 --- a/server/src/qwen4exp/qwen4exp_seq_engine.cpp +++ b/server/src/qwen4exp/qwen4exp_seq_engine.cpp @@ -355,12 +355,17 @@ int Qwen4ExpSeqEngine::prefill_segment(int slot, int max_tokens) const { } // One granule per prompt forward, whatever the step carries. While slots -// decode, one granule in total per step bounds the hold on the live streams. +// decode, one granule in total per step bounds the hold on the live streams, +// and beside a prompt of more than 8192 new tokens (a document, not an agent's +// next step) they then decode alone for as long as the granule took: a granule +// pays the whole expert read, so a shorter one would not let them run faster, +// and a long prompt would otherwise give them one token per granule. StepPlanLimits Qwen4ExpSeqEngine::step_plan_limits(int decode_rows) const { const int prefill_slots = slot_count() - std::clamp(decode_rows, 0, slot_count()); if (decode_rows > 0) { const int sequences = std::min(prefill_slots, 1); - return {sequences, prefill_granule_, sequences * prefill_granule_, prefill_granule_}; + return {sequences, prefill_granule_, sequences * prefill_granule_, prefill_granule_, + /*prefill_time_share=*/0.5, /*time_share_prompt_tokens=*/8192}; } return {prefill_slots, prefill_granule_, prefill_slots * prefill_granule_, prefill_granule_}; } diff --git a/server/src/server/chat_template.cpp b/server/src/server/chat_template.cpp index d3c021ed6..921ea3466 100644 --- a/server/src/server/chat_template.cpp +++ b/server/src/server/chat_template.cpp @@ -1046,7 +1046,7 @@ static std::shared_ptr get_or_parse(const std::string & template } // namespace -std::string render_chat_template_jinja( +static std::string render_jinja( const std::string & template_src, const std::vector & messages, const std::string & bos_token, @@ -1055,7 +1055,8 @@ std::string render_chat_template_jinja( bool enable_thinking, const std::string & tools_json, const std::string & reasoning_effort, - int preserve_thinking) + int preserve_thinking, + bool late_system_as_user) { if (template_src.empty()) { throw std::runtime_error("render_chat_template_jinja: template_src is empty"); @@ -1077,7 +1078,9 @@ std::string render_chat_template_jinja( "normalize with tool-memory replay disabled"); } nlohmann::ordered_json mj; - mj["role"] = m.role; + const bool late_system = late_system_as_user && !messages_j.empty() && + (m.role == "system" || m.role == "developer"); + mj["role"] = late_system ? std::string("user") : m.role; mj["content"] = m.content; if (!m.tool_call_id.empty()) { mj["tool_call_id"] = m.tool_call_id; @@ -1164,4 +1167,33 @@ std::string render_chat_template_jinja( } } +std::string render_chat_template_jinja( + const std::string & template_src, + const std::vector & messages, + const std::string & bos_token, + const std::string & eos_token, + bool add_generation_prompt, + bool enable_thinking, + const std::string & tools_json, + const std::string & reasoning_effort, + int preserve_thinking) +{ + try { + return render_jinja(template_src, messages, bos_token, eos_token, add_generation_prompt, + enable_thinking, tools_json, reasoning_effort, preserve_thinking, false); + } catch (const std::runtime_error &) { + // Some templates take a system message only first (Qwen's raise). + // Claude Code sends its environment as a system message after the + // user's first one; for such a template it reads as a user turn + // there, in place, so the prompt before it stays the same and cached. + const bool late_system = messages.size() > 1 && + std::any_of(messages.begin() + 1, messages.end(), [](const ChatMessage & m) { + return m.role == "system" || m.role == "developer"; + }); + if (!late_system) throw; + return render_jinja(template_src, messages, bos_token, eos_token, add_generation_prompt, + enable_thinking, tools_json, reasoning_effort, preserve_thinking, true); + } +} + } // namespace luce::common diff --git a/server/src/server/http_server.cpp b/server/src/server/http_server.cpp index c2793a269..e71d34f10 100644 --- a/server/src/server/http_server.cpp +++ b/server/src/server/http_server.cpp @@ -1470,9 +1470,9 @@ HttpServer::HttpServer(luce::engine::LuceEngine & engine, curl_global_init(CURL_GLOBAL_DEFAULT); #endif prefix_cache_.init_full_cache(config.prefill_cache_cap); - // Single-sequence commits prune superseded snapshots (see - // trim_snapshots_after_commit); the paged scheduler does not. - prefix_cache_.set_prunes_superseded(!config.concurrent_prefix_cache); + // Commits prune the snapshots they supersede: single-sequence serving in + // trim_snapshots_after_commit, the concurrent scheduler as each capture lands. + prefix_cache_.set_prunes_superseded(true); // Fold model+config identity into the layout fingerprint BEFORE init() // so compute_layout_id sees it on every learn/verify call. Prevents stale // KV hits when the server restarts over the same --kv-cache-dir with a @@ -2415,7 +2415,9 @@ void apply_request_reasoning( } if (body.contains("thinking") && body["thinking"].is_object()) { const auto & thinking = body["thinking"]; - if (thinking.contains("type")) { + // "adaptive" (Claude Code's default) leaves thinking to the model: the + // same as no thinking field. + if (thinking.contains("type") && thinking.value("type", "") != "adaptive") { const bool enabled = thinking.value("type", "") == "enabled"; enable_thinking = enabled; req.thinking_opt_in = enabled; @@ -2560,7 +2562,9 @@ bool HttpServer::render_and_tokenize_request( if (!render_messages_to_text(chat_messages, req, /*add_generation_prompt=*/true, req.rendered_prompt, error)) { - send_error(fd, 500, error); + // The template rejected these messages: a client error, which a + // client must not retry as it would a server failure. + send_error(fd, 400, error); return false; } req.started_in_thinking = prompt_ends_in_open_think(req.rendered_prompt); @@ -2973,6 +2977,9 @@ json build_openai_completion_response( {"completion_tokens_details", { {"reasoning_tokens", counts.reasoning}, }}, + {"prompt_tokens_details", { + {"cached_tokens", timings.cached_prefix_tokens}, + }}, {"timings", build_timings_json(timings, counts.total)}, {"accept_rate", result.accept_rate}, {"spec_decode_ran", result.spec_decode_ran}, @@ -3073,6 +3080,9 @@ json build_responses_api_response( {"input_tokens", prompt_tokens}, {"output_tokens", counts.total}, {"total_tokens", prompt_tokens + counts.total}, + {"input_tokens_details", { + {"cached_tokens", timings.cached_prefix_tokens}, + }}, {"timings", build_timings_json(timings, counts.total)}, {"accept_rate", result.accept_rate}, {"spec_decode_ran", result.spec_decode_ran}, @@ -3667,10 +3677,9 @@ bool HttpServer::forward_upstream( } std::vector HttpServer::request_restore_points(const std::vector & prompt, - bool ends_with_tool_result, std::initializer_list cuts) const { return prefix_restore_points(prompt, prefix_cache_.chat_markers(), cuts, - /*drop_last_boundary=*/!(backend_.prefill_cuts_at_restore_points() && ends_with_tool_result), + /*drop_last_boundary=*/!backend_.prefill_cuts_at_restore_points(), backend_.restore_point_spacing()); } @@ -4033,8 +4042,7 @@ HttpServer::GenerationCacheState HttpServer::prepare_generation_cache( [this](int target_cut) { return backend_.snapshot_bytes_estimate(target_cut); }, - req.ends_with_tool_result, reachable_from, - backend_.restore_point_spacing()); + reachable_from, backend_.restore_point_spacing()); cache.snap_slot = cache.snap_reservation.slot(); cache.snap_cut = cache.snap_reservation.target_cut(); }; @@ -4114,7 +4122,7 @@ HttpServer::GenerationCacheState HttpServer::prepare_generation_cache( if (!prefix_cache_.disabled() || !disk_cache_.disabled()) { generate_request.restore_points = request_restore_points( - effective_prompt, req.ends_with_tool_result, + effective_prompt, {forced_cut, selected_boundary, cache.snap_prepared ? cache.snap_cut : 0}); } diff --git a/server/src/server/http_server.h b/server/src/server/http_server.h index 59135f4bf..0e9d9a029 100644 --- a/server/src/server/http_server.h +++ b/server/src/server/http_server.h @@ -378,8 +378,8 @@ struct ParsedRequest { DiskPrefixCachePolicy disk_cache_policy; // PPP: stable pin cut for tool-heavy requests (0 = use default boundary). int pin_end_token = 0; - // The prompt ends with tool results. Clients only append after them, so - // the inline snapshot may cover them (see select_inline_snapshot_boundary). + // The prompt ends with tool results: an agent's step, which clients only + // append to. bool ends_with_tool_result = false; }; @@ -542,10 +542,9 @@ class HttpServer { // A request's restore points (prefix_restore_points) on both serving // paths: its chat boundaries spaced for the backend, plus `cuts`. The - // generation prompt's own boundary stays only for a tool-result request - // on a backend that cuts prefill at every restore point. + // generation prompt's own boundary, where the request's snapshot goes, + // stays only on a backend that cuts prefill at every restore point. std::vector request_restore_points(const std::vector & prompt, - bool ends_with_tool_result, std::initializer_list cuts) const; GenerationCacheState prepare_generation_cache( const ParsedRequest & req, PreparedPrompt & prepared, diff --git a/server/src/server/prefix_cache.cpp b/server/src/server/prefix_cache.cpp index 6752480b2..c2ee5bc62 100644 --- a/server/src/server/prefix_cache.cpp +++ b/server/src/server/prefix_cache.cpp @@ -240,6 +240,19 @@ int select_inline_evict_victim(const std::vector *> & } return false; }; + auto has_ancestor = [&](int i) { + for (int j = 0; j < n; j++) { + if (j == i) continue; + if (is_strict_prefix(*ids_lru[j], *ids_lru[i])) return true; + } + return false; + }; + // A turn kept between a conversation's head and its newest snapshot only + // serves an edit of the last message: it goes first. + for (int i = 0; i < n; i++) { + if (i == skip_index || is_protected(i)) continue; + if (is_ancestor(i) && has_ancestor(i)) return i; + } // Oldest-first scan: prefer an unprotected leaf so sticky tools pins // survive. skip_index (the in-flight restore source) is never a victim. int oldest_protected_leaf = -1; @@ -286,25 +299,21 @@ int select_inline_evict_victim(const std::vector> & ids_lru int select_inline_snapshot_boundary(const std::vector & boundaries, int restored_prefix_len, - bool prefer_tools_boundary, - bool include_last_message, + bool prefer_head, int reachable_from) { if (boundaries.empty()) return 0; const auto usable = [&](int cut) { return cut > restored_prefix_len && cut >= reachable_from; }; - // Tool-heavy cold path: pin the system+tools head (first marker) before - // deepening into conversation turns. Matches Python thin-pin semantics. - // A head the backend cannot save is skipped: it is cheap to prefill, and - // waiting for it would keep the cache from ever saving a deeper cut. - if (prefer_tools_boundary) { - const int tools_cut = boundaries.front(); - if (usable(tools_cut)) return tools_cut; - } - const int target = boundaries.size() >= 2 && !include_last_message - ? boundaries[boundaries.size() - 2] - : boundaries.back(); - return usable(target) ? target : 0; + // Save the system/tools head (first marker) before deepening into + // conversation turns. A head the backend cannot save is skipped: it is + // cheap to prefill, and waiting for it would keep the cache from ever + // saving a deeper cut. + if (prefer_head) { + const int head = boundaries.front(); + if (usable(head)) return head; + } + return usable(boundaries.back()) ? boundaries.back() : 0; } std::vector spaced_restore_points(const std::vector & boundaries, int min_gap) { @@ -589,11 +598,27 @@ void PrefixCache::InlineReservation::take(InlineReservation && other) { } bool PrefixCache::inline_reservation_active(uint64_t id) const { - return id != 0 && active_inline_reservation_ == id; + if (id == 0) return false; + for (const auto & active : active_reservations_) { + if (active.id == id) return true; + } + return false; } void PrefixCache::release_inline_reservation(uint64_t id) { - if (inline_reservation_active(id)) active_inline_reservation_ = 0; + for (size_t i = 0; i < active_reservations_.size(); ++i) { + if (active_reservations_[i].id == id) { + active_reservations_.erase(active_reservations_.begin() + (std::ptrdiff_t)i); + return; + } + } +} + +bool PrefixCache::slot_reserved(int slot) const { + for (const auto & active : active_reservations_) { + if (active.slot == slot) return true; + } + return false; } namespace { @@ -620,27 +645,29 @@ PrefixCache::InlineReservation PrefixCache::reserve_inline_snap( int forced_cut, int restore_source_slot, InlineSnapshotSize estimate_bytes, - bool include_last_message, int reachable_from, int restore_point_spacing) { - if (disabled_ || active_inline_reservation_ != 0) return {}; + if (disabled_) return {}; const auto candidates = find_all_boundaries(prompt_ids, markers_); - // A long tail past a short system/tools head (an agent's first turn, cold - // or with only the shared head restored): snapshot at the last message - // boundary instead of the head, or the first follow-up re-prefills the - // whole turn. That boundary is the end of the prompt unless the template - // appends tokens after it. Only a short head gives up its own pin this - // way, so a new conversation that shares it re-prefills at most - // LUCE_PC_DEEP_FIRST_MAX_HEAD tokens (default 2048). - // LUCE_PC_DEEP_FIRST_MIN is the tail length that triggers it (default - // 4096 tokens; 0 keeps the head pin). Values that are not a whole - // non-negative integer keep the defaults. - // Only the tool-request path pins the head, so only it changes. A forced - // pin still ahead of the restored prefix (an explicit pin_end, or PPP's - // pin of a head seen before) wins and keeps its protection, and a prompt - // the resident budget could never hold keeps the head pin instead of - // saving nothing. + // A conversation's first turn (one message past a head it did not + // restore) saves its system head, which every new conversation on that + // head restores; each later turn saves its end. Tool requests keep their + // system/tools head pinned until it is restored (prefer_tools_boundary). + bool prefer_head = prefer_tools_boundary || + (candidates.size() == 2 && restored_prefix_len < candidates.front()); + // A long tail past a short head (a first turn, cold or with only the + // shared head restored): snapshot at the last message boundary instead of + // the head, or the first follow-up re-prefills the whole turn. Only a + // short head gives up its own pin this way, so a new conversation that + // shares it re-prefills at most LUCE_PC_DEEP_FIRST_MAX_HEAD tokens + // (default 2048). LUCE_PC_DEEP_FIRST_MIN is the tail length that triggers + // it (default 4096 tokens; 0 keeps the head pin). Values that are not a + // whole non-negative integer keep the defaults. + // A forced pin still ahead of the restored prefix (an explicit pin_end, or + // PPP's pin of a head seen before) wins and keeps its protection, and a + // prompt the resident budget could never hold keeps the head pin instead + // of saving nothing. static const int deep_first_min = env_nonneg_int("LUCE_PC_DEEP_FIRST_MIN", 4096); static const int deep_first_max_head = env_nonneg_int("LUCE_PC_DEEP_FIRST_MAX_HEAD", 2048); const auto whole_prompt_fits = [&]() { @@ -648,14 +675,14 @@ PrefixCache::InlineReservation PrefixCache::reserve_inline_snap( const size_t bytes = estimate_bytes ? estimate_bytes((int) prompt_ids.size()) : 0; return bytes != 0 && bytes <= max_resident_bytes_; }; - if (deep_first_min > 0 && prefer_tools_boundary && !candidates.empty() && + if (deep_first_min > 0 && prefer_head && !candidates.empty() && restored_prefix_len <= candidates.front() && forced_cut <= restored_prefix_len && candidates.front() <= deep_first_max_head && (int) prompt_ids.size() - candidates.front() >= deep_first_min && whole_prompt_fits()) { + prefer_head = false; prefer_tools_boundary = false; - include_last_message = true; } int target_cut = 0; bool forced = false; @@ -666,10 +693,10 @@ PrefixCache::InlineReservation PrefixCache::reserve_inline_snap( forced = true; } else { target_cut = select_inline_snapshot_boundary( - candidates, restored_prefix_len, prefer_tools_boundary, - include_last_message, reachable_from); + candidates, restored_prefix_len, prefer_head, reachable_from); } - if (restore_point_spacing > 0 && target_cut > 0) { + if (restore_point_spacing > 0 && target_cut > 0 && + target_cut != candidates.back()) { int point = 0; for (int p : spaced_restore_points(candidates, restore_point_spacing)) if (p <= target_cut) point = p; @@ -695,13 +722,23 @@ PrefixCache::InlineReservation PrefixCache::reserve_inline_snap( const auto key = hash_prefix(prompt_ids.data(), target_cut); if (find_entry(key) >= 0) return {}; + for (const auto & active : active_reservations_) { + if (active.key == key) return {}; + } const bool protect = prefer_tools_boundary && (forced || (!candidates.empty() && target_cut == candidates.front())); PrefixHash victim_key{}; bool has_victim = false; int slot = -1; - if ((int)entries_.size() >= cap_) { + // Captures in flight hold their slots and budget until they resolve. + int reserved_free_slots = 0; + size_t pending_bytes = 0; + for (const auto & active : active_reservations_) { + if (!active.has_victim) ++reserved_free_slots; + pending_bytes += active.bytes; + } + if ((int)entries_.size() + reserved_free_slots >= cap_) { // At capacity — reserve a slot without evicting yet. Prefix-aware: prefer // the oldest leaf so shared ancestor prefixes (reused by later branches) // stay resident. Skip protected tools pins when an unprotected leaf @@ -734,20 +771,20 @@ PrefixCache::InlineReservation PrefixCache::reserve_inline_snap( // Skip the in-flight restore source too, so the new snapshot lands // in a different slot (the http_server/agent-replay guards would // cancel an unlucky collision, leaving the restore point pinned). - slot = next_slot_; for (int step = 0; step < cap_; ++step) { const int candidate = (next_slot_ + step) % cap_; if (candidate == restore_source_slot && cap_ > 1) continue; - if (find_slot_entry(candidate) >= 0) continue; + if (find_slot_entry(candidate) >= 0 || slot_reserved(candidate)) continue; slot = candidate; break; } + if (slot < 0) return {}; next_slot_ = (slot + 1) % cap_; } + size_t estimated_bytes = 0; if (max_resident_bytes_ > 0) { - const size_t estimated_bytes = estimate_bytes - ? estimate_bytes(target_cut) : 0; + estimated_bytes = estimate_bytes ? estimate_bytes(target_cut) : 0; // When the caller prunes after every commit, the entries this // capture supersedes (including its restore source) are freed as it // lands, so they do not compete with it for the budget. A failed @@ -759,6 +796,7 @@ PrefixCache::InlineReservation PrefixCache::reserve_inline_snap( if (prunes_superseded_) { for (const int i : superseded_entries( prompt_ids.data(), (size_t)target_cut)) { + if (slot_reserved(entries_[(size_t)i].slot)) continue; reclaimed[(size_t)i] = true; reclaimed_bytes += entries_[(size_t)i].resident_bytes; } @@ -770,8 +808,8 @@ PrefixCache::InlineReservation PrefixCache::reserve_inline_snap( if (victim_idx >= 0 && !reclaimed[(size_t)victim_idx]) { freed += entries_[(size_t)victim_idx].resident_bytes; } - const size_t after_free = freed <= resident_bytes_ - ? resident_bytes_ - freed : 0; + const size_t held = resident_bytes_ + pending_bytes; + const size_t after_free = freed <= held ? held - freed : 0; return after_free <= max_resident_bytes_ && estimated_bytes <= max_resident_bytes_ - after_free; }; @@ -789,15 +827,32 @@ PrefixCache::InlineReservation PrefixCache::reserve_inline_snap( } return true; }; + const auto is_middle = [&](int candidate) { + if (is_leaf(candidate)) return false; + for (int i = 0; i < (int)entries_.size(); ++i) { + if (i != candidate && + is_strict_prefix(entries_[(size_t)i].ids, + entries_[(size_t)candidate].ids)) { + return true; + } + } + return false; + }; + // A turn kept behind a conversation's newest snapshot first (see + // select_inline_evict_victim), then leaves, then the rest; + // unprotected entries before protected ones. int victim = -1; - for (int pass = 0; pass < 4 && victim < 0; ++pass) { - const bool require_leaf = pass < 2; - const bool allow_protected = pass == 1 || pass == 3; + for (int pass = 0; pass < 5 && victim < 0; ++pass) { + const bool require_middle = pass == 0; + const bool require_leaf = pass == 1 || pass == 2; + const bool allow_protected = pass == 2 || pass == 4; for (int i = 0; i < (int)entries_.size(); ++i) { if (restore_source_slot >= 0 && entries_[(size_t)i].slot == restore_source_slot) continue; + if (slot_reserved(entries_[(size_t)i].slot)) continue; if (!fits(i)) continue; + if (require_middle && !is_middle(i)) continue; if (require_leaf && !is_leaf(i)) continue; if (!allow_protected && entries_[(size_t)i].protect) continue; @@ -832,7 +887,8 @@ PrefixCache::InlineReservation PrefixCache::reserve_inline_snap( uint64_t id = next_inline_reservation_++; if (id == 0) id = next_inline_reservation_++; - active_inline_reservation_ = id; + active_reservations_.push_back( + {id, slot, key, victim_key, has_victim, estimated_bytes}); return InlineReservation( this, id, slot, target_cut, victim_key, has_victim, protect); } @@ -907,7 +963,7 @@ void PrefixCache::confirm_inline_snap( bool protect, size_t resident_bytes) { if (disabled_ || slot < 0 || target_cut <= 0 || target_cut > (int)prompt_ids.size()) return; - if (active_inline_reservation_ != 0) { + if (!active_reservations_.empty()) { std::fprintf(stderr, "[pc] direct commit refused while a reservation is active\n"); return; @@ -924,24 +980,31 @@ void PrefixCache::invalidate_inline_snap(int slot) { } int PrefixCache::pick_evict_victim(int skip_index) const { + // Every slot can be held by captures in flight, with nothing committed yet. + if (entries_.empty()) return -1; std::vector *> ids_lru; std::vector protected_lru; ids_lru.reserve(entries_.size()); protected_lru.reserve(entries_.size()); for (const auto & entry : entries_) { ids_lru.push_back(&entry.ids); - protected_lru.push_back(entry.protect); + // A capture in flight writes over its victim's slot: never a second victim. + protected_lru.push_back(entry.protect || slot_reserved(entry.slot)); } - return select_inline_evict_victim(ids_lru, &protected_lru, skip_index); + const int victim = select_inline_evict_victim(ids_lru, &protected_lru, skip_index); + return victim >= 0 && victim < (int)entries_.size() && + slot_reserved(entries_[(size_t)victim].slot) ? -1 : victim; } std::vector PrefixCache::superseded_entries( const int32_t * ids, size_t len) const { - // An agent conversation saves one snapshot per turn, each a strict - // prefix of the next. Restores only ever use the deepest one, so the - // intermediate turns are dead weight. Keep the shallowest ancestor (the - // system/tools head that a compacted or sibling conversation reuses) - // and any protected pin. + // A conversation saves one snapshot per turn, each a strict prefix of the + // next. Restores use the deepest one, or the one before it when a person + // edits their last message; the turns before those are dead weight. Keep + // the shallowest ancestor (the system/tools head that a compacted or + // sibling conversation reuses) and any protected pin. The previous turn + // is in the list: prune_superseded_ancestors keeps it while the budget + // holds it, and a capture's budget counts it as freed. std::vector out; int shallowest = -1; for (int i = 0; i < (int)entries_.size(); ++i) { @@ -957,20 +1020,38 @@ std::vector PrefixCache::superseded_entries( return out; } -std::vector PrefixCache::prune_superseded_ancestors(int slot) { +std::vector PrefixCache::prune_superseded_ancestors( + int slot, std::vector * pruned_tokens) { std::vector pruned; if (disabled_) return pruned; const int newest = find_slot_entry(slot); if (newest < 0) return pruned; const auto & newest_ids = entries_[(size_t)newest].ids; const size_t newest_len = newest_ids.size(); - const auto superseded = superseded_entries(newest_ids.data(), newest_len); + auto superseded = superseded_entries(newest_ids.data(), newest_len); + // Keep the previous turn, the deepest of them, if the budget holds it + // once the others are gone. + int previous = -1; + size_t others = 0; + for (const int i : superseded) { + if (previous < 0 || + entries_[(size_t)i].ids.size() > entries_[(size_t)previous].ids.size()) { + previous = i; + } + others += entries_[(size_t)i].resident_bytes; + } + if (previous >= 0) others -= entries_[(size_t)previous].resident_bytes; + if (previous >= 0 && + (max_resident_bytes_ == 0 || resident_bytes_ - others <= max_resident_bytes_)) { + superseded.erase(std::find(superseded.begin(), superseded.end(), previous)); + } size_t freed = 0; // Erase from the back so earlier indices stay valid. for (auto it = superseded.rbegin(); it != superseded.rend(); ++it) { const auto & entry = entries_[(size_t)*it]; freed += entry.resident_bytes; pruned.push_back(entry.slot); + if (pruned_tokens) pruned_tokens->push_back((int)entry.ids.size()); erase_inline_entry(*it); } if (!pruned.empty()) { @@ -1040,7 +1121,7 @@ void PrefixCache::mark_all_cleared() { resident_bytes_ = 0; resident_bytes_count_.store(0, std::memory_order_relaxed); next_slot_ = 0; - active_inline_reservation_ = 0; + active_reservations_.clear(); std::fprintf(stderr, "[pc] all-cleared — dropped %d LRU entries\n", n); } diff --git a/server/src/server/prefix_cache.h b/server/src/server/prefix_cache.h index 55200e90a..73ce911aa 100644 --- a/server/src/server/prefix_cache.h +++ b/server/src/server/prefix_cache.h @@ -52,8 +52,11 @@ PrefixHash hash_prefix(const int32_t * ids, int count); // Prefix-aware inline eviction: given cached prefixes in LRU order (0 = oldest), // return the index of the oldest "leaf" — an entry that is not a strict prefix -// of any other — so shared ancestors stay resident. Pointer overload is the -// core (no token copies); the value overload is for tests. +// of any other — so shared ancestors stay resident. A turn kept between a +// conversation's head and its newest snapshot (see +// PrefixCache::prune_superseded_ancestors) goes before any leaf: it only +// serves an edit of the last message. Pointer overload is the core (no token +// copies); the value overload is for tests. // // protected_lru (optional, same size): entries marked true are skipped. // Without skip_index, if every leaf is protected, the oldest protected leaf @@ -73,23 +76,21 @@ int select_inline_evict_victim(const std::vector> & ids_lru const std::vector * protected_lru = nullptr, int skip_index = -1); -// Pick the inline snapshot boundary for a request. -// Default: boundary before the current user turn (second-to-last marker), -// only when it advances past an already-restored prefix. A person may edit -// or resend that turn, so the snapshot stops before it. -// include_last_message cuts at the last marker instead, the start of the -// generation prompt. Use it when the prompt ends with tool results: clients -// append to them and never rewrite them, so the next request reuses them. -// When prefer_tools_boundary is set (tool-heavy agent requests), prefer the -// first marker (system+tools head) until that cut is already restored — this -// is the sticky "thin pin" Python tool-split used to keep under multi-chat -// eviction. Returns 0 when there is no useful new boundary. +// Pick the inline snapshot boundary for a request: the last marker, the end +// of the prompt's last message and the start of the generation prompt, only +// when it advances past an already-restored prefix. That is the point the +// next turn extends: a person's follow-up, or an agent's next step after its +// tool results. A person who edits that message resumes from the previous +// turn's snapshot instead (PrefixCache::prune_superseded_ancestors keeps it). +// When prefer_head is set (a conversation's first turn, or a tool-heavy agent +// request), prefer the first marker (the system/tools head) until that cut is +// already restored, so new conversations on the same head restore it. +// Returns 0 when there is no useful new boundary. // Cuts below `reachable_from` (the first position the backend can save at, // see ModelBackend::snapshot_granularity) are skipped. int select_inline_snapshot_boundary(const std::vector & boundaries, int restored_prefix_len = 0, - bool prefer_tools_boundary = false, - bool include_last_message = false, + bool prefer_head = false, int reachable_from = 0); // Restore points spaced at least `min_gap` tokens apart: the first boundary @@ -216,22 +217,23 @@ class PrefixCache { } // Select a boundary, destination, and optional budget victim as one owned - // operation. At most one reservation can be live; destroying it cancels + // operation. Each prefilling request can hold one; a live reservation's + // slot and bytes are out of reach of the others. Destroying it cancels // without changing committed metadata. // `restored_prefix_len` prevents reserving a slot for a boundary already - // covered by the restored snapshot. `prefer_tools_boundary` selects the - // system/tools head first (see select_inline_snapshot_boundary). When + // covered by the restored snapshot. `prefer_tools_boundary` pins the + // system/tools head first (see select_inline_snapshot_boundary), as a + // conversation's first turn saves it too. When // `forced_cut` > restored, that cut is used instead (PPP pin_end, // including mid-message LCP cuts). `restore_source_slot` (default -1) is // the slot this request restores from; it is never chosen as the // eviction or budget victim and the free-slot path skips it, so the new // snapshot lands in a different slot and the restore point can slide - // forward past the deepest slot. `include_last_message` and - // `reachable_from` are forwarded to select_inline_snapshot_boundary; a - // forced cut below `reachable_from` is not forced. With a - // `restore_point_spacing` (ModelBackend::restore_point_spacing), the cut - // moves down to the deepest spaced restore point, and nothing is reserved - // when that adds no prefix. + // forward past the deepest slot. `reachable_from` is forwarded to + // select_inline_snapshot_boundary; a forced cut below it is not forced. + // With a `restore_point_spacing` (ModelBackend::restore_point_spacing), a + // cut other than the end of the last message moves down to the deepest + // spaced restore point, and nothing is reserved when that adds no prefix. InlineReservation reserve_inline_snap( const std::vector & prompt_ids, int restored_prefix_len = 0, @@ -239,7 +241,6 @@ class PrefixCache { int forced_cut = 0, int restore_source_slot = -1, InlineSnapshotSize estimate_bytes = {}, - bool include_last_message = false, int reachable_from = 0, int restore_point_spacing = 0); @@ -254,9 +255,13 @@ class PrefixCache { void invalidate_inline_snap(int slot); // Drop the entries that the entry committed in `slot` supersedes: its - // strict-prefix ancestors, except the shallowest one and protected pins. - // Returns the dropped slots; the caller frees their snapshot payloads. - std::vector prune_superseded_ancestors(int slot); + // strict-prefix ancestors, except the shallowest one, protected pins and, + // while the resident budget holds it, the deepest one: the previous turn, + // where an edit of the last message resumes. + // Returns the dropped slots (and, with `pruned_tokens`, each one's prefix + // length); the caller frees their snapshot payloads. + std::vector prune_superseded_ancestors( + int slot, std::vector * pruned_tokens = nullptr); // Evicts committed entries, never `keep_slot` or protected pins, until the // resident bytes fit the budget again. Returns the evicted slots; the @@ -344,7 +349,18 @@ class PrefixCache { }; std::vector entries_; int next_slot_ = 0; - uint64_t active_inline_reservation_ = 0; + // Captures in flight, one per prefilling request. Each holds its slot (a + // free one, or its budget victim's) and its estimated bytes until it + // commits, aborts or cancels; no other capture selects that slot. + struct ActiveReservation { + uint64_t id; + int slot; + PrefixHash key; + PrefixHash victim; + bool has_victim; + size_t bytes; + }; + std::vector active_reservations_; uint64_t next_inline_reservation_ = 1; size_t max_resident_bytes_ = 0; bool prunes_superseded_ = false; @@ -406,6 +422,8 @@ class PrefixCache { bool record_hit); bool inline_reservation_active(uint64_t id) const; void release_inline_reservation(uint64_t id); + // A capture in flight writes this slot. + bool slot_reserved(int slot) const; bool commit_inline_reservation(InlineReservation & reservation, const std::vector & prompt_ids, int committed_cut, size_t resident_bytes, diff --git a/server/src/server/scheduler.cpp b/server/src/server/scheduler.cpp index 3c27ed83d..0d228e306 100644 --- a/server/src/server/scheduler.cpp +++ b/server/src/server/scheduler.cpp @@ -17,6 +17,7 @@ #include #include #include +#include #include #include @@ -36,6 +37,12 @@ struct SchedSlot { SocketHandle fd = kInvalidSocket; std::unique_ptr emitter; bool prefilling = false; + // Prompt tokens still to prefill, less each step's grant (a step can end + // its chunk early, so it may run low), and the new tokens the prompt + // brought (StepPlanLimits::time_share_prompt_tokens). A recomputed + // request counts as long. + int prompt_left = 0; + int prompt_new = 0; uint64_t admission_order = 0; std::chrono::steady_clock::time_point started_at{}; std::chrono::steady_clock::time_point decode_started_at{}; @@ -472,8 +479,7 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { [&engine](int target_cut) { return engine.estimate_prefix_store_bytes(target_cut); }, - req.ends_with_tool_result, /*reachable_from=*/0, - backend_.restore_point_spacing()); + /*reachable_from=*/0, backend_.restore_point_spacing()); if (capture_reservation.active()) { const uint64_t capture_id = next_prefix_capture_id++; if (next_prefix_capture_id == 0) @@ -493,7 +499,7 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { // The same cuts on hits and misses, whatever the cache holds (as the // single-sequence path passes GenerateRequest::restore_points). prefix_plan.restore_points = request_restore_points( - req.prompt_tokens, req.ends_with_tool_result, + req.prompt_tokens, {prefix_plan.restore.tokens, prefix_plan.capture.checkpoint.tokens}); } if (prefix_plan.capture.valid()) { @@ -663,6 +669,8 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { s.started_at = started_at; s.decode_started_at = started_at; // sane on prefill failure s.cached_prefix_tokens = prefix.restored.tokens; + s.prompt_left = (int)req.prompt_tokens.size() - prefix.restored.tokens; + s.prompt_new = s.prompt_left; s.cache_capture = std::move(prepared_capture); s.n_gen_cap = std::min( n_gen_cap, @@ -703,6 +711,10 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { std::vector prefill_candidates; prefill_candidates.reserve((size_t)n_slots); size_t prefill_round_robin_start = 0; + // The decoding requests' turn beside a long prompt + // (StepPlanLimits::prefill_time_share): no prompt chunk joins a step that + // has them until it ends. + std::chrono::steady_clock::time_point decode_only_until{}; while (true) { // Phase 1 — Admission: deferred job first (FIFO), then the queue. @@ -823,6 +835,8 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { auto & s = slots[(size_t)oldest]; if (oldest_state.recompute) { s.prefilling = true; + s.prompt_left = std::numeric_limits::max(); + s.prompt_new = std::numeric_limits::max(); s.pending_tok = -1; std::fprintf(stderr, "[parallel] slot %d resumed via re-prefill\n", @@ -876,9 +890,24 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { prefill_candidates.push_back({i, s.admission_order}); } } - step_plan.prefills = plan_prefill_slices(prefill_candidates, - engine.step_plan_limits((int)step_plan.decode.size()), - prefill_round_robin_start); + const bool decoders_turn = !step_plan.decode.empty() && + std::chrono::steady_clock::now() < decode_only_until; + const auto limits = engine.step_plan_limits((int)step_plan.decode.size()); + // Prompts that finish in one step go first, without the longer + // ones: a question asked beside a long document answers after one + // short forward, not after the document's next chunk. + const auto fits_one_step = [&](const PrefillCandidate & c) { + return slots[(size_t)c.slot].prompt_left <= limits.max_prefill_tokens_per_sequence; + }; + if (std::any_of(prefill_candidates.begin(), prefill_candidates.end(), fits_one_step) && + !std::all_of(prefill_candidates.begin(), prefill_candidates.end(), fits_one_step)) { + prefill_candidates.erase(std::remove_if(prefill_candidates.begin(), prefill_candidates.end(), + [&](const PrefillCandidate & c) { return !fits_one_step(c); }), + prefill_candidates.end()); + } + step_plan.prefills = decoders_turn + ? std::vector{} + : plan_prefill_slices(prefill_candidates, limits, prefill_round_robin_start); }; build_plan(); while (!engine.reserve_decode(step_plan)) { @@ -944,7 +973,31 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { if (!prefill_candidates.empty()) ++prefill_round_robin_start; if (step_plan.decode.empty() && step_plan.prefills.empty()) continue; + const auto step_started = std::chrono::steady_clock::now(); SeqEngine::StepResult step_result = engine.step(step_plan); + // A long prompt shares the time with the decoding requests, from its + // first chunk to its last; a short one (an agent's next turn, even when + // a prefix capture splits it in two forwards) does not hold them up. + const auto step_limits = engine.step_plan_limits((int)step_plan.decode.size()); + const int long_prompt_tokens = step_limits.time_share_prompt_tokens > 0 + ? step_limits.time_share_prompt_tokens + : step_limits.max_prefill_tokens_per_sequence; + bool long_prompt = false; + for (const PrefillSlice & slice : step_plan.prefills) { + SchedSlot & s = slots[(size_t)slice.slot]; + long_prompt |= s.prompt_new > long_prompt_tokens; + if (s.prompt_left != std::numeric_limits::max()) + s.prompt_left = std::max(0, s.prompt_left - slice.max_tokens); + } + if (long_prompt && !step_plan.decode.empty()) { + const double share = step_limits.prefill_time_share; + if (share > 0.0 && share < 1.0) { + const auto now = std::chrono::steady_clock::now(); + decode_only_until = now + std::chrono::duration_cast< + std::chrono::steady_clock::duration>( + (now - step_started) * ((1.0 - share) / share)); + } + } const std::string protocol_error = validate_step_result(step_plan, step_result, n_slots); if (!protocol_error.empty()) { @@ -997,7 +1050,17 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { using Resolution = PrefixCaptureTxn::Resolution; const Resolution resolution = s.cache_capture.resolve( out.prefix_store, s.job->req.prompt_tokens); - if (resolution == Resolution::failed) { + if (resolution == Resolution::saved) { + // A conversation's checkpoints are prefixes of each other; + // restores use only the deepest, so the ones it supersedes + // would hold memory the other conversations need. + std::vector tokens; + const auto pruned = prefix_cache_.prune_superseded_ancestors( + (int)out.prefix_store.ticket.checkpoint.id - 1, &tokens); + for (size_t i = 0; i < pruned.size(); ++i) { + engine.discard_prefix_store({(uint64_t)pruned[i] + 1, tokens[i]}); + } + } else if (resolution == Resolution::failed) { std::fprintf(stderr, "[parallel-pc] capture failed checkpoint=%llu: %s\n", (unsigned long long) diff --git a/server/src/server/sse_emitter.cpp b/server/src/server/sse_emitter.cpp index 40d5129f2..895fc53c6 100644 --- a/server/src/server/sse_emitter.cpp +++ b/server/src/server/sse_emitter.cpp @@ -455,6 +455,17 @@ std::vector SseEmitter::emit_token(const std::string & raw_piece) { } // mode_ == StreamMode::CONTENT + // The reply starts at its first line of text: the newlines a template + // puts after are not part of it (apps would show them as + // blank lines on top). Spaces stay: they can indent its first line. + if (accumulated_content_.empty()) { + const size_t first = window_.find_first_not_of("\r\n"); + if (first == std::string::npos) { + window_.clear(); + break; + } + window_.erase(0, first); + } // Look for , , or supported tool-call starts. size_t think_idx = window_.find(THINK_OPEN); size_t think_close_idx = window_.find(THINK_CLOSE); @@ -964,6 +975,8 @@ std::vector SseEmitter::emit_finish(int completion_tokens, {"total_tokens", prompt_tokens_ + completion_tokens} }; if (timings) { + // OpenAI's field for the prompt tokens a cache served. + usage_body["prompt_tokens_details"] = {{"cached_tokens", timings->cached_prefix_tokens}}; usage_body["timings"] = build_timings_json(*timings, completion_tokens); } json usage = { @@ -1066,6 +1079,7 @@ std::vector SseEmitter::emit_finish(int completion_tokens, {"total_tokens", prompt_tokens_ + completion_tokens} }; if (timings) { + resp_usage["input_tokens_details"] = {{"cached_tokens", timings->cached_prefix_tokens}}; resp_usage["timings"] = build_timings_json(*timings, completion_tokens); } json shell = { diff --git a/server/test/unit/test_server_unit.cpp b/server/test/unit/test_server_unit.cpp index f74e2af70..6e765595f 100644 --- a/server/test/unit/test_server_unit.cpp +++ b/server/test/unit/test_server_unit.cpp @@ -66,7 +66,10 @@ #include #include #include +#include +#include #include +#include #include #include #if !defined(_WIN32) @@ -307,6 +310,27 @@ TEST_CASE(ServerUnitFixture, test_pflash_scorer_uses_user_query_before_chat_suff TEST_ASSERT((int)rendered.size() - window.end == 10); } +// Claude Code sends its environment as a system message after the user's +// first one; a template that takes a system message only first gets it as a +// user turn in place, and one that takes it anywhere keeps it. +TEST_CASE(ServerUnitFixture, test_template_gets_late_system_message_as_user_turn) { + const std::string strict = + "{%- for m in messages -%}" + "{%- if m.role == 'system' and not loop.first -%}" + "{{- raise_exception('System message must be at the beginning.') -}}" + "{%- endif -%}" + "[{{ m.role }}:{{ m.content }}]" + "{%- endfor -%}"; + const std::string permissive = + "{%- for m in messages -%}[{{ m.role }}:{{ m.content }}]{%- endfor -%}"; + const std::vector msgs = { + {"system", "agent"}, {"user", "hi"}, {"system", "env"}, {"developer", "brief"}}; + TEST_ASSERT(render_chat_template_jinja(strict, msgs, "", "", false, false, "") == + "[system:agent][user:hi][user:env][user:brief]"); + TEST_ASSERT(render_chat_template_jinja(permissive, msgs, "", "", false, false, "") == + "[system:agent][user:hi][system:env][developer:brief]"); +} + TEST_CASE(ServerUnitFixture, test_pflash_scorer_accepts_responses_string_input) { ToolMemory tool_memory; const auto messages = normalize_chat_messages( @@ -2749,7 +2773,8 @@ TEST_CASE(ServerUnitFixture, test_emitter_funcname_tool_buffer_detection) { TEST_ASSERT(args["path"] == "/tmp/tool-input.md"); } TEST_ASSERT(em.accumulated_text().find("") == std::string::npos); - TEST_ASSERT(em.accumulated_text() == "\n\n"); + // The newlines before the call are not reply text. + TEST_ASSERT(em.accumulated_text().empty()); TEST_ASSERT(wire.find("\"finish_reason\":\"tool_calls\"") != std::string::npos); } @@ -3395,17 +3420,18 @@ TEST_CASE(ServerUnitFixture, test_resolve_deepseek_chat_markers) { }; TEST_ASSERT(find_all_boundaries(prompt, markers) == std::vector({3, 5, 8, 10})); - // The default snapshot cut stays before the current user turn. + // The snapshot cut is the end of the last message, which the next turn + // extends. TEST_ASSERT(select_inline_snapshot_boundary( - find_all_boundaries(prompt, markers)) == 8); + find_all_boundaries(prompt, markers)) == 10); - // A first turn snapshots its system text, so a new session on the same + // A first turn prefers its system text, so a new session on the same // system prompt restores it instead of prefilling it again. const std::vector first_turn = {1, 100, 3, 101, 4}; TEST_ASSERT(find_all_boundaries(first_turn, markers) == std::vector({3, 5})); TEST_ASSERT(select_inline_snapshot_boundary( - find_all_boundaries(first_turn, markers)) == 3); + find_all_boundaries(first_turn, markers), 0, true) == 3); // Tool-heavy requests pin the same system head. TEST_ASSERT(select_inline_snapshot_boundary( find_all_boundaries(prompt, markers), 0, true) == 3); @@ -3524,14 +3550,18 @@ TEST_CASE(ServerUnitFixture, test_prefix_cache_long_first_turn_snapshots_whole_p r.cancel(); } - // A request without tools never pinned the head; it keeps its usual cut - // (the start of the last message) rather than a whole-prompt snapshot. + // A first turn without tools saves the same way: its long message whole, + // and the head when the message is short. { PrefixCache cache(2, tokenizer); auto r = cache.reserve_inline_snap( short_head, /*restored_prefix_len=*/0, /*prefer_tools_boundary=*/false); TEST_ASSERT(r.active()); + TEST_ASSERT(r.target_cut() == (int) short_head.size()); + r.cancel(); + r = cache.reserve_inline_snap(short_tail); + TEST_ASSERT(r.active()); TEST_ASSERT(r.target_cut() == 3); r.cancel(); } @@ -3612,10 +3642,12 @@ TEST_CASE(ServerUnitFixture, test_restore_invalidation_preserves_pending_pin) { /*prefer_tools_boundary=*/true, /*forced_cut=*/2); TEST_ASSERT(prepared.slot() == 0); TEST_ASSERT(prepared.target_cut() == 2); - auto blocked = cache.reserve_inline_snap( + // Another request's capture can run beside it, never in its slot. + auto beside = cache.reserve_inline_snap( next, /*restored_prefix_len=*/0, /*prefer_tools_boundary=*/false, /*forced_cut=*/2); - TEST_ASSERT(!blocked.active()); + TEST_ASSERT(beside.active() && beside.slot() == 1); + beside.cancel(); cache.invalidate_inline_snap(/*slot=*/1); TEST_ASSERT(prepared.commit(pinned)); @@ -4019,28 +4051,17 @@ TEST_CASE(ServerUnitFixture, test_find_boundaries_clean_prompt_unchanged) { TEST_CASE(ServerUnitFixture, test_inline_snapshot_boundary_advances_past_restore) { const std::vector boundaries = {100, 240, 380, 520}; - // Second-to-last is the boundary before the current user turn. - TEST_ASSERT(select_inline_snapshot_boundary(boundaries) == 380); - TEST_ASSERT(select_inline_snapshot_boundary(boundaries, 240) == 380); + // The last boundary, the end of the last message: a person's follow-up + // and an agent's next step after its tool results both extend it. + TEST_ASSERT(select_inline_snapshot_boundary(boundaries) == 520); + TEST_ASSERT(select_inline_snapshot_boundary(boundaries, 380) == 520); // Do not reserve a snapshot when the restore already covers that point. - TEST_ASSERT(select_inline_snapshot_boundary(boundaries, 380) == 0); - TEST_ASSERT(select_inline_snapshot_boundary(boundaries, 500) == 0); + TEST_ASSERT(select_inline_snapshot_boundary(boundaries, 520) == 0); + TEST_ASSERT(select_inline_snapshot_boundary(boundaries, 600) == 0); TEST_ASSERT(select_inline_snapshot_boundary({}, 0) == 0); TEST_ASSERT(select_inline_snapshot_boundary({100}, 0) == 100); } -TEST_CASE(ServerUnitFixture, test_inline_snapshot_boundary_includes_trailing_tool_result) { - const std::vector boundaries = {100, 240, 380, 520}; - // A prompt ending in tool results is snapshotted through them, up to - // the generation prompt, because the next agent turn only appends. - TEST_ASSERT(select_inline_snapshot_boundary(boundaries, 0, false, true) == 520); - TEST_ASSERT(select_inline_snapshot_boundary(boundaries, 380, false, true) == 520); - TEST_ASSERT(select_inline_snapshot_boundary(boundaries, 520, false, true) == 0); - // The cold tools head still comes first. - TEST_ASSERT(select_inline_snapshot_boundary(boundaries, 0, true, true) == 100); - TEST_ASSERT(select_inline_snapshot_boundary(boundaries, 100, true, true) == 520); -} - // Families with a trailing role marker can snapshot through the tool results. // Gemma instead continues the open model turn; its continuation is tested below. TEST_CASE(ServerUnitFixture, test_tool_result_prompt_prefix_survives_next_agent_turn) { @@ -4091,43 +4112,73 @@ TEST_CASE(ServerUnitFixture, test_spaced_restore_points) { // Prefix-stable: a shorter prompt keeps the longer prompt's points up to its end. TEST_ASSERT((spaced_restore_points({300, 900, 1500, 2600}, 2048) == std::vector{300, 2600})); - // A capture moves down to a spaced point, and none is reserved when that adds no prefix. + // A capture is the end of the last message, the point the next turn + // extends; with spacing, a cut short of it moves down to a spaced point, + // and none is reserved when that adds no prefix. const std::string path = write_deepseek_marker_tokenizer_fixture(); Tokenizer tokenizer; TEST_ASSERT(tokenizer.load_from_gguf(path.c_str())); PrefixCache cache(4, tokenizer); const std::vector prompt = {1, 100, 3, 101, 4, 102, 3}; const std::vector boundaries = find_all_boundaries(prompt, cache.chat_markers()); - TEST_ASSERT(boundaries.size() >= 2); + TEST_ASSERT(boundaries.size() >= 3); const int head = boundaries.front(); - auto plain = cache.reserve_inline_snap(prompt, 0, false, 0, -1, {}, /*include_last_message=*/true); - TEST_ASSERT(plain.active() && plain.target_cut() > head); + auto plain = cache.reserve_inline_snap(prompt); + TEST_ASSERT(plain.active() && plain.target_cut() == boundaries.back()); plain.cancel(); - // A spacing wider than the prompt keeps only the head. - auto moved = cache.reserve_inline_snap(prompt, 0, false, 0, -1, {}, true, 0, /*restore_point_spacing=*/64); + // A spacing wider than the whole prompt still keeps the last message's end. + auto last = cache.reserve_inline_snap(prompt, 0, false, 0, -1, {}, 0, /*restore_point_spacing=*/64); + TEST_ASSERT(last.active() && last.target_cut() == boundaries.back()); + last.cancel(); + // A pinned cut in the middle moves down to the head, the only spaced point. + auto moved = cache.reserve_inline_snap(prompt, 0, false, /*forced_cut=*/boundaries[1], -1, {}, 0, 64); TEST_ASSERT(moved.active() && moved.target_cut() == head); moved.cancel(); - auto none = cache.reserve_inline_snap(prompt, head, false, 0, -1, {}, true, 0, 64); + auto none = cache.reserve_inline_snap(prompt, boundaries.back(), false, 0, -1, {}, 0, 64); TEST_ASSERT(!none.active()); } +// Concurrent requests each capture their own prefix at once: a capture in +// flight holds its slot and bytes, and the same prefix is not captured twice. +TEST_CASE(ServerUnitFixture, test_inline_captures_in_flight_hold_their_slots) { + const std::string path = write_deepseek_marker_tokenizer_fixture(); + Tokenizer tokenizer; + TEST_ASSERT(tokenizer.load_from_gguf(path.c_str())); + PrefixCache cache(2, tokenizer, /*max_resident_bytes=*/1000); + const std::vector a = {1, 100, 3, 101, 4}; + const std::vector b = {1, 200, 3, 201, 4}; + const std::vector c = {1, 300, 3, 301, 4}; + const auto bytes = [](int) { return (size_t)400; }; + auto first = cache.reserve_inline_snap(a, 0, false, 0, -1, bytes); + auto twin = cache.reserve_inline_snap(a, 0, false, 0, -1, bytes); + auto second = cache.reserve_inline_snap(b, 0, false, 0, -1, bytes); + TEST_ASSERT(first.active() && !twin.active() && second.active()); + TEST_ASSERT(first.slot() != second.slot()); + // Both slots are held, and a third capture would not fit the budget anyway. + auto third = cache.reserve_inline_snap(c, 0, false, 0, -1, bytes); + TEST_ASSERT(!third.active()); + TEST_ASSERT(second.commit(b, 400)); + TEST_ASSERT(first.commit(a, 400)); + TEST_ASSERT(cache.lookup_candidate({1, 100, 3, 101, 4, 9}, 5).second > 0); + TEST_ASSERT(cache.lookup_candidate({1, 200, 3, 201, 4, 9}, 5).second > 0); + unlink(path.c_str()); +} + TEST_CASE(ServerUnitFixture, test_inline_snapshot_skips_unreachable_cuts) { const std::vector boundaries = {266, 700, 1300, 1800}; TEST_ASSERT(select_inline_snapshot_boundary( - boundaries, 0, true, false, /*reachable_from=*/512) == 1300); - TEST_ASSERT(select_inline_snapshot_boundary( - boundaries, 0, true, true, 512) == 1800); + boundaries, 0, true, /*reachable_from=*/512) == 1800); // A reachable head is still pinned first. TEST_ASSERT(select_inline_snapshot_boundary( - boundaries, 0, true, false, 1) == 266); + boundaries, 0, true, 1) == 266); // Nothing reachable yet: no reservation, prefill stays cold. TEST_ASSERT(select_inline_snapshot_boundary( - {100, 300}, 0, true, true, 512) == 0); + {100, 300}, 0, true, 512) == 0); // Restored at 1024 with 512-token chunks: 1300 lies in the first chunk. TEST_ASSERT(select_inline_snapshot_boundary( - boundaries, 1024, false, false, 1024 + 512) == 0); + {266, 1300}, 1024, false, 1024 + 512) == 0); TEST_ASSERT(select_inline_snapshot_boundary( - boundaries, 1024, false, true, 1024 + 512) == 1800); + boundaries, 1024, false, 1024 + 512) == 1800); // A PPP pin the backend cannot save is not forced. TEST_ASSERT(!should_force_inline_snapshot_boundary( boundaries, 2000, 0, true, 266, 512)); @@ -4142,7 +4193,7 @@ TEST_CASE(ServerUnitFixture, test_inline_snapshot_skips_unreachable_cuts) { const std::vector prompt = {1, 100, 3, 101, 4, 102, 3}; auto capture = cache.reserve_inline_snap( prompt, 0, /*prefer_tools_boundary=*/true, /*forced_cut=*/2, -1, {}, - /*include_last_message=*/false, /*reachable_from=*/3); + /*reachable_from=*/3); TEST_ASSERT(capture.active()); TEST_ASSERT(capture.target_cut() > 2); capture.cancel(); @@ -4163,9 +4214,9 @@ TEST_CASE(ServerUnitFixture, test_inline_snapshot_prefers_tools_boundary_until_r const std::vector boundaries = {100, 240, 380, 520}; // Cold tool-heavy: pin system+tools head (first marker), not deepen cut. TEST_ASSERT(select_inline_snapshot_boundary(boundaries, 0, true) == 100); - // After tools head is restored, deepen to second-to-last. - TEST_ASSERT(select_inline_snapshot_boundary(boundaries, 100, true) == 380); - TEST_ASSERT(select_inline_snapshot_boundary(boundaries, 380, true) == 0); + // After tools head is restored, deepen to the end of the last message. + TEST_ASSERT(select_inline_snapshot_boundary(boundaries, 100, true) == 520); + TEST_ASSERT(select_inline_snapshot_boundary(boundaries, 520, true) == 0); TEST_ASSERT(select_inline_snapshot_boundary({100}, 0, true) == 100); TEST_ASSERT(select_inline_snapshot_boundary({100}, 100, true) == 0); } @@ -4181,7 +4232,7 @@ TEST_CASE(ServerUnitFixture, test_forced_tools_pin_yields_to_deepen_after_restor // later conversation boundary instead of forcing the nearby pin again. TEST_ASSERT(!should_force_inline_snapshot_boundary( boundaries, 600, 100, true, 110)); - TEST_ASSERT(select_inline_snapshot_boundary(boundaries, 100, true) == 380); + TEST_ASSERT(select_inline_snapshot_boundary(boundaries, 100, true) == 520); // Without tools preference, a still-unrestored forced cut retains its // original behavior. @@ -4372,10 +4423,17 @@ TEST_CASE(ServerUnitFixture, test_evict_single_is_zero) { } TEST_CASE(ServerUnitFixture, test_evict_chain_keeps_ancestors) { - // Oldest-first chain: [s] < [s,a] < [s,a,b]. Only the longest is a leaf, so - // the short shared ancestors are kept and the victim is the deepest entry. + // Oldest-first chain: [s] < [s,a] < [s,a,b]. [s,a] sits between the head + // and the newest turn, kept only for an edit of the last message: it goes + // first. Then the leaf, and the shared head last. std::vector> ids = {{9}, {9, 1}, {9, 1, 2}}; - TEST_ASSERT(select_inline_evict_victim(ids) == 2); + TEST_ASSERT(select_inline_evict_victim(ids) == 1); + ids = {{9}, {9, 1, 2}}; + TEST_ASSERT(select_inline_evict_victim(ids) == 1); + // A protected turn in between stays; the leaf goes. + ids = {{9}, {9, 1}, {9, 1, 2}}; + std::vector protect = {false, true, false}; + TEST_ASSERT(select_inline_evict_victim(ids, &protect) == 2); } TEST_CASE(ServerUnitFixture, test_evict_unrelated_falls_back_to_lru) { @@ -4481,11 +4539,12 @@ TEST_CASE(ServerUnitFixture, test_slide_restore_source_never_evicted) { const int victim = select_inline_evict_victim(ids, nullptr, skip); TEST_ASSERT(victim >= 0 && victim != skip); } - // Linear chain whose only leaf is the restore source: evict the - // shallowest unprotected ancestor instead of cancelling everything. - TEST_ASSERT(select_inline_evict_victim(ids, nullptr, 3) == 0); - // With a free leaf the usual leaf preference still applies. - TEST_ASSERT(select_inline_evict_victim(ids, nullptr, 0) == 3); + // Linear chain whose only leaf is the restore source: evict a turn kept + // between the head and the newest instead of cancelling everything. + TEST_ASSERT(select_inline_evict_victim(ids, nullptr, 3) == 1); + // Otherwise the shallowest unprotected ancestor. + ids = {{9}, {9, 1}}; + TEST_ASSERT(select_inline_evict_victim(ids, nullptr, 1) == 0); } // (c) The protected tools pin is never evicted, even when it is the @@ -4659,25 +4718,40 @@ TEST_CASE(ServerUnitFixture, test_prune_superseded_ancestors) { cache.confirm_inline_snap(3, (int)sibling.size(), sibling, false, 20); cache.confirm_inline_snap(4, (int)t3.size(), t3, false, 40); + // The head and the previous turn stay; the turn before that goes. auto pruned = cache.prune_superseded_ancestors(4); - std::sort(pruned.begin(), pruned.end()); - TEST_ASSERT(pruned == std::vector({1, 2})); - TEST_ASSERT(cache.stats().in_use == 3); - TEST_ASSERT(cache.stats().resident_bytes == 70); + TEST_ASSERT(pruned == std::vector({1})); + TEST_ASSERT(cache.stats().in_use == 4); + TEST_ASSERT(cache.stats().resident_bytes == 100); TEST_ASSERT(cache.lookup(t3).first == 4); TEST_ASSERT(cache.lookup(sibling).first == 3); + // A person who edits the last message resumes from the previous turn. + std::vector edited = t2; + edited.insert(edited.end(), {3, 104}); + TEST_ASSERT(cache.lookup(edited).first == 2); // Pruning again finds nothing left to drop. TEST_ASSERT(cache.prune_superseded_ancestors(4).empty()); // A slot without an entry prunes nothing. TEST_ASSERT(cache.prune_superseded_ancestors(7).empty()); + // The previous turn goes too when the resident budget cannot hold it. + PrefixCache tight(8, tokenizer, /*max_resident_bytes=*/70); + tight.confirm_inline_snap(0, (int)head.size(), head, false, 10); + tight.confirm_inline_snap(1, (int)t2.size(), t2, false, 30); + tight.confirm_inline_snap(2, (int)t3.size(), t3, false, 40); + TEST_ASSERT(tight.prune_superseded_ancestors(2) == std::vector({1})); + TEST_ASSERT(tight.stats().resident_bytes == 50); + // A protected pin that is not the shallowest ancestor also survives. + std::vector t4 = t3; + t4.insert(t4.end(), {4, 104}); PrefixCache pinned(8, tokenizer); pinned.confirm_inline_snap(0, (int)head.size(), head); pinned.confirm_inline_snap(1, (int)t1.size(), t1, /*protect=*/true); pinned.confirm_inline_snap(2, (int)t2.size(), t2); pinned.confirm_inline_snap(3, (int)t3.size(), t3); - TEST_ASSERT(pinned.prune_superseded_ancestors(3) == std::vector({2})); + pinned.confirm_inline_snap(4, (int)t4.size(), t4); + TEST_ASSERT(pinned.prune_superseded_ancestors(4) == std::vector({2})); unlink(path.c_str()); } @@ -5104,6 +5178,10 @@ TEST_CASE(ServerUnitFixture, test_qwen4exp_thinks_by_default) { TEST_ASSERT(by_default.thinking_enabled && by_default.thinking_opt_in); const ParsedRequest off = resolve_qwen4exp_reasoning({{"thinking", {{"type", "disabled"}}}}); TEST_ASSERT(!off.thinking_enabled && !off.thinking_opt_in); + // Claude Code's "adaptive" leaves it to the model: the default. + const ParsedRequest adaptive = resolve_qwen4exp_reasoning({{"thinking", {{"type", "adaptive"}}}}); + TEST_ASSERT(adaptive.thinking_enabled && adaptive.thinking_opt_in); + TEST_ASSERT(!resolve_qwen_reasoning({{"thinking", {{"type", "adaptive"}}}}).thinking_enabled); const ParsedRequest kw_off = resolve_qwen4exp_reasoning({{"chat_template_kwargs", {{"enable_thinking", false}}}}); TEST_ASSERT(!kw_off.thinking_enabled && !kw_off.thinking_opt_in); TEST_ASSERT(!resolve_qwen_reasoning(json::object()).thinking_enabled); @@ -7513,24 +7591,29 @@ TEST_CASE(ServerUnitFixture, test_inline_snapshot_finalization_prunes_chain) { turn1.insert(turn1.end(), {3, 101}); std::vector turn2 = turn1; turn2.insert(turn2.end(), {4, 102}); + std::vector turn3 = turn2; + turn3.insert(turn3.end(), {3, 103}); cache.confirm_inline_snap(0, (int)head.size(), head); cache.confirm_inline_snap(1, (int)turn1.size(), turn1); + cache.confirm_inline_snap(2, (int)turn2.size(), turn2); auto reservation = cache.reserve_inline_snap( - turn2, (int)turn1.size(), /*prefer_tools_boundary=*/false, - /*forced_cut=*/(int)turn2.size(), /*restore_source_slot=*/1); + turn3, (int)turn2.size(), /*prefer_tools_boundary=*/false, + /*forced_cut=*/(int)turn3.size(), /*restore_source_slot=*/2); TEST_ASSERT(reservation.active()); const int slot = reservation.slot(); - TEST_ASSERT(slot != 0 && slot != 1); + TEST_ASSERT(slot != 0 && slot != 1 && slot != 2); backend.saved_slot = slot; - backend.saved_position = (int)turn2.size(); + backend.saved_position = (int)turn3.size(); SchedulerTestHarness::finalize_inline_snapshot( - server, turn2, std::move(reservation), slot, - /*requested_cut=*/(int)turn2.size()); + server, turn3, std::move(reservation), slot, + /*requested_cut=*/(int)turn3.size()); + // The head and the previous turn stay. TEST_ASSERT(backend.freed == std::vector({1})); - TEST_ASSERT(cache.stats().in_use == 2); - TEST_ASSERT(cache.lookup(turn2).first == slot); + TEST_ASSERT(cache.stats().in_use == 3); + TEST_ASSERT(cache.lookup(turn3).first == slot); + TEST_ASSERT(cache.lookup(turn2).first == 2); std::vector other = head; other.insert(other.end(), {3, 999}); TEST_ASSERT(cache.lookup(other).first == 0); @@ -8099,6 +8182,104 @@ struct SchedulerPrefixBackend : MockBackend { SeqEngine * seq_engine() override { return &engine; } }; +// Concurrent serving: two requests admitted together both capture, and a +// conversation's new checkpoint frees the ones it supersedes (keeping the +// shallowest, the shared head, and the previous turn for an edit). +TEST_CASE(ServerUnitFixture, + test_scheduler_captures_concurrently_and_prunes_superseded) { + const std::string path = write_deepseek_marker_tokenizer_fixture(); + Tokenizer tokenizer; + TEST_ASSERT(tokenizer.load_from_gguf(path.c_str())); + + auto backend_owner = std::make_unique(); + SchedulerPrefixBackend & backend = *backend_owner; + LuceEngine engine(std::move(backend_owner)); + ServerConfig config; + config.arch = "qwen35"; + config.max_ctx = 64; + config.prefix_cache_cap = 8; + config.concurrent_prefix_cache_max_bytes = 4096; + config.concurrent_prefix_cache = true; + config.admission_coalesce_ms = 0; + HttpServer server(engine, tokenizer, config); + PrefixCache & cache = SchedulerTestHarness::prefix_cache(server); + + std::vector> jobs; + std::vector fds; + const auto make_job = [&](std::vector prompt, const char * id) { + int sockets[2] = {-1, -1}; + TEST_ASSERT(socketpair(AF_UNIX, SOCK_STREAM, 0, sockets) == 0); + fds.push_back(sockets[0]); + fds.push_back(sockets[1]); + auto job = std::make_unique(); + job->fd = sockets[0]; + job->req.format = ApiFormat::OPENAI_CHAT; + job->req.prompt_tokens = std::move(prompt); + job->req.max_output = 1; + job->req.stream = false; + job->req.model = "scheduler-test"; + job->req.response_id = id; + jobs.push_back(std::move(job)); + return jobs.back().get(); + }; + const auto wait_done = [](ServerJob * job) { + std::unique_lock lock(job->mu); + return job->cv.wait_for(lock, std::chrono::seconds(5), [&] { return job->done; }); + }; + + // Admitted in one pass: both capture their head. + ServerJob * a = make_job({1, 100, 3, 101, 4}, "a"); + ServerJob * b = make_job({1, 200, 3, 201, 4}, "b"); + SchedulerTestHarness::enqueue(server, a); + SchedulerTestHarness::enqueue(server, b); + std::thread scheduler([&] { + SchedulerTestHarness::run(server, backend.engine); + }); + const bool ab_done = wait_done(a) && wait_done(b); + + // One conversation, three more turns, each saved at its end: the fourth + // turn's checkpoint supersedes the second's, which is freed; the head and + // the third turn stay. + const std::vector turn2 = {1, 100, 3, 101, 4, 102, 2, 3, 103, 4}; + std::vector turn3 = turn2; + turn3.insert(turn3.end(), {104, 2, 3, 105, 4}); + std::vector turn4 = turn3; + turn4.insert(turn4.end(), {106, 2, 3, 107, 4}); + ServerJob * second = make_job(turn2, "turn2"); + SchedulerTestHarness::enqueue(server, second); + const bool second_done = wait_done(second); + ServerJob * third = make_job(turn3, "turn3"); + SchedulerTestHarness::enqueue(server, third); + const bool third_done = wait_done(third); + ServerJob * fourth = make_job(turn4, "turn4"); + SchedulerTestHarness::enqueue(server, fourth); + const bool fourth_done = wait_done(fourth); + + SchedulerTestHarness::stop(server); + scheduler.join(); + for (const int fd : fds) close(fd); + unlink(path.c_str()); + + TEST_ASSERT(ab_done && second_done && third_done && fourth_done); + // Both heads were captured although their requests prefilled together. + TEST_ASSERT(cache.lookup_candidate({1, 100, 3, 101, 4, 9}, 5).second == 3); + TEST_ASSERT(cache.lookup_candidate({1, 200, 3, 201, 4, 9}, 5).second == 3); + std::vector next = turn4; + next.push_back(9); + TEST_ASSERT(cache.lookup_candidate(next, (int)next.size() - 1).second == 20); + // An edit of the last message resumes from the third turn's end. + std::vector edited = turn3; + edited.insert(edited.end(), {106, 2, 3, 999, 4}); + TEST_ASSERT(cache.lookup_candidate(edited, (int)edited.size() - 1).second == 15); + // The turn-2 checkpoint is gone from the cache and freed in the engine. + std::vector after_turn2 = turn2; + after_turn2.push_back(9); + TEST_ASSERT(cache.lookup_candidate(after_turn2, 10).second == 3); + bool freed_turn2 = false; + for (const auto & ref : backend.engine.discarded) freed_turn2 |= ref.tokens == 10; + TEST_ASSERT(freed_turn2); +} + TEST_CASE(ServerUnitFixture, test_scheduler_counts_restore_only_after_engine_attempts_it) { const std::string path = write_deepseek_marker_tokenizer_fixture(); @@ -8235,6 +8416,196 @@ TEST_CASE(ServerUnitFixture, TEST_ASSERT(stats.restore_stall_us_total == 2500); } +// A long prompt shares the time with a decoding request: after each of its +// chunks the decoder steps alone for as long as the chunk took. Prompt steps +// here take 40 ms and decode steps 2 ms. +class TimeShareEngine final : public SeqEngine { +public: + explicit TimeShareEngine(int long_prompt_tokens = 0) + : long_prompt_tokens_(long_prompt_tokens) {} + int slot_count() const override { return 2; } + int max_context() const override { return 4096; } + bool token_is_eos(int32_t) const override { return false; } + StepPlanLimits step_plan_limits(int decode_rows) const override { + return {decode_rows > 0 ? 1 : 2, 16, 32, 16, /*prefill_time_share=*/0.5, + long_prompt_tokens_}; + } + AdmitResult admit(uint64_t, const std::vector & prompt, const SamplerCfg &) override { + AdmitResult result; + const int slot = !slots_[0].active ? 0 : !slots_[1].active ? 1 : -1; + if (slot < 0) { + result.status = AdmitResult::Status::busy; + return result; + } + slots_[(size_t)slot] = {true, (int)prompt.size()}; + result.status = AdmitResult::Status::admitted; + result.slot = slot; + return result; + } + StepResult step(const StepPlan & plan) override { + StepResult result; + for (const StepInput & input : plan.decode) result.decode.push_back({input.slot, 7, false, {}}); + for (const PrefillSlice & slice : plan.prefills) { + Slot & slot = slots_[(size_t)slice.slot]; + slot.remaining -= std::min(slot.remaining, slice.max_tokens); + PrefillOutput out; + out.slot = slice.slot; + out.status = slot.remaining == 0 ? PrefillOutput::Status::completed + : PrefillOutput::Status::advanced; + if (slot.remaining == 0) out.token = 7; + result.prefills.push_back(out); + } + std::this_thread::sleep_for(std::chrono::milliseconds(plan.prefills.empty() ? 2 : 40)); + std::lock_guard lock(mu); + steps.push_back({(int)plan.decode.size(), (int)plan.prefills.size()}); + std::vector chunk_slots; + for (const PrefillSlice & slice : plan.prefills) chunk_slots.push_back(slice.slot); + prompt_slots.push_back(std::move(chunk_slots)); + return result; + } + void retire(int slot) override { slots_[(size_t)slot] = {}; } + + std::mutex mu; + std::vector> steps; // (decode rows, prompt chunks) + std::vector> prompt_slots; // the slots each step prefilled + +private: + struct Slot { bool active = false; int remaining = 0; }; + std::array slots_{}; + int long_prompt_tokens_; +}; + +struct TimeShareBackend : MockBackend { + explicit TimeShareBackend(int long_prompt_tokens = 0) : engine(long_prompt_tokens) {} + TimeShareEngine engine; + SeqEngine * seq_engine() override { return &engine; } +}; + +// A chat that decodes, then a 128-token prompt (8 chunks) beside it. Returns +// the most decode-only steps the chat took between two of the prompt's chunks. +static int decode_only_steps_between_chunks(int long_prompt_tokens) { + const std::string path = write_deepseek_marker_tokenizer_fixture(); + Tokenizer tokenizer; + TEST_ASSERT(tokenizer.load_from_gguf(path.c_str())); + auto backend_owner = std::make_unique(long_prompt_tokens); + TimeShareBackend & backend = *backend_owner; + LuceEngine engine(std::move(backend_owner)); + ServerConfig config; + config.arch = "qwen35"; + config.max_ctx = 4096; + config.prefix_cache_cap = 0; + config.admission_coalesce_ms = 0; + HttpServer server(engine, tokenizer, config); + + int fds[4] = {-1, -1, -1, -1}; + TEST_ASSERT(socketpair(AF_UNIX, SOCK_STREAM, 0, fds) == 0); + TEST_ASSERT(socketpair(AF_UNIX, SOCK_STREAM, 0, fds + 2) == 0); + const auto make_job = [](ServerJob & job, int fd, int prompt, int output, const char * id) { + job.fd = fd; + job.req.format = ApiFormat::OPENAI_CHAT; + job.req.prompt_tokens.assign((size_t)prompt, 100); + job.req.max_output = output; + job.req.stream = false; + job.req.model = "scheduler-test"; + job.req.response_id = id; + }; + ServerJob chat, document; + make_job(chat, fds[0], 4, 400, "chat"); + make_job(document, fds[2], 128, 1, "document"); + SchedulerTestHarness::enqueue(server, &chat); + std::thread scheduler([&] { SchedulerTestHarness::run(server, backend.engine); }); + std::this_thread::sleep_for(std::chrono::milliseconds(30)); + SchedulerTestHarness::enqueue(server, &document); + const auto wait_done = [](ServerJob & job) { + std::unique_lock lock(job.mu); + return job.cv.wait_for(lock, std::chrono::seconds(10), [&] { return job.done; }); + }; + const bool done = wait_done(document) && wait_done(chat); + SchedulerTestHarness::stop(server); + scheduler.join(); + for (const int fd : fds) close(fd); + unlink(path.c_str()); + TEST_ASSERT(done); + + int chunks_beside = 0, alone_between = 0, alone_run = 0; + bool after_chunk = false; + for (const auto & [decode, chunks] : backend.engine.steps) { + if (chunks > 0 && decode > 0) { + if (after_chunk) alone_between = std::max(alone_between, alone_run); + ++chunks_beside; + after_chunk = true; + alone_run = 0; + } else if (chunks == 0 && decode > 0 && after_chunk) { + ++alone_run; + } + } + TEST_ASSERT(chunks_beside >= 2); + return alone_between; +} + +TEST_CASE(ServerUnitFixture, test_scheduler_gives_decoders_time_beside_a_long_prompt) { + // Between two chunks of the long prompt that ran beside the chat, the chat + // stepped alone: several decode-only steps, not none. + TEST_ASSERT(decode_only_steps_between_chunks(/*long_prompt_tokens=*/0) >= 5); + // A prompt under the engine's long-prompt size keeps every step. + TEST_ASSERT(decode_only_steps_between_chunks(/*long_prompt_tokens=*/256) == 0); +} + +// A prompt that fits in one step runs before a longer one, alone. +TEST_CASE(ServerUnitFixture, test_scheduler_runs_a_short_prompt_before_a_long_one) { + const std::string path = write_deepseek_marker_tokenizer_fixture(); + Tokenizer tokenizer; + TEST_ASSERT(tokenizer.load_from_gguf(path.c_str())); + auto backend_owner = std::make_unique(); + TimeShareBackend & backend = *backend_owner; + LuceEngine engine(std::move(backend_owner)); + ServerConfig config; + config.arch = "qwen35"; + config.max_ctx = 4096; + config.prefix_cache_cap = 0; + config.admission_coalesce_ms = 0; + HttpServer server(engine, tokenizer, config); + + int fds[4] = {-1, -1, -1, -1}; + TEST_ASSERT(socketpair(AF_UNIX, SOCK_STREAM, 0, fds) == 0); + TEST_ASSERT(socketpair(AF_UNIX, SOCK_STREAM, 0, fds + 2) == 0); + ServerJob document, question; + for (auto [job, fd, prompt, id] : {std::tuple{&document, fds[0], 128, "document"}, + std::tuple{&question, fds[2], 4, "question"}}) { + job->fd = fd; + job->req.format = ApiFormat::OPENAI_CHAT; + job->req.prompt_tokens.assign((size_t)prompt, 100); + job->req.max_output = 1; + job->req.stream = false; + job->req.model = "scheduler-test"; + job->req.response_id = id; + } + SchedulerTestHarness::enqueue(server, &document); + std::thread scheduler([&] { SchedulerTestHarness::run(server, backend.engine); }); + std::this_thread::sleep_for(std::chrono::milliseconds(60)); + SchedulerTestHarness::enqueue(server, &question); + const auto wait_done = [](ServerJob & job) { + std::unique_lock lock(job.mu); + return job.cv.wait_for(lock, std::chrono::seconds(10), [&] { return job.done; }); + }; + const bool done = wait_done(question) && wait_done(document); + SchedulerTestHarness::stop(server); + scheduler.join(); + for (const int fd : fds) close(fd); + unlink(path.c_str()); + TEST_ASSERT(done); + + // The question's step carried only the question; the document went on after. + int question_step = -1, last_document_step = -1; + for (int i = 0; i < (int)backend.engine.prompt_slots.size(); ++i) { + const auto & chunk_slots = backend.engine.prompt_slots[(size_t)i]; + if (chunk_slots == std::vector{1}) question_step = i; + if (std::find(chunk_slots.begin(), chunk_slots.end(), 0) != chunk_slots.end()) last_document_step = i; + } + TEST_ASSERT(question_step > 0); + TEST_ASSERT(last_document_step > question_step); +} + TEST_CASE(ServerUnitFixture, test_scheduler_discards_unrequested_restored_checkpoint) { const std::string path = write_deepseek_marker_tokenizer_fixture(); @@ -10130,6 +10501,18 @@ TEST_CASE(ServerUnitFixture, test_usage_timings_reports_prefix_cache_work) { TEST_ASSERT(j["prefilled_tokens"].get() == 64); TEST_ASSERT(j["effective_prompt_tokens"].get() == 8256); TEST_ASSERT(j["agent_turn_cache_hit"].get()); + + // Apps read the cached share where OpenAI puts it. + for (const auto format : {ApiFormat::OPENAI_CHAT, ApiFormat::RESPONSES}) { + auto em = make_emitter(format); + em.emit_start(); + em.emit_token("x"); + const std::string wire = concat(em.emit_finish(1, &t)); + const char * field = format == ApiFormat::OPENAI_CHAT + ? "\"prompt_tokens_details\":{\"cached_tokens\":8192}" + : "\"input_tokens_details\":{\"cached_tokens\":8192}"; + TEST_ASSERT(wire.find(field) != std::string::npos); + } } TEST_CASE(ServerUnitFixture, test_usage_timings_omitted_when_null) {