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_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/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..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; @@ -1462,6 +1453,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; + 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)); + } ggml_tensor * wk_src = st.head_loaded ? st.head_wk : L.wk; session->keys_trained = st.head_loaded; @@ -2224,15 +2225,13 @@ 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; - 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; - }), stored.end()); - stored.push_back({window.begin, window.end, window.rows}); - if (stored.size() > 8) stored.erase(stored.begin()); + // 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; 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/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/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; 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 {