From 8222fb1d063d49351fb7737d8334c23fb750706d Mon Sep 17 00:00:00 2001 From: mrciffa Date: Sat, 10 Oct 2026 17:28:35 +0200 Subject: [PATCH 1/9] fix(server): the prefix cache keeps each turn's end, for every request at once Three things kept conversations from reusing their own prefix. - qwen4exp spaces its restore points 4096 tokens apart (#828), and a capture moved down to the deepest spaced point. A chat shorter than that never cached past its system prompt, an agent re-read up to 4K tokens on every turn, and a follow-up question on a 130K-token document re-read all of it (103 s). A backend with spaced restore points now saves the end of the prompt's last message, the point the next turn extends; the extra cut costs one forward over the generation prompt. Other cuts still move down to a spaced point. - The prefix cache held one capture reservation at a time, so the requests prefilling beside a capturing one saved nothing. Each prefilling request can now hold one: a capture in flight keeps its slot and its bytes out of the other selections, and the same prefix is not captured twice. - Concurrent serving never pruned the snapshots a new one supersedes. It now prunes them as each capture lands, as single-sequence serving does, and frees their payloads. Co-Authored-By: Claude Opus 5.5 --- server/src/server/http_server.cpp | 6 +- server/src/server/prefix_cache.cpp | 77 ++++++++++---- server/src/server/prefix_cache.h | 31 ++++-- server/src/server/scheduler.cpp | 12 ++- server/test/unit/test_server_unit.cpp | 139 ++++++++++++++++++++++++-- 5 files changed, 228 insertions(+), 37 deletions(-) diff --git a/server/src/server/http_server.cpp b/server/src/server/http_server.cpp index c2793a269..bcfc7272c 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 diff --git a/server/src/server/prefix_cache.cpp b/server/src/server/prefix_cache.cpp index 6752480b2..a8c22a476 100644 --- a/server/src/server/prefix_cache.cpp +++ b/server/src/server/prefix_cache.cpp @@ -589,11 +589,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 { @@ -623,9 +639,15 @@ PrefixCache::InlineReservation PrefixCache::reserve_inline_snap( 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 backend that cuts prefill at its restore points saves the prompt up to + // the end of its last message, the point the next turn extends: a person's + // follow-up, or the agent's next step after a tool result. The cut adds + // only a forward over the generation prompt after it, where the turn + // before it would make the next one prefill again up to `spacing` tokens. + if (restore_point_spacing > 0) include_last_message = true; // 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 @@ -669,7 +691,8 @@ PrefixCache::InlineReservation PrefixCache::reserve_inline_snap( candidates, restored_prefix_len, prefer_tools_boundary, include_last_message, 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 +718,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 +767,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 +792,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 +804,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; }; @@ -797,6 +831,7 @@ PrefixCache::InlineReservation PrefixCache::reserve_inline_snap( 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_leaf && !is_leaf(i)) continue; if (!allow_protected && entries_[(size_t)i].protect) @@ -832,7 +867,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 +943,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,15 +960,20 @@ 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( @@ -957,7 +998,8 @@ 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); @@ -971,6 +1013,7 @@ std::vector PrefixCache::prune_superseded_ancestors(int slot) { 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 +1083,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..69b0175eb 100644 --- a/server/src/server/prefix_cache.h +++ b/server/src/server/prefix_cache.h @@ -216,7 +216,8 @@ 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 @@ -229,9 +230,10 @@ class PrefixCache { // 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. + // `restore_point_spacing` (ModelBackend::restore_point_spacing), the cut is + // the end of the last message, the point the next turn extends; any other + // cut 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, @@ -255,8 +257,10 @@ class PrefixCache { // 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); + // 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 +348,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 +421,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..77f3fb339 100644 --- a/server/src/server/scheduler.cpp +++ b/server/src/server/scheduler.cpp @@ -997,7 +997,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/test/unit/test_server_unit.cpp b/server/test/unit/test_server_unit.cpp index f74e2af70..098311eb9 100644 --- a/server/test/unit/test_server_unit.cpp +++ b/server/test/unit/test_server_unit.cpp @@ -3612,10 +3612,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)); @@ -4091,26 +4093,58 @@ 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. + // With spacing, a capture is the end of the last message, the point the + // next turn extends, even for a chat (no tool result); 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, 0, false, 0, -1, {}, /*include_last_message=*/false); + TEST_ASSERT(plain.active() && plain.target_cut() == boundaries[boundaries.size() - 2]); 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, {}, false, 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, {}, false, 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, {}, false, 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, true); + auto twin = cache.reserve_inline_snap(a, 0, false, 0, -1, bytes, true); + auto second = cache.reserve_inline_snap(b, 0, false, 0, -1, bytes, true); + 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, true); + 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( @@ -8099,6 +8133,93 @@ 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 one it supersedes (keeping the +// shallowest, the shared head). +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 = 4; + 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, two more turns: the second turn's checkpoint is + // superseded by the third's and freed; the head stays. + 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}); + 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); + + SchedulerTestHarness::stop(server); + scheduler.join(); + for (const int fd : fds) close(fd); + unlink(path.c_str()); + + TEST_ASSERT(ab_done && second_done && third_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 = turn3; + next.push_back(9); + TEST_ASSERT(cache.lookup_candidate(next, (int)next.size() - 1).second == 13); + // The turn-2 checkpoint is gone from the cache and freed in the engine. + TEST_ASSERT(cache.lookup_candidate( + {1, 100, 3, 101, 4, 102, 2, 3, 103, 4, 9}, 10).second == 3); + bool freed_turn2 = false; + for (const auto & ref : backend.engine.discarded) freed_turn2 |= ref.tokens == 8; + 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(); From b31fc4972a0d14d6d72218b9c4f2374e465112fe Mon Sep 17 00:00:00 2001 From: mrciffa Date: Sat, 10 Oct 2026 17:28:35 +0200 Subject: [PATCH 2/9] feat(server): requests decoding beside a long prompt keep half the time While one request prefilled a 130K-token document, a person chatting beside it got one token per 4096-token chunk: a 300-token reply took about 200 s. - StepPlanLimits::prefill_time_share: while requests decode, a prompt longer than one step's chunk takes at most that share of the time; after each of its chunks the decoding requests step alone for the rest. qwen4exp sets 0.5, since a chunk pays the whole expert read and a shorter one would not let them run faster. A prompt that fits in one chunk (an agent's next turn) is not held back. - A prompt that fits in one step runs before longer ones, alone, so a question asked beside a long document waits for one short forward instead of the document's next chunk. Co-Authored-By: Claude Opus 5.5 --- server/src/common/concurrency/seq_engine.h | 4 + server/src/qwen4exp/qwen4exp_seq_engine.cpp | 9 +- server/src/server/scheduler.cpp | 51 +++++- server/test/unit/test_server_unit.cpp | 181 ++++++++++++++++++++ 4 files changed, 240 insertions(+), 5 deletions(-) diff --git a/server/src/common/concurrency/seq_engine.h b/server/src/common/concurrency/seq_engine.h index 33c93744d..21387d2cd 100644 --- a/server/src/common/concurrency/seq_engine.h +++ b/server/src/common/concurrency/seq_engine.h @@ -78,6 +78,10 @@ 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 prompt that needs more + // than one step 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; }; // Select the oldest eligible sequences, then distribute the step's token diff --git a/server/src/qwen4exp/qwen4exp_seq_engine.cpp b/server/src/qwen4exp/qwen4exp_seq_engine.cpp index 5c62b3cff..fd35407dc 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 longer than that 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}; } return {prefill_slots, prefill_granule_, prefill_slots * prefill_granule_, prefill_granule_}; } diff --git a/server/src/server/scheduler.cpp b/server/src/server/scheduler.cpp index 77f3fb339..a01c1e5f2 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,9 @@ 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). A recomputed request counts as long. + int prompt_left = 0; uint64_t admission_order = 0; std::chrono::steady_clock::time_point started_at{}; std::chrono::steady_clock::time_point decode_started_at{}; @@ -663,6 +667,7 @@ 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.cache_capture = std::move(prepared_capture); s.n_gen_cap = std::min( n_gen_cap, @@ -703,6 +708,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 +832,7 @@ 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.pending_tok = -1; std::fprintf(stderr, "[parallel] slot %d resumed via re-prefill\n", @@ -876,9 +886,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 +969,27 @@ 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 prompt longer than one step's chunk shares the time with the + // decoding requests; one that fits (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()); + bool long_prompt = false; + for (const PrefillSlice & slice : step_plan.prefills) { + int & left = slots[(size_t)slice.slot].prompt_left; + long_prompt |= left > step_limits.max_prefill_tokens_per_sequence; + if (left != std::numeric_limits::max()) left = std::max(0, 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()) { diff --git a/server/test/unit/test_server_unit.cpp b/server/test/unit/test_server_unit.cpp index 098311eb9..a6e782107 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) @@ -8356,6 +8359,184 @@ TEST_CASE(ServerUnitFixture, TEST_ASSERT(stats.restore_stall_us_total == 2500); } +// A prompt that takes several steps 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: + 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}; + } + 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_{}; +}; + +struct TimeShareBackend : MockBackend { + TimeShareEngine engine; + SeqEngine * seq_engine() override { return &engine; } +}; + +TEST_CASE(ServerUnitFixture, test_scheduler_gives_decoders_time_beside_a_long_prompt) { + 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); + 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; + }; + // A chat that decodes, then a prompt of 8 chunks beside it. + 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); + + // Between two chunks of the long prompt that ran beside the chat, the chat + // stepped alone: several decode-only steps, not none. + 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); + TEST_ASSERT(alone_between >= 5); +} + +// 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(); From 569f4dab30e0b9140a2b0cf87c72ded44c11526a Mon Sep 17 00:00:00 2001 From: mrciffa Date: Sat, 10 Oct 2026 17:28:35 +0200 Subject: [PATCH 3/9] fix(server): Claude Code requests no longer fail, replies start without blank lines - Claude Code 2.1 sends its environment as a system message after the user's first one. Qwen3.8-Flash-Next's template raises on a system message that is not first, so every Claude Code request failed with HTTP 500 and was retried 10 times. A template now gets a late system or developer message as a user turn, in place, so the prompt before it stays the same and cached. - A template that rejects the messages answers 400, which clients do not retry like a server failure. - Claude Code asks for thinking type "adaptive", which turned thinking off. It now means the model's default, as when the field is absent. - A reply no longer starts with the newlines a template puts after ; apps showed them as blank lines on top. Spaces stay, since they can indent the first line. Co-Authored-By: Claude Opus 5.5 --- server/src/server/chat_template.cpp | 8 +++++++- server/src/server/http_server.cpp | 8 ++++++-- server/src/server/sse_emitter.cpp | 11 +++++++++++ server/test/unit/test_server_unit.cpp | 24 +++++++++++++++++++++++- 4 files changed, 47 insertions(+), 4 deletions(-) diff --git a/server/src/server/chat_template.cpp b/server/src/server/chat_template.cpp index d3c021ed6..4363f7e7c 100644 --- a/server/src/server/chat_template.cpp +++ b/server/src/server/chat_template.cpp @@ -1077,7 +1077,13 @@ std::string render_chat_template_jinja( "normalize with tool-memory replay disabled"); } nlohmann::ordered_json mj; - mj["role"] = m.role; + // Templates take a system message only first (Qwen's raise otherwise). + // Claude Code sends its environment as a system message after the + // user's first one; it reads as a user turn there, in place, so the + // prompt before it stays the same and cached. + const bool late_system = !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; diff --git a/server/src/server/http_server.cpp b/server/src/server/http_server.cpp index bcfc7272c..9fcc7f11e 100644 --- a/server/src/server/http_server.cpp +++ b/server/src/server/http_server.cpp @@ -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); diff --git a/server/src/server/sse_emitter.cpp b/server/src/server/sse_emitter.cpp index 40d5129f2..788ab9740 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); diff --git a/server/test/unit/test_server_unit.cpp b/server/test/unit/test_server_unit.cpp index a6e782107..b9e51a999 100644 --- a/server/test/unit/test_server_unit.cpp +++ b/server/test/unit/test_server_unit.cpp @@ -310,6 +310,23 @@ 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. +TEST_CASE(ServerUnitFixture, test_template_gets_late_system_message_as_user_turn) { + const std::string tmpl = + "{%- 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::vector msgs = { + {"system", "agent"}, {"user", "hi"}, {"system", "env"}, {"developer", "brief"}}; + TEST_ASSERT(render_chat_template_jinja(tmpl, msgs, "", "", false, false, "") == + "[system:agent][user:hi][user:env][user:brief]"); +} + TEST_CASE(ServerUnitFixture, test_pflash_scorer_accepts_responses_string_input) { ToolMemory tool_memory; const auto messages = normalize_chat_messages( @@ -2752,7 +2769,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); } @@ -5141,6 +5159,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); From 980e3f4326ce46d63e3734c519096d0600485005 Mon Sep 17 00:00:00 2001 From: mrciffa Date: Sat, 10 Oct 2026 17:38:28 +0200 Subject: [PATCH 4/9] feat(server): usage reports the prompt tokens the prefix cache served Codex showed 0 cached input tokens for every turn. Chat Completions usage now carries prompt_tokens_details.cached_tokens and Responses usage input_tokens_details.cached_tokens, where OpenAI's clients read them, from the same count as usage.timings.cached_prefix_tokens. Anthropic usage is unchanged: its input_tokens and cache fields add up to the context, which Claude Code would count twice. Co-Authored-By: Claude Opus 5.5 --- server/src/server/http_server.cpp | 6 ++++++ server/src/server/sse_emitter.cpp | 3 +++ server/test/unit/test_server_unit.cpp | 12 ++++++++++++ 3 files changed, 21 insertions(+) diff --git a/server/src/server/http_server.cpp b/server/src/server/http_server.cpp index 9fcc7f11e..f20c138e9 100644 --- a/server/src/server/http_server.cpp +++ b/server/src/server/http_server.cpp @@ -2977,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}, @@ -3077,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}, diff --git a/server/src/server/sse_emitter.cpp b/server/src/server/sse_emitter.cpp index 788ab9740..895fc53c6 100644 --- a/server/src/server/sse_emitter.cpp +++ b/server/src/server/sse_emitter.cpp @@ -975,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 = { @@ -1077,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 b9e51a999..666f1b9c6 100644 --- a/server/test/unit/test_server_unit.cpp +++ b/server/test/unit/test_server_unit.cpp @@ -10454,6 +10454,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) { From de01afd7ed2772c0410298204fedc73c4011a784 Mon Sep 17 00:00:00 2001 From: mrciffa Date: Sat, 10 Oct 2026 17:53:04 +0200 Subject: [PATCH 5/9] fix(server): a late system message stays one where the template takes it The Claude Code fix rendered every system or developer message after the first as a user turn, for every Jinja template. A template that accepts system messages anywhere now gets them as they are; only when the template rejects the messages and a late system message is among them is it rendered again with those as user turns. Co-Authored-By: Claude Opus 5.5 --- server/src/server/chat_template.cpp | 40 ++++++++++++++++++++++----- server/test/unit/test_server_unit.cpp | 10 +++++-- 2 files changed, 40 insertions(+), 10 deletions(-) diff --git a/server/src/server/chat_template.cpp b/server/src/server/chat_template.cpp index 4363f7e7c..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,11 +1078,7 @@ std::string render_chat_template_jinja( "normalize with tool-memory replay disabled"); } nlohmann::ordered_json mj; - // Templates take a system message only first (Qwen's raise otherwise). - // Claude Code sends its environment as a system message after the - // user's first one; it reads as a user turn there, in place, so the - // prompt before it stays the same and cached. - const bool late_system = !messages_j.empty() && + 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; @@ -1170,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/test/unit/test_server_unit.cpp b/server/test/unit/test_server_unit.cpp index 666f1b9c6..1d62254fc 100644 --- a/server/test/unit/test_server_unit.cpp +++ b/server/test/unit/test_server_unit.cpp @@ -312,19 +312,23 @@ TEST_CASE(ServerUnitFixture, test_pflash_scorer_uses_user_query_before_chat_suff // 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. +// 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 tmpl = + 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(tmpl, msgs, "", "", false, false, "") == + 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) { From d70e30dd448c63e8e5813cc31406f93797bcd87f Mon Sep 17 00:00:00 2001 From: mrciffa Date: Sun, 11 Oct 2026 12:38:06 +0200 Subject: [PATCH 6/9] fix(server): every model keeps each turn's end, and the turn before it for edits Only Qwen3.8-Flash-Next saved a turn at the end of its last message. Every other model saved before it, so a follow-up question re-read the message it followed: Qwen 3.8 27B read a 118K-token document again in 532 s, and now restores it and answers in 4 s. - Every backend saves the prompt up to the end of its last message, the point the next turn extends. A conversation's first turn still saves its system head, so new conversations on the same system prompt restore it, unless the head is short beside a long message (as tool requests do). - A turn's new snapshot keeps the previous one while the resident budget holds it: a person who edits their last message resumes there. Eviction drops these kept turns before any conversation's newest one. Co-Authored-By: Claude Opus 5.5 --- server/src/common/model_backend.h | 6 +- server/src/server/http_server.cpp | 8 +- server/src/server/http_server.h | 9 +- server/src/server/prefix_cache.cpp | 136 +++++++++++------- server/src/server/prefix_cache.h | 53 ++++---- server/src/server/scheduler.cpp | 5 +- server/test/unit/test_server_unit.cpp | 189 +++++++++++++++----------- 7 files changed, 236 insertions(+), 170 deletions(-) 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/server/http_server.cpp b/server/src/server/http_server.cpp index f20c138e9..e71d34f10 100644 --- a/server/src/server/http_server.cpp +++ b/server/src/server/http_server.cpp @@ -3677,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()); } @@ -4043,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(); }; @@ -4124,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 a8c22a476..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) { @@ -636,33 +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_) return {}; const auto candidates = find_all_boundaries(prompt_ids, markers_); - // A backend that cuts prefill at its restore points saves the prompt up to - // the end of its last message, the point the next turn extends: a person's - // follow-up, or the agent's next step after a tool result. The cut adds - // only a forward over the generation prompt after it, where the turn - // before it would make the next one prefill again up to `spacing` tokens. - if (restore_point_spacing > 0) include_last_message = true; - // 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 = [&]() { @@ -670,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; @@ -688,8 +693,7 @@ 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 && target_cut != candidates.back()) { @@ -823,16 +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; @@ -978,11 +998,13 @@ int PrefixCache::pick_evict_victim(int skip_index) const { 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) { @@ -1006,7 +1028,23 @@ std::vector PrefixCache::prune_superseded_ancestors( 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) { diff --git a/server/src/server/prefix_cache.h b/server/src/server/prefix_cache.h index 69b0175eb..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 @@ -220,20 +221,19 @@ class PrefixCache { // 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 is - // the end of the last message, the point the next turn extends; any other - // 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, @@ -241,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); @@ -256,7 +255,9 @@ 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. + // 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( diff --git a/server/src/server/scheduler.cpp b/server/src/server/scheduler.cpp index a01c1e5f2..823a26684 100644 --- a/server/src/server/scheduler.cpp +++ b/server/src/server/scheduler.cpp @@ -476,8 +476,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) @@ -497,7 +496,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()) { diff --git a/server/test/unit/test_server_unit.cpp b/server/test/unit/test_server_unit.cpp index 1d62254fc..0dc17fd65 100644 --- a/server/test/unit/test_server_unit.cpp +++ b/server/test/unit/test_server_unit.cpp @@ -3420,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); @@ -3549,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(); } @@ -4046,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) { @@ -4118,9 +4112,9 @@ 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})); - // With spacing, a capture is the end of the last message, the point the - // next turn extends, even for a chat (no tool result); a cut short of it - // 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())); @@ -4129,18 +4123,18 @@ TEST_CASE(ServerUnitFixture, test_spaced_restore_points) { const std::vector boundaries = find_all_boundaries(prompt, cache.chat_markers()); TEST_ASSERT(boundaries.size() >= 3); const int head = boundaries.front(); - auto plain = cache.reserve_inline_snap(prompt, 0, false, 0, -1, {}, /*include_last_message=*/false); - TEST_ASSERT(plain.active() && plain.target_cut() == boundaries[boundaries.size() - 2]); + auto plain = cache.reserve_inline_snap(prompt); + TEST_ASSERT(plain.active() && plain.target_cut() == boundaries.back()); plain.cancel(); // 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, {}, false, 0, /*restore_point_spacing=*/64); + 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, {}, false, 0, 64); + 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, boundaries.back(), false, 0, -1, {}, false, 0, 64); + auto none = cache.reserve_inline_snap(prompt, boundaries.back(), false, 0, -1, {}, 0, 64); TEST_ASSERT(!none.active()); } @@ -4155,13 +4149,13 @@ TEST_CASE(ServerUnitFixture, test_inline_captures_in_flight_hold_their_slots) { 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, true); - auto twin = cache.reserve_inline_snap(a, 0, false, 0, -1, bytes, true); - auto second = cache.reserve_inline_snap(b, 0, false, 0, -1, bytes, true); + 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, true); + 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)); @@ -4173,20 +4167,18 @@ TEST_CASE(ServerUnitFixture, test_inline_captures_in_flight_hold_their_slots) { 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)); @@ -4201,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(); @@ -4222,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); } @@ -4240,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. @@ -4431,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) { @@ -4540,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 @@ -4718,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()); } @@ -7576,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); @@ -8163,8 +8183,8 @@ struct SchedulerPrefixBackend : MockBackend { }; // Concurrent serving: two requests admitted together both capture, and a -// conversation's new checkpoint frees the one it supersedes (keeping the -// shallowest, the shared head). +// 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(); @@ -8177,7 +8197,7 @@ TEST_CASE(ServerUnitFixture, ServerConfig config; config.arch = "qwen35"; config.max_ctx = 64; - config.prefix_cache_cap = 4; + config.prefix_cache_cap = 8; config.concurrent_prefix_cache_max_bytes = 4096; config.concurrent_prefix_cache = true; config.admission_coalesce_ms = 0; @@ -8217,35 +8237,46 @@ TEST_CASE(ServerUnitFixture, }); const bool ab_done = wait_done(a) && wait_done(b); - // One conversation, two more turns: the second turn's checkpoint is - // superseded by the third's and freed; the head stays. + // 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); + 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 = turn3; + std::vector next = turn4; next.push_back(9); - TEST_ASSERT(cache.lookup_candidate(next, (int)next.size() - 1).second == 13); + 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. - TEST_ASSERT(cache.lookup_candidate( - {1, 100, 3, 101, 4, 102, 2, 3, 103, 4, 9}, 10).second == 3); + 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 == 8; + for (const auto & ref : backend.engine.discarded) freed_turn2 |= ref.tokens == 10; TEST_ASSERT(freed_turn2); } From a3836a837118cf6ab5b50436e3662ad682f72ac4 Mon Sep 17 00:00:00 2001 From: mrciffa Date: Sun, 11 Oct 2026 13:20:11 +0200 Subject: [PATCH 7/9] fix(ggml-cuda): paged attention keeps its partials within 256 MiB Paged attention split the context into one partition per 1024 tokens whatever the batch, and each partition keeps partials for every output row. A 2048-row prompt chunk 118K tokens deep asked for more memory than the R9700 had left beside Qwen 3.8 27B's weights and KV cache, and the server aborted: a long document sent to Qwen 3.8 27B on its own crashed the engine 79 s in. The partials now take at most 256 MiB. A chunk with that many rows covers the context with fewer, longer partitions, its rows giving it the parallelism: the same document now prefills in 524 s, as in 512-row chunks (530 s). Co-Authored-By: Claude Opus 5.5 --- .../ggml/src/ggml-cuda/paged-attn.cu | 27 +++++++++++++++++-- 1 file changed, 25 insertions(+), 2 deletions(-) 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. From c103960c91797ebeb11729fec3af18098a46bf62 Mon Sep 17 00:00:00 2001 From: mrciffa Date: Sun, 11 Oct 2026 13:20:12 +0200 Subject: [PATCH 8/9] feat(qwen35): requests decoding beside a long prompt keep half the time A step that carries a prompt chunk decodes without the DFlash drafter, one token per request. Beside a 118K-token document, a person chatting with Qwen 3.8 27B got 0.6 tok/s for the 9 minutes the document took. - While requests decode, a prompt of more than 8192 new tokens takes half the time (StepPlanLimits::prefill_time_share, as Qwen3.8-Flash-Next does), and the steps between its chunks decode with the drafter: the person gets 46 to 57 tok/s, and the document answers after 543 s instead of 530 s. - A prompt is long by the new tokens it brought (StepPlanLimits::time_share_prompt_tokens; by default more than one step's chunk), so it shares the time from its first chunk to its last. An agent's next step stays under 8192 and keeps every step: four agents at once finish in 117 s, against 153 s when their steps shared the time. Co-Authored-By: Claude Opus 5.5 --- server/src/common/concurrency/seq_engine.h | 9 ++++-- .../qwen35/concurrency/qwen35_seq_engine.h | 6 ++++ server/src/server/scheduler.cpp | 23 +++++++++---- server/test/unit/test_server_unit.cpp | 32 +++++++++++++------ 4 files changed, 50 insertions(+), 20 deletions(-) diff --git a/server/src/common/concurrency/seq_engine.h b/server/src/common/concurrency/seq_engine.h index 21387d2cd..f58de81d4 100644 --- a/server/src/common/concurrency/seq_engine.h +++ b/server/src/common/concurrency/seq_engine.h @@ -78,10 +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 prompt that needs more - // than one step may take: after each of its chunks, the decoding requests - // step alone for the rest. 1: a prompt chunk rides every step. + // 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/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/server/scheduler.cpp b/server/src/server/scheduler.cpp index 823a26684..0d228e306 100644 --- a/server/src/server/scheduler.cpp +++ b/server/src/server/scheduler.cpp @@ -38,8 +38,11 @@ struct SchedSlot { 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). A recomputed request counts as long. + // 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{}; @@ -667,6 +670,7 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { 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, @@ -832,6 +836,7 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { 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", @@ -970,15 +975,19 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { const auto step_started = std::chrono::steady_clock::now(); SeqEngine::StepResult step_result = engine.step(step_plan); - // A prompt longer than one step's chunk shares the time with the - // decoding requests; one that fits (an agent's next turn, even when a - // prefix capture splits it in two forwards) does not hold them up. + // 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) { - int & left = slots[(size_t)slice.slot].prompt_left; - long_prompt |= left > step_limits.max_prefill_tokens_per_sequence; - if (left != std::numeric_limits::max()) left = std::max(0, left - slice.max_tokens); + 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; diff --git a/server/test/unit/test_server_unit.cpp b/server/test/unit/test_server_unit.cpp index 0dc17fd65..6e765595f 100644 --- a/server/test/unit/test_server_unit.cpp +++ b/server/test/unit/test_server_unit.cpp @@ -8416,16 +8416,19 @@ TEST_CASE(ServerUnitFixture, TEST_ASSERT(stats.restore_stall_us_total == 2500); } -// A prompt that takes several steps 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. +// 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}; + 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; @@ -8469,18 +8472,22 @@ class TimeShareEngine final : public SeqEngine { 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; } }; -TEST_CASE(ServerUnitFixture, test_scheduler_gives_decoders_time_beside_a_long_prompt) { +// 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(); + auto backend_owner = std::make_unique(long_prompt_tokens); TimeShareBackend & backend = *backend_owner; LuceEngine engine(std::move(backend_owner)); ServerConfig config; @@ -8502,7 +8509,6 @@ TEST_CASE(ServerUnitFixture, test_scheduler_gives_decoders_time_beside_a_long_pr job.req.model = "scheduler-test"; job.req.response_id = id; }; - // A chat that decodes, then a prompt of 8 chunks beside it. ServerJob chat, document; make_job(chat, fds[0], 4, 400, "chat"); make_job(document, fds[2], 128, 1, "document"); @@ -8521,8 +8527,6 @@ TEST_CASE(ServerUnitFixture, test_scheduler_gives_decoders_time_beside_a_long_pr unlink(path.c_str()); TEST_ASSERT(done); - // Between two chunks of the long prompt that ran beside the chat, the chat - // stepped alone: several decode-only steps, not none. int chunks_beside = 0, alone_between = 0, alone_run = 0; bool after_chunk = false; for (const auto & [decode, chunks] : backend.engine.steps) { @@ -8536,7 +8540,15 @@ TEST_CASE(ServerUnitFixture, test_scheduler_gives_decoders_time_beside_a_long_pr } } TEST_ASSERT(chunks_beside >= 2); - TEST_ASSERT(alone_between >= 5); + 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. From aa2b7d259a20f072134ecc6752508bff7e2b227c Mon Sep 17 00:00:00 2001 From: mrciffa Date: Sun, 11 Oct 2026 14:31:45 +0200 Subject: [PATCH 9/9] perf(qwen4exp): an agent's next step no longer gives up time to decoding requests A prompt beside decoding requests shared the time while more than one granule (4096 tokens) of it was left, so an agent's next step after a large tool result waited for the others' replies. It now shares the time only when it brings more than 8192 new tokens, as on Qwen 3.8 27B: a document, not an agent's step. Four agents at once on Qwen3.8-Flash-Next: time to first token 3.29 s mean and 5.95 s at most, against 3.79 s and 11.37 s; the person chatting beside a 130K-token document still decodes at 19 tok/s. Co-Authored-By: Claude Opus 5.5 --- server/src/qwen4exp/qwen4exp_seq_engine.cpp | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/server/src/qwen4exp/qwen4exp_seq_engine.cpp b/server/src/qwen4exp/qwen4exp_seq_engine.cpp index fd35407dc..714cab679 100644 --- a/server/src/qwen4exp/qwen4exp_seq_engine.cpp +++ b/server/src/qwen4exp/qwen4exp_seq_engine.cpp @@ -356,16 +356,16 @@ 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, -// and beside a prompt longer than that 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. +// 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_, - /*prefill_time_share=*/0.5}; + /*prefill_time_share=*/0.5, /*time_share_prompt_tokens=*/8192}; } return {prefill_slots, prefill_granule_, prefill_slots * prefill_granule_, prefill_granule_}; }