From f7cc9ebe4e0f58c08063a7bd662defbe454aa2c9 Mon Sep 17 00:00:00 2001 From: Graffioh Date: Wed, 14 Oct 2026 18:01:44 +0000 Subject: [PATCH 1/2] fix(pflash): drafter resumes after a tool output A tool session's prompt ends after the tool output, so the latest user question's history window (the last token of its reply's header) is never scored and its rows are never stored. The next request scores that window as a history query, finds no stored rows before the resume point, and re-scores the whole prompt from 0. The server now passes that window as a keep span; the drafter computes its rows from the hidden states it already has and stores them with the session. Selections are unchanged. Co-Authored-By: Claude Opus 5.5 --- server/src/common/model_backend.h | 3 + server/src/pflash/pflash_drafter.cpp | 9 ++- server/src/pflash/pflash_drafter.h | 6 +- server/src/pflash/pflash_selection.h | 6 ++ server/src/pflash/qwen35_drafter.cpp | 24 ++++-- server/src/qwen35/qwen35_backend.cpp | 2 +- .../src/qwen35/qwen35_layer_split_adapter.cpp | 3 +- server/src/server/http_server.cpp | 21 +++++ server/test/unit/test_server_unit.cpp | 76 +++++++++++++++++++ 9 files changed, 140 insertions(+), 10 deletions(-) diff --git a/server/src/common/model_backend.h b/server/src/common/model_backend.h index e19bd49ae..8d8f7de33 100644 --- a/server/src/common/model_backend.h +++ b/server/src/common/model_backend.h @@ -490,6 +490,9 @@ struct ModelBackend { // Earlier user questions (their scorer windows), most recent first: // they score the context alongside the query at halving weights. std::vector history_query_spans; + // Windows a later request scores as history queries, never scored + // here: the scorer keeps their rows so that request can resume. + std::vector keep_query_spans; // The latest user turn's tail, a second query window at full weight // (prompt-end chat queries; {-1, -1} otherwise). PFlashTokenSpan turn_query_span{-1, -1}; diff --git a/server/src/pflash/pflash_drafter.cpp b/server/src/pflash/pflash_drafter.cpp index e7ea23740..aa56438bf 100644 --- a/server/src/pflash/pflash_drafter.cpp +++ b/server/src/pflash/pflash_drafter.cpp @@ -114,7 +114,8 @@ std::vector drafter_score_and_compress( const std::vector & history_queries, PFlashTokenSpan turn_query, const std::vector & extra_queries, - const PFlashRequestMessages * request_messages) { + const PFlashRequestMessages * request_messages, + const std::vector & keep_queries) { pflash_clear_kept_spans(); if (!ctx.loaded) { set_last_error("drafter not loaded"); @@ -141,6 +142,12 @@ std::vector drafter_score_and_compress( experiment.history_queries.push_back(window); } } + for (const auto & window : keep_queries) { + if (window.begin >= 0 && window.end > window.begin && + window.end <= (int) ids.size()) { + experiment.keep_queries.push_back(window); + } + } if (turn_query.begin >= 0 && turn_query.end > turn_query.begin && turn_query.end <= (int) ids.size()) { experiment.turn_query = turn_query; diff --git a/server/src/pflash/pflash_drafter.h b/server/src/pflash/pflash_drafter.h index 13fc5876c..cbb8abe3b 100644 --- a/server/src/pflash/pflash_drafter.h +++ b/server/src/pflash/pflash_drafter.h @@ -100,6 +100,9 @@ void free_drafter_weights(DrafterContext & ctx); // messages' content spans the request probe reads // (nullptr: none, the request falls back to today's // windows) +// keep_queries strict selection only: windows a later request scores as +// history queries; never scored here, their rows are +// kept so that request can resume the scoring session // // On failure returns empty vector + sets last_error. std::vector drafter_score_and_compress( @@ -116,6 +119,7 @@ std::vector drafter_score_and_compress( const std::vector & history_queries = {}, PFlashTokenSpan turn_query = {-1, -1}, const std::vector & extra_queries = {}, - const PFlashRequestMessages * request_messages = nullptr); + const PFlashRequestMessages * request_messages = nullptr, + const std::vector & keep_queries = {}); } // namespace luce::common diff --git a/server/src/pflash/pflash_selection.h b/server/src/pflash/pflash_selection.h index c259859fe..785042492 100644 --- a/server/src/pflash/pflash_selection.h +++ b/server/src/pflash/pflash_selection.h @@ -139,6 +139,12 @@ struct PFlashSelectionConfig { // first. The head scores the context against each and mixes the masses // with the query's at weights 1/2, 1/4, ... (multi-turn chats). std::vector history_queries; + // Per request: windows a later request of the same chat will score as + // history queries but this one does not (the latest user question's + // counterpart in a tool session, whose prompt ends after tool output). + // Never scored: the scorer keeps their rows in its session, so the next + // request can resume instead of scoring from token 0. + std::vector keep_queries; // Per request: the tail of the latest user turn, scored as a second // query window at full weight next to the prompt-end query. The last // token reads the whole request; the user's own tokens match literal diff --git a/server/src/pflash/qwen35_drafter.cpp b/server/src/pflash/qwen35_drafter.cpp index 5dd861edb..ba7b2f804 100644 --- a/server/src/pflash/qwen35_drafter.cpp +++ b/server/src/pflash/qwen35_drafter.cpp @@ -1462,6 +1462,16 @@ static std::vector strict_score_impl( } } } + // Rows a later request scores as history queries: kept with the session, + // never scored here. + std::vector keep_rows; + for (const auto & span : experiment.keep_queries) { + if (drop_history || span.end > query_start || span.end - span.begin < 1) continue; + Qwen35ScoringSession::QueryRows keep; + keep.begin = span.begin; + keep.end = span.end; + if (rows_for(span.begin, span.end, keep.rows)) keep_rows.push_back(std::move(keep)); + } ggml_tensor * wk_src = st.head_loaded ? st.head_wk : L.wk; session->keys_trained = st.head_loaded; @@ -2224,16 +2234,18 @@ static std::vector strict_score_impl( for (const auto & [tokens, rows] : qrows) { session->query_sets.put(tokens, rows.data(), hidden, exact_rows_end); } - for (const auto & window : windows) { - if (window.end > exact_rows_end) continue; + const auto store_rows = [&] (int begin, int end, const std::vector & rows) { + if (end > exact_rows_end) return; auto & stored = session->query_windows; stored.erase(std::remove_if(stored.begin(), stored.end(), - [&window] (const Qwen35ScoringSession::QueryRows & old) { - return old.begin == window.begin && old.end == window.end; + [begin, end] (const Qwen35ScoringSession::QueryRows & old) { + return old.begin == begin && old.end == end; }), stored.end()); - stored.push_back({window.begin, window.end, window.rows}); + stored.push_back({begin, end, rows}); if (stored.size() > 8) stored.erase(stored.begin()); - } + }; + for (const auto & keep : keep_rows) store_rows(keep.begin, keep.end, keep.rows); + for (const auto & window : windows) store_rows(window.begin, window.end, window.rows); session->last_used = ++st.session_clock; // Block-15 NoPE Q/K scoring, once per query window: softmax over the diff --git a/server/src/qwen35/qwen35_backend.cpp b/server/src/qwen35/qwen35_backend.cpp index 5fdea8a0c..a78332d3c 100644 --- a/server/src/qwen35/qwen35_backend.cpp +++ b/server/src/qwen35/qwen35_backend.cpp @@ -1513,7 +1513,7 @@ std::vector Qwen35Backend::compress_batch( score_query_end, request.required_instruction_spans, request.query_suffix_candidates, request.history_query_spans, request.turn_query_span, request.extra_query_spans, - &request.request_messages); + &request.request_messages, request.keep_query_spans); result.ok = !result.compressed_ids.empty(); if (result.ok) result.kept_spans = pflash_last_kept_spans(); if (result.ok) { diff --git a/server/src/qwen35/qwen35_layer_split_adapter.cpp b/server/src/qwen35/qwen35_layer_split_adapter.cpp index e709d3b67..77cee4184 100644 --- a/server/src/qwen35/qwen35_layer_split_adapter.cpp +++ b/server/src/qwen35/qwen35_layer_split_adapter.cpp @@ -1396,7 +1396,8 @@ Qwen35LayerSplitAdapter::compress(const ModelBackend::CompressRequest & req) { /*chunk_size=*/32, req.score_query_tokens, /*pool_kernel=*/13, score_query_end, req.required_instruction_spans, req.query_suffix_candidates, req.history_query_spans, - req.turn_query_span, req.extra_query_spans, &req.request_messages); + req.turn_query_span, req.extra_query_spans, &req.request_messages, + req.keep_query_spans); result.ok = !result.compressed_ids.empty(); if (result.ok) result.kept_spans = pflash_last_kept_spans(); if (result.ok) { diff --git a/server/src/server/http_server.cpp b/server/src/server/http_server.cpp index 3ff542b0b..76074d045 100644 --- a/server/src/server/http_server.cpp +++ b/server/src/server/http_server.cpp @@ -5008,6 +5008,9 @@ std::string HttpServer::apply_pflash_compression( // Earlier user questions of a multi-turn chat, most recent first: they // score the context alongside the current query at halving weights. std::vector history_query_spans; + // The window a later request of this chat scores for the latest question + // as a history query, when this request does not score it (below). + std::vector keep_query_spans; // The latest user turn's tail (PFLASH_SELECT_QUERY_TOKENS), a second // query window next to the prompt-end query: literal strings in the // question (an identifier, a described function) match their passage. @@ -5399,6 +5402,22 @@ std::string HttpServer::apply_pflash_compression( {window.end - window.tokens, window.end}); } } + // A tool session's prompt ends after tool output, so the + // latest question's history window (the header of the + // reply that followed it) is no query here. The next + // question scores it; the drafter keeps its rows so that + // request resumes the scoring session instead of + // scoring the whole prompt again. + const size_t reply = (size_t) chat_turn.query_turn + 1; + if (prompt_end_query && history_queries > 0 && + chat_turn.query_turn >= 0 && + reply < chat_turn.turns.size() && + chat_turn.turns[reply].role != "user" && + chat_turn.turns[reply].content_begin > + chat_turn.turns[reply].role_begin) { + const int end = chat_turn.turns[reply].content_begin; + keep_query_spans.push_back({end - 1, end}); + } } // The last non-system turns before the generation prompt // (latest first), from the rendered prompt's chat markers: @@ -5835,6 +5854,7 @@ std::string HttpServer::apply_pflash_compression( http_detail::canonicalize_pflash_token_spans(std::move(pins)); unified_request.query_suffix_candidates = query_suffix_candidates; unified_request.history_query_spans = history_query_spans; + unified_request.keep_query_spans = keep_query_spans; unified_request.turn_query_span = turn_query_span; unified_request.extra_query_spans = extra_query_spans; unified_request.request_messages = request_probe_messages; @@ -6306,6 +6326,7 @@ std::string HttpServer::apply_pflash_compression( std::move(required_instruction_spans); compress_request.query_suffix_candidates = query_suffix_candidates; compress_request.history_query_spans = history_query_spans; + compress_request.keep_query_spans = keep_query_spans; compress_request.turn_query_span = turn_query_span; compress_request.extra_query_spans = extra_query_spans; compress_request.request_messages = request_probe_messages; diff --git a/server/test/unit/test_server_unit.cpp b/server/test/unit/test_server_unit.cpp index 8a5d1e6a1..e6f6caae3 100644 --- a/server/test/unit/test_server_unit.cpp +++ b/server/test/unit/test_server_unit.cpp @@ -11368,9 +11368,85 @@ TEST_CASE(ServerUnitFixture, TEST_ASSERT_MSG(tokenizer.decode({ids.begin() + history.end, ids.begin() + history.end + 2}) == "Sure.", text_of(history)); + // A chat whose prompt ends at the latest question's reply has nothing to + // keep: that question's window is the query itself. + TEST_ASSERT(request.keep_query_spans.empty()); unlink(path.c_str()); } +// A tool session's prompt ends after tool output: the latest question's +// history window (its reply's header) is kept for the next request, which +// scores exactly that window once a new question arrives. +TEST_CASE(ServerUnitFixture, test_pflash_tool_session_keeps_next_history_query) { + luce_test::ScopedEnvVar mode{"PFLASH_SELECT_MODE", "budget_only"}; + luce_test::ScopedEnvVar skeleton{"PFLASH_CHAT_SKELETON_TOKENS", nullptr}; + + std::string material; + for (int i = 0; i < 80; ++i) material += "filler words here. "; + std::vector messages{ + {"system", "You are helpful.", ""}, + {"user", "What is the first answer?", ""}, + {"assistant", "Sure.", ""}, + {"tool", "Here is text: " + material, ""}, + }; + const auto run = [&] (const std::vector & chat, + ModelBackend::CompressRequest & out, + std::vector & ids_out) { + const std::string rendered = render_chat_template( + chat, ChatFormat::QWEN3, /*add_generation_prompt=*/true, + /*enable_thinking=*/true); + const std::string path = write_pflash_bpe_tokenizer_fixture( + {"What", " is", " the", " answer", " first", "?", "user", "assistant", + "system", "tool", "\n", "Sure", ".", " filler", " words", " here", + "Here", " text", ":", " Next", " question"}, + rendered); + Tokenizer tokenizer; + TEST_ASSERT(tokenizer.load_from_gguf(path.c_str())); + auto backend_owner = std::make_unique(); + MockPflashCompressBackend & backend = *backend_owner; + LuceEngine engine(std::move(backend_owner)); + ServerConfig config; + config.pflash_keep_ratio = 1.0f; + config.prefix_cache_cap = 0; + config.prefill_cache_cap = 0; + { + HttpServer server(engine, tokenizer, config); + server.set_drafter_tokenizer(&tokenizer); + ParsedRequest request; + request.format = ApiFormat::OPENAI_CHAT; + json wire = json::array(); + for (const auto & message : chat) { + wire.push_back({{"role", message.role}, {"content", message.content}}); + } + request.messages = wire; + request.prompt_tokens = tokenizer.encode(rendered); + const std::string error = + HttpServerTestAccess::apply_pflash_compression(server, request); + TEST_ASSERT_MSG(error.empty(), error); + } + TEST_ASSERT(backend.compress_calls == 1); + out = backend.last_request; + ids_out = out.input_ids; + unlink(path.c_str()); + }; + ModelBackend::CompressRequest first; + std::vector first_ids; + run(messages, first, first_ids); + TEST_ASSERT(first.history_query_spans.empty()); + TEST_ASSERT(first.keep_query_spans.size() == 1); + + messages.push_back({"assistant", "Sure.", ""}); + messages.push_back({"user", "Next question?", ""}); + ModelBackend::CompressRequest second; + std::vector second_ids; + run(messages, second, second_ids); + TEST_ASSERT(second.history_query_spans.size() == 1); + const auto kept = first.keep_query_spans[0]; + const auto history = second.history_query_spans[0]; + TEST_ASSERT(kept.begin == history.begin && kept.end == history.end); + TEST_ASSERT(std::equal(first_ids.begin(), first_ids.begin() + kept.end, second_ids.begin())); +} + // Two turns over one document through prepare_prompt; returns the served // prompts, the view modes and how often the compressor ran. struct PflashTwoTurnRun { From d8b0b1ae4d32407ccd7ffa1b8a97fde6d427c90a Mon Sep 17 00:00:00 2001 From: Graffioh Date: Wed, 14 Oct 2026 18:16:38 +0000 Subject: [PATCH 2/2] fix(pflash): drop stored query windows past a prompt divergence The scoring session's query windows now use the exact-row sets the query step already keeps: a window is dropped once the prompt diverges at or before it, so a later request never scores with rows from another prompt. The windows the next request needs are stored last, so the eight-entry cap evicts this call's windows first. Co-Authored-By: Claude Opus 5.5 --- server/src/pflash/pflash_scoring_resume.cpp | 6 +++ server/src/pflash/pflash_scoring_resume.h | 3 ++ server/src/pflash/qwen35_drafter.cpp | 45 +++++++------------ server/src/pflash/qwen35_drafter.h | 8 +--- .../pflash/test_pflash_scoring_resume.cpp | 16 +++++++ 5 files changed, 43 insertions(+), 35 deletions(-) diff --git a/server/src/pflash/pflash_scoring_resume.cpp b/server/src/pflash/pflash_scoring_resume.cpp index 78aadba23..4e005ad21 100644 --- a/server/src/pflash/pflash_scoring_resume.cpp +++ b/server/src/pflash/pflash_scoring_resume.cpp @@ -24,6 +24,12 @@ int pflash_scoring_resume(int checkpoint, int shared, int S, int cap) { return checkpoint; } +std::vector pflash_span_tokens(int begin, int end) { + std::vector tokens; + for (int t = begin; t < end; ++t) tokens.push_back(t); + return tokens; +} + bool PFlashExactRowSets::put(const std::vector & tokens, const float * rows, int hidden, int exact_end) { if (tokens.empty() || hidden <= 0 || !rows || tokens.front() < 0 || tokens.back() >= exact_end) return false; for (size_t i = 1; i < tokens.size(); ++i) { diff --git a/server/src/pflash/pflash_scoring_resume.h b/server/src/pflash/pflash_scoring_resume.h index bbf1ef2f0..b8466494f 100644 --- a/server/src/pflash/pflash_scoring_resume.h +++ b/server/src/pflash/pflash_scoring_resume.h @@ -45,6 +45,9 @@ int pflash_session_capacity(int S, int ctx_limit); // the two that are > 0, or 0 when neither is. int pflash_session_ctx_limit(int drafter_ctx, int engine_max_ctx); +// The token positions of the window [begin, end), ascending; empty when end <= begin. +std::vector pflash_span_tokens(int begin, int end); + // Exact F32 block-14 rows of query token sets (the query step's candidate and END sets), kept per session for later // turns whose activations are gone. Only sets whose last token is below the storing call's exact-rows end are put. class PFlashExactRowSets { diff --git a/server/src/pflash/qwen35_drafter.cpp b/server/src/pflash/qwen35_drafter.cpp index ba7b2f804..32038ba21 100644 --- a/server/src/pflash/qwen35_drafter.cpp +++ b/server/src/pflash/qwen35_drafter.cpp @@ -684,13 +684,9 @@ Qwen35ScoringSession * acquire_scoring_session( // Stored windows hold exact rows only (see the store below): one is // reusable while it sits in the shared prefix. const auto stored_window = [&](const Qwen35ScoringSession & session, int begin, int end, int shared) { - const size_t floats = (size_t) w.n_embd * (size_t) (end - begin); - for (const auto & window : session.query_windows) { - if (window.begin == begin && window.end == end && end <= shared && window.rows.size() == floats) { - return true; - } - } - return false; + std::vector rows; + return end <= shared && + session.query_windows.lookup(luce::pflash::pflash_span_tokens(begin, end), w.n_embd, rows); }; Qwen35ScoringSession * best = nullptr; int best_shared = 0; @@ -1205,6 +1201,7 @@ static std::vector strict_score_impl( const int request_rows_invalidated = request_probe ? session->request_rows.invalidate_from(shared_prefix) : 0; session->query_sets.invalidate_from(shared_prefix); + session->query_windows.invalidate_from(shared_prefix); // A session is released (freed or kept) on every exit below. struct SessionExit { std::unique_ptr & scratch; @@ -1411,14 +1408,8 @@ static std::vector strict_score_impl( rows.size() * sizeof(float)); return true; } - for (const auto & stored : session->query_windows) { - if (stored.begin == begin && stored.end == end && - end <= shared_prefix && stored.rows.size() == rows.size()) { - rows = stored.rows; - return true; - } - } - return false; + return end <= shared_prefix && + session->query_windows.lookup(luce::pflash::pflash_span_tokens(begin, end), hidden, rows); }; { ScoredWindow query; @@ -1464,10 +1455,10 @@ static std::vector strict_score_impl( } // Rows a later request scores as history queries: kept with the session, // never scored here. - std::vector keep_rows; + std::vector keep_rows; for (const auto & span : experiment.keep_queries) { if (drop_history || span.end > query_start || span.end - span.begin < 1) continue; - Qwen35ScoringSession::QueryRows keep; + ScoredWindow keep; keep.begin = span.begin; keep.end = span.end; if (rows_for(span.begin, span.end, keep.rows)) keep_rows.push_back(std::move(keep)); @@ -2234,18 +2225,14 @@ static std::vector strict_score_impl( for (const auto & [tokens, rows] : qrows) { session->query_sets.put(tokens, rows.data(), hidden, exact_rows_end); } - const auto store_rows = [&] (int begin, int end, const std::vector & rows) { - if (end > exact_rows_end) return; - auto & stored = session->query_windows; - stored.erase(std::remove_if(stored.begin(), stored.end(), - [begin, end] (const Qwen35ScoringSession::QueryRows & old) { - return old.begin == begin && old.end == end; - }), stored.end()); - stored.push_back({begin, end, rows}); - if (stored.size() > 8) stored.erase(stored.begin()); - }; - for (const auto & keep : keep_rows) store_rows(keep.begin, keep.end, keep.rows); - for (const auto & window : windows) store_rows(window.begin, window.end, window.rows); + // The kept windows go last: the next request needs them, so the cap + // evicts this call's windows first. + for (const auto * list : {&windows, &keep_rows}) { + for (const auto & window : *list) { + session->query_windows.put(luce::pflash::pflash_span_tokens(window.begin, window.end), + window.rows.data(), hidden, exact_rows_end); + } + } session->last_used = ++st.session_clock; // Block-15 NoPE Q/K scoring, once per query window: softmax over the diff --git a/server/src/pflash/qwen35_drafter.h b/server/src/pflash/qwen35_drafter.h index 2293a0bcb..d269efef9 100644 --- a/server/src/pflash/qwen35_drafter.h +++ b/server/src/pflash/qwen35_drafter.h @@ -63,12 +63,8 @@ struct Qwen35ScoringSession { // questions), reused while they sit in the shared prefix: an agent step // appends tool output after the same user turn, and a new turn's history // queries are earlier turns' queries. Most recent last, at most 8. - struct QueryRows { - int begin = -1; - int end = -1; - std::vector rows; // [hidden, end - begin] - }; - std::vector query_windows; + // Keyed by the window's token positions. + luce::pflash::PFlashExactRowSets query_windows; // The query step's (PFLASH_QUERY_CANDIDATES / PFLASH_QUERY_PIN) query token sets, F32 rows exact for later // prompts that share them (never the BF16 request_rows). luce::pflash::PFlashExactRowSets query_sets; diff --git a/server/test/pflash/test_pflash_scoring_resume.cpp b/server/test/pflash/test_pflash_scoring_resume.cpp index 72d331600..db3e48b11 100644 --- a/server/test/pflash/test_pflash_scoring_resume.cpp +++ b/server/test/pflash/test_pflash_scoring_resume.cpp @@ -123,6 +123,22 @@ TEST_CASE(PFlashScoringResumeFixture, query_sets_store_exact_rows_only) { CHECK(sets.size() == 0); } +// Query windows are row sets keyed by their positions: a window stored before +// a divergence point stays, one that reaches it is dropped. +TEST_CASE(PFlashScoringResumeFixture, query_windows_are_position_sets) { + CHECK(pflash_span_tokens(5, 8) == std::vector({5, 6, 7})); + CHECK(pflash_span_tokens(8, 8).empty()); + PFlashExactRowSets windows; + const int hidden = 1; + const float row[] = {1}; + std::vector out; + CHECK(windows.put(pflash_span_tokens(447, 448), row, hidden, 65536)); + CHECK(windows.put(pflash_span_tokens(12000, 12001), row, hidden, 65536)); + CHECK(windows.invalidate_from(11000) == 1); + CHECK(windows.lookup(pflash_span_tokens(447, 448), hidden, out)); + CHECK(!windows.lookup(pflash_span_tokens(12000, 12001), hidden, out)); +} + // Session capacity: headroom S + S/2 + 4096, clamped to the context limit, never below S padded to 256. TEST_CASE(PFlashScoringResumeFixture, session_capacity_is_headroom_clamped_to_the_context_limit) { using luce::pflash::pflash_session_capacity;