From c56de9cc6eec2741e000ba69a16a7850b2da1b4c Mon Sep 17 00:00:00 2001 From: Graffioh Date: Sat, 19 Sep 2026 14:08:57 +0000 Subject: [PATCH 01/26] feat(pflash): Qwen3.5-0.8B drafter scorer with segment probe and scoring head Teach the PFlash drafter to score with a Qwen3.5-0.8B hybrid prefix instead of only the Qwen3-0.6B tail-attention scorer. The new drafter lives in qwen35_drafter.cpp behind qwen35_drafter_score_and_compress, its GGUF and optional companion files load through qwen35_loader.cpp, and the pieces both architectures share moved into qwen3_drafter_common.cpp. An optional scoring head replaces the block-15 Q/K projections and an optional segment probe proposes variable-length candidates instead of fixed chunks; both are GGUF files validated against an explicit contract and fail closed when it is not met. pflash_selection.{cpp,h} turns the PFLASH_SELECT_* environment into a strict budget selector that ranks candidates, honours structurally required spans and stops at the token budget. Co-Authored-By: Claude Fable 5.1 --- server/CMakeLists.txt | 4 + server/src/common/model_backend.h | 13 + server/src/common/pflash_types.h | 27 + server/src/flashprefill_q8.cpp | 11 +- server/src/qwen3/pflash_selection.cpp | 549 ++++++++++ server/src/qwen3/pflash_selection.h | 147 +++ server/src/qwen3/qwen35_drafter.cpp | 997 ++++++++++++++++++ server/src/qwen3/qwen35_drafter.h | 108 ++ server/src/qwen3/qwen35_loader.cpp | 376 +++++++ server/src/qwen3/qwen3_backend.cpp | 7 +- server/src/qwen3/qwen3_drafter.cpp | 582 ++-------- server/src/qwen3/qwen3_drafter.h | 6 +- server/src/qwen3/qwen3_drafter_common.cpp | 326 ++++++ server/src/qwen3/qwen3_drafter_common.h | 77 ++ server/src/qwen3/qwen3_drafter_model.h | 22 + server/src/qwen3/qwen3_graph.cpp | 177 +++- server/src/qwen3/qwen3_loader.cpp | 118 +++ server/src/qwen35/gguf_target_loader.cpp | 4 +- server/src/qwen35/qwen35_backend.cpp | 2 +- .../src/qwen35/qwen35_layer_split_adapter.cpp | 2 +- 20 files changed, 3013 insertions(+), 542 deletions(-) create mode 100644 server/src/common/pflash_types.h create mode 100644 server/src/qwen3/pflash_selection.cpp create mode 100644 server/src/qwen3/pflash_selection.h create mode 100644 server/src/qwen3/qwen35_drafter.cpp create mode 100644 server/src/qwen3/qwen35_drafter.h create mode 100644 server/src/qwen3/qwen35_loader.cpp create mode 100644 server/src/qwen3/qwen3_drafter_common.cpp create mode 100644 server/src/qwen3/qwen3_drafter_common.h diff --git a/server/CMakeLists.txt b/server/CMakeLists.txt index 48ba81573..e5b1b01f4 100644 --- a/server/CMakeLists.txt +++ b/server/CMakeLists.txt @@ -474,7 +474,11 @@ add_library(dflash_common STATIC src/draft/draft_safetensors_loader.cpp src/draft/draft_graph.cpp src/qwen3/anchor_scan.cpp + src/qwen3/pflash_selection.cpp src/qwen3/qwen3_drafter.cpp + src/qwen3/qwen3_drafter_common.cpp + src/qwen3/qwen35_drafter.cpp + src/qwen3/qwen35_loader.cpp src/qwen3/qwen3_kvflash_scorer.cpp src/qwen3/qwen3_loader.cpp src/qwen3/qwen3_graph.cpp diff --git a/server/src/common/model_backend.h b/server/src/common/model_backend.h index 2082d2f03..0f98bc9a9 100644 --- a/server/src/common/model_backend.h +++ b/server/src/common/model_backend.h @@ -10,6 +10,8 @@ #pragma once +#include "pflash_types.h" + #include #include #include @@ -266,6 +268,9 @@ struct ModelBackend { // that knob controls lexical anchors, not neural scorer Q rows. int score_query_end = -1; int score_query_tokens = 8; + // Role-derived instruction structure in drafter-token coordinates. + // Empty is a valid instruction-free or legacy request. + std::vector required_instruction_spans; std::string drafter_path; // GGUF path (for lazy-load) int drafter_gpu = 0; // backend-local GPU for PFlash drafter bool skip_park = false; // true on >=32GB GPUs @@ -275,6 +280,14 @@ struct ModelBackend { struct CompressResult { bool ok = false; std::vector compressed_ids; // surviving token IDs + + static CompressResult from_compressed_ids( + std::vector ids) { + CompressResult result; + result.compressed_ids = std::move(ids); + result.ok = !result.compressed_ids.empty(); + return result; + } }; // Typed compress API (preferred for in-process callers). diff --git a/server/src/common/pflash_types.h b/server/src/common/pflash_types.h new file mode 100644 index 000000000..afd7902ec --- /dev/null +++ b/server/src/common/pflash_types.h @@ -0,0 +1,27 @@ +#pragma once + +#include + +namespace dflash::common { + +inline constexpr size_t kPFlashMaxInstructionSpans = 64; + +// Half-open token range in the drafter-tokenized prompt. +struct PFlashTokenSpan { + int begin = 0; + int end = 0; +}; + +inline bool operator==( + const PFlashTokenSpan & left, + const PFlashTokenSpan & right) noexcept { + return left.begin == right.begin && left.end == right.end; +} + +inline bool operator!=( + const PFlashTokenSpan & left, + const PFlashTokenSpan & right) noexcept { + return !(left == right); +} + +} // namespace dflash::common diff --git a/server/src/flashprefill_q8.cpp b/server/src/flashprefill_q8.cpp index e3e2737b2..df8d07dd6 100644 --- a/server/src/flashprefill_q8.cpp +++ b/server/src/flashprefill_q8.cpp @@ -156,7 +156,16 @@ int flash_prefill_forward_q8( (size_t)kv_len * cl * sizeof(uint16_t)); } - ggml_backend_graph_compute(backend, gf); + const ggml_status compute_status = + ggml_backend_graph_compute(backend, gf); + if (compute_status != GGML_STATUS_SUCCESS) { + std::fprintf(stderr, + "[flashprefill_q8] graph compute failed at cs=%d: %s\n", + cs, ggml_status_to_string(compute_status)); + ggml_free(ctx); + ggml_gallocr_free(galloc); + return -1; + } ggml_backend_synchronize(backend); ggml_free(ctx); } diff --git a/server/src/qwen3/pflash_selection.cpp b/server/src/qwen3/pflash_selection.cpp new file mode 100644 index 000000000..48901a814 --- /dev/null +++ b/server/src/qwen3/pflash_selection.cpp @@ -0,0 +1,549 @@ +#include "pflash_selection.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace dflash::qwen3 { + +namespace { + +constexpr const char * kModeEnv = "PFLASH_SELECT_MODE"; +constexpr const char * kChunkEnv = "PFLASH_SELECT_CHUNK_SIZE"; +constexpr const char * kQueryEnv = "PFLASH_SELECT_QUERY_TOKENS"; +constexpr const char * kQueryParserEnv = "PFLASH_SELECT_QUERY_PARSER"; +constexpr const char * kTopPEnv = "PFLASH_SELECT_TOP_P"; +constexpr const char * kSegmentsEnv = "PFLASH_SELECT_SEGMENTS"; +constexpr const char * kSelectEnv = "PFLASH_SELECT_SCORE"; +constexpr const char * kScorerEnv = "PFLASH_SELECT_SCORER"; +constexpr const char * kSplitEnv = "PFLASH_SELECT_SPLIT"; + +PFlashSelectionResult invalid_result(std::string error) { + PFlashSelectionResult result; + result.error = std::move(error); + return result; +} + +bool parse_int(const char * raw, int & out) { + if (!raw || !*raw) return false; + errno = 0; + char * end = nullptr; + const long value = std::strtol(raw, &end, 10); + if (errno == ERANGE || end == raw || *end != '\0' || + value < INT_MIN || value > INT_MAX) { + return false; + } + out = static_cast(value); + return true; +} + +bool parse_double(const char * raw, double & out) { + if (!raw || !*raw) return false; + errno = 0; + char * end = nullptr; + const double value = std::strtod(raw, &end); + if (errno == ERANGE || end == raw || *end != '\0' || + !std::isfinite(value)) { + return false; + } + out = value; + return true; +} + +int scheduled_chunk_size(int input_tokens) { + if (input_tokens < 500) return 128; + if (input_tokens < 3000) return 512; + return 1024; +} + +} // namespace + +bool has_pflash_selection_environment() noexcept { + return std::getenv(kModeEnv) != nullptr || + std::getenv(kChunkEnv) != nullptr || + std::getenv(kQueryEnv) != nullptr || + std::getenv(kQueryParserEnv) != nullptr || + std::getenv(kTopPEnv) != nullptr || + std::getenv(kSegmentsEnv) != nullptr || + std::getenv(kSelectEnv) != nullptr || + std::getenv(kScorerEnv) != nullptr || + std::getenv(kSplitEnv) != nullptr; +} + +bool pflash_chunk_is_structurally_required( + int begin, + int end, + int query_begin, + int query_end, + int input_tokens, + const std::vector & + required_instruction_spans) noexcept { + if (begin < 0 || end <= begin || query_begin < 0 || + query_end < query_begin || input_tokens < query_end || + end > input_tokens) { + return false; + } + const bool query_chunk = begin < query_end && end > query_begin; + const bool structural_suffix_chunk = + begin < input_tokens && end > query_end; + if (query_chunk || structural_suffix_chunk) return true; + for (const auto & span : required_instruction_spans) { + if (begin < span.end && end > span.begin) return true; + } + return false; +} + +bool validate_pflash_instruction_spans( + const std::vector & spans, + int input_tokens, + std::string & error) noexcept { + error.clear(); + if (input_tokens < 0) { + error = "PFlash input token count must not be negative"; + return false; + } + if (spans.size() > dflash::common::kPFlashMaxInstructionSpans) { + error = "PFlash has too many instruction spans"; + return false; + } + int previous_end = 0; + for (const auto & span : spans) { + if (span.begin < 0 || span.end <= span.begin || + span.end > input_tokens) { + error = "PFlash instruction span is outside the input"; + return false; + } + if (span.begin < previous_end) { + error = "PFlash instruction spans must be ordered and non-overlapping"; + return false; + } + previous_end = span.end; + } + return true; +} + +PFlashSelectionResult select_pflash_candidates( + const std::vector & candidates, + const PFlashSelectionPolicy & policy, + PFlashSelectionMode mode) { + if (mode == PFlashSelectionMode::Legacy) { + return invalid_result("legacy mode does not use strict PFlash selection"); + } + if (policy.token_budget <= 0) { + return invalid_result("PFlash token budget must be positive"); + } + if (!std::isfinite(policy.top_p) || policy.top_p <= 0.0 || policy.top_p > 1.0) { + return invalid_result("PFlash top_p must be finite and in (0, 1]"); + } + + std::vector source_ranges; + source_ranges.reserve(candidates.size()); + std::vector ordinals; + ordinals.reserve(candidates.size()); + for (const auto & candidate : candidates) { + if (candidate.begin < 0 || candidate.end <= candidate.begin) { + return invalid_result("PFlash candidate range is invalid"); + } + if (!std::isfinite(candidate.score)) { + return invalid_result("PFlash candidate score must be finite"); + } + source_ranges.push_back(&candidate); + ordinals.push_back(candidate.ordinal); + } + + std::sort(ordinals.begin(), ordinals.end()); + if (std::adjacent_find(ordinals.begin(), ordinals.end()) != ordinals.end()) { + return invalid_result("PFlash candidate ordinals must be unique"); + } + std::sort(source_ranges.begin(), source_ranges.end(), + [](const auto * left, const auto * right) { + if (left->begin != right->begin) return left->begin < right->begin; + return left->end < right->end; + }); + for (size_t index = 1; index < source_ranges.size(); ++index) { + if (source_ranges[index - 1]->end > source_ranges[index]->begin) { + return invalid_result("PFlash candidate ranges must not overlap"); + } + } + + PFlashSelectionResult result; + result.ok = true; + result.stop = PFlashSelectionStop::CandidatesExhausted; + std::vector selected_candidates; + selected_candidates.reserve(candidates.size()); + + std::vector optional; + optional.reserve(candidates.size()); + for (const auto & candidate : candidates) { + if (candidate.mandatory) { + const int length = candidate.end - candidate.begin; + if (length > policy.token_budget - result.retained_tokens) { + result = {}; + result.stop = PFlashSelectionStop::MandatoryQueryExceedsBudget; + result.error = "mandatory PFlash retention tokens exceed the token budget"; + return result; + } + selected_candidates.push_back(&candidate); + result.retained_tokens += length; + } else { + optional.push_back(&candidate); + } + } + + std::sort(optional.begin(), optional.end(), + [](const auto * left, const auto * right) { + const double left_score = std::max(0.0, left->score); + const double right_score = std::max(0.0, right->score); + if (left_score != right_score) return left_score > right_score; + return left->ordinal < right->ordinal; + }); + + double max_score = 0.0; + for (const auto * candidate : optional) { + max_score = std::max(max_score, std::max(0.0, candidate->score)); + } + double scaled_total = 0.0; + if (max_score > 0.0) { + for (const auto * candidate : optional) { + scaled_total += std::max(0.0, candidate->score) / max_score; + } + } + + for (const auto * candidate : optional) { + if (mode == PFlashSelectionMode::CumulativeTopP && + result.retained_mass >= policy.top_p) { + result.stop = PFlashSelectionStop::TopPReached; + break; + } + + const int length = candidate->end - candidate->begin; + if (length > policy.token_budget - result.retained_tokens) { + result.stop = PFlashSelectionStop::BudgetReached; + if (policy.skip_oversized) continue; + break; + } + + selected_candidates.push_back(candidate); + result.retained_tokens += length; + if (!optional.empty()) { + result.retained_mass += max_score > 0.0 + ? (std::max(0.0, candidate->score) / max_score) / scaled_total + : 1.0 / static_cast(optional.size()); + } + } + + std::sort(selected_candidates.begin(), selected_candidates.end(), + [](const auto * left, const auto * right) { + if (left->begin != right->begin) return left->begin < right->begin; + return left->end < right->end; + }); + result.ordinals.reserve(selected_candidates.size()); + for (const auto * candidate : selected_candidates) { + result.ordinals.push_back(candidate->ordinal); + } + return result; +} + +const char * pflash_selection_mode_name(PFlashSelectionMode mode) noexcept { + switch (mode) { + case PFlashSelectionMode::Legacy: return "legacy"; + case PFlashSelectionMode::BudgetOnly: return "budget_only"; + case PFlashSelectionMode::CumulativeTopP: return "top_p"; + } + return "unknown"; +} + +const char * pflash_selection_stop_name(PFlashSelectionStop stop) noexcept { + switch (stop) { + case PFlashSelectionStop::TopPReached: return "top_p_reached"; + case PFlashSelectionStop::BudgetReached: return "budget_reached"; + case PFlashSelectionStop::CandidatesExhausted: return "candidates_exhausted"; + case PFlashSelectionStop::InvalidInput: return "invalid_input"; + case PFlashSelectionStop::MandatoryQueryExceedsBudget: + return "mandatory_query_exceeds_budget"; + } + return "unknown"; +} + +const char * pflash_query_parser_name(PFlashQueryParser parser) noexcept { + switch (parser) { + case PFlashQueryParser::SemanticUser: return "latest_user"; + case PFlashQueryParser::ArbitraryTail: return "arbitrary_tail"; + } + return "unknown"; +} + +bool resolve_pflash_selection( + int input_tokens, + int legacy_chunk_size, + PFlashSelectionConfig & out, + std::string & error) { + error.clear(); + if (input_tokens < 0) { + error = "PFlash input token count must not be negative"; + return false; + } + if (legacy_chunk_size <= 0) { + error = "PFlash legacy chunk size must be positive"; + return false; + } + + const char * mode_raw = std::getenv(kModeEnv); + const char * chunk_raw = std::getenv(kChunkEnv); + const char * query_raw = std::getenv(kQueryEnv); + const char * query_parser_raw = std::getenv(kQueryParserEnv); + const char * top_p_raw = std::getenv(kTopPEnv); + const char * segments_raw = std::getenv(kSegmentsEnv); + const char * select_raw = std::getenv(kSelectEnv); + const char * scorer_raw = std::getenv(kScorerEnv); + const char * split_raw = std::getenv(kSplitEnv); + + PFlashSelectionConfig config; + config.configured = mode_raw || chunk_raw || query_raw || + query_parser_raw || top_p_raw || segments_raw || select_raw || + scorer_raw || split_raw; + if (scorer_raw) { + if (std::strcmp(scorer_raw, "head") == 0) { + config.scorer = PFlashScorer::Head; + } else if (std::strcmp(scorer_raw, "legacy") == 0) { + config.scorer = PFlashScorer::Legacy; + } else if (std::strcmp(scorer_raw, "split") == 0) { + config.scorer = PFlashScorer::Split; + } else { + error = std::string(kScorerEnv) + " must be head, legacy or split"; + return false; + } + } + if (split_raw) { + char * end = nullptr; + errno = 0; + const double value = std::strtod(split_raw, &end); + if (errno != 0 || end == split_raw || *end != '\0' || !(value > 0.0 && value < 1.0)) { + error = std::string(kSplitEnv) + " must be a fraction in (0, 1)"; + return false; + } + config.split_fraction = value; + } + if (segments_raw) { + if (std::strcmp(segments_raw, "fixed") == 0) { + config.segmentation = PFlashSegmentation::Fixed; + } else if (std::strcmp(segments_raw, "probe") == 0) { + config.segmentation = PFlashSegmentation::Probe; + } else if (std::strcmp(segments_raw, "auto") != 0) { + error = std::string(kSegmentsEnv) + " must be auto, fixed or probe"; + return false; + } + } + if (select_raw) { + if (std::strcmp(select_raw, "sum") == 0) { + config.candidate_score = PFlashCandidateScore::Sum; + } else if (std::strcmp(select_raw, "density") == 0) { + config.candidate_score = PFlashCandidateScore::Density; + } else if (std::strcmp(select_raw, "auto") != 0) { + error = std::string(kSelectEnv) + " must be auto, sum or density"; + return false; + } + } + config.chunk_size = legacy_chunk_size; + + if (mode_raw) { + if (std::strcmp(mode_raw, "budget_only") == 0) { + config.mode = PFlashSelectionMode::BudgetOnly; + } else if (std::strcmp(mode_raw, "top_p") == 0) { + config.mode = PFlashSelectionMode::CumulativeTopP; + } else { + error = std::string(kModeEnv) + " must be budget_only or top_p"; + return false; + } + } + config.selection_active = config.mode != PFlashSelectionMode::Legacy; + + if (chunk_raw) { + if (!parse_int(chunk_raw, config.chunk_size) || config.chunk_size <= 0) { + error = std::string(kChunkEnv) + " must be a positive integer"; + return false; + } + } else if (config.selection_active) { + config.chunk_size = scheduled_chunk_size(input_tokens); + } + + if (query_raw && + (!parse_int(query_raw, config.query_tokens) || + config.query_tokens < 1 || config.query_tokens > 512)) { + error = std::string(kQueryEnv) + " must be an integer in [1, 512]"; + return false; + } + + if (query_parser_raw) { + if (std::strcmp(query_parser_raw, "latest_user") == 0) { + config.query_parser = PFlashQueryParser::SemanticUser; + } else if (std::strcmp(query_parser_raw, "arbitrary_tail") == 0) { + config.query_parser = PFlashQueryParser::ArbitraryTail; + } else { + error = std::string(kQueryParserEnv) + + " must be latest_user or arbitrary_tail"; + return false; + } + } + + if (top_p_raw && + (!parse_double(top_p_raw, config.top_p) || + config.top_p <= 0.0 || config.top_p > 1.0)) { + error = std::string(kTopPEnv) + " must be finite and in (0, 1]"; + return false; + } + + out = config; + return true; +} + +std::vector pflash_probe_segments( + const std::vector & boundary_scores, + int input_tokens, + float threshold, + int min_segment, + int max_segment, + const std::vector & forced_cuts, + const std::vector & split_scores) { + using dflash::common::PFlashTokenSpan; + std::vector spans; + if (input_tokens <= 0 || (int) boundary_scores.size() < input_tokens || + min_segment < 1 || max_segment < min_segment) { + return spans; + } + // Sub-unit scores feed only the oversize interior argmax; + // the boundary threshold and merge floor always read the unit scores. + const std::vector & interior = + (int) split_scores.size() >= input_tokens ? split_scores : boundary_scores; + std::vector forced((size_t) input_tokens + 1, 0); + for (int cut : forced_cuts) { + if (cut > 0 && cut < input_tokens) forced[(size_t) cut] = 1; + } + std::vector cuts; + cuts.push_back(0); + for (int token = 1; token < input_tokens; ++token) { + const bool wanted = forced[(size_t) token] || + (std::isfinite(boundary_scores[(size_t) token]) && + boundary_scores[(size_t) token] > threshold); + if (!wanted) continue; + if (!forced[(size_t) token] && token - cuts.back() < min_segment) continue; + cuts.push_back(token); + } + cuts.push_back(input_tokens); + for (size_t index = 1; index < cuts.size(); ++index) { + int begin = cuts[index - 1]; + const int end = cuts[index]; + while (end - begin > max_segment) { + // Split at the best-scoring interior token in the second half of + // the next max_segment piece (the distance guard keeps the split + // off the near edge), honoring the min_segment margins on both + // sides; else on a fixed grid. + const int lo = std::max(begin + min_segment, begin + max_segment / 2); + const int hi = std::min(end - min_segment, begin + max_segment); + int best = -1; + float best_score = 0.0f; + for (int token = lo; token <= hi; ++token) { + const float score = interior[(size_t) token]; + if (std::isfinite(score) && score > best_score) { + best_score = score; + best = token; + } + } + if (best < 0) best = begin + max_segment; + spans.push_back({begin, best}); + begin = best; + } + spans.push_back({begin, end}); + } + return spans; +} + +PFlashSelectionResult select_pflash_split( + const std::vector & head, + const std::vector & other, + const PFlashSelectionPolicy & policy, + double head_fraction, + PFlashSelectionMode mode) { + PFlashSelectionResult result; + if (head.size() != other.size() || !(head_fraction > 0.0 && head_fraction < 1.0)) { + result.stop = PFlashSelectionStop::InvalidInput; + result.error = "split selection needs matching candidate lists and a fraction in (0, 1)"; + return result; + } + for (size_t i = 0; i < head.size(); ++i) { + if (head[i].ordinal != other[i].ordinal || head[i].begin != other[i].begin || + head[i].end != other[i].end || head[i].mandatory != other[i].mandatory) { + result.stop = PFlashSelectionStop::InvalidInput; + result.error = "split selection candidate lists describe different spans"; + return result; + } + } + PFlashSelectionPolicy first = policy; + first.token_budget = static_cast(policy.token_budget * head_fraction); + // Mandatory spans must fit even when the head's share is small. + int mandatory = 0; + for (const auto & c : head) if (c.mandatory) mandatory += c.end - c.begin; + first.token_budget = (std::max)(first.token_budget, (std::min)(mandatory, policy.token_budget)); + const PFlashSelectionResult pass1 = select_pflash_candidates(head, first, mode); + if (!pass1.ok) return pass1; + std::vector taken(head.size(), 0); + for (size_t ordinal : pass1.ordinals) { + for (size_t i = 0; i < head.size(); ++i) if (head[i].ordinal == ordinal) taken[i] = 1; + } + std::vector rest; + for (size_t i = 0; i < other.size(); ++i) { + if (taken[i]) continue; + PFlashSelectionCandidate c = other[i]; + c.mandatory = false; // mandatory spans were charged in pass 1 + rest.push_back(c); + } + PFlashSelectionPolicy second = policy; + second.token_budget = policy.token_budget - pass1.retained_tokens; + const PFlashSelectionResult pass2 = second.token_budget > 0 + ? select_pflash_candidates(rest, second, mode) + : PFlashSelectionResult{}; + result.ok = true; + result.ordinals = pass1.ordinals; + result.ordinals.insert(result.ordinals.end(), pass2.ordinals.begin(), pass2.ordinals.end()); + std::sort(result.ordinals.begin(), result.ordinals.end()); + result.retained_tokens = pass1.retained_tokens + pass2.retained_tokens; + result.retained_mass = pass1.retained_mass; // the head's normalised mass share + result.stop = second.token_budget > 0 ? pass2.stop : pass1.stop; + return result; +} + +const char * pflash_scorer_name(PFlashScorer scorer) noexcept { + switch (scorer) { + case PFlashScorer::Head: return "head"; + case PFlashScorer::Legacy: return "legacy"; + case PFlashScorer::Split: return "split"; + } + return "unknown"; +} + +const char * pflash_segmentation_name(PFlashSegmentation segmentation) noexcept { + switch (segmentation) { + case PFlashSegmentation::Auto: return "auto"; + case PFlashSegmentation::Fixed: return "fixed"; + case PFlashSegmentation::Probe: return "probe"; + } + return "unknown"; +} + +const char * pflash_candidate_score_name(PFlashCandidateScore score) noexcept { + switch (score) { + case PFlashCandidateScore::Auto: return "auto"; + case PFlashCandidateScore::Sum: return "sum"; + case PFlashCandidateScore::Density: return "density"; + } + return "unknown"; +} + +} // namespace dflash::qwen3 diff --git a/server/src/qwen3/pflash_selection.h b/server/src/qwen3/pflash_selection.h new file mode 100644 index 000000000..9d4bfc4c0 --- /dev/null +++ b/server/src/qwen3/pflash_selection.h @@ -0,0 +1,147 @@ +#pragma once + +#include "common/pflash_types.h" + +#include +#include +#include + +namespace dflash::qwen3 { + +enum class PFlashSelectionMode { + Legacy, + BudgetOnly, + CumulativeTopP, +}; + +enum class PFlashQueryParser { + SemanticUser, + ArbitraryTail, +}; + +enum class PFlashSelectionStop { + TopPReached, + BudgetReached, + CandidatesExhausted, + InvalidInput, + MandatoryQueryExceedsBudget, +}; + +struct PFlashSelectionCandidate { + size_t ordinal = 0; + int begin = 0; + int end = 0; + double score = 0.0; + bool mandatory = false; +}; + +struct PFlashSelectionPolicy { + int token_budget = 0; + double top_p = 0.95; + // Variable-length segments: a candidate that does not fit the remaining + // budget is skipped instead of ending the fill, so smaller segments + // ranked below it can still be kept. + bool skip_oversized = false; +}; + +struct PFlashSelectionResult { + bool ok = false; + std::vector ordinals; + int retained_tokens = 0; + double retained_mass = 0.0; + PFlashSelectionStop stop = PFlashSelectionStop::InvalidInput; + std::string error; +}; + +bool pflash_chunk_is_structurally_required( + int begin, + int end, + int query_begin, + int query_end, + int input_tokens, + const std::vector & + required_instruction_spans = {}) noexcept; + +bool validate_pflash_instruction_spans( + const std::vector & spans, + int input_tokens, + std::string & error) noexcept; + +PFlashSelectionResult select_pflash_candidates( + const std::vector & candidates, + const PFlashSelectionPolicy & policy, + PFlashSelectionMode mode); + +const char * pflash_selection_mode_name(PFlashSelectionMode mode) noexcept; +const char * pflash_selection_stop_name(PFlashSelectionStop stop) noexcept; +const char * pflash_query_parser_name(PFlashQueryParser parser) noexcept; + +// How the context is cut into candidates and how a candidate is scored. +// ``Auto`` resolves at scoring time: probe segments when a segment probe is +// loaded, fixed chunks otherwise; density with probe segments, sum otherwise. +enum class PFlashSegmentation { Auto, Fixed, Probe }; +enum class PFlashCandidateScore { Auto, Sum, Density }; +// Which scorer ranks the candidates: the block-15 attention-mass head, the +// original all-layer running-max scorer, or both with a split budget (the +// head fills ``split_fraction`` of the budget first, the other scorer the rest). +enum class PFlashScorer { Head, Legacy, Split }; + +struct PFlashSelectionConfig { + PFlashSelectionMode mode = PFlashSelectionMode::Legacy; + PFlashQueryParser query_parser = PFlashQueryParser::SemanticUser; + int chunk_size = 0; + int query_tokens = 8; + double top_p = 0.95; + PFlashSegmentation segmentation = PFlashSegmentation::Auto; + PFlashCandidateScore candidate_score = PFlashCandidateScore::Auto; + PFlashScorer scorer = PFlashScorer::Head; + double split_fraction = 0.5; + bool configured = false; + bool selection_active = false; +}; + +// Segment probe: cut the context before every token whose boundary score is +// above ``threshold``; ``forced_cuts`` (query start, instruction span edges) +// are always cut; a cut closer than ``min_segment`` tokens to the previous +// accepted cut is dropped unless forced; a span longer than ``max_segment`` +// is split at its best-scoring interior token, or evenly when no interior +// token scores above zero. ``split_scores`` (the sub-unit logit when the +// probe artifact carries one) feeds only the oversize interior argmax; when +// empty the unit boundary scores are used. Returns contiguous spans covering +// [0, input_tokens), or an empty vector on invalid input. +std::vector pflash_probe_segments( + const std::vector & boundary_scores, + int input_tokens, + float threshold, + int min_segment, + int max_segment, + const std::vector & forced_cuts, + const std::vector & split_scores = {}); + +// Two-scorer selection: ``head`` candidates fill ``head_fraction`` of the +// budget (mandatory candidates first, charged once), then ``other`` +// candidates (same spans and ordinals, scored by the other scorer) fill what +// remains, skipping ordinals already selected. Both lists must describe the +// same spans in the same order. +PFlashSelectionResult select_pflash_split( + const std::vector & head, + const std::vector & other, + const PFlashSelectionPolicy & policy, + double head_fraction, + PFlashSelectionMode mode); + +const char * pflash_scorer_name(PFlashScorer scorer) noexcept; +const char * pflash_segmentation_name(PFlashSegmentation segmentation) noexcept; +const char * pflash_candidate_score_name(PFlashCandidateScore score) noexcept; + +// Presence, rather than validity, gates cache and continuation policy so an +// empty or invalid experiment variable cannot silently fall back to legacy. +bool has_pflash_selection_environment() noexcept; + +bool resolve_pflash_selection( + int input_tokens, + int legacy_chunk_size, + PFlashSelectionConfig & out, + std::string & error); + +} // namespace dflash::qwen3 diff --git a/server/src/qwen3/qwen35_drafter.cpp b/server/src/qwen3/qwen35_drafter.cpp new file mode 100644 index 000000000..7a3e1582c --- /dev/null +++ b/server/src/qwen3/qwen35_drafter.cpp @@ -0,0 +1,997 @@ +// Qwen3.5-0.8B drafter scoring for pflash speculative prefill. +// +// Two scorers share these weights: +// - qwen35_score_and_compress : the original all-layer running-max +// scorer, on the Qwen3.5 architecture +// - qwen35_strict_score_and_compress : blocks 0..14 plus the block-15 NoPE +// Q/K scoring head, under strict +// budget selection +// +// Loading lives in qwen35_loader.cpp; qwen3_drafter.cpp dispatches into +// qwen35_drafter_score_and_compress on DrafterArch::Qwen35_0p8b. + +#include "qwen35_drafter.h" + +#include "qwen3_drafter.h" +#include "qwen3_drafter_common.h" +#include "pflash_selection.h" +#include "common/gguf_inspect.h" +#include "qwen3/anchor_params.h" +#include "internal.h" + +#include "ggml.h" +#include "ggml-alloc.h" +#include "ggml-backend.h" +#include "gguf.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace dflash::common { + +namespace { + +static constexpr uint16_t F16_ZERO = 0x0000; +static constexpr uint16_t F16_NEG_INF = 0xFC00; + +static int align_up_i(int x, int a) { return ((x + a - 1) / a) * a; } + +static void build_causal_mask_f16(std::vector & out, int kv_len, int n_tokens, int kv_start) { + const int kv_pad = align_up_i(kv_len, 32); + const int q_pad = align_up_i(n_tokens, 32); + out.assign((size_t)kv_pad * q_pad, F16_NEG_INF); + static_assert(F16_ZERO == 0, "visible mask entries are zero-filled with memset"); + for (int q = 0; q < n_tokens; ++q) { + const int visible = std::min(kv_len, kv_start + q + 1); + if (visible > 0) { + std::memset(out.data() + (size_t)q * kv_pad, 0, (size_t)visible * sizeof(uint16_t)); + } + } +} + +// create_target_cache honours DFLASH27B_KV_TQ3; the drafter cache never wants +// the TurboQuant rotation, so force it off while the cache is created. +struct ScopedKvTq3Off { + ScopedKvTq3Off() { +#if defined(_WIN32) + char * raw = nullptr; + size_t len = 0; + _dupenv_s(&raw, &len, "DFLASH27B_KV_TQ3"); + had_ = raw != nullptr; + old_ = had_ ? raw : ""; + free(raw); + _putenv_s("DFLASH27B_KV_TQ3", "0"); +#else + const char * raw = std::getenv("DFLASH27B_KV_TQ3"); + had_ = raw != nullptr; + old_ = had_ ? raw : ""; + setenv("DFLASH27B_KV_TQ3", "0", 1); +#endif + } + ~ScopedKvTq3Off() { +#if defined(_WIN32) + // _putenv_s with empty value removes the variable on MSVCRT. + _putenv_s("DFLASH27B_KV_TQ3", had_ ? old_.c_str() : ""); +#else + if (had_) setenv("DFLASH27B_KV_TQ3", old_.c_str(), 1); + else unsetenv("DFLASH27B_KV_TQ3"); +#endif + } + bool had_ = false; + std::string old_; +}; + +} // namespace + +std::vector qwen35_score_and_compress( + TargetWeights & w, + const std::vector & ids, + float keep_ratio, + int chunk_size, + int n_lookahead, + int pool_kernel, + int score_query_end, + const dflash::qwen3::PFlashSelectionConfig & experiment, + const std::vector & required_instruction_spans, + std::vector * token_scores_out) { + + const int S = (int)ids.size(); + const int hidden = w.n_embd; + if (S < n_lookahead + 1) return ids; + const int query_end = score_query_end < 0 ? S : score_query_end; + if (n_lookahead < 1 || query_end < n_lookahead || query_end > S) { + set_last_error("qwen35 scorer query window out of range"); + return {}; + } + const int query_start = query_end - n_lookahead; + + auto t0 = std::chrono::steady_clock::now(); + std::vector running_max((size_t)n_lookahead * S, -INFINITY); + + TargetCache cache; + { + ScopedKvTq3Off tq3_off; + if (!create_target_cache(w, S, 0, w.backend, cache, true)) { + return {}; + } + } + + ggml_init_params act_ip{}; + act_ip.mem_size = (size_t)8 * ggml_tensor_overhead() + 4096; + act_ip.no_alloc = true; + ggml_context * act_ctx = ggml_init(act_ip); + if (!act_ctx) { + free_target_cache(cache); + set_last_error("qwen35 drafter activation ctx init failed"); + return {}; + } + ggml_tensor * act_in = ggml_new_tensor_2d(act_ctx, GGML_TYPE_F32, hidden, S); + ggml_tensor * act_out = ggml_new_tensor_2d(act_ctx, GGML_TYPE_F32, hidden, S); + ggml_backend_buffer_t act_buf = ggml_backend_alloc_ctx_tensors(act_ctx, w.backend); + if (!act_buf) { + ggml_free(act_ctx); + free_target_cache(cache); + set_last_error("qwen35 drafter activation allocation failed"); + return {}; + } + + { + const int batch = 2048; + std::vector emb((size_t)hidden * batch); + for (int i = 0; i < S; i += batch) { + const int n = std::min(batch, S - i); + if (!w.embedder.embed(ids.data() + i, n, emb.data())) { + ggml_backend_buffer_free(act_buf); ggml_free(act_ctx); free_target_cache(cache); + set_last_error("qwen35 drafter embedding failed"); + return {}; + } + ggml_backend_tensor_set(act_in, emb.data(), (size_t)i * act_in->nb[1], (size_t)hidden * n * sizeof(float)); + } + } + + ggml_gallocr_t alloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(w.backend)); + const int ubatch = 1024; + for (int il = 0; il < w.n_layer; ++il) { + const bool is_attn = (((il + 1) % w.full_attention_interval) == 0); + int fa_idx = 0; + if (is_attn) { + for (int k = 0; k < il; ++k) if (((k + 1) % w.full_attention_interval) == 0) ++fa_idx; + } + for (int start = 0; start < S; start += ubatch) { + const int n = std::min(ubatch, S - start); + const int kv_len = start + n; + + ggml_init_params ip{}; + ip.mem_size = 512 * 1024 * 1024; + ip.no_alloc = true; + ggml_context * ctx = ggml_init(ip); + if (!ctx) { + ggml_gallocr_free(alloc); ggml_backend_buffer_free(act_buf); ggml_free(act_ctx); free_target_cache(cache); + set_last_error("qwen35 drafter layer graph ctx init failed"); + return {}; + } + ggml_cgraph * gf = ggml_new_graph_custom(ctx, 16384, false); + ggml_tensor * inp = ggml_view_2d(ctx, act_in, hidden, n, act_in->nb[1], (size_t)start * act_in->nb[1]); + ggml_tensor * pos = nullptr; + ggml_tensor * mask = nullptr; + if (is_attn) { + pos = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 4 * n); + ggml_set_input(pos); + mask = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, align_up_i(kv_len, 32), align_up_i(n, 32)); + ggml_set_input(mask); + } + ggml_tensor * out = build_qwen35_layer(ctx, gf, w, cache, il, inp, pos, mask, start, n, false, 0); + ggml_tensor * dst = ggml_view_2d(ctx, act_out, hidden, n, act_out->nb[1], (size_t)start * act_out->nb[1]); + if (ggml_nelements(out) != ggml_nelements(dst)) { + std::fprintf(stderr, + "[qwen35-drafter] layer output shape mismatch il=%d start=%d out=[%lld,%lld,%lld,%lld] dst=[%lld,%lld,%lld,%lld]\n", + il, start, + (long long)out->ne[0], (long long)out->ne[1], (long long)out->ne[2], (long long)out->ne[3], + (long long)dst->ne[0], (long long)dst->ne[1], (long long)dst->ne[2], (long long)dst->ne[3]); + ggml_free(ctx); ggml_gallocr_free(alloc); ggml_backend_buffer_free(act_buf); ggml_free(act_ctx); free_target_cache(cache); + set_last_error("qwen35 layer output shape mismatch"); + return {}; + } + ggml_build_forward_expand(gf, ggml_cpy(ctx, out, dst)); + if (!ggml_gallocr_alloc_graph(alloc, gf)) { + ggml_free(ctx); ggml_gallocr_free(alloc); ggml_backend_buffer_free(act_buf); ggml_free(act_ctx); free_target_cache(cache); + set_last_error("qwen35 drafter graph allocation failed"); + return {}; + } + if (is_attn) { + std::vector p4((size_t)4 * n, 0); + for (int i = 0; i < n; ++i) { + int p = start + i; + p4[(size_t)0 * n + i] = p; + p4[(size_t)1 * n + i] = p; + p4[(size_t)2 * n + i] = p; + } + ggml_backend_tensor_set(pos, p4.data(), 0, p4.size() * sizeof(int32_t)); + std::vector m; + build_causal_mask_f16(m, kv_len, n, start); + ggml_backend_tensor_set(mask, m.data(), 0, m.size() * sizeof(uint16_t)); + } + auto st = ggml_backend_graph_compute(w.backend, gf); + ggml_free(ctx); + if (st != GGML_STATUS_SUCCESS) { + ggml_gallocr_free(alloc); ggml_backend_buffer_free(act_buf); ggml_free(act_ctx); free_target_cache(cache); + set_last_error("qwen35 drafter graph compute failed"); + return {}; + } + } + + if (is_attn) { + ggml_init_params sip{}; + sip.mem_size = ggml_tensor_overhead() * 32 + ggml_graph_overhead_custom(1024, false) + 64 * 1024; + sip.no_alloc = true; + ggml_context * sctx = ggml_init(sip); + if (!sctx) { + ggml_gallocr_free(alloc); ggml_backend_buffer_free(act_buf); ggml_free(act_ctx); free_target_cache(cache); + set_last_error("qwen35 score graph ctx allocation failed"); + return {}; + } + ggml_cgraph * sgf = ggml_new_graph_custom(sctx, 1024, false); + const int K_len = (int) cache.attn_k[(size_t)fa_idx]->ne[1]; + ggml_tensor * mask_tail = ggml_new_tensor_2d(sctx, GGML_TYPE_F32, K_len, n_lookahead); + ggml_tensor * K_f32 = ggml_new_tensor_3d(sctx, GGML_TYPE_F32, w.n_embd_head_k, K_len, w.n_head_kv); + ggml_tensor * K_cast = ggml_cpy(sctx, cache.attn_k[(size_t)fa_idx], K_f32); + ggml_tensor * K_score = nullptr; + if (w.n_head != w.n_head_kv) { + const int gqa = w.n_head / w.n_head_kv; + ggml_tensor * K_4d = ggml_reshape_4d(sctx, K_cast, w.n_embd_head_k, K_len, 1, w.n_head_kv); + ggml_tensor * K_tpl = ggml_new_tensor_4d(sctx, GGML_TYPE_F32, w.n_embd_head_k, K_len, gqa, w.n_head_kv); + ggml_tensor * K_rep = ggml_repeat(sctx, K_4d, K_tpl); + K_score = ggml_reshape_3d(sctx, K_rep, w.n_embd_head_k, K_len, w.n_head); + } else { + K_score = K_cast; + } + const TargetLayer & L = w.layers[il]; + ggml_tensor * inp_tail = ggml_view_2d(sctx, act_in, hidden, n_lookahead, + act_in->nb[1], (size_t)query_start * act_in->nb[1]); + ggml_tensor * q_cur = ggml_rms_norm(sctx, inp_tail, w.rms_eps); + q_cur = ggml_mul(sctx, q_cur, L.attn_norm); + ggml_tensor * QG = ggml_mul_mat(sctx, L.wq, q_cur); + QG = ggml_reshape_3d(sctx, QG, w.n_embd_head_k * 2, w.n_head, n_lookahead); + ggml_tensor * Q = ggml_view_3d(sctx, QG, + w.n_embd_head_k, w.n_head, n_lookahead, + ggml_element_size(QG) * w.n_embd_head_k * 2, + ggml_element_size(QG) * w.n_embd_head_k * 2 * w.n_head, + 0); + Q = ggml_rms_norm(sctx, Q, w.rms_eps); + Q = ggml_mul(sctx, Q, L.q_norm); + ggml_tensor * pos_tail = ggml_new_tensor_1d(sctx, GGML_TYPE_I32, 4 * n_lookahead); + int sections[4]; + for (int k = 0; k < 4; ++k) sections[k] = w.rope_sections[k]; + Q = ggml_rope_multi(sctx, Q, pos_tail, nullptr, + w.rope_dimension_count, sections, GGML_ROPE_TYPE_MROPE, + 0, w.rope_theta, 1.0f, + 0.0f, 1.0f, 0.0f, 0.0f); + ggml_tensor * Q_tail_perm = ggml_cont(sctx, ggml_permute(sctx, Q, 0, 2, 1, 3)); + ggml_tensor * attn_score = ggml_mul_mat(sctx, K_score, Q_tail_perm); + ggml_tensor * probs = ggml_soft_max_ext(sctx, attn_score, mask_tail, 1.0f / std::sqrt((float)w.n_embd_head_k), 0.0f); + ggml_set_output(probs); + ggml_build_forward_expand(sgf, probs); + ggml_gallocr_t salloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(w.backend)); + if (!ggml_gallocr_alloc_graph(salloc, sgf)) { + ggml_gallocr_free(salloc); ggml_free(sctx); ggml_gallocr_free(alloc); ggml_backend_buffer_free(act_buf); ggml_free(act_ctx); free_target_cache(cache); + set_last_error("qwen35 score graph allocation failed"); + return {}; + } + std::vector pos4((size_t)4 * n_lookahead, 0); + for (int i = 0; i < n_lookahead; ++i) { + const int p = query_start + i; + pos4[(size_t)0 * n_lookahead + i] = p; + pos4[(size_t)1 * n_lookahead + i] = p; + pos4[(size_t)2 * n_lookahead + i] = p; + } + ggml_backend_tensor_set(pos_tail, pos4.data(), 0, pos4.size() * sizeof(int32_t)); + std::vector mask((size_t)n_lookahead * K_len, 0.0f); + for (int t = 0; t < n_lookahead; ++t) { + const int visible_end = query_start + t + 1; + for (int j = 0; j < K_len; ++j) { + mask[(size_t)t * K_len + j] = (j < visible_end) ? 0.0f : -INFINITY; + } + } + ggml_backend_tensor_set(mask_tail, mask.data(), 0, mask.size() * sizeof(float)); + auto st = ggml_backend_graph_compute(w.backend, sgf); + if (st != GGML_STATUS_SUCCESS) { + ggml_gallocr_free(salloc); ggml_free(sctx); ggml_gallocr_free(alloc); ggml_backend_buffer_free(act_buf); ggml_free(act_ctx); free_target_cache(cache); + set_last_error("qwen35 score graph compute failed"); + return {}; + } + std::vector tmp((size_t)K_len * n_lookahead * w.n_head); + ggml_backend_tensor_get(probs, tmp.data(), 0, tmp.size() * sizeof(float)); + const size_t nonfinite = + count_nonfinite_scores(tmp.data(), tmp.size()); + if (nonfinite != 0) { + const std::string message = + "non-finite Qwen3.5 PFlash scores at layer " + + std::to_string(il) + ": " + std::to_string(nonfinite) + + "/" + std::to_string(tmp.size()); + std::fprintf(stderr, "[pflash] ERROR: %s\n", message.c_str()); + std::fflush(stderr); + ggml_gallocr_free(salloc); ggml_free(sctx); + ggml_gallocr_free(alloc); ggml_backend_buffer_free(act_buf); + ggml_free(act_ctx); free_target_cache(cache); + set_last_error(message); + return {}; + } + for (int h = 0; h < w.n_head; ++h) { + for (int t = 0; t < n_lookahead; ++t) { + for (int j = 0; j < S; ++j) { + const size_t src = (size_t)h * K_len * n_lookahead + (size_t)t * K_len + j; + const size_t dst = (size_t)t * S + j; + running_max[dst] = std::max(running_max[dst], tmp[src]); + } + } + } + ggml_gallocr_free(salloc); + ggml_free(sctx); + } + std::swap(act_in, act_out); + } + ggml_gallocr_free(alloc); + ggml_backend_buffer_free(act_buf); + ggml_free(act_ctx); + free_target_cache(cache); + + std::vector score((size_t)S, 0.0f); + for (int j = 0; j < S; ++j) { + float s = 0.0f; + for (int t = 0; t < n_lookahead; ++t) s += running_max[(size_t)t * S + j]; + score[(size_t)j] = s / (float)n_lookahead; + } + + const int n_chunks = (S + chunk_size - 1) / chunk_size; + const int n_keep = std::max(1, (int)((float)n_chunks * keep_ratio)); + + std::vector smooth_score = score; + // Caller pool_kernel takes precedence; if zero/negative, fall back to env or 5. + const int pk = (pool_kernel > 0) + ? pool_kernel + : std::max(3, env_int("DFLASH_COMPRESS_POOL_KERNEL", 5)); + std::vector smoothed((size_t)S, 0.0f); + int half = pk / 2; + for (int j = 0; j < S; ++j) { + int lo = std::max(0, j - half); + int hi = std::min(S - 1, j + half); + float s = 0.0f; + int n = 0; + for (int k = lo; k <= hi; ++k) { s += score[(size_t)k]; ++n; } + smoothed[(size_t)j] = (n > 0) ? (s / (float)n) : 0.0f; + } + smooth_score.swap(smoothed); + + if (token_scores_out) { + // Scoring only (two-scorer selection): hand the smoothed per-token + // scores back and let the caller select. + *token_scores_out = smooth_score; + return ids; + } + + if (experiment.selection_active) { + return select_pflash_chunks( + ids, smooth_score, keep_ratio, n_lookahead, score_query_end, + pk, experiment, required_instruction_spans, false, true); + } + + std::vector> chunk_means; + for (int c = 0; c < n_chunks; ++c) { + int lo = c * chunk_size, hi = std::min(S, lo + chunk_size); + float s = 0.0f; + for (int j = lo; j < hi; ++j) s += smooth_score[(size_t)j]; + chunk_means.push_back({s / std::max(1, hi - lo), c}); + } + std::sort(chunk_means.begin(), chunk_means.end(), [](auto a, auto b) { return a.first > b.first; }); + + std::vector selected((size_t)n_chunks, 0); + int count = 0; + // Scale head/tail forced chunks so they don't crowd out top-K scoring. + { + const int h_raw = env_int("DFLASH_COMPRESS_HEAD_CHUNKS", 8); + const int t_raw = env_int("DFLASH_COMPRESS_TAIL_CHUNKS", 24); + int h_n = h_raw, t_n = t_raw; + if (h_n + t_n >= n_keep) { + const int budget = std::max(1, n_keep - 1); + h_n = std::max(0, h_raw * budget / (h_raw + t_raw)); + t_n = std::max(0, budget - h_n); + } + for (int c = 0; c < std::min(n_chunks, h_n); ++c) { selected[(size_t)c] = 1; ++count; } + for (int c = std::max(0, n_chunks - t_n); c < n_chunks; ++c) if (!selected[(size_t)c]) { selected[(size_t)c] = 1; ++count; } + } + + const int query_tokens = env_int("DFLASH_COMPRESS_QUERY_TOKENS", 96); + const auto ap = resolve_anchor_params(n_chunks, + env_int("PFLASH_COMPRESS_ANCHOR_RADIUS", -1), + env_int("PFLASH_COMPRESS_MAX_ANCHOR_HITS", -1), + env_int("DFLASH_COMPRESS_ANCHOR_RADIUS", -1), + env_int("DFLASH_COMPRESS_MAX_ANCHOR_HITS", -1)); + const int anchor_radius = ap.radius; + const int max_anchor_hits = ap.max_hits; + std::vector forced((size_t)n_chunks, 0); + + const int q0 = std::max(0, S - query_tokens); + constexpr int NGRAM = 4; + for (int q = q0; q + NGRAM <= S; ++q) { + int hits = 0; + std::vector hit_pos(max_anchor_hits); + const int search_end = std::max(0, q0 - NGRAM); + for (int p = 0; p <= search_end && hits <= max_anchor_hits; ++p) { + bool same = true; + for (int k = 0; k < NGRAM; ++k) { + if (ids[(size_t)p + k] != ids[(size_t)q + k]) { same = false; break; } + } + if (same) { + if (hits < max_anchor_hits) hit_pos[hits] = p; + ++hits; + } + } + if (hits > 0 && hits <= max_anchor_hits) { + for (int i = 0; i < hits && i < max_anchor_hits; ++i) { + force_chunk_neighborhood(forced, n_chunks, hit_pos[i] / chunk_size, anchor_radius); + } + } + } + + for (int c = 0; c < n_chunks; ++c) { + if (forced[(size_t)c] && !selected[(size_t)c]) { + selected[(size_t)c] = 1; + ++count; + } + } + + // Global aggregation tasks often depend on repeated rare tokens that do + // not appear in the final query. Preserve high-frequency-but-not-filler + // token chunks before filling with model-score top-K. + const int repeat_min = env_int("DFLASH_COMPRESS_REPEAT_MIN", 4); + const int repeat_max = env_int("DFLASH_COMPRESS_REPEAT_MAX", 32); + const int repeat_limit = env_int("DFLASH_COMPRESS_REPEAT_CHUNKS", n_keep); + if (repeat_min > 1 && count < repeat_limit) { + std::unordered_map freq; + freq.reserve((size_t)S); + const int repeat_scan_end = std::max(0, S - query_tokens); + for (int j = 0; j < repeat_scan_end; ++j) { + ++freq[ids[(size_t)j]]; + } + std::vector> repeated; + repeated.reserve(freq.size()); + for (const auto & kv : freq) { + if (kv.second >= repeat_min && kv.second <= repeat_max) { + repeated.push_back({kv.second, kv.first}); + } + } + std::sort(repeated.begin(), repeated.end(), [](const auto & a, const auto & b) { + if (a.first != b.first) return a.first > b.first; + return a.second < b.second; + }); + for (const auto & rp : repeated) { + if (count >= repeat_limit) break; + const int32_t tok = rp.second; + for (int j = 0; j < repeat_scan_end && count < repeat_limit; ++j) { + if (ids[(size_t)j] != tok) continue; + const int c = j / chunk_size; + if (!selected[(size_t)c]) { + selected[(size_t)c] = 1; + ++count; + } + } + } + } + + for (auto [_, c] : chunk_means) { + if (count >= n_keep) break; + if (!selected[(size_t)c]) { selected[(size_t)c] = 1; ++count; } + } + + std::vector out_ids; + std::vector selected_chunks; + for (int c = 0; c < n_chunks; ++c) { + if (selected[(size_t)c]) selected_chunks.push_back(c); + } + int span_start = -1, span_end = -1; + for (int c : selected_chunks) { + int s_ = c * chunk_size; + int e_ = std::min(S, (c + 1) * chunk_size); + if (span_start < 0) { + span_start = s_; span_end = e_; + } else if (s_ == span_end) { + span_end = e_; + } else { + for (int j = span_start; j < span_end; ++j) out_ids.push_back(ids[j]); + span_start = s_; span_end = e_; + } + } + if (span_start >= 0) { + for (int j = span_start; j < span_end; ++j) out_ids.push_back(ids[j]); + } + + auto t1 = std::chrono::steady_clock::now(); + std::fprintf(stderr, "[qwen35-drafter] forward+compress %.2fs S=%d kept=%zu (%d/%d chunks)\n", + std::chrono::duration(t1 - t0).count(), S, out_ids.size(), count, n_chunks); + std::fflush(stderr); + return out_ids; +} + +// Scoring-head selection for the Qwen3.5-0.8B drafter: run blocks 0..14, then +// score every context token against the query window with block 15's NoPE +// Q/K (or a trained replacement) and select chunks by attention mass. This +// is the runtime counterpart of the Python retention screen (trial 0075). +std::vector qwen35_strict_score_and_compress( + Qwen35DrafterState & st, + const std::vector & ids, + float keep_ratio, + int n_lookahead, + int score_query_end, + const dflash::qwen3::PFlashSelectionConfig & experiment, + const std::vector & required_instruction_spans, + std::vector * token_mass_out, + std::vector * segments_out, + bool * density_out) { + + TargetWeights & w = st.weights; + const int S = (int)ids.size(); + const int hidden = w.n_embd; + const int H = w.n_head; + const int Hk = w.n_head_kv; + const int D = w.n_embd_head_k; + std::string block_error; + if (!qwen35_head_block_available(w, block_error)) { + set_last_error(block_error); + return {}; + } + if (n_lookahead < 1 || S < n_lookahead + 1) { + set_last_error("qwen35 scoring head input is too short"); + return {}; + } + const int query_end = score_query_end < 0 ? S : score_query_end; + if (query_end < n_lookahead || query_end > S) { + set_last_error("qwen35 scoring head query window out of range"); + return {}; + } + const int query_start = query_end - n_lookahead; + const TargetLayer & L = w.layers[(size_t)kQwen35HeadBlock]; + + auto t0 = std::chrono::steady_clock::now(); + TargetCache cache; + { + ScopedKvTq3Off tq3_off; + if (!create_target_cache(w, S, 0, w.backend, cache, true)) { + return {}; + } + } + + ggml_init_params act_ip{}; + act_ip.mem_size = (size_t)8 * ggml_tensor_overhead() + 4096; + act_ip.no_alloc = true; + ggml_context * act_ctx = ggml_init(act_ip); + if (!act_ctx) { + free_target_cache(cache); + set_last_error("qwen35 drafter activation ctx init failed"); + return {}; + } + ggml_tensor * act_in = ggml_new_tensor_2d(act_ctx, GGML_TYPE_F32, hidden, S); + ggml_tensor * act_out = ggml_new_tensor_2d(act_ctx, GGML_TYPE_F32, hidden, S); + ggml_backend_buffer_t act_buf = ggml_backend_alloc_ctx_tensors(act_ctx, w.backend); + if (!act_buf) { + ggml_free(act_ctx); + free_target_cache(cache); + set_last_error("qwen35 drafter activation allocation failed"); + return {}; + } + auto cleanup = [&]() { + ggml_backend_buffer_free(act_buf); + ggml_free(act_ctx); + free_target_cache(cache); + }; + + { + const int batch = 2048; + std::vector emb((size_t)hidden * batch); + for (int i = 0; i < S; i += batch) { + const int n = std::min(batch, S - i); + if (!w.embedder.embed(ids.data() + i, n, emb.data())) { + cleanup(); + set_last_error("qwen35 drafter embedding failed"); + return {}; + } + ggml_backend_tensor_set(act_in, emb.data(), (size_t)i * act_in->nb[1], + (size_t)hidden * n * sizeof(float)); + } + } + + ggml_gallocr_t alloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(w.backend)); + const int ubatch = 1024; + std::vector mask_bits; + for (int il = 0; il < kQwen35HeadBlock; ++il) { + const bool is_attn = (((il + 1) % w.full_attention_interval) == 0); + for (int start = 0; start < S; start += ubatch) { + const int n = std::min(ubatch, S - start); + const int kv_len = start + n; + ggml_init_params ip{}; + ip.mem_size = 512 * 1024 * 1024; + ip.no_alloc = true; + ggml_context * ctx = ggml_init(ip); + if (!ctx) { + ggml_gallocr_free(alloc); cleanup(); + set_last_error("qwen35 drafter layer graph ctx init failed"); + return {}; + } + ggml_cgraph * gf = ggml_new_graph_custom(ctx, 16384, false); + ggml_tensor * inp = ggml_view_2d(ctx, act_in, hidden, n, act_in->nb[1], + (size_t)start * act_in->nb[1]); + ggml_tensor * pos = nullptr; + ggml_tensor * mask = nullptr; + if (is_attn) { + pos = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 4 * n); + ggml_set_input(pos); + mask = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, + align_up_i(kv_len, 32), align_up_i(n, 32)); + ggml_set_input(mask); + } + ggml_tensor * out = build_qwen35_layer(ctx, gf, w, cache, il, inp, pos, mask, + start, n, false, 0); + ggml_tensor * dst = ggml_view_2d(ctx, act_out, hidden, n, act_out->nb[1], + (size_t)start * act_out->nb[1]); + if (ggml_nelements(out) != ggml_nelements(dst)) { + ggml_free(ctx); ggml_gallocr_free(alloc); cleanup(); + set_last_error("qwen35 layer output shape mismatch"); + return {}; + } + ggml_build_forward_expand(gf, ggml_cpy(ctx, out, dst)); + if (!ggml_gallocr_alloc_graph(alloc, gf)) { + ggml_free(ctx); ggml_gallocr_free(alloc); cleanup(); + set_last_error("qwen35 drafter graph allocation failed"); + return {}; + } + if (is_attn) { + std::vector p4((size_t)4 * n, 0); + for (int i = 0; i < n; ++i) { + const int p = start + i; + p4[(size_t)0 * n + i] = p; + p4[(size_t)1 * n + i] = p; + p4[(size_t)2 * n + i] = p; + } + ggml_backend_tensor_set(pos, p4.data(), 0, p4.size() * sizeof(int32_t)); + build_causal_mask_f16(mask_bits, kv_len, n, start); + ggml_backend_tensor_set(mask, mask_bits.data(), 0, + mask_bits.size() * sizeof(uint16_t)); + } + const auto status = ggml_backend_graph_compute(w.backend, gf); + ggml_free(ctx); + if (status != GGML_STATUS_SUCCESS) { + ggml_gallocr_free(alloc); cleanup(); + set_last_error("qwen35 drafter graph compute failed"); + return {}; + } + } + std::swap(act_in, act_out); + } + ggml_gallocr_free(alloc); + auto t1 = std::chrono::steady_clock::now(); + + // Block-15 NoPE Q/K scoring: softmax over keys before the query window, + // then mean over heads and query tokens. The query never scores itself. + // Keys are projected in chunks so no intermediate tensor puts the + // sequence length into a HIP grid y/z dimension (65,535 limit); the + // logits land in one [S, n_lookahead, H] buffer for a single softmax. + const int key_chunk = 8192; + const int n_key_chunks = (S + key_chunk - 1) / key_chunk; + ggml_init_params lip{}; + lip.mem_size = (size_t)8 * ggml_tensor_overhead() + 4096; + lip.no_alloc = true; + ggml_context * lctx = ggml_init(lip); + if (!lctx) { + cleanup(); + set_last_error("qwen35 score buffer ctx allocation failed"); + return {}; + } + ggml_tensor * logits = ggml_new_tensor_3d(lctx, GGML_TYPE_F32, S, n_lookahead, H); + ggml_tensor * mask = ggml_new_tensor_2d(lctx, GGML_TYPE_F32, S, n_lookahead); + const bool use_probe = st.probe_loaded && + experiment.segmentation != dflash::qwen3::PFlashSegmentation::Fixed; + ggml_tensor * probe_logits = use_probe + ? ggml_new_tensor_1d(lctx, GGML_TYPE_F32, S) : nullptr; + ggml_tensor * subunit_logits = use_probe && st.probe_sub_fc2_w + ? ggml_new_tensor_1d(lctx, GGML_TYPE_F32, S) : nullptr; + ggml_backend_buffer_t lbuf = ggml_backend_alloc_ctx_tensors(lctx, w.backend); + if (!lbuf) { + ggml_free(lctx); cleanup(); + set_last_error("qwen35 score buffer allocation failed"); + return {}; + } + { + std::vector m((size_t)n_lookahead * S, -INFINITY); + for (int t = 0; t < n_lookahead; ++t) { + std::fill_n(m.begin() + (size_t)t * S, (size_t)query_start, 0.0f); + } + ggml_backend_tensor_set(mask, m.data(), 0, m.size() * sizeof(float)); + } + ggml_init_params sip{}; + sip.mem_size = ggml_tensor_overhead() * (size_t)(64 + 24 * n_key_chunks) + + ggml_graph_overhead_custom(4096, false) + 64 * 1024; + sip.no_alloc = true; + ggml_context * sctx = ggml_init(sip); + if (!sctx) { + ggml_backend_buffer_free(lbuf); ggml_free(lctx); cleanup(); + set_last_error("qwen35 score graph ctx allocation failed"); + return {}; + } + ggml_cgraph * sgf = ggml_new_graph_custom(sctx, 4096, false); + ggml_tensor * wk_src = st.head_loaded ? st.head_wk : L.wk; + ggml_tensor * x_q = ggml_view_2d(sctx, act_in, hidden, n_lookahead, act_in->nb[1], + (size_t)query_start * act_in->nb[1]); + ggml_tensor * q_in = ggml_mul(sctx, ggml_rms_norm(sctx, x_q, w.rms_eps), L.attn_norm); + ggml_tensor * Q = nullptr; + if (st.head_loaded) { + Q = ggml_reshape_3d(sctx, ggml_mul_mat(sctx, st.head_wq, q_in), D, H, n_lookahead); + } else { + // Native block 15 packs query and gate rows per head; keep the query half. + ggml_tensor * QG = ggml_reshape_3d(sctx, ggml_mul_mat(sctx, L.wq, q_in), + D * 2, H, n_lookahead); + Q = ggml_view_3d(sctx, QG, D, H, n_lookahead, + ggml_element_size(QG) * D * 2, + ggml_element_size(QG) * D * 2 * H, 0); + } + Q = ggml_mul(sctx, ggml_rms_norm(sctx, Q, w.rms_eps), L.q_norm); + ggml_tensor * Q_perm = ggml_cont(sctx, ggml_permute(sctx, Q, 0, 2, 1, 3)); // [D, n_lookahead, H] + for (int b = 0; b < S; b += key_chunk) { + const int n = std::min(key_chunk, S - b); + ggml_tensor * x_c = ggml_view_2d(sctx, act_in, hidden, n, act_in->nb[1], + (size_t)b * act_in->nb[1]); + ggml_tensor * x_norm = ggml_mul(sctx, ggml_rms_norm(sctx, x_c, w.rms_eps), L.attn_norm); + ggml_tensor * K = ggml_reshape_3d(sctx, ggml_mul_mat(sctx, wk_src, x_norm), D, Hk, n); + K = ggml_mul(sctx, ggml_rms_norm(sctx, K, w.rms_eps), L.k_norm); + K = ggml_cont(sctx, ggml_permute(sctx, K, 0, 2, 1, 3)); // [D, n, Hk] + ggml_tensor * K_score = K; + if (H != Hk) { + const int gqa = H / Hk; + ggml_tensor * K_4d = ggml_reshape_4d(sctx, K, D, n, 1, Hk); + ggml_tensor * K_tpl = ggml_new_tensor_4d(sctx, GGML_TYPE_F32, D, n, gqa, Hk); + K_score = ggml_reshape_3d(sctx, ggml_repeat(sctx, K_4d, K_tpl), D, n, H); + } + ggml_tensor * part = ggml_mul_mat(sctx, K_score, Q_perm); // [n, n_lookahead, H] + ggml_tensor * dst = ggml_view_3d(sctx, logits, n, n_lookahead, H, + logits->nb[1], logits->nb[2], + (size_t)b * logits->nb[0]); + ggml_build_forward_expand(sgf, ggml_cpy(sctx, part, dst)); + if (use_probe) { + // Segment probe on the same tap: LayerNorm -> fc1 -> GELU trunk, + // then one fc2 row per head (unit always; sub-unit when shipped). + ggml_tensor * p = ggml_norm(sctx, x_c, st.probe_norm_eps); + p = ggml_add(sctx, ggml_mul(sctx, p, st.probe_norm_w), st.probe_norm_b); + p = ggml_gelu(sctx, ggml_add(sctx, ggml_mul_mat(sctx, st.probe_fc1_w, p), + st.probe_fc1_b)); // [width, n] + ggml_tensor * unit = ggml_add(sctx, ggml_mul_mat(sctx, st.probe_fc2_w, p), + st.probe_fc2_b); // [1, n] + ggml_tensor * p_dst = ggml_view_1d(sctx, probe_logits, n, + (size_t)b * ggml_element_size(probe_logits)); + ggml_build_forward_expand(sgf, ggml_cpy(sctx, ggml_reshape_1d(sctx, unit, n), p_dst)); + if (subunit_logits) { + ggml_tensor * sub = ggml_add(sctx, + ggml_mul_mat(sctx, st.probe_sub_fc2_w, p), st.probe_sub_fc2_b); + ggml_tensor * s_dst = ggml_view_1d(sctx, subunit_logits, n, + (size_t)b * ggml_element_size(subunit_logits)); + ggml_build_forward_expand(sgf, + ggml_cpy(sctx, ggml_reshape_1d(sctx, sub, n), s_dst)); + } + } + } + ggml_tensor * probs = ggml_soft_max_ext(sctx, logits, mask, + 1.0f / std::sqrt((float)D), 0.0f); + ggml_set_output(probs); + ggml_build_forward_expand(sgf, probs); + ggml_gallocr_t salloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(w.backend)); + if (!ggml_gallocr_alloc_graph(salloc, sgf)) { + ggml_gallocr_free(salloc); ggml_free(sctx); + ggml_backend_buffer_free(lbuf); ggml_free(lctx); cleanup(); + set_last_error("qwen35 score graph allocation failed"); + return {}; + } + const auto score_status = ggml_backend_graph_compute(w.backend, sgf); + if (score_status != GGML_STATUS_SUCCESS) { + ggml_gallocr_free(salloc); ggml_free(sctx); + ggml_backend_buffer_free(lbuf); ggml_free(lctx); cleanup(); + set_last_error("qwen35 score graph compute failed"); + return {}; + } + std::vector probs_h((size_t)S * n_lookahead * H); + ggml_backend_tensor_get(probs, probs_h.data(), 0, probs_h.size() * sizeof(float)); + std::vector probe_raw; + std::vector subunit_raw; + if (use_probe) { + probe_raw.resize((size_t) S); + ggml_backend_tensor_get(probe_logits, probe_raw.data(), 0, probe_raw.size() * sizeof(float)); + if (subunit_logits) { + subunit_raw.resize((size_t) S); + ggml_backend_tensor_get(subunit_logits, subunit_raw.data(), 0, + subunit_raw.size() * sizeof(float)); + } + } + ggml_gallocr_free(salloc); + ggml_free(sctx); + ggml_backend_buffer_free(lbuf); + ggml_free(lctx); + cleanup(); + const size_t nonfinite = count_nonfinite_scores(probs_h.data(), probs_h.size()); + if (nonfinite != 0) { + const std::string message = + "non-finite Qwen3.5 scoring-head scores: " + std::to_string(nonfinite) + + "/" + std::to_string(probs_h.size()); + std::fprintf(stderr, "[pflash] ERROR: %s\n", message.c_str()); + std::fflush(stderr); + set_last_error(message); + return {}; + } + std::vector token_mass; + scoring_head_mean_token_mass(probs_h.data(), S, n_lookahead, H, token_mass); + auto t2 = std::chrono::steady_clock::now(); + std::fprintf(stderr, + "[qwen35-scorer] forward %.2fs (blocks 0-%d, S=%d) score %.2fs " + "total %.2fs head=%s\n", + std::chrono::duration(t1 - t0).count(), kQwen35HeadBlock - 1, S, + std::chrono::duration(t2 - t1).count(), + std::chrono::duration(t2 - t0).count(), + st.head_loaded ? "trained" : "native-block15"); + std::fflush(stderr); + + std::vector segments; + bool density = experiment.candidate_score == dflash::qwen3::PFlashCandidateScore::Density; + if (use_probe) { + // Tap-count smoothing over the raw logits (torch Conv1d, symmetric + // padding) plus the residual logit, then sigmoid: the boundary score + // per token. + const auto smooth = [&](const std::vector & raw, + const std::vector & conv_w, float conv_b, + std::vector & out) { + const int taps = (int) conv_w.size(); + const int radius = taps / 2; + out.assign((size_t) S, 0.0f); + for (int t = 0; t < S; ++t) { + float acc = raw[(size_t) t] + conv_b; + for (int k = 0; k < taps; ++k) { + const int u = t + k - radius; + if (u >= 0 && u < S) acc += conv_w[(size_t) k] * raw[(size_t) u]; + } + out[(size_t) t] = 1.0f / (1.0f + std::exp(-acc)); + } + }; + std::vector boundary; + smooth(probe_raw, st.probe_conv_w, st.probe_conv_b, boundary); + // Sub-unit scores feed only the oversize interior argmax below. + std::vector split_scores; + if (!subunit_raw.empty()) { + smooth(subunit_raw, st.probe_sub_conv_w, st.probe_sub_conv_b, split_scores); + } + const int query_end = score_query_end < 0 ? S : score_query_end; + const int query_begin = query_end - std::min(n_lookahead, query_end); + std::vector forced{query_begin, query_end}; + for (const auto & span : required_instruction_spans) { + forced.push_back(span.begin); + forced.push_back(span.end); + } + int boundaries_in_context = 0; + for (int t = 1; t < query_begin; ++t) { + if (boundary[(size_t) t] > st.probe_threshold) ++boundaries_in_context; + } + const bool forced_probe = + experiment.segmentation == dflash::qwen3::PFlashSegmentation::Probe; + if (boundaries_in_context >= 4 || forced_probe) { + segments = dflash::qwen3::pflash_probe_segments( + boundary, S, st.probe_threshold, st.probe_min_segment, + st.probe_max_segment, forced, split_scores); + } + if (segments.empty()) { + std::fprintf(stderr, + "[qwen35-segment-probe] %d boundaries in the context, " + "falling back to fixed %d-token chunks\n", + boundaries_in_context, experiment.chunk_size); + } else { + if (experiment.candidate_score == dflash::qwen3::PFlashCandidateScore::Auto) { + density = true; + } + std::fprintf(stderr, + "[qwen35-segment-probe] %d boundaries in the context -> %zu segments " + "(threshold %.2f, %d-%d tokens), score=%s\n", + boundaries_in_context, segments.size(), st.probe_threshold, + st.probe_min_segment, st.probe_max_segment, + density ? "density" : "sum"); + } + std::fflush(stderr); + } + + if (token_mass_out) { + // Scoring only (two-scorer selection): return the per-token mass, + // the probe segments and the ranking rule; the caller selects. + *token_mass_out = token_mass; + if (segments_out) *segments_out = segments; + if (density_out) *density_out = density; + return ids; + } + + return select_pflash_chunks( + ids, token_mass, keep_ratio, n_lookahead, score_query_end, + /*pool_kernel=*/1, experiment, required_instruction_spans, + /*direct_mass=*/true, /*write_trace=*/true, + segments.empty() ? nullptr : &segments, density); +} + +std::vector qwen35_drafter_score_and_compress( + DrafterContext & ctx, + const std::vector & ids, + float keep_ratio, + int chunk_size, + int n_lookahead, + int pool_kernel, + int score_query_end, + const dflash::qwen3::PFlashSelectionConfig & experiment, + const std::vector & required_instruction_spans) { + if (!ctx.arch_state) { + set_last_error("qwen35 drafter state missing"); + return {}; + } + auto * st = static_cast(ctx.arch_state); + // Strict budget selection scores with the block-15 head; the + // legacy all-layer running-max scorer stays available for legacy + // selection or when PFLASH_QWEN35_LEGACY_SCORER=1 forces it. + const char * legacy_scorer = std::getenv("PFLASH_QWEN35_LEGACY_SCORER"); + const bool force_legacy = (legacy_scorer && std::string(legacy_scorer) == "1") || + experiment.scorer == dflash::qwen3::PFlashScorer::Legacy; + if (experiment.selection_active && + experiment.scorer == dflash::qwen3::PFlashScorer::Split) { + // Two scorers, one budget: the block-15 head ranks (and segments) + // first, the all-layer running-max scorer fills the remainder. + std::vector head_mass; + std::vector head_segments; + bool head_density = false; + if (qwen35_strict_score_and_compress( + *st, ids, keep_ratio, n_lookahead, score_query_end, experiment, + required_instruction_spans, &head_mass, &head_segments, + &head_density).empty()) { + return {}; + } + std::vector other_scores; + if (qwen35_score_and_compress(st->weights, ids, keep_ratio, chunk_size, + n_lookahead, pool_kernel, score_query_end, + experiment, required_instruction_spans, + &other_scores).empty()) { + return {}; + } + if (other_scores.size() != head_mass.size()) { + set_last_error("two-scorer selection: score lengths differ"); + return {}; + } + std::fprintf(stderr, + "[pflash-select] two-scorer selection: head fraction %.2f, " + "segments=%s\n", experiment.split_fraction, + head_segments.empty() ? "fixed" : "probe"); + std::fflush(stderr); + return select_pflash_chunks( + ids, head_mass, keep_ratio, n_lookahead, score_query_end, + /*pool_kernel=*/1, experiment, required_instruction_spans, + /*direct_mass=*/true, /*write_trace=*/true, + head_segments.empty() ? nullptr : &head_segments, head_density, + &other_scores, experiment.split_fraction); + } + if (experiment.selection_active && !force_legacy) { + return qwen35_strict_score_and_compress( + *st, ids, keep_ratio, n_lookahead, score_query_end, experiment, + required_instruction_spans); + } + if (st->head_loaded && !experiment.selection_active) { + set_last_error("Qwen3.5 scoring head requires strict selection"); + return {}; + } + return qwen35_score_and_compress(st->weights, ids, keep_ratio, chunk_size, + n_lookahead, pool_kernel, score_query_end, + experiment, + required_instruction_spans); +} + +} // namespace dflash::common diff --git a/server/src/qwen3/qwen35_drafter.h b/server/src/qwen3/qwen35_drafter.h new file mode 100644 index 000000000..2a391118c --- /dev/null +++ b/server/src/qwen3/qwen35_drafter.h @@ -0,0 +1,108 @@ +// Internal interface of the Qwen3.5-0.8B drafter. +// +// The Qwen3.5-0.8B scorer runs on the Qwen3.5 target architecture +// (TargetWeights, build_qwen35_layer) rather than the Qwen3-0.6B drafter +// graph, so it lives in its own translation units: qwen35_loader.cpp loads +// the GGUF, the scoring head and the segment probe; qwen35_drafter.cpp runs +// the two scorers. qwen3_drafter.cpp dispatches here on DrafterArch. + +#pragma once + +#include "qwen3_drafter.h" +#include "pflash_selection.h" +#include "common/pflash_types.h" +#include "internal.h" + +#include "ggml.h" +#include "ggml-backend.h" + +#include +#include + +namespace dflash::common { + +// Qwen3.5-0.8B scoring head. Features are the residual entering +// full-attention block 15 after the first 15 blocks (twelve GatedDeltaNet and +// three full-attention blocks). Block 15's own Q/K projections score the +// context without RoPE, exactly like the Qwen3-0.6B block-13 head; an +// optional trained head replaces those two projections. +static constexpr int kQwen35HeadBlock = 15; + +struct Qwen35DrafterState { + TargetWeights weights; + std::string gguf_sha256; + ggml_context * head_ctx = nullptr; + ggml_backend_buffer_t head_buf = nullptr; + ggml_tensor * head_wq = nullptr; // [hidden, n_head * head_dim], query rows only + ggml_tensor * head_wk = nullptr; // [hidden, n_head_kv * head_dim] + bool head_loaded = false; + // Segment probe: per-token boundary scores from the same block-14 tap. + ggml_context * probe_ctx = nullptr; + ggml_backend_buffer_t probe_buf = nullptr; + ggml_tensor * probe_norm_w = nullptr; // [hidden] + ggml_tensor * probe_norm_b = nullptr; // [hidden] + ggml_tensor * probe_fc1_w = nullptr; // [hidden, probe_width] + ggml_tensor * probe_fc1_b = nullptr; // [probe_width] + ggml_tensor * probe_fc2_w = nullptr; // [probe_width, 1] + ggml_tensor * probe_fc2_b = nullptr; // [1] + ggml_tensor * probe_sub_fc2_w = nullptr; // optional sub-unit head (oversize split only) + ggml_tensor * probe_sub_fc2_b = nullptr; + std::vector probe_conv_w; // taps from the GGUF tensor, applied on the CPU + float probe_conv_b = 0.0f; + std::vector probe_sub_conv_w; + float probe_sub_conv_b = 0.0f; + float probe_norm_eps = 1e-5f; + float probe_threshold = 0.9f; + int probe_min_segment = 1; + int probe_max_segment = 2048; + int probe_width = 0; + bool probe_loaded = false; +}; + +// Defined in qwen35_loader.cpp. +bool qwen35_head_block_available(const TargetWeights & w, std::string & error); +bool load_qwen35_drafter(const std::string & gguf_path, DrafterArch arch, + DrafterContext & out); +void free_qwen35_drafter_state(DrafterContext & ctx); + +// Defined in qwen35_drafter.cpp. +// +// The legacy all-layer running-max scorer, on the Qwen3.5 architecture. +std::vector qwen35_score_and_compress( + TargetWeights & w, + const std::vector & ids, + float keep_ratio, + int chunk_size, + int n_lookahead, + int pool_kernel, + int score_query_end, + const dflash::qwen3::PFlashSelectionConfig & experiment, + const std::vector & required_instruction_spans, + std::vector * token_scores_out = nullptr); + +// The block-15 scoring head under strict budget selection. +std::vector qwen35_strict_score_and_compress( + Qwen35DrafterState & st, + const std::vector & ids, + float keep_ratio, + int n_lookahead, + int score_query_end, + const dflash::qwen3::PFlashSelectionConfig & experiment, + const std::vector & required_instruction_spans, + std::vector * token_mass_out = nullptr, + std::vector * segments_out = nullptr, + bool * density_out = nullptr); + +// Arch dispatch target of drafter_score_and_compress. +std::vector qwen35_drafter_score_and_compress( + DrafterContext & ctx, + const std::vector & ids, + float keep_ratio, + int chunk_size, + int n_lookahead, + int pool_kernel, + int score_query_end, + const dflash::qwen3::PFlashSelectionConfig & experiment, + const std::vector & required_instruction_spans); + +} // namespace dflash::common diff --git a/server/src/qwen3/qwen35_loader.cpp b/server/src/qwen3/qwen35_loader.cpp new file mode 100644 index 000000000..099c87312 --- /dev/null +++ b/server/src/qwen3/qwen35_loader.cpp @@ -0,0 +1,376 @@ +// Qwen3.5-0.8B drafter loading: the drafter GGUF, the optional trained +// block-15 scoring head and the optional segment probe. +// +// The Qwen3-0.6B drafter has qwen3_loader.cpp; this is its counterpart for +// the Qwen3.5-0.8B scorer, which is built on the Qwen3.5 target weights +// (load_target_gguf_partial) instead of the Qwen3-0.6B drafter weights. + +#include "qwen35_drafter.h" + +#include "qwen3_drafter.h" +#include "common/gguf_inspect.h" +#include "internal.h" + +#include "ggml.h" +#include "ggml-backend.h" +#include "gguf.h" + +#include +#include +#include +#include + +namespace dflash::common { + +bool qwen35_head_block_available(const TargetWeights & w, std::string & error) { + if (w.n_layer <= kQwen35HeadBlock || (size_t)kQwen35HeadBlock >= w.layers.size()) { + error = "qwen35 scoring head needs at least 16 blocks"; + return false; + } + const TargetLayer & L = w.layers[(size_t)kQwen35HeadBlock]; + if (((kQwen35HeadBlock + 1) % w.full_attention_interval) != 0 || + !L.wq || !L.wk || !L.attn_norm || !L.q_norm || !L.k_norm) { + error = "qwen35 scoring head block 15 is not a full-attention block"; + return false; + } + return true; +} + +namespace { + +static constexpr const char * kQwen35HeadSchema = "qwen3_5_0_8b_nope_qk_mass_v1"; +static constexpr const char * kQwen35HeadBaseModel = "Qwen/Qwen3.5-0.8B"; +static constexpr const char * kQwen35HeadFeatureTap = + "post_block14_residual_before_block15"; + +static void free_qwen35_head(Qwen35DrafterState & st) { + if (st.head_buf) { ggml_backend_buffer_free(st.head_buf); st.head_buf = nullptr; } + if (st.head_ctx) { ggml_free(st.head_ctx); st.head_ctx = nullptr; } + st.head_wq = st.head_wk = nullptr; + st.head_loaded = false; +} + +static void free_qwen35_segment_probe(Qwen35DrafterState & st) { + if (st.probe_buf) { ggml_backend_buffer_free(st.probe_buf); st.probe_buf = nullptr; } + if (st.probe_ctx) { ggml_free(st.probe_ctx); st.probe_ctx = nullptr; } + st.probe_norm_w = st.probe_norm_b = st.probe_fc1_w = st.probe_fc1_b = + st.probe_fc2_w = st.probe_fc2_b = + st.probe_sub_fc2_w = st.probe_sub_fc2_b = nullptr; + st.probe_conv_w.clear(); + st.probe_sub_conv_w.clear(); + st.probe_loaded = false; +} + +static bool qwen35_metadata_equals(gguf_context * g, const char * key, + const std::string & expected) { + const int id = gguf_find_key(g, key); + return id >= 0 && gguf_get_kv_type(g, id) == GGUF_TYPE_STRING && + expected == gguf_get_val_str(g, id); +} + +// Optional trained head for the block-15 tap. Fails closed on any contract +// mismatch, mirroring the Qwen3-0.6B head loader. +static bool load_qwen35_scoring_head(const std::string & path, + Qwen35DrafterState & st) { + const TargetWeights & w = st.weights; + std::string block_error; + if (!qwen35_head_block_available(w, block_error)) { + set_last_error(block_error); + return false; + } + if (st.gguf_sha256.empty()) { + set_last_error("scoring head requires the drafter GGUF identity hash"); + return false; + } + ggml_context * data_ctx = nullptr; + gguf_init_params params{ /*no_alloc=*/ false, /*ctx=*/ &data_ctx }; + gguf_context * g = gguf_init_from_file(path.c_str(), params); + if (!g) { + set_last_error("scoring head GGUF could not be opened: " + path); + return false; + } + auto fail = [&](const std::string & message) { + free_qwen35_head(st); + gguf_free(g); + if (data_ctx) ggml_free(data_ctx); + set_last_error(message); + return false; + }; + // GGUF contract of a scoring-head file: architecture `pflash_scoring_head`, + // metadata and tensors under `scoringhead.*`. + if (!qwen35_metadata_equals(g, "general.architecture", "pflash_scoring_head") || + !qwen35_metadata_equals(g, "scoringhead.schema", kQwen35HeadSchema) || + !qwen35_metadata_equals(g, "scoringhead.base_model", kQwen35HeadBaseModel) || + !qwen35_metadata_equals(g, "scoringhead.runtime_gguf_sha256", st.gguf_sha256) || + !qwen35_metadata_equals(g, "scoringhead.feature_tap", kQwen35HeadFeatureTap)) { + return fail("scoring head metadata does not match the loaded Qwen3.5-0.8B drafter"); + } + struct Contract { + const char * name; + int64_t ne0; + int64_t ne1; + ggml_tensor ** destination; + }; + const Contract contracts[] = { + {"scoringhead.attn_q.weight", (int64_t)w.n_embd, + (int64_t)w.n_head * w.n_embd_head_k, &st.head_wq}, + {"scoringhead.attn_k.weight", (int64_t)w.n_embd, + (int64_t)w.n_head_kv * w.n_embd_head_k, &st.head_wk}, + }; + ggml_init_params head_params{}; + head_params.mem_size = 4 * ggml_tensor_overhead(); + head_params.no_alloc = true; + st.head_ctx = ggml_init(head_params); + if (!st.head_ctx) return fail("scoring head context allocation failed"); + for (const auto & contract : contracts) { + ggml_tensor * source = data_ctx ? ggml_get_tensor(data_ctx, contract.name) : nullptr; + if (!source || source->type != GGML_TYPE_F32 || ggml_n_dims(source) != 2 || + source->ne[0] != contract.ne0 || source->ne[1] != contract.ne1) { + return fail(std::string("scoring head tensor contract mismatch: ") + + contract.name); + } + *contract.destination = + ggml_new_tensor_2d(st.head_ctx, GGML_TYPE_F32, contract.ne0, contract.ne1); + ggml_set_name(*contract.destination, contract.name); + } + st.head_buf = ggml_backend_alloc_ctx_tensors(st.head_ctx, w.backend); + if (!st.head_buf) return fail("scoring head buffer allocation failed"); + for (const auto & contract : contracts) { + ggml_tensor * source = ggml_get_tensor(data_ctx, contract.name); + ggml_backend_tensor_set(*contract.destination, source->data, 0, ggml_nbytes(source)); + } + gguf_free(g); + ggml_free(data_ctx); + st.head_loaded = true; + std::fprintf(stderr, "[qwen35-drafter] loaded scoring head: %s\n", path.c_str()); + std::fflush(stderr); + return true; +} + +static constexpr const char * kQwen35ProbeSchema = "qwen3_5_0_8b_segment_probe_v1"; +static constexpr const char * kQwen35ProbeSchemaV2 = "qwen3_5_0_8b_segment_probe_v2"; + +static bool qwen35_metadata_f32(gguf_context * g, const char * key, float & out) { + const int id = gguf_find_key(g, key); + if (id < 0 || gguf_get_kv_type(g, id) != GGUF_TYPE_FLOAT32) return false; + out = gguf_get_val_f32(g, id); + return true; +} + +static bool qwen35_metadata_u32(gguf_context * g, const char * key, int & out) { + const int id = gguf_find_key(g, key); + if (id < 0 || gguf_get_kv_type(g, id) != GGUF_TYPE_UINT32) return false; + out = (int) gguf_get_val_u32(g, id); + return true; +} + +// Optional segment probe for the block-14 tap: LayerNorm -> Linear -> GELU -> +// Linear on the GPU, a 5-tap smoothing on the CPU, sigmoid, cut above the +// threshold. Fails closed on any contract mismatch, like the head loader. +static bool load_qwen35_segment_probe(const std::string & path, + Qwen35DrafterState & st) { + const TargetWeights & w = st.weights; + if (st.gguf_sha256.empty()) { + set_last_error("segment probe requires the drafter GGUF identity hash"); + return false; + } + ggml_context * data_ctx = nullptr; + gguf_init_params params{ /*no_alloc=*/ false, /*ctx=*/ &data_ctx }; + gguf_context * g = gguf_init_from_file(path.c_str(), params); + if (!g) { + set_last_error("segment probe GGUF could not be opened: " + path); + return false; + } + auto fail = [&](const std::string & message) { + free_qwen35_segment_probe(st); + gguf_free(g); + if (data_ctx) ggml_free(data_ctx); + set_last_error(message); + return false; + }; + if (!qwen35_metadata_equals(g, "general.architecture", "segmentprobe") || + (!qwen35_metadata_equals(g, "segmentprobe.schema", kQwen35ProbeSchema) && + !qwen35_metadata_equals(g, "segmentprobe.schema", kQwen35ProbeSchemaV2)) || + !qwen35_metadata_equals(g, "segmentprobe.base_model", kQwen35HeadBaseModel) || + !qwen35_metadata_equals(g, "segmentprobe.runtime_gguf_sha256", st.gguf_sha256) || + !qwen35_metadata_equals(g, "segmentprobe.feature_tap", kQwen35HeadFeatureTap)) { + return fail("segment probe metadata does not match the loaded Qwen3.5-0.8B drafter"); + } + if (!qwen35_metadata_f32(g, "segmentprobe.threshold", st.probe_threshold) || + !qwen35_metadata_f32(g, "segmentprobe.norm_eps", st.probe_norm_eps) || + !qwen35_metadata_u32(g, "segmentprobe.min_segment", st.probe_min_segment) || + !qwen35_metadata_u32(g, "segmentprobe.max_segment", st.probe_max_segment) || + !(st.probe_threshold > 0.0f && st.probe_threshold < 1.0f) || + st.probe_min_segment < 1 || st.probe_max_segment < st.probe_min_segment) { + return fail("segment probe parameters are missing or out of range"); + } + ggml_tensor * fc1 = data_ctx ? ggml_get_tensor(data_ctx, "segmentprobe.fc1.weight") : nullptr; + if (!fc1 || fc1->type != GGML_TYPE_F32 || ggml_n_dims(fc1) != 2 || + fc1->ne[0] != w.n_embd || fc1->ne[1] < 1) { + return fail("segment probe tensor contract mismatch: segmentprobe.fc1.weight"); + } + st.probe_width = (int) fc1->ne[1]; + struct Contract { + const char * name; + int n_dims; + int64_t ne0; + int64_t ne1; + ggml_tensor ** destination; + }; + const Contract contracts[] = { + {"segmentprobe.norm.weight", 1, (int64_t) w.n_embd, 1, &st.probe_norm_w}, + {"segmentprobe.norm.bias", 1, (int64_t) w.n_embd, 1, &st.probe_norm_b}, + {"segmentprobe.fc1.weight", 2, (int64_t) w.n_embd, (int64_t) st.probe_width, &st.probe_fc1_w}, + {"segmentprobe.fc1.bias", 1, (int64_t) st.probe_width, 1, &st.probe_fc1_b}, + // ggml drops trailing unit dimensions: the [width, 1] output row is 1-D. + {"segmentprobe.fc2.weight", 1, (int64_t) st.probe_width, 1, &st.probe_fc2_w}, + {"segmentprobe.fc2.bias", 1, 1, 1, &st.probe_fc2_b}, + }; + ggml_init_params probe_params{}; + probe_params.mem_size = 8 * ggml_tensor_overhead(); + probe_params.no_alloc = true; + st.probe_ctx = ggml_init(probe_params); + if (!st.probe_ctx) return fail("segment probe context allocation failed"); + for (const auto & contract : contracts) { + ggml_tensor * source = ggml_get_tensor(data_ctx, contract.name); + if (!source || source->type != GGML_TYPE_F32 || + ggml_n_dims(source) != contract.n_dims || + source->ne[0] != contract.ne0 || + (contract.n_dims == 2 && source->ne[1] != contract.ne1)) { + return fail(std::string("segment probe tensor contract mismatch: ") + contract.name); + } + *contract.destination = contract.n_dims == 1 + ? ggml_new_tensor_1d(st.probe_ctx, GGML_TYPE_F32, contract.ne0) + : ggml_new_tensor_2d(st.probe_ctx, GGML_TYPE_F32, contract.ne0, contract.ne1); + ggml_set_name(*contract.destination, contract.name); + } + ggml_tensor * conv_w = ggml_get_tensor(data_ctx, "segmentprobe.conv.weight"); + ggml_tensor * conv_b = ggml_get_tensor(data_ctx, "segmentprobe.conv.bias"); + if (!conv_w || conv_w->type != GGML_TYPE_F32 || ggml_n_dims(conv_w) != 1 || + conv_w->ne[0] < 1 || conv_w->ne[0] % 2 != 1 || + !conv_b || conv_b->type != GGML_TYPE_F32 || ggml_n_dims(conv_b) != 1 || conv_b->ne[0] != 1) { + return fail("segment probe tensor contract mismatch: segmentprobe.conv"); + } + st.probe_conv_w.assign((const float *) conv_w->data, + (const float *) conv_w->data + conv_w->ne[0]); + st.probe_conv_b = ((const float *) conv_b->data)[0]; + // Optional sub-unit head (schema v2): scores feed only the oversize + // split rule's interior argmax. All four tensors ship together or none. + ggml_tensor * sub_src = ggml_get_tensor(data_ctx, "segmentprobe.subunit.fc2.weight"); + ggml_tensor * sub_b_src = nullptr; + if (sub_src) { + sub_b_src = ggml_get_tensor(data_ctx, "segmentprobe.subunit.fc2.bias"); + ggml_tensor * sub_cw_src = ggml_get_tensor(data_ctx, "segmentprobe.subunit.conv.weight"); + ggml_tensor * sub_cb_src = ggml_get_tensor(data_ctx, "segmentprobe.subunit.conv.bias"); + if (sub_src->type != GGML_TYPE_F32 || ggml_n_dims(sub_src) != 1 || + sub_src->ne[0] != (int64_t) st.probe_width || + !sub_b_src || sub_b_src->type != GGML_TYPE_F32 || + ggml_n_dims(sub_b_src) != 1 || sub_b_src->ne[0] != 1 || + !sub_cw_src || sub_cw_src->type != GGML_TYPE_F32 || + ggml_n_dims(sub_cw_src) != 1 || sub_cw_src->ne[0] != conv_w->ne[0] || + !sub_cb_src || sub_cb_src->type != GGML_TYPE_F32 || + ggml_n_dims(sub_cb_src) != 1 || sub_cb_src->ne[0] != 1) { + return fail("segment probe tensor contract mismatch: segmentprobe.subunit"); + } + st.probe_sub_fc2_w = ggml_new_tensor_1d(st.probe_ctx, GGML_TYPE_F32, st.probe_width); + ggml_set_name(st.probe_sub_fc2_w, "segmentprobe.subunit.fc2.weight"); + st.probe_sub_fc2_b = ggml_new_tensor_1d(st.probe_ctx, GGML_TYPE_F32, 1); + ggml_set_name(st.probe_sub_fc2_b, "segmentprobe.subunit.fc2.bias"); + st.probe_sub_conv_w.assign((const float *) sub_cw_src->data, + (const float *) sub_cw_src->data + sub_cw_src->ne[0]); + st.probe_sub_conv_b = ((const float *) sub_cb_src->data)[0]; + } + st.probe_buf = ggml_backend_alloc_ctx_tensors(st.probe_ctx, w.backend); + if (!st.probe_buf) return fail("segment probe buffer allocation failed"); + for (const auto & contract : contracts) { + ggml_tensor * source = ggml_get_tensor(data_ctx, contract.name); + ggml_backend_tensor_set(*contract.destination, source->data, 0, ggml_nbytes(source)); + } + if (st.probe_sub_fc2_w) { + ggml_backend_tensor_set(st.probe_sub_fc2_w, sub_src->data, 0, ggml_nbytes(sub_src)); + ggml_backend_tensor_set(st.probe_sub_fc2_b, sub_b_src->data, 0, ggml_nbytes(sub_b_src)); + } + gguf_free(g); + ggml_free(data_ctx); + st.probe_loaded = true; + std::fprintf(stderr, + "[qwen35-drafter] loaded segment probe: %s (width %d, threshold %.3f, " + "segments %d-%d tokens, %zu conv taps%s)\n", + path.c_str(), st.probe_width, st.probe_threshold, + st.probe_min_segment, st.probe_max_segment, st.probe_conv_w.size(), + st.probe_sub_fc2_w ? ", sub-unit head" : ""); + std::fflush(stderr); + return true; +} + +} // namespace + +bool load_qwen35_drafter(const std::string & gguf_path, DrafterArch arch, + DrafterContext & out) { + auto * st = new Qwen35DrafterState(); + // The scorer never needs logits, and tied-embedding Qwen3.5-0.8B + // exports omit output.weight, so skip the lm_head entirely. + TargetLoadPlan plan; + plan.load_output = false; + if (!load_target_gguf_partial(gguf_path, out.backend, plan, st->weights)) { + delete st; + return false; + } + const char * head_path = std::getenv("PFLASH_SCORING_HEAD_GGUF"); + const char * probe_path = std::getenv("PFLASH_SEGMENT_PROBE_GGUF"); + if (head_path || probe_path) { + const auto identity = read_gguf_metadata(gguf_path, /*compute_sha256=*/ true); + st->gguf_sha256 = identity.ok ? identity.sha256 : std::string(); + } + if (probe_path) { + if (!*probe_path || !load_qwen35_segment_probe(probe_path, *st)) { + if (!*probe_path) { + set_last_error("PFLASH_SEGMENT_PROBE_GGUF is empty"); + } + std::fprintf(stderr, + "[qwen35-drafter] ERROR: segment probe load failed, " + "refusing to serve without it\n"); + std::fflush(stderr); + free_target_weights(st->weights); + delete st; + return false; + } + } + if (head_path) { + if (!*head_path || !load_qwen35_scoring_head(head_path, *st)) { + if (!*head_path) { + set_last_error("PFLASH_SCORING_HEAD_GGUF is empty"); + } + std::fprintf(stderr, + "[qwen35-drafter] ERROR: scoring head load failed, " + "refusing to serve without it\n"); + std::fflush(stderr); + free_target_weights(st->weights); + delete st; + return false; + } + } + out.arch_state = st; + out.loaded = true; + out.arch = arch; + std::fprintf(stderr, + "[drafter] loaded %s qwen35: n_layer=%d n_head=%d n_head_kv=%d " + "n_embd=%d n_ff=%d head_dim=%d vocab=%d gpu=%d\n", + drafter_arch_name(arch), + st->weights.n_layer, st->weights.n_head, st->weights.n_head_kv, + st->weights.n_embd, st->weights.n_ff, st->weights.n_embd_head_k, + st->weights.n_vocab, out.gpu); + std::fflush(stderr); + return true; +} + +void free_qwen35_drafter_state(DrafterContext & ctx) { + auto * st = static_cast(ctx.arch_state); + free_qwen35_head(*st); + free_qwen35_segment_probe(*st); + free_target_weights(st->weights); + delete st; + ctx.arch_state = nullptr; +} + +} // namespace dflash::common diff --git a/server/src/qwen3/qwen3_backend.cpp b/server/src/qwen3/qwen3_backend.cpp index 31b9a97fd..6d28cba89 100644 --- a/server/src/qwen3/qwen3_backend.cpp +++ b/server/src/qwen3/qwen3_backend.cpp @@ -969,11 +969,10 @@ ModelBackend::CompressResult Qwen3Backend::compress(const CompressRequest & req) drafter_loaded_ = true; } - result.compressed_ids = drafter_score_and_compress( + result = CompressResult::from_compressed_ids(drafter_score_and_compress( drafter_ctx_, req.input_ids, req.keep_ratio, /*chunk_size=*/32, req.score_query_tokens, /*pool_kernel=*/13, - req.score_query_end); - result.ok = !result.compressed_ids.empty(); + req.score_query_end, req.required_instruction_spans)); if (req.residency_action == DraftResidencyAction::ReleaseAfterUse) { free_drafter(); @@ -1038,7 +1037,7 @@ bool Qwen3Backend::handle_compress(const std::string & line, const DaemonIO & io for (int32_t t : compressed) io.emit(t); io.emit(-1); - return true; + return !compressed.empty(); } void Qwen3Backend::free_drafter() { diff --git a/server/src/qwen3/qwen3_drafter.cpp b/server/src/qwen3/qwen3_drafter.cpp index 15b557a78..da83be911 100644 --- a/server/src/qwen3/qwen3_drafter.cpp +++ b/server/src/qwen3/qwen3_drafter.cpp @@ -3,7 +3,11 @@ // Wires three pieces: // - qwen3_loader.cpp : mmap GGUF + populate ggml tensors on backend // - qwen3_graph.cpp : custom forward (per-layer ggml + FP CUDA kernel) -// - chunk-top-K + span merge (this file) +// - qwen3_drafter_common.cpp : chunk-top-K + span merge, shared with the +// Qwen3.5-0.8B drafter +// +// The Qwen3.5-0.8B drafter lives in qwen35_loader.cpp / qwen35_drafter.cpp; +// this file dispatches to it on DrafterArch. // // Single-pass forward at full S using a custom Qwen3-0.6B graph with the // FlashPrefill block-sparse attention kernel (or BSA when enabled). Tail @@ -16,14 +20,19 @@ #include "qwen3_drafter.h" #include "common/dspark_head.h" #include "qwen3_drafter_model.h" +#include "qwen3_drafter_common.h" +#include "qwen35_drafter.h" +#include "pflash_selection.h" #include "qwen3/anchor_params.h" #include "common/backend_precision.h" +#include "common/gguf_inspect.h" #include "internal.h" #include "anchor_scan.h" #include "ggml.h" #include "ggml-alloc.h" #include "ggml-backend.h" +#include "gguf.h" #include #include @@ -31,6 +40,7 @@ #include #include #include +#include #include #include @@ -38,49 +48,6 @@ namespace dflash::common { namespace { -static constexpr uint16_t F16_ZERO = 0x0000; -static constexpr uint16_t F16_NEG_INF = 0xFC00; - -static int align_up_i(int x, int a) { return ((x + a - 1) / a) * a; } - -static void build_causal_mask_f16(std::vector & out, int kv_len, int n_tokens, int kv_start) { - const int kv_pad = align_up_i(kv_len, 32); - const int q_pad = align_up_i(n_tokens, 32); - out.assign((size_t)kv_pad * q_pad, F16_NEG_INF); - for (int q = 0; q < n_tokens; ++q) { - const int abs_q = kv_start + q; - for (int k = 0; k <= abs_q && k < kv_len; ++k) { - out[(size_t)q * kv_pad + k] = F16_ZERO; - } - } -} - -struct Qwen35DrafterState { - TargetWeights weights; -}; - -static int env_int(const char * name, int fallback) { - if (const char * v = std::getenv(name)) { - int x = std::atoi(v); - if (x >= 0) return x; - } - return fallback; -} - -static float env_float(const char * name, float def) { - if (const char * v = std::getenv(name)) { - try { return std::stof(v); } catch (...) {} - } - return def; -} - -static void force_chunk_neighborhood(std::vector & forced, int n_chunks, - int chunk, int radius) { - int lo = std::max(0, chunk - radius); - int hi = std::min(n_chunks - 1, chunk + radius); - for (int c = lo; c <= hi; ++c) forced[(size_t)c] = 1; -} - #if defined(DFLASH27B_BACKEND_HIP) bool prewarm_drafter_once(const Qwen3DrafterWeights & w) { static bool warmed = false; @@ -196,23 +163,7 @@ bool load_drafter(const std::string & gguf_path, int /*gpu_layers*/, } if (arch == DrafterArch::Qwen35_0p8b) { - auto * st = new Qwen35DrafterState(); - if (!load_target_gguf(gguf_path, out.backend, st->weights)) { - delete st; - return false; - } - out.arch_state = st; - out.loaded = true; - out.arch = arch; - std::fprintf(stderr, - "[drafter] loaded %s qwen35: n_layer=%d n_head=%d n_head_kv=%d " - "n_embd=%d n_ff=%d head_dim=%d vocab=%d gpu=%d\n", - drafter_arch_name(arch), - st->weights.n_layer, st->weights.n_head, st->weights.n_head_kv, - st->weights.n_embd, st->weights.n_ff, st->weights.n_embd_head_k, - st->weights.n_vocab, out.gpu); - std::fflush(stderr); - return true; + return load_qwen35_drafter(gguf_path, arch, out); } if (!load_qwen3_drafter_model(gguf_path, out.backend, out.weights)) { @@ -255,10 +206,7 @@ void free_drafter(DrafterContext & ctx) { void free_drafter_weights(DrafterContext & ctx) { if (ctx.arch == DrafterArch::Qwen35_0p8b && ctx.arch_state) { - auto * st = static_cast(ctx.arch_state); - free_target_weights(st->weights); - delete st; - ctx.arch_state = nullptr; + free_qwen35_drafter_state(ctx); } if (ctx.loaded) { if (ctx.arch == DrafterArch::Qwen3_0p6b) { @@ -268,464 +216,70 @@ void free_drafter_weights(DrafterContext & ctx) { ctx.loaded = false; } -static std::vector qwen35_score_and_compress( - TargetWeights & w, +std::vector drafter_score_and_compress( + DrafterContext & ctx, const std::vector & ids, float keep_ratio, int chunk_size, int n_lookahead, int pool_kernel, - int score_query_end) { - - const int S = (int)ids.size(); - const int hidden = w.n_embd; - if (S < n_lookahead + 1) return ids; - const int query_end = score_query_end; - if (n_lookahead < 1 || query_end < n_lookahead || query_end > S) { - set_last_error("qwen35 scorer query window out of range"); - return {}; - } - const int query_start = query_end - n_lookahead; - - auto t0 = std::chrono::steady_clock::now(); - std::vector running_max((size_t)n_lookahead * S, -INFINITY); - - TargetCache cache; -#if defined(_WIN32) - char * old_tq3_raw = nullptr; - size_t old_tq3_len = 0; - _dupenv_s(&old_tq3_raw, &old_tq3_len, "DFLASH27B_KV_TQ3"); - const bool had_old_tq3 = (old_tq3_raw != nullptr); - std::string old_tq3_s = had_old_tq3 ? old_tq3_raw : ""; - free(old_tq3_raw); - _putenv_s("DFLASH27B_KV_TQ3", "0"); - auto restore_tq3 = [&]() { - // _putenv_s with empty value removes the variable on MSVCRT. - _putenv_s("DFLASH27B_KV_TQ3", had_old_tq3 ? old_tq3_s.c_str() : ""); - }; -#else - const char * old_tq3 = std::getenv("DFLASH27B_KV_TQ3"); - std::string old_tq3_s = old_tq3 ? old_tq3 : ""; - const bool had_old_tq3 = (old_tq3 != nullptr); - setenv("DFLASH27B_KV_TQ3", "0", 1); - auto restore_tq3 = [&]() { - if (had_old_tq3) setenv("DFLASH27B_KV_TQ3", old_tq3_s.c_str(), 1); - else unsetenv("DFLASH27B_KV_TQ3"); - }; -#endif - if (!create_target_cache(w, S, 0, w.backend, cache, true)) { - restore_tq3(); + int score_query_end, + const std::vector & required_instruction_spans) { + if (!ctx.loaded) { + set_last_error("drafter not loaded"); return {}; } - restore_tq3(); - ggml_init_params act_ip{}; - act_ip.mem_size = (size_t)8 * ggml_tensor_overhead() + 4096; - act_ip.no_alloc = true; - ggml_context * act_ctx = ggml_init(act_ip); - if (!act_ctx) { - free_target_cache(cache); - set_last_error("qwen35 drafter activation ctx init failed"); + dflash::qwen3::PFlashSelectionConfig experiment; + std::string experiment_error; + if (!dflash::qwen3::resolve_pflash_selection( + (int) ids.size(), chunk_size, experiment, experiment_error)) { + set_last_error("invalid PFlash strict selection config: " + experiment_error); + std::fprintf(stderr, "[pflash-select] ERROR config: %s\n", + experiment_error.c_str()); + std::fflush(stderr); return {}; } - ggml_tensor * act_in = ggml_new_tensor_2d(act_ctx, GGML_TYPE_F32, hidden, S); - ggml_tensor * act_out = ggml_new_tensor_2d(act_ctx, GGML_TYPE_F32, hidden, S); - ggml_backend_buffer_t act_buf = ggml_backend_alloc_ctx_tensors(act_ctx, w.backend); - if (!act_buf) { - ggml_free(act_ctx); - free_target_cache(cache); - set_last_error("qwen35 drafter activation allocation failed"); + chunk_size = experiment.chunk_size; + if (!experiment.selection_active && !required_instruction_spans.empty()) { + set_last_error( + "PFlash instruction spans require strict budget selection"); + std::fprintf(stderr, + "[pflash-select] ERROR instruction spans require strict selection\n"); + std::fflush(stderr); return {}; } - - { - const int batch = 2048; - std::vector emb((size_t)hidden * batch); - for (int i = 0; i < S; i += batch) { - const int n = std::min(batch, S - i); - if (!w.embedder.embed(ids.data() + i, n, emb.data())) { - ggml_backend_buffer_free(act_buf); ggml_free(act_ctx); free_target_cache(cache); - set_last_error("qwen35 drafter embedding failed"); - return {}; - } - ggml_backend_tensor_set(act_in, emb.data(), (size_t)i * act_in->nb[1], (size_t)hidden * n * sizeof(float)); - } - } - - ggml_gallocr_t alloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(w.backend)); - const int ubatch = 1024; - for (int il = 0; il < w.n_layer; ++il) { - const bool is_attn = (((il + 1) % w.full_attention_interval) == 0); - int fa_idx = 0; - if (is_attn) { - for (int k = 0; k < il; ++k) if (((k + 1) % w.full_attention_interval) == 0) ++fa_idx; - } - for (int start = 0; start < S; start += ubatch) { - const int n = std::min(ubatch, S - start); - const int kv_len = start + n; - - ggml_init_params ip{}; - ip.mem_size = 512 * 1024 * 1024; - ip.no_alloc = true; - ggml_context * ctx = ggml_init(ip); - if (!ctx) { - ggml_gallocr_free(alloc); ggml_backend_buffer_free(act_buf); ggml_free(act_ctx); free_target_cache(cache); - set_last_error("qwen35 drafter layer graph ctx init failed"); - return {}; - } - ggml_cgraph * gf = ggml_new_graph_custom(ctx, 16384, false); - ggml_tensor * inp = ggml_view_2d(ctx, act_in, hidden, n, act_in->nb[1], (size_t)start * act_in->nb[1]); - ggml_tensor * pos = nullptr; - ggml_tensor * mask = nullptr; - if (is_attn) { - pos = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 4 * n); - ggml_set_input(pos); - mask = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, align_up_i(kv_len, 32), align_up_i(n, 32)); - ggml_set_input(mask); - } - ggml_tensor * out = build_qwen35_layer(ctx, gf, w, cache, il, inp, pos, mask, start, n, false, 0); - ggml_tensor * dst = ggml_view_2d(ctx, act_out, hidden, n, act_out->nb[1], (size_t)start * act_out->nb[1]); - if (ggml_nelements(out) != ggml_nelements(dst)) { - std::fprintf(stderr, - "[qwen35-drafter] layer output shape mismatch il=%d start=%d out=[%lld,%lld,%lld,%lld] dst=[%lld,%lld,%lld,%lld]\n", - il, start, - (long long)out->ne[0], (long long)out->ne[1], (long long)out->ne[2], (long long)out->ne[3], - (long long)dst->ne[0], (long long)dst->ne[1], (long long)dst->ne[2], (long long)dst->ne[3]); - ggml_free(ctx); ggml_gallocr_free(alloc); ggml_backend_buffer_free(act_buf); ggml_free(act_ctx); free_target_cache(cache); - set_last_error("qwen35 layer output shape mismatch"); - return {}; - } - ggml_build_forward_expand(gf, ggml_cpy(ctx, out, dst)); - if (!ggml_gallocr_alloc_graph(alloc, gf)) { - ggml_free(ctx); ggml_gallocr_free(alloc); ggml_backend_buffer_free(act_buf); ggml_free(act_ctx); free_target_cache(cache); - set_last_error("qwen35 drafter graph allocation failed"); - return {}; - } - if (is_attn) { - std::vector p4((size_t)4 * n, 0); - for (int i = 0; i < n; ++i) { - int p = start + i; - p4[(size_t)0 * n + i] = p; - p4[(size_t)1 * n + i] = p; - p4[(size_t)2 * n + i] = p; - } - ggml_backend_tensor_set(pos, p4.data(), 0, p4.size() * sizeof(int32_t)); - std::vector m; - build_causal_mask_f16(m, kv_len, n, start); - ggml_backend_tensor_set(mask, m.data(), 0, m.size() * sizeof(uint16_t)); - } - auto st = ggml_backend_graph_compute(w.backend, gf); - ggml_free(ctx); - if (st != GGML_STATUS_SUCCESS) { - ggml_gallocr_free(alloc); ggml_backend_buffer_free(act_buf); ggml_free(act_ctx); free_target_cache(cache); - set_last_error("qwen35 drafter graph compute failed"); - return {}; - } - } - - if (is_attn) { - ggml_init_params sip{}; - sip.mem_size = ggml_tensor_overhead() * 32 + ggml_graph_overhead_custom(1024, false) + 64 * 1024; - sip.no_alloc = true; - ggml_context * sctx = ggml_init(sip); - if (!sctx) { - ggml_gallocr_free(alloc); ggml_backend_buffer_free(act_buf); ggml_free(act_ctx); free_target_cache(cache); - set_last_error("qwen35 score graph ctx allocation failed"); - return {}; - } - ggml_cgraph * sgf = ggml_new_graph_custom(sctx, 1024, false); - const int K_len = (int) cache.attn_k[(size_t)fa_idx]->ne[1]; - ggml_tensor * mask_tail = ggml_new_tensor_2d(sctx, GGML_TYPE_F32, K_len, n_lookahead); - ggml_tensor * K_f32 = ggml_new_tensor_3d(sctx, GGML_TYPE_F32, w.n_embd_head_k, K_len, w.n_head_kv); - ggml_tensor * K_cast = ggml_cpy(sctx, cache.attn_k[(size_t)fa_idx], K_f32); - ggml_tensor * K_score = nullptr; - if (w.n_head != w.n_head_kv) { - const int gqa = w.n_head / w.n_head_kv; - ggml_tensor * K_4d = ggml_reshape_4d(sctx, K_cast, w.n_embd_head_k, K_len, 1, w.n_head_kv); - ggml_tensor * K_tpl = ggml_new_tensor_4d(sctx, GGML_TYPE_F32, w.n_embd_head_k, K_len, gqa, w.n_head_kv); - ggml_tensor * K_rep = ggml_repeat(sctx, K_4d, K_tpl); - K_score = ggml_reshape_3d(sctx, K_rep, w.n_embd_head_k, K_len, w.n_head); - } else { - K_score = K_cast; - } - const TargetLayer & L = w.layers[il]; - ggml_tensor * inp_tail = ggml_view_2d(sctx, act_in, hidden, n_lookahead, - act_in->nb[1], (size_t)query_start * act_in->nb[1]); - ggml_tensor * q_cur = ggml_rms_norm(sctx, inp_tail, w.rms_eps); - q_cur = ggml_mul(sctx, q_cur, L.attn_norm); - ggml_tensor * QG = ggml_mul_mat(sctx, L.wq, q_cur); - QG = ggml_reshape_3d(sctx, QG, w.n_embd_head_k * 2, w.n_head, n_lookahead); - ggml_tensor * Q = ggml_view_3d(sctx, QG, - w.n_embd_head_k, w.n_head, n_lookahead, - ggml_element_size(QG) * w.n_embd_head_k * 2, - ggml_element_size(QG) * w.n_embd_head_k * 2 * w.n_head, - 0); - Q = ggml_rms_norm(sctx, Q, w.rms_eps); - Q = ggml_mul(sctx, Q, L.q_norm); - ggml_tensor * pos_tail = ggml_new_tensor_1d(sctx, GGML_TYPE_I32, 4 * n_lookahead); - int sections[4]; - for (int k = 0; k < 4; ++k) sections[k] = w.rope_sections[k]; - Q = ggml_rope_multi(sctx, Q, pos_tail, nullptr, - w.rope_dimension_count, sections, GGML_ROPE_TYPE_MROPE, - 0, w.rope_theta, 1.0f, - 0.0f, 1.0f, 0.0f, 0.0f); - ggml_tensor * Q_tail_perm = ggml_cont(sctx, ggml_permute(sctx, Q, 0, 2, 1, 3)); - ggml_tensor * attn_score = ggml_mul_mat(sctx, K_score, Q_tail_perm); - ggml_tensor * probs = ggml_soft_max_ext(sctx, attn_score, mask_tail, 1.0f / std::sqrt((float)w.n_embd_head_k), 0.0f); - ggml_set_output(probs); - ggml_build_forward_expand(sgf, probs); - ggml_gallocr_t salloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(w.backend)); - if (!ggml_gallocr_alloc_graph(salloc, sgf)) { - ggml_gallocr_free(salloc); ggml_free(sctx); ggml_gallocr_free(alloc); ggml_backend_buffer_free(act_buf); ggml_free(act_ctx); free_target_cache(cache); - set_last_error("qwen35 score graph allocation failed"); - return {}; - } - std::vector pos4((size_t)4 * n_lookahead, 0); - for (int i = 0; i < n_lookahead; ++i) { - const int p = query_start + i; - pos4[(size_t)0 * n_lookahead + i] = p; - pos4[(size_t)1 * n_lookahead + i] = p; - pos4[(size_t)2 * n_lookahead + i] = p; - } - ggml_backend_tensor_set(pos_tail, pos4.data(), 0, pos4.size() * sizeof(int32_t)); - std::vector mask((size_t)n_lookahead * K_len, 0.0f); - for (int t = 0; t < n_lookahead; ++t) { - const int visible_end = query_start + t + 1; - for (int j = 0; j < K_len; ++j) { - mask[(size_t)t * K_len + j] = (j < visible_end) ? 0.0f : -INFINITY; - } - } - ggml_backend_tensor_set(mask_tail, mask.data(), 0, mask.size() * sizeof(float)); - auto st = ggml_backend_graph_compute(w.backend, sgf); - if (st != GGML_STATUS_SUCCESS) { - ggml_gallocr_free(salloc); ggml_free(sctx); ggml_gallocr_free(alloc); ggml_backend_buffer_free(act_buf); ggml_free(act_ctx); free_target_cache(cache); - set_last_error("qwen35 score graph compute failed"); - return {}; - } - std::vector tmp((size_t)K_len * n_lookahead * w.n_head); - ggml_backend_tensor_get(probs, tmp.data(), 0, tmp.size() * sizeof(float)); - const size_t nonfinite = - count_nonfinite_scores(tmp.data(), tmp.size()); - if (nonfinite != 0) { - const std::string message = - "non-finite Qwen3.5 PFlash scores at layer " + - std::to_string(il) + ": " + std::to_string(nonfinite) + - "/" + std::to_string(tmp.size()); - std::fprintf(stderr, "[pflash] ERROR: %s\n", message.c_str()); - std::fflush(stderr); - ggml_gallocr_free(salloc); ggml_free(sctx); - ggml_gallocr_free(alloc); ggml_backend_buffer_free(act_buf); - ggml_free(act_ctx); free_target_cache(cache); - set_last_error(message); - return {}; - } - for (int h = 0; h < w.n_head; ++h) { - for (int t = 0; t < n_lookahead; ++t) { - for (int j = 0; j < S; ++j) { - const size_t src = (size_t)h * K_len * n_lookahead + (size_t)t * K_len + j; - const size_t dst = (size_t)t * S + j; - running_max[dst] = std::max(running_max[dst], tmp[src]); - } - } - } - ggml_gallocr_free(salloc); - ggml_free(sctx); - } - std::swap(act_in, act_out); - } - ggml_gallocr_free(alloc); - ggml_backend_buffer_free(act_buf); - ggml_free(act_ctx); - free_target_cache(cache); - - std::vector score((size_t)S, 0.0f); - for (int j = 0; j < S; ++j) { - float s = 0.0f; - for (int t = 0; t < n_lookahead; ++t) s += running_max[(size_t)t * S + j]; - score[(size_t)j] = s / (float)n_lookahead; - } - - const int n_chunks = (S + chunk_size - 1) / chunk_size; - const int n_keep = std::max(1, (int)((float)n_chunks * keep_ratio)); - - std::vector smooth_score = score; - // Caller pool_kernel takes precedence; if zero/negative, fall back to env or 5. - const int pk = (pool_kernel > 0) - ? pool_kernel - : std::max(3, env_int("DFLASH_COMPRESS_POOL_KERNEL", 5)); - std::vector smoothed((size_t)S, 0.0f); - int half = pk / 2; - for (int j = 0; j < S; ++j) { - int lo = std::max(0, j - half); - int hi = std::min(S - 1, j + half); - float s = 0.0f; - int n = 0; - for (int k = lo; k <= hi; ++k) { s += score[(size_t)k]; ++n; } - smoothed[(size_t)j] = (n > 0) ? (s / (float)n) : 0.0f; - } - smooth_score.swap(smoothed); - - std::vector> chunk_means; - for (int c = 0; c < n_chunks; ++c) { - int lo = c * chunk_size, hi = std::min(S, lo + chunk_size); - float s = 0.0f; - for (int j = lo; j < hi; ++j) s += smooth_score[(size_t)j]; - chunk_means.push_back({s / std::max(1, hi - lo), c}); - } - std::sort(chunk_means.begin(), chunk_means.end(), [](auto a, auto b) { return a.first > b.first; }); - - std::vector selected((size_t)n_chunks, 0); - int count = 0; - // Scale head/tail forced chunks so they don't crowd out top-K scoring. - { - const int h_raw = env_int("DFLASH_COMPRESS_HEAD_CHUNKS", 8); - const int t_raw = env_int("DFLASH_COMPRESS_TAIL_CHUNKS", 24); - int h_n = h_raw, t_n = t_raw; - if (h_n + t_n >= n_keep) { - const int budget = std::max(1, n_keep - 1); - h_n = std::max(0, h_raw * budget / (h_raw + t_raw)); - t_n = std::max(0, budget - h_n); - } - for (int c = 0; c < std::min(n_chunks, h_n); ++c) { selected[(size_t)c] = 1; ++count; } - for (int c = std::max(0, n_chunks - t_n); c < n_chunks; ++c) if (!selected[(size_t)c]) { selected[(size_t)c] = 1; ++count; } - } - - const int query_tokens = env_int("DFLASH_COMPRESS_QUERY_TOKENS", 96); - const auto ap = resolve_anchor_params(n_chunks, - env_int("PFLASH_COMPRESS_ANCHOR_RADIUS", -1), - env_int("PFLASH_COMPRESS_MAX_ANCHOR_HITS", -1), - env_int("DFLASH_COMPRESS_ANCHOR_RADIUS", -1), - env_int("DFLASH_COMPRESS_MAX_ANCHOR_HITS", -1)); - const int anchor_radius = ap.radius; - const int max_anchor_hits = ap.max_hits; - std::vector forced((size_t)n_chunks, 0); - - const int q0 = std::max(0, S - query_tokens); - constexpr int NGRAM = 4; - for (int q = q0; q + NGRAM <= S; ++q) { - int hits = 0; - std::vector hit_pos(max_anchor_hits); - const int search_end = std::max(0, q0 - NGRAM); - for (int p = 0; p <= search_end && hits <= max_anchor_hits; ++p) { - bool same = true; - for (int k = 0; k < NGRAM; ++k) { - if (ids[(size_t)p + k] != ids[(size_t)q + k]) { same = false; break; } - } - if (same) { - if (hits < max_anchor_hits) hit_pos[hits] = p; - ++hits; - } - } - if (hits > 0 && hits <= max_anchor_hits) { - for (int i = 0; i < hits && i < max_anchor_hits; ++i) { - force_chunk_neighborhood(forced, n_chunks, hit_pos[i] / chunk_size, anchor_radius); - } - } - } - - for (int c = 0; c < n_chunks; ++c) { - if (forced[(size_t)c] && !selected[(size_t)c]) { - selected[(size_t)c] = 1; - ++count; - } - } - - // Global aggregation tasks often depend on repeated rare tokens that do - // not appear in the final query. Preserve high-frequency-but-not-filler - // token chunks before filling with model-score top-K. - const int repeat_min = env_int("DFLASH_COMPRESS_REPEAT_MIN", 4); - const int repeat_max = env_int("DFLASH_COMPRESS_REPEAT_MAX", 32); - const int repeat_limit = env_int("DFLASH_COMPRESS_REPEAT_CHUNKS", n_keep); - if (repeat_min > 1 && count < repeat_limit) { - std::unordered_map freq; - freq.reserve((size_t)S); - const int repeat_scan_end = std::max(0, S - query_tokens); - for (int j = 0; j < repeat_scan_end; ++j) { - ++freq[ids[(size_t)j]]; - } - std::vector> repeated; - repeated.reserve(freq.size()); - for (const auto & kv : freq) { - if (kv.second >= repeat_min && kv.second <= repeat_max) { - repeated.push_back({kv.second, kv.first}); - } - } - std::sort(repeated.begin(), repeated.end(), [](const auto & a, const auto & b) { - if (a.first != b.first) return a.first > b.first; - return a.second < b.second; - }); - for (const auto & rp : repeated) { - if (count >= repeat_limit) break; - const int32_t tok = rp.second; - for (int j = 0; j < repeat_scan_end && count < repeat_limit; ++j) { - if (ids[(size_t)j] != tok) continue; - const int c = j / chunk_size; - if (!selected[(size_t)c]) { - selected[(size_t)c] = 1; - ++count; - } - } - } - } - - for (auto [_, c] : chunk_means) { - if (count >= n_keep) break; - if (!selected[(size_t)c]) { selected[(size_t)c] = 1; ++count; } - } - - std::vector out_ids; - std::vector selected_chunks; - for (int c = 0; c < n_chunks; ++c) { - if (selected[(size_t)c]) selected_chunks.push_back(c); - } - int span_start = -1, span_end = -1; - for (int c : selected_chunks) { - int s_ = c * chunk_size; - int e_ = std::min(S, (c + 1) * chunk_size); - if (span_start < 0) { - span_start = s_; span_end = e_; - } else if (s_ == span_end) { - span_end = e_; - } else { - for (int j = span_start; j < span_end; ++j) out_ids.push_back(ids[j]); - span_start = s_; span_end = e_; + if (experiment.selection_active) { + std::string span_error; + if (!dflash::qwen3::validate_pflash_instruction_spans( + required_instruction_spans, (int) ids.size(), span_error)) { + set_last_error("invalid PFlash instruction spans: " + span_error); + std::fprintf(stderr, + "[pflash-select] ERROR instruction spans: %s\n", + span_error.c_str()); + std::fflush(stderr); + return {}; } } - if (span_start >= 0) { - for (int j = span_start; j < span_end; ++j) out_ids.push_back(ids[j]); - } - - auto t1 = std::chrono::steady_clock::now(); - std::fprintf(stderr, "[qwen35-drafter] forward+compress %.2fs S=%d kept=%zu (%d/%d chunks)\n", - std::chrono::duration(t1 - t0).count(), S, out_ids.size(), count, n_chunks); - std::fflush(stderr); - return out_ids; -} - -std::vector drafter_score_and_compress( - DrafterContext & ctx, - const std::vector & ids, - float keep_ratio, - int chunk_size, - int n_lookahead, - int pool_kernel, - int score_query_end) { - if (!ctx.loaded) { - set_last_error("drafter not loaded"); - return {}; + if (experiment.configured) { + std::fprintf(stderr, + "[pflash-select] config mode=%s active=%d chunk=%d " + "query_parser=%s query_cap=%d query_actual=%d top_p=%.9g " + "input=%zu\n", + dflash::qwen3::pflash_selection_mode_name(experiment.mode), + (int) experiment.selection_active, experiment.chunk_size, + dflash::qwen3::pflash_query_parser_name(experiment.query_parser), + experiment.query_tokens, n_lookahead, experiment.top_p, ids.size()); + std::fflush(stderr); } if (ctx.arch == DrafterArch::Qwen35_0p8b) { if (score_query_end < 0) { set_last_error("qwen35 scorer query window out of range"); return {}; } - if (!ctx.arch_state) { - set_last_error("qwen35 drafter state missing"); - return {}; - } - auto * st = static_cast(ctx.arch_state); - return qwen35_score_and_compress(st->weights, ids, keep_ratio, chunk_size, - n_lookahead, pool_kernel, score_query_end); + return qwen35_drafter_score_and_compress( + ctx, ids, keep_ratio, chunk_size, n_lookahead, pool_kernel, + score_query_end, experiment, required_instruction_spans); } const int S = (int)ids.size(); if (S < n_lookahead + 1) { @@ -767,6 +321,15 @@ std::vector drafter_score_and_compress( smooth[j] = (n > 0) ? (s / (float)n) : 0.0f; } + if (experiment.selection_active) { + return select_pflash_chunks( + ids, ctx.weights.scoring_head_loaded ? score : smooth, + keep_ratio, n_lookahead, score_query_end, + ctx.weights.scoring_head_loaded ? 1 : pool_kernel, + experiment, required_instruction_spans, + ctx.weights.scoring_head_loaded, true); + } + // ── 4. Chunk-top-K + span merge ─────────────────────────────────── int n_chunks = (S + chunk_size - 1) / chunk_size; int n_keep = std::max(1, (int)((float)n_chunks * keep_ratio)); @@ -881,6 +444,19 @@ std::vector drafter_score_and_compress( S, out.size(), (int)selected.size(), n_chunks, forced_count); std::fflush(stderr); + const int query_end = score_query_end < 0 ? S : score_query_end; + const int query_begin = query_end - n_lookahead; + const int token_budget = (int) std::floor( + (double) S * (double) keep_ratio); + const PFlashTraceFields trace_fields{ + &ids, query_begin, query_end, experiment.mode, + experiment.query_parser, token_budget, + dflash::qwen3::PFlashSelectionStop::InvalidInput, (int) out.size(), + 0.0, nullptr}; + write_compression_trace(S, keep_ratio, chunk_size, n_lookahead, + pool_kernel, n_keep, chunk_means, selected_mask, forced, out, + &trace_fields); + return out; } diff --git a/server/src/qwen3/qwen3_drafter.h b/server/src/qwen3/qwen3_drafter.h index 2edac6306..b4a18bfd1 100644 --- a/server/src/qwen3/qwen3_drafter.h +++ b/server/src/qwen3/qwen3_drafter.h @@ -13,6 +13,8 @@ #pragma once +#include "common/pflash_types.h" + #include #include #include @@ -82,6 +84,8 @@ std::vector drafter_score_and_compress( int chunk_size = 32, int n_lookahead = 8, int pool_kernel = 13, - int score_query_end = -1); + int score_query_end = -1, + const std::vector & + required_instruction_spans = {}); } // namespace dflash::common diff --git a/server/src/qwen3/qwen3_drafter_common.cpp b/server/src/qwen3/qwen3_drafter_common.cpp new file mode 100644 index 000000000..769599b73 --- /dev/null +++ b/server/src/qwen3/qwen3_drafter_common.cpp @@ -0,0 +1,326 @@ +// Helpers shared by the Qwen3-0.6B and Qwen3.5-0.8B drafter paths. +// See qwen3_drafter_common.h. + +#include "qwen3_drafter_common.h" + +#include "qwen3_drafter.h" +#include "pflash_selection.h" +#include "internal.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace dflash::common { + +int env_int(const char * name, int fallback) { + if (const char * v = std::getenv(name)) { + int x = std::atoi(v); + if (x >= 0) return x; + } + return fallback; +} + +float env_float(const char * name, float def) { + if (const char * v = std::getenv(name)) { + try { return std::stof(v); } catch (...) {} + } + return def; +} + +void force_chunk_neighborhood(std::vector & forced, int n_chunks, + int chunk, int radius) { + int lo = std::max(0, chunk - radius); + int hi = std::min(n_chunks - 1, chunk + radius); + for (int c = lo; c <= hi; ++c) forced[(size_t)c] = 1; +} + +void write_compression_trace( + int input_tokens, + float keep_ratio, + int chunk_size, + int n_lookahead, + int pool_kernel, + int n_keep, + const std::vector> & chunk_means, + const std::vector & selected, + const std::vector & forced, + const std::vector & compressed_ids, + const PFlashTraceFields * trace_fields) { + const char * path = std::getenv("DFLASH_PFLASH_TRACE_PATH"); + if (!path || !*path) return; + + FILE * file = std::fopen(path, "a"); + if (!file) { + std::fprintf(stderr, "[pflash-trace] cannot append %s\n", path); + return; + } + + std::vector scores(selected.size(), 0.0f); + for (const auto & chunk : chunk_means) { + scores[(size_t)chunk.second] = chunk.first; + } + const bool has_exact_scores = trace_fields && + trace_fields->exact_chunk_scores && + trace_fields->exact_chunk_scores->size() == scores.size(); + if (trace_fields && + trace_fields->selector_mode != + dflash::qwen3::PFlashSelectionMode::Legacy && + !has_exact_scores) { + std::fclose(file); + std::fprintf(stderr, "[pflash-trace] exact strict scores unavailable\n"); + return; + } + + std::fprintf(file, + "{\"schema_version\":%d,\"input_tokens\":%d,\"keep_ratio\":%.9g", + trace_fields ? 3 : 1, input_tokens, keep_ratio); + if (trace_fields) { + std::fputs(",\"input_ids\":[", file); + for (size_t index = 0; index < trace_fields->input_ids->size(); ++index) { + if (index) std::fputc(',', file); + std::fprintf(file, "%d", (*trace_fields->input_ids)[index]); + } + std::fprintf(file, + "],\"query_begin\":%d,\"query_end\":%d," + "\"selector_mode\":\"%s\",\"query_parser\":\"%s\"," + "\"token_budget\":%d," + "\"retained_tokens\":%d", + trace_fields->query_begin, trace_fields->query_end, + dflash::qwen3::pflash_selection_mode_name( + trace_fields->selector_mode), + dflash::qwen3::pflash_query_parser_name(trace_fields->query_parser), + trace_fields->token_budget, trace_fields->retained_tokens); + std::fputs(",\"required_instruction_spans\":[", file); + if (trace_fields->required_instruction_spans) { + for (size_t index = 0; + index < trace_fields->required_instruction_spans->size(); + ++index) { + if (index) std::fputc(',', file); + const auto & span = + (*trace_fields->required_instruction_spans)[index]; + std::fprintf(file, "[%d,%d]", span.begin, span.end); + } + } + std::fputc(']', file); + if (trace_fields->selector_mode == + dflash::qwen3::PFlashSelectionMode::Legacy) { + std::fputs(",\"stop_reason\":null,\"retained_mass\":null", file); + } else { + std::fprintf(file, + ",\"stop_reason\":\"%s\",\"retained_mass\":%.17g", + dflash::qwen3::pflash_selection_stop_name(trace_fields->stop), + trace_fields->retained_mass); + } + } + if (trace_fields) { + std::fprintf(file, ",\"segmentation\":\"%s\",\"candidate_score\":\"%s\",\"scorer\":\"%s\",\"split_fraction\":%.4f", + trace_fields->segmentation, trace_fields->candidate_score, + trace_fields->scorer, trace_fields->split_fraction); + if (trace_fields->other_chunk_scores) { + std::fputs(",\"other_chunk_scores\":[", file); + for (size_t index = 0; index < trace_fields->other_chunk_scores->size(); ++index) { + const double score = (*trace_fields->other_chunk_scores)[index]; + if (index) std::fputc(',', file); + if (std::isfinite(score)) std::fprintf(file, "%.9g", score); else std::fputs("null", file); + } + std::fputc(']', file); + } + if (trace_fields->segments) { + std::fputs(",\"segments\":[", file); + for (size_t index = 0; index < trace_fields->segments->size(); ++index) { + const auto & span = (*trace_fields->segments)[index]; + std::fprintf(file, "%s[%d,%d]", index ? "," : "", span.begin, span.end); + } + std::fputc(']', file); + } + } + std::fprintf(file, + ",\"chunk_size\":%d,\"n_lookahead\":%d,\"pool_kernel\":%d," + "\"n_keep\":%d,\"chunk_scores\":[", + chunk_size, n_lookahead, pool_kernel, n_keep); + for (size_t index = 0; index < scores.size(); ++index) { + if (index) std::fputc(',', file); + const double score = has_exact_scores + ? (*trace_fields->exact_chunk_scores)[index] + : (double) scores[index]; + if (std::isfinite(score)) { + std::fprintf(file, has_exact_scores ? "%.17g" : "%.9g", score); + } else { + std::fputs("null", file); + } + } + std::fputs("],\"selected_chunks\":[", file); + bool first = true; + for (size_t index = 0; index < selected.size(); ++index) { + if (!selected[index]) continue; + if (!first) std::fputc(',', file); + std::fprintf(file, "%zu", index); + first = false; + } + std::fputs("],\"forced_chunks\":[", file); + first = true; + for (size_t index = 0; index < forced.size(); ++index) { + if (!forced[index]) continue; + if (!first) std::fputc(',', file); + std::fprintf(file, "%zu", index); + first = false; + } + std::fputs("],\"compressed_ids\":[", file); + for (size_t index = 0; index < compressed_ids.size(); ++index) { + if (index) std::fputc(',', file); + std::fprintf(file, "%d", compressed_ids[index]); + } + std::fputs("]}\n", file); + std::fclose(file); +} + +std::vector select_pflash_chunks( + const std::vector & ids, + const std::vector & token_scores, + float keep_ratio, + int n_lookahead, + int score_query_end, + int pool_kernel, + const dflash::qwen3::PFlashSelectionConfig & config, + const std::vector & required_instruction_spans, + bool direct_mass, + bool write_trace, + const std::vector * segments, + bool density, + const std::vector * other_token_scores, + double split_fraction) { + const int input_tokens = (int) ids.size(); + const int query_end = score_query_end < 0 ? input_tokens : score_query_end; + const int query_tokens = std::min(n_lookahead, query_end); + const int query_begin = query_end - query_tokens; + const int selector_budget = (int) std::floor( + (double) input_tokens * (double) keep_ratio); + // Fixed grid unless the caller provides variable-length segments. + const int n_chunks = segments + ? (int) segments->size() + : (input_tokens + config.chunk_size - 1) / config.chunk_size; + + std::vector candidates; + std::vector> chunk_means; + std::vector exact_chunk_scores; + candidates.reserve((size_t) n_chunks); + chunk_means.reserve((size_t) n_chunks); + exact_chunk_scores.reserve((size_t) n_chunks); + for (int chunk = 0; chunk < n_chunks; ++chunk) { + const int begin = segments ? (*segments)[(size_t) chunk].begin : chunk * config.chunk_size; + const int end = segments ? (*segments)[(size_t) chunk].end + : std::min(input_tokens, begin + config.chunk_size); + double score = 0.0; + for (int token = begin; token < end; ++token) { + score += token_scores[(size_t) token]; + } + if (!direct_mass || density) { + score /= (double) std::max(1, end - begin); + } + const bool mandatory = + dflash::qwen3::pflash_chunk_is_structurally_required( + begin, end, query_begin, query_end, input_tokens, + required_instruction_spans); + candidates.push_back({(size_t) chunk, begin, end, score, mandatory}); + chunk_means.push_back({(float) score, chunk}); + exact_chunk_scores.push_back(score); + } + // Two-scorer selection: the other scorer's mean per-token score over the + // same spans (its native ranking rule). + std::vector other_candidates; + std::vector other_scores; + const bool split = other_token_scores != nullptr && split_fraction > 0.0; + if (split) { + for (const auto & candidate : candidates) { + double score = 0.0; + for (int token = candidate.begin; token < candidate.end; ++token) { + score += (*other_token_scores)[(size_t) token]; + } + score /= (double) std::max(1, candidate.end - candidate.begin); + other_candidates.push_back({candidate.ordinal, candidate.begin, candidate.end, score, candidate.mandatory}); + other_scores.push_back(score); + } + } + + const dflash::qwen3::PFlashSelectionPolicy policy{selector_budget, config.top_p, + /*skip_oversized=*/ segments != nullptr}; + const auto selected = split + ? dflash::qwen3::select_pflash_split(candidates, other_candidates, policy, split_fraction, config.mode) + : dflash::qwen3::select_pflash_candidates(candidates, policy, config.mode); + if (!selected.ok) { + set_last_error("PFlash selection failed: " + selected.error); + std::fprintf(stderr, + "[pflash-select] ERROR mode=%s budget=%d stop=%s: %s\n", + dflash::qwen3::pflash_selection_mode_name(config.mode), + selector_budget, + dflash::qwen3::pflash_selection_stop_name(selected.stop), + selected.error.c_str()); + std::fflush(stderr); + return {}; + } + + std::vector selected_mask((size_t) n_chunks, 0); + std::vector mandatory_mask((size_t) n_chunks, 0); + for (const auto & candidate : candidates) { + if (candidate.mandatory) mandatory_mask[candidate.ordinal] = 1; + } + for (size_t ordinal : selected.ordinals) { + if (ordinal >= selected_mask.size()) { + set_last_error("PFlash selector returned an invalid ordinal"); + return {}; + } + selected_mask[ordinal] = 1; + } + + std::vector output; + output.reserve((size_t) selected.retained_tokens); + for (const auto & candidate : candidates) { + if (!selected_mask[candidate.ordinal]) continue; + output.insert(output.end(), + ids.begin() + candidate.begin, + ids.begin() + candidate.end); + } + + std::fprintf(stderr, + "[pflash-select] selected mode=%s scorer=%s segments=%s score=%s chunk=%d query=%d " + "budget=%d selected_tokens=%zu chunks=%zu/%d stop=%s mass=%.9g\n", + dflash::qwen3::pflash_selection_mode_name(config.mode), + split ? "split" : "single", + segments ? "probe" : "fixed", density ? "density" : "sum", + segments ? 0 : config.chunk_size, query_tokens, selector_budget, output.size(), + selected.ordinals.size(), n_chunks, + dflash::qwen3::pflash_selection_stop_name(selected.stop), + selected.retained_mass); + std::fflush(stderr); + + if (write_trace) { + const int trace_chunk = segments ? 0 : config.chunk_size; + const int n_keep_approx = segments + ? (int) selected.ordinals.size() + : std::max(1, (selector_budget + config.chunk_size - 1) / config.chunk_size); + PFlashTraceFields strict_fields{ + &ids, query_begin, query_end, config.mode, config.query_parser, + selector_budget, + selected.stop, selected.retained_tokens, selected.retained_mass, + &exact_chunk_scores, &required_instruction_spans}; + strict_fields.segments = segments; + strict_fields.segmentation = segments ? "probe" : "fixed"; + strict_fields.candidate_score = density ? "density" : "sum"; + strict_fields.scorer = split ? "split" : dflash::qwen3::pflash_scorer_name(config.scorer); + strict_fields.split_fraction = split ? split_fraction : 0.0; + strict_fields.other_chunk_scores = split ? &other_scores : nullptr; + write_compression_trace( + input_tokens, keep_ratio, trace_chunk, query_tokens, + pool_kernel, n_keep_approx, chunk_means, selected_mask, + mandatory_mask, output, &strict_fields); + } + return output; +} + +} // namespace dflash::common diff --git a/server/src/qwen3/qwen3_drafter_common.h b/server/src/qwen3/qwen3_drafter_common.h new file mode 100644 index 000000000..743c8a3fc --- /dev/null +++ b/server/src/qwen3/qwen3_drafter_common.h @@ -0,0 +1,77 @@ +// Helpers shared by the Qwen3-0.6B and Qwen3.5-0.8B drafter paths. +// +// Moved verbatim out of qwen3_drafter.cpp so qwen35_drafter.cpp can use the +// same selector, trace writer and environment readers without a second copy. + +#pragma once + +#include "pflash_selection.h" +#include "common/pflash_types.h" + +#include +#include +#include +#include + +namespace dflash::common { + +int env_int(const char * name, int fallback); +float env_float(const char * name, float def); +void force_chunk_neighborhood(std::vector & forced, int n_chunks, + int chunk, int radius); + +struct PFlashTraceFields { + const std::vector * input_ids = nullptr; + int query_begin = -1; + int query_end = -1; + dflash::qwen3::PFlashSelectionMode selector_mode = + dflash::qwen3::PFlashSelectionMode::Legacy; + dflash::qwen3::PFlashQueryParser query_parser = + dflash::qwen3::PFlashQueryParser::SemanticUser; + int token_budget = 0; + dflash::qwen3::PFlashSelectionStop stop = + dflash::qwen3::PFlashSelectionStop::InvalidInput; + int retained_tokens = 0; + double retained_mass = 0.0; + const std::vector * exact_chunk_scores = nullptr; + const std::vector * required_instruction_spans = nullptr; + // Variable-length candidates (segment probe): spans in candidate order. + const std::vector * segments = nullptr; + const char * segmentation = "fixed"; + const char * candidate_score = "sum"; + // Two-scorer selection: the other scorer's candidate scores, same order. + const char * scorer = "head"; + double split_fraction = 0.0; + const std::vector * other_chunk_scores = nullptr; +}; + +void write_compression_trace( + int input_tokens, + float keep_ratio, + int chunk_size, + int n_lookahead, + int pool_kernel, + int n_keep, + const std::vector> & chunk_means, + const std::vector & selected, + const std::vector & forced, + const std::vector & compressed_ids, + const PFlashTraceFields * trace_fields = nullptr); + +std::vector select_pflash_chunks( + const std::vector & ids, + const std::vector & token_scores, + float keep_ratio, + int n_lookahead, + int score_query_end, + int pool_kernel, + const dflash::qwen3::PFlashSelectionConfig & config, + const std::vector & required_instruction_spans, + bool direct_mass, + bool write_trace, + const std::vector * segments = nullptr, + bool density = false, + const std::vector * other_token_scores = nullptr, + double split_fraction = 0.0); + +} // namespace dflash::common diff --git a/server/src/qwen3/qwen3_drafter_model.h b/server/src/qwen3/qwen3_drafter_model.h index c3f7c235d..d7063ec36 100644 --- a/server/src/qwen3/qwen3_drafter_model.h +++ b/server/src/qwen3/qwen3_drafter_model.h @@ -65,6 +65,7 @@ struct Qwen3DrafterWeights { int n_vocab = 151936; int n_ctx_max = 40960; float rope_theta = 1000000.0f; + bool scoring_head_loaded = false; }; bool load_qwen3_drafter_model(const std::string & gguf_path, @@ -126,4 +127,25 @@ inline size_t count_nonfinite_scores(const float * values, size_t count) { return nonfinite; } +// Scoring-head token mass: mean over heads and query tokens of softmax +// probabilities laid out as ggml [n_keys, n_queries, n_heads] (ne0 fastest). +inline void scoring_head_mean_token_mass( + const float * probs, + int n_keys, + int n_queries, + int n_heads, + std::vector & out) { + out.assign((size_t) n_keys, 0.0f); + if (n_keys <= 0 || n_queries <= 0 || n_heads <= 0) return; + std::vector sum((size_t) n_keys, 0.0); + for (int h = 0; h < n_heads; ++h) { + for (int t = 0; t < n_queries; ++t) { + const float * row = probs + ((size_t) h * n_queries + t) * n_keys; + for (int j = 0; j < n_keys; ++j) sum[(size_t) j] += row[j]; + } + } + const double denominator = (double) n_heads * (double) n_queries; + for (int j = 0; j < n_keys; ++j) out[(size_t) j] = (float) (sum[(size_t) j] / denominator); +} + } // namespace dflash::common diff --git a/server/src/qwen3/qwen3_graph.cpp b/server/src/qwen3/qwen3_graph.cpp index 3e40b08f5..07066d952 100644 --- a/server/src/qwen3/qwen3_graph.cpp +++ b/server/src/qwen3/qwen3_graph.cpp @@ -178,10 +178,12 @@ bool build_hip_chunk_graph_b(const Qwen3DrafterLayer & L, return true; } -void warm_hip_chunk_graph_b_once(ggml_backend_t backend, HipChunkGraphB & out) { +bool warm_hip_chunk_graph_b_once(ggml_backend_t backend, + HipChunkGraphB & out, + std::string & error) { static bool warmed = false; if (warmed) { - return; + return true; } struct ggml_tensor * warm_tensors[] = { @@ -190,13 +192,27 @@ void warm_hip_chunk_graph_b_once(ggml_backend_t backend, HipChunkGraphB & out) { for (ggml_tensor * t : warm_tensors) { cudaError_t e = cudaMemset(t->data, 0, ggml_nbytes(t)); if (e != cudaSuccess) { - return; + error = std::string("memset failed: ") + cudaGetErrorString(e); + return false; } } - ggml_backend_graph_compute(backend, out.gf_proj_add); - ggml_backend_graph_compute(backend, out.gf_ffn); + const ggml_status proj_status = + ggml_backend_graph_compute(backend, out.gf_proj_add); + if (proj_status != GGML_STATUS_SUCCESS) { + error = std::string("projection graph failed: ") + + ggml_status_to_string(proj_status); + return false; + } + const ggml_status ffn_status = + ggml_backend_graph_compute(backend, out.gf_ffn); + if (ffn_status != GGML_STATUS_SUCCESS) { + error = std::string("FFN graph failed: ") + + ggml_status_to_string(ffn_status); + return false; + } warmed = true; + return true; } #endif @@ -246,10 +262,11 @@ bool forward_qwen3_drafter_model( const float rope_b = w.rope_theta; // Pre-RoPE tail scoring: removes RoPE distance decay from the score signal. // Default ON; set DFLASH_FP_NOPE_TAIL=0 to disable (saves ~K_curr_v memory). - static const bool nope_tail = []() -> bool { + static const bool configured_nope_tail = []() -> bool { const char * e = std::getenv("DFLASH_FP_NOPE_TAIL"); return e == nullptr || std::string(e) != "0"; }(); + const bool nope_tail = w.scoring_head_loaded || configured_nope_tail; if (n_lookahead < 1 || S < n_lookahead + 1) { set_last_error("forward_qwen3_drafter_model: S too small"); @@ -275,9 +292,13 @@ bool forward_qwen3_drafter_model( if (e) { int v = std::atoi(e); if (v > 0) return v; } return -1; }(); - const int fwd_layer_limit_pre = (early_exit_pre > 0 && early_exit_pre < w.n_layer) - ? early_exit_pre : w.n_layer; - const ScoreRange pre_range = compute_score_range(w.n_layer, score_layers_pre, fwd_layer_limit_pre); + const int fwd_layer_limit_pre = w.scoring_head_loaded + ? 14 + : ((early_exit_pre > 0 && early_exit_pre < w.n_layer) + ? early_exit_pre : w.n_layer); + const ScoreRange pre_range = w.scoring_head_loaded + ? ScoreRange{13, 14} + : compute_score_range(w.n_layer, score_layers_pre, fwd_layer_limit_pre); const int score_layer_start_pre = pre_range.start; const int n_score_layers = pre_range.count(); // K_norope/Q_norope sized to this, not n_layer @@ -361,7 +382,9 @@ bool forward_qwen3_drafter_model( { std::vector m((size_t)n_lookahead * S, 0.0f); for (int t = 0; t < n_lookahead; ++t) { - const int visible_end = query_start + t + 1; + const int visible_end = w.scoring_head_loaded + ? query_start + : query_start + t + 1; for (int j = 0; j < S; ++j) { m[(size_t)t * S + j] = (j < visible_end) ? 0.0f : -INFINITY; } @@ -393,14 +416,22 @@ bool forward_qwen3_drafter_model( return false; } ggml_backend_tensor_set(t_ids, ids.data(), 0, (size_t)S * sizeof(int32_t)); - ggml_backend_graph_compute(w.backend, gf); + const ggml_status embed_status = + ggml_backend_graph_compute(w.backend, gf); + if (embed_status != GGML_STATUS_SUCCESS) { + set_last_error(std::string("embed graph compute failed: ") + + ggml_status_to_string(embed_status)); + ggml_gallocr_free(galloc); + if (in_buf) ggml_backend_buffer_free(in_buf); + ggml_free(gctx); + cleanup_all(); + return false; + } ggml_gallocr_free(galloc); if (in_buf) ggml_backend_buffer_free(in_buf); ggml_free(gctx); } - const int & early_exit_n = early_exit_pre; // alias for readability in loop below - // Per-layer A→FA→B loop. ggml_gallocr_t galloc = ggml_gallocr_new( ggml_backend_get_default_buffer_type(w.backend)); @@ -421,8 +452,7 @@ bool forward_qwen3_drafter_model( double t_b_warm = 0.0, t_b_setup = 0.0, t_b_alloc = 0.0, t_b_copy_in = 0.0, t_b_norm = 0.0, t_compute_b = 0.0, t_b_copy_out = 0.0; double t_fp = 0.0; - const int fwd_layer_limit = (early_exit_n > 0 && early_exit_n < w.n_layer) - ? early_exit_n : w.n_layer; + const int fwd_layer_limit = fwd_layer_limit_pre; for (int il = 0; il < fwd_layer_limit; ++il) { const auto & L = w.layers[il]; @@ -558,10 +588,19 @@ bool forward_qwen3_drafter_model( auto tA_alloc1 = std::chrono::steady_clock::now(); t_a_alloc += std::chrono::duration(tA_alloc1 - tA_alloc0).count(); auto tA0 = std::chrono::steady_clock::now(); - ggml_backend_graph_compute(w.backend, gfA); + const ggml_status graph_a_status = + ggml_backend_graph_compute(w.backend, gfA); ggml_backend_synchronize(w.backend); auto tA1 = std::chrono::steady_clock::now(); t_compute_a += std::chrono::duration(tA1 - tA0).count(); + if (graph_a_status != GGML_STATUS_SUCCESS) { + set_last_error(std::string("graph A compute failed at layer ") + + std::to_string(il) + " chunk " + + std::to_string(cs) + ": " + + ggml_status_to_string(graph_a_status)); + ggml_free(gA); + ggml_gallocr_free(galloc); cleanup_all(); return false; + } if (debug_first_layer) { std::fprintf(stderr, "[qwen3-0.6b-fp dbg] layer0 chunk A done setup=%.3fs alloc=%.3fs compute=%.3fs\n", @@ -573,6 +612,10 @@ bool forward_qwen3_drafter_model( ggml_free(gA); } + if (w.scoring_head_loaded && il == 13) { + continue; + } + // ── Attention dispatch ── auto tF0 = std::chrono::steady_clock::now(); int rc = flashprefill::flash_prefill_forward( @@ -588,7 +631,20 @@ bool forward_qwen3_drafter_model( set_last_error("flash_prefill_forward failed at layer " + std::to_string(il)); ggml_gallocr_free(galloc); cleanup_all(); return false; } - cudaDeviceSynchronize(); + cudaError_t fp_launch_e = cudaGetLastError(); + if (fp_launch_e != cudaSuccess) { + set_last_error(std::string("flash_prefill launch failed at layer ") + + std::to_string(il) + ": " + + cudaGetErrorString(fp_launch_e)); + ggml_gallocr_free(galloc); cleanup_all(); return false; + } + cudaError_t fp_sync_e = cudaDeviceSynchronize(); + if (fp_sync_e != cudaSuccess) { + set_last_error(std::string("flash_prefill synchronization failed at layer ") + + std::to_string(il) + ": " + + cudaGetErrorString(fp_sync_e)); + ggml_gallocr_free(galloc); cleanup_all(); return false; + } auto tF1 = std::chrono::steady_clock::now(); t_fp += std::chrono::duration(tF1 - tF0).count(); if (debug_first_layer) { @@ -615,7 +671,13 @@ bool forward_qwen3_drafter_model( } auto tB_warm0 = std::chrono::steady_clock::now(); - warm_hip_chunk_graph_b_once(w.backend, gb); + std::string warm_error; + if (!warm_hip_chunk_graph_b_once(w.backend, gb, warm_error)) { + set_last_error(std::string("graph B warmup failed at layer ") + + std::to_string(il) + ": " + warm_error); + free_hip_chunk_graph_b(gb); + ggml_gallocr_free(galloc); cleanup_all(); return false; + } auto tB_warm1 = std::chrono::steady_clock::now(); t_b_warm += std::chrono::duration(tB_warm1 - tB_warm0).count(); if (debug_first_layer) { @@ -652,6 +714,16 @@ bool forward_qwen3_drafter_model( free_hip_chunk_graph_b(gb); ggml_gallocr_free(galloc); cleanup_all(); return false; } + // HIP D2D hipMemcpy can return before its null-stream copy finishes. + // GGML uses a nonblocking stream, so make the copy-in dependency explicit. + cudaError_t copy_in_sync_e = cudaStreamSynchronize(nullptr); + if (copy_in_sync_e != cudaSuccess) { + set_last_error(std::string("graph B copy-in synchronization failed at layer ") + + std::to_string(il) + " chunk " + std::to_string(cs) + ": " + + cudaGetErrorString(copy_in_sync_e)); + free_hip_chunk_graph_b(gb); + ggml_gallocr_free(galloc); cleanup_all(); return false; + } auto tB_copy_in1 = std::chrono::steady_clock::now(); t_b_copy_in += std::chrono::duration(tB_copy_in1 - tB_copy_in0).count(); if (debug_first_layer) { @@ -689,11 +761,21 @@ bool forward_qwen3_drafter_model( double proj_s = 0, ffn_s = 0; auto one = [&](ggml_cgraph * gf, double & acc) { auto ts0 = std::chrono::steady_clock::now(); - ggml_backend_graph_compute(w.backend, gf); + const ggml_status status = + ggml_backend_graph_compute(w.backend, gf); auto ts1 = std::chrono::steady_clock::now(); acc = std::chrono::duration(ts1 - ts0).count(); + return status; }; - one(gb.gf_proj_add, proj_s); + const ggml_status proj_status = one(gb.gf_proj_add, proj_s); + if (proj_status != GGML_STATUS_SUCCESS) { + set_last_error(std::string("graph B projection compute failed at layer ") + + std::to_string(il) + " chunk " + + std::to_string(cs) + ": " + + ggml_status_to_string(proj_status)); + free_hip_chunk_graph_b(gb); + ggml_gallocr_free(galloc); cleanup_all(); return false; + } auto tB_norm0 = std::chrono::steady_clock::now(); launch_rms_norm_mul_w_f32( @@ -702,11 +784,36 @@ bool forward_qwen3_drafter_model( (float *)gb.hf->data, cl, hidden, eps, /*stream=*/nullptr); - cudaDeviceSynchronize(); + cudaError_t rms_launch_e = cudaGetLastError(); + if (rms_launch_e != cudaSuccess) { + set_last_error(std::string("graph B RMSNorm launch failed at layer ") + + std::to_string(il) + " chunk " + + std::to_string(cs) + ": " + + cudaGetErrorString(rms_launch_e)); + free_hip_chunk_graph_b(gb); + ggml_gallocr_free(galloc); cleanup_all(); return false; + } + cudaError_t rms_sync_e = cudaDeviceSynchronize(); + if (rms_sync_e != cudaSuccess) { + set_last_error(std::string("graph B RMSNorm synchronization failed at layer ") + + std::to_string(il) + " chunk " + + std::to_string(cs) + ": " + + cudaGetErrorString(rms_sync_e)); + free_hip_chunk_graph_b(gb); + ggml_gallocr_free(galloc); cleanup_all(); return false; + } auto tB_norm1 = std::chrono::steady_clock::now(); t_b_norm += std::chrono::duration(tB_norm1 - tB_norm0).count(); - one(gb.gf_ffn, ffn_s); + const ggml_status ffn_status = one(gb.gf_ffn, ffn_s); + if (ffn_status != GGML_STATUS_SUCCESS) { + set_last_error(std::string("graph B FFN compute failed at layer ") + + std::to_string(il) + " chunk " + + std::to_string(cs) + ": " + + ggml_status_to_string(ffn_status)); + free_hip_chunk_graph_b(gb); + ggml_gallocr_free(galloc); cleanup_all(); return false; + } auto tB1 = std::chrono::steady_clock::now(); t_compute_b += std::chrono::duration(tB1 - tB0).count(); @@ -877,15 +984,25 @@ bool forward_qwen3_drafter_model( for (int t = 0; t < n_lookahead; ++t) { for (int j = 0; j < S; ++j) { - float m = -INFINITY; - for (int h = 0; h < H; ++h) { - float v = probs_h[(size_t)j - + (size_t)t * S - + (size_t)h * S * n_lookahead]; - if (v > m) m = v; - } size_t idx = (size_t)t * S + j; - if (m > running_max[idx]) running_max[idx] = m; + if (w.scoring_head_loaded) { + float sum = 0.0f; + for (int h = 0; h < H; ++h) { + sum += probs_h[(size_t)j + + (size_t)t * S + + (size_t)h * S * n_lookahead]; + } + running_max[idx] = sum / (float)H; + } else { + float m = -INFINITY; + for (int h = 0; h < H; ++h) { + float v = probs_h[(size_t)j + + (size_t)t * S + + (size_t)h * S * n_lookahead]; + if (v > m) m = v; + } + if (m > running_max[idx]) running_max[idx] = m; + } } } } diff --git a/server/src/qwen3/qwen3_loader.cpp b/server/src/qwen3/qwen3_loader.cpp index 583261992..af08b8133 100644 --- a/server/src/qwen3/qwen3_loader.cpp +++ b/server/src/qwen3/qwen3_loader.cpp @@ -24,6 +24,7 @@ #include "qwen3_drafter_model.h" #include "common/backend_precision.h" +#include "common/gguf_inspect.h" #include "common/gguf_mmap.h" #include "internal.h" @@ -91,6 +92,20 @@ bool copy_tensor_from_file(gguf_context * gctx, const char * name, return true; } + if (src_type == GGML_TYPE_F32 && dst_type == GGML_TYPE_BF16) { + std::vector tmp_bf16((size_t)n); + ggml_fp32_to_bf16_row((const float *)src, tmp_bf16.data(), n); + ggml_backend_tensor_set(dst, tmp_bf16.data(), 0, ggml_nbytes(dst)); + return true; + } + + if (src_type == GGML_TYPE_F32 && dst_type == GGML_TYPE_F16) { + std::vector tmp_f16((size_t)n); + ggml_fp32_to_fp16_row((const float *)src, tmp_f16.data(), n); + ggml_backend_tensor_set(dst, tmp_f16.data(), 0, ggml_nbytes(dst)); + return true; + } + std::fprintf(stderr, "[qwen3-0.6b] unsupported tensor conversion for %s: %s -> %s\n", name, ggml_type_name(src_type), ggml_type_name(dst_type)); return false; @@ -108,6 +123,91 @@ float get_f32(gguf_context * g, const char * key, float def) { return gguf_get_val_f32(g, k); } +bool metadata_equals(gguf_context * g, const char * key, const char * expected) { + const int id = gguf_find_key(g, key); + return id >= 0 && std::string(gguf_get_val_str(g, id)) == expected; +} + +bool load_scoring_head( + const std::string & path, + const std::string & drafter_sha256, + Qwen3DrafterWeights & out) { + ggml_context * tensor_ctx = nullptr; + gguf_init_params iparams{ /*no_alloc=*/ true, /*ctx=*/ &tensor_ctx }; + gguf_context * gctx = gguf_init_from_file(path.c_str(), iparams); + if (!gctx) { + set_last_error("scoring head GGUF could not be opened: " + path); + return false; + } + auto fail = [&](const std::string & message) { + gguf_free(gctx); + if (tensor_ctx) ggml_free(tensor_ctx); + set_last_error(message); + return false; + }; + // GGUF contract of a scoring-head file: architecture `pflash_scoring_head`, + // metadata and tensors under `scoringhead.*`. + const bool metadata_ok = + metadata_equals(gctx, "general.architecture", "pflash_scoring_head") && + metadata_equals(gctx, "scoringhead.schema", "qwen3_0_6b_nope_qk_mass_v1") && + metadata_equals(gctx, "scoringhead.base_model", "Qwen/Qwen3-0.6B") && + metadata_equals( + gctx, + "scoringhead.runtime_gguf_sha256", + drafter_sha256.c_str()) && + metadata_equals( + gctx, + "scoringhead.feature_tap", + "post_block12_residual_before_block13"); + if (!metadata_ok) { + return fail("scoring head metadata does not match the loaded Qwen3-0.6B drafter"); + } + struct TensorContract { + const char * name; + ggml_tensor * destination; + }; + const TensorContract contracts[] = { + {"scoringhead.attn_q.weight", out.layers[13].wq}, + {"scoringhead.attn_k.weight", out.layers[13].wk}, + }; + for (const auto & contract : contracts) { + const int64_t id = gguf_find_tensor(gctx, contract.name); + ggml_tensor * source = tensor_ctx + ? ggml_get_tensor(tensor_ctx, contract.name) + : nullptr; + if (id < 0 || gguf_get_tensor_type(gctx, id) != GGML_TYPE_F32 || + !source || !ggml_are_same_shape(source, contract.destination) || + gguf_get_tensor_size(gctx, id) != + (size_t)ggml_nelements(contract.destination) * sizeof(float)) { + return fail(std::string("scoring head tensor contract mismatch: ") + + contract.name); + } + } + const size_t data_offset = gguf_get_data_offset(gctx); + GgufMmap mmap; + std::string mmap_error; + if (!mmap.open(path, mmap_error)) { + return fail(mmap_error); + } + for (const auto & contract : contracts) { + const int64_t id = gguf_find_tensor(gctx, contract.name); + const size_t offset = gguf_get_tensor_offset(gctx, id); + const size_t size = gguf_get_tensor_size(gctx, id); + if (data_offset > mmap.size() || offset > mmap.size() - data_offset || + size > mmap.size() - data_offset - offset || + !copy_tensor_from_file( + gctx, contract.name, mmap.data(), data_offset, contract.destination)) { + return fail(std::string("scoring head tensor load failed: ") + + contract.name); + } + } + gguf_free(gctx); + if (tensor_ctx) ggml_free(tensor_ctx); + out.scoring_head_loaded = true; + std::fprintf(stderr, "[qwen3-0.6b] loaded scoring head: %s\n", path.c_str()); + return true; +} + } // namespace bool load_qwen3_drafter_model(const std::string & path, @@ -287,6 +387,23 @@ bool load_qwen3_drafter_model(const std::string & path, out.ctx = nullptr; return false; } + if (const char * head_path = std::getenv("PFLASH_SCORING_HEAD_GGUF")) { + constexpr const char * expected_drafter_sha256 = + "f9c9f1d3c1e21755b82d4e165f88dbbbd4355646d632fb5d6cef7c66ed4ee04e"; + const auto drafter_identity = read_gguf_metadata(path, true); + if (!*head_path || out.n_layer < 14 || !drafter_identity.ok || + drafter_identity.sha256 != expected_drafter_sha256 || + !load_scoring_head(head_path, drafter_identity.sha256, out)) { + if (drafter_identity.sha256 != expected_drafter_sha256) { + set_last_error("scoring head requires the pinned Qwen3-0.6B drafter GGUF"); + } + ggml_backend_buffer_free(out.buf); + ggml_free(out.ctx); + out.buf = nullptr; + out.ctx = nullptr; + return false; + } + } return true; } @@ -294,6 +411,7 @@ void free_qwen3_drafter_model(Qwen3DrafterWeights & w) { if (w.buf) { ggml_backend_buffer_free(w.buf); w.buf = nullptr; } if (w.ctx) { ggml_free(w.ctx); w.ctx = nullptr; } w.layers.clear(); + w.scoring_head_loaded = false; w.tok_embd = w.out_norm = w.output = nullptr; w.backend = nullptr; } diff --git a/server/src/qwen35/gguf_target_loader.cpp b/server/src/qwen35/gguf_target_loader.cpp index 3a869b2f6..15f2d9a9f 100644 --- a/server/src/qwen35/gguf_target_loader.cpp +++ b/server/src/qwen35/gguf_target_loader.cpp @@ -562,7 +562,9 @@ bool load_target_gguf_partial(const std::string & path, out.tok_embd = g("token_embd.weight"); out.out_norm = g("output_norm.weight"); out.output = g("output.weight"); - if (!out.tok_embd || !out.out_norm || !out.output) { + // Tied-embedding exports omit output.weight; that is only fatal when the + // load plan needs the lm_head (the PFlash drafter never computes logits). + if (!out.tok_embd || !out.out_norm || (!out.output && plan.load_output)) { set_last_error("missing top-level tensors (token_embd/output_norm/output)"); gguf_free(gctx); return false; diff --git a/server/src/qwen35/qwen35_backend.cpp b/server/src/qwen35/qwen35_backend.cpp index d86a5bfc0..4e1150bed 100644 --- a/server/src/qwen35/qwen35_backend.cpp +++ b/server/src/qwen35/qwen35_backend.cpp @@ -1197,7 +1197,7 @@ std::vector Qwen35Backend::compress_batch( result.compressed_ids = drafter_score_and_compress( drafter_ctx_, request.input_ids, request.keep_ratio, /*chunk_size=*/32, request.score_query_tokens, /*pool_kernel=*/13, - request.score_query_end); + request.score_query_end, request.required_instruction_spans); result.ok = !result.compressed_ids.empty(); if (result.ok) { std::fprintf(stderr, "[compress] %zu -> %zu tokens\n", diff --git a/server/src/qwen35/qwen35_layer_split_adapter.cpp b/server/src/qwen35/qwen35_layer_split_adapter.cpp index 4cee40651..3189647ec 100644 --- a/server/src/qwen35/qwen35_layer_split_adapter.cpp +++ b/server/src/qwen35/qwen35_layer_split_adapter.cpp @@ -1388,7 +1388,7 @@ Qwen35LayerSplitAdapter::compress(const ModelBackend::CompressRequest & req) { result.compressed_ids = drafter_score_and_compress( pflash_drafter_, req.input_ids, req.keep_ratio, /*chunk_size=*/32, req.score_query_tokens, /*pool_kernel=*/13, - req.score_query_end); + req.score_query_end, req.required_instruction_spans); result.ok = !result.compressed_ids.empty(); if (result.ok) { std::fprintf(stderr, "[target-split][compress] %zu -> %zu tokens\n", From 2877e58ba7374d47eaac8fe709146ca02da758fa Mon Sep 17 00:00:00 2001 From: Graffioh Date: Sat, 19 Sep 2026 14:11:44 +0000 Subject: [PATCH 02/26] feat(pflash): strict budget selection, explicit scorer query and compression trace Wire the request path to the strict selector. A request may now carry an explicit `pflash_query`, mapped to drafter tokens by the parser rules, and `required_text` spans that must survive compression; tool definitions and the instruction structure of the chat template are mapped the same way and passed to the drafter as required spans. Compression failures fail closed instead of silently falling back, and a trace records the resolved config, the mapped spans and the retained budget. The adaptive keep-ratio controller now seeds a new session from the configured curve for its prompt length rather than a fixed default. The drafter IPC gains compress2/compress3 so a remote drafter receives the unquantized keep ratio, the query window and the required spans, with one formatter and one parser shared by both ends. The server unit tests move with the query-mapping API they exercise, since the older helpers this replaces had no other callers. Co-Authored-By: Claude Fable 5.1 --- server/src/common/pflash_drafter_ipc.cpp | 228 ++++- server/src/common/pflash_drafter_ipc.h | 38 +- .../src/common/pflash_drafter_ipc_daemon.cpp | 28 +- server/src/server/adaptive_keep_ratio.h | 17 +- server/src/server/http_server.cpp | 821 +++++++++++++++--- server/src/server/http_server.h | 119 ++- server/test/test_server_unit.cpp | 709 ++++++++++++++- 7 files changed, 1795 insertions(+), 165 deletions(-) diff --git a/server/src/common/pflash_drafter_ipc.cpp b/server/src/common/pflash_drafter_ipc.cpp index 724717e49..a36d71d26 100644 --- a/server/src/common/pflash_drafter_ipc.cpp +++ b/server/src/common/pflash_drafter_ipc.cpp @@ -3,11 +3,219 @@ #include "pflash_drafter_ipc.h" #include +#include +#include #include #include +#include +#include namespace dflash::common { +namespace { + +bool parse_int_token(const std::string & raw, int & out) { + if (raw.empty()) return false; + errno = 0; + char * end = nullptr; + const long value = std::strtol(raw.c_str(), &end, 10); + if (errno == ERANGE || end == raw.c_str() || *end != '\0' || + value < INT_MIN || value > INT_MAX) { + return false; + } + out = (int) value; + return true; +} + +bool parse_float_token(const std::string & raw, float & out) { + if (raw.empty()) return false; + errno = 0; + char * end = nullptr; + const float value = std::strtof(raw.c_str(), &end); + if (errno == ERANGE || end == raw.c_str() || *end != '\0' || + !std::isfinite(value)) { + return false; + } + out = value; + return true; +} + +bool validate_request_fields( + float keep_ratio, + int score_query_tokens, + const std::vector & instruction_spans, + const std::string & path, + std::string & error) { + if (!std::isfinite(keep_ratio) || keep_ratio < 0.0f || keep_ratio > 1.0f) { + error = "PFlash IPC keep_ratio must be finite and in [0, 1]"; + return false; + } + if (score_query_tokens < 1) { + error = "PFlash IPC score_query_tokens must be positive"; + return false; + } + if (instruction_spans.size() > kPFlashMaxInstructionSpans) { + error = "PFlash IPC has too many instruction spans"; + return false; + } + int previous_end = 0; + for (const auto & span : instruction_spans) { + if (span.begin < 0 || span.end <= span.begin) { + error = "PFlash IPC instruction span is invalid"; + return false; + } + if (span.begin < previous_end) { + error = "PFlash IPC instruction spans must be ordered and non-overlapping"; + return false; + } + previous_end = span.end; + } + if (path.empty()) { + error = "PFlash IPC token path must not be empty"; + return false; + } + return true; +} + +} // namespace + +bool format_pflash_drafter_ipc_compress_command( + float keep_ratio, + int score_query_end, + int score_query_tokens, + const std::string & path, + std::string & out, + std::string & error) { + out.clear(); + error.clear(); + if (!validate_request_fields( + keep_ratio, score_query_tokens, {}, path, error)) { + return false; + } + + char keep_text[64]; + std::snprintf(keep_text, sizeof(keep_text), "%.9g", keep_ratio); + + std::ostringstream line; + line << "compress2 " << keep_text << ' ' << score_query_end << ' ' + << score_query_tokens << ' ' << path; + out = line.str(); + return true; +} + +bool format_pflash_drafter_ipc_compress_command( + float keep_ratio, + int score_query_end, + int score_query_tokens, + const std::vector & required_instruction_spans, + const std::string & path, + std::string & out, + std::string & error) { + out.clear(); + error.clear(); + if (required_instruction_spans.empty()) { + return format_pflash_drafter_ipc_compress_command( + keep_ratio, score_query_end, score_query_tokens, + path, out, error); + } + if (!validate_request_fields( + keep_ratio, score_query_tokens, + required_instruction_spans, path, error)) { + return false; + } + + char keep_text[64]; + std::snprintf(keep_text, sizeof(keep_text), "%.9g", keep_ratio); + std::ostringstream line; + line << "compress3 " << keep_text << ' ' << score_query_end << ' ' + << score_query_tokens << ' ' << required_instruction_spans.size(); + for (const auto & span : required_instruction_spans) { + line << ' ' << span.begin << ' ' << span.end; + } + line << ' ' << path; + out = line.str(); + return true; +} + +bool parse_pflash_drafter_ipc_compress_command( + const std::string & line, + PFlashDrafterIpcCompressCommand & out, + std::string & error) { + out = {}; + error.clear(); + + std::istringstream iss(line); + std::string command; + if (!(iss >> command)) { + error = "PFlash IPC command is empty"; + return false; + } + + std::string keep_raw; + std::string query_end_raw; + std::string query_tokens_raw; + if (!(iss >> keep_raw >> query_end_raw >> query_tokens_raw)) { + error = "PFlash IPC compress command is missing fields"; + return false; + } + if (!parse_int_token(query_end_raw, out.score_query_end) || + !parse_int_token(query_tokens_raw, out.score_query_tokens)) { + error = "PFlash IPC query fields must be integers"; + return false; + } + if (command == "compress3") { + if (!parse_float_token(keep_raw, out.keep_ratio)) { + error = "PFlash IPC keep_ratio must be a float"; + return false; + } + std::string count_raw; + int span_count = -1; + if (!(iss >> count_raw) || !parse_int_token(count_raw, span_count) || + span_count < 0 || + (size_t) span_count > kPFlashMaxInstructionSpans) { + error = "PFlash IPC instruction span count is invalid"; + return false; + } + out.required_instruction_spans.reserve((size_t) span_count); + for (int index = 0; index < span_count; ++index) { + std::string begin_raw; + std::string end_raw; + PFlashTokenSpan span; + if (!(iss >> begin_raw >> end_raw) || + !parse_int_token(begin_raw, span.begin) || + !parse_int_token(end_raw, span.end)) { + error = "PFlash IPC instruction span fields must be integers"; + return false; + } + out.required_instruction_spans.push_back(span); + } + out.path = read_line_tail(iss); + } else if (command == "compress2") { + if (!parse_float_token(keep_raw, out.keep_ratio)) { + error = "PFlash IPC keep_ratio must be a float"; + return false; + } + out.path = read_line_tail(iss); + } else if (command == "compress") { + int keep_x1000 = 0; + if (!parse_int_token(keep_raw, keep_x1000) || + keep_x1000 < 0 || keep_x1000 > 1000) { + error = "PFlash IPC legacy keep_x1000 must be in [0, 1000]"; + return false; + } + out.legacy_quantized_ratio = true; + out.keep_ratio = (float) keep_x1000 / 1000.0f; + out.path = read_line_tail(iss); + } else { + error = "unknown PFlash IPC command"; + return false; + } + + return validate_request_fields( + out.keep_ratio, out.score_query_tokens, + out.required_instruction_spans, out.path, error); +} + bool PFlashDrafterIpcClient::start( const std::string & bin, const std::string & drafter_path, @@ -42,10 +250,12 @@ bool PFlashDrafterIpcClient::compress( float keep_ratio, std::vector & compressed_ids, int score_query_end, - int score_query_tokens) { + int score_query_tokens, + const std::vector & required_instruction_spans) { #if defined(_WIN32) (void)input_ids; (void)keep_ratio; (void)compressed_ids; (void)score_query_end; (void)score_query_tokens; + (void)required_instruction_spans; return false; #else compressed_ids.clear(); @@ -58,12 +268,16 @@ bool PFlashDrafterIpcClient::compress( std::fprintf(stderr, "pflash-ipc write tokens failed: %s\n", path.c_str()); return false; } - int keep_x1000 = (int)std::lround(std::max(0.0f, keep_ratio) * 1000.0f); - keep_x1000 = std::max(0, std::min(1000, keep_x1000)); - - std::fprintf(cmd, "compress %d %d %d %s\n", - keep_x1000, score_query_end, score_query_tokens, - path.c_str()); + std::string line; + std::string error; + if (!format_pflash_drafter_ipc_compress_command( + keep_ratio, score_query_end, score_query_tokens, + required_instruction_spans, path, line, error)) { + std::fprintf(stderr, "pflash-ipc bad compress request: %s\n", error.c_str()); + std::remove(path.c_str()); + return false; + } + std::fprintf(cmd, "%s\n", line.c_str()); std::fflush(cmd); int32_t status = -1; diff --git a/server/src/common/pflash_drafter_ipc.h b/server/src/common/pflash_drafter_ipc.h index 0d91460f2..3dba15c88 100644 --- a/server/src/common/pflash_drafter_ipc.h +++ b/server/src/common/pflash_drafter_ipc.h @@ -8,6 +8,7 @@ #include "backend_ipc.h" #include "io_utils.h" +#include "pflash_types.h" #include #include @@ -16,9 +17,36 @@ namespace dflash::common { -inline bool valid_pflash_score_query_tokens(int score_query_tokens) { - return score_query_tokens >= 1 && score_query_tokens <= 8; -} +struct PFlashDrafterIpcCompressCommand { + bool legacy_quantized_ratio = false; + float keep_ratio = 0.0f; + int score_query_end = -1; + int score_query_tokens = 8; + std::vector required_instruction_spans; + std::string path; +}; + +bool format_pflash_drafter_ipc_compress_command( + float keep_ratio, + int score_query_end, + int score_query_tokens, + const std::string & path, + std::string & out, + std::string & error); + +bool format_pflash_drafter_ipc_compress_command( + float keep_ratio, + int score_query_end, + int score_query_tokens, + const std::vector & required_instruction_spans, + const std::string & path, + std::string & out, + std::string & error); + +bool parse_pflash_drafter_ipc_compress_command( + const std::string & line, + PFlashDrafterIpcCompressCommand & out, + std::string & error); class PFlashDrafterIpcClient { public: @@ -36,7 +64,9 @@ class PFlashDrafterIpcClient { float keep_ratio, std::vector & compressed_ids, int score_query_end = -1, - int score_query_tokens = 8); + int score_query_tokens = 8, + const std::vector & + required_instruction_spans = {}); bool active() const { return active_; } void close(); diff --git a/server/src/common/pflash_drafter_ipc_daemon.cpp b/server/src/common/pflash_drafter_ipc_daemon.cpp index dc27d0ffd..f63299013 100644 --- a/server/src/common/pflash_drafter_ipc_daemon.cpp +++ b/server/src/common/pflash_drafter_ipc_daemon.cpp @@ -45,31 +45,27 @@ int run_pflash_drafter_ipc_daemon(const char * drafter_path, std::string cmd; iss >> cmd; if (cmd == "quit" || cmd == "exit") break; - if (cmd == "compress") { - int keep_x1000 = 0; - int score_query_end = -1; - int score_query_tokens = 8; - iss >> keep_x1000 >> score_query_end >> score_query_tokens; - std::string path = read_line_tail(iss); - if (keep_x1000 < 0 || keep_x1000 > 1000 || - !valid_pflash_score_query_tokens(score_query_tokens) || - path.empty()) { - std::fprintf(stderr, "[pflash-ipc-daemon] bad compress: %s\n", - line.c_str()); + if (cmd == "compress" || cmd == "compress2" || cmd == "compress3") { + PFlashDrafterIpcCompressCommand request; + std::string parse_error; + if (!parse_pflash_drafter_ipc_compress_command(line, request, parse_error)) { + std::fprintf(stderr, "[pflash-ipc-daemon] bad compress: %s (%s)\n", + line.c_str(), parse_error.c_str()); stream_status(stream_fd, -1); continue; } - auto input_ids = read_int32_file(path); + auto input_ids = read_int32_file(request.path); if (input_ids.empty()) { std::fprintf(stderr, "[pflash-ipc-daemon] read tokens failed: %s\n", - path.c_str()); + request.path.c_str()); stream_status(stream_fd, -1); continue; } - const float keep = (float)keep_x1000 / 1000.0f; auto compressed = drafter_score_and_compress( - ctx, input_ids, keep, /*chunk_size=*/32, score_query_tokens, - /*pool_kernel=*/13, score_query_end); + ctx, input_ids, request.keep_ratio, /*chunk_size=*/32, + request.score_query_tokens, /*pool_kernel=*/13, + request.score_query_end, + request.required_instruction_spans); if (compressed.empty()) { std::fprintf(stderr, "[pflash-ipc-daemon] compress returned empty\n"); stream_status(stream_fd, -1); diff --git a/server/src/server/adaptive_keep_ratio.h b/server/src/server/adaptive_keep_ratio.h index 959b87bce..a2ca3c55d 100644 --- a/server/src/server/adaptive_keep_ratio.h +++ b/server/src/server/adaptive_keep_ratio.h @@ -53,23 +53,32 @@ inline AdaptiveKeepRatioState step_adaptive_keep_ratio( // Prevents memory exhaustion from unbounded unique-session insertion. class HttpServerSessions { public: - void update(const std::string& session_id, float observed_accept) { + // ``seed_keep`` is the ratio a brand-new session adapts from: the server + // passes the configured (curve) ratio so a session starts at the real-use + // budget for its prompt length rather than the fixed default, and the + // controller then moves it by acceptance feedback within the same bounds. + void update(const std::string& session_id, float observed_accept, + float seed_keep = AdaptiveKeepRatioState{}.last_keep) { std::lock_guard lock(mu_); auto it = map_.find(session_id); if (it == map_.end()) { evict_if_full_locked(); lru_.push_front(session_id); - map_.emplace(session_id, Entry{step_adaptive_keep_ratio({}, observed_accept), lru_.begin()}); + AdaptiveKeepRatioState seed; + seed.last_keep = std::clamp(seed_keep, kBanditKeepMin, kBanditKeepMax); + map_.emplace(session_id, Entry{step_adaptive_keep_ratio(seed, observed_accept), lru_.begin()}); } else { it->second.state = step_adaptive_keep_ratio(it->second.state, observed_accept); lru_.splice(lru_.begin(), lru_, it->second.lru_it); } } - float get_keep_ratio(const std::string& session_id) const { + // A session with no feedback yet reports ``fallback`` (the configured ratio). + float get_keep_ratio(const std::string& session_id, + float fallback = AdaptiveKeepRatioState{}.last_keep) const { std::lock_guard lock(mu_); auto it = map_.find(session_id); - if (it == map_.end()) return AdaptiveKeepRatioState{}.last_keep; + if (it == map_.end()) return fallback; lru_.splice(lru_.begin(), lru_, it->second.lru_it); return it->second.state.last_keep; } diff --git a/server/src/server/http_server.cpp b/server/src/server/http_server.cpp index d9fc9da29..933a19ad7 100644 --- a/server/src/server/http_server.cpp +++ b/server/src/server/http_server.cpp @@ -27,6 +27,7 @@ #include "pin_friendly_prompt.h" #include "common/kv_rotation.h" #include "common/sha1.h" +#include "qwen3/pflash_selection.h" #include "freeze_history.h" #ifdef DFLASH_HAS_CURL @@ -37,6 +38,7 @@ #include #include #include +#include #include #include #include @@ -189,38 +191,245 @@ HeartbeatSendResult try_send_sse_heartbeat( return HeartbeatSendResult::Complete; } -std::string pflash_user_query_text( - const std::vector & messages) { - for (auto it = messages.rbegin(); it != messages.rend(); ++it) { - if (it->role == "user") return it->content; - } - return {}; -} - PflashQueryWindow find_pflash_query_window( const std::vector & prompt, const std::vector & query, + int max_tokens, int search_end, - int max_tokens) { - if (prompt.empty() || query.empty() || search_end < 1 || - search_end > (int) prompt.size() || max_tokens < 1) { - return {}; - } - + int search_begin, + bool anchored) { + PflashQueryWindow result; + if (prompt.empty() || query.empty() || max_tokens < 1) return result; + + const int limit = search_end < 0 + ? (int) prompt.size() + : (std::min)((int) prompt.size(), search_end); + if (limit < 1 || search_begin < 0 || search_begin >= limit) return result; const int widest = (std::min)(max_tokens, (int) query.size()); // For a short query, require all available tokens. For a normal query, // four matching suffix tokens are enough to tolerate a BPE boundary // difference without accidentally selecting a lone punctuation token. const int narrowest = (std::min)(4, widest); - const auto prompt_end = prompt.begin() + search_end; - for (int width = widest; width >= narrowest; --width) { - const auto match = std::find_end( - prompt.begin(), prompt_end, query.end() - width, query.end()); - if (match != prompt_end) { - return {(int) (match - prompt.begin()) + width, width}; + // A configured semantic boundary is the exact content end. Shorten only + // the suffix width there; accepting an earlier occurrence would silently + // turn a failed latest-user mapping into preceding context. + if (search_end >= 0 && anchored) { + const int bounded_widest = (std::min)(widest, limit - search_begin); + for (int width = bounded_widest; width >= narrowest; --width) { + const auto query_begin = query.end() - width; + if (std::equal(query_begin, query.end(), + prompt.begin() + limit - width)) { + result.end = limit; + result.tokens = width; + return result; + } + } + return result; + } + // Unanchored: the latest occurrence of the (suffix-trimmed) query inside + // [search_begin, limit). An explicit query may sit before trailing + // instructions, so it is not required to end at the content boundary. + // Its last token may also merge with what follows it in the prompt + // (trailing whitespace before a newline), so up to two trailing query + // tokens may be dropped; the untrimmed query is always preferred. + const int floor = (std::max)(0, search_begin); + const int max_trailing = (std::min)(2, (int) query.size() - narrowest); + for (int drop = 0; drop <= max_trailing; ++drop) { + const auto query_end = query.end() - drop; + const int available = (int) query.size() - drop; + const int drop_widest = (std::min)(max_tokens, available); + for (int width = drop_widest; width >= narrowest; --width) { + const auto query_begin = query_end - width; + for (int end = limit; end - width >= floor; --end) { + if (std::equal(query_begin, query_end, + prompt.begin() + end - width)) { + result.end = end; + result.tokens = width; + result.trailing_trimmed = drop; + return result; + } + } } } - return {}; + return result; +} + +PflashQueryWindow pflash_tail_query_window( + const std::vector & prompt, + int max_tokens, + int query_end, + int query_begin) noexcept { + PflashQueryWindow result; + if (prompt.empty() || max_tokens < 1) return result; + result.end = query_end < 0 + ? static_cast(prompt.size()) + : query_end; + if (result.end < 1 || result.end > static_cast(prompt.size()) || + query_begin < 0 || query_begin >= result.end) { + return {}; + } + result.tokens = (std::min)(max_tokens, result.end - query_begin); + return result; +} + +PFlashTokenSpan pflash_decoded_text_span( + const Tokenizer & tokenizer, + const std::vector & prompt, + int begin, + int end, + const std::string & needle) { + if (needle.empty() || begin < 0 || end <= begin || + end > (int) prompt.size()) { + return {-1, -1}; + } + std::string decoded; + std::vector offsets; + offsets.reserve((size_t)(end - begin)); + for (int index = begin; index < end; ++index) { + offsets.push_back(decoded.size()); + decoded += tokenizer.token_text(prompt[(size_t) index]); + } + const size_t pos = decoded.rfind(needle); + if (pos == std::string::npos) { + return {-1, -1}; + } + const size_t needle_end = pos + needle.size(); + int first = begin; + while (first + 1 < end && offsets[(size_t)(first + 1 - begin)] <= pos) { + ++first; + } + int after = first; + while (after < end && offsets[(size_t)(after - begin)] < needle_end) { + ++after; + } + return {first, after}; +} + +int pflash_query_search_end_from_sentinel( + const std::vector & original, + const std::vector & sentinel) noexcept { + size_t common_suffix = 0; + while (common_suffix < original.size() && + common_suffix < sentinel.size() && + original[original.size() - common_suffix - 1] == + sentinel[sentinel.size() - common_suffix - 1]) { + ++common_suffix; + } + if (common_suffix == 0 || common_suffix >= original.size()) { + return -1; + } + return static_cast(original.size() - common_suffix); +} +int pflash_query_search_begin_from_sentinel( + const std::vector & original, + const std::vector & sentinel) noexcept { + size_t common_prefix = 0; + while (common_prefix < original.size() && + common_prefix < sentinel.size() && + original[common_prefix] == sentinel[common_prefix]) { + ++common_prefix; + } + if (common_prefix >= original.size()) return -1; + return static_cast(common_prefix); +} + +PflashInstructionMessagePlan plan_pflash_instruction_messages( + const std::vector & messages) { + PflashInstructionMessagePlan plan; + for (size_t index = 0; index < messages.size(); ++index) { + const auto & message = messages[index]; + const bool instruction_role = + message.role == "system" || message.role == "developer"; + if (instruction_role && !message.content.empty()) { + plan.instruction_messages.push_back(index); + } + } + return plan; +} + +PFlashTokenSpan pflash_changed_token_span( + const std::vector & original, + const std::vector & variant) noexcept { + size_t common_prefix = 0; + while (common_prefix < original.size() && + common_prefix < variant.size() && + original[common_prefix] == variant[common_prefix]) { + ++common_prefix; + } + + size_t common_suffix = 0; + while (common_suffix < original.size() - common_prefix && + common_suffix < variant.size() - common_prefix && + original[original.size() - common_suffix - 1] == + variant[variant.size() - common_suffix - 1]) { + ++common_suffix; + } + + const int begin = static_cast(common_prefix); + const int end = static_cast(original.size() - common_suffix); + return end > begin ? PFlashTokenSpan{begin, end} + : PFlashTokenSpan{-1, -1}; +} + +std::vector canonicalize_pflash_token_spans( + std::vector spans) { + std::sort(spans.begin(), spans.end(), [] ( + const PFlashTokenSpan & left, + const PFlashTokenSpan & right) { + return left.begin < right.begin || + (left.begin == right.begin && left.end < right.end); + }); + std::vector result; + for (const PFlashTokenSpan & span : spans) { + if (!result.empty() && span.begin <= result.back().end) { + result.back().end = std::max(result.back().end, span.end); + } else { + result.push_back(span); + } + } + return result; +} + +std::string pflash_token_fingerprint( + const std::vector & ids) { + uint64_t hash = UINT64_C(14695981039346656037); + for (int32_t token : ids) { + const uint32_t value = static_cast(token); + for (int shift = 0; shift < 32; shift += 8) { + hash ^= static_cast(value >> shift); + hash *= UINT64_C(1099511628211); + } + } + + char encoded[17]; + std::snprintf(encoded, sizeof(encoded), "%016llx", + static_cast(hash)); + return encoded; +} + + +bool pflash_full_cache_restore_allowed( + bool selection_environment_present) noexcept { + return !selection_environment_present; +} + +bool pflash_continuation_must_fail_closed( + bool selection_environment_present) noexcept { + return selection_environment_present; +} + +int pflash_target_token_ceiling( + int original_target_tokens, double keep_ratio) noexcept { + if (original_target_tokens < 0 || !std::isfinite(keep_ratio) || + keep_ratio <= 0.0) { + return -1; + } + const double ceiling = std::floor( + static_cast(original_target_tokens) * keep_ratio); + if (!std::isfinite(ceiling) || ceiling < 0.0 || ceiling > INT_MAX) { + return -1; + } + return static_cast(ceiling); } } // namespace http_detail @@ -280,9 +489,11 @@ bool flowkv_should_activate(const ServerConfig & config, float resolve_pflash_keep_ratio(float configured_ratio, const std::string & session_id, const HttpServerSessions & sessions) { + // A session adapts from the configured (curve) ratio: until its first + // acceptance feedback it keeps that ratio, afterwards the controller's. return session_id.empty() ? configured_ratio - : sessions.get_keep_ratio(session_id); + : sessions.get_keep_ratio(session_id, configured_ratio); } bool should_clamp_flowkv_disk_cache( @@ -2436,6 +2647,8 @@ bool HttpServer::handle_model_request(SocketHandle fd, ParsedRequest & req, apply_request_reasoning(body, config_, req); // Bandit: parse session_id from extra_body (opt-in adaptive keep_ratio). req.session_id = parse_session_id_from_body(body); + req.pflash_query = parse_pflash_query_from_body(body); + req.pflash_required = parse_pflash_required_from_body(body); // PPP rearrange (optional): peel ephemeral system banners into a // following system message so the first chat boundary is stable. @@ -3112,8 +3325,11 @@ void HttpServer::apply_flowkv_compression( std::string HttpServer::apply_pflash_compression( const ParsedRequest & req, PreparedPrompt & prepared) { + const bool selection_environment = + dflash::qwen3::has_pflash_selection_environment(); auto [full_slot, full_len] = prefix_cache_.lookup_full(req.prompt_tokens); - if (full_slot >= 0) { + if (http_detail::pflash_full_cache_restore_allowed( + selection_environment) && full_slot >= 0) { std::fprintf(stderr, "[pflash] full-cache hit slot=%d — skipping compress\n", full_slot); @@ -3130,45 +3346,368 @@ std::string HttpServer::apply_pflash_compression( const std::string prompt_text = tokenizer_.decode(req.prompt_tokens); auto drafter_ids = drafter_tokenizer_->encode(prompt_text); - const std::vector chat_messages = normalize_chat_messages( - req.messages, req.format, tool_memory_); - std::string rendered_messages; - std::string render_error; - if (!render_messages_to_text( - chat_messages, req, /*add_generation_prompt=*/false, - rendered_messages, render_error)) { - std::fprintf(stderr, - "[pflash] ERROR: scorer query boundary render failed; " - "refusing compression\n"); - return "PFlash scorer query boundary render failed"; - } - const std::string normalized_messages = tokenizer_.decode( - tokenizer_.encode(rendered_messages)); - const auto rendered_message_ids = drafter_tokenizer_->encode( - normalized_messages); - const auto shared_end = std::mismatch( - drafter_ids.begin(), drafter_ids.end(), - rendered_message_ids.begin(), rendered_message_ids.end()).first; - const int query_search_end = (int) (shared_end - drafter_ids.begin()); - - const std::string last_user_text = - http_detail::pflash_user_query_text(chat_messages); - const auto query_ids = last_user_text.empty() - ? std::vector{} - : drafter_tokenizer_->encode(last_user_text); - const auto query_window = http_detail::find_pflash_query_window( - drafter_ids, query_ids, query_search_end); if (drafter_ids.empty()) { return "PFlash drafter tokenizer produced an empty prompt"; } + dflash::qwen3::PFlashSelectionConfig experiment; + std::string experiment_error; + if (!dflash::qwen3::resolve_pflash_selection( + (int) drafter_ids.size(), 32, experiment, experiment_error)) { + return "invalid PFlash strict selection config: " + experiment_error; + } + if (!experiment.selection_active && !req.pflash_required.empty()) { + return "PFlash pflash_required needs strict budget selection"; + } + + const bool messages_input = + req.messages.is_array() && !req.messages.empty(); + const bool raw_text_input = req.messages.is_string(); + const char * parser_input_kind = messages_input + ? "messages" : (raw_text_input ? "raw_text" : "unsupported"); + std::string parser_selection_rule; + std::string last_user_text; + int query_content_begin = -1; + int query_content_end = -1; + // Complete token span of an explicit pflash_query inside the boundary + // content, when it was mapped against the decoded token text. The strict + // selector keeps the whole span mandatory; the scorer window is its tail. + PFlashTokenSpan explicit_query_span{-1, -1}; + std::vector required_instruction_spans; + if (experiment.configured) { + if (!messages_input && !raw_text_input) { + return "PFlash strict selection input has no parseable text"; + } + try { + auto messages = + normalize_chat_messages(req.messages, req.format, tool_memory_); + if (messages.empty()) { + return "PFlash strict selection normalized messages are empty"; + } + + int last_user_index = -1; + for (int index = (int) messages.size() - 1; index >= 0; --index) { + if (messages[(size_t) index].role == "user") { + last_user_index = index; + break; + } + } + if (last_user_index >= 0) { + last_user_text = messages[(size_t) last_user_index].content; + } + + int boundary_index = (int) messages.size() - 1; + if (!raw_text_input && + experiment.query_parser == + dflash::qwen3::PFlashQueryParser::SemanticUser) { + boundary_index = last_user_index; + } + if (boundary_index < 0 || + (experiment.query_parser == + dflash::qwen3::PFlashQueryParser::SemanticUser && + !raw_text_input && last_user_text.empty())) { + return "PFlash strict selection latest-user boundary is unavailable"; + } + + static constexpr const char * kContentBegin = + "__DFLASH_PFLASH_CONTENT_BEGIN_02C47F91__"; + static constexpr const char * kContentEnd = + "__DFLASH_PFLASH_CONTENT_END_6E6B61A8__"; + const auto map_message_content = [&] ( + size_t message_index, + int & content_begin, + int & content_end, + std::string & boundary_error) -> bool { + auto begin_messages = messages; + begin_messages[message_index].content = + std::string(kContentBegin) + + begin_messages[message_index].content; + auto end_messages = messages; + end_messages[message_index].content += kContentEnd; + + std::string begin_rendered; + std::string end_rendered; + if (!render_messages_to_text( + begin_messages, req, /*add_generation_prompt=*/true, + begin_rendered, boundary_error)) { + boundary_error = "content-start render failed: " + + boundary_error; + return false; + } + boundary_error.clear(); + if (!render_messages_to_text( + end_messages, req, /*add_generation_prompt=*/true, + end_rendered, boundary_error)) { + boundary_error = "content-end render failed: " + + boundary_error; + return false; + } + + const auto begin_ids = drafter_tokenizer_->encode( + tokenizer_.decode(tokenizer_.encode(begin_rendered))); + const auto end_ids = drafter_tokenizer_->encode( + tokenizer_.decode(tokenizer_.encode(end_rendered))); + content_begin = + http_detail::pflash_query_search_begin_from_sentinel( + drafter_ids, begin_ids); + content_end = + http_detail::pflash_query_search_end_from_sentinel( + drafter_ids, end_ids); + if (content_begin < 0 || content_end <= content_begin || + content_end >= (int) drafter_ids.size()) { + boundary_error = "content boundary mapping failed"; + return false; + } + return true; + }; + const auto map_rendered_message = [&] ( + size_t message_index, + PFlashTokenSpan & message_span, + std::string & boundary_error) -> bool { + auto without_message = messages; + without_message.erase(without_message.begin() + message_index); + + std::string without_message_rendered; + if (!render_messages_to_text( + without_message, req, /*add_generation_prompt=*/true, + without_message_rendered, boundary_error)) { + boundary_error = "message-removal render failed: " + + boundary_error; + return false; + } + const auto without_message_ids = drafter_tokenizer_->encode( + tokenizer_.decode(tokenizer_.encode( + without_message_rendered))); + message_span = http_detail::pflash_changed_token_span( + drafter_ids, without_message_ids); + if (message_span.begin < 0) { + boundary_error = + "complete rendered message did not map to a prompt span"; + return false; + } + return true; + }; + + std::string boundary_error; + if (!map_message_content( + (size_t) boundary_index, + query_content_begin, query_content_end, + boundary_error)) { + return "PFlash strict selection " + boundary_error; + } + if (query_content_begin < 0 || + query_content_end <= query_content_begin || + query_content_end >= (int) drafter_ids.size()) { + return "PFlash strict selection content boundary mapping failed"; + } + + if (experiment.selection_active) { + const auto instruction_plan = + http_detail::plan_pflash_instruction_messages(messages); + for (size_t instruction_index : + instruction_plan.instruction_messages) { + PFlashTokenSpan instruction_span; + boundary_error.clear(); + if (!map_rendered_message( + instruction_index, instruction_span, + boundary_error)) { + return "PFlash strict selection instruction mapping failed: " + + boundary_error; + } + required_instruction_spans.push_back(instruction_span); + } + + if (!req.tools.is_null() && !req.tools.empty()) { + ParsedRequest tool_free_req = req; + tool_free_req.tools = json::array(); + std::string tool_free_rendered; + boundary_error.clear(); + if (!render_messages_to_text( + messages, tool_free_req, + /*add_generation_prompt=*/true, + tool_free_rendered, boundary_error)) { + return "PFlash strict selection tool mapping failed: " + + boundary_error; + } + const auto tool_free_ids = drafter_tokenizer_->encode( + tokenizer_.decode(tokenizer_.encode(tool_free_rendered))); + const PFlashTokenSpan tool_span = + http_detail::pflash_changed_token_span( + drafter_ids, tool_free_ids); + if (tool_span.begin < 0) { + return "PFlash strict selection tool mapping failed: " + "tools did not produce a retained prompt span"; + } + required_instruction_spans.push_back(tool_span); + } + + // Client-declared literal text that must survive compression + // (e.g. an answer-format directive embedded in the user + // message). Each string maps to its last occurrence inside + // the boundary content; mapping failure fails the request + // rather than silently keeping a shorter span. + for (const auto & required : req.pflash_required) { + if (required.empty()) continue; + const PFlashTokenSpan required_span = + http_detail::pflash_decoded_text_span( + *drafter_tokenizer_, drafter_ids, + query_content_begin, query_content_end, required); + if (required_span.begin < 0) { + return "PFlash strict selection required-text mapping " + "failed: a pflash_required string does not occur " + "in the latest user content"; + } + required_instruction_spans.push_back(required_span); + } + // An explicit scorer query also pins its complete span: the + // whole question is mandatory even though the scorer only + // consumes its bounded tail. Mapping against the decoded + // content text (not a standalone encoding) keeps BPE boundary + // merges like " What" inside the span. + if (!req.pflash_query.empty() && + experiment.query_parser == + dflash::qwen3::PFlashQueryParser::SemanticUser) { + explicit_query_span = + http_detail::pflash_decoded_text_span( + *drafter_tokenizer_, drafter_ids, + query_content_begin, query_content_end, + req.pflash_query); + if (explicit_query_span.begin < 0) { + return "PFlash strict selection explicit query mapping " + "failed: pflash_query does not occur in the " + "latest user content"; + } + required_instruction_spans.push_back(explicit_query_span); + } + required_instruction_spans = + http_detail::canonicalize_pflash_token_spans( + std::move(required_instruction_spans)); + std::string instruction_error; + if (!dflash::qwen3::validate_pflash_instruction_spans( + required_instruction_spans, + (int) drafter_ids.size(), instruction_error)) { + return "PFlash strict selection instruction mapping failed: " + + instruction_error; + } + } + } catch (const std::exception & error) { + return std::string("PFlash retention normalization failed: ") + + error.what(); + } + } else if (raw_text_input) { + last_user_text = req.messages.get(); + } else if (req.messages.is_array()) { + for (int index = (int) req.messages.size() - 1; index >= 0; --index) { + if (req.messages[index].value("role", "") != "user") continue; + const auto & content = req.messages[index]["content"]; + if (content.is_string()) { + last_user_text = content.get(); + } else if (content.is_array()) { + for (const auto & part : content) { + const std::string type = part.value("type", ""); + if (type == "text" || type == "input_text" || + type == "output_text") { + last_user_text += part.value("text", ""); + } + } + } + break; + } + } + + std::vector semantic_query_ids; + std::vector expected_query_ids; + http_detail::PflashQueryWindow query_window; + if (!req.pflash_query.empty()) { + // An explicit query replaces the message-tail heuristic; it must occur + // inside the latest user content so the window maps onto real tokens. + last_user_text = req.pflash_query; + parser_selection_rule = "explicit_query"; + } + if (!last_user_text.empty()) { + semantic_query_ids = drafter_tokenizer_->encode(last_user_text); + } + if (experiment.configured && raw_text_input) { + parser_selection_rule = "content_tail"; + query_window = http_detail::pflash_tail_query_window( + drafter_ids, experiment.query_tokens, + query_content_end, query_content_begin); + } else if (experiment.configured && + experiment.query_parser == + dflash::qwen3::PFlashQueryParser::ArbitraryTail) { + parser_selection_rule = "prompt_tail"; + query_window = http_detail::pflash_tail_query_window( + drafter_ids, experiment.query_tokens, query_content_end); + } else if (explicit_query_span.begin >= 0) { + // The explicit query was already mapped against the decoded content + // text and pinned as a mandatory span. The scorer consumes the span's + // bounded tail window; the complete span stays in the target prompt. + parser_selection_rule = "explicit_query_span"; + query_window.end = explicit_query_span.end; + query_window.tokens = (std::min)( + experiment.query_tokens, + explicit_query_span.end - explicit_query_span.begin); + expected_query_ids.assign( + drafter_ids.begin() + (query_window.end - query_window.tokens), + drafter_ids.begin() + query_window.end); + } else if (!semantic_query_ids.empty()) { + if (experiment.configured) parser_selection_rule = "semantic_suffix"; + query_window = http_detail::find_pflash_query_window( + drafter_ids, semantic_query_ids, + experiment.configured ? experiment.query_tokens : 8, + query_content_end, + experiment.configured ? query_content_begin : 0, + /*anchored=*/ req.pflash_query.empty()); + if (experiment.configured && query_window.valid()) { + const auto matched_end = + semantic_query_ids.end() - query_window.trailing_trimmed; + expected_query_ids.assign( + matched_end - query_window.tokens, matched_end); + } + } + ModelBackend::CompressRequest compress_request; compress_request.input_ids = std::move(drafter_ids); + compress_request.required_instruction_spans = + std::move(required_instruction_spans); compress_request.keep_ratio = http_detail::resolve_pflash_keep_ratio( pflash_keep_ratio(config_, prompt_tokens), req.session_id, sessions_); if (query_window.valid()) { compress_request.score_query_end = query_window.end; compress_request.score_query_tokens = query_window.tokens; + if (experiment.configured) { + const int query_begin = query_window.end - query_window.tokens; + json instruction_spans = json::array(); + for (const auto & span : + compress_request.required_instruction_spans) { + instruction_spans.push_back({span.begin, span.end}); + } + const json provenance = { + {"schema_version", 1}, + {"input_kind", parser_input_kind}, + {"selection_rule", parser_selection_rule}, + {"query_parser", + dflash::qwen3::pflash_query_parser_name( + experiment.query_parser)}, + {"input_tokens", (int) compress_request.input_ids.size()}, + {"input_fingerprint_fnv1a64", + http_detail::pflash_token_fingerprint( + compress_request.input_ids)}, + {"content_begin", query_content_begin}, + {"content_end", query_content_end}, + {"query_begin", query_begin}, + {"query_end", query_window.end}, + {"query_span_begin", explicit_query_span.begin}, + {"query_span_end", explicit_query_span.end}, + {"requested_query_tokens", experiment.query_tokens}, + {"required_text_count", req.pflash_required.size()}, + {"expected_query_ids", expected_query_ids}, + {"required_instruction_spans", instruction_spans}, + }; + std::fprintf(stderr, "[pflash-parser] %s\n", + provenance.dump().c_str()); + std::fflush(stderr); + } std::fprintf(stderr, "[pflash] scorer query mapped to drafter tokens [%d,%d); " "rendered suffix=%zu tokens\n", @@ -3191,68 +3730,119 @@ std::string HttpServer::apply_pflash_compression( }); compress_request.residency_action = residency; + // The selector budget is counted in drafter tokens, but the ceiling is + // enforced on the re-encoded target-token prompt. Vocabulary mismatches + // between the two tokenizers can inflate the re-encode past the ceiling, + // so retry with a tightened keep ratio; fail closed if it persists. + const int target_ceiling = experiment.selection_active + ? http_detail::pflash_target_token_ceiling( + prompt_tokens, compress_request.keep_ratio) + : -1; + if (experiment.selection_active && target_ceiling < 0) { + return "PFlash strict selection target-token ceiling is invalid"; + } + const float requested_keep_ratio = compress_request.keep_ratio; + ModelBackend::CompressResult result; - if (config_.pflash_remote_drafter) { - if (!pflash_remote_.active() && - !pflash_remote_.start(config_.pflash_remote.ipc_bin, - config_.pflash_drafter_path, - config_.pflash_drafter_gpu, - config_.pflash_remote.work_dir)) { - return "remote PFlash drafter start failed"; - } - result.ok = pflash_remote_.compress( - compress_request.input_ids, compress_request.keep_ratio, - result.compressed_ids, - compress_request.score_query_end, - compress_request.score_query_tokens); - if (residency == DraftResidencyAction::ReleaseAfterUse) { - pflash_remote_.close(); + std::vector final_tokens; + for (int attempt = 0; ; ++attempt) { + result = {}; + if (config_.pflash_remote_drafter) { + if (!pflash_remote_.active() && + !pflash_remote_.start(config_.pflash_remote.ipc_bin, + config_.pflash_drafter_path, + config_.pflash_drafter_gpu, + config_.pflash_remote.work_dir)) { + return "remote PFlash drafter start failed"; + } + result.ok = pflash_remote_.compress( + compress_request.input_ids, compress_request.keep_ratio, + result.compressed_ids, + compress_request.score_query_end, + compress_request.score_query_tokens, + compress_request.required_instruction_spans); + if (residency == DraftResidencyAction::ReleaseAfterUse) { + pflash_remote_.close(); + } + } else { + result = backend_.compress(compress_request); } - } else { - result = backend_.compress(compress_request); - } - if (!result.ok || result.compressed_ids.empty()) { - return config_.pflash_remote_drafter - ? "remote PFlash drafter compression failed" - : "PFlash compression failed"; - } + if (!result.ok || result.compressed_ids.empty()) { + return config_.pflash_remote_drafter + ? "remote PFlash drafter compression failed" + : "PFlash compression failed"; + } - std::string compressed_text = - drafter_tokenizer_->decode(result.compressed_ids); + std::string compressed_text = + drafter_tokenizer_->decode(result.compressed_ids); - // Compression is allowed to be lossy, but the active user query must - // survive. Re-append short queries when fewer than 80% of their tokens do. - if (!last_user_text.empty()) { - int query_kept = 0; - if (!query_ids.empty()) { - int query_index = (int) query_ids.size() - 1; - for (int kept_index = (int) result.compressed_ids.size() - 1; - kept_index >= 0 && query_index >= 0; --kept_index) { - if (result.compressed_ids[kept_index] == query_ids[query_index]) { - ++query_kept; - --query_index; + // Compression is allowed to be lossy, but the active user query must + // survive. Re-append short queries when fewer than 80% of their tokens do. + if (!experiment.selection_active && !last_user_text.empty()) { + int query_kept = 0; + if (!semantic_query_ids.empty()) { + int query_index = (int) semantic_query_ids.size() - 1; + for (int kept_index = (int) result.compressed_ids.size() - 1; + kept_index >= 0 && query_index >= 0; --kept_index) { + if (result.compressed_ids[kept_index] == semantic_query_ids[query_index]) { + ++query_kept; + --query_index; + } } } + const float survival = (float) query_kept / + (std::max)(1, (int) semantic_query_ids.size()); + std::fprintf(stderr, + "[pflash] query survival: %d/%d (%.0f%%)\n", + query_kept, (int) semantic_query_ids.size(), survival * 100.0f); + if (survival < 0.80f && (int) semantic_query_ids.size() < 1000) { + compressed_text += "\n" + last_user_text; + std::fprintf(stderr, + "[pflash] query below 80%% — re-appended full query (%d tokens)\n", + (int) semantic_query_ids.size()); + } else if (survival < 0.80f) { + std::fprintf(stderr, + "[pflash] query below 80%% but too large to re-append (%d tokens)\n", + (int) semantic_query_ids.size()); + } + } + + final_tokens = tokenizer_.encode(compressed_text); + if (!experiment.selection_active || + (int) final_tokens.size() <= target_ceiling) { + break; + } + const int overflow = (int) final_tokens.size() - target_ceiling; + const int tightened = target_ceiling - overflow - 1; + if (attempt >= 2 || tightened <= 0) { + break; + } + compress_request.keep_ratio = + requested_keep_ratio * (float) tightened / + (float) std::max(1, target_ceiling); + if (compress_request.keep_ratio <= 0.0f) { + break; } - const float survival = (float) query_kept / - (std::max)(1, (int) query_ids.size()); std::fprintf(stderr, - "[pflash] query survival: %d/%d (%.0f%%)\n", - query_kept, (int) query_ids.size(), survival * 100.0f); - if (survival < 0.80f && (int) query_ids.size() < 1000) { - compressed_text += "\n" + last_user_text; - std::fprintf(stderr, - "[pflash] query below 80%% — re-appended full query (%d tokens)\n", - (int) query_ids.size()); - } else if (survival < 0.80f) { - std::fprintf(stderr, - "[pflash] query below 80%% but too large to re-append (%d tokens)\n", - (int) query_ids.size()); + "[pflash-select] final prompt %zu exceeds ceiling %d " + "(re-encode overshoot); retrying with keep_ratio %.6f\n", + final_tokens.size(), target_ceiling, + (double) compress_request.keep_ratio); + std::fflush(stderr); + } + if (experiment.selection_active) { + std::fprintf(stderr, + "[pflash-select] final target tokens=%zu ceiling=%d\n", + final_tokens.size(), target_ceiling); + std::fflush(stderr); + if ((int) final_tokens.size() > target_ceiling) { + return "PFlash strict selection final prompt exceeds target-token ceiling " + "(" + std::to_string(final_tokens.size()) + " > " + + std::to_string(target_ceiling) + ")"; } } - - prepared.tokens = tokenizer_.encode(compressed_text); + prepared.tokens = std::move(final_tokens); prepared.compressed = true; std::fprintf(stderr, "[pflash] %d -> %d -> %d tokens (%.1f%% kept)\n", @@ -3276,6 +3866,27 @@ HttpServer::PreparedPrompt HttpServer::prepare_prompt( prompt_tokens >= config_.pflash_threshold); const bool continuation = should_compress && is_continuation_request(req.messages); + const bool selection_environment = + dflash::qwen3::has_pflash_selection_environment(); + if (should_compress && selection_environment) { + dflash::qwen3::PFlashSelectionConfig experiment; + std::string experiment_error; + if (!dflash::qwen3::resolve_pflash_selection( + 0, 32, experiment, experiment_error)) { + prepared.error_status = 500; + prepared.error = "invalid PFlash strict selection config: " + + experiment_error; + return prepared; + } + if (http_detail::pflash_continuation_must_fail_closed( + selection_environment) && + (continuation || req.disk_cache_policy.compress)) { + prepared.error_status = 500; + prepared.error = + "PFlash strict selection does not support continuation or FlowKV compression"; + return prepared; + } + } if (should_compress && continuation && req.messages.is_array()) { // FlowKV owns continuation compression automatically. Falling @@ -4420,9 +5031,11 @@ void HttpServer::process_job(ServerJob * job) { // Bandit: update when spec decode actually ran — including 0-accept case, // which signals the current keep_ratio is too low. if (result.ok() && !req.session_id.empty() && result.spec_decode_ran) { - float old_keep = sessions_.get_keep_ratio(req.session_id); + const float configured_keep = + pflash_keep_ratio(config_, (int) req.prompt_tokens.size()); + float old_keep = sessions_.get_keep_ratio(req.session_id, configured_keep); int old_turn = sessions_.turn_count(req.session_id); - sessions_.update(req.session_id, result.accept_rate); + sessions_.update(req.session_id, result.accept_rate, configured_keep); float new_keep = sessions_.get_keep_ratio(req.session_id); float ema = sessions_.get_ema(req.session_id); std::fprintf(stderr, diff --git a/server/src/server/http_server.h b/server/src/server/http_server.h index 3ee26ab86..c7fad0165 100644 --- a/server/src/server/http_server.h +++ b/server/src/server/http_server.h @@ -281,22 +281,81 @@ bool canonical_assistant_content( struct PflashQueryWindow { int end = -1; // exclusive token offset in the rendered prompt int tokens = 0; // width of the matching query suffix + int trailing_trimmed = 0; // query tokens dropped from its end to match bool valid() const { return end >= tokens && tokens > 0; } }; -// Select the final normalized user message as the scorer query. Public for -// model-free coverage of every request shape accepted by prompt rendering. -std::string pflash_user_query_text( +struct PflashInstructionMessagePlan { + std::vector instruction_messages; +}; + +PflashInstructionMessagePlan plan_pflash_instruction_messages( const std::vector & messages); -// Find the last sufficiently-specific suffix of the user query before the -// assistant-generation suffix. Public for model-free regression tests. +// Return the conservative token interval in `original` changed by rendering +// a request variant. Used to retain tool definitions independently of where +// an arbitrary chat template places them. +PFlashTokenSpan pflash_changed_token_span( + const std::vector & original, + const std::vector & variant) noexcept; + +// Sort and merge overlapping/adjacent mapped spans before selector validation. +std::vector canonicalize_pflash_token_spans( + std::vector spans); + +// Find the last sufficiently-specific suffix of the user query inside the +// rendered drafter-tokenized prompt. Public for model-free regression tests. PflashQueryWindow find_pflash_query_window( const std::vector & prompt, const std::vector & query, - int search_end, - int max_tokens = 8); + int max_tokens = 8, + int search_end = -1, + int search_begin = 0, + bool anchored = true); + +PflashQueryWindow pflash_tail_query_window( + const std::vector & prompt, + int max_tokens, + int query_end = -1, + int query_begin = 0) noexcept; + +// Map the last occurrence of `needle` inside the decoded token text of +// `prompt[begin, end)` to the token span covering it. Searching the joined +// per-token text (not a standalone encoding of `needle`) keeps the mapping +// correct at BPE boundary merges: a token that spans the needle's first or +// last character (e.g. " What" after "Question:") is included in the span. +// Returns {-1, -1} when the needle is absent or the range is invalid. +PFlashTokenSpan pflash_decoded_text_span( + const Tokenizer & tokenizer, + const std::vector & prompt, + int begin, + int end, + const std::string & needle); + +// Return the original prompt offset immediately before the stable trailing +// suffix shared with a version whose latest user message carries a sentinel. +// Invalid when no such bounded suffix can establish the semantic boundary. +int pflash_query_search_end_from_sentinel( + const std::vector & original, + const std::vector & sentinel) noexcept; + +// Return the first token offset affected by a version whose selected message +// content carries a leading sentinel. This is a conservative lower bound for +// the selected content in the original rendered prompt. +int pflash_query_search_begin_from_sentinel( + const std::vector & original, + const std::vector & sentinel) noexcept; + +std::string pflash_token_fingerprint( + const std::vector & ids); + +bool pflash_full_cache_restore_allowed( + bool selection_environment_present) noexcept; +bool pflash_continuation_must_fail_closed( + bool selection_environment_present) noexcept; +int pflash_target_token_ceiling( + int original_target_tokens, double keep_ratio) noexcept; } // namespace http_detail @@ -343,6 +402,11 @@ struct ParsedRequest { std::vector stop_sequences; // Bandit: per-session adaptive keep_ratio opt-in std::string session_id; + std::string pflash_query; // explicit scorer query text (optional request field) + // Literal strings inside the boundary message that must survive + // compression (e.g. an answer-format directive embedded in a user + // message). Each occurrence is mapped and retained as a mandatory span. + std::vector pflash_required; DiskPrefixCachePolicy disk_cache_policy; // PPP: stable pin cut for tool-heavy requests (0 = use default boundary). int pin_end_token = 0; @@ -416,6 +480,7 @@ class HttpServer { private: friend struct SchedulerTestHarness; + friend struct HttpServerTestAccess; // Client thread: read HTTP request, parse, enqueue job, wait. void handle_client(SocketHandle fd); @@ -703,6 +768,46 @@ struct ServerJob { // ─── Parse session_id from a chat-completion JSON body ────────────────── // Returns empty string when session_id is absent or not a string (int/null/array). // Checks extra_body.session_id first, then top-level session_id. +// PFlash: an explicit scorer query. The compressor scores context against +// this text instead of the last user message's tail, so a caller that knows +// its question (a benchmark, a RAG layer) can hand it over. Accepted at the top +// level or under extra_body, like session_id. +inline std::string parse_pflash_query_from_body(const json & body) { + if (body.contains("extra_body")) { + const auto & eb = body["extra_body"]; + if (eb.is_object() && eb.contains("pflash_query") && eb["pflash_query"].is_string()) { + return eb["pflash_query"].get(); + } + } + if (body.contains("pflash_query") && body["pflash_query"].is_string()) { + return body["pflash_query"].get(); + } + return {}; +} + +// PFlash: literal strings that must survive compression. Each string must +// occur inside the boundary message's content; the compressor marks its last +// occurrence there as a mandatory retention span. Accepted at the top level +// or under extra_body, like pflash_query. +inline std::vector parse_pflash_required_from_body(const json & body) { + const json * field = nullptr; + if (body.contains("extra_body")) { + const auto & eb = body["extra_body"]; + if (eb.is_object() && eb.contains("pflash_required") && eb["pflash_required"].is_array()) { + field = &eb["pflash_required"]; + } + } + if (!field && body.contains("pflash_required") && body["pflash_required"].is_array()) { + field = &body["pflash_required"]; + } + std::vector result; + if (!field) return result; + for (const auto & entry : *field) { + if (entry.is_string()) result.push_back(entry.get()); + } + return result; +} + inline std::string parse_session_id_from_body(const json & body) { if (body.contains("extra_body")) { const auto & eb = body["extra_body"]; diff --git a/server/test/test_server_unit.cpp b/server/test/test_server_unit.cpp index 7edb22428..5743f4712 100644 --- a/server/test/test_server_unit.cpp +++ b/server/test/test_server_unit.cpp @@ -7,6 +7,7 @@ // Run: ./test_server_unit #include "CppUnitTestFramework.hpp" +#include "scoped_env.h" #include "server/sse_emitter.h" #include "server/tool_parser.h" @@ -51,6 +52,7 @@ #include #include +#include #include #include #include @@ -151,6 +153,14 @@ struct SchedulerTestHarness { return server.slot_tokens_.at(slot); } }; + +struct HttpServerTestAccess { + static std::string apply_pflash_compression( + HttpServer & server, const ParsedRequest & req) { + HttpServer::PreparedPrompt prepared; + return server.apply_pflash_compression(req, prepared); + } +}; } namespace { @@ -199,26 +209,33 @@ TEST_CASE(ServerUnitFixture, test_pflash_scorer_uses_user_query_before_chat_suff }; const std::vector rendered{ 1, 2, 100, 101, 102, 103, 104, 105, 106, 107, - 200, 201, 100, 101, 102, 103, 104, 105, 106, 107, + 200, 201, 202, 203, 204, 205, 206, 207, }; - const auto window = http_detail::find_pflash_query_window( - rendered, query, /*search_end=*/12); + const auto window = http_detail::find_pflash_query_window(rendered, query); TEST_ASSERT(window.valid()); TEST_ASSERT(window.tokens == 8); TEST_ASSERT(window.end == 10); - TEST_ASSERT((int)rendered.size() - window.end == 10); + TEST_ASSERT((int)rendered.size() - window.end == 8); } -TEST_CASE(ServerUnitFixture, test_pflash_scorer_accepts_responses_string_input) { - ToolMemory tool_memory; - const auto messages = normalize_chat_messages( - json("Which token is the answer?"), ApiFormat::RESPONSES, tool_memory); +TEST_CASE(ServerUnitFixture, test_pflash_scorer_maps_last_128_user_tokens) { + std::vector query; + for (int token = 0; token < 160; ++token) { + query.push_back(1000 + token); + } + std::vector rendered{1, 2}; + rendered.insert(rendered.end(), query.begin(), query.end()); + rendered.insert(rendered.end(), {200, 201, 202, 203}); - TEST_ASSERT( - http_detail::pflash_user_query_text(messages) == - "Which token is the answer?"); + const auto window = + http_detail::find_pflash_query_window(rendered, query, 128); + + TEST_ASSERT(window.valid()); + TEST_ASSERT(window.tokens == 128); + TEST_ASSERT(window.end == 162); + TEST_ASSERT((int)rendered.size() - window.end == 4); } TEST_CASE(ServerUnitFixture, test_pflash_query_mapping_tolerates_one_bpe_boundary_token) { @@ -227,25 +244,597 @@ TEST_CASE(ServerUnitFixture, test_pflash_query_mapping_tolerates_one_bpe_boundar 1, 2, 999, 11, 12, 13, 14, 15, 16, 17, 200, 201, }; - const auto window = http_detail::find_pflash_query_window( - rendered, query, /*search_end=*/10); + const auto window = http_detail::find_pflash_query_window(rendered, query); TEST_ASSERT(window.valid()); TEST_ASSERT(window.tokens == 7); TEST_ASSERT(window.end == 10); } +TEST_CASE(ServerUnitFixture, test_pflash_query_mapping_stays_before_latest_user_boundary) { + const std::vector query{10, 11, 12, 13, 14, 15, 16, 17}; + const std::vector rendered{ + 1, 2, 999, 11, 12, 13, 14, 15, 16, 17, + 200, 201, 10, 11, 12, 13, 14, 15, 16, 17, 202, + }; + // The sentinel changes the token at the user end, but the later template + // and assistant tokens remain stable. The common suffix establishes the + // semantic boundary after the original complete user suffix. + const std::vector sentinel_rendered{ + 1, 2, 999, 11, 12, 13, 14, 15, 16, 9999, + 200, 201, 10, 11, 12, 13, 14, 15, 16, 17, 202, + }; + + const int search_end = http_detail::pflash_query_search_end_from_sentinel( + rendered, sentinel_rendered); + const auto bounded = + http_detail::find_pflash_query_window(rendered, query, 8, search_end); + const auto unbounded = http_detail::find_pflash_query_window(rendered, query); + + TEST_ASSERT(search_end == 10); + TEST_ASSERT(bounded.valid()); + TEST_ASSERT(bounded.tokens == 7); + TEST_ASSERT(bounded.end == 10); + TEST_ASSERT(unbounded.valid()); + TEST_ASSERT(unbounded.tokens == 8); + TEST_ASSERT(unbounded.end == 20); + TEST_ASSERT(http_detail::pflash_query_search_end_from_sentinel( + rendered, std::vector{42}) < 0); +} + +TEST_CASE(ServerUnitFixture, test_pflash_explicit_query_matches_before_trailing_content) { + // An explicit query sits before trailing instructions inside the latest + // user content; anchored matching (message tail) must fail, unanchored + // matching must find its latest occurrence and tolerate a leading-token + // boundary difference from tokenizing the query on its own. + const std::vector query{5, 10, 11, 12, 13, 14}; // 5 = boundary drift + const std::vector rendered{ + 1, 2, 3, 10, 11, 12, 13, 14, 300, 301, 302, 303, 304, 305, 306, 307, + }; + const auto anchored = http_detail::find_pflash_query_window(rendered, query, 8, 16, 1); + const auto explicit_window = http_detail::find_pflash_query_window(rendered, query, 8, 16, 1, false); + TEST_ASSERT(!anchored.valid()); + TEST_ASSERT(explicit_window.valid()); + TEST_ASSERT(explicit_window.end == 8); + TEST_ASSERT(explicit_window.tokens == 5); + const auto out_of_window = http_detail::find_pflash_query_window(rendered, query, 8, 16, 9, false); + TEST_ASSERT(!out_of_window.valid()); +} + +TEST_CASE(ServerUnitFixture, test_pflash_explicit_query_tolerates_merged_trailing_token) { + // The query's last token merges with the prompt's following newline + // (" " alone vs " \n" in context): the unanchored match drops that + // trailing token, the untrimmed match still wins when it exists, and + // anchored (message tail) matching keeps its exact-suffix contract. + const std::vector query{10, 11, 12, 13, 14, 77}; // 77 = " " + const std::vector rendered{ + 1, 2, 3, 10, 11, 12, 13, 14, 78, 300, 301, 302, // 78 = " \n" + }; + const auto trimmed = http_detail::find_pflash_query_window(rendered, query, 8, 12, 1, false); + TEST_ASSERT(trimmed.valid()); + TEST_ASSERT(trimmed.end == 8); + TEST_ASSERT(trimmed.tokens == 5); + TEST_ASSERT(trimmed.trailing_trimmed == 1); + const std::vector exact{1, 10, 11, 12, 13, 14, 77, 2, 10, 11, 12, 13, 14, 78}; + const auto untrimmed = http_detail::find_pflash_query_window(exact, query, 8, 14, 0, false); + TEST_ASSERT(untrimmed.valid()); + TEST_ASSERT(untrimmed.end == 7); + TEST_ASSERT(untrimmed.tokens == 6); + TEST_ASSERT(untrimmed.trailing_trimmed == 0); + const auto anchored = http_detail::find_pflash_query_window(rendered, query, 8, 12, 1); + TEST_ASSERT(!anchored.valid()); + const std::vector short_query{10, 11, 12, 77}; + const auto too_short = http_detail::find_pflash_query_window(rendered, short_query, 8, 12, 1, false); + TEST_ASSERT(!too_short.valid()); +} + +TEST_CASE(ServerUnitFixture, test_pflash_bounded_mapping_prefers_latest_shortened_suffix) { + const std::vector query{10, 11, 12, 13, 14, 15, 16, 17}; + const std::vector rendered{ + 1, 10, 11, 12, 13, 14, 15, 16, 17, + 900, 11, 12, 13, 14, 15, 16, 17, 200, 201, + }; + + const auto bounded = + http_detail::find_pflash_query_window(rendered, query, 8, 17); + const auto unbounded = http_detail::find_pflash_query_window(rendered, query, 8); + + TEST_ASSERT(bounded.valid()); + TEST_ASSERT(bounded.end == 17); + TEST_ASSERT(bounded.tokens == 7); + // The compatibility path remains width-first when there is no semantic + // boundary, so the earlier exact duplicate still wins there. + TEST_ASSERT(unbounded.valid()); + TEST_ASSERT(unbounded.end == 9); + TEST_ASSERT(unbounded.tokens == 8); +} + +TEST_CASE(ServerUnitFixture, test_pflash_bounded_mapping_rejects_earlier_duplicate) { + const std::vector query{10, 11, 12, 13, 14, 15, 16, 17}; + const std::vector rendered{ + 1, 10, 11, 12, 13, 14, 15, 16, 17, + 900, 11, 12, 13, 14, 15, 16, 999, 200, + }; + + const auto bounded = + http_detail::find_pflash_query_window(rendered, query, 8, 18); + const auto unbounded = + http_detail::find_pflash_query_window(rendered, query, 8); + + TEST_ASSERT(!bounded.valid()); + TEST_ASSERT(unbounded.valid()); + TEST_ASSERT(unbounded.end == 9); + TEST_ASSERT(unbounded.tokens == 8); +} + +TEST_CASE(ServerUnitFixture, test_pflash_bounded_mapping_stays_inside_content_start) { + const std::vector query{10, 11, 12, 13, 14, 15, 16, 17}; + const std::vector rendered{ + 1, 2, 3, 4, 10, 11, 12, 13, 14, 15, 16, 17, 200, + }; + + const auto bounded = http_detail::find_pflash_query_window( + rendered, query, 8, 12, 5); + + TEST_ASSERT(bounded.valid()); + TEST_ASSERT(bounded.end == 12); + TEST_ASSERT(bounded.tokens == 7); + TEST_ASSERT(bounded.end - bounded.tokens == 5); +} + +TEST_CASE(ServerUnitFixture, test_pflash_maps_content_start_from_leading_sentinel) { + const std::vector rendered{1, 2, 10, 11, 12, 13, 20}; + const std::vector sentinel_rendered{ + 1, 2, 999, 10, 11, 12, 13, 20, + }; + + TEST_ASSERT(http_detail::pflash_query_search_begin_from_sentinel( + rendered, sentinel_rendered) == 2); + TEST_ASSERT(http_detail::pflash_query_search_begin_from_sentinel( + rendered, rendered) < 0); +} + +// ─── Explicit query / required-text span mapping ──────────────────────── +// A minimal GPT-2 byte-BPE tokenizer whose vocab reproduces the runtime +// failure shape: " What" and "What" are distinct tokens, so a standalone +// query encoding cannot match a prompt where the question's first character +// merges with the preceding space. + +static std::string test_gpt2_encode(const std::string & text) { + static const auto fwd = []() { + std::array table{}; + int n = 0; + for (int b = 0; b < 256; ++b) { + const bool printable = + (b >= 33 && b <= 126) || (b >= 161 && b <= 172) || + (b >= 174 && b <= 255); + table[b] = printable ? (uint32_t) b : (uint32_t) (256 + n++); + } + return table; + }(); + std::string out; + for (char ch : text) { + const uint32_t cp = fwd[(uint8_t) ch]; + if (cp < 0x80) { + out.push_back((char) cp); + } else { + out.push_back((char) (0xC0 | (cp >> 6))); + out.push_back((char) (0x80 | (cp & 0x3F))); + } + } + return out; +} + +static std::string write_pflash_bpe_tokenizer_fixture( + const std::vector & raw_tokens, + const std::string & byte_cover) { + std::vector tokens{"<|im_start|>", "<|im_end|>"}; + std::vector types{3, 3}; + const auto add = [&](const std::string & encoded, uint32_t type) { + if (std::find(tokens.begin(), tokens.end(), encoded) == tokens.end()) { + tokens.push_back(encoded); + types.push_back(type); + } + }; + for (const auto & raw : raw_tokens) add(test_gpt2_encode(raw), 1); + for (char ch : byte_cover) add(test_gpt2_encode(std::string(1, ch)), 1); + + std::vector token_ptrs; + for (const auto & token : tokens) token_ptrs.push_back(token.c_str()); + gguf_context * g = gguf_init_empty(); + gguf_set_arr_str(g, "tokenizer.ggml.tokens", token_ptrs.data(), + (int32_t) tokens.size()); + gguf_set_arr_data(g, "tokenizer.ggml.token_type", GGUF_TYPE_UINT32, + types.data(), (int32_t) types.size()); + gguf_set_val_str(g, "tokenizer.ggml.model", "gpt2"); + gguf_set_val_str(g, "tokenizer.ggml.pre", "qwen35"); + gguf_set_val_u32(g, "tokenizer.ggml.bos_token_id", 0); + gguf_set_val_u32(g, "tokenizer.ggml.eos_token_id", 1); + static int fixture_serial = 0; + const std::string path = "/tmp/dflash_test_pflash_bpe_" + + std::to_string(++fixture_serial) + ".gguf"; + gguf_write_to_file(g, path.c_str(), /*only_meta=*/false); + gguf_free(g); + return path; +} + +TEST_CASE(ServerUnitFixture, + test_pflash_decoded_span_covers_bpe_merged_first_token) { + const std::string content = + "See docs.\n\nQuestion: What is the answer?\n Answer:"; + const std::string rendered = "<|im_start|>user\n" + content + + "<|im_end|>\n<|im_start|>assistant\n"; + const std::string path = write_pflash_bpe_tokenizer_fixture( + {"Question", ":", " What", "What", " is", " the", " answer", "?", + "\n", "\n\n", " Answer", "user", "assistant", "See", " docs", "."}, + rendered + "What is the answer?"); + Tokenizer tok; + TEST_ASSERT(tok.load_from_gguf(path.c_str())); + + const auto prompt = tok.encode(rendered); + const std::string needle = "What is the answer?"; + // The standalone query encoding starts with a bare "What" token; the + // prompt merged the preceding space into " What". The id-suffix matcher + // therefore accepts a shortened window that misses the first word — the + // regression this span mapping fixes. + const auto query_ids = tok.encode(needle); + const auto legacy = http_detail::find_pflash_query_window( + prompt, query_ids, 64, -1, 0, /*anchored=*/false); + TEST_ASSERT(legacy.valid()); + TEST_ASSERT(tok.decode({prompt.begin() + (legacy.end - legacy.tokens), + prompt.begin() + legacy.end}) != needle); + + const auto span = http_detail::pflash_decoded_text_span( + tok, prompt, 0, (int) prompt.size(), needle); + TEST_ASSERT(span.begin >= 0); + const std::string covered = tok.decode( + {prompt.begin() + span.begin, prompt.begin() + span.end}); + TEST_ASSERT(covered.find(needle) != std::string::npos); + TEST_ASSERT(span.end == legacy.end); + TEST_ASSERT(span.begin == legacy.end - legacy.tokens - 1); + unlink(path.c_str()); +} + +TEST_CASE(ServerUnitFixture, + test_pflash_decoded_span_prefers_last_occurrence) { + const std::string content = + "Ask: What is up? Then again: What is up?"; + const std::string rendered = "<|im_start|>user\n" + content + + "<|im_end|>\n"; + const std::string path = write_pflash_bpe_tokenizer_fixture( + {"Ask", ":", " What", " is", " up", "?", " Then", " again", + "user", "\n"}, + rendered + "What is up?"); + Tokenizer tok; + TEST_ASSERT(tok.load_from_gguf(path.c_str())); + + const auto prompt = tok.encode(rendered); + const std::string needle = "What is up?"; + const auto span = http_detail::pflash_decoded_text_span( + tok, prompt, 0, (int) prompt.size(), needle); + TEST_ASSERT(span.begin >= 0); + const size_t last = rendered.rfind(needle); + const auto tail = tok.encode(rendered.substr(0, last)); + // The needle's leading space merges into " What" in the prompt, but the + // standalone-encoded prefix keeps it as its own " " token, so the merged + // token sits at tail.size() - 1. + TEST_ASSERT(span.begin == (int) tail.size() - 1); + TEST_ASSERT(tok.decode({prompt.begin() + span.begin, + prompt.begin() + span.end}) + .find(needle) != std::string::npos); + unlink(path.c_str()); +} + +TEST_CASE(ServerUnitFixture, + test_pflash_decoded_span_handles_unicode_and_content_end) { + const std::string content = "Discuss the café Über Alles? now"; + const std::string rendered = "<|im_start|>user\n" + content + "<|im_end|>\n"; + const std::string path = write_pflash_bpe_tokenizer_fixture( + {"Discuss", " the", " café", "café", " Über", " Alles", "?", " now", + "user", "\n"}, + rendered + "café Über Alles?"); + Tokenizer tok; + TEST_ASSERT(tok.load_from_gguf(path.c_str())); + + const auto prompt = tok.encode(rendered); + const std::string needle = "café Über Alles?"; + const auto span = http_detail::pflash_decoded_text_span( + tok, prompt, 0, (int) prompt.size(), needle); + TEST_ASSERT(span.begin >= 0); + TEST_ASSERT(tok.decode({prompt.begin() + span.begin, + prompt.begin() + span.end}) + .find(needle) != std::string::npos); + // A needle that never occurs maps to no span — callers must fail closed. + const auto missing = http_detail::pflash_decoded_text_span( + tok, prompt, 0, (int) prompt.size(), "never-present question"); + TEST_ASSERT(missing.begin < 0); + // An empty needle and an empty range are invalid rather than a + // degenerate zero-width span. + TEST_ASSERT(http_detail::pflash_decoded_text_span( + tok, prompt, 0, (int) prompt.size(), "").begin < 0); + TEST_ASSERT(http_detail::pflash_decoded_text_span( + tok, prompt, 4, 4, needle).begin < 0); + unlink(path.c_str()); +} + +TEST_CASE(ServerUnitFixture, test_pflash_required_parses_string_arrays) { + TEST_ASSERT(parse_pflash_required_from_body({}).empty()); + TEST_ASSERT(parse_pflash_required_from_body( + {{"pflash_required", "not-an-array"}}).empty()); + const auto top = parse_pflash_required_from_body( + {{"pflash_required", + {"answer briefly.", "Question:", 7, nullptr}}}); + TEST_ASSERT(top.size() == 2); + TEST_ASSERT(top[0] == "answer briefly."); + const auto nested = parse_pflash_required_from_body( + {{"extra_body", {{"pflash_required", {"keep me"}}}}}); + TEST_ASSERT(nested.size() == 1 && nested[0] == "keep me"); + // extra_body wins over the top-level field, like session_id. + const auto both = parse_pflash_required_from_body( + {{"pflash_required", {"outer"}}, + {"extra_body", {{"pflash_required", {"inner"}}}}}); + TEST_ASSERT(both.size() == 1 && both[0] == "inner"); +} + +TEST_CASE(ServerUnitFixture, test_pflash_instruction_plan_covers_tools_and_late_developer_roles) { + const std::vector messages{ + {"system", "system instruction", ""}, + {"developer", "leading developer instruction", ""}, + {"user", "document text", ""}, + {"assistant", "prior answer", ""}, + {"developer", "late developer instruction", ""}, + {"tool", "tool result is history", "call-1"}, + {"user", "latest query", ""}, + }; + const auto plan = http_detail::plan_pflash_instruction_messages(messages); + + TEST_ASSERT(plan.instruction_messages == + std::vector({0, 1, 4})); +} + +TEST_CASE(ServerUnitFixture, test_pflash_instruction_plan_handles_empty_instructions) { + const std::vector tool_only{ + {"user", "use the tool", ""}, + }; + const auto tool_plan = + http_detail::plan_pflash_instruction_messages(tool_only); + TEST_ASSERT(tool_plan.instruction_messages.empty()); + + const std::vector empty_instruction{ + {"system", "", ""}, + {"user", "plain query", ""}, + }; + const auto empty_plan = + http_detail::plan_pflash_instruction_messages(empty_instruction); + TEST_ASSERT(empty_plan.instruction_messages.empty()); +} + +TEST_CASE(ServerUnitFixture, test_pflash_tool_span_follows_arbitrary_jinja_placement) { + static const char TPL[] = + "{%- for m in messages -%}{{ m.content }}{%- endfor -%}" + "{%- if tools -%}|TOOLS:{{ tools[0].function.name }}{%- endif -%}"; + const std::vector messages{{"user", "query-first", ""}}; + const std::string tools = + R"([{"type":"function","function":{"name":"late_lookup"}}])"; + const std::string with_tools = render_chat_template_jinja( + TPL, messages, "", "", true, false, tools); + const std::string without_tools = render_chat_template_jinja( + TPL, messages, "", "", true, false, "[]"); + const std::vector original(with_tools.begin(), with_tools.end()); + const std::vector variant(without_tools.begin(), without_tools.end()); + + const PFlashTokenSpan span = + http_detail::pflash_changed_token_span(original, variant); + TEST_ASSERT(span.begin >= (int) messages[0].content.size()); + TEST_ASSERT(span.end == (int) original.size()); + TEST_ASSERT(with_tools.substr( + (size_t) span.begin, (size_t) (span.end - span.begin)).find( + "late_lookup") != std::string::npos); +} + +TEST_CASE(ServerUnitFixture, test_pflash_instruction_span_follows_reordered_jinja_message) { + static const char TPL[] = + "{%- for m in messages -%}{%- if m.role == 'user' -%}" + "<{{ m.role }}>{{ m.content }}" + "{%- endif -%}{%- endfor -%}" + "{%- for m in messages -%}{%- if m.role == 'system' -%}" + "<{{ m.role }}>{{ m.content }}" + "{%- endif -%}{%- endfor -%}"; + const std::vector messages{ + {"system", "retain this rule", ""}, + {"user", "question first", ""}, + }; + const std::string rendered = render_chat_template_jinja( + TPL, messages, "", "", true, false); + auto without_system = messages; + without_system.erase(without_system.begin()); + const std::string variant_rendered = render_chat_template_jinja( + TPL, without_system, "", "", true, false); + const std::vector original(rendered.begin(), rendered.end()); + const std::vector variant( + variant_rendered.begin(), variant_rendered.end()); + + const PFlashTokenSpan span = + http_detail::pflash_changed_token_span(original, variant); + TEST_ASSERT(span.begin > (int) messages[1].content.size()); + const std::string retained = rendered.substr( + (size_t) span.begin, (size_t) (span.end - span.begin)); + TEST_ASSERT(retained.find("") != std::string::npos); + TEST_ASSERT(retained.find(messages[0].content) != std::string::npos); + TEST_ASSERT(retained.find("") != std::string::npos); +} + +TEST_CASE(ServerUnitFixture, test_pflash_instruction_spans_are_canonicalized) { + const auto spans = http_detail::canonicalize_pflash_token_spans( + {{12, 20}, {0, 4}, {3, 8}, {20, 24}}); + TEST_ASSERT(spans == std::vector({{0, 8}, {12, 24}})); +} + +TEST_CASE(ServerUnitFixture, test_pflash_qwen_tool_prefix_boundary_covers_schema) { + const std::vector messages{{"user", "find weather", ""}}; + const std::string tools = + R"([{"type":"function","function":{"name":"lookup_weather"}}])"; + const std::string sentinel = "__PFLASH_BEGIN_02C47F91__"; + const std::string rendered = render_chat_template( + messages, ChatFormat::QWEN3, true, false, tools); + auto marked_messages = messages; + marked_messages[0].content = sentinel + marked_messages[0].content; + const std::string marked = render_chat_template( + marked_messages, ChatFormat::QWEN3, true, false, tools); + const std::vector rendered_ids(rendered.begin(), rendered.end()); + const std::vector marked_ids(marked.begin(), marked.end()); + + const int prefix_end = http_detail::pflash_query_search_begin_from_sentinel( + rendered_ids, marked_ids); + TEST_ASSERT(prefix_end > 0); + TEST_ASSERT(rendered.substr(0, (size_t) prefix_end).find("lookup_weather") != + std::string::npos); + TEST_ASSERT(rendered.substr(0, (size_t) prefix_end).find("find weather") == + std::string::npos); +} + +TEST_CASE(ServerUnitFixture, test_pflash_qwen_late_developer_span_covers_role_envelope) { + const std::vector messages{ + {"user", std::string(930, 'u'), ""}, + {"assistant", "history", ""}, + {"developer", std::string(180, 'd'), ""}, + {"user", "latest query", ""}, + }; + const std::string rendered = render_chat_template( + messages, ChatFormat::QWEN3, true, false); + auto without_developer = messages; + without_developer.erase(without_developer.begin() + 2); + const std::string without_developer_rendered = render_chat_template( + without_developer, ChatFormat::QWEN3, true, false); + const std::vector ids(rendered.begin(), rendered.end()); + const std::vector variant( + without_developer_rendered.begin(), without_developer_rendered.end()); + + const PFlashTokenSpan span = + http_detail::pflash_changed_token_span(ids, variant); + const size_t content_begin = rendered.find(messages[2].content); + TEST_ASSERT(content_begin != std::string::npos); + TEST_ASSERT(span.begin >= 0); + TEST_ASSERT((size_t) span.begin < content_begin); + TEST_ASSERT((size_t) span.end > content_begin + messages[2].content.size()); + TEST_ASSERT(span.begin < 1024); + TEST_ASSERT(span.end > 1024); + TEST_ASSERT(rendered.substr( + (size_t) span.begin, + (size_t) (span.end - span.begin)).find("developer") != + std::string::npos); +} + +TEST_CASE(ServerUnitFixture, test_pflash_responses_string_tails_only_raw_content) { + ToolMemory tool_memory; + const auto normalized = normalize_chat_messages( + json("raw completion input"), ApiFormat::RESPONSES, tool_memory); + TEST_ASSERT(normalized.size() == 1); + TEST_ASSERT(normalized[0].role == "user"); + TEST_ASSERT(normalized[0].content == "raw completion input"); + + const std::vector rendered{ + 1, 2, 10, 11, 12, 13, 14, 200, 201, + }; + const auto window = http_detail::pflash_tail_query_window( + rendered, 128, /*query_end=*/7, /*query_begin=*/2); + TEST_ASSERT(window.valid()); + TEST_ASSERT(window.end == 7); + TEST_ASSERT(window.tokens == 5); + TEST_ASSERT(window.end - window.tokens == 2); +} + +TEST_CASE(ServerUnitFixture, test_compress_result_fails_closed_on_empty_ids) { + const auto failed = ModelBackend::CompressResult::from_compressed_ids({}); + const auto succeeded = + ModelBackend::CompressResult::from_compressed_ids({11, 12}); + + TEST_ASSERT(!failed.ok); + TEST_ASSERT(failed.compressed_ids.empty()); + TEST_ASSERT(succeeded.ok); + TEST_ASSERT(succeeded.compressed_ids == std::vector({11, 12})); +} + +TEST_CASE(ServerUnitFixture, test_pflash_parser_fingerprint_has_fixed_encoding) { + const std::vector ids{1, -2, 2147483647}; + TEST_ASSERT(http_detail::pflash_token_fingerprint(ids) == + "45409a0b0f44c5fd"); +} + +TEST_CASE(ServerUnitFixture, test_pflash_tail_query_window) { + const auto empty = http_detail::pflash_tail_query_window({}, 128); + TEST_ASSERT(!empty.valid()); + + const std::vector short_prompt{1, 2, 3}; + const auto short_tail = + http_detail::pflash_tail_query_window(short_prompt, 128); + TEST_ASSERT(short_tail.valid()); + TEST_ASSERT(short_tail.end == 3); + TEST_ASSERT(short_tail.tokens == 3); + + std::vector long_prompt(200); + const auto capped_tail = + http_detail::pflash_tail_query_window(long_prompt, 128); + TEST_ASSERT(capped_tail.valid()); + TEST_ASSERT(capped_tail.end == 200); + TEST_ASSERT(capped_tail.tokens == 128); + const auto bounded_tail = + http_detail::pflash_tail_query_window(long_prompt, 128, 150); + TEST_ASSERT(bounded_tail.valid()); + TEST_ASSERT(bounded_tail.end == 150); + TEST_ASSERT(bounded_tail.tokens == 128); + TEST_ASSERT(!http_detail::pflash_tail_query_window(long_prompt, 0).valid()); + TEST_ASSERT(!http_detail::pflash_tail_query_window(long_prompt, 128, 0).valid()); + TEST_ASSERT(!http_detail::pflash_tail_query_window(long_prompt, 128, 201).valid()); +} + +TEST_CASE(ServerUnitFixture, test_pflash_normalizes_multipart_latest_user_for_reverse_lookup) { + ToolMemory tool_memory; + const json messages = json::array({ + {{"role", "user"}, {"content", "older user"}}, + {{"role", "assistant"}, {"content", "assistant before latest"}}, + {{"role", "user"}, {"content", json::array({ + {{"type", "input_text"}, {"text", "input-"}}, + {{"type", "input_image"}, {"image_url", "ignored"}}, + {{"type", "text"}, {"text", "text"}} + })}}, + {{"role", "assistant"}, {"content", "assistant after latest"}}, + {{"role", "tool"}, {"content", "tool after latest"}} + }); + + const auto normalized = normalize_chat_messages( + messages, ApiFormat::OPENAI_CHAT, tool_memory); + const auto latest_user = std::find_if( + normalized.rbegin(), normalized.rend(), + [](const ChatMessage & message) { return message.role == "user"; }); + TEST_ASSERT(latest_user != normalized.rend()); + if (latest_user != normalized.rend()) { + TEST_ASSERT(latest_user->content == "input-text"); + } +} + +TEST_CASE(ServerUnitFixture, test_pflash_selection_cache_and_continuation_policy) { + TEST_ASSERT(http_detail::pflash_full_cache_restore_allowed(false)); + TEST_ASSERT(!http_detail::pflash_full_cache_restore_allowed(true)); + TEST_ASSERT(!http_detail::pflash_continuation_must_fail_closed(false)); + TEST_ASSERT(http_detail::pflash_continuation_must_fail_closed(true)); +} + +TEST_CASE(ServerUnitFixture, test_pflash_target_token_ceiling_floors) { + TEST_ASSERT(http_detail::pflash_target_token_ceiling(7, 0.5) == 3); + TEST_ASSERT(http_detail::pflash_target_token_ceiling(120000, 16384.0 / 120000.0) == 16384); + TEST_ASSERT(http_detail::pflash_target_token_ceiling(-1, 0.5) < 0); +} + TEST_CASE(ServerUnitFixture, test_pflash_query_mapping_rejects_weak_punctuation_match) { const std::vector query{10, 11, 12, 13, 14, 15, 16, 17}; const std::vector rendered{1, 2, 15, 16, 17, 200, 201}; - TEST_ASSERT(!http_detail::find_pflash_query_window( - rendered, query, /*search_end=*/7).valid()); + TEST_ASSERT(!http_detail::find_pflash_query_window(rendered, query).valid()); const std::vector short_query{30, 31, 32}; const std::vector short_rendered{1, 30, 31, 32, 200}; const auto short_window = - http_detail::find_pflash_query_window( - short_rendered, short_query, /*search_end=*/4); + http_detail::find_pflash_query_window(short_rendered, short_query); TEST_ASSERT(short_window.valid()); TEST_ASSERT(short_window.tokens == 3); TEST_ASSERT(short_window.end == 4); @@ -279,12 +868,24 @@ TEST_CASE(ServerUnitFixture, test_qwen35_pflash_rejects_missing_query_window) { } TEST_CASE(ServerUnitFixture, test_pflash_ipc_rejects_unsupported_query_widths) { - TEST_ASSERT(valid_pflash_score_query_tokens(1)); - TEST_ASSERT(valid_pflash_score_query_tokens(8)); - TEST_ASSERT(!valid_pflash_score_query_tokens(0)); - TEST_ASSERT(!valid_pflash_score_query_tokens(9)); - TEST_ASSERT(!valid_pflash_score_query_tokens( - (std::numeric_limits::max)())); + const auto accepted = [](int score_query_tokens) { + std::string line; + std::string error; + const bool formatted = format_pflash_drafter_ipc_compress_command( + 0.5f, 16, score_query_tokens, "/tmp/pflash_ids.bin", line, error); + if (!formatted) return false; + PFlashDrafterIpcCompressCommand parsed; + return parse_pflash_drafter_ipc_compress_command(line, parsed, error) && + parsed.score_query_tokens == score_query_tokens; + }; + // The explicit scorer query sizes its own window, so widths above the + // former eight-token cap round-trip; non-positive widths still fail closed. + TEST_ASSERT(accepted(1)); + TEST_ASSERT(accepted(8)); + TEST_ASSERT(accepted(128)); + TEST_ASSERT(!accepted(0)); + TEST_ASSERT(!accepted(-1)); + TEST_ASSERT(!accepted((std::numeric_limits::min)())); } TEST_CASE(ServerUnitFixture, test_pflash_query_capture_splits_across_chunks) { @@ -5612,6 +6213,55 @@ TEST_CASE(ServerUnitFixture, } #endif +struct MockPflashCompressBackend : MockBackend { + int compress_calls = 0; + CompressRequest last_request; + + CompressResult compress(const CompressRequest & request) override { + ++compress_calls; + last_request = request; + return CompressResult::from_compressed_ids(request.input_ids); + } +}; + +TEST_CASE(ServerUnitFixture, test_pflash_default_raw_text_maps_user_query) { + luce_test::ScopedEnvVar mode{"PFLASH_SELECT_MODE", nullptr}; + luce_test::ScopedEnvVar chunk{"PFLASH_SELECT_CHUNK_SIZE", nullptr}; + luce_test::ScopedEnvVar query{"PFLASH_SELECT_QUERY_TOKENS", nullptr}; + luce_test::ScopedEnvVar parser{"PFLASH_SELECT_QUERY_PARSER", nullptr}; + luce_test::ScopedEnvVar top_p{"PFLASH_SELECT_TOP_P", nullptr}; + + const std::string tokenizer_path = + write_deepseek_marker_tokenizer_fixture(); + Tokenizer tokenizer; + TEST_ASSERT(tokenizer.load_from_gguf(tokenizer_path.c_str())); + + auto backend_owner = std::make_unique(); + MockPflashCompressBackend & backend = *backend_owner; + LuceEngine engine(std::move(backend_owner)); + ServerConfig config; + 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::RESPONSES; + request.messages = "x"; + request.prompt_tokens = tokenizer.encode("x"); + + const std::string error = + HttpServerTestAccess::apply_pflash_compression(server, request); + TEST_ASSERT_MSG(error.empty(), error); + } + + TEST_ASSERT(backend.compress_calls == 1); + TEST_ASSERT(backend.last_request.score_query_end == 1); + TEST_ASSERT(backend.last_request.score_query_tokens == 1); + unlink(tokenizer_path.c_str()); +} + struct MockBatchCompressBackend : MockBackend { int compress_calls = 0; @@ -8006,6 +8656,19 @@ TEST_CASE(ServerUnitFixture, test_flowkv_session_keep_ratio_override) { TEST_ASSERT(std::fabs(static_ratio - configured_ratio) < 1e-6f); TEST_ASSERT(std::fabs(adaptive_ratio - 0.09f) < 1e-6f); + + // A session without feedback keeps the configured (curve) ratio, and its + // first feedback adapts from that ratio rather than the fixed default: + // 0.1875 (the 16K real-use budget) minus one small step at high acceptance. + const float curve_ratio = 0.1875f; + TEST_ASSERT(std::fabs(http_detail::resolve_pflash_keep_ratio( + curve_ratio, "fresh", sessions) - curve_ratio) < 1e-6f); + sessions.update("fresh", 0.95f, curve_ratio); + TEST_ASSERT(std::fabs(http_detail::resolve_pflash_keep_ratio( + curve_ratio, "fresh", sessions) - (curve_ratio - 0.01f)) < 1e-6f); + // Seeds are clamped to the controller's range like every later step. + sessions.update("wide", 0.50f, 0.30f); + TEST_ASSERT(std::fabs(sessions.get_keep_ratio("wide") - 0.20f) < 1e-6f); } // ═══════════════════════════════════════════════════════════════════════ From 38144573ef5e48e5badf8f63ce02d99c527bf63e Mon Sep 17 00:00:00 2001 From: Graffioh Date: Sat, 19 Sep 2026 14:12:43 +0000 Subject: [PATCH 03/26] test(pflash): selection and IPC coverage; README MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Cover the strict selector's budget, ordering and required-span rules in test_pflash_selection.cpp, and the compress2/compress3 wire format — its round trips, the legacy quantized parser and its malformed inputs — in test_pflash_drafter_ipc.cpp, including a case that shows a remote drafter selects exactly what the local one would. Add parser coverage for the explicit `pflash_query` request field at both the top level and inside extra_body. Document the Qwen3.5-0.8B drafter, the two optional GGUF files and the PFLASH_SELECT_* contract in the server README. Co-Authored-By: Claude Fable 5.1 --- server/CMakeLists.txt | 2 + server/README.md | 42 +- server/test/test_bandit_integration.cpp | 11 + server/test/test_pflash_drafter_ipc.cpp | 187 +++++++ server/test/test_pflash_selection.cpp | 675 ++++++++++++++++++++++++ 5 files changed, 916 insertions(+), 1 deletion(-) create mode 100644 server/test/test_pflash_drafter_ipc.cpp create mode 100644 server/test/test_pflash_selection.cpp diff --git a/server/CMakeLists.txt b/server/CMakeLists.txt index e5b1b01f4..90501557b 100644 --- a/server/CMakeLists.txt +++ b/server/CMakeLists.txt @@ -1947,6 +1947,8 @@ if(DFLASH27B_TESTS) test/test_drafter_tail_capture_guard.cpp test/test_drafter_warm_path_regression.cpp test/test_qwen3_buffer_plan.cpp + test/test_pflash_drafter_ipc.cpp + test/test_pflash_selection.cpp test/test_model_test_paths.cpp test/test_gguf_mmap.cpp test/test_kv_quant.cpp diff --git a/server/README.md b/server/README.md index 77eff1619..cf11ed246 100644 --- a/server/README.md +++ b/server/README.md @@ -371,12 +371,52 @@ the whole request's device footprint. `/status/json` reports | `--prefill-threshold ` | `32000` | Token threshold used by auto mode. | | `--prefill-keep-ratio ` | `0.05` | Fraction of source tokens kept. | | `--prefill-curve T:R [T:R ...]` | none | Piecewise keep-ratio curve; overrides the flat ratio. | -| `--prefill-drafter ` | none | PFlash drafter GGUF. | +| `--prefill-drafter ` | none | PFlash drafter GGUF: Qwen3-0.6B, or Qwen3.5-0.8B when the file name contains `qwen3.5`/`qwen35`. | | `--prefill-skip-park` | off | Keep target and decode draft resident while PFlash runs. | | `--prefill-upstream-base ` | none | Enable compression-proxy mode. | | `--prefill-upstream-key ` | none | Bearer token for the upstream. | | `--prefill-upstream-model ` | none | Model name forwarded upstream. | +With a Qwen3.5-0.8B drafter and strict budget selection +(`PFLASH_SELECT_MODE=budget_only`, `PFLASH_SELECT_CHUNK_SIZE`, +`PFLASH_SELECT_QUERY_TOKENS`), the drafter runs only its first fifteen +blocks and scores the context with block 15's NoPE Q/K projections, the same +attention-mass scorer the Qwen3-0.6B block-13 head uses. Its 262K native +context covers inputs the Qwen3-0.6B drafter cannot score within its 32K +window. `PFLASH_SCORING_HEAD_GGUF` accepts a trained block-15 head +(schema `qwen3_5_0_8b_nope_qk_mass_v1`); `PFLASH_QWEN35_LEGACY_SCORER=1` +restores the previous all-layer running-max scorer. The Qwen3.5 attention +runs dense (`ggml_flash_attn_ext`); the block-sparse FlashPrefill kernels +still dispatch head dimension 128 only. + +`PFLASH_SEGMENT_PROBE_GGUF` loads a segment probe (schema +`qwen3_5_0_8b_segment_probe_v1`): a 264K-parameter network on the same block-14 +tap that scores every token for "a new unit of text starts here". With it +loaded, the context is cut at every boundary above the probe's threshold +(the query start and instruction-span edges are always cut; minimum and +maximum segment lengths come from the GGUF metadata) and the strict selector +ranks the resulting whole functions, classes, files or paragraphs by mass +density, skipping segments that do not fit the remaining budget, so a kept +piece is never a definition cut in half. It falls back to fixed chunks when +the probe finds fewer than four boundaries in a context. +`PFLASH_SELECT_SEGMENTS=auto|fixed|probe` and +`PFLASH_SELECT_SCORE=auto|sum|density` override the defaults (auto = +probe segments and density when a probe is loaded, fixed chunks and mass sum +otherwise); the compression trace records `segmentation`, `candidate_score` +and the segment spans. + +The per-session adaptive keep ratio applies to this path unchanged: a request +carrying a `session_id` retains the session's ratio, the strict selector fills +its token budget from it, and the ratio is updated from the smoothed DFlash +acceptance rate after every turn where speculative decoding ran (below 75% +acceptance retain more, above 85% retain less, 0.5-1 point per turn, bounded +to 2.5-20%; `server/src/server/adaptive_keep_ratio.h`). A new session starts +from the configured ratio for its prompt length (`--prefill-keep-ratio` or +`--prefill-curve`), so the controller adapts around the real-use budget +instead of a fixed 10%. Acceptance is a proxy for compression quality: it +does not detect a dropped answer document directly, so the ratio curve and +the retention benchmarks remain the quality reference. + ### Reasoning and MoE controls | Option | Default | Purpose | diff --git a/server/test/test_bandit_integration.cpp b/server/test/test_bandit_integration.cpp index 0453db70c..9f9299f32 100644 --- a/server/test/test_bandit_integration.cpp +++ b/server/test/test_bandit_integration.cpp @@ -101,3 +101,14 @@ TEST_CASE(BanditIntegrationFixture, non_string_session_id_array_extra_body) { std::string sid = parse_session_id_from_body(body); CHECK(sid.empty()); } + +TEST_CASE(BanditIntegrationFixture, pflash_query_top_level_and_extra_body) { + json top = {{"pflash_query", "Which function has the deliberate error?"}}; + CHECK(parse_pflash_query_from_body(top) == "Which function has the deliberate error?"); + json nested = {{"extra_body", {{"pflash_query", "What is the major tributary of the Rhine?"}}}}; + CHECK(parse_pflash_query_from_body(nested) == "What is the major tributary of the Rhine?"); + json absent = {{"messages", json::array()}}; + CHECK(parse_pflash_query_from_body(absent).empty()); + json wrong_type = {{"pflash_query", 7}}; + CHECK(parse_pflash_query_from_body(wrong_type).empty()); +} diff --git a/server/test/test_pflash_drafter_ipc.cpp b/server/test/test_pflash_drafter_ipc.cpp new file mode 100644 index 000000000..e1d45518d --- /dev/null +++ b/server/test/test_pflash_drafter_ipc.cpp @@ -0,0 +1,187 @@ +#include "CppUnitTestFramework.hpp" + +#include "common/pflash_drafter_ipc.h" +#include "common/model_backend.h" +#include "qwen3/pflash_selection.h" + +#include +#include + +using namespace dflash::common; + +namespace { + +struct PFlashDrafterIpcFixture : CppUnitTestFramework::CommonFixture { + using CppUnitTestFramework::CommonFixture::CommonFixture; +}; + +} // namespace + +TEST_CASE(PFlashDrafterIpcFixture, compress2_round_trips_paper_ratio_budget) { + const float keep = 16384.0f / 120000.0f; + std::string line; + std::string error; + + REQUIRE(format_pflash_drafter_ipc_compress_command( + keep, 119900, 128, "/tmp/pflash ids.bin", line, error)); + REQUIRE(error.empty()); + + PFlashDrafterIpcCompressCommand parsed; + REQUIRE(parse_pflash_drafter_ipc_compress_command(line, parsed, error)); + REQUIRE(error.empty()); + REQUIRE(!parsed.legacy_quantized_ratio); + REQUIRE(parsed.keep_ratio == keep); + REQUIRE((int) std::floor(120000.0 * (double) parsed.keep_ratio) == 16384); + REQUIRE(parsed.score_query_end == 119900); + REQUIRE(parsed.score_query_tokens == 128); + REQUIRE(parsed.required_instruction_spans.empty()); + REQUIRE(parsed.path == "/tmp/pflash ids.bin"); +} + +TEST_CASE(PFlashDrafterIpcFixture, compress3_round_trips_ordered_instruction_spans) { + const float keep = 16384.0f / 120000.0f; + const std::vector instructions{{0, 384}, {4096, 4352}}; + std::string line; + std::string error; + + REQUIRE(format_pflash_drafter_ipc_compress_command( + keep, 119900, 128, instructions, + "/tmp/pflash ids.bin", line, error)); + REQUIRE(error.empty()); + REQUIRE(line.rfind("compress3 ", 0) == 0); + + PFlashDrafterIpcCompressCommand parsed; + REQUIRE(parse_pflash_drafter_ipc_compress_command(line, parsed, error)); + REQUIRE(error.empty()); + REQUIRE(parsed.keep_ratio == keep); + REQUIRE(parsed.score_query_end == 119900); + REQUIRE(parsed.score_query_tokens == 128); + REQUIRE(parsed.required_instruction_spans == instructions); + REQUIRE(parsed.path == "/tmp/pflash ids.bin"); +} + +TEST_CASE(PFlashDrafterIpcFixture, compress3_matches_the_local_selector_contract) { + ModelBackend::CompressRequest local; + local.input_ids.resize(32); + local.keep_ratio = 0.5f; + local.score_query_end = 32; + local.score_query_tokens = 4; + local.required_instruction_spans = {{0, 4}, {12, 14}}; + + std::string line; + std::string error; + REQUIRE(format_pflash_drafter_ipc_compress_command( + local.keep_ratio, local.score_query_end, local.score_query_tokens, + local.required_instruction_spans, "/tmp/ids.bin", line, error)); + PFlashDrafterIpcCompressCommand remote; + REQUIRE(parse_pflash_drafter_ipc_compress_command(line, remote, error)); + REQUIRE(remote.keep_ratio == local.keep_ratio); + REQUIRE(remote.score_query_end == local.score_query_end); + REQUIRE(remote.score_query_tokens == local.score_query_tokens); + REQUIRE(remote.required_instruction_spans == local.required_instruction_spans); + + const auto select = [&] (const std::vector & spans) { + std::vector candidates; + constexpr double scores[]{0.0, 9.0, 1.0, 0.0, 2.0, 3.0, 4.0, 0.0}; + for (int chunk = 0; chunk < 8; ++chunk) { + const int begin = chunk * 4; + const int end = begin + 4; + candidates.push_back({ + (size_t) chunk, begin, end, scores[chunk], + dflash::qwen3::pflash_chunk_is_structurally_required( + begin, end, 28, 32, 32, spans), + }); + } + return dflash::qwen3::select_pflash_candidates( + candidates, {16, 0.95}, + dflash::qwen3::PFlashSelectionMode::BudgetOnly); + }; + const auto local_result = select(local.required_instruction_spans); + const auto remote_result = select(remote.required_instruction_spans); + REQUIRE(local_result.ok); + REQUIRE(remote_result.ok); + REQUIRE(local_result.ordinals == remote_result.ordinals); + REQUIRE(local_result.retained_tokens == remote_result.retained_tokens); + REQUIRE(local_result.stop == remote_result.stop); + + const std::vector out_of_range{{0, 33}}; + std::string local_error; + std::string remote_error; + REQUIRE(!dflash::qwen3::validate_pflash_instruction_spans( + out_of_range, (int) local.input_ids.size(), local_error)); + REQUIRE(format_pflash_drafter_ipc_compress_command( + local.keep_ratio, local.score_query_end, local.score_query_tokens, + out_of_range, "/tmp/ids.bin", line, error)); + REQUIRE(parse_pflash_drafter_ipc_compress_command(line, remote, error)); + REQUIRE(!dflash::qwen3::validate_pflash_instruction_spans( + remote.required_instruction_spans, (int) local.input_ids.size(), + remote_error)); + REQUIRE(local_error == remote_error); +} + +TEST_CASE(PFlashDrafterIpcFixture, legacy_x1000_parser_is_supported_but_quantized) { + PFlashDrafterIpcCompressCommand parsed; + std::string error; + + REQUIRE(parse_pflash_drafter_ipc_compress_command( + "compress 137 119900 128 /tmp/pflash_ids.bin", parsed, error)); + REQUIRE(error.empty()); + REQUIRE(parsed.legacy_quantized_ratio); + REQUIRE(parsed.keep_ratio == 0.137f); + REQUIRE((int) std::floor(120000.0 * (double) parsed.keep_ratio) > 16384); + REQUIRE(parsed.score_query_end == 119900); + REQUIRE(parsed.score_query_tokens == 128); + REQUIRE(parsed.path == "/tmp/pflash_ids.bin"); +} + +TEST_CASE(PFlashDrafterIpcFixture, parser_rejects_malformed_values) { + const char * bad_lines[] = { + "compress2 nan 10 8 /tmp/ids.bin", + "compress2 inf 10 8 /tmp/ids.bin", + "compress2 -0.1 10 8 /tmp/ids.bin", + "compress2 1.1 10 8 /tmp/ids.bin", + "compress2 0.5 10 0 /tmp/ids.bin", + "compress2 0.5 10 8", + "compress 1001 10 8 /tmp/ids.bin", + "compress -1 10 8 /tmp/ids.bin", + "compress 500 10 0 /tmp/ids.bin", + "compress3 0.5 10 8 -1 /tmp/ids.bin", + "compress3 0.5 10 8 1 0 /tmp/ids.bin", + "compress3 0.5 10 8 1 4 4 /tmp/ids.bin", + "compress3 0.5 10 8 2 0 4 3 6 /tmp/ids.bin", + "compress3 0.5 10 8 65 /tmp/ids.bin", + "unknown 0.5 10 8 /tmp/ids.bin", + }; + + for (const char * line : bad_lines) { + PFlashDrafterIpcCompressCommand parsed; + std::string error; + REQUIRE(!parse_pflash_drafter_ipc_compress_command( + line, parsed, error)); + REQUIRE(!error.empty()); + } +} + +TEST_CASE(PFlashDrafterIpcFixture, formatter_fails_closed_on_invalid_values) { + struct BadFormatInput { + float keep_ratio; + int score_query_tokens; + const char * path; + }; + const BadFormatInput bad_inputs[] = { + {-0.1f, 8, "/tmp/ids.bin"}, + {1.1f, 8, "/tmp/ids.bin"}, + {0.5f, 0, "/tmp/ids.bin"}, + {0.5f, 8, ""}, + }; + + for (const auto & input : bad_inputs) { + std::string line; + std::string error; + REQUIRE(!format_pflash_drafter_ipc_compress_command( + input.keep_ratio, 10, input.score_query_tokens, + input.path, line, error)); + REQUIRE(line.empty()); + REQUIRE(!error.empty()); + } +} diff --git a/server/test/test_pflash_selection.cpp b/server/test/test_pflash_selection.cpp new file mode 100644 index 000000000..5539bc640 --- /dev/null +++ b/server/test/test_pflash_selection.cpp @@ -0,0 +1,675 @@ +#include "CppUnitTestFramework.hpp" + +#include "qwen3/pflash_selection.h" +#include "qwen3/qwen3_drafter_model.h" +#include "scoped_env.h" + +#include +#include +#include +#include +#include +#include +#include + +using namespace dflash::qwen3; + +namespace { + +constexpr const char * kModeEnv = "PFLASH_SELECT_MODE"; +constexpr const char * kChunkEnv = "PFLASH_SELECT_CHUNK_SIZE"; +constexpr const char * kQueryEnv = "PFLASH_SELECT_QUERY_TOKENS"; +constexpr const char * kQueryParserEnv = "PFLASH_SELECT_QUERY_PARSER"; +constexpr const char * kTopPEnv = "PFLASH_SELECT_TOP_P"; + +struct CleanPFlashEnv { + luce_test::ScopedEnvVar mode{kModeEnv, nullptr}; + luce_test::ScopedEnvVar chunk{kChunkEnv, nullptr}; + luce_test::ScopedEnvVar query{kQueryEnv, nullptr}; + luce_test::ScopedEnvVar query_parser{kQueryParserEnv, nullptr}; + luce_test::ScopedEnvVar top_p{kTopPEnv, nullptr}; +}; + +void set_env(const char * name, const char * value) { +#if defined(_WIN32) + _putenv_s(name, value ? value : ""); +#else + if (value) { + setenv(name, value, 1); + } else { + unsetenv(name); + } +#endif +} + +PFlashSelectionCandidate candidate( + size_t ordinal, + int begin, + int end, + double score, + bool mandatory = false) { + return {ordinal, begin, end, score, mandatory}; +} + +void require_ordinals( + const PFlashSelectionResult & result, + const std::vector & expected) { + if (result.ordinals.size() != expected.size()) { + throw std::runtime_error("unexpected selected ordinal count"); + } + for (size_t index = 0; index < expected.size(); ++index) { + if (result.ordinals[index] != expected[index]) { + throw std::runtime_error("unexpected selected ordinal"); + } + } +} + +PFlashSelectionConfig resolve_or_fail(int input_tokens, int legacy_chunk) { + PFlashSelectionConfig config; + std::string error; + if (!resolve_pflash_selection( + input_tokens, legacy_chunk, config, error)) { + throw std::runtime_error(error); + } + if (!error.empty()) throw std::runtime_error(error); + return config; +} + +struct PFlashSelectionFixture : CppUnitTestFramework::CommonFixture { + using CppUnitTestFramework::CommonFixture::CommonFixture; +}; + +} // namespace + +TEST_CASE(PFlashSelectionFixture, structural_suffix_only_chunk_is_mandatory_and_charged) { + constexpr int input_tokens = 101; + constexpr int query_begin = 80; + constexpr int query_end = 90; + REQUIRE(!pflash_chunk_is_structurally_required( + 0, 64, query_begin, query_end, input_tokens)); + REQUIRE(pflash_chunk_is_structurally_required( + 64, 96, query_begin, query_end, input_tokens)); + REQUIRE(pflash_chunk_is_structurally_required( + 96, 101, query_begin, query_end, input_tokens)); + + const std::vector candidates{ + candidate(0, 0, 64, 100.0), + candidate(1, 64, 96, 0.0, + pflash_chunk_is_structurally_required( + 64, 96, query_begin, query_end, input_tokens)), + candidate(2, 96, 101, 0.0, + pflash_chunk_is_structurally_required( + 96, 101, query_begin, query_end, input_tokens)), + }; + const auto result = select_pflash_candidates( + candidates, {37, 0.95}, PFlashSelectionMode::BudgetOnly); + + REQUIRE(result.ok); + REQUIRE(result.retained_tokens == 37); + REQUIRE(result.stop == PFlashSelectionStop::BudgetReached); + require_ordinals(result, {1, 2}); +} + +TEST_CASE(PFlashSelectionFixture, instruction_overlap_is_mandatory_without_changing_optional_ranking) { + constexpr int input_tokens = 120000; + constexpr int query_begin = 119872; + constexpr int query_end = 120000; + const std::vector instructions{ + {0, 384}, + {4096, 4352}, + }; + std::string error; + REQUIRE(validate_pflash_instruction_spans( + instructions, input_tokens, error)); + REQUIRE(error.empty()); + REQUIRE(pflash_chunk_is_structurally_required( + 0, 1024, query_begin, query_end, input_tokens, instructions)); + REQUIRE(pflash_chunk_is_structurally_required( + 4096, 5120, query_begin, query_end, input_tokens, instructions)); + REQUIRE(!pflash_chunk_is_structurally_required( + 1024, 2048, query_begin, query_end, input_tokens, instructions)); + + const std::vector candidates{ + candidate(0, 0, 1024, 0.0, true), + candidate(1, 1024, 2048, 10.0), + candidate(2, 2048, 3072, 1.0), + candidate(3, 4096, 5120, 0.0, true), + candidate(4, 119872, 120000, 0.0, true), + }; + const auto result = select_pflash_candidates( + candidates, {3200, 0.95}, PFlashSelectionMode::BudgetOnly); + + REQUIRE(result.ok); + REQUIRE(result.stop == PFlashSelectionStop::BudgetReached); + REQUIRE(result.retained_tokens == 3200); + require_ordinals(result, {0, 1, 3, 4}); +} + +TEST_CASE(PFlashSelectionFixture, invalid_instruction_spans_fail_closed) { + constexpr int input_tokens = 32; + const std::vector> invalid{ + {{-1, 2}}, + {{4, 4}}, + {{4, 3}}, + {{0, 4}, {3, 6}}, + {{8, 12}, {0, 4}}, + {{0, 33}}, + }; + for (const auto & spans : invalid) { + std::string error; + REQUIRE(!validate_pflash_instruction_spans( + spans, input_tokens, error)); + REQUIRE(!error.empty()); + } + std::vector too_many(65, {0, 1}); + std::string error; + REQUIRE(!validate_pflash_instruction_spans( + too_many, input_tokens, error)); + REQUIRE(!error.empty()); +} + +TEST_CASE(PFlashSelectionFixture, cumulative_top_p_is_scale_invariant_and_keeps_crossing_chunk) { + const std::vector base{ + candidate(0, 0, 4, 6.0), + candidate(1, 4, 8, 3.0), + candidate(2, 8, 12, 1.0), + }; + auto scaled = base; + for (auto & item : scaled) item.score *= 100.0; + + const PFlashSelectionPolicy policy{12, 0.8}; + const auto base_result = select_pflash_candidates( + base, policy, PFlashSelectionMode::CumulativeTopP); + const auto scaled_result = select_pflash_candidates( + scaled, policy, PFlashSelectionMode::CumulativeTopP); + + REQUIRE(base_result.ok); + REQUIRE(scaled_result.ok); + REQUIRE(base_result.stop == PFlashSelectionStop::TopPReached); + REQUIRE(scaled_result.stop == PFlashSelectionStop::TopPReached); + require_ordinals(base_result, {0, 1}); + require_ordinals(scaled_result, {0, 1}); + REQUIRE(std::abs(base_result.retained_mass - 0.9) < 1e-12); + REQUIRE(std::abs(scaled_result.retained_mass - 0.9) < 1e-12); +} + +TEST_CASE(PFlashSelectionFixture, zero_and_negative_scores_use_equal_mass_and_ordinal_ties) { + const std::vector candidates{ + candidate(2, 8, 12, -8.0), + candidate(0, 0, 4, -2.0), + candidate(1, 4, 8, 0.0), + }; + + const auto result = select_pflash_candidates( + candidates, {12, 0.5}, PFlashSelectionMode::CumulativeTopP); + + REQUIRE(result.ok); + REQUIRE(result.stop == PFlashSelectionStop::TopPReached); + require_ordinals(result, {0, 1}); + REQUIRE(std::abs(result.retained_mass - 2.0 / 3.0) < 1e-12); +} + +TEST_CASE(PFlashSelectionFixture, budget_only_disables_only_the_mass_stop) { + const std::vector candidates{ + candidate(0, 0, 4, 6.0), + candidate(1, 4, 8, 3.0), + candidate(2, 8, 12, 1.0), + }; + + const auto result = select_pflash_candidates( + candidates, {12, 0.1}, PFlashSelectionMode::BudgetOnly); + + REQUIRE(result.ok); + REQUIRE(result.stop == PFlashSelectionStop::CandidatesExhausted); + require_ordinals(result, {0, 1, 2}); + REQUIRE(result.retained_tokens == 12); + REQUIRE(std::abs(result.retained_mass - 1.0) < 1e-12); +} + +TEST_CASE(PFlashSelectionFixture, mandatory_scores_do_not_enter_optional_mass) { + const std::vector candidates{ + candidate(0, 0, 1, 1.0e30, true), + candidate(1, 1, 2, 6.0), + candidate(2, 2, 3, 3.0), + candidate(3, 3, 4, 1.0), + }; + + const auto result = select_pflash_candidates( + candidates, {4, 0.8}, PFlashSelectionMode::CumulativeTopP); + + REQUIRE(result.ok); + REQUIRE(result.stop == PFlashSelectionStop::TopPReached); + require_ordinals(result, {0, 1, 2}); + REQUIRE(std::abs(result.retained_mass - 0.9) < 1e-12); +} + +TEST_CASE(PFlashSelectionFixture, real_ranges_charge_a_short_final_chunk) { + const std::vector candidates{ + candidate(0, 0, 4, 2.0), + candidate(1, 4, 6, 1.0), + }; + + const auto result = select_pflash_candidates( + candidates, {6, 1.0}, PFlashSelectionMode::BudgetOnly); + + REQUIRE(result.ok); + REQUIRE(result.stop == PFlashSelectionStop::CandidatesExhausted); + require_ordinals(result, {0, 1}); + REQUIRE(result.retained_tokens == 6); +} + +TEST_CASE(PFlashSelectionFixture, mandatory_overflow_has_a_distinct_failure) { + const std::vector candidates{ + candidate(0, 0, 4, 0.0, true), + candidate(1, 4, 6, 1.0), + }; + + const auto result = select_pflash_candidates( + candidates, {3, 0.95}, PFlashSelectionMode::CumulativeTopP); + + REQUIRE(!result.ok); + REQUIRE(result.stop == PFlashSelectionStop::MandatoryQueryExceedsBudget); + REQUIRE(result.ordinals.empty()); + REQUIRE(result.retained_tokens == 0); + REQUIRE(!result.error.empty()); +} + +TEST_CASE(PFlashSelectionFixture, budget_stop_does_not_skip_to_a_smaller_candidate) { + const std::vector candidates{ + candidate(0, 0, 2, 0.0, true), + candidate(1, 2, 6, 10.0), + candidate(2, 6, 9, 1.0), + }; + + const auto result = select_pflash_candidates( + candidates, {5, 1.0}, PFlashSelectionMode::BudgetOnly); + + REQUIRE(result.ok); + REQUIRE(result.stop == PFlashSelectionStop::BudgetReached); + require_ordinals(result, {0}); + REQUIRE(result.retained_tokens == 2); +} + +TEST_CASE(PFlashSelectionFixture, output_ordinals_are_in_source_order) { + const std::vector candidates{ + candidate(42, 8, 12, 10.0), + candidate(7, 0, 4, 1.0), + candidate(99, 4, 8, 99.0, true), + }; + + const auto result = select_pflash_candidates( + candidates, {12, 1.0}, PFlashSelectionMode::BudgetOnly); + + REQUIRE(result.ok); + require_ordinals(result, {7, 99, 42}); +} + +TEST_CASE(PFlashSelectionFixture, invalid_selector_inputs_fail_closed) { + const PFlashSelectionPolicy valid_policy{16, 0.95}; + const auto nan_result = select_pflash_candidates( + {candidate(0, 0, 4, std::numeric_limits::quiet_NaN())}, + valid_policy, + PFlashSelectionMode::CumulativeTopP); + REQUIRE(!nan_result.ok); + REQUIRE(nan_result.stop == PFlashSelectionStop::InvalidInput); + + const auto inf_result = select_pflash_candidates( + {candidate(0, 0, 4, std::numeric_limits::infinity())}, + valid_policy, + PFlashSelectionMode::CumulativeTopP); + REQUIRE(!inf_result.ok); + + const auto overlap_result = select_pflash_candidates( + {candidate(0, 0, 4, 1.0), candidate(1, 3, 6, 2.0)}, + valid_policy, + PFlashSelectionMode::CumulativeTopP); + REQUIRE(!overlap_result.ok); + + const auto duplicate_result = select_pflash_candidates( + {candidate(0, 0, 4, 1.0), candidate(0, 4, 8, 2.0)}, + valid_policy, + PFlashSelectionMode::CumulativeTopP); + REQUIRE(!duplicate_result.ok); + + REQUIRE(!select_pflash_candidates( + {}, {0, 0.95}, PFlashSelectionMode::BudgetOnly).ok); + REQUIRE(!select_pflash_candidates( + {}, {1, 0.0}, PFlashSelectionMode::BudgetOnly).ok); + REQUIRE(!select_pflash_candidates( + {}, {1, 1.01}, PFlashSelectionMode::BudgetOnly).ok); +} + +TEST_CASE(PFlashSelectionFixture, resolver_defaults_to_legacy_arguments) { + CleanPFlashEnv env; + const auto config = resolve_or_fail(500, 32); + + REQUIRE(!config.configured); + REQUIRE(!config.selection_active); + REQUIRE(config.mode == PFlashSelectionMode::Legacy); + REQUIRE(config.query_parser == PFlashQueryParser::SemanticUser); + REQUIRE(config.chunk_size == 32); + REQUIRE(config.query_tokens == 8); + REQUIRE(std::abs(config.top_p - 0.95) < 1e-12); +} + +TEST_CASE(PFlashSelectionFixture, resolver_applies_chunk_and_query_without_enabling_selection) { + CleanPFlashEnv env; + set_env(kChunkEnv, "64"); + set_env(kQueryEnv, "32"); + + const auto config = resolve_or_fail(500, 32); + REQUIRE(config.configured); + REQUIRE(!config.selection_active); + REQUIRE(config.mode == PFlashSelectionMode::Legacy); + REQUIRE(config.chunk_size == 64); + REQUIRE(config.query_tokens == 32); +} + +TEST_CASE(PFlashSelectionFixture, any_selection_environment_is_observable_before_resolution) { + CleanPFlashEnv env; + REQUIRE(!has_pflash_selection_environment()); + + const std::pair values[] = { + {kModeEnv, "budget_only"}, + {kChunkEnv, "1024"}, + {kQueryEnv, "128"}, + {kQueryParserEnv, "arbitrary_tail"}, + {kTopPEnv, "0.95"}, + }; + for (const auto & [name, value] : values) { + set_env(name, value); + REQUIRE(has_pflash_selection_environment()); + PFlashSelectionConfig config; + std::string error; + REQUIRE(resolve_pflash_selection(120000, 32, config, error)); + REQUIRE(config.configured); + set_env(name, nullptr); + } + + set_env(kQueryEnv, ""); + REQUIRE(has_pflash_selection_environment()); + PFlashSelectionConfig config; + std::string error; + REQUIRE(!resolve_pflash_selection(120000, 32, config, error)); +} + +TEST_CASE(PFlashSelectionFixture, resolver_selects_explicit_query_parser) { + CleanPFlashEnv env; + set_env(kQueryParserEnv, "arbitrary_tail"); + const auto arbitrary = resolve_or_fail(120000, 32); + REQUIRE(arbitrary.configured); + REQUIRE(arbitrary.query_parser == PFlashQueryParser::ArbitraryTail); + + set_env(kQueryParserEnv, "latest_user"); + const auto latest_user = resolve_or_fail(120000, 32); + REQUIRE(latest_user.query_parser == PFlashQueryParser::SemanticUser); +} + +TEST_CASE(PFlashSelectionFixture, strict_mode_uses_length_schedule_without_chunk_override) { + CleanPFlashEnv env; + set_env(kModeEnv, "top_p"); + + REQUIRE(resolve_or_fail(499, 32).chunk_size == 128); + REQUIRE(resolve_or_fail(500, 32).chunk_size == 512); + REQUIRE(resolve_or_fail(2999, 32).chunk_size == 512); + const auto large = resolve_or_fail(3000, 32); + REQUIRE(large.chunk_size == 1024); + REQUIRE(large.configured); + REQUIRE(large.selection_active); + REQUIRE(large.mode == PFlashSelectionMode::CumulativeTopP); + + set_env(kModeEnv, "budget_only"); + const auto budget = resolve_or_fail(3000, 32); + REQUIRE(budget.selection_active); + REQUIRE(budget.mode == PFlashSelectionMode::BudgetOnly); + + set_env(kChunkEnv, "256"); + REQUIRE(resolve_or_fail(3000, 32).chunk_size == 256); +} + +TEST_CASE(PFlashSelectionFixture, resolver_rejects_invalid_environment_values) { + CleanPFlashEnv env; + struct InvalidValue { + const char * name; + const char * value; + }; + const InvalidValue invalid_values[] = { + {kModeEnv, "legacy"}, + {kModeEnv, "TOP_P"}, + {kChunkEnv, "0"}, + {kChunkEnv, "12x"}, + {kQueryEnv, "0"}, + {kQueryEnv, "513"}, + {kQueryParserEnv, "last_128"}, + {kTopPEnv, "0"}, + {kTopPEnv, "1.01"}, + {kTopPEnv, "nan"}, + }; + + for (const auto & invalid : invalid_values) { + set_env(invalid.name, invalid.value); + PFlashSelectionConfig config; + std::string error; + REQUIRE(!resolve_pflash_selection(500, 32, config, error)); + REQUIRE(!error.empty()); + set_env(invalid.name, nullptr); + } + + set_env(kTopPEnv, "1"); + REQUIRE(std::abs(resolve_or_fail(500, 32).top_p - 1.0) < 1e-12); +} + +TEST_CASE(PFlashSelectionFixture, mode_and_stop_names_are_stable) { + REQUIRE(std::string(pflash_selection_mode_name(PFlashSelectionMode::Legacy)) == "legacy"); + REQUIRE(std::string(pflash_selection_mode_name(PFlashSelectionMode::BudgetOnly)) == "budget_only"); + REQUIRE(std::string(pflash_selection_mode_name(PFlashSelectionMode::CumulativeTopP)) == "top_p"); + REQUIRE(std::string(pflash_selection_stop_name(PFlashSelectionStop::TopPReached)) == "top_p_reached"); + REQUIRE(std::string(pflash_selection_stop_name(PFlashSelectionStop::BudgetReached)) == "budget_reached"); + REQUIRE(std::string(pflash_selection_stop_name(PFlashSelectionStop::CandidatesExhausted)) == "candidates_exhausted"); + REQUIRE(std::string(pflash_selection_stop_name(PFlashSelectionStop::InvalidInput)) == "invalid_input"); + REQUIRE(std::string(pflash_selection_stop_name(PFlashSelectionStop::MandatoryQueryExceedsBudget)) == + "mandatory_query_exceeds_budget"); + REQUIRE(std::string(pflash_query_parser_name(PFlashQueryParser::SemanticUser)) == "latest_user"); + REQUIRE(std::string(pflash_query_parser_name(PFlashQueryParser::ArbitraryTail)) == "arbitrary_tail"); +} + +TEST_CASE(PFlashSelectionFixture, scoring_head_token_mass_averages_heads_and_queries) { + // ggml layout [n_keys=3, n_queries=2, n_heads=2]: key index fastest. + const std::vector probs = { + 0.2f, 0.3f, 0.5f, // head 0, query 0 + 0.6f, 0.4f, 0.0f, // head 0, query 1 + 0.0f, 0.0f, 1.0f, // head 1, query 0 + 1.0f, 0.0f, 0.0f, // head 1, query 1 + }; + std::vector mass; + dflash::common::scoring_head_mean_token_mass(probs.data(), 3, 2, 2, mass); + REQUIRE(mass.size() == 3u); + CHECK(std::fabs(mass[0] - 0.45f) < 1e-6f); + CHECK(std::fabs(mass[1] - 0.175f) < 1e-6f); + CHECK(std::fabs(mass[2] - 0.375f) < 1e-6f); + double total = 0.0; + for (float value : mass) total += value; + CHECK(std::fabs(total - 1.0) < 1e-6); + + dflash::common::scoring_head_mean_token_mass(probs.data(), 0, 2, 2, mass); + CHECK(mass.empty()); +} + +// ═════════════════════════════════════════════════ +// Segment probe: variable-length candidates from per-token boundary scores +// ═════════════════════════════════════════════════ + +TEST_CASE(PFlashSelectionFixture, probe_segments_cut_at_scores_and_forced_edges) { + // 20 tokens; boundary scores above 0.9 at 5 and 12; the query starts at 17. + std::vector scores(20, 0.0f); + scores[5] = 0.95f; + scores[12] = 0.99f; + scores[13] = 0.97f; // too close to 12 for min_segment 3: dropped + const auto spans = pflash_probe_segments(scores, 20, 0.9f, 3, 100, {17}); + REQUIRE(spans.size() == 4); + REQUIRE(spans[0].begin == 0 && spans[0].end == 5); + REQUIRE(spans[1].begin == 5 && spans[1].end == 12); + REQUIRE(spans[2].begin == 12 && spans[2].end == 17); + REQUIRE(spans[3].begin == 17 && spans[3].end == 20); + // A forced cut is kept even inside the minimum distance. + const auto forced = pflash_probe_segments(scores, 20, 0.9f, 3, 100, {13}); + REQUIRE(forced.size() == 4 && forced[2].begin == 12 && forced[2].end == 13); +} + +TEST_CASE(PFlashSelectionFixture, probe_segments_split_oversized_spans_at_best_interior_score) { + std::vector scores(30, 0.0f); + scores[9] = 0.4f; // below threshold, but the best interior candidate + scores[20] = 0.3f; + const auto spans = pflash_probe_segments(scores, 30, 0.9f, 2, 12, {}); + // [0,30) exceeds 12: split at 9 -> [0,9), then [9,30) exceeds 12: split at 20 -> [9,20), [20,30) + REQUIRE(spans.size() == 3); + REQUIRE(spans[0].end == 9 && spans[1].end == 20 && spans[2].end == 30); + // No interior score: fixed grid of max_segment. + const auto grid = pflash_probe_segments(std::vector(30, 0.0f), 30, 0.9f, 2, 12, {}); + REQUIRE(grid.size() == 3 && grid[0].end == 12 && grid[1].end == 24 && grid[2].end == 30); + // Invalid input fails closed. + REQUIRE(pflash_probe_segments(scores, 40, 0.9f, 2, 12, {}).empty()); + REQUIRE(pflash_probe_segments(scores, 30, 0.9f, 0, 12, {}).empty()); +} + +TEST_CASE(PFlashSelectionFixture, probe_segments_oversize_split_keeps_off_the_near_edge) { + // The distance guard searches [begin + max/2, begin + max]: a score near + // the span start cannot produce a tiny leading fragment. + std::vector scores(30, 0.0f); + scores[3] = 0.5f; // below threshold and inside max_segment/2: ignored + const auto spans = pflash_probe_segments(scores, 30, 0.9f, 2, 12, {}); + REQUIRE(spans.size() == 3); + REQUIRE(spans[0].begin == 0 && spans[0].end == 12); + REQUIRE(spans[1].begin == 12 && spans[1].end == 24); + REQUIRE(spans[2].begin == 24 && spans[2].end == 30); + // An interior score in the guarded window still wins over the grid. + std::vector interior(30, 0.0f); + interior[10] = 0.5f; + const auto guarded = pflash_probe_segments(interior, 30, 0.9f, 2, 12, {}); + REQUIRE(guarded.size() == 3); + REQUIRE(guarded[0].end == 10); + // The emitted piece never exceeds max_segment even when the best score + // sits at the window edge. + std::vector edge(40, 0.0f); + edge[11] = 0.7f; // at begin + max_segment - 1: still inside the window + edge[30] = 0.9f; + const auto capped = pflash_probe_segments(edge, 40, 0.9f, 2, 12, {30}); + REQUIRE(!capped.empty()); + for (const auto & span : capped) { + REQUIRE(span.end - span.begin <= 12); + } +} + +TEST_CASE(PFlashSelectionFixture, probe_segments_split_scores_feed_only_the_interior_argmax) { + // The sub-unit score vector steers the oversize interior split without + // touching the boundary threshold or merge floor. + std::vector boundary(30, 0.0f); + boundary[10] = 0.5f; // unit score in the window — must NOT be used + std::vector subunit(30, 0.0f); + subunit[8] = 0.6f; // sub-unit score wins the argmax instead + const auto spans = pflash_probe_segments( + boundary, 30, 0.9f, 2, 12, {}, subunit); + REQUIRE(spans.size() == 3); + REQUIRE(spans[0].begin == 0 && spans[0].end == 8); + REQUIRE(spans[1].begin == 8 && spans[1].end == 20); + REQUIRE(spans[2].begin == 20 && spans[2].end == 30); + // An empty split vector falls back to unit scores (v1 artifacts). + const auto fallback = pflash_probe_segments(boundary, 30, 0.9f, 2, 12, {}, {}); + REQUIRE(fallback.size() == 3); + REQUIRE(fallback[0].end == 10); + // Sub-unit scores never create boundaries below the unit threshold. + std::vector flat(30, 0.0f); + std::vector hot(30, 0.0f); + hot[5] = 1.0f; // inside max/2: only reachable via the argmax + const auto no_new_cuts = pflash_probe_segments(flat, 30, 0.9f, 2, 12, {}, hot); + REQUIRE(no_new_cuts.size() == 3); + for (const auto & span : no_new_cuts) { + REQUIRE(span.end - span.begin <= 12); + } +} + +TEST_CASE(PFlashSelectionFixture, skip_oversized_keeps_filling_with_smaller_segments) { + // Ranked by score: a 600-token segment first, then two 200-token ones. Budget 500. + const std::vector candidates = { + candidate(0, 0, 600, 0.9), + candidate(1, 600, 800, 0.5), + candidate(2, 800, 1000, 0.4), + }; + const auto strict = select_pflash_candidates( + candidates, PFlashSelectionPolicy{500, 0.95, false}, PFlashSelectionMode::BudgetOnly); + REQUIRE(strict.ok && strict.ordinals.empty() && + strict.stop == PFlashSelectionStop::BudgetReached); + const auto skipping = select_pflash_candidates( + candidates, PFlashSelectionPolicy{500, 0.95, true}, PFlashSelectionMode::BudgetOnly); + REQUIRE(skipping.ok); + require_ordinals(skipping, {1, 2}); + REQUIRE(skipping.retained_tokens == 400); +} + +TEST_CASE(PFlashSelectionFixture, segmentation_and_score_environment_resolve_or_fail) { + CleanPFlashEnv clean; + luce_test::ScopedEnvVar segments{"PFLASH_SELECT_SEGMENTS", nullptr}; + luce_test::ScopedEnvVar select{"PFLASH_SELECT_SCORE", nullptr}; + set_env(kModeEnv, "budget_only"); + auto config = resolve_or_fail(4096, 1024); + REQUIRE(config.segmentation == PFlashSegmentation::Auto); + REQUIRE(config.candidate_score == PFlashCandidateScore::Auto); + set_env("PFLASH_SELECT_SEGMENTS", "probe"); + set_env("PFLASH_SELECT_SCORE", "density"); + config = resolve_or_fail(4096, 1024); + REQUIRE(config.segmentation == PFlashSegmentation::Probe); + REQUIRE(config.candidate_score == PFlashCandidateScore::Density); + set_env("PFLASH_SELECT_SEGMENTS", "sentences"); + PFlashSelectionConfig invalid; + std::string error; + REQUIRE(!resolve_pflash_selection(4096, 1024, invalid, error)); + REQUIRE(error.find("PFLASH_SELECT_SEGMENTS") != std::string::npos); +} + +// ═════════════════════════════════════════════════ +// Two-scorer split selection +// ═════════════════════════════════════════════════ + +TEST_CASE(PFlashSelectionFixture, split_selection_fills_head_share_then_other_scorer_without_duplicates) { + // Six 100-token chunks; the head ranks 0 > 1 > 2 ..., the other scorer ranks 5 > 4 > 3 ...; chunk 2 mandatory. + std::vector head, other; + for (size_t i = 0; i < 6; ++i) { + const int begin = (int) i * 100; + head.push_back(candidate(i, begin, begin + 100, 6.0 - (double) i, i == 2)); + other.push_back(candidate(i, begin, begin + 100, (double) i, i == 2)); + } + // Budget 400, head fraction 0.5: pass 1 keeps mandatory 2 and the head's top 0 (200 tokens); + // pass 2 fills 200 tokens from the other scorer's order skipping 2 and 0 -> 5, 4. + const auto result = select_pflash_split(head, other, PFlashSelectionPolicy{400, 0.95, false}, 0.5, PFlashSelectionMode::BudgetOnly); + REQUIRE(result.ok); + require_ordinals(result, {0, 2, 4, 5}); + REQUIRE(result.retained_tokens == 400); + // Mismatched lists fail closed. + std::vector shifted = other; + shifted[1].begin += 1; + REQUIRE(!select_pflash_split(head, shifted, PFlashSelectionPolicy{400, 0.95, false}, 0.5, PFlashSelectionMode::BudgetOnly).ok); + REQUIRE(!select_pflash_split(head, other, PFlashSelectionPolicy{400, 0.95, false}, 1.5, PFlashSelectionMode::BudgetOnly).ok); + // A tiny head share still charges the mandatory span once and the other scorer gets the rest. + const auto tiny = select_pflash_split(head, other, PFlashSelectionPolicy{300, 0.95, false}, 0.01, PFlashSelectionMode::BudgetOnly); + REQUIRE(tiny.ok); + require_ordinals(tiny, {2, 4, 5}); +} + +TEST_CASE(PFlashSelectionFixture, scorer_and_split_environment_resolve_or_fail) { + CleanPFlashEnv clean; + luce_test::ScopedEnvVar scorer{"PFLASH_SELECT_SCORER", nullptr}; + luce_test::ScopedEnvVar split{"PFLASH_SELECT_SPLIT", nullptr}; + set_env(kModeEnv, "budget_only"); + auto config = resolve_or_fail(4096, 1024); + REQUIRE(config.scorer == PFlashScorer::Head); + set_env("PFLASH_SELECT_SCORER", "split"); + set_env("PFLASH_SELECT_SPLIT", "0.6"); + config = resolve_or_fail(4096, 1024); + REQUIRE(config.scorer == PFlashScorer::Split); + REQUIRE(std::fabs(config.split_fraction - 0.6) < 1e-9); + set_env("PFLASH_SELECT_SPLIT", "1.0"); + PFlashSelectionConfig invalid; + std::string error; + REQUIRE(!resolve_pflash_selection(4096, 1024, invalid, error)); + REQUIRE(error.find("PFLASH_SELECT_SPLIT") != std::string::npos); +} From 9e2ed0454b410544e822e65dfb0bcbe8b09b30a4 Mon Sep 17 00:00:00 2001 From: Graffioh Date: Sun, 20 Sep 2026 10:43:52 +0000 Subject: [PATCH 04/26] refactor(pflash): move Qwen3.5 drafter to src/pflash, drop Qwen3-0.6B scorer src/qwen3/ had become a catch-all mixing three things: the Qwen3-0.6B standalone inference backend, the legacy Qwen3-0.6B PFlash drafter, and the Qwen3.5-0.8B PFlash scorer that replaced it. Split along the real boundary: - src/pflash/ now owns the whole PFlash pipeline: pflash_drafter (was qwen3_drafter), pflash_compress (was qwen3_drafter_common), pflash_selection, qwen35_drafter/qwen35_loader, kvflash_drafter_scorer (was qwen3_kvflash_scorer), anchor_scan/anchor_params. - src/qwen3/ keeps only standalone Qwen3-0.6B inference, with the drafter terminology removed: qwen3_model.h (Qwen3Weights, load/free_qwen3_model). - The Qwen3-0.6B PFlash path is deleted outright: qwen3_graph.cpp, qwen3_buffer_plan.h, score_range.h, DrafterArch dispatch, the 0.6B scoring-head loader, and the tests that covered them. Qwen3.5-0.8B is now the only scorer; the optional drafter_arch command arg is still accepted but ignored. - PFlash selection symbols move from dflash::qwen3 to dflash::pflash. - score_query_end < 0 ("legacy tail window") is still on the CompressRequest/IPC wire formats; backends and the IPC daemon now translate it to input_ids.size() at the boundary since the qwen35 scorer requires an explicit end. - Default/located drafter paths, help text, docs and scripts point at Qwen3.5-0.8B-BF16.gguf. Generated with [Devin](https://devin.ai) Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- optimizations/kvflash/DESIGN.md | 6 +- optimizations/kvflash/README.md | 6 +- optimizations/pflash/README.md | 32 +- optimizations/pflash/pflash/dflash_client.py | 2 +- optimizations/pflash/tests/bench_niah_cpp.py | 8 +- optimizations/pflash/tests/niah_gen.py | 2 +- server/CMakeLists.txt | 36 +- server/README.md | 18 +- server/docs/DS4.md | 2 +- server/docs/ENVIRONMENT.md | 10 +- server/docs/SPEC_PREFILL.md | 13 +- server/docs/laguna_integration_plan.md | 4 +- server/scripts/laguna_pflash_niah.py | 8 +- server/scripts/phase_split_dual_gpu.py | 6 +- server/scripts/quality_ab_simple.py | 2 +- server/scripts/quality_humaneval_plus.py | 2 +- server/scripts/test_full_compress_cache.py | 4 +- server/src/common/kvflash_pager.h | 10 +- server/src/common/kvflash_scorer.h | 4 +- .../src/common/pflash_drafter_ipc_daemon.cpp | 8 +- server/src/common/score_range.h | 31 - server/src/deepseek4/deepseek4_backend.cpp | 8 +- server/src/deepseek4/deepseek4_backend.h | 2 +- server/src/flashprefill.h | 2 +- server/src/gemma4/gemma4_backend.cpp | 12 +- server/src/gemma4/gemma4_backend.h | 4 +- .../src/gemma4/gemma4_layer_split_adapter.cpp | 4 +- .../src/gemma4/gemma4_layer_split_adapter.h | 2 +- server/src/laguna/laguna_backend.cpp | 13 +- server/src/laguna/laguna_backend.h | 4 +- .../src/laguna/laguna_layer_split_adapter.cpp | 4 +- .../src/laguna/laguna_layer_split_adapter.h | 2 +- server/src/{qwen3 => pflash}/anchor_params.h | 0 server/src/{qwen3 => pflash}/anchor_scan.cpp | 4 +- server/src/{qwen3 => pflash}/anchor_scan.h | 4 +- .../kvflash_drafter_scorer.cpp} | 39 +- .../kvflash_drafter_scorer.h} | 14 +- .../pflash_compress.cpp} | 40 +- server/src/pflash/pflash_compress.h | 137 +++ server/src/pflash/pflash_drafter.cpp | 162 +++ .../pflash_drafter.h} | 46 +- .../{qwen3 => pflash}/pflash_selection.cpp | 4 +- .../src/{qwen3 => pflash}/pflash_selection.h | 4 +- .../src/{qwen3 => pflash}/qwen35_drafter.cpp | 34 +- server/src/{qwen3 => pflash}/qwen35_drafter.h | 26 +- .../src/{qwen3 => pflash}/qwen35_loader.cpp | 20 +- server/src/qwen3/qwen3_backend.cpp | 22 +- server/src/qwen3/qwen3_backend.h | 12 +- server/src/qwen3/qwen3_buffer_plan.h | 29 - server/src/qwen3/qwen3_drafter.cpp | 463 -------- server/src/qwen3/qwen3_drafter_common.h | 77 -- server/src/qwen3/qwen3_drafter_model.h | 151 --- server/src/qwen3/qwen3_graph.cpp | 1023 ----------------- server/src/qwen3/qwen3_loader.cpp | 121 +- server/src/qwen3/qwen3_model.h | 74 ++ server/src/qwen35/qwen35_backend.cpp | 14 +- server/src/qwen35/qwen35_backend.h | 2 +- .../src/qwen35/qwen35_layer_split_adapter.cpp | 14 +- .../src/qwen35/qwen35_layer_split_adapter.h | 2 +- server/src/server/http_server.cpp | 26 +- server/src/server/http_server.h | 2 +- server/src/server/server_main.cpp | 2 +- server/test/bench_laguna_pflash.cpp | 10 +- server/test/pflash_daemon.cpp | 11 +- server/test/smoke_qwen3_forward.cpp | 16 +- server/test/test_anchor_params.cpp | 2 +- server/test/test_anchor_transitive.cpp | 30 +- server/test/test_dflash.cpp | 33 +- .../test_drafter_early_exit_score_range.cpp | 90 -- .../test/test_drafter_tail_capture_guard.cpp | 118 -- .../test_drafter_warm_path_regression.cpp | 170 --- server/test/test_kvflash.cpp | 16 +- server/test/test_pflash_drafter_ipc.cpp | 14 +- server/test/test_pflash_selection.cpp | 6 +- server/test/test_qwen3_buffer_plan.cpp | 71 -- server/test/test_server_unit.cpp | 34 +- 76 files changed, 744 insertions(+), 2716 deletions(-) delete mode 100644 server/src/common/score_range.h rename server/src/{qwen3 => pflash}/anchor_params.h (100%) rename server/src/{qwen3 => pflash}/anchor_scan.cpp (98%) rename server/src/{qwen3 => pflash}/anchor_scan.h (95%) rename server/src/{qwen3/qwen3_kvflash_scorer.cpp => pflash/kvflash_drafter_scorer.cpp} (88%) rename server/src/{qwen3/qwen3_kvflash_scorer.h => pflash/kvflash_drafter_scorer.h} (83%) rename server/src/{qwen3/qwen3_drafter_common.cpp => pflash/pflash_compress.cpp} (89%) create mode 100644 server/src/pflash/pflash_compress.h create mode 100644 server/src/pflash/pflash_drafter.cpp rename server/src/{qwen3/qwen3_drafter.h => pflash/pflash_drafter.h} (58%) rename server/src/{qwen3 => pflash}/pflash_selection.cpp (99%) rename server/src/{qwen3 => pflash}/pflash_selection.h (98%) rename server/src/{qwen3 => pflash}/qwen35_drafter.cpp (97%) rename server/src/{qwen3 => pflash}/qwen35_drafter.h (80%) rename server/src/{qwen3 => pflash}/qwen35_loader.cpp (96%) delete mode 100644 server/src/qwen3/qwen3_buffer_plan.h delete mode 100644 server/src/qwen3/qwen3_drafter.cpp delete mode 100644 server/src/qwen3/qwen3_drafter_common.h delete mode 100644 server/src/qwen3/qwen3_drafter_model.h delete mode 100644 server/src/qwen3/qwen3_graph.cpp create mode 100644 server/src/qwen3/qwen3_model.h delete mode 100644 server/test/test_drafter_early_exit_score_range.cpp delete mode 100644 server/test/test_drafter_tail_capture_guard.cpp delete mode 100644 server/test/test_drafter_warm_path_regression.cpp delete mode 100644 server/test/test_qwen3_buffer_plan.cpp diff --git a/optimizations/kvflash/DESIGN.md b/optimizations/kvflash/DESIGN.md index c4d349a05..bf1b51535 100644 --- a/optimizations/kvflash/DESIGN.md +++ b/optimizations/kvflash/DESIGN.md @@ -106,7 +106,7 @@ FA span traffic is bandwidth-realistic: ## Full LSA loop (drafter as Memory Indexer) — measured Test run F implements the paper's complete inference paradigm with the -pflash drafter (Qwen3-0.6B, `/opt/lucebox/models/drafter/`) standing in +pflash drafter (Qwen3.5-0.8B, `/opt/lucebox/models/drafter/`) standing in for the trained indexer: prompt (2048) larger than the pool (1024) so prefill itself evicts, then every τ=64 decoded tokens the drafter rescores the full sequence (tail attention = indexer query, chunk means @@ -176,7 +176,7 @@ The pool is wired into the qwen35 backend behind `--kvflash ` pool (live LRU eviction mid-request). Coherent story end to end, 36.9 tok/s, clean finish. Second request (per-request pager reset) ok. 2. WITH pflash: `--kvflash 2048 --prefill-compression always - --prefill-threshold 256 --prefill-drafter `. Compression + --prefill-threshold 256 --prefill-drafter `. Compression 1468 -> 60 tokens, then `[kvflash] drafter scorer attached (tau=64)` automatically; 400 coherent tokens answering from the compressed context. Same binary, zero pflash-specific configuration on the pool. @@ -246,7 +246,7 @@ and masks through it. What differs per arch: 3.09, identical text). Policy: drafter-scored residency is the default on all four archs. The -server probes for the Qwen3-0.6B next to the model (or --prefill-drafter) +server probes for the Qwen3.5-0.8B next to the model (or --prefill-drafter) and lazy-loads it at the first reselect; `--kvflash-policy lru` opts out. qwen35/qwen35moe feed the drafter target ids directly; laguna/gemma4 use KvFlashCrossTokScorer (detokenize -> re-tokenize -> score -> map back by diff --git a/optimizations/kvflash/README.md b/optimizations/kvflash/README.md index 191f1c9cc..d94046004 100644 --- a/optimizations/kvflash/README.md +++ b/optimizations/kvflash/README.md @@ -41,7 +41,7 @@ does not fit at all.) # recommended: drafter-scored residency, pool auto-sized from VRAM. # pass --prefill-drafter so the drafter is guaranteed (no silent LRU fallback). dflash_server model.gguf --max-ctx 32768 --kvflash auto \ - --prefill-drafter /opt/lucebox/models/drafter/Qwen3-0.6B-BF16.gguf + --prefill-drafter /opt/lucebox/models/drafter/Qwen3.5-0.8B-BF16.gguf # drop the path to auto-probe (model dir, drafter/, draft/, /opt/lucebox/models/drafter/); # falls back to LRU if none is found, so check the banner reads policy=drafter @@ -52,7 +52,7 @@ dflash_server model.gguf --max-ctx 32768 --kvflash 8192 --kvflash-policy lru ``` Drafter-scored residency is the DEFAULT policy on every model family: -the server probes for `Qwen3-0.6B-BF16.gguf` next to the model (same +the server probes for `Qwen3.5-0.8B-BF16.gguf` next to the model (same dir, `drafter/`, `draft/`, then `/opt/lucebox/models/drafter/`) and lazy-loads it on the first reselect; `--prefill-drafter` overrides the location, prefill compression can stay off either way. Qwen-family @@ -147,7 +147,7 @@ Env: `DFLASH_KVFLASH_POLICY=qk`. Bench: `test_kvflash --qkbench`. - `server/src/common/kvflash_qk.h` — target-QK scorer: pure scoring math (unit-tested in `server/test/test_kvflash_qk.cpp`), seal-time key pooling, `KvFlashTargetQkScorer` -- `server/src/qwen3/qwen3_kvflash_scorer.{h,cpp}` — pflash-drafter scorer +- `server/src/pflash/kvflash_drafter_scorer.{h,cpp}` — pflash-drafter scorer (tail attention; bisects on allocation pressure) - `server/src/qwen35/*` — cache `ctx_alloc`, masked pooled decode, slot-mapped spec verify, daemon flags diff --git a/optimizations/pflash/README.md b/optimizations/pflash/README.md index 2281564c8..48cce0215 100644 --- a/optimizations/pflash/README.md +++ b/optimizations/pflash/README.md @@ -39,7 +39,7 @@ Long-context prefill is O(S²): vanilla llama.cpp on a single RTX 3090 takes **~ **What was missing:** no implementation that sits in front of a quantized GGUF target on a 24 GB card without dragging Python+Triton into the runtime path. PFlash is that: - C++/CUDA daemon-resident drafter + scoring + target generation, all in one process, one ggml allocator. -- Custom Qwen3-0.6B BF16 forward (`qwen3_0p6b_loader.cpp` + `qwen3_0p6b_graph.cpp`) — no libllama. +- Custom Qwen3.5-0.8B BF16 forward (`src/pflash/qwen35_loader.cpp` + the qwen35 target graph) — no libllama. - 4 CUDA kernels for the FlashPrefill `mean_K → score → select → sparse_fwd` algorithm (`flashprefill_kernels.cu`). - BSA ([mit-han-lab/Block-Sparse-Attention](https://github.com/mit-han-lab/Block-Sparse-Attention), FA-2 derived, sm_80+) for the long-context drafter forward, wired without `libtorch` via 3 ATen/c10 header stubs (`server/deps/bsa_stubs/`). - 128K → 2.6K span selection at `keep_ratio=0.05`, NIAH retrieved at every measured context, decode ~74 tok/s downstream. @@ -72,13 +72,13 @@ cmake --build server/build --target test_dflash test_flashprefill_kernels -j # 2. fetch weights (target + spec-decode draft + drafter scorer) uv run hf download unsloth/Qwen3.6-27B-GGUF Qwen3.6-27B-Q4_K_M.gguf --local-dir server/models/ -uv run hf download Qwen/Qwen3-0.6B model.safetensors tokenizer.json --local-dir server/models/drafter/ +uv run hf download Qwen/Qwen3.5-0.8B model.safetensors tokenizer.json --local-dir server/models/drafter/ uv run hf download z-lab/Qwen3.6-27B-DFlash model.safetensors --local-dir server/models/draft/ -# 2b. convert the drafter (Qwen3-0.6B HF) to a BF16 GGUF for the C++ scorer. +# 2b. convert the drafter (Qwen3.5-0.8B HF) to a BF16 GGUF for the C++ scorer. # The submodule already vendors llama.cpp at deps/llama.cpp. uv run python server/deps/llama.cpp/convert_hf_to_gguf.py server/models/drafter \ - --outtype bf16 --outfile server/models/Qwen3-0.6B-BF16.gguf + --outtype bf16 --outfile server/models/Qwen3.5-0.8B-BF16.gguf # 3. generate NIAH cases + run head-to-head bench against the C++ daemon uv run --directory pflash python tests/niah_gen.py --n 1 --ctx 131072 --out /tmp/niah_128k.jsonl @@ -86,7 +86,7 @@ uv run --directory pflash python tests/bench_niah_cpp.py \ --bin ../server/build/test_dflash \ --target ../server/models/Qwen3.6-27B-Q4_K_M.gguf \ --draft-spec ../server/models/draft/model.safetensors \ - --drafter-gguf ../server/models/Qwen3-0.6B-BF16.gguf \ + --drafter-gguf ../server/models/Qwen3.5-0.8B-BF16.gguf \ --cases /tmp/niah_128k.jsonl --keep-ratio 0.05 --n-gen 256 ``` @@ -99,8 +99,8 @@ For an OpenAI-compatible server with transparent compression on long prompts, ru | `--prefill-compression` | `off` / `auto` / `always` | `off` | When to run pflash. `auto` compresses when total prompt ≥ threshold; `always` compresses every request. | | `--prefill-threshold` | int (tokens) | `32000` | Token threshold for `auto` mode. | | `--prefill-keep-ratio` | float `(0, 1]` | `0.05` | Fraction of source tokens to keep after compression. `0.02` for 128K, `0.10` for 32K. | -| `--prefill-drafter` | path to `.gguf` | required when not `off` | Drafter weights (Qwen3-0.6B BF16 GGUF). | -| `--prefill-drafter-tokenizer` | HF repo id | `Qwen/Qwen3-0.6B` | HF tokenizer for the drafter vocab. | +| `--prefill-drafter` | path to `.gguf` | required when not `off` | Drafter weights (Qwen3.5-0.8B BF16 GGUF). | +| `--prefill-drafter-tokenizer` | HF repo id | `Qwen/Qwen3.5-0.8B` | HF tokenizer for the drafter vocab. | When `--prefill-compression != off`, the server auto-sets `DFLASH27B_LM_HEAD_FIX=0` and `DFLASH27B_FA_WINDOW=0` (matching the bench harness — needed so the post-compress draft graph fits on a 24 GB card without OOM). @@ -111,7 +111,7 @@ When `--prefill-compression != off`, the server auto-sets `DFLASH27B_LM_HEAD_FIX --prefill-compression auto \ --prefill-threshold 4096 \ --prefill-keep-ratio 0.02 \ - --prefill-drafter server/models/Qwen3-0.6B-BF16.gguf + --prefill-drafter server/models/Qwen3.5-0.8B-BF16.gguf ``` Below the threshold the server runs the standard target generate (no compression). Above it, the server transparently runs `compress` on the daemon, swaps the prompt for the compressed text, and continues the normal `/v1/chat/completions` flow. Tool-calling requests (`req.tools` non-empty) skip compression so JSON tool definitions stay intact. @@ -155,7 +155,7 @@ prompt (≤ 128K tokens) ▼ ┌──────────────────────────────────────────────┐ │ drafter (in-process) │ -│ custom Qwen3-0.6B BF16 forward in ggml │ +│ custom Qwen3.5-0.8B BF16 forward in ggml │ │ FlashPrefill block-sparse via BSA (≥ 32K) │ │ tail-attention scoring → score [S] │ │ chunk(128) + alpha-threshold → top blocks │ @@ -177,7 +177,7 @@ prompt (≤ 128K tokens) └──────────────────────────────────────────────┘ ``` -**Drafter forward.** Custom Qwen3-0.6B graph (`qwen3_0p6b_graph.cpp`) per-layer A/FP/B blocks: dense attention up to ~32K source, FlashPrefill sparse attention at and above. The 4 FP kernels live in `flashprefill_kernels.cu`; BSA dispatch is in `bsa_launcher.cu` + `bsa_fwd_inst.cu`. +**Drafter forward.** The Qwen3.5-0.8B drafter (`src/pflash/qwen35_drafter.cpp` on the qwen35 target graph) runs the model's first fifteen blocks and scores the context with block 15's NoPE Q/K attention-mass head; `PFLASH_QWEN35_LEGACY_SCORER=1` selects the all-layer running-max scorer instead. **Scoring + selection.** Tail attention `Q[-N:] @ K^T / sqrt(d)` per layer/head, max over (L, H), mean over the tail window. Block-level threshold by `alpha * mean(scores)` selects which K-blocks each Q-block attends to. Configurable via `DFLASH_FP_ALPHA`. @@ -205,14 +205,14 @@ What we built: - C++/CUDA port of the FlashPrefill algorithm: 4 kernels (`mean_K / score / select / sparse_fwd`), no Triton dependency. - BSA ([mit-han-lab/Block-Sparse-Attention](https://github.com/mit-han-lab/Block-Sparse-Attention)) wired without `libtorch` via 3 ATen/c10 header stubs (`server/deps/bsa_stubs/`). -- Custom Qwen3-0.6B BF16 forward so the drafter runs through the same ggml allocator as the 27B target. +- Custom Qwen3.5-0.8B BF16 forward so the drafter runs through the same ggml allocator as the 27B target. - Daemon stdin protocol (`compress` / `generate` / `park` / `unpark` / `free drafter`) so target + drafter coexist on a 24 GB card. - NIAH harness against `llama-bench` for end-to-end validation. ## Scope and limits - **Single 24 GB GPU** target (RTX 3090 reference). On 32+ GB cards, drafter + target can coexist and the park/unpark dance disappears. -- **Qwen3.6-27B Q4_K_M target + Qwen3-0.6B drafter** is the validated pair. Other targets/drafters need keep_ratio + alpha re-calibration. +- **Qwen3.6-27B Q4_K_M target + Qwen3.5-0.8B drafter** is the validated pair. Other targets/drafters need keep_ratio + alpha re-calibration. - **NIAH single-needle** is the only retrieval task validated end-to-end. Multi-doc QA, long-form code retrieval, etc. still TBD. - **sm_80+** required for BSA (RTX 3090 sm_86 is the reference). On sm_75 (Turing) the build auto-disables BSA and falls back to the WMMA path; expect a slower drafter forward at long ctx. @@ -242,16 +242,16 @@ These are operator-side flags on the launcher; they do not change PFlash semantics. A short prompt lane should keep the original defaults. -### Drafter selection: BF16 Qwen3-0.6B for compress +### Drafter selection: BF16 Qwen3.5-0.8B for compress PFlash compress benefits from a small, fast drafter. The validated -choice is **Qwen3-0.6B** in **BF16 safetensors** with ~5 attention +choice is **Qwen3.5-0.8B** in **BF16 safetensors** with ~5 attention layers. The DFlash drafter for the same target works correctly during decode-after-unpark but is heavier than ideal for compress. Practical guidance: -- Use Qwen3-0.6B BF16 for `compress` (PFlash side). +- Use Qwen3.5-0.8B BF16 for `compress` (PFlash side). - Reuse the larger DFlash drafter for `decode` after unpark (DFlash side). @@ -262,7 +262,7 @@ simultaneously on a 24 GB GPU. Reproducible comparison vs Ollama native `/api/chat` on the same 64K unique-prompt summary task, RTX 6000 Ada sm_89, -Qwen3.6-27B-Q4_K_M, FA_WINDOW=0. Drafter setup: Qwen3-0.6B BF16 +Qwen3.6-27B-Q4_K_M, FA_WINDOW=0. Drafter setup: Qwen3.5-0.8B BF16 GGUF for the PFlash compress path (see "Drafter selection" above); the larger DFlash drafter on the dflash daemon side ran as FP16 safetensors during decode-after-unpark on this run. Feel free to diff --git a/optimizations/pflash/pflash/dflash_client.py b/optimizations/pflash/pflash/dflash_client.py index 7116bc099..293160fe3 100644 --- a/optimizations/pflash/pflash/dflash_client.py +++ b/optimizations/pflash/pflash/dflash_client.py @@ -230,7 +230,7 @@ def park_target(self): self._send("park target\n") def unpark_target(self): self._send("unpark target\n") def compress(self, prompt_ids: list[int], keep_ratio: float, drafter_gguf: str, - drafter_arch: str = "qwen3-0.6b") -> list[int]: + drafter_arch: str = "qwen35-0.8b") -> list[int]: """C++ drafter score+compress via daemon. Returns compressed token ids. Daemon command: compress diff --git a/optimizations/pflash/tests/bench_niah_cpp.py b/optimizations/pflash/tests/bench_niah_cpp.py index 7aa396ee0..a0861cd43 100644 --- a/optimizations/pflash/tests/bench_niah_cpp.py +++ b/optimizations/pflash/tests/bench_niah_cpp.py @@ -27,12 +27,12 @@ def main(): ap.add_argument("--target", default="/opt/lucebox/models/Qwen3.6-27B-Q4_K_M.gguf") ap.add_argument("--draft-spec", default="/home/lucebox/lucebox-hub/dflash/models/draft/model.safetensors", help="draft model used for spec decoding (NOT drafter scorer)") - ap.add_argument("--drafter-gguf", default="/home/lucebox/lucebox-hub/dflash/models/Qwen3-0.6B-BF16.gguf", - help="C++ drafter scorer GGUF (Qwen3-0.6B BF16)") - ap.add_argument("--drafter-arch", default="qwen3-0.6b", choices=["qwen3-0.6b", "qwen35-0.8b"], + ap.add_argument("--drafter-gguf", default="/home/lucebox/lucebox-hub/dflash/models/Qwen3.5-0.8B-BF16.gguf", + help="C++ drafter scorer GGUF (Qwen3.5-0.8B BF16)") + ap.add_argument("--drafter-arch", default="qwen35-0.8b", choices=["qwen35-0.8b"], help="C++ drafter architecture selector") ap.add_argument("--target-tokenizer", default="Qwen/Qwen3.6-27B") - ap.add_argument("--drafter-tokenizer", default="Qwen/Qwen3-0.6B") + ap.add_argument("--drafter-tokenizer", default="Qwen/Qwen3.5-0.8B") ap.add_argument("--max-ctx", type=int, default=16384, help="daemon KV cache max ctx; sized for compressed prompt+gen, NOT source") ap.add_argument("--keep-ratio", type=float, default=0.020) diff --git a/optimizations/pflash/tests/niah_gen.py b/optimizations/pflash/tests/niah_gen.py index 39db4f1cb..c307e5588 100644 --- a/optimizations/pflash/tests/niah_gen.py +++ b/optimizations/pflash/tests/niah_gen.py @@ -110,7 +110,7 @@ def main(): # Default matches bench_niah_cpp.py's --drafter-tokenizer, since that is # the tokenizer the downstream NIAH bench uses to size case["prompt"] # for the drafter forward. Override for any other harness. - ap.add_argument("--tokenizer", default="Qwen/Qwen3-0.6B") + ap.add_argument("--tokenizer", default="Qwen/Qwen3.5-0.8B") args = ap.parse_args() tok = AutoTokenizer.from_pretrained(args.tokenizer) diff --git a/server/CMakeLists.txt b/server/CMakeLists.txt index 90501557b..fba1c9715 100644 --- a/server/CMakeLists.txt +++ b/server/CMakeLists.txt @@ -136,9 +136,9 @@ endif() # ─── ggml (vendored from llama.cpp) ────────────────────────────────── # -# We use only ggml from the vendored llama.cpp snapshot. The drafter is -# loaded via our own custom Qwen3-0.6B forward -# (src/qwen3/qwen3_loader.cpp + src/qwen3/qwen3_graph.cpp) +# We use only ggml from the vendored llama.cpp snapshot. The PFlash drafter +# is loaded via our own custom Qwen3.5-0.8B forward +# (src/pflash/qwen35_loader.cpp + src/qwen35/qwen35_target_graph.cpp) # rather than libllama, so libllama is not built. # # No BLAS, no Metal, no Vulkan, no examples/tests/tools. @@ -457,6 +457,7 @@ set(DFLASH27B_SRC_INCLUDE_DIRS ${CMAKE_CURRENT_SOURCE_DIR}/src/bailingmoe3 ${CMAKE_CURRENT_SOURCE_DIR}/src/laguna ${CMAKE_CURRENT_SOURCE_DIR}/src/qwen3 + ${CMAKE_CURRENT_SOURCE_DIR}/src/pflash ${CMAKE_CURRENT_SOURCE_DIR}/src/gemma4 ${CMAKE_CURRENT_SOURCE_DIR}/src/deepseek4 ${CMAKE_CURRENT_SOURCE_DIR}/src/server @@ -473,15 +474,14 @@ add_library(dflash_common STATIC src/draft/draft_gguf_loader.cpp src/draft/draft_safetensors_loader.cpp src/draft/draft_graph.cpp - src/qwen3/anchor_scan.cpp - src/qwen3/pflash_selection.cpp - src/qwen3/qwen3_drafter.cpp - src/qwen3/qwen3_drafter_common.cpp - src/qwen3/qwen35_drafter.cpp - src/qwen3/qwen35_loader.cpp - src/qwen3/qwen3_kvflash_scorer.cpp + src/pflash/anchor_scan.cpp + src/pflash/pflash_selection.cpp + src/pflash/pflash_drafter.cpp + src/pflash/pflash_compress.cpp + src/pflash/qwen35_drafter.cpp + src/pflash/qwen35_loader.cpp + src/pflash/kvflash_drafter_scorer.cpp src/qwen3/qwen3_loader.cpp - src/qwen3/qwen3_graph.cpp src/qwen3/qwen3_backend.cpp src/qwen3/qwen3_daemon.cpp src/gemma4/gemma4_loader.cpp @@ -666,7 +666,7 @@ endif() # - CUDA sm_60–sm_69 (Pascal): scalar F16, no tensor cores — flashprefill_scalar.cu # - HIP Phase 1 (default): ggml q8 fallback, no custom kernels. # - HIP Phase 2 (DFLASH27B_HIP_SM80_EQUIV=ON): rocWMMA-native kernels. -# The dispatch in qwen3_graph.cpp checks buffer type at runtime: +# The dispatch in flashprefill.h checks buffer type at runtime: # BF16 buffers → bf16 WMMA kernel; F16 buffers → f16 WMMA kernel; else → ggml FA. if(DFLASH27B_GPU_BACKEND STREQUAL "hip") # rms_norm_hip.cu is needed by the HIP chunk-B graph path regardless of SM80_EQUIV. @@ -1320,12 +1320,6 @@ if(DFLASH27B_TESTS) target_link_libraries(test_turbo_wht_warp PRIVATE CUDA::cudart) list(APPEND _raw_unit_test_targets test_turbo_wht_warp) endif() - if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_drafter_tail_capture_guard.cpp") - # RED phase binary: same source WITHOUT the fix flag — documents the bug. - add_executable(test_drafter_tail_capture_guard_red - test/test_unit_main.cpp - test/test_drafter_tail_capture_guard.cpp) - endif() if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_draft_vs_reference.cpp") add_executable(test_draft_vs_reference test/test_draft_vs_reference.cpp) target_link_libraries(test_draft_vs_reference PRIVATE dflash_common) @@ -1943,10 +1937,6 @@ if(DFLASH27B_TESTS) test/test_chain_rollback_policy.cpp test/test_ddtree_tau.cpp test/test_anchor_transitive.cpp - test/test_drafter_early_exit_score_range.cpp - test/test_drafter_tail_capture_guard.cpp - test/test_drafter_warm_path_regression.cpp - test/test_qwen3_buffer_plan.cpp test/test_pflash_drafter_ipc.cpp test/test_pflash_selection.cpp test/test_model_test_paths.cpp @@ -1972,7 +1962,7 @@ if(DFLASH27B_TESTS) src/server/scheduler.cpp src/server/model_card.cpp src/server/prompt_normalize.cpp - src/qwen3/anchor_scan.cpp) + src/pflash/anchor_scan.cpp) # Keep the GREEN formula target-local: the separate RED regression # executable compiles the same source without this definition. target_compile_definitions(test_server_unit PRIVATE diff --git a/server/README.md b/server/README.md index cf11ed246..7e11d8780 100644 --- a/server/README.md +++ b/server/README.md @@ -194,7 +194,7 @@ server is byte-identical to local-inference mode. ```bash ./build/dflash_server models/Qwen3.6-27B-Q4_K_M.gguf \ --prefill-compression auto --prefill-threshold 10000 \ - --prefill-drafter models/Qwen3-0.6B-BF16.gguf \ + --prefill-drafter models/Qwen3.5-0.8B-BF16.gguf \ --prefill-curve 10000:0.5 40000:0.2 100000:0.1 \ --prefill-upstream-base http://127.0.0.1:8099 \ --prefill-upstream-model my-upstream-model \ @@ -371,7 +371,7 @@ the whole request's device footprint. `/status/json` reports | `--prefill-threshold ` | `32000` | Token threshold used by auto mode. | | `--prefill-keep-ratio ` | `0.05` | Fraction of source tokens kept. | | `--prefill-curve T:R [T:R ...]` | none | Piecewise keep-ratio curve; overrides the flat ratio. | -| `--prefill-drafter ` | none | PFlash drafter GGUF: Qwen3-0.6B, or Qwen3.5-0.8B when the file name contains `qwen3.5`/`qwen35`. | +| `--prefill-drafter ` | none | PFlash drafter GGUF (Qwen3.5-0.8B). | | `--prefill-skip-park` | off | Keep target and decode draft resident while PFlash runs. | | `--prefill-upstream-base ` | none | Enable compression-proxy mode. | | `--prefill-upstream-key ` | none | Bearer token for the upstream. | @@ -380,10 +380,8 @@ the whole request's device footprint. `/status/json` reports With a Qwen3.5-0.8B drafter and strict budget selection (`PFLASH_SELECT_MODE=budget_only`, `PFLASH_SELECT_CHUNK_SIZE`, `PFLASH_SELECT_QUERY_TOKENS`), the drafter runs only its first fifteen -blocks and scores the context with block 15's NoPE Q/K projections, the same -attention-mass scorer the Qwen3-0.6B block-13 head uses. Its 262K native -context covers inputs the Qwen3-0.6B drafter cannot score within its 32K -window. `PFLASH_SCORING_HEAD_GGUF` accepts a trained block-15 head +blocks and scores the context with block 15's NoPE Q/K projections as an +attention-mass scorer. Its 262K native context covers very long inputs. `PFLASH_SCORING_HEAD_GGUF` accepts a trained block-15 head (schema `qwen3_5_0_8b_nope_qk_mass_v1`); `PFLASH_QWEN35_LEGACY_SCORER=1` restores the previous all-layer running-max scorer. The Qwen3.5 attention runs dense (`ggml_flash_attn_ext`); the block-sparse FlashPrefill kernels @@ -560,9 +558,9 @@ drives both arches end-to-end. The only thing the user changes is the model path ```bash cmake --build build --target test_dflash test_laguna_daemon pflash_daemon -j -# 19 GB Q4_K_M target + 1.2 GB Qwen3-0.6B BF16 drafter + tokenizers +# 19 GB Q4_K_M target + ~1.6 GB Qwen3.5-0.8B BF16 drafter + tokenizers hf download Lucebox/Laguna-XS.2-GGUF laguna-xs2-Q4_K_M.gguf --local-dir models/ -hf download unsloth/Qwen3-0.6B-GGUF Qwen3-0.6B-BF16.gguf --local-dir models/ +hf download unsloth/Qwen3.5-0.8B-GGUF Qwen3.5-0.8B-BF16.gguf --local-dir models/ hf download poolside/Laguna-XS.2 --local-dir models/Laguna-XS-2 \ --include 'tokenizer*' '*.json' @@ -584,9 +582,9 @@ DFLASH_KV_TYPE=q4_0 ./build/bench_laguna_ttft models/laguna-xs2-Q4_K_M.gguf '409 # standalone test_laguna_daemon binary so it can run without dflash_server. python3 scripts/laguna_pflash_niah.py \ --target models/laguna-xs2-Q4_K_M.gguf \ - --drafter models/Qwen3-0.6B-BF16.gguf \ + --drafter models/Qwen3.5-0.8B-BF16.gguf \ --laguna-tok models/Laguna-XS-2 \ - --drafter-tok Qwen/Qwen3-0.6B \ + --drafter-tok Qwen/Qwen3.5-0.8B \ --pflash-bin ./build/pflash_daemon \ --laguna-bin ./build/test_laguna_daemon \ --ctx 131072 --depth 0.5 --keep 0.10 --target-kv q4_0 diff --git a/server/docs/DS4.md b/server/docs/DS4.md index 95520faa7..e879ddb6d 100644 --- a/server/docs/DS4.md +++ b/server/docs/DS4.md @@ -978,7 +978,7 @@ park a target while it owns live sequence state. ./server/build-hip/dflash_server /path/to/deepseek4-target.gguf \ --target-device hip:0 \ --prefill-compression auto \ - --prefill-drafter /path/to/Qwen3-0.6B-BF16.gguf \ + --prefill-drafter /path/to/Qwen3.5-0.8B-BF16.gguf \ --prefill-skip-park ``` diff --git a/server/docs/ENVIRONMENT.md b/server/docs/ENVIRONMENT.md index c8b759477..21b640ec1 100644 --- a/server/docs/ENVIRONMENT.md +++ b/server/docs/ENVIRONMENT.md @@ -96,7 +96,7 @@ consolidation of this list into CLI flags is tracked as follow-up work. - `DFLASH27B_KV_F16` - kv_quant.cpp - `DFLASH27B_KV_K` - kv_quant.cpp, laguna_backend.cpp - `DFLASH27B_KV_Q4` - kv_quant.cpp -- `DFLASH27B_KV_TQ3` - kv_quant.cpp, qwen3_drafter.cpp +- `DFLASH27B_KV_TQ3` - kv_quant.cpp, pflash/qwen35_drafter.cpp - `DFLASH27B_KV_V` - kv_quant.cpp, laguna_backend.cpp - `DFLASH27B_LM_HEAD_FIX` - http_server.cpp - `DFLASH27B_PAGED_WMMA` - paged-attn.cu (ggml-cuda) (=1 routes paged full-attention layers to the WMMA kernel; RDNA4 only, F16/Q8_0/Q4_0, non-tree) @@ -199,14 +199,10 @@ consolidation of this list into CLI flags is tracked as follow-up work. - `DFLASH_FAST_ROLLBACK_THRESHOLD` - chain_rollback_policy.h - `DFLASH_FEATURE_DTYPE` - dflash_feature_ring.cpp - `DFLASH_KV_ROTATE` - qwen35_target_graph.cpp (set to 1 to force FWHT K rotation on; off by default for f16/q8_0 caches, on for narrower types) -- `DFLASH_FP_ALPHA` - http_server.cpp, qwen3_graph.cpp, server_main.cpp -- `DFLASH_FP_CHUNK_S` - qwen3_graph.cpp -- `DFLASH_FP_DEBUG_LAYER0` - qwen3_graph.cpp +- `DFLASH_FP_ALPHA` - http_server.cpp, server_main.cpp - `DFLASH_FP_DUMP_COUNTS` - flashprefill.cpp - `DFLASH_FP_HIP_ROW` - flashprefill_kernels.cu -- `DFLASH_FP_NOPE_TAIL` - qwen3_graph.cpp - `DFLASH_FP_PROFILE` - flashprefill.cpp -- `DFLASH_FP_SKIP_PREWARM` - qwen3_drafter.cpp - `DFLASH_FP_USE_BSA` - flashprefill.cpp, http_server.cpp, server_main.cpp - `DFLASH_G4_BSA_CHUNK` - gemma4_graph.cpp - `DFLASH_GEMMA4_LAYER_SPLIT_UBATCH` - gemma4_layer_split_adapter.cpp @@ -369,7 +365,5 @@ consolidation of this list into CLI flags is tracked as follow-up work. - `LUCE_MMVQ_MAX_NCOLS` - deepseek4_backend.cpp - `LUCE_QK_FUSE_LAYERS` - laguna_target_graph.cpp - `LUCE_QK_FUSE_MODE` - laguna_target_graph.cpp -- `PFLASH_DRAFTER_EARLY_EXIT_N` - qwen3_graph.cpp -- `PFLASH_DRAFTER_SCORE_LAYERS` - qwen3_graph.cpp - `PFLASH_FREEZE_HOT_WINDOW` - http_server.cpp - `TMPDIR` - backend_ipc.cpp, moe_expert_compute_ipc.cpp diff --git a/server/docs/SPEC_PREFILL.md b/server/docs/SPEC_PREFILL.md index 0859c50ee..99b019003 100644 --- a/server/docs/SPEC_PREFILL.md +++ b/server/docs/SPEC_PREFILL.md @@ -60,7 +60,7 @@ PFlash phase or DFlash draft-process boundary. See ## Performance NIAH single-needle end-to-end on RTX 3090 (Qwen3.6-27B Q4_K_M target, -Qwen3-0.6B drafter, in-process daemon, `DFLASH_FP_USE_BSA=1`, +Qwen3.5-0.8B drafter, in-process daemon, `DFLASH_FP_USE_BSA=1`, `DFLASH_FP_ALPHA=0.85`, `keep_ratio=0.05`): | Source S | dflash TTFT | llama.cpp baseline | Speedup | NIAH | @@ -81,10 +81,15 @@ src/ flashprefill_select.cpp Host fallback for block_select (rarely used) bsa_launcher.cu BSA launcher: blockmask conversion + Flash_fwd_params bsa_fwd_inst.cu Single-TU instantiation of BSA's hdim128 kernel - qwen3/ Qwen3-0.6B drafter model code + pflash/ PFlash drafter (Qwen3.5-0.8B scorer) code + qwen35_loader.cpp GGUF → Qwen3.5-0.8B weights + scoring head + probe + qwen35_drafter.{h,cpp} block-15 head scorer + legacy running-max scorer + pflash_drafter.{h,cpp} drafter_score_and_compress() entry point + pflash_compress.{h,cpp} scores → strict selection → compressed ids + pflash_selection.{h,cpp} strict budget selection + segment probing + qwen3/ Qwen3-0.6B standalone inference (not the drafter) qwen3_loader.cpp GGUF → Qwen3-0.6B BF16 weight tensors - qwen3_graph.cpp Custom Qwen3-0.6B forward (per-layer A/FP/B graphs) - qwen3_drafter.{h,cpp} drafter_score_and_compress() entry point + qwen3_backend.{h,cpp} step forward + ModelBackend qwen35/ Qwen3.5/3.6 target + DFlash draft model code qwen35_target_graph.cpp Qwen3.5/3.6 target graph (ggml) gguf_target_loader.cpp Qwen3.5 target GGUF loader diff --git a/server/docs/laguna_integration_plan.md b/server/docs/laguna_integration_plan.md index 60423e922..ab24a6b63 100644 --- a/server/docs/laguna_integration_plan.md +++ b/server/docs/laguna_integration_plan.md @@ -4,7 +4,7 @@ Status: scaffolding. PR #115 in lucebox-hub bumps llama.cpp submodule to `luce-d ## Context -- `pflash_daemon` (test/pflash_daemon.cpp): drafter-only stdin compressor, loads Qwen3-0.6B via dflash's own loader, emits compressed token IDs in DRAFTER vocab. Already model-agnostic on the target side. **No change needed.** +- `pflash_daemon` (test/pflash_daemon.cpp): drafter-only stdin compressor, loads Qwen3.5-0.8B via dflash's own loader, emits compressed token IDs in DRAFTER vocab. Already model-agnostic on the target side. **No change needed.** - `test_dflash` (test/test_dflash.cpp 190 KB): main target runner. Hand-rolled CUDA forward graph for qwen35 hybrid. Loads via `load_target_gguf` which hardcodes `arch == "qwen35"`. **Hard-blocked on Laguna.** - `qwen35_target_graph.cpp` (60 KB): hand-rolled CUDA forward, builds full-attn + delta-net + FFN. Uses `flash_prefill_forward_bf16` for sparse prefill. - `flashprefill.{h,cpp}` + `flashprefill_kernels.cu`: model-agnostic block-sparse FA. Takes Q/K/V tensors, returns O. Already works for any GQA arch with head_dim 128. **Reusable as-is.** @@ -57,7 +57,7 @@ No libllama dependency in dflash runtime. Keep ggml-only stack. (libllama+LAGUNA - Detect arch from loaded weights - For Laguna arch, use `LagunaTargetCache` + `build_laguna_graph` instead of qwen35 equivalents - Adjust per-layer-head-count in attention buffer sizing - - PFlash drafter call unchanged (drafter is Qwen3-0.6B regardless of target) + - PFlash drafter call unchanged (drafter is Qwen3.5-0.8B regardless of target) - Cross-tokenizer mapping (Qwen3 IDs → Laguna IDs): byte-level round-trip via existing optimizations/pflash/ Python module OR port to C++ helper ## Phasing diff --git a/server/scripts/laguna_pflash_niah.py b/server/scripts/laguna_pflash_niah.py index b14fe5dfd..512acca92 100644 --- a/server/scripts/laguna_pflash_niah.py +++ b/server/scripts/laguna_pflash_niah.py @@ -23,9 +23,9 @@ Usage: python3 laguna_pflash_niah.py \\ --target /path/to/laguna-xs2-Q4_K_M.gguf \\ - --drafter /path/to/Qwen3-0.6B-BF16.gguf \\ + --drafter /path/to/Qwen3.5-0.8B-BF16.gguf \\ --laguna-tok /path/to/Laguna-XS.2 \\ - --drafter-tok /path/to/Qwen3-0.6B \\ + --drafter-tok /path/to/Qwen3.5-0.8B \\ --pflash-bin /path/to/pflash_daemon \\ --laguna-bin /path/to/test_laguna_daemon \\ --ctx 16384 --depth 0.5 --keep 0.10 @@ -308,9 +308,9 @@ def close(self): def main(): ap = argparse.ArgumentParser() ap.add_argument("--target", required=True, type=Path, help="Laguna GGUF") - ap.add_argument("--drafter", required=True, type=Path, help="Qwen3-0.6B drafter GGUF") + ap.add_argument("--drafter", required=True, type=Path, help="Qwen3.5-0.8B drafter GGUF") ap.add_argument("--laguna-tok", required=True, type=Path, help="Laguna HF dir with tokenizer.json") - ap.add_argument("--drafter-tok", required=True, type=Path, help="Qwen3 HF dir with tokenizer.json") + ap.add_argument("--drafter-tok", required=True, type=Path, help="Qwen3.5 HF dir with tokenizer.json") ap.add_argument("--pflash-bin", required=True, type=Path) ap.add_argument("--laguna-bin", required=True, type=Path) ap.add_argument("--ctx", type=int, default=16384) diff --git a/server/scripts/phase_split_dual_gpu.py b/server/scripts/phase_split_dual_gpu.py index bbfcd53ef..988948ea7 100644 --- a/server/scripts/phase_split_dual_gpu.py +++ b/server/scripts/phase_split_dual_gpu.py @@ -1,7 +1,7 @@ #!/usr/bin/env python3 """Run PFlash prefill through a persistent daemon, optionally followed by target generation. -This phase-split harness is intentionally PFlash-only. It keeps the Qwen3-0.6B +This phase-split harness is intentionally PFlash-only. It keeps the Qwen3.5-0.8B PFlash drafter resident in `pflash_daemon`, optionally on a different CUDA or HIP backend from the later target run. The cross-backend boundary is host-side token/text data; target layer split remains inside one backend binary. @@ -35,8 +35,8 @@ def env_path(name: str, default: Path) -> Path: DEFAULT_BUILD = env_path("PFLASH_PHASE_BUILD_DIR", ROOT / "build") -DEFAULT_DRAFTER = env_path("PFLASH_PHASE_DRAFTER", ROOT / "models" / "Qwen3-0.6B-BF16.gguf") -DEFAULT_TOKENIZER = os.environ.get("PFLASH_PHASE_TOKENIZER", "Qwen/Qwen3-0.6B") +DEFAULT_DRAFTER = env_path("PFLASH_PHASE_DRAFTER", ROOT / "models" / "Qwen3.5-0.8B-BF16.gguf") +DEFAULT_TOKENIZER = os.environ.get("PFLASH_PHASE_TOKENIZER", "Qwen/Qwen3.5-0.8B") DEFAULT_TARGET = env_path("DFLASH_TARGET", ROOT / "models" / "Qwen3.6-27B-Q4_K_M.gguf") DEFAULT_TARGET_DRAFT = env_path("DFLASH_DRAFT", ROOT / "models" / "draft") DEFAULT_TARGET_TOKENIZER = os.environ.get("PFLASH_PHASE_TARGET_TOKENIZER", "Qwen/Qwen3.6-27B") diff --git a/server/scripts/quality_ab_simple.py b/server/scripts/quality_ab_simple.py index 19e9ee85a..9bee25d43 100644 --- a/server/scripts/quality_ab_simple.py +++ b/server/scripts/quality_ab_simple.py @@ -65,7 +65,7 @@ TARGET = os.environ.get("PFLASH_TARGET", "/home/peppi/models/qwen3.6-27b/Qwen3.6-27B-UD-Q4_K_XL.gguf") DRAFT = os.environ.get("PFLASH_DRAFT", "/home/peppi/models/qwen3.6-27b-dflash/model.safetensors") SERVER_BIN = os.environ.get("DFLASH_SERVER_BIN", "dflash/build/dflash_server") -DRAFTER = os.environ.get("PFLASH_DRAFTER", str(Path.home() / "models/Qwen3-0.6B-BF16.gguf")) +DRAFTER = os.environ.get("PFLASH_DRAFTER", str(Path.home() / "models/Qwen3.5-0.8B-BF16.gguf")) def chat_post(payload, timeout=120): diff --git a/server/scripts/quality_humaneval_plus.py b/server/scripts/quality_humaneval_plus.py index 24a5cfa1a..2d51eac46 100644 --- a/server/scripts/quality_humaneval_plus.py +++ b/server/scripts/quality_humaneval_plus.py @@ -61,7 +61,7 @@ TARGET = os.environ.get("PFLASH_TARGET", "/home/peppi/models/qwen3.6-27b/Qwen3.6-27B-UD-Q4_K_XL.gguf") DRAFT = os.environ.get("PFLASH_DRAFT", "/home/peppi/models/qwen3.6-27b-dflash/model.safetensors") SERVER_BIN = os.environ.get("DFLASH_SERVER_BIN", str(PROJECT_ROOT / "dflash/build/dflash_server")) -DRAFTER = os.environ.get("PFLASH_DRAFTER", str(Path.home() / "models/Qwen3-0.6B-BF16.gguf")) +DRAFTER = os.environ.get("PFLASH_DRAFTER", str(Path.home() / "models/Qwen3.5-0.8B-BF16.gguf")) # Canonical EvalPlus chat-mode prompt (evalplus/codegen.py:222-223) INSTRUCTION_PREFIX = ( diff --git a/server/scripts/test_full_compress_cache.py b/server/scripts/test_full_compress_cache.py index 715d50f91..fc708d01a 100644 --- a/server/scripts/test_full_compress_cache.py +++ b/server/scripts/test_full_compress_cache.py @@ -14,7 +14,7 @@ Skipped automatically if any prerequisite is missing: - target GGUF - draft (drafter) safetensors dir or GGUF - - Qwen3-0.6B-BF16 drafter GGUF + - Qwen3.5-0.8B-BF16 drafter GGUF - test_dflash binary """ import os @@ -33,7 +33,7 @@ ROOT = Path(__file__).resolve().parent.parent.parent TARGET = Path.home() / "models/qwen3.6-27b/Qwen3.6-27B-UD-Q4_K_XL.gguf" DRAFT = Path.home() / "models/qwen3.6-27b-dflash" -DRAFTER_GGUF = Path.home() / "models/Qwen3-0.6B-BF16.gguf" +DRAFTER_GGUF = Path.home() / "models/Qwen3.5-0.8B-BF16.gguf" SERVER_BIN = ROOT / "dflash/build/dflash_server" for p, label in [ diff --git a/server/src/common/kvflash_pager.h b/server/src/common/kvflash_pager.h index 0f80b2274..7c67dbc1b 100644 --- a/server/src/common/kvflash_pager.h +++ b/server/src/common/kvflash_pager.h @@ -700,7 +700,7 @@ inline bool kvflash_policy_is_qk() { return env && std::strcmp(env, "qk") == 0; } -// Locate the Qwen3-0.6B residency drafter: the explicit override +// Locate the Qwen3.5-0.8B residency drafter: the explicit override // (DFLASH_KVFLASH_DRAFTER, set from --prefill-drafter), then the // well-known locations next to the target model, then the appliance path. // Returns "" when nothing is readable (callers fall back to LRU, loudly). @@ -712,10 +712,10 @@ inline std::string kvflash_find_drafter(const char * target_path) { const size_t slash = dir.find_last_of('/'); dir = (slash == std::string::npos) ? "." : dir.substr(0, slash); const std::string candidates[] = { - dir + "/Qwen3-0.6B-BF16.gguf", - dir + "/drafter/Qwen3-0.6B-BF16.gguf", - dir + "/draft/Qwen3-0.6B-BF16.gguf", - "/opt/lucebox/models/drafter/Qwen3-0.6B-BF16.gguf", + dir + "/Qwen3.5-0.8B-BF16.gguf", + dir + "/drafter/Qwen3.5-0.8B-BF16.gguf", + dir + "/draft/Qwen3.5-0.8B-BF16.gguf", + "/opt/lucebox/models/drafter/Qwen3.5-0.8B-BF16.gguf", }; for (const std::string & c : candidates) { if (std::FILE * f = std::fopen(c.c_str(), "rb")) { diff --git a/server/src/common/kvflash_scorer.h b/server/src/common/kvflash_scorer.h index 407d94c6d..303b01bef 100644 --- a/server/src/common/kvflash_scorer.h +++ b/server/src/common/kvflash_scorer.h @@ -8,8 +8,8 @@ // // Implementations: // - (none) pure LRU + recency, zero dependencies -// - KvFlashDrafterScorer qwen3/qwen3_kvflash_scorer.h — pflash drafter tail -// attention (shared with pflash compression) +// - KvFlashDrafterScorer pflash/kvflash_drafter_scorer.h — pflash drafter +// tail attention (shared with pflash compression) #pragma once diff --git a/server/src/common/pflash_drafter_ipc_daemon.cpp b/server/src/common/pflash_drafter_ipc_daemon.cpp index f63299013..972ac4a8b 100644 --- a/server/src/common/pflash_drafter_ipc_daemon.cpp +++ b/server/src/common/pflash_drafter_ipc_daemon.cpp @@ -4,7 +4,7 @@ #include "dflash27b.h" #include "dflash_draft_ipc.h" -#include "qwen3/qwen3_drafter.h" +#include "pflash/pflash_drafter.h" #include #include @@ -61,10 +61,14 @@ int run_pflash_drafter_ipc_daemon(const char * drafter_path, stream_status(stream_fd, -1); continue; } + // The IPC protocol uses score_query_end < 0 for "tail"; the + // qwen35 scorer requires an explicit end, so translate here. + const int score_query_end = request.score_query_end >= 0 + ? request.score_query_end : (int)input_ids.size(); auto compressed = drafter_score_and_compress( ctx, input_ids, request.keep_ratio, /*chunk_size=*/32, request.score_query_tokens, /*pool_kernel=*/13, - request.score_query_end, + score_query_end, request.required_instruction_spans); if (compressed.empty()) { std::fprintf(stderr, "[pflash-ipc-daemon] compress returned empty\n"); diff --git a/server/src/common/score_range.h b/server/src/common/score_range.h deleted file mode 100644 index eb4a581a4..000000000 --- a/server/src/common/score_range.h +++ /dev/null @@ -1,31 +0,0 @@ -// Compute [score_layer_start, score_layer_end) for tail-attention scoring. -// SCORE_LAYERS counts from the END of [0, fwd_layer_limit); -1 = all computed layers. -#pragma once - -#include - -namespace dflash::common { - -struct ScoreRange { - int start; // inclusive - int end; // exclusive - int count() const { return end - start; } - bool empty() const { return start >= end; } -}; - -// Returns scoring layer range within [0, fwd_layer_limit). -inline ScoreRange compute_score_range(int n_layer, int score_layers, int fwd_layer_limit) { - const int effective_n = fwd_layer_limit; - int start; - if (score_layers > 0 && score_layers < n_layer) { - int want = std::min(score_layers, effective_n); - start = effective_n - want; - } else { - start = 0; - } - int end = fwd_layer_limit; - if (start > end) start = end; - return { start, end }; -} - -} // namespace dflash::common diff --git a/server/src/deepseek4/deepseek4_backend.cpp b/server/src/deepseek4/deepseek4_backend.cpp index 3dff0db77..f0fbcac16 100644 --- a/server/src/deepseek4/deepseek4_backend.cpp +++ b/server/src/deepseek4/deepseek4_backend.cpp @@ -3174,8 +3174,14 @@ std::vector DeepSeek4Backend::compress_batch( const CompressRequest & request = requests[index]; if (!valid_request(request)) continue; CompressResult & result = results[index]; + // score_query_end < 0 is the legacy "tail window" request value; + // the qwen35 scorer requires an explicit end. + const int score_query_end = request.score_query_end >= 0 + ? request.score_query_end : (int)request.input_ids.size(); result.compressed_ids = drafter_score_and_compress( - pflash_drafter_ctx_, request.input_ids, request.keep_ratio); + pflash_drafter_ctx_, request.input_ids, request.keep_ratio, + /*chunk_size=*/32, request.score_query_tokens, /*pool_kernel=*/13, + score_query_end, request.required_instruction_spans); result.ok = !result.compressed_ids.empty(); } diff --git a/server/src/deepseek4/deepseek4_backend.h b/server/src/deepseek4/deepseek4_backend.h index 64628e8cc..98a0a2c5c 100644 --- a/server/src/deepseek4/deepseek4_backend.h +++ b/server/src/deepseek4/deepseek4_backend.h @@ -15,7 +15,7 @@ #include "../common/moe_hybrid_stream.h" #include "deepseek4_internal.h" #include "deepseek4_dspark.h" -#include "qwen3/qwen3_drafter.h" +#include "pflash/pflash_drafter.h" #include "deepseek4_seq_engine.h" #include "ggml.h" diff --git a/server/src/flashprefill.h b/server/src/flashprefill.h index fd0c64e04..e3ea64ba5 100644 --- a/server/src/flashprefill.h +++ b/server/src/flashprefill.h @@ -1,5 +1,5 @@ // Public C++ entry point for the FlashPrefill block-sparse attention used by -// the in-process Qwen3-0.6B drafter (speculative prefill scoring). +// the in-process Qwen3.5-0.8B drafter (speculative prefill scoring). // // Wraps kernels 1-4 + GPU block_select into one call. Call signature mirrors // the upstream `flash_prefill` from qhfan/FlashPrefill (arXiv:2603.06199). diff --git a/server/src/gemma4/gemma4_backend.cpp b/server/src/gemma4/gemma4_backend.cpp index 83c5a7782..d691699b1 100644 --- a/server/src/gemma4/gemma4_backend.cpp +++ b/server/src/gemma4/gemma4_backend.cpp @@ -6,7 +6,7 @@ #include "gemma4_backend.h" #include "dflash27b.h" -#include "../qwen3/qwen3_kvflash_scorer.h" +#include "pflash/kvflash_drafter_scorer.h" #include "common/sampler.h" #include "common/io_utils.h" #include "common/dflash_feature_ring.h" @@ -190,7 +190,7 @@ void Gemma4Backend::kvflash_read_config() { } // Drafter rescore + repage (FlashMemory tau loop) with the cross-tokenizer -// scorer: gemma ids are detokenized and re-scored through the Qwen3-0.6B +// scorer: gemma ids are detokenized and re-scored through the Qwen3.5-0.8B // drafter. Lazy: the drafter + tokenizers load on the first reselect that // needs them, never on a request's first tokens. void Gemma4Backend::kvflash_maybe_reselect(int generated) { @@ -257,7 +257,7 @@ bool Gemma4Backend::kvflash_attach() { cache_.swa_size, !kvflash_drafter_path_.empty() ? "drafter/cross-tok (attaches on first reselect)" - : "lru (recency-only: no Qwen3-0.6B drafter found)"); + : "lru (recency-only: no Qwen3.5-0.8B drafter found)"); std::fflush(stdout); return true; } @@ -1203,7 +1203,7 @@ bool Gemma4Backend::handle_compress(const std::string & line, const char * dpath = (n >= 3 && drafter_path[0]) ? drafter_path - : "/opt/lucebox/models/drafter/Qwen3-0.6B-BF16.gguf"; + : "/opt/lucebox/models/drafter/Qwen3.5-0.8B-BF16.gguf"; // Park target to free VRAM for the drafter (unless skip_park). const bool was_parked = parked_; @@ -1232,7 +1232,9 @@ bool Gemma4Backend::handle_compress(const std::string & line, bool ok = false; if (!tokens.empty()) { const float keep = (float)keep_x1000 / 1000.0f; - auto compressed = drafter_score_and_compress(drafter_ctx_, tokens, keep); + auto compressed = drafter_score_and_compress(drafter_ctx_, tokens, keep, + /*chunk_size=*/32, /*n_lookahead=*/8, /*pool_kernel=*/13, + (int)tokens.size()); ok = !compressed.empty(); if (ok) { std::fprintf(stderr, "[compress] %zu -> %zu tokens\n", diff --git a/server/src/gemma4/gemma4_backend.h b/server/src/gemma4/gemma4_backend.h index 689237c83..3832d7529 100644 --- a/server/src/gemma4/gemma4_backend.h +++ b/server/src/gemma4/gemma4_backend.h @@ -14,7 +14,7 @@ #include "common/sampler.h" #include "../common/kvflash_pager.h" #include "../common/kvflash_scorer.h" -#include "../qwen3/qwen3_drafter.h" +#include "pflash/pflash_drafter.h" #include "ggml.h" #include "ggml-backend.h" @@ -106,7 +106,7 @@ class Gemma4Backend : public ModelBackend { // Pools the FULL-attention layers only (SWA layers already ring-buffer). // Drafter-scored residency by default via the cross-tokenizer bridge // (KvFlashCrossTokScorer: gemma ids are detokenized and re-scored by - // the Qwen3-0.6B drafter); LRU is the fallback when no drafter is + // the Qwen3.5-0.8B drafter); LRU is the fallback when no drafter is // found or --kvflash-policy lru. KvFlashPager kvflash_pager_; std::unique_ptr kvflash_scorer_; diff --git a/server/src/gemma4/gemma4_layer_split_adapter.cpp b/server/src/gemma4/gemma4_layer_split_adapter.cpp index 839c123d6..3e16946ae 100644 --- a/server/src/gemma4/gemma4_layer_split_adapter.cpp +++ b/server/src/gemma4/gemma4_layer_split_adapter.cpp @@ -10,7 +10,7 @@ #include "common/target_shard_ipc_daemon.h" #include "dflash27b.h" #include "placement/placement_backend.h" -#include "qwen3/qwen3_kvflash_scorer.h" +#include "pflash/kvflash_drafter_scorer.h" #include "ggml-cuda.h" @@ -410,7 +410,7 @@ bool Gemma4LayerSplitAdapter::kvflash_attach() { kvflash_tau_, !kvflash_drafter_path_.empty() ? "drafter/cross-tok (attaches on first reselect)" - : "lru (recency-only: no Qwen3-0.6B drafter found)"); + : "lru (recency-only: no Qwen3.5-0.8B drafter found)"); std::fflush(stdout); return true; } diff --git a/server/src/gemma4/gemma4_layer_split_adapter.h b/server/src/gemma4/gemma4_layer_split_adapter.h index 5b1a2b1a3..7466284fd 100644 --- a/server/src/gemma4/gemma4_layer_split_adapter.h +++ b/server/src/gemma4/gemma4_layer_split_adapter.h @@ -11,7 +11,7 @@ #include "gemma4_internal.h" #include "placement/placement_config.h" #include "placement/remote_target_shard_config.h" -#include "qwen3/qwen3_drafter.h" +#include "pflash/pflash_drafter.h" #include "ggml-backend.h" diff --git a/server/src/laguna/laguna_backend.cpp b/server/src/laguna/laguna_backend.cpp index c32617b04..f3a974fc6 100644 --- a/server/src/laguna/laguna_backend.cpp +++ b/server/src/laguna/laguna_backend.cpp @@ -8,7 +8,7 @@ #include "laguna_backend.h" #include "laguna_internal.h" -#include "qwen3/qwen3_kvflash_scorer.h" +#include "pflash/kvflash_drafter_scorer.h" #include "dflash27b.h" #include "common/ddtree.h" #include "common/domino_head.h" @@ -270,7 +270,7 @@ void LagunaBackend::kvflash_read_config() { } // Drafter rescore + repage (FlashMemory tau loop) with the cross-tokenizer -// scorer: laguna ids are detokenized and re-scored through the Qwen3-0.6B +// scorer: laguna ids are detokenized and re-scored through the Qwen3.5-0.8B // drafter (relevance is text-level, so the tokenizer gap is bridged by // re-tokenization). Lazy: the drafter + tokenizers load on the first // reselect that needs them, never on a request's first tokens. @@ -341,7 +341,7 @@ bool LagunaBackend::kvflash_attach() { kvflash_tokens_, args_.max_ctx, !kvflash_drafter_path_.empty() ? "drafter/cross-tok (attaches on first reselect)" - : "lru (recency-only: no Qwen3-0.6B drafter found)", + : "lru (recency-only: no Qwen3.5-0.8B drafter found)", pc.tail_window_chunks); std::fflush(stdout); return true; @@ -1864,13 +1864,14 @@ bool LagunaBackend::handle_compress(const std::string & line, return true; } drafter_loaded_ = true; - std::printf("[drafter] loaded %s vocab=%d\n", - drafter_path, drafter_ctx_.weights.n_vocab); + std::printf("[drafter] loaded %s\n", drafter_path); std::fflush(stdout); } const float keep = (float)keep_x1000 / 1000.0f; - auto compressed = drafter_score_and_compress(drafter_ctx_, src_ids, keep); + auto compressed = drafter_score_and_compress(drafter_ctx_, src_ids, keep, + /*chunk_size=*/32, /*n_lookahead=*/8, /*pool_kernel=*/13, + (int)src_ids.size()); std::printf("[compress] %zu -> %zu tokens (keep_ratio=%.3f)\n", src_ids.size(), compressed.size(), keep); std::fflush(stdout); diff --git a/server/src/laguna/laguna_backend.h b/server/src/laguna/laguna_backend.h index 939795eaa..8f4b5a602 100644 --- a/server/src/laguna/laguna_backend.h +++ b/server/src/laguna/laguna_backend.h @@ -13,7 +13,7 @@ #include "common/dflash_draft_graph.h" #include "common/dflash_draft_kv.h" #include "placement/placement_config.h" -#include "qwen3_drafter.h" +#include "pflash/pflash_drafter.h" #include "kvflash_pager.h" #include "kvflash_scorer.h" #include "../common/moe_hybrid_ffn_eval.h" @@ -141,7 +141,7 @@ class LagunaBackend : public ModelBackend { bool ensure_slot(int slot); // ── kvflash (bounded KV residency; see common/kvflash_pager.h) ── - // Drafter-scored residency by default: the Qwen3-0.6B drafter scores + // Drafter-scored residency by default: the Qwen3.5-0.8B drafter scores // chunks through the cross-tokenizer bridge (KvFlashCrossTokScorer — // relevance is text-level, so the target's ids are detokenized and // re-tokenized for the drafter). LRU is the fallback when no drafter is diff --git a/server/src/laguna/laguna_layer_split_adapter.cpp b/server/src/laguna/laguna_layer_split_adapter.cpp index 3c5515cba..06f63bb54 100644 --- a/server/src/laguna/laguna_layer_split_adapter.cpp +++ b/server/src/laguna/laguna_layer_split_adapter.cpp @@ -11,7 +11,7 @@ #include "common/target_shard_ipc_daemon.h" #include "dflash27b.h" #include "placement/placement_backend.h" -#include "qwen3/qwen3_kvflash_scorer.h" +#include "pflash/kvflash_drafter_scorer.h" #include "ggml-cuda.h" #include "ggml-cpu.h" @@ -329,7 +329,7 @@ bool LagunaLayerSplitAdapter::kvflash_attach() { kvflash_tau_, !kvflash_drafter_path_.empty() ? "drafter/cross-tok (attaches on first reselect)" - : "lru (recency-only: no Qwen3-0.6B drafter found)", + : "lru (recency-only: no Qwen3.5-0.8B drafter found)", pc.tail_window_chunks); std::fflush(stdout); return true; diff --git a/server/src/laguna/laguna_layer_split_adapter.h b/server/src/laguna/laguna_layer_split_adapter.h index 12085231c..5eb6c552b 100644 --- a/server/src/laguna/laguna_layer_split_adapter.h +++ b/server/src/laguna/laguna_layer_split_adapter.h @@ -11,7 +11,7 @@ #include "laguna_internal.h" #include "placement/placement_config.h" #include "placement/remote_target_shard_config.h" -#include "qwen3/qwen3_drafter.h" +#include "pflash/pflash_drafter.h" #include "ggml-backend.h" diff --git a/server/src/qwen3/anchor_params.h b/server/src/pflash/anchor_params.h similarity index 100% rename from server/src/qwen3/anchor_params.h rename to server/src/pflash/anchor_params.h diff --git a/server/src/qwen3/anchor_scan.cpp b/server/src/pflash/anchor_scan.cpp similarity index 98% rename from server/src/qwen3/anchor_scan.cpp rename to server/src/pflash/anchor_scan.cpp index 1c1592caf..9a1f710bf 100644 --- a/server/src/qwen3/anchor_scan.cpp +++ b/server/src/pflash/anchor_scan.cpp @@ -5,7 +5,7 @@ #include #include -namespace dflash::qwen3 { +namespace dflash::pflash { // Force chunk and its radius-neighborhood into `forced`. static void force_neighborhood(std::vector& forced, int n_chunks, @@ -161,4 +161,4 @@ void scan_and_force_transitive( } } -} // namespace dflash::qwen3 +} // namespace dflash::pflash diff --git a/server/src/qwen3/anchor_scan.h b/server/src/pflash/anchor_scan.h similarity index 95% rename from server/src/qwen3/anchor_scan.h rename to server/src/pflash/anchor_scan.h index 8f75a0855..82c3fefa0 100644 --- a/server/src/qwen3/anchor_scan.h +++ b/server/src/pflash/anchor_scan.h @@ -6,7 +6,7 @@ #include #include -namespace dflash::qwen3 { +namespace dflash::pflash { struct AnchorScanCfg { int chunk_size; @@ -39,4 +39,4 @@ void scan_and_force_transitive( std::vector& forced ); -} // namespace dflash::qwen3 +} // namespace dflash::pflash diff --git a/server/src/qwen3/qwen3_kvflash_scorer.cpp b/server/src/pflash/kvflash_drafter_scorer.cpp similarity index 88% rename from server/src/qwen3/qwen3_kvflash_scorer.cpp rename to server/src/pflash/kvflash_drafter_scorer.cpp index 4dc00c7c9..e08e25024 100644 --- a/server/src/qwen3/qwen3_kvflash_scorer.cpp +++ b/server/src/pflash/kvflash_drafter_scorer.cpp @@ -1,6 +1,6 @@ -#include "qwen3_kvflash_scorer.h" +#include "kvflash_drafter_scorer.h" -#include "qwen3_drafter_model.h" +#include "qwen35_drafter.h" #include "server/tokenizer.h" #include @@ -15,30 +15,25 @@ constexpr int kLookahead = 8; constexpr int kPoolKernel = 13; constexpr int kMinSegment = 4096; -// Tail-attention token scores for `ids`: mean over the lookahead window of -// the drafter's running-max, then AvgPool smoothing. Same math as -// drafter_score_and_compress. +// Tail-attention token scores for `ids` from the Qwen3.5-0.8B drafter: +// the all-layer running-max scorer with AvgPool smoothing. Same math as +// drafter_score_and_compress with the legacy scorer. bool score_tokens_direct(DrafterContext & ctx, const std::vector & ids, std::vector & out) { - const int S = (int)ids.size(); - std::vector running_max; - if (!forward_qwen3_drafter_model(ctx.weights, ids, kLookahead, running_max)) { + if (!ctx.state) return false; + const dflash::pflash::PFlashSelectionConfig experiment; + std::vector scores; + if (qwen35_score_and_compress(ctx.state->weights, ids, + /*keep_ratio=*/1.0f, /*chunk_size=*/64, + kLookahead, kPoolKernel, + /*score_query_end=*/-1, + experiment, + /*required_instruction_spans=*/{}, + &scores).empty() || + scores.size() != ids.size()) { return false; } - std::vector score((size_t)S, 0.0f); - for (int j = 0; j < S; j++) { - float s = 0.0f; - for (int t = 0; t < kLookahead; t++) s += running_max[(size_t)t * S + j]; - score[j] = s / kLookahead; - } - out.assign((size_t)S, 0.0f); - const int half = kPoolKernel / 2; - for (int j = 0; j < S; j++) { - const int lo = std::max(0, j - half), hi = std::min(S - 1, j + half); - float s = 0.0f; - for (int k = lo; k <= hi; k++) s += score[k]; - out[j] = s / (hi - lo + 1); - } + out = std::move(scores); return true; } diff --git a/server/src/qwen3/qwen3_kvflash_scorer.h b/server/src/pflash/kvflash_drafter_scorer.h similarity index 83% rename from server/src/qwen3/qwen3_kvflash_scorer.h rename to server/src/pflash/kvflash_drafter_scorer.h index e0fda5074..2451716b9 100644 --- a/server/src/qwen3/qwen3_kvflash_scorer.h +++ b/server/src/pflash/kvflash_drafter_scorer.h @@ -1,15 +1,15 @@ // KvFlashDrafterScorer — pflash drafter as the KV pager's Memory Indexer. // -// Scores 64-token chunks with the same Liu Q-hook tail attention that -// pflash compression uses (forward_qwen3_drafter_model), but returns the -// per-chunk relevance scores instead of a compressed token list. The -// DrafterContext is borrowed: the daemon shares its pflash drafter; the -// pager itself never depends on this file (see common/kvflash_scorer.h). +// Scores 64-token chunks with the same tail-attention scoring that pflash +// compression uses (the Qwen3.5-0.8B drafter), but returns the per-chunk +// relevance scores instead of a compressed token list. The DrafterContext +// is borrowed: the daemon shares its pflash drafter; the pager itself never +// depends on this file (see common/kvflash_scorer.h). #pragma once #include "kvflash_scorer.h" -#include "qwen3_drafter.h" +#include "pflash_drafter.h" #include @@ -19,7 +19,7 @@ class KvFlashDrafterScorer : public KvFlashScorer { public: // `vocab_clamp`: ids >= clamp are folded into the drafter's vocab range // before scoring. Needed when the target vocabulary is a superset of - // the drafter's (e.g. Qwen3.6 target + Qwen3-0.6B drafter); prompt ids + // the drafter's (e.g. Qwen3.6 target + Qwen3.5-0.8B drafter); prompt ids // tokenized for the target may be unembeddable by the drafter. explicit KvFlashDrafterScorer(DrafterContext * ctx, int32_t vocab_clamp = 100000) : ctx_(ctx), vocab_clamp_(vocab_clamp) {} diff --git a/server/src/qwen3/qwen3_drafter_common.cpp b/server/src/pflash/pflash_compress.cpp similarity index 89% rename from server/src/qwen3/qwen3_drafter_common.cpp rename to server/src/pflash/pflash_compress.cpp index 769599b73..07b7410db 100644 --- a/server/src/qwen3/qwen3_drafter_common.cpp +++ b/server/src/pflash/pflash_compress.cpp @@ -1,9 +1,7 @@ -// Helpers shared by the Qwen3-0.6B and Qwen3.5-0.8B drafter paths. -// See qwen3_drafter_common.h. +// PFlash scoring pipeline glue. See pflash_compress.h. -#include "qwen3_drafter_common.h" +#include "pflash_compress.h" -#include "qwen3_drafter.h" #include "pflash_selection.h" #include "internal.h" @@ -70,7 +68,7 @@ void write_compression_trace( trace_fields->exact_chunk_scores->size() == scores.size(); if (trace_fields && trace_fields->selector_mode != - dflash::qwen3::PFlashSelectionMode::Legacy && + dflash::pflash::PFlashSelectionMode::Legacy && !has_exact_scores) { std::fclose(file); std::fprintf(stderr, "[pflash-trace] exact strict scores unavailable\n"); @@ -92,9 +90,9 @@ void write_compression_trace( "\"token_budget\":%d," "\"retained_tokens\":%d", trace_fields->query_begin, trace_fields->query_end, - dflash::qwen3::pflash_selection_mode_name( + dflash::pflash::pflash_selection_mode_name( trace_fields->selector_mode), - dflash::qwen3::pflash_query_parser_name(trace_fields->query_parser), + dflash::pflash::pflash_query_parser_name(trace_fields->query_parser), trace_fields->token_budget, trace_fields->retained_tokens); std::fputs(",\"required_instruction_spans\":[", file); if (trace_fields->required_instruction_spans) { @@ -109,12 +107,12 @@ void write_compression_trace( } std::fputc(']', file); if (trace_fields->selector_mode == - dflash::qwen3::PFlashSelectionMode::Legacy) { + dflash::pflash::PFlashSelectionMode::Legacy) { std::fputs(",\"stop_reason\":null,\"retained_mass\":null", file); } else { std::fprintf(file, ",\"stop_reason\":\"%s\",\"retained_mass\":%.17g", - dflash::qwen3::pflash_selection_stop_name(trace_fields->stop), + dflash::pflash::pflash_selection_stop_name(trace_fields->stop), trace_fields->retained_mass); } } @@ -187,7 +185,7 @@ std::vector select_pflash_chunks( int n_lookahead, int score_query_end, int pool_kernel, - const dflash::qwen3::PFlashSelectionConfig & config, + const dflash::pflash::PFlashSelectionConfig & config, const std::vector & required_instruction_spans, bool direct_mass, bool write_trace, @@ -206,7 +204,7 @@ std::vector select_pflash_chunks( ? (int) segments->size() : (input_tokens + config.chunk_size - 1) / config.chunk_size; - std::vector candidates; + std::vector candidates; std::vector> chunk_means; std::vector exact_chunk_scores; candidates.reserve((size_t) n_chunks); @@ -224,7 +222,7 @@ std::vector select_pflash_chunks( score /= (double) std::max(1, end - begin); } const bool mandatory = - dflash::qwen3::pflash_chunk_is_structurally_required( + dflash::pflash::pflash_chunk_is_structurally_required( begin, end, query_begin, query_end, input_tokens, required_instruction_spans); candidates.push_back({(size_t) chunk, begin, end, score, mandatory}); @@ -233,7 +231,7 @@ std::vector select_pflash_chunks( } // Two-scorer selection: the other scorer's mean per-token score over the // same spans (its native ranking rule). - std::vector other_candidates; + std::vector other_candidates; std::vector other_scores; const bool split = other_token_scores != nullptr && split_fraction > 0.0; if (split) { @@ -248,18 +246,18 @@ std::vector select_pflash_chunks( } } - const dflash::qwen3::PFlashSelectionPolicy policy{selector_budget, config.top_p, + const dflash::pflash::PFlashSelectionPolicy policy{selector_budget, config.top_p, /*skip_oversized=*/ segments != nullptr}; const auto selected = split - ? dflash::qwen3::select_pflash_split(candidates, other_candidates, policy, split_fraction, config.mode) - : dflash::qwen3::select_pflash_candidates(candidates, policy, config.mode); + ? dflash::pflash::select_pflash_split(candidates, other_candidates, policy, split_fraction, config.mode) + : dflash::pflash::select_pflash_candidates(candidates, policy, config.mode); if (!selected.ok) { set_last_error("PFlash selection failed: " + selected.error); std::fprintf(stderr, "[pflash-select] ERROR mode=%s budget=%d stop=%s: %s\n", - dflash::qwen3::pflash_selection_mode_name(config.mode), + dflash::pflash::pflash_selection_mode_name(config.mode), selector_budget, - dflash::qwen3::pflash_selection_stop_name(selected.stop), + dflash::pflash::pflash_selection_stop_name(selected.stop), selected.error.c_str()); std::fflush(stderr); return {}; @@ -290,12 +288,12 @@ std::vector select_pflash_chunks( std::fprintf(stderr, "[pflash-select] selected mode=%s scorer=%s segments=%s score=%s chunk=%d query=%d " "budget=%d selected_tokens=%zu chunks=%zu/%d stop=%s mass=%.9g\n", - dflash::qwen3::pflash_selection_mode_name(config.mode), + dflash::pflash::pflash_selection_mode_name(config.mode), split ? "split" : "single", segments ? "probe" : "fixed", density ? "density" : "sum", segments ? 0 : config.chunk_size, query_tokens, selector_budget, output.size(), selected.ordinals.size(), n_chunks, - dflash::qwen3::pflash_selection_stop_name(selected.stop), + dflash::pflash::pflash_selection_stop_name(selected.stop), selected.retained_mass); std::fflush(stderr); @@ -312,7 +310,7 @@ std::vector select_pflash_chunks( strict_fields.segments = segments; strict_fields.segmentation = segments ? "probe" : "fixed"; strict_fields.candidate_score = density ? "density" : "sum"; - strict_fields.scorer = split ? "split" : dflash::qwen3::pflash_scorer_name(config.scorer); + strict_fields.scorer = split ? "split" : dflash::pflash::pflash_scorer_name(config.scorer); strict_fields.split_fraction = split ? split_fraction : 0.0; strict_fields.other_chunk_scores = split ? &other_scores : nullptr; write_compression_trace( diff --git a/server/src/pflash/pflash_compress.h b/server/src/pflash/pflash_compress.h new file mode 100644 index 000000000..d97fcc024 --- /dev/null +++ b/server/src/pflash/pflash_compress.h @@ -0,0 +1,137 @@ +// PFlash scoring pipeline glue: everything between per-token scores and the +// compressed id list. +// +// - env_int / env_float process knobs used across the drafter +// - count_nonfinite_scores / scoring_head_mean_token_mass +// score post-processing helpers +// - PFlashTraceFields / write_compression_trace +// JSONL compression trace (DFLASH_PFLASH_TRACE_PATH) +// - select_pflash_chunks per-token scores -> candidates -> strict +// selection -> merged output ids (+ trace) + +#pragma once + +#include "pflash_selection.h" +#include "common/pflash_types.h" + +#include +#include +#include +#include +#include + +namespace dflash::common { + +int env_int(const char * name, int fallback); +float env_float(const char * name, float def); +void force_chunk_neighborhood(std::vector & forced, int n_chunks, + int chunk, int radius); + +struct QueryCaptureSlice { + int chunk_offset = 0; + int query_offset = 0; + int tokens = 0; + + bool valid() const { return tokens > 0; } +}; + +inline QueryCaptureSlice query_capture_slice( + int query_start, + int query_end, + int chunk_start, + int chunk_tokens) { + const int chunk_end = chunk_start + chunk_tokens; + const int overlap_start = query_start > chunk_start ? query_start : chunk_start; + const int overlap_end = query_end < chunk_end ? query_end : chunk_end; + if (overlap_start >= overlap_end) return {}; + return { + overlap_start - chunk_start, + overlap_start - query_start, + overlap_end - overlap_start, + }; +} + +inline size_t count_nonfinite_scores(const float * values, size_t count) { + size_t nonfinite = 0; + for (size_t index = 0; index < count; ++index) { + if (!std::isfinite(values[index])) ++nonfinite; + } + return nonfinite; +} + +// Scoring-head token mass: mean over heads and query tokens of softmax +// probabilities laid out as ggml [n_keys, n_queries, n_heads] (ne0 fastest). +inline void scoring_head_mean_token_mass( + const float * probs, + int n_keys, + int n_queries, + int n_heads, + std::vector & out) { + out.assign((size_t) n_keys, 0.0f); + if (n_keys <= 0 || n_queries <= 0 || n_heads <= 0) return; + std::vector sum((size_t) n_keys, 0.0); + for (int h = 0; h < n_heads; ++h) { + for (int t = 0; t < n_queries; ++t) { + const float * row = probs + ((size_t) h * n_queries + t) * n_keys; + for (int j = 0; j < n_keys; ++j) sum[(size_t) j] += row[j]; + } + } + const double denominator = (double) n_heads * (double) n_queries; + for (int j = 0; j < n_keys; ++j) out[(size_t) j] = (float) (sum[(size_t) j] / denominator); +} + +struct PFlashTraceFields { + const std::vector * input_ids = nullptr; + int query_begin = -1; + int query_end = -1; + dflash::pflash::PFlashSelectionMode selector_mode = + dflash::pflash::PFlashSelectionMode::Legacy; + dflash::pflash::PFlashQueryParser query_parser = + dflash::pflash::PFlashQueryParser::SemanticUser; + int token_budget = 0; + dflash::pflash::PFlashSelectionStop stop = + dflash::pflash::PFlashSelectionStop::InvalidInput; + int retained_tokens = 0; + double retained_mass = 0.0; + const std::vector * exact_chunk_scores = nullptr; + const std::vector * required_instruction_spans = nullptr; + // Variable-length candidates (segment probe): spans in candidate order. + const std::vector * segments = nullptr; + const char * segmentation = "fixed"; + const char * candidate_score = "sum"; + // Two-scorer selection: the other scorer's candidate scores, same order. + const char * scorer = "head"; + double split_fraction = 0.0; + const std::vector * other_chunk_scores = nullptr; +}; + +void write_compression_trace( + int input_tokens, + float keep_ratio, + int chunk_size, + int n_lookahead, + int pool_kernel, + int n_keep, + const std::vector> & chunk_means, + const std::vector & selected, + const std::vector & forced, + const std::vector & compressed_ids, + const PFlashTraceFields * trace_fields = nullptr); + +std::vector select_pflash_chunks( + const std::vector & ids, + const std::vector & token_scores, + float keep_ratio, + int n_lookahead, + int score_query_end, + int pool_kernel, + const dflash::pflash::PFlashSelectionConfig & config, + const std::vector & required_instruction_spans, + bool direct_mass, + bool write_trace, + const std::vector * segments = nullptr, + bool density = false, + const std::vector * other_token_scores = nullptr, + double split_fraction = 0.0); + +} // namespace dflash::common diff --git a/server/src/pflash/pflash_drafter.cpp b/server/src/pflash/pflash_drafter.cpp new file mode 100644 index 000000000..484ede122 --- /dev/null +++ b/server/src/pflash/pflash_drafter.cpp @@ -0,0 +1,162 @@ +// PFlash drafter entry points: load/free the Qwen3.5-0.8B scorer and run +// drafter_score_and_compress. +// +// Wires three pieces: +// - qwen35_loader.cpp : mmap GGUF + populate ggml tensors on backend, +// plus the optional scoring head and segment probe +// - qwen35_drafter.cpp : the block-15 head scorer and the all-layer +// running-max scorer +// - pflash_compress.cpp : score -> candidate -> strict selection + trace +// +// Single-pass forward over the first fifteen blocks on the Qwen3.5 target +// architecture (build_qwen35_layer); the block-15 NoPE Q/K projections score +// the context against the request's explicit query window. + +#include "pflash_drafter.h" + +#include "qwen35_drafter.h" +#include "pflash_selection.h" +#include "common/dspark_head.h" +#include "internal.h" + +#include "ggml.h" +#include "ggml-backend.h" + +#include +#include +#include + +namespace dflash::common { + +bool load_drafter(const std::string & gguf_path, int /*gpu_layers*/, + DrafterContext & out) { + return load_drafter(gguf_path, /*gpu_layers=*/999, /*gpu=*/0, out); +} + +bool load_drafter(const std::string & gguf_path, int /*gpu_layers*/, + int gpu, DrafterContext & out) { + if (gpu < 0) { + set_last_error("load_drafter: negative GPU index"); + return false; + } + if (out.loaded) { + set_last_error("drafter already loaded"); + return false; + } + if (out.backend && out.gpu >= 0 && out.gpu != gpu) { + set_last_error("load_drafter: backend already bound to a different GPU"); + return false; + } + + // If caller didn't supply a backend, spin up our own GPU backend. Sharing + // would be ideal but we don't have a handle to the daemon's backend + // through this API. Same-process GPU pools coexist fine; fragmentation is + // the only cost, and we free everything in free_drafter. + if (!out.backend) { + size_t n_dev = ggml_backend_dev_count(); + int seen_gpu = 0; + for (size_t i = 0; i < n_dev; ++i) { + ggml_backend_dev_t dev = ggml_backend_dev_get(i); + if (ggml_backend_dev_type(dev) == GGML_BACKEND_DEVICE_TYPE_GPU) { + if (seen_gpu == gpu) { + out.backend = ggml_backend_dev_init(dev, nullptr); + break; + } + seen_gpu++; + } + } + if (!out.backend) { + set_last_error("load_drafter: requested GPU backend unavailable"); + return false; + } + out.gpu = gpu; + } else if (out.gpu < 0) { + out.gpu = gpu; + } + + return load_qwen35_drafter(gguf_path, out); +} + +void free_drafter(DrafterContext & ctx) { + dspark_note_drafter_lifecycle(); + free_drafter_weights(ctx); + if (ctx.backend) { + ggml_backend_free(ctx.backend); + ctx.backend = nullptr; + } + ctx.gpu = -1; +} + +void free_drafter_weights(DrafterContext & ctx) { + if (ctx.state) { + free_qwen35_drafter_state(ctx); + } + ctx.loaded = false; +} + +std::vector drafter_score_and_compress( + DrafterContext & ctx, + const std::vector & ids, + float keep_ratio, + int chunk_size, + int n_lookahead, + int pool_kernel, + int score_query_end, + const std::vector & required_instruction_spans) { + if (!ctx.loaded) { + set_last_error("drafter not loaded"); + return {}; + } + + dflash::pflash::PFlashSelectionConfig experiment; + std::string experiment_error; + if (!dflash::pflash::resolve_pflash_selection( + (int) ids.size(), chunk_size, experiment, experiment_error)) { + set_last_error("invalid PFlash strict selection config: " + experiment_error); + std::fprintf(stderr, "[pflash-select] ERROR config: %s\n", + experiment_error.c_str()); + std::fflush(stderr); + return {}; + } + chunk_size = experiment.chunk_size; + if (!experiment.selection_active && !required_instruction_spans.empty()) { + set_last_error( + "PFlash instruction spans require strict budget selection"); + std::fprintf(stderr, + "[pflash-select] ERROR instruction spans require strict selection\n"); + std::fflush(stderr); + return {}; + } + if (experiment.selection_active) { + std::string span_error; + if (!dflash::pflash::validate_pflash_instruction_spans( + required_instruction_spans, (int) ids.size(), span_error)) { + set_last_error("invalid PFlash instruction spans: " + span_error); + std::fprintf(stderr, + "[pflash-select] ERROR instruction spans: %s\n", + span_error.c_str()); + std::fflush(stderr); + return {}; + } + } + if (experiment.configured) { + std::fprintf(stderr, + "[pflash-select] config mode=%s active=%d chunk=%d " + "query_parser=%s query_cap=%d query_actual=%d top_p=%.9g " + "input=%zu\n", + dflash::pflash::pflash_selection_mode_name(experiment.mode), + (int) experiment.selection_active, experiment.chunk_size, + dflash::pflash::pflash_query_parser_name(experiment.query_parser), + experiment.query_tokens, n_lookahead, experiment.top_p, ids.size()); + std::fflush(stderr); + } + if (score_query_end < 0) { + set_last_error("qwen35 scorer query window out of range"); + return {}; + } + return qwen35_drafter_score_and_compress( + ctx, ids, keep_ratio, chunk_size, n_lookahead, pool_kernel, + score_query_end, experiment, required_instruction_spans); +} + +} // namespace dflash::common diff --git a/server/src/qwen3/qwen3_drafter.h b/server/src/pflash/pflash_drafter.h similarity index 58% rename from server/src/qwen3/qwen3_drafter.h rename to server/src/pflash/pflash_drafter.h index b4a18bfd1..7749e437e 100644 --- a/server/src/qwen3/qwen3_drafter.h +++ b/server/src/pflash/pflash_drafter.h @@ -1,11 +1,14 @@ -// In-process Qwen3-0.6B drafter for pflash speculative prefill. +// In-process PFlash drafter for speculative prefill. +// +// The drafter is the Qwen3.5-0.8B scorer (qwen35_drafter.cpp + +// qwen35_loader.cpp): it runs the model's first fifteen blocks and scores +// the context with block 15's NoPE Q/K attention-mass head, with the +// all-layer running-max scorer kept as an opt-in alternative +// (PFLASH_QWEN35_LEGACY_SCORER=1 or the PFLASH scorer config). // // Hosted in the SAME process / SAME ggml allocator as the dflash target, so // we never pay the cross-process VRAM contention that broke the Python -// subprocess integration. Drafter uses our custom Qwen3-0.6B forward -// (qwen3_graph.cpp + qwen3_loader.cpp) which calls our FlashPrefill -// CUDA kernels for the attention compute, replacing libllama. This removes -// the dense O(S²) FA cost that made libllama 3+ minutes at 140K. +// subprocess integration. // // Public entry point: drafter_score_and_compress() takes raw input token IDs, // runs the full pflash compression pipeline in C++, returns the surviving @@ -20,31 +23,21 @@ #include #include -#include "qwen3_drafter_model.h" - struct ggml_backend; typedef struct ggml_backend * ggml_backend_t; namespace dflash::common { -enum class DrafterArch { - Qwen3_0p6b, - Qwen35_0p8b, -}; - -bool parse_drafter_arch(const std::string & name, DrafterArch & out); -const char * drafter_arch_name(DrafterArch arch); +struct Qwen35DrafterState; struct DrafterContext { ggml_backend_t backend = nullptr; // owned (created in load_drafter) - Qwen3DrafterWeights weights; // weights on the selected backend - DrafterArch arch = DrafterArch::Qwen3_0p6b; - void * arch_state = nullptr; - int gpu = -1; + Qwen35DrafterState * state = nullptr; // owned scorer weights + heads + int gpu = -1; bool loaded = false; }; -// Load the drafter GGUF (e.g. /opt/lucebox/models/drafter/Qwen3-0.6B-BF16.gguf). +// Load the drafter GGUF (a Qwen3.5-0.8B GGUF). // Creates a fresh GPU backend if `backend` is null. Otherwise uses the // caller-provided backend (so the drafter shares the daemon's allocator). // @@ -54,10 +47,6 @@ bool load_drafter(const std::string & gguf_path, int gpu_layers, DrafterContext & out); bool load_drafter(const std::string & gguf_path, int gpu_layers, int gpu, DrafterContext & out); -bool load_drafter(const std::string & gguf_path, int gpu_layers, - DrafterArch arch, DrafterContext & out); -bool load_drafter(const std::string & gguf_path, int gpu_layers, - DrafterArch arch, int gpu, DrafterContext & out); void free_drafter(DrafterContext & ctx); @@ -65,16 +54,17 @@ void free_drafter(DrafterContext & ctx); // Avoids repeated ggml backend create/destroy during daemon reuse. void free_drafter_weights(DrafterContext & ctx); -// Score importance per token via Liu Q-hook tail attention, then chunk-top-K -// span merge. Returns surviving token IDs (drafter vocab). +// Score the context with the block-15 scoring head, then run strict budget +// selection (or the configured selection mode). Returns surviving token IDs +// (drafter vocab). // // ids input token IDs of length S -// keep_ratio fraction of `chunk_size`-token chunks to keep +// keep_ratio fraction of the token budget to keep // chunk_size span granularity (default 32) // n_lookahead Q tokens used for scorer attention (default 8) // pool_kernel AvgPool kernel for score smoothing (default 13) -// score_query_end exclusive end of Q window in ids; negative means tail -// for Qwen3 and is rejected for Qwen3.5 +// score_query_end exclusive end of the scorer query window in ids; +// required (negative values are rejected) // // On failure returns empty vector + sets last_error. std::vector drafter_score_and_compress( diff --git a/server/src/qwen3/pflash_selection.cpp b/server/src/pflash/pflash_selection.cpp similarity index 99% rename from server/src/qwen3/pflash_selection.cpp rename to server/src/pflash/pflash_selection.cpp index 48901a814..f6d3fe6bf 100644 --- a/server/src/qwen3/pflash_selection.cpp +++ b/server/src/pflash/pflash_selection.cpp @@ -11,7 +11,7 @@ #include #include -namespace dflash::qwen3 { +namespace dflash::pflash { namespace { @@ -546,4 +546,4 @@ const char * pflash_candidate_score_name(PFlashCandidateScore score) noexcept { return "unknown"; } -} // namespace dflash::qwen3 +} // namespace dflash::pflash diff --git a/server/src/qwen3/pflash_selection.h b/server/src/pflash/pflash_selection.h similarity index 98% rename from server/src/qwen3/pflash_selection.h rename to server/src/pflash/pflash_selection.h index 9d4bfc4c0..a88e4c6f0 100644 --- a/server/src/qwen3/pflash_selection.h +++ b/server/src/pflash/pflash_selection.h @@ -6,7 +6,7 @@ #include #include -namespace dflash::qwen3 { +namespace dflash::pflash { enum class PFlashSelectionMode { Legacy, @@ -144,4 +144,4 @@ bool resolve_pflash_selection( PFlashSelectionConfig & out, std::string & error); -} // namespace dflash::qwen3 +} // namespace dflash::pflash diff --git a/server/src/qwen3/qwen35_drafter.cpp b/server/src/pflash/qwen35_drafter.cpp similarity index 97% rename from server/src/qwen3/qwen35_drafter.cpp rename to server/src/pflash/qwen35_drafter.cpp index 7a3e1582c..7f34d8b4c 100644 --- a/server/src/qwen3/qwen35_drafter.cpp +++ b/server/src/pflash/qwen35_drafter.cpp @@ -7,16 +7,16 @@ // Q/K scoring head, under strict // budget selection // -// Loading lives in qwen35_loader.cpp; qwen3_drafter.cpp dispatches into -// qwen35_drafter_score_and_compress on DrafterArch::Qwen35_0p8b. +// Loading lives in qwen35_loader.cpp; pflash_drafter.cpp dispatches into +// qwen35_drafter_score_and_compress. #include "qwen35_drafter.h" -#include "qwen3_drafter.h" -#include "qwen3_drafter_common.h" +#include "pflash_drafter.h" +#include "pflash_compress.h" #include "pflash_selection.h" #include "common/gguf_inspect.h" -#include "qwen3/anchor_params.h" +#include "anchor_params.h" #include "internal.h" #include "ggml.h" @@ -98,7 +98,7 @@ std::vector qwen35_score_and_compress( int n_lookahead, int pool_kernel, int score_query_end, - const dflash::qwen3::PFlashSelectionConfig & experiment, + const dflash::pflash::PFlashSelectionConfig & experiment, const std::vector & required_instruction_spans, std::vector * token_scores_out) { @@ -529,7 +529,7 @@ std::vector qwen35_strict_score_and_compress( float keep_ratio, int n_lookahead, int score_query_end, - const dflash::qwen3::PFlashSelectionConfig & experiment, + const dflash::pflash::PFlashSelectionConfig & experiment, const std::vector & required_instruction_spans, std::vector * token_mass_out, std::vector * segments_out, @@ -695,7 +695,7 @@ std::vector qwen35_strict_score_and_compress( ggml_tensor * logits = ggml_new_tensor_3d(lctx, GGML_TYPE_F32, S, n_lookahead, H); ggml_tensor * mask = ggml_new_tensor_2d(lctx, GGML_TYPE_F32, S, n_lookahead); const bool use_probe = st.probe_loaded && - experiment.segmentation != dflash::qwen3::PFlashSegmentation::Fixed; + experiment.segmentation != dflash::pflash::PFlashSegmentation::Fixed; ggml_tensor * probe_logits = use_probe ? ggml_new_tensor_1d(lctx, GGML_TYPE_F32, S) : nullptr; ggml_tensor * subunit_logits = use_probe && st.probe_sub_fc2_w @@ -842,7 +842,7 @@ std::vector qwen35_strict_score_and_compress( std::fflush(stderr); std::vector segments; - bool density = experiment.candidate_score == dflash::qwen3::PFlashCandidateScore::Density; + bool density = experiment.candidate_score == dflash::pflash::PFlashCandidateScore::Density; if (use_probe) { // Tap-count smoothing over the raw logits (torch Conv1d, symmetric // padding) plus the residual logit, then sigmoid: the boundary score @@ -881,9 +881,9 @@ std::vector qwen35_strict_score_and_compress( if (boundary[(size_t) t] > st.probe_threshold) ++boundaries_in_context; } const bool forced_probe = - experiment.segmentation == dflash::qwen3::PFlashSegmentation::Probe; + experiment.segmentation == dflash::pflash::PFlashSegmentation::Probe; if (boundaries_in_context >= 4 || forced_probe) { - segments = dflash::qwen3::pflash_probe_segments( + segments = dflash::pflash::pflash_probe_segments( boundary, S, st.probe_threshold, st.probe_min_segment, st.probe_max_segment, forced, split_scores); } @@ -893,7 +893,7 @@ std::vector qwen35_strict_score_and_compress( "falling back to fixed %d-token chunks\n", boundaries_in_context, experiment.chunk_size); } else { - if (experiment.candidate_score == dflash::qwen3::PFlashCandidateScore::Auto) { + if (experiment.candidate_score == dflash::pflash::PFlashCandidateScore::Auto) { density = true; } std::fprintf(stderr, @@ -930,21 +930,21 @@ std::vector qwen35_drafter_score_and_compress( int n_lookahead, int pool_kernel, int score_query_end, - const dflash::qwen3::PFlashSelectionConfig & experiment, + const dflash::pflash::PFlashSelectionConfig & experiment, const std::vector & required_instruction_spans) { - if (!ctx.arch_state) { + if (!ctx.state) { set_last_error("qwen35 drafter state missing"); return {}; } - auto * st = static_cast(ctx.arch_state); + auto * st = static_cast(ctx.state); // Strict budget selection scores with the block-15 head; the // legacy all-layer running-max scorer stays available for legacy // selection or when PFLASH_QWEN35_LEGACY_SCORER=1 forces it. const char * legacy_scorer = std::getenv("PFLASH_QWEN35_LEGACY_SCORER"); const bool force_legacy = (legacy_scorer && std::string(legacy_scorer) == "1") || - experiment.scorer == dflash::qwen3::PFlashScorer::Legacy; + experiment.scorer == dflash::pflash::PFlashScorer::Legacy; if (experiment.selection_active && - experiment.scorer == dflash::qwen3::PFlashScorer::Split) { + experiment.scorer == dflash::pflash::PFlashScorer::Split) { // Two scorers, one budget: the block-15 head ranks (and segments) // first, the all-layer running-max scorer fills the remainder. std::vector head_mass; diff --git a/server/src/qwen3/qwen35_drafter.h b/server/src/pflash/qwen35_drafter.h similarity index 80% rename from server/src/qwen3/qwen35_drafter.h rename to server/src/pflash/qwen35_drafter.h index 2a391118c..434353612 100644 --- a/server/src/qwen3/qwen35_drafter.h +++ b/server/src/pflash/qwen35_drafter.h @@ -1,14 +1,13 @@ -// Internal interface of the Qwen3.5-0.8B drafter. +// Internal interface of the Qwen3.5-0.8B drafter — the only pflash scorer. // -// The Qwen3.5-0.8B scorer runs on the Qwen3.5 target architecture -// (TargetWeights, build_qwen35_layer) rather than the Qwen3-0.6B drafter -// graph, so it lives in its own translation units: qwen35_loader.cpp loads -// the GGUF, the scoring head and the segment probe; qwen35_drafter.cpp runs -// the two scorers. qwen3_drafter.cpp dispatches here on DrafterArch. +// The scorer runs on the Qwen3.5 target architecture (TargetWeights, +// build_qwen35_layer): qwen35_loader.cpp loads the GGUF, the optional +// scoring head and the segment probe; qwen35_drafter.cpp runs the two +// scorers. pflash_drafter.cpp owns the public entry points. #pragma once -#include "qwen3_drafter.h" +#include "pflash_drafter.h" #include "pflash_selection.h" #include "common/pflash_types.h" #include "internal.h" @@ -24,8 +23,8 @@ namespace dflash::common { // Qwen3.5-0.8B scoring head. Features are the residual entering // full-attention block 15 after the first 15 blocks (twelve GatedDeltaNet and // three full-attention blocks). Block 15's own Q/K projections score the -// context without RoPE, exactly like the Qwen3-0.6B block-13 head; an -// optional trained head replaces those two projections. +// context without RoPE; an optional trained head replaces those two +// projections. static constexpr int kQwen35HeadBlock = 15; struct Qwen35DrafterState { @@ -61,8 +60,7 @@ struct Qwen35DrafterState { // Defined in qwen35_loader.cpp. bool qwen35_head_block_available(const TargetWeights & w, std::string & error); -bool load_qwen35_drafter(const std::string & gguf_path, DrafterArch arch, - DrafterContext & out); +bool load_qwen35_drafter(const std::string & gguf_path, DrafterContext & out); void free_qwen35_drafter_state(DrafterContext & ctx); // Defined in qwen35_drafter.cpp. @@ -76,7 +74,7 @@ std::vector qwen35_score_and_compress( int n_lookahead, int pool_kernel, int score_query_end, - const dflash::qwen3::PFlashSelectionConfig & experiment, + const dflash::pflash::PFlashSelectionConfig & experiment, const std::vector & required_instruction_spans, std::vector * token_scores_out = nullptr); @@ -87,7 +85,7 @@ std::vector qwen35_strict_score_and_compress( float keep_ratio, int n_lookahead, int score_query_end, - const dflash::qwen3::PFlashSelectionConfig & experiment, + const dflash::pflash::PFlashSelectionConfig & experiment, const std::vector & required_instruction_spans, std::vector * token_mass_out = nullptr, std::vector * segments_out = nullptr, @@ -102,7 +100,7 @@ std::vector qwen35_drafter_score_and_compress( int n_lookahead, int pool_kernel, int score_query_end, - const dflash::qwen3::PFlashSelectionConfig & experiment, + const dflash::pflash::PFlashSelectionConfig & experiment, const std::vector & required_instruction_spans); } // namespace dflash::common diff --git a/server/src/qwen3/qwen35_loader.cpp b/server/src/pflash/qwen35_loader.cpp similarity index 96% rename from server/src/qwen3/qwen35_loader.cpp rename to server/src/pflash/qwen35_loader.cpp index 099c87312..a7012021b 100644 --- a/server/src/qwen3/qwen35_loader.cpp +++ b/server/src/pflash/qwen35_loader.cpp @@ -1,13 +1,11 @@ // Qwen3.5-0.8B drafter loading: the drafter GGUF, the optional trained // block-15 scoring head and the optional segment probe. // -// The Qwen3-0.6B drafter has qwen3_loader.cpp; this is its counterpart for -// the Qwen3.5-0.8B scorer, which is built on the Qwen3.5 target weights -// (load_target_gguf_partial) instead of the Qwen3-0.6B drafter weights. +// The scorer is built on the Qwen3.5 target weights +// (load_target_gguf_partial). #include "qwen35_drafter.h" -#include "qwen3_drafter.h" #include "common/gguf_inspect.h" #include "internal.h" @@ -69,7 +67,7 @@ static bool qwen35_metadata_equals(gguf_context * g, const char * key, } // Optional trained head for the block-15 tap. Fails closed on any contract -// mismatch, mirroring the Qwen3-0.6B head loader. +// mismatch. static bool load_qwen35_scoring_head(const std::string & path, Qwen35DrafterState & st) { const TargetWeights & w = st.weights; @@ -305,7 +303,7 @@ static bool load_qwen35_segment_probe(const std::string & path, } // namespace -bool load_qwen35_drafter(const std::string & gguf_path, DrafterArch arch, +bool load_qwen35_drafter(const std::string & gguf_path, DrafterContext & out) { auto * st = new Qwen35DrafterState(); // The scorer never needs logits, and tied-embedding Qwen3.5-0.8B @@ -350,13 +348,11 @@ bool load_qwen35_drafter(const std::string & gguf_path, DrafterArch arch, return false; } } - out.arch_state = st; + out.state = st; out.loaded = true; - out.arch = arch; std::fprintf(stderr, - "[drafter] loaded %s qwen35: n_layer=%d n_head=%d n_head_kv=%d " + "[drafter] loaded qwen35: n_layer=%d n_head=%d n_head_kv=%d " "n_embd=%d n_ff=%d head_dim=%d vocab=%d gpu=%d\n", - drafter_arch_name(arch), st->weights.n_layer, st->weights.n_head, st->weights.n_head_kv, st->weights.n_embd, st->weights.n_ff, st->weights.n_embd_head_k, st->weights.n_vocab, out.gpu); @@ -365,12 +361,12 @@ bool load_qwen35_drafter(const std::string & gguf_path, DrafterArch arch, } void free_qwen35_drafter_state(DrafterContext & ctx) { - auto * st = static_cast(ctx.arch_state); + auto * st = static_cast(ctx.state); free_qwen35_head(*st); free_qwen35_segment_probe(*st); free_target_weights(st->weights); delete st; - ctx.arch_state = nullptr; + ctx.state = nullptr; } } // namespace dflash::common diff --git a/server/src/qwen3/qwen3_backend.cpp b/server/src/qwen3/qwen3_backend.cpp index 6d28cba89..42d276ca3 100644 --- a/server/src/qwen3/qwen3_backend.cpp +++ b/server/src/qwen3/qwen3_backend.cpp @@ -6,7 +6,7 @@ // After all layers, out_norm + lm_head produces logits for the last token. #include "qwen3_backend.h" -#include "qwen3_drafter.h" +#include "pflash/pflash_drafter.h" #include "dflash27b.h" #include "common/sampler.h" #include "common/io_utils.h" @@ -25,7 +25,7 @@ namespace dflash::common { // ── Cache management ─────────────────────────────────────────────────── -bool create_qwen3_cache(ggml_backend_t backend, const Qwen3DrafterWeights & w, +bool create_qwen3_cache(ggml_backend_t backend, const Qwen3Weights & w, int max_ctx, Qwen3Cache & out) { const int n_layer = w.n_layer; const int D = w.head_dim; @@ -91,7 +91,7 @@ bool Qwen3Backend::init() { return false; } - if (!load_qwen3_drafter_model(cfg_.model_path, backend_, w_)) { + if (!load_qwen3_model(cfg_.model_path, backend_, w_)) { std::fprintf(stderr, "[qwen3] model load failed: %s\n", dflash27b_last_error()); return false; } @@ -145,8 +145,8 @@ bool Qwen3Backend::unpark(ParkTarget target) { if (target == ParkTarget::TargetModel || target == ParkTarget::All) { if (parked_) { // Reload weights - Qwen3DrafterWeights w_new; - if (!load_qwen3_drafter_model(cfg_.model_path, backend_, w_new)) { + Qwen3Weights w_new; + if (!load_qwen3_model(cfg_.model_path, backend_, w_new)) { std::fprintf(stderr, "[qwen3] unpark reload failed\n"); return false; } @@ -969,10 +969,14 @@ ModelBackend::CompressResult Qwen3Backend::compress(const CompressRequest & req) drafter_loaded_ = true; } + // score_query_end < 0 is the legacy "tail window" request value; the + // qwen35 scorer requires an explicit end. + const int score_query_end = req.score_query_end >= 0 + ? req.score_query_end : (int)req.input_ids.size(); result = CompressResult::from_compressed_ids(drafter_score_and_compress( drafter_ctx_, req.input_ids, req.keep_ratio, /*chunk_size=*/32, req.score_query_tokens, /*pool_kernel=*/13, - req.score_query_end, req.required_instruction_spans)); + score_query_end, req.required_instruction_spans)); if (req.residency_action == DraftResidencyAction::ReleaseAfterUse) { free_drafter(); @@ -1029,7 +1033,9 @@ bool Qwen3Backend::handle_compress(const std::string & line, const DaemonIO & io } const float keep = (float)keep_x1000 / 1000.0f; - auto compressed = drafter_score_and_compress(drafter_ctx_, src_ids, keep); + auto compressed = drafter_score_and_compress(drafter_ctx_, src_ids, keep, + /*chunk_size=*/32, /*n_lookahead=*/8, /*pool_kernel=*/13, + (int)src_ids.size()); std::printf("[compress] %zu -> %zu tokens\n", src_ids.size(), compressed.size()); std::fflush(stdout); @@ -1063,7 +1069,7 @@ void Qwen3Backend::shutdown() { } free_qwen3_cache(cache_); if (!parked_) { - free_qwen3_drafter_model(w_); + free_qwen3_model(w_); } if (backend_) { ggml_backend_free(backend_); diff --git a/server/src/qwen3/qwen3_backend.h b/server/src/qwen3/qwen3_backend.h index d8d4c9668..ba1edd1b3 100644 --- a/server/src/qwen3/qwen3_backend.h +++ b/server/src/qwen3/qwen3_backend.h @@ -1,10 +1,10 @@ // Qwen3Backend — ModelBackend for the Qwen3-0.6B model used as a standalone -// inference backend (not just as a pflash drafter). +// inference backend. // // Architecture: 28-layer transformer, 16 heads (8 KV), hidden=1024, vocab=151936. // Sliding-window attention (FA_WINDOW=512), standard RoPE. // -// This backend reuses the Qwen3DrafterWeights loader but adds: +// This backend reuses the Qwen3Weights loader but adds: // - Persistent KV cache for incremental decode // - Step-based forward (prefill chunks + single-token decode) // - Logits output via out_norm + lm_head @@ -13,8 +13,8 @@ #include "common/model_backend.h" #include "placement/placement_config.h" -#include "qwen3_drafter_model.h" -#include "qwen3_drafter.h" +#include "qwen3_model.h" +#include "pflash/pflash_drafter.h" #include "common/sampler.h" #include "ggml.h" @@ -48,7 +48,7 @@ struct Qwen3Cache { ggml_backend_buffer_t buf = nullptr; }; -bool create_qwen3_cache(ggml_backend_t backend, const Qwen3DrafterWeights & w, +bool create_qwen3_cache(ggml_backend_t backend, const Qwen3Weights & w, int max_ctx, Qwen3Cache & out); void free_qwen3_cache(Qwen3Cache & c); @@ -110,7 +110,7 @@ class Qwen3Backend : public ModelBackend { private: Qwen3BackendConfig cfg_; ggml_backend_t backend_ = nullptr; - Qwen3DrafterWeights w_; + Qwen3Weights w_; Qwen3Cache cache_; bool parked_ = false; diff --git a/server/src/qwen3/qwen3_buffer_plan.h b/server/src/qwen3/qwen3_buffer_plan.h deleted file mode 100644 index a3d632520..000000000 --- a/server/src/qwen3/qwen3_buffer_plan.h +++ /dev/null @@ -1,29 +0,0 @@ -#pragma once - -#include - -namespace dflash::common { - -struct Qwen3DrafterBufferPlan { - std::size_t rope_k_buffers; - std::size_t value_buffers; - std::size_t rope_q_tail_buffers; - bool reuse_current_layer_kv; - - std::size_t layer_cache_index(int layer) const { - return reuse_current_layer_kv ? 0u : static_cast(layer); - } -}; - -inline Qwen3DrafterBufferPlan qwen3_drafter_buffer_plan( - bool nope_tail, int n_layer) { - const std::size_t layers = n_layer > 0 ? (std::size_t)n_layer : 0u; - return { - nope_tail ? (layers > 0 ? 1u : 0u) : layers, - layers > 0 ? 1u : 0u, - nope_tail ? 0u : layers, - nope_tail, - }; -} - -} // namespace dflash::common diff --git a/server/src/qwen3/qwen3_drafter.cpp b/server/src/qwen3/qwen3_drafter.cpp deleted file mode 100644 index da83be911..000000000 --- a/server/src/qwen3/qwen3_drafter.cpp +++ /dev/null @@ -1,463 +0,0 @@ -// Qwen3-0.6B drafter for pflash speculative prefill, hosted in-process. -// -// Wires three pieces: -// - qwen3_loader.cpp : mmap GGUF + populate ggml tensors on backend -// - qwen3_graph.cpp : custom forward (per-layer ggml + FP CUDA kernel) -// - qwen3_drafter_common.cpp : chunk-top-K + span merge, shared with the -// Qwen3.5-0.8B drafter -// -// The Qwen3.5-0.8B drafter lives in qwen35_loader.cpp / qwen35_drafter.cpp; -// this file dispatches to it on DrafterArch. -// -// Single-pass forward at full S using a custom Qwen3-0.6B graph with the -// FlashPrefill block-sparse attention kernel (or BSA when enabled). Tail -// attention scoring runs in a separate post-forward graph using saved Q_last -// and K_curr per layer. -// -// Result running_max [n_lookahead, S] f32 is reduced to per-token scores via -// mean-over-lookahead, smoothed with AvgPool, scored per chunk, top-K kept. - -#include "qwen3_drafter.h" -#include "common/dspark_head.h" -#include "qwen3_drafter_model.h" -#include "qwen3_drafter_common.h" -#include "qwen35_drafter.h" -#include "pflash_selection.h" -#include "qwen3/anchor_params.h" -#include "common/backend_precision.h" -#include "common/gguf_inspect.h" -#include "internal.h" -#include "anchor_scan.h" - -#include "ggml.h" -#include "ggml-alloc.h" -#include "ggml-backend.h" -#include "gguf.h" - -#include -#include -#include -#include -#include -#include -#include -#include -#include - -namespace dflash::common { - -namespace { - -#if defined(DFLASH27B_BACKEND_HIP) -bool prewarm_drafter_once(const Qwen3DrafterWeights & w) { - static bool warmed = false; - if (warmed || std::getenv("DFLASH_FP_SKIP_PREWARM")) { - return true; - } - - const int warm_tokens = 1024; - const int n_lookahead = 8; - std::vector ids((size_t)warm_tokens, 0); - std::vector running_max; - - auto t0 = std::chrono::steady_clock::now(); - bool ok = forward_qwen3_drafter_model(w, ids, n_lookahead, running_max); - auto t1 = std::chrono::steady_clock::now(); - if (!ok) { - return false; - } - - std::fprintf(stderr, "[drafter] HIP prewarm %.2fs (%d tokens)\n", - std::chrono::duration(t1 - t0).count(), warm_tokens); - std::fflush(stderr); - warmed = true; - return true; -} -#endif - -} // namespace - -bool parse_drafter_arch(const std::string & name, DrafterArch & out) { - if (name == "qwen3-0.6b" || name == "qwen3_0p6b" || name == "qwen3") { - out = DrafterArch::Qwen3_0p6b; - return true; - } - if (name == "qwen35-0.8b" || name == "qwen3.5-0.8b" || name == "qwen35_0p8b" || name == "qwen35") { - out = DrafterArch::Qwen35_0p8b; - return true; - } - return false; -} - -const char * drafter_arch_name(DrafterArch arch) { - switch (arch) { - case DrafterArch::Qwen3_0p6b: return "qwen3-0.6b"; - case DrafterArch::Qwen35_0p8b: return "qwen35-0.8b"; - } - return "unknown"; -} - -bool load_drafter(const std::string & gguf_path, int /*gpu_layers*/, - DrafterContext & out) { - return load_drafter(gguf_path, /*gpu_layers=*/999, /*gpu=*/0, out); -} - -bool load_drafter(const std::string & gguf_path, int /*gpu_layers*/, - int gpu, DrafterContext & out) { - DrafterArch arch = DrafterArch::Qwen3_0p6b; - { - std::string lower = gguf_path; - for (auto & c : lower) c = (char)std::tolower((unsigned char)c); - if (lower.find("qwen3.5") != std::string::npos || - lower.find("qwen35") != std::string::npos) { - arch = DrafterArch::Qwen35_0p8b; - } - } - return load_drafter(gguf_path, /*gpu_layers=*/999, arch, gpu, out); -} - -bool load_drafter(const std::string & gguf_path, int /*gpu_layers*/, - DrafterArch arch, DrafterContext & out) { - return load_drafter(gguf_path, /*gpu_layers=*/999, arch, /*gpu=*/0, out); -} - -bool load_drafter(const std::string & gguf_path, int /*gpu_layers*/, - DrafterArch arch, int gpu, DrafterContext & out) { - if (gpu < 0) { - set_last_error("load_drafter: negative GPU index"); - return false; - } - if (out.loaded) { - set_last_error("drafter already loaded"); - return false; - } - if (out.backend && out.gpu >= 0 && out.gpu != gpu) { - set_last_error("load_drafter: backend already bound to a different GPU"); - return false; - } - - // If caller didn't supply a backend, spin up our own GPU backend. Sharing - // would be ideal but we don't have a handle to the daemon's backend - // through this API. Same-process GPU pools coexist fine; fragmentation is - // the only cost, and we free everything in free_drafter. - if (!out.backend) { - size_t n_dev = ggml_backend_dev_count(); - int seen_gpu = 0; - for (size_t i = 0; i < n_dev; ++i) { - ggml_backend_dev_t dev = ggml_backend_dev_get(i); - if (ggml_backend_dev_type(dev) == GGML_BACKEND_DEVICE_TYPE_GPU) { - if (seen_gpu == gpu) { - out.backend = ggml_backend_dev_init(dev, nullptr); - break; - } - seen_gpu++; - } - } - if (!out.backend) { - set_last_error("load_drafter: requested GPU backend unavailable"); - return false; - } - out.gpu = gpu; - } else if (out.gpu < 0) { - out.gpu = gpu; - } - - if (arch == DrafterArch::Qwen35_0p8b) { - return load_qwen35_drafter(gguf_path, arch, out); - } - - if (!load_qwen3_drafter_model(gguf_path, out.backend, out.weights)) { - // last_error already set by loader - return false; - } - - out.loaded = true; - out.arch = arch; - std::fprintf(stderr, - "[drafter] loaded %s weights=%s compute=%s: n_layer=%d n_head=%d n_kv=%d " - "n_embd=%d n_ff=%d head_dim=%d vocab=%d gpu=%d\n", - drafter_arch_name(arch), - backend_precision_type_name(out.weights.weight_type), - backend_precision_type_name(out.weights.compute_type), - out.weights.n_layer, out.weights.n_head, out.weights.n_head_kv, - out.weights.n_embd, out.weights.n_ff, out.weights.head_dim, - out.weights.n_vocab, out.gpu); - std::fflush(stderr); - -#if defined(DFLASH27B_BACKEND_HIP) - if (!prewarm_drafter_once(out.weights)) { - free_drafter(out); - return false; - } -#endif - - return true; -} - -void free_drafter(DrafterContext & ctx) { - dspark_note_drafter_lifecycle(); - free_drafter_weights(ctx); - if (ctx.backend) { - ggml_backend_free(ctx.backend); - ctx.backend = nullptr; - } - ctx.gpu = -1; -} - -void free_drafter_weights(DrafterContext & ctx) { - if (ctx.arch == DrafterArch::Qwen35_0p8b && ctx.arch_state) { - free_qwen35_drafter_state(ctx); - } - if (ctx.loaded) { - if (ctx.arch == DrafterArch::Qwen3_0p6b) { - free_qwen3_drafter_model(ctx.weights); - } - } - ctx.loaded = false; -} - -std::vector drafter_score_and_compress( - DrafterContext & ctx, - const std::vector & ids, - float keep_ratio, - int chunk_size, - int n_lookahead, - int pool_kernel, - int score_query_end, - const std::vector & required_instruction_spans) { - if (!ctx.loaded) { - set_last_error("drafter not loaded"); - return {}; - } - - dflash::qwen3::PFlashSelectionConfig experiment; - std::string experiment_error; - if (!dflash::qwen3::resolve_pflash_selection( - (int) ids.size(), chunk_size, experiment, experiment_error)) { - set_last_error("invalid PFlash strict selection config: " + experiment_error); - std::fprintf(stderr, "[pflash-select] ERROR config: %s\n", - experiment_error.c_str()); - std::fflush(stderr); - return {}; - } - chunk_size = experiment.chunk_size; - if (!experiment.selection_active && !required_instruction_spans.empty()) { - set_last_error( - "PFlash instruction spans require strict budget selection"); - std::fprintf(stderr, - "[pflash-select] ERROR instruction spans require strict selection\n"); - std::fflush(stderr); - return {}; - } - if (experiment.selection_active) { - std::string span_error; - if (!dflash::qwen3::validate_pflash_instruction_spans( - required_instruction_spans, (int) ids.size(), span_error)) { - set_last_error("invalid PFlash instruction spans: " + span_error); - std::fprintf(stderr, - "[pflash-select] ERROR instruction spans: %s\n", - span_error.c_str()); - std::fflush(stderr); - return {}; - } - } - if (experiment.configured) { - std::fprintf(stderr, - "[pflash-select] config mode=%s active=%d chunk=%d " - "query_parser=%s query_cap=%d query_actual=%d top_p=%.9g " - "input=%zu\n", - dflash::qwen3::pflash_selection_mode_name(experiment.mode), - (int) experiment.selection_active, experiment.chunk_size, - dflash::qwen3::pflash_query_parser_name(experiment.query_parser), - experiment.query_tokens, n_lookahead, experiment.top_p, ids.size()); - std::fflush(stderr); - } - if (ctx.arch == DrafterArch::Qwen35_0p8b) { - if (score_query_end < 0) { - set_last_error("qwen35 scorer query window out of range"); - return {}; - } - return qwen35_drafter_score_and_compress( - ctx, ids, keep_ratio, chunk_size, n_lookahead, pool_kernel, - score_query_end, experiment, required_instruction_spans); - } - const int S = (int)ids.size(); - if (S < n_lookahead + 1) { - // Too short to score — return as-is. - return ids; - } - - // ── 1. Custom forward + GPU tail-attention scoring ──────────────── - auto t0 = std::chrono::steady_clock::now(); - std::vector running_max; - if (!forward_qwen3_drafter_model( - ctx.weights, ids, n_lookahead, running_max, score_query_end)) { - return {}; - } - auto t1 = std::chrono::steady_clock::now(); - std::fprintf(stderr, "[drafter] forward+score in %.2fs S=%d\n", - std::chrono::duration(t1 - t0).count(), S); - std::fflush(stderr); - - // ── 2. Mean over lookahead → per-token score [S] ────────────────── - std::vector score((size_t)S, 0.0f); - for (int j = 0; j < S; ++j) { - float s = 0.0f; - for (int t = 0; t < n_lookahead; ++t) { - s += running_max[(size_t)t * S + j]; - } - score[j] = s / (float)n_lookahead; - } - - // ── 3. AvgPool 1D smoothing ─────────────────────────────────────── - std::vector smooth((size_t)S, 0.0f); - int half = pool_kernel / 2; - for (int j = 0; j < S; ++j) { - int lo = std::max(0, j - half); - int hi = std::min(S - 1, j + half); - float s = 0.0f; - int n = 0; - for (int k = lo; k <= hi; ++k) { s += score[k]; ++n; } - smooth[j] = (n > 0) ? (s / (float)n) : 0.0f; - } - - if (experiment.selection_active) { - return select_pflash_chunks( - ids, ctx.weights.scoring_head_loaded ? score : smooth, - keep_ratio, n_lookahead, score_query_end, - ctx.weights.scoring_head_loaded ? 1 : pool_kernel, - experiment, required_instruction_spans, - ctx.weights.scoring_head_loaded, true); - } - - // ── 4. Chunk-top-K + span merge ─────────────────────────────────── - int n_chunks = (S + chunk_size - 1) / chunk_size; - int n_keep = std::max(1, (int)((float)n_chunks * keep_ratio)); - std::vector> chunk_means; - chunk_means.reserve((size_t)n_chunks); - for (int c = 0; c < n_chunks; ++c) { - int s_ = c * chunk_size; - int e_ = std::min(S, (c + 1) * chunk_size); - float m = 0.0f; - for (int j = s_; j < e_; ++j) m += smooth[j]; - m /= std::max(1, e_ - s_); - chunk_means.push_back({m, c}); - } - std::sort(chunk_means.begin(), chunk_means.end(), - [](auto a, auto b) { return a.first > b.first; }); - - // Retrieval tasks often repeat a rare key in the final query and in the - // needle span. Exact scores alone can keep the query while dropping the - // neighboring answer chunk, so force a small token-only anchor neighborhood. - // Head/tail forced chunks scale with n_keep so top-K scoring always gets slots. - const int h_raw = env_int("DFLASH_COMPRESS_HEAD_CHUNKS", 8); - const int t_raw = env_int("DFLASH_COMPRESS_TAIL_CHUNKS", 24); - int head_chunks = h_raw, tail_chunks = t_raw; - if (head_chunks + tail_chunks >= n_keep) { - const int budget = std::max(1, n_keep - 1); - head_chunks = std::max(0, h_raw * budget / (h_raw + t_raw)); - tail_chunks = std::max(0, budget - head_chunks); - } - const int query_tokens = env_int("DFLASH_COMPRESS_QUERY_TOKENS", 96); - const auto ap = resolve_anchor_params(n_chunks, - env_int("PFLASH_COMPRESS_ANCHOR_RADIUS", -1), - env_int("PFLASH_COMPRESS_MAX_ANCHOR_HITS", -1), - env_int("DFLASH_COMPRESS_ANCHOR_RADIUS", -1), - env_int("DFLASH_COMPRESS_MAX_ANCHOR_HITS", -1)); - const int anchor_radius = ap.radius; - const int max_anchor_hits = ap.max_hits; - std::vector selected_mask((size_t)n_chunks, 0); - std::vector forced((size_t)n_chunks, 0); - for (int c = 0; c < std::min(n_chunks, head_chunks); ++c) forced[(size_t)c] = 1; - for (int c = std::max(0, n_chunks - tail_chunks); c < n_chunks; ++c) forced[(size_t)c] = 1; - - const int q0 = std::max(0, S - query_tokens); - constexpr int NGRAM = 4; - for (int q = q0; q + NGRAM <= S; ++q) { - int hits = 0; - std::vector hit_pos(max_anchor_hits); - const int search_end = std::max(0, q0 - NGRAM); - for (int p = 0; p <= search_end && hits <= max_anchor_hits; ++p) { - bool same = true; - for (int k = 0; k < NGRAM; ++k) { - if (ids[(size_t)p + k] != ids[(size_t)q + k]) { same = false; break; } - } - if (same) { - if (hits < max_anchor_hits) hit_pos[hits] = p; - ++hits; - } - } - if (hits > 0 && hits <= max_anchor_hits) { - for (int i = 0; i < hits && i < max_anchor_hits; ++i) { - force_chunk_neighborhood(forced, n_chunks, hit_pos[i] / chunk_size, anchor_radius); - } - } - } - - int selected_count = 0; - int forced_count = 0; - for (int c = 0; c < n_chunks; ++c) { - if (forced[(size_t)c]) { - selected_mask[(size_t)c] = 1; - ++selected_count; - ++forced_count; - } - } - for (const auto & cm : chunk_means) { - if (selected_count >= n_keep) break; - int c = cm.second; - if (!selected_mask[(size_t)c]) { - selected_mask[(size_t)c] = 1; - ++selected_count; - } - } - - std::vector selected; - selected.reserve((size_t)selected_count); - for (int c = 0; c < n_chunks; ++c) { - if (selected_mask[(size_t)c]) selected.push_back(c); - } - - std::vector out; - out.reserve((size_t)n_keep * chunk_size + 16); - int span_start = -1, span_end = -1; - for (int c : selected) { - int s_ = c * chunk_size; - int e_ = std::min(S, (c + 1) * chunk_size); - if (span_start < 0) { - span_start = s_; span_end = e_; - } else if (s_ == span_end) { - span_end = e_; - } else { - for (int j = span_start; j < span_end; ++j) out.push_back(ids[j]); - span_start = s_; span_end = e_; - } - } - if (span_start >= 0) { - for (int j = span_start; j < span_end; ++j) out.push_back(ids[j]); - } - - auto t2 = std::chrono::steady_clock::now(); - std::fprintf(stderr, - "[drafter] score_and_compress total %.2fs S=%d kept=%zu (%d/%d chunks, forced=%d)\n", - std::chrono::duration(t2 - t0).count(), - S, out.size(), (int)selected.size(), n_chunks, forced_count); - std::fflush(stderr); - - const int query_end = score_query_end < 0 ? S : score_query_end; - const int query_begin = query_end - n_lookahead; - const int token_budget = (int) std::floor( - (double) S * (double) keep_ratio); - const PFlashTraceFields trace_fields{ - &ids, query_begin, query_end, experiment.mode, - experiment.query_parser, token_budget, - dflash::qwen3::PFlashSelectionStop::InvalidInput, (int) out.size(), - 0.0, nullptr}; - write_compression_trace(S, keep_ratio, chunk_size, n_lookahead, - pool_kernel, n_keep, chunk_means, selected_mask, forced, out, - &trace_fields); - - return out; -} - -} // namespace dflash::common diff --git a/server/src/qwen3/qwen3_drafter_common.h b/server/src/qwen3/qwen3_drafter_common.h deleted file mode 100644 index 743c8a3fc..000000000 --- a/server/src/qwen3/qwen3_drafter_common.h +++ /dev/null @@ -1,77 +0,0 @@ -// Helpers shared by the Qwen3-0.6B and Qwen3.5-0.8B drafter paths. -// -// Moved verbatim out of qwen3_drafter.cpp so qwen35_drafter.cpp can use the -// same selector, trace writer and environment readers without a second copy. - -#pragma once - -#include "pflash_selection.h" -#include "common/pflash_types.h" - -#include -#include -#include -#include - -namespace dflash::common { - -int env_int(const char * name, int fallback); -float env_float(const char * name, float def); -void force_chunk_neighborhood(std::vector & forced, int n_chunks, - int chunk, int radius); - -struct PFlashTraceFields { - const std::vector * input_ids = nullptr; - int query_begin = -1; - int query_end = -1; - dflash::qwen3::PFlashSelectionMode selector_mode = - dflash::qwen3::PFlashSelectionMode::Legacy; - dflash::qwen3::PFlashQueryParser query_parser = - dflash::qwen3::PFlashQueryParser::SemanticUser; - int token_budget = 0; - dflash::qwen3::PFlashSelectionStop stop = - dflash::qwen3::PFlashSelectionStop::InvalidInput; - int retained_tokens = 0; - double retained_mass = 0.0; - const std::vector * exact_chunk_scores = nullptr; - const std::vector * required_instruction_spans = nullptr; - // Variable-length candidates (segment probe): spans in candidate order. - const std::vector * segments = nullptr; - const char * segmentation = "fixed"; - const char * candidate_score = "sum"; - // Two-scorer selection: the other scorer's candidate scores, same order. - const char * scorer = "head"; - double split_fraction = 0.0; - const std::vector * other_chunk_scores = nullptr; -}; - -void write_compression_trace( - int input_tokens, - float keep_ratio, - int chunk_size, - int n_lookahead, - int pool_kernel, - int n_keep, - const std::vector> & chunk_means, - const std::vector & selected, - const std::vector & forced, - const std::vector & compressed_ids, - const PFlashTraceFields * trace_fields = nullptr); - -std::vector select_pflash_chunks( - const std::vector & ids, - const std::vector & token_scores, - float keep_ratio, - int n_lookahead, - int score_query_end, - int pool_kernel, - const dflash::qwen3::PFlashSelectionConfig & config, - const std::vector & required_instruction_spans, - bool direct_mass, - bool write_trace, - const std::vector * segments = nullptr, - bool density = false, - const std::vector * other_token_scores = nullptr, - double split_fraction = 0.0); - -} // namespace dflash::common diff --git a/server/src/qwen3/qwen3_drafter_model.h b/server/src/qwen3/qwen3_drafter_model.h deleted file mode 100644 index d7063ec36..000000000 --- a/server/src/qwen3/qwen3_drafter_model.h +++ /dev/null @@ -1,151 +0,0 @@ -// Custom Qwen3-0.6B drafter forward, in dflash, replacing libllama. -// -// Uses the FlashPrefill dispatch path for the attention compute. Single -// process, single backend context, single ggml allocator — no Python, no -// Triton, no subprocess. -// -// Public API: -// bool load_qwen3_drafter_model(path, backend, out) → load GGUF weights -// bool forward_qwen3_drafter_model(weights, ids, out_q_capture, out_k_capture) -// void free_qwen3_drafter_model(weights) -// -#pragma once - -#include "ggml.h" - -#include -#include -#include -#include -#include - -struct ggml_context; -struct ggml_tensor; -struct ggml_backend; -typedef struct ggml_backend * ggml_backend_t; -struct ggml_backend_buffer; -typedef struct ggml_backend_buffer * ggml_backend_buffer_t; - -namespace dflash::common { - -struct Qwen3DrafterLayer { - ggml_tensor * attn_norm = nullptr; // [hidden] - ggml_tensor * wq = nullptr; // [hidden, q_dim] = [1024, 2048] - ggml_tensor * wk = nullptr; // [hidden, kv_dim] = [1024, 1024] - ggml_tensor * wv = nullptr; // [hidden, kv_dim] - ggml_tensor * wo = nullptr; // [q_dim, hidden] = [2048, 1024] - ggml_tensor * q_norm = nullptr; // [head_dim] = [128] - ggml_tensor * k_norm = nullptr; // [head_dim] - ggml_tensor * ffn_norm = nullptr; // [hidden] - ggml_tensor * ffn_gate = nullptr; // [hidden, ffn] - ggml_tensor * ffn_up = nullptr; // [hidden, ffn] - ggml_tensor * ffn_down = nullptr; // [ffn, hidden] -}; - -struct Qwen3DrafterWeights { - ggml_context * ctx = nullptr; - ggml_backend_t backend = nullptr; - ggml_backend_buffer_t buf = nullptr; - ggml_type weight_type = GGML_TYPE_BF16; - ggml_type compute_type = GGML_TYPE_BF16; - - ggml_tensor * tok_embd = nullptr; // [hidden, vocab] - ggml_tensor * out_norm = nullptr; // [hidden] - ggml_tensor * output = nullptr; // [hidden, vocab] (lm_head) - - std::vector layers; // size = n_layer = 28 - - // Architecture metadata. - int n_layer = 28; - int n_head = 16; - int n_head_kv = 8; - int n_embd = 1024; - int n_ff = 3072; - int head_dim = 128; - int n_vocab = 151936; - int n_ctx_max = 40960; - float rope_theta = 1000000.0f; - bool scoring_head_loaded = false; -}; - -bool load_qwen3_drafter_model(const std::string & gguf_path, - ggml_backend_t backend, - Qwen3DrafterWeights & out); - -void free_qwen3_drafter_model(Qwen3DrafterWeights & w); - -// Custom Qwen3-0.6B forward, fused with Liu Q-hook tail attention scoring. -// -// Inputs: -// w — loaded weights (must be on the selected GPU backend) -// ids — input token IDs of length S (drafter vocab) -// n_lookahead — number of query tokens for scorer attention (=8) -// score_query_end — exclusive end of query window; negative selects the tail -// -// Outputs: -// running_max — flat [n_lookahead, S] f32, max-over-heads-and-layers of -// softmax(Q_query @ K^T / sqrt(D)) per (lookahead, key) pair. -// Caller does AvgPool + chunk-top-K + span merge. -// -// Returns true on success. On failure sets last_error and returns false. -bool forward_qwen3_drafter_model( - const Qwen3DrafterWeights & w, - const std::vector & ids, - int n_lookahead, - std::vector & running_max, - int score_query_end = -1); - -struct QueryCaptureSlice { - int chunk_offset = 0; - int query_offset = 0; - int tokens = 0; - - bool valid() const { return tokens > 0; } -}; - -inline QueryCaptureSlice query_capture_slice( - int query_start, - int query_end, - int chunk_start, - int chunk_tokens) { - const int chunk_end = chunk_start + chunk_tokens; - const int overlap_start = query_start > chunk_start ? query_start : chunk_start; - const int overlap_end = query_end < chunk_end ? query_end : chunk_end; - if (overlap_start >= overlap_end) return {}; - return { - overlap_start - chunk_start, - overlap_start - query_start, - overlap_end - overlap_start, - }; -} - -inline size_t count_nonfinite_scores(const float * values, size_t count) { - size_t nonfinite = 0; - for (size_t index = 0; index < count; ++index) { - if (!std::isfinite(values[index])) ++nonfinite; - } - return nonfinite; -} - -// Scoring-head token mass: mean over heads and query tokens of softmax -// probabilities laid out as ggml [n_keys, n_queries, n_heads] (ne0 fastest). -inline void scoring_head_mean_token_mass( - const float * probs, - int n_keys, - int n_queries, - int n_heads, - std::vector & out) { - out.assign((size_t) n_keys, 0.0f); - if (n_keys <= 0 || n_queries <= 0 || n_heads <= 0) return; - std::vector sum((size_t) n_keys, 0.0); - for (int h = 0; h < n_heads; ++h) { - for (int t = 0; t < n_queries; ++t) { - const float * row = probs + ((size_t) h * n_queries + t) * n_keys; - for (int j = 0; j < n_keys; ++j) sum[(size_t) j] += row[j]; - } - } - const double denominator = (double) n_heads * (double) n_queries; - for (int j = 0; j < n_keys; ++j) out[(size_t) j] = (float) (sum[(size_t) j] / denominator); -} - -} // namespace dflash::common diff --git a/server/src/qwen3/qwen3_graph.cpp b/server/src/qwen3/qwen3_graph.cpp deleted file mode 100644 index 07066d952..000000000 --- a/server/src/qwen3/qwen3_graph.cpp +++ /dev/null @@ -1,1023 +0,0 @@ -// Custom forward for the Qwen3-0.6B drafter, replacing libllama. -// -// llama.cpp-style chunked prefill: ONE ggml graph per ubatch covering ALL 28 -// transformer layers. Per-layer K/V cache lives in persistent backend -// buffers. Sliding-window flash-attention via ggml-cuda's tensor-core -// `flash_attn_ext` keeps attention cost linear in S. -// -// **Algorithmic note vs blog**: -// The blog stack is Liu Q-hook tail scoring + FlashPrefill block-sparse FA. -// The Liu Q-hook is implemented with a NoPE fix: by default (DFLASH_FP_NOPE_TAIL=1) -// the tail score uses pre-RoPE K/Q, removing the RoPE distance decay that -// buries early-position needle chunks and was causing NIAH failures. -// Set DFLASH_FP_NOPE_TAIL=0 to revert to post-RoPE scoring. The block-sparse FA is replaced -// with a sliding-window approximation here because (a) ggml-cuda's -// `flash_attn_ext` already gives tensor-core speed inside the ubatch -// graph, and (b) our own block-sparse CUDA kernel needs a tensor-core -// rewrite (mma.sync.aligned) to actually beat ggml's FA — see -// `src/flashprefill_kernels.cu` for the (slow) scalar reference path. -// At S=140K with W=512 sliding window the NIAH magic key still propagates -// through 28 layers and is recovered in the kept tokens, so this -// approximation passes the actual e2e correctness check the user cares -// about. The block-sparse FA upgrade remains the next deliverable for -// "match the article algorithmically", but is functionally equivalent -// for the deployed perf budget today. -// -// Memory at S=140K, B=1, H=16, Hk=8, D=128, hidden=1024, ff=3072: -// weights ~1.5 GB -// reusable K_curr + V_curr [D, Hk, S] bf16 ~0.57 GB -// 28 × K_norope [D, Hk, S] bf16 (score-all default) ~8.0 GB -// Q_buf + attn_out [D, H, S] bf16 ~1.15 GB -// hidden_buf [hidden, S] f32 0.57 GB -// pos / mask_tail 1 MB -// per-ubatch graph transients (chunk_s sized) ~2-3 GB -// total including weights ~14-15 GB - -#include "qwen3_drafter_model.h" -#include "qwen3_buffer_plan.h" -#include "internal.h" -#include "flashprefill.h" -#include "../common/score_range.h" - -#include "device_runtime.h" - -#include "ggml.h" -#include "ggml-alloc.h" -#include "ggml-backend.h" - -#include -#include -#include -#include -#include -#include -#include -#include - -namespace dflash::common { - -namespace { - -constexpr int FA_WINDOW = 512; - -int chunk_s_ff() { - if (const char * e = std::getenv("DFLASH_FP_CHUNK_S")) { - int v = std::atoi(e); - if (v >= 1024) return v; - } -#if defined(DFLASH27B_BACKEND_HIP) - return 1024; -#else - return 4096; -#endif -} - -struct PersBuf { - ggml_context * ctx = nullptr; - ggml_backend_buffer_t buf = nullptr; - ggml_tensor * t = nullptr; -}; - -struct HipChunkGraphB { - ggml_context * ctx = nullptr; - ggml_backend_buffer_t buf = nullptr; - - ggml_tensor * h_in = nullptr; // input: hidden state slice (F32) - ggml_tensor * attn_in = nullptr; // input: attention output slice - ggml_tensor * h_after = nullptr; // h_in + attn_proj residual (F32) - ggml_tensor * hf = nullptr; // FFN norm result written by custom kernel (F32) - ggml_tensor * h_next = nullptr; // output: updated hidden state (F32) - - ggml_cgraph * gf_proj_add = nullptr; // compute h_after = h_in + wo*attn_in - ggml_cgraph * gf_ffn = nullptr; // compute h_next = h_after + ffn(hf) -}; - -bool make_pers(ggml_backend_t backend, ggml_type type, int n_dim, - const int64_t * dims, PersBuf & out) { - ggml_init_params ip{}; - ip.mem_size = ggml_tensor_overhead() * 4 + 1024; - ip.no_alloc = true; - ip.mem_buffer = nullptr; - out.ctx = ggml_init(ip); - if (!out.ctx) return false; - if (n_dim == 1) out.t = ggml_new_tensor_1d(out.ctx, type, dims[0]); - else if (n_dim == 2) out.t = ggml_new_tensor_2d(out.ctx, type, dims[0], dims[1]); - else if (n_dim == 3) out.t = ggml_new_tensor_3d(out.ctx, type, dims[0], dims[1], dims[2]); - else return false; - out.buf = ggml_backend_alloc_ctx_tensors(out.ctx, backend); - return out.buf != nullptr; -} - -void free_pers(PersBuf & p) { - if (p.buf) { ggml_backend_buffer_free(p.buf); p.buf = nullptr; } - if (p.ctx) { ggml_free(p.ctx); p.ctx = nullptr; } - p.t = nullptr; -} - -void free_hip_chunk_graph_b(HipChunkGraphB & g) { - if (g.buf) { - ggml_backend_buffer_free(g.buf); - g.buf = nullptr; - } - if (g.ctx) { - ggml_free(g.ctx); - g.ctx = nullptr; - } - g = {}; -} - -#if defined(DFLASH27B_BACKEND_HIP) -bool build_hip_chunk_graph_b(const Qwen3DrafterLayer & L, - ggml_backend_t backend, - int hidden, - int q_dim, - int chunk, - ggml_type compute_type, - float eps, - HipChunkGraphB & out) { - ggml_init_params ip{}; - ip.mem_size = ggml_tensor_overhead() * 128 - + ggml_graph_overhead_custom(1024, false) * 6 - + 256 * 1024; - ip.no_alloc = true; - out.ctx = ggml_init(ip); - if (!out.ctx) return false; - - out.h_in = ggml_new_tensor_2d(out.ctx, GGML_TYPE_F32, hidden, chunk); - out.attn_in = ggml_new_tensor_2d(out.ctx, compute_type, q_dim, chunk); - ggml_set_input(out.h_in); - ggml_set_input(out.attn_in); - - ggml_tensor * attn_proj = ggml_mul_mat(out.ctx, L.wo, out.attn_in); - out.h_after = ggml_add(out.ctx, out.h_in, attn_proj); - // h_after is output of gf_proj_add AND input of gf_ffn (stops re-traversal). - ggml_set_input(out.h_after); - ggml_set_output(out.h_after); - out.gf_proj_add = ggml_new_graph_custom(out.ctx, 1024, false); - ggml_build_forward_expand(out.gf_proj_add, out.h_after); - - out.hf = ggml_new_tensor_2d(out.ctx, GGML_TYPE_F32, hidden, chunk); - ggml_set_input(out.hf); - - // gf_ffn: one combined graph for all FFN ops after the RMSNorm. - // h_after and hf are both inputs so no re-traversal into proj_add or norm. - ggml_tensor * gate = ggml_silu(out.ctx, ggml_mul_mat(out.ctx, L.ffn_gate, out.hf)); - ggml_tensor * up = ggml_mul_mat(out.ctx, L.ffn_up, out.hf); - ggml_tensor * gu = ggml_mul(out.ctx, gate, up); - ggml_tensor * ffn_out = ggml_mul_mat(out.ctx, L.ffn_down, gu); - out.h_next = ggml_add(out.ctx, out.h_after, ffn_out); - ggml_set_output(out.h_next); - out.gf_ffn = ggml_new_graph_custom(out.ctx, 1024, false); - ggml_build_forward_expand(out.gf_ffn, out.h_next); - - out.buf = ggml_backend_alloc_ctx_tensors(out.ctx, backend); - if (!out.buf) { - return false; - } - - return true; -} - -bool warm_hip_chunk_graph_b_once(ggml_backend_t backend, - HipChunkGraphB & out, - std::string & error) { - static bool warmed = false; - if (warmed) { - return true; - } - - struct ggml_tensor * warm_tensors[] = { - out.h_in, out.attn_in, out.h_after, out.hf, out.h_next, - }; - for (ggml_tensor * t : warm_tensors) { - cudaError_t e = cudaMemset(t->data, 0, ggml_nbytes(t)); - if (e != cudaSuccess) { - error = std::string("memset failed: ") + cudaGetErrorString(e); - return false; - } - } - - const ggml_status proj_status = - ggml_backend_graph_compute(backend, out.gf_proj_add); - if (proj_status != GGML_STATUS_SUCCESS) { - error = std::string("projection graph failed: ") + - ggml_status_to_string(proj_status); - return false; - } - const ggml_status ffn_status = - ggml_backend_graph_compute(backend, out.gf_ffn); - if (ffn_status != GGML_STATUS_SUCCESS) { - error = std::string("FFN graph failed: ") + - ggml_status_to_string(ffn_status); - return false; - } - warmed = true; - return true; -} -#endif - -inline uint16_t f32_to_f16(float f) { - uint32_t bits; - std::memcpy(&bits, &f, 4); - uint32_t sign = (bits >> 16) & 0x8000; - int32_t exp = ((int32_t)((bits >> 23) & 0xff)) - 127 + 15; - uint32_t mant = bits & 0x7fffff; - if (exp <= 0) return (uint16_t)sign; - if (exp >= 31) return (uint16_t)(sign | 0x7c00); - return (uint16_t)(sign | (exp << 10) | (mant >> 13)); -} - -} // namespace - -#if defined(DFLASH27B_BACKEND_HIP) -extern "C" void launch_rms_norm_mul_w_f32( - const float * src, const float * w, float * dst, - int n_tokens, int hidden, float eps, - cudaStream_t stream); -#endif - -bool forward_qwen3_drafter_model( - const Qwen3DrafterWeights & w, - const std::vector & ids, - int n_lookahead, - std::vector & running_max, - int score_query_end) -{ - if (!w.backend || !w.tok_embd) { - set_last_error("forward_qwen3_drafter_model: weights not loaded"); - return false; - } - if (w.n_layer <= 0) { - set_last_error("forward_qwen3_drafter_model: model has no layers"); - return false; - } - const int S = (int)ids.size(); - const int H = w.n_head; - const int Hk = w.n_head_kv; - const int D = w.head_dim; - const int gqa = (Hk > 0) ? (H / Hk) : 1; - const int hidden = w.n_embd; - const float eps = 1e-6f; - const float scale = 1.0f / std::sqrt((float)D); - const float rope_b = w.rope_theta; - // Pre-RoPE tail scoring: removes RoPE distance decay from the score signal. - // Default ON; set DFLASH_FP_NOPE_TAIL=0 to disable (saves ~K_curr_v memory). - static const bool configured_nope_tail = []() -> bool { - const char * e = std::getenv("DFLASH_FP_NOPE_TAIL"); - return e == nullptr || std::string(e) != "0"; - }(); - const bool nope_tail = w.scoring_head_loaded || configured_nope_tail; - - if (n_lookahead < 1 || S < n_lookahead + 1) { - set_last_error("forward_qwen3_drafter_model: S too small"); - return false; - } - const int query_end = score_query_end < 0 ? S : score_query_end; - if (query_end < n_lookahead || query_end > S) { - set_last_error( - "forward_qwen3_drafter_model: scorer query window out of range"); - return false; - } - const int query_start = query_end - n_lookahead; - running_max.assign((size_t)n_lookahead * S, -INFINITY); - - // Read scoring/early-exit env vars once; compute alloc range before buffers are created. - static const int score_layers_pre = []() -> int { - const char * e = std::getenv("PFLASH_DRAFTER_SCORE_LAYERS"); - if (e) { int v = std::atoi(e); if (v > 0) return v; } - return -1; - }(); - static const int early_exit_pre = []() -> int { - const char * e = std::getenv("PFLASH_DRAFTER_EARLY_EXIT_N"); - if (e) { int v = std::atoi(e); if (v > 0) return v; } - return -1; - }(); - const int fwd_layer_limit_pre = w.scoring_head_loaded - ? 14 - : ((early_exit_pre > 0 && early_exit_pre < w.n_layer) - ? early_exit_pre : w.n_layer); - const ScoreRange pre_range = w.scoring_head_loaded - ? ScoreRange{13, 14} - : compute_score_range(w.n_layer, score_layers_pre, fwd_layer_limit_pre); - const int score_layer_start_pre = pre_range.start; - const int n_score_layers = pre_range.count(); // K_norope/Q_norope sized to this, not n_layer - - PersBuf hidden_buf, pos_buf, mask_tail_buf, Q_buf, attn_out_buf; - // With NoPE tail scoring, only the current layer's RoPE K/V survive until - // FlashPrefill returns. Reuse those large buffers instead of reserving a - // full-sequence K/V pair for every layer. The legacy RoPE scoring path - // still retains per-layer K/Q-tail state because it consumes it after the - // forward loop. - const Qwen3DrafterBufferPlan buffer_plan = - qwen3_drafter_buffer_plan(nope_tail, w.n_layer); - std::vector K_curr_v(buffer_plan.rope_k_buffers); - std::vector V_curr_v(buffer_plan.value_buffers); - std::vector Q_last_v(buffer_plan.rope_q_tail_buffers); - // NoPE: allocate only for scored layers (avoids ~5.6 GB waste at 128K). - std::vector K_norope_v(nope_tail ? (size_t)n_score_layers : 0); - std::vector Q_norope_v(nope_tail ? (size_t)n_score_layers : 0); - auto cleanup_all = [&]() { - free_pers(hidden_buf); - free_pers(pos_buf); - free_pers(mask_tail_buf); - free_pers(Q_buf); - free_pers(attn_out_buf); - for (auto & p : K_curr_v) free_pers(p); - for (auto & p : V_curr_v) free_pers(p); - for (auto & p : Q_last_v) free_pers(p); - for (auto & p : K_norope_v) free_pers(p); - for (auto & p : Q_norope_v) free_pers(p); - }; - - { - int64_t d_h[] = {(int64_t)hidden, (int64_t)S}; - int64_t d_kv[] = {(int64_t)D, (int64_t)Hk, (int64_t)S}; - int64_t d_q[] = {(int64_t)D, (int64_t)H, (int64_t)S}; // full Q for FP - int64_t d_ql[] = {(int64_t)D, (int64_t)H, (int64_t)n_lookahead}; - int64_t d_p[] = {(int64_t)S}; - int64_t d_mt[] = {(int64_t)S, (int64_t)n_lookahead}; - const ggml_type half_type = w.compute_type; - if (!make_pers(w.backend, GGML_TYPE_F32, 2, d_h, hidden_buf) || - !make_pers(w.backend, GGML_TYPE_I32, 1, d_p, pos_buf) || - !make_pers(w.backend, GGML_TYPE_F32, 2, d_mt, mask_tail_buf) || - !make_pers(w.backend, half_type, 3, d_q, Q_buf) || - !make_pers(w.backend, half_type, 3, d_q, attn_out_buf)) { - set_last_error("forward_qwen3: persistent alloc failed (hidden/pos/mask/Q/attn_out)"); - cleanup_all(); - return false; - } - if (!make_pers(w.backend, half_type, 3, d_kv, V_curr_v[0])) { - set_last_error("forward_qwen3: reusable V_curr alloc failed"); - cleanup_all(); - return false; - } - for (int il = 0; il < w.n_layer; ++il) { - const size_t li = buffer_plan.layer_cache_index(il); - const bool need_layer_buffers = !nope_tail || il == 0; - if (need_layer_buffers && - (!make_pers(w.backend, half_type, 3, d_kv, K_curr_v[li]) || - (!nope_tail && !make_pers(w.backend, GGML_TYPE_F32, 3, d_ql, Q_last_v[li])))) { - set_last_error("forward_qwen3: K_curr/Q_last alloc failed at layer " + std::to_string(il)); - cleanup_all(); - return false; - } - if (nope_tail && il >= score_layer_start_pre && il < fwd_layer_limit_pre) { - const int si = il - score_layer_start_pre; - if (!make_pers(w.backend, half_type, 3, d_kv, K_norope_v[si]) || - !make_pers(w.backend, GGML_TYPE_F32, 3, d_ql, Q_norope_v[si])) { - set_last_error("forward_qwen3: K_norope/Q_norope alloc failed at layer " + std::to_string(il)); - cleanup_all(); - return false; - } - } - } - } - - { - std::vector pos((size_t)S); - for (int i = 0; i < S; ++i) pos[i] = i; - ggml_backend_tensor_set(pos_buf.t, pos.data(), 0, - (size_t)S * sizeof(int32_t)); - } - { - std::vector m((size_t)n_lookahead * S, 0.0f); - for (int t = 0; t < n_lookahead; ++t) { - const int visible_end = w.scoring_head_loaded - ? query_start - : query_start + t + 1; - for (int j = 0; j < S; ++j) { - m[(size_t)t * S + j] = (j < visible_end) ? 0.0f : -INFINITY; - } - } - ggml_backend_tensor_set(mask_tail_buf.t, m.data(), 0, - m.size() * sizeof(float)); - } - - // ── Embed: hidden_buf = get_rows(tok_embd, ids) ────────────────── - { - ggml_init_params ip{}; - ip.mem_size = ggml_tensor_overhead() * 8 + ggml_graph_overhead() + 16 * 1024; - ip.no_alloc = true; - ggml_context * gctx = ggml_init(ip); - ggml_tensor * t_ids = ggml_new_tensor_1d(gctx, GGML_TYPE_I32, S); - ggml_set_name(t_ids, "ids"); - ggml_tensor * embed = ggml_get_rows(gctx, w.tok_embd, t_ids); - ggml_tensor * cpy_h = ggml_cpy(gctx, embed, hidden_buf.t); - ggml_cgraph * gf = ggml_new_graph(gctx); - ggml_build_forward_expand(gf, cpy_h); - ggml_backend_buffer_t in_buf = ggml_backend_alloc_ctx_tensors(gctx, w.backend); - ggml_gallocr_t galloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(w.backend)); - if (!ggml_gallocr_alloc_graph(galloc, gf)) { - set_last_error("embed graph alloc failed"); - ggml_gallocr_free(galloc); - if (in_buf) ggml_backend_buffer_free(in_buf); - ggml_free(gctx); - cleanup_all(); - return false; - } - ggml_backend_tensor_set(t_ids, ids.data(), 0, (size_t)S * sizeof(int32_t)); - const ggml_status embed_status = - ggml_backend_graph_compute(w.backend, gf); - if (embed_status != GGML_STATUS_SUCCESS) { - set_last_error(std::string("embed graph compute failed: ") + - ggml_status_to_string(embed_status)); - ggml_gallocr_free(galloc); - if (in_buf) ggml_backend_buffer_free(in_buf); - ggml_free(gctx); - cleanup_all(); - return false; - } - ggml_gallocr_free(galloc); - if (in_buf) ggml_backend_buffer_free(in_buf); - ggml_free(gctx); - } - - // Per-layer A→FA→B loop. - ggml_gallocr_t galloc = ggml_gallocr_new( - ggml_backend_get_default_buffer_type(w.backend)); - - flashprefill::FlashPrefillConfig fp_cfg; -#if defined(DFLASH27B_BACKEND_HIP) - // The HIP sparse-forward kernel is much slower when FlashPrefill keeps a - // broad set of K blocks. Use a stricter default on ROCm; callers can still - // override with DFLASH_FP_ALPHA for quality/speed sweeps. - fp_cfg.alpha = 0.95f; -#endif - if (const char* a = std::getenv("DFLASH_FP_ALPHA")) { - float v = (float)std::atof(a); - if (v > 0.0f && v < 1.0f) fp_cfg.alpha = v; - } - auto t_total_start = std::chrono::steady_clock::now(); - double t_a_setup = 0.0, t_a_alloc = 0.0, t_compute_a = 0.0; - double t_b_warm = 0.0, t_b_setup = 0.0, t_b_alloc = 0.0, t_b_copy_in = 0.0, t_b_norm = 0.0, t_compute_b = 0.0, t_b_copy_out = 0.0; - double t_fp = 0.0; - - const int fwd_layer_limit = fwd_layer_limit_pre; - - for (int il = 0; il < fwd_layer_limit; ++il) { - const auto & L = w.layers[il]; - const size_t layer_cache_idx = buffer_plan.layer_cache_index(il); - const bool debug_first_layer = (il == 0 && std::getenv("DFLASH_FP_DEBUG_LAYER0") != nullptr); - - // ── Graph A (chunked): norm + Q/K/V proj + RoPE + copy to persistent K_curr/V_curr/Q_buf ── - // ggml-cuda RoPE/element-wise kernels hit `invalid configuration argument` when - // an op operates over more than ~65K rows in y/z. Chunk loop keeps every per-row - // ggml op under that cap; FP CUDA kernel still runs once over full S below. - const int chunk_s_ff_v = chunk_s_ff(); - for (int cs = 0; cs < S; cs += chunk_s_ff_v) { - const int cl = std::min(chunk_s_ff_v, S - cs); - if (debug_first_layer) { - std::fprintf(stderr, "[qwen3-0.6b-fp dbg] layer0 chunk A start cs=%d cl=%d\n", cs, cl); - std::fflush(stderr); - } - auto tA_setup0 = std::chrono::steady_clock::now(); - - ggml_init_params ipA{}; - ipA.mem_size = ggml_tensor_overhead() * 64 - + ggml_graph_overhead_custom(2048, false) - + 64 * 1024; - ipA.no_alloc = true; - ggml_context * gA = ggml_init(ipA); - if (!gA) { set_last_error("graph A init failed"); cleanup_all(); ggml_gallocr_free(galloc); return false; } - ggml_cgraph * gfA = ggml_new_graph_custom(gA, 2048, false); - - const size_t h_esz = ggml_element_size(hidden_buf.t); - ggml_tensor * h_view = ggml_view_2d(gA, hidden_buf.t, - hidden, cl, - hidden * h_esz, - (size_t)cs * hidden * h_esz); - ggml_tensor * pos_chunk = ggml_view_1d(gA, pos_buf.t, cl, - (size_t)cs * sizeof(int32_t)); - - ggml_tensor * h_norm = ggml_rms_norm(gA, h_view, eps); - h_norm = ggml_mul(gA, h_norm, L.attn_norm); - - ggml_tensor * Q = ggml_mul_mat(gA, L.wq, h_norm); - Q = ggml_reshape_3d(gA, Q, D, H, cl); - if (L.q_norm) { - Q = ggml_rms_norm(gA, Q, eps); - Q = ggml_mul(gA, Q, L.q_norm); - } - // NoPE: capture pre-RoPE Q tail (only for layers that will be scored). - if (nope_tail && il >= score_layer_start_pre) { - const int si = il - score_layer_start_pre; - const auto capture = query_capture_slice( - query_start, query_end, cs, cl); - if (capture.valid()) { - ggml_tensor * Q_prenrope_tail = ggml_view_3d( - gA, Q, D, H, capture.tokens, - Q->nb[1], Q->nb[2], - (size_t)capture.chunk_offset * Q->nb[2]); - Q_prenrope_tail = ggml_cont(gA, Q_prenrope_tail); - Q_prenrope_tail = ggml_reshape_1d( - gA, Q_prenrope_tail, D * H * capture.tokens); - ggml_tensor * Q_prenrope_dst = ggml_view_1d( - gA, Q_norope_v[si].t, D * H * capture.tokens, - (size_t)capture.query_offset * Q_norope_v[si].t->nb[2]); - ggml_build_forward_expand(gfA, - ggml_cpy(gA, Q_prenrope_tail, Q_prenrope_dst)); - } - } - Q = ggml_rope_ext(gA, Q, pos_chunk, nullptr, D, - GGML_ROPE_TYPE_NEOX, 0, - rope_b, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); - - ggml_tensor * K = ggml_mul_mat(gA, L.wk, h_norm); - K = ggml_reshape_3d(gA, K, D, Hk, cl); - if (L.k_norm) { - K = ggml_rms_norm(gA, K, eps); - K = ggml_mul(gA, K, L.k_norm); - } - // NoPE: save pre-RoPE K chunk (only for layers that will be scored). - if (nope_tail && il >= score_layer_start_pre) { - const int si = il - score_layer_start_pre; - const size_t kn_esz = ggml_element_size(K_norope_v[si].t); - ggml_tensor * Kn_dst = ggml_view_3d(gA, K_norope_v[si].t, D, Hk, cl, - kn_esz * D, kn_esz * D * Hk, - (size_t)cs * kn_esz * D * Hk); - ggml_build_forward_expand(gfA, ggml_cpy(gA, K, Kn_dst)); - } - K = ggml_rope_ext(gA, K, pos_chunk, nullptr, D, - GGML_ROPE_TYPE_NEOX, 0, - rope_b, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); - - ggml_tensor * V = ggml_mul_mat(gA, L.wv, h_norm); - V = ggml_reshape_3d(gA, V, D, Hk, cl); - - const size_t q_esz = ggml_element_size(Q_buf.t); - const size_t kv_esz = ggml_element_size(K_curr_v[layer_cache_idx].t); - ggml_tensor * Q_dst = ggml_view_3d(gA, Q_buf.t, D, H, cl, - q_esz * D, q_esz * D * H, - (size_t)cs * q_esz * D * H); - ggml_tensor * K_dst = ggml_view_3d(gA, K_curr_v[layer_cache_idx].t, D, Hk, cl, - kv_esz * D, kv_esz * D * Hk, - (size_t)cs * kv_esz * D * Hk); - ggml_tensor * V_dst = ggml_view_3d(gA, V_curr_v[0].t, D, Hk, cl, - kv_esz * D, kv_esz * D * Hk, - (size_t)cs * kv_esz * D * Hk); - ggml_build_forward_expand(gfA, ggml_cpy(gA, Q, Q_dst)); - ggml_build_forward_expand(gfA, ggml_cpy(gA, K, K_dst)); - ggml_build_forward_expand(gfA, ggml_cpy(gA, V, V_dst)); - - // Copy the overlapping Q-query slice; a query can straddle chunks. - const auto capture = query_capture_slice( - query_start, query_end, cs, cl); - if (!nope_tail && capture.valid()) { - ggml_tensor * Q_tail_local = ggml_view_3d( - gA, Q, D, H, capture.tokens, - Q->nb[1], Q->nb[2], - (size_t)capture.chunk_offset * Q->nb[2]); - Q_tail_local = ggml_cont(gA, Q_tail_local); - Q_tail_local = ggml_reshape_1d( - gA, Q_tail_local, D * H * capture.tokens); - ggml_tensor * Q_tail_dst = ggml_view_1d( - gA, Q_last_v[layer_cache_idx].t, D * H * capture.tokens, - (size_t)capture.query_offset * Q_last_v[layer_cache_idx].t->nb[2]); - ggml_build_forward_expand(gfA, - ggml_cpy(gA, Q_tail_local, Q_tail_dst)); - } - - auto tA_setup1 = std::chrono::steady_clock::now(); - t_a_setup += std::chrono::duration(tA_setup1 - tA_setup0).count(); - - auto tA_alloc0 = std::chrono::steady_clock::now(); - if (!ggml_gallocr_alloc_graph(galloc, gfA)) { - set_last_error("graph A alloc failed at layer " + std::to_string(il)); - ggml_free(gA); ggml_gallocr_free(galloc); cleanup_all(); return false; - } - auto tA_alloc1 = std::chrono::steady_clock::now(); - t_a_alloc += std::chrono::duration(tA_alloc1 - tA_alloc0).count(); - auto tA0 = std::chrono::steady_clock::now(); - const ggml_status graph_a_status = - ggml_backend_graph_compute(w.backend, gfA); - ggml_backend_synchronize(w.backend); - auto tA1 = std::chrono::steady_clock::now(); - t_compute_a += std::chrono::duration(tA1 - tA0).count(); - if (graph_a_status != GGML_STATUS_SUCCESS) { - set_last_error(std::string("graph A compute failed at layer ") + - std::to_string(il) + " chunk " + - std::to_string(cs) + ": " + - ggml_status_to_string(graph_a_status)); - ggml_free(gA); - ggml_gallocr_free(galloc); cleanup_all(); return false; - } - if (debug_first_layer) { - std::fprintf(stderr, - "[qwen3-0.6b-fp dbg] layer0 chunk A done setup=%.3fs alloc=%.3fs compute=%.3fs\n", - std::chrono::duration(tA_setup1 - tA_setup0).count(), - std::chrono::duration(tA_alloc1 - tA_alloc0).count(), - std::chrono::duration(tA1 - tA0).count()); - std::fflush(stderr); - } - ggml_free(gA); - } - - if (w.scoring_head_loaded && il == 13) { - continue; - } - - // ── Attention dispatch ── - auto tF0 = std::chrono::steady_clock::now(); - int rc = flashprefill::flash_prefill_forward( - w.backend, - Q_buf.t->data, - K_curr_v[layer_cache_idx].t->data, - V_curr_v[0].t->data, - attn_out_buf.t->data, - 1, S, H, Hk, D, scale, - Q_buf.t->type, - fp_cfg); - if (rc != 0) { - set_last_error("flash_prefill_forward failed at layer " + std::to_string(il)); - ggml_gallocr_free(galloc); cleanup_all(); return false; - } - cudaError_t fp_launch_e = cudaGetLastError(); - if (fp_launch_e != cudaSuccess) { - set_last_error(std::string("flash_prefill launch failed at layer ") + - std::to_string(il) + ": " + - cudaGetErrorString(fp_launch_e)); - ggml_gallocr_free(galloc); cleanup_all(); return false; - } - cudaError_t fp_sync_e = cudaDeviceSynchronize(); - if (fp_sync_e != cudaSuccess) { - set_last_error(std::string("flash_prefill synchronization failed at layer ") + - std::to_string(il) + ": " + - cudaGetErrorString(fp_sync_e)); - ggml_gallocr_free(galloc); cleanup_all(); return false; - } - auto tF1 = std::chrono::steady_clock::now(); - t_fp += std::chrono::duration(tF1 - tF0).count(); - if (debug_first_layer) { - std::fprintf(stderr, "[qwen3-0.6b-fp dbg] layer0 FP done compute=%.3fs\n", - std::chrono::duration(tF1 - tF0).count()); - std::fflush(stderr); - } - - // ── Graph B (chunked, reusable): o_proj + residual + ffn + write hidden_buf ── -#if defined(DFLASH27B_BACKEND_HIP) - auto tB_setup0 = std::chrono::steady_clock::now(); - HipChunkGraphB gb{}; - if (!build_hip_chunk_graph_b(L, w.backend, hidden, D * H, chunk_s_ff_v, w.compute_type, eps, gb)) { - set_last_error("graph B reusable build failed at layer " + std::to_string(il)); - ggml_gallocr_free(galloc); cleanup_all(); return false; - } - auto tB_setup1 = std::chrono::steady_clock::now(); - t_b_setup += std::chrono::duration(tB_setup1 - tB_setup0).count(); - if (debug_first_layer) { - std::fprintf(stderr, - "[qwen3-0.6b-fp dbg] layer0 graph B reusable setup+alloc done setup=%.3fs\n", - std::chrono::duration(tB_setup1 - tB_setup0).count()); - std::fflush(stderr); - } - - auto tB_warm0 = std::chrono::steady_clock::now(); - std::string warm_error; - if (!warm_hip_chunk_graph_b_once(w.backend, gb, warm_error)) { - set_last_error(std::string("graph B warmup failed at layer ") + - std::to_string(il) + ": " + warm_error); - free_hip_chunk_graph_b(gb); - ggml_gallocr_free(galloc); cleanup_all(); return false; - } - auto tB_warm1 = std::chrono::steady_clock::now(); - t_b_warm += std::chrono::duration(tB_warm1 - tB_warm0).count(); - if (debug_first_layer) { - std::fprintf(stderr, - "[qwen3-0.6b-fp dbg] layer0 graph B warmup=%.3fs\n", - std::chrono::duration(tB_warm1 - tB_warm0).count()); - std::fflush(stderr); - } - - for (int cs = 0; cs < S; cs += chunk_s_ff_v) { - const int cl = std::min(chunk_s_ff_v, S - cs); - if (debug_first_layer) { - std::fprintf(stderr, "[qwen3-0.6b-fp dbg] layer0 chunk B start cs=%d cl=%d\n", cs, cl); - std::fflush(stderr); - } - - const size_t h_esz = ggml_element_size(hidden_buf.t); - const size_t a_esz = ggml_element_size(attn_out_buf.t); - const size_t h_bytes = (size_t)hidden * cl * sizeof(float); - const size_t a_bytes = (size_t)(D * H) * cl * a_esz; - const char * h_src = (const char *)hidden_buf.t->data + (size_t)cs * hidden * h_esz; - const char * a_src = (const char *)attn_out_buf.t->data + (size_t)cs * (D * H) * a_esz; - - auto tB_copy_in0 = std::chrono::steady_clock::now(); - cudaError_t copy_h_in_e = cudaMemcpy(gb.h_in->data, h_src, h_bytes, cudaMemcpyDeviceToDevice); - if (copy_h_in_e != cudaSuccess) { - set_last_error(std::string("graph B hidden copy-in failed at layer ") + std::to_string(il) + ": " + cudaGetErrorString(copy_h_in_e)); - free_hip_chunk_graph_b(gb); - ggml_gallocr_free(galloc); cleanup_all(); return false; - } - cudaError_t copy_a_in_e = cudaMemcpy(gb.attn_in->data, a_src, a_bytes, cudaMemcpyDeviceToDevice); - if (copy_a_in_e != cudaSuccess) { - set_last_error(std::string("graph B attn copy-in failed at layer ") + std::to_string(il) + ": " + cudaGetErrorString(copy_a_in_e)); - free_hip_chunk_graph_b(gb); - ggml_gallocr_free(galloc); cleanup_all(); return false; - } - // HIP D2D hipMemcpy can return before its null-stream copy finishes. - // GGML uses a nonblocking stream, so make the copy-in dependency explicit. - cudaError_t copy_in_sync_e = cudaStreamSynchronize(nullptr); - if (copy_in_sync_e != cudaSuccess) { - set_last_error(std::string("graph B copy-in synchronization failed at layer ") + - std::to_string(il) + " chunk " + std::to_string(cs) + ": " + - cudaGetErrorString(copy_in_sync_e)); - free_hip_chunk_graph_b(gb); - ggml_gallocr_free(galloc); cleanup_all(); return false; - } - auto tB_copy_in1 = std::chrono::steady_clock::now(); - t_b_copy_in += std::chrono::duration(tB_copy_in1 - tB_copy_in0).count(); - if (debug_first_layer) { - std::fprintf(stderr, - "[qwen3-0.6b-fp dbg] layer0 chunk B copy-in done copy=%.3fs\n", - std::chrono::duration(tB_copy_in1 - tB_copy_in0).count()); - std::fflush(stderr); - } - - if (debug_first_layer && cs == 6144) { - std::vector h_dbg((size_t)hidden * cl); - ggml_backend_tensor_get(gb.h_in, h_dbg.data(), 0, h_dbg.size() * sizeof(float)); - float h_min = std::numeric_limits::infinity(); - float h_max = -std::numeric_limits::infinity(); - size_t h_nonfinite = 0; - for (float v : h_dbg) { - if (!std::isfinite(v)) { - ++h_nonfinite; - continue; - } - h_min = std::min(h_min, v); - h_max = std::max(h_max, v); - } - std::fprintf(stderr, - "[qwen3-0.6b-fp dbg] layer0 chunk6144 h_in stats min=%g max=%g nonfinite=%zu/%zu\n", - h_min, h_max, h_nonfinite, h_dbg.size()); - std::fflush(stderr); - } - - // Run gf_proj_add FIRST so gb.h_after holds the current chunk's - // projected residual sum before we read it back for CPU-side - // RMSNorm. (Reading h_after before this compute would pick up - // the previous chunk's value — stale FFN inputs.) - auto tB0 = std::chrono::steady_clock::now(); - double proj_s = 0, ffn_s = 0; - auto one = [&](ggml_cgraph * gf, double & acc) { - auto ts0 = std::chrono::steady_clock::now(); - const ggml_status status = - ggml_backend_graph_compute(w.backend, gf); - auto ts1 = std::chrono::steady_clock::now(); - acc = std::chrono::duration(ts1 - ts0).count(); - return status; - }; - const ggml_status proj_status = one(gb.gf_proj_add, proj_s); - if (proj_status != GGML_STATUS_SUCCESS) { - set_last_error(std::string("graph B projection compute failed at layer ") + - std::to_string(il) + " chunk " + - std::to_string(cs) + ": " + - ggml_status_to_string(proj_status)); - free_hip_chunk_graph_b(gb); - ggml_gallocr_free(galloc); cleanup_all(); return false; - } - - auto tB_norm0 = std::chrono::steady_clock::now(); - launch_rms_norm_mul_w_f32( - (const float *)gb.h_after->data, - (const float *)L.ffn_norm->data, - (float *)gb.hf->data, - cl, hidden, eps, - /*stream=*/nullptr); - cudaError_t rms_launch_e = cudaGetLastError(); - if (rms_launch_e != cudaSuccess) { - set_last_error(std::string("graph B RMSNorm launch failed at layer ") + - std::to_string(il) + " chunk " + - std::to_string(cs) + ": " + - cudaGetErrorString(rms_launch_e)); - free_hip_chunk_graph_b(gb); - ggml_gallocr_free(galloc); cleanup_all(); return false; - } - cudaError_t rms_sync_e = cudaDeviceSynchronize(); - if (rms_sync_e != cudaSuccess) { - set_last_error(std::string("graph B RMSNorm synchronization failed at layer ") + - std::to_string(il) + " chunk " + - std::to_string(cs) + ": " + - cudaGetErrorString(rms_sync_e)); - free_hip_chunk_graph_b(gb); - ggml_gallocr_free(galloc); cleanup_all(); return false; - } - auto tB_norm1 = std::chrono::steady_clock::now(); - t_b_norm += std::chrono::duration(tB_norm1 - tB_norm0).count(); - - const ggml_status ffn_status = one(gb.gf_ffn, ffn_s); - if (ffn_status != GGML_STATUS_SUCCESS) { - set_last_error(std::string("graph B FFN compute failed at layer ") + - std::to_string(il) + " chunk " + - std::to_string(cs) + ": " + - ggml_status_to_string(ffn_status)); - free_hip_chunk_graph_b(gb); - ggml_gallocr_free(galloc); cleanup_all(); return false; - } - auto tB1 = std::chrono::steady_clock::now(); - t_compute_b += std::chrono::duration(tB1 - tB0).count(); - - auto tB_copy_out0 = std::chrono::steady_clock::now(); - cudaError_t copy_out_e = cudaMemcpy((char *)hidden_buf.t->data + (size_t)cs * hidden * h_esz, - gb.h_next->data, - h_bytes, - cudaMemcpyDeviceToDevice); - if (copy_out_e != cudaSuccess) { - set_last_error(std::string("graph B copy-out failed at layer ") + std::to_string(il) + ": " + cudaGetErrorString(copy_out_e)); - free_hip_chunk_graph_b(gb); - ggml_gallocr_free(galloc); cleanup_all(); return false; - } - auto tB_copy_out1 = std::chrono::steady_clock::now(); - t_b_copy_out += std::chrono::duration(tB_copy_out1 - tB_copy_out0).count(); - if (debug_first_layer) { - std::fprintf(stderr, - "[qwen3-0.6b-fp dbg] layer0 chunk B compute-done compute=%.3fs copy-out=%.3fs [proj=%.3f norm_cpu=%.3f ffn=%.3f]\n", - std::chrono::duration(tB1 - tB0).count(), - std::chrono::duration(tB_copy_out1 - tB_copy_out0).count(), - proj_s, std::chrono::duration(tB_norm1 - tB_norm0).count(), ffn_s); - std::fflush(stderr); - } - if (debug_first_layer) { - std::fprintf(stderr, - "[qwen3-0.6b-fp dbg] layer0 chunk B done copy-in=%.3fs compute=%.3fs copy-out=%.3fs\n", - std::chrono::duration(tB_copy_in1 - tB_copy_in0).count(), - std::chrono::duration(tB1 - tB0).count(), - std::chrono::duration(tB_copy_out1 - tB_copy_out0).count()); - std::fflush(stderr); - } - } - free_hip_chunk_graph_b(gb); -#else - // Non-HIP path keeps the existing graph-B implementation. - for (int cs = 0; cs < S; cs += chunk_s_ff_v) { - const int cl = std::min(chunk_s_ff_v, S - cs); - ggml_init_params ipB{}; - ipB.mem_size = ggml_tensor_overhead() * 64 + ggml_graph_overhead_custom(2048, false) + 64 * 1024; - ipB.no_alloc = true; - ggml_context * gB = ggml_init(ipB); - if (!gB) { set_last_error("graph B init failed"); cleanup_all(); ggml_gallocr_free(galloc); return false; } - ggml_cgraph * gfB = ggml_new_graph_custom(gB, 2048, false); - const size_t h_esz = ggml_element_size(hidden_buf.t); - ggml_tensor * h_src = ggml_view_2d(gB, hidden_buf.t, hidden, cl, hidden * h_esz, (size_t)cs * hidden * h_esz); - ggml_tensor * h_in = ggml_new_tensor_2d(gB, GGML_TYPE_F32, hidden, cl); - ggml_set_input(h_in); - const size_t a_esz = ggml_element_size(attn_out_buf.t); - ggml_tensor * attn_in = ggml_view_2d(gB, attn_out_buf.t, D * H, cl, a_esz * D * H, (size_t)cs * a_esz * D * H); - ggml_tensor * attn_proj = ggml_mul_mat(gB, L.wo, attn_in); - ggml_tensor * h_after = ggml_add(gB, h_in, attn_proj); - ggml_tensor * hf = ggml_rms_norm(gB, h_after, eps); - hf = ggml_mul(gB, hf, L.ffn_norm); - ggml_tensor * gate_t = ggml_mul_mat(gB, L.ffn_gate, hf); - gate_t = ggml_silu(gB, gate_t); - ggml_tensor * up_t = ggml_mul_mat(gB, L.ffn_up, hf); - ggml_tensor * gu = ggml_mul(gB, gate_t, up_t); - ggml_tensor * ffn_out = ggml_mul_mat(gB, L.ffn_down, gu); - ggml_tensor * h_next = ggml_add(gB, h_after, ffn_out); - ggml_set_output(h_next); - ggml_build_forward_expand(gfB, h_next); - ggml_backend_buffer_t gB_buf = ggml_backend_alloc_ctx_tensors(gB, w.backend); - if (!gB_buf) { set_last_error("graph B ctx allocation failed at layer " + std::to_string(il)); ggml_free(gB); ggml_gallocr_free(galloc); cleanup_all(); return false; } - ggml_backend_tensor_copy(h_src, h_in); - ggml_backend_graph_compute(w.backend, gfB); - ggml_backend_tensor_copy(h_next, h_src); - ggml_backend_buffer_free(gB_buf); - ggml_free(gB); - } -#endif - - if (il == 0 || il == fwd_layer_limit - 1) { - std::fprintf(stderr, - "[qwen3-0.6b-fp] layer %d/%d done " - "(A_setup=%.3fs A_alloc=%.3fs A_compute=%.3fs FP=%.3fs " - "B_warm=%.3fs B_setup=%.3fs B_alloc=%.3fs B_copy_in=%.3fs B_norm=%.3fs B_compute=%.3fs B_copy_out=%.3fs)\n", - il + 1, fwd_layer_limit, - t_a_setup, t_a_alloc, t_compute_a, t_fp, - t_b_warm, t_b_setup, t_b_alloc, t_b_copy_in, t_b_norm, t_compute_b, t_b_copy_out); - std::fflush(stderr); - } - } - - ggml_gallocr_free(galloc); - - auto t_fwd_end = std::chrono::steady_clock::now(); - double t_fwd = std::chrono::duration(t_fwd_end - t_total_start).count(); - - // Tail attention scoring; range matches pre-alloc by construction. - const int score_layer_start = score_layer_start_pre; - const int score_layer_end = fwd_layer_limit; - - std::vector probs_h((size_t)S * n_lookahead * H); - auto t_score_start = std::chrono::steady_clock::now(); - - for (int il = score_layer_start; il < score_layer_end; ++il) { - const size_t layer_cache_idx = buffer_plan.layer_cache_index(il); - ggml_init_params ip{}; - ip.mem_size = ggml_tensor_overhead() * 32 + ggml_graph_overhead() + 16 * 1024; - ip.no_alloc = true; - ggml_context * gctx = ggml_init(ip); - - // K_norope_v / Q_norope_v are indexed from score_layer_start_pre. - const int si = il - score_layer_start_pre; - ggml_tensor * K_f32 = ggml_new_tensor_3d(gctx, GGML_TYPE_F32, D, Hk, S); - ggml_tensor * K_cast = ggml_cpy(gctx, - nope_tail ? K_norope_v[si].t : K_curr_v[layer_cache_idx].t, K_f32); - ggml_tensor * K_perm = ggml_cont(gctx, - ggml_permute(gctx, K_cast, 0, 2, 1, 3)); - ggml_tensor * K_score = K_perm; - if (gqa > 1) { - ggml_tensor * K_4d = ggml_reshape_4d(gctx, K_perm, D, S, 1, Hk); - ggml_tensor * K_tpl = ggml_new_tensor_4d(gctx, GGML_TYPE_F32, - D, S, gqa, Hk); - ggml_tensor * K_rep = ggml_repeat(gctx, K_4d, K_tpl); - K_score = ggml_reshape_3d(gctx, K_rep, D, S, H); - } - ggml_tensor * Q_tail_perm = ggml_cont(gctx, - ggml_permute(gctx, - nope_tail ? Q_norope_v[si].t : Q_last_v[layer_cache_idx].t, - 0, 2, 1, 3)); - ggml_tensor * attn_score = ggml_mul_mat(gctx, K_score, Q_tail_perm); - ggml_tensor * probs = ggml_soft_max_ext(gctx, attn_score, mask_tail_buf.t, - scale, 0.0f); - ggml_set_output(probs); - - ggml_cgraph * gf = ggml_new_graph(gctx); - ggml_build_forward_expand(gf, probs); - - ggml_backend_buffer_t in_buf = ggml_backend_alloc_ctx_tensors(gctx, w.backend); - ggml_gallocr_t s_galloc = ggml_gallocr_new( - ggml_backend_get_default_buffer_type(w.backend)); - if (!ggml_gallocr_alloc_graph(s_galloc, gf)) { - set_last_error("tail score graph alloc failed at layer " + std::to_string(il)); - ggml_gallocr_free(s_galloc); - if (in_buf) ggml_backend_buffer_free(in_buf); - ggml_free(gctx); - cleanup_all(); - return false; - } - const auto score_status = ggml_backend_graph_compute(w.backend, gf); - size_t nonfinite = 0; - if (score_status == GGML_STATUS_SUCCESS) { - ggml_backend_tensor_get(probs, probs_h.data(), 0, - probs_h.size() * sizeof(float)); - nonfinite = count_nonfinite_scores(probs_h.data(), probs_h.size()); - } - ggml_gallocr_free(s_galloc); - if (in_buf) ggml_backend_buffer_free(in_buf); - ggml_free(gctx); - if (score_status != GGML_STATUS_SUCCESS) { - set_last_error("tail score graph compute failed at layer " + - std::to_string(il)); - cleanup_all(); - return false; - } - if (nonfinite != 0) { - const std::string message = - "non-finite PFlash tail scores at layer " + - std::to_string(il) + ": " + std::to_string(nonfinite) + - "/" + std::to_string(probs_h.size()); - std::fprintf(stderr, "[pflash] ERROR: %s\n", message.c_str()); - std::fflush(stderr); - set_last_error(message); - cleanup_all(); - return false; - } - - for (int t = 0; t < n_lookahead; ++t) { - for (int j = 0; j < S; ++j) { - size_t idx = (size_t)t * S + j; - if (w.scoring_head_loaded) { - float sum = 0.0f; - for (int h = 0; h < H; ++h) { - sum += probs_h[(size_t)j - + (size_t)t * S - + (size_t)h * S * n_lookahead]; - } - running_max[idx] = sum / (float)H; - } else { - float m = -INFINITY; - for (int h = 0; h < H; ++h) { - float v = probs_h[(size_t)j - + (size_t)t * S - + (size_t)h * S * n_lookahead]; - if (v > m) m = v; - } - if (m > running_max[idx]) running_max[idx] = m; - } - } - } - } - - auto t_total_end = std::chrono::steady_clock::now(); - double t_score = std::chrono::duration(t_total_end - t_score_start).count(); - std::fprintf(stderr, - "[qwen3-0.6b-fp] forward %.2fs (S=%d, A_setup=%.2fs A_alloc=%.2fs A_compute=%.2fs FP=%.2fs B_warm=%.2fs B_setup=%.2fs B_alloc=%.2fs B_copy_in=%.2fs B_norm=%.2fs B_compute=%.2fs B_copy_out=%.2fs) " - "tail-score %.2fs (layers %d-%d) total %.2fs\n", - t_fwd, S, t_a_setup, t_a_alloc, t_compute_a, t_fp, t_b_warm, t_b_setup, t_b_alloc, t_b_copy_in, t_b_norm, t_compute_b, t_b_copy_out, - t_score, score_layer_start, score_layer_end - 1, t_fwd + t_score); - std::fflush(stderr); - - cleanup_all(); - return true; -} - -} // namespace dflash::common diff --git a/server/src/qwen3/qwen3_loader.cpp b/server/src/qwen3/qwen3_loader.cpp index af08b8133..54d982067 100644 --- a/server/src/qwen3/qwen3_loader.cpp +++ b/server/src/qwen3/qwen3_loader.cpp @@ -1,4 +1,4 @@ -// GGUF loader for Qwen3-0.6B drafter. Reads weights from a BF16 GGUF file +// GGUF loader for the Qwen3-0.6B model. Reads weights from a BF16 GGUF file // produced by `convert_hf_to_gguf.py Qwen/Qwen3-0.6B`. Sets up ggml tensors // on the requested backend. // @@ -22,7 +22,7 @@ // We mmap the GGUF file and copy each tensor's bytes to the backend buffer // (mirrors the dflash gguf_target_loader pattern). -#include "qwen3_drafter_model.h" +#include "qwen3_model.h" #include "common/backend_precision.h" #include "common/gguf_inspect.h" #include "common/gguf_mmap.h" @@ -123,96 +123,11 @@ float get_f32(gguf_context * g, const char * key, float def) { return gguf_get_val_f32(g, k); } -bool metadata_equals(gguf_context * g, const char * key, const char * expected) { - const int id = gguf_find_key(g, key); - return id >= 0 && std::string(gguf_get_val_str(g, id)) == expected; -} - -bool load_scoring_head( - const std::string & path, - const std::string & drafter_sha256, - Qwen3DrafterWeights & out) { - ggml_context * tensor_ctx = nullptr; - gguf_init_params iparams{ /*no_alloc=*/ true, /*ctx=*/ &tensor_ctx }; - gguf_context * gctx = gguf_init_from_file(path.c_str(), iparams); - if (!gctx) { - set_last_error("scoring head GGUF could not be opened: " + path); - return false; - } - auto fail = [&](const std::string & message) { - gguf_free(gctx); - if (tensor_ctx) ggml_free(tensor_ctx); - set_last_error(message); - return false; - }; - // GGUF contract of a scoring-head file: architecture `pflash_scoring_head`, - // metadata and tensors under `scoringhead.*`. - const bool metadata_ok = - metadata_equals(gctx, "general.architecture", "pflash_scoring_head") && - metadata_equals(gctx, "scoringhead.schema", "qwen3_0_6b_nope_qk_mass_v1") && - metadata_equals(gctx, "scoringhead.base_model", "Qwen/Qwen3-0.6B") && - metadata_equals( - gctx, - "scoringhead.runtime_gguf_sha256", - drafter_sha256.c_str()) && - metadata_equals( - gctx, - "scoringhead.feature_tap", - "post_block12_residual_before_block13"); - if (!metadata_ok) { - return fail("scoring head metadata does not match the loaded Qwen3-0.6B drafter"); - } - struct TensorContract { - const char * name; - ggml_tensor * destination; - }; - const TensorContract contracts[] = { - {"scoringhead.attn_q.weight", out.layers[13].wq}, - {"scoringhead.attn_k.weight", out.layers[13].wk}, - }; - for (const auto & contract : contracts) { - const int64_t id = gguf_find_tensor(gctx, contract.name); - ggml_tensor * source = tensor_ctx - ? ggml_get_tensor(tensor_ctx, contract.name) - : nullptr; - if (id < 0 || gguf_get_tensor_type(gctx, id) != GGML_TYPE_F32 || - !source || !ggml_are_same_shape(source, contract.destination) || - gguf_get_tensor_size(gctx, id) != - (size_t)ggml_nelements(contract.destination) * sizeof(float)) { - return fail(std::string("scoring head tensor contract mismatch: ") + - contract.name); - } - } - const size_t data_offset = gguf_get_data_offset(gctx); - GgufMmap mmap; - std::string mmap_error; - if (!mmap.open(path, mmap_error)) { - return fail(mmap_error); - } - for (const auto & contract : contracts) { - const int64_t id = gguf_find_tensor(gctx, contract.name); - const size_t offset = gguf_get_tensor_offset(gctx, id); - const size_t size = gguf_get_tensor_size(gctx, id); - if (data_offset > mmap.size() || offset > mmap.size() - data_offset || - size > mmap.size() - data_offset - offset || - !copy_tensor_from_file( - gctx, contract.name, mmap.data(), data_offset, contract.destination)) { - return fail(std::string("scoring head tensor load failed: ") + - contract.name); - } - } - gguf_free(gctx); - if (tensor_ctx) ggml_free(tensor_ctx); - out.scoring_head_loaded = true; - std::fprintf(stderr, "[qwen3-0.6b] loaded scoring head: %s\n", path.c_str()); - return true; -} - } // namespace -bool load_qwen3_drafter_model(const std::string & path, - ggml_backend_t backend, - Qwen3DrafterWeights & out) { +bool load_qwen3_model(const std::string & path, + ggml_backend_t backend, + Qwen3Weights & out) { out.backend = backend; const BackendPrecisionPolicy precision = select_drafter_precision_policy(backend); out.weight_type = precision.weight_type; @@ -294,7 +209,7 @@ bool load_qwen3_drafter_model(const std::string & path, out.buf = ggml_backend_alloc_ctx_tensors(out.ctx, backend); if (!out.buf) { - set_last_error("ggml_backend_alloc_ctx_tensors failed for Qwen3-0.6B drafter"); + set_last_error("ggml_backend_alloc_ctx_tensors failed for Qwen3-0.6B"); gguf_free(gctx); ggml_free(out.ctx); out.ctx = nullptr; @@ -336,10 +251,10 @@ bool load_qwen3_drafter_model(const std::string & path, const size_t off = gguf_get_tensor_offset(gctx, i); // relative to data_off const size_t sz = gguf_get_tensor_size(gctx, i); if (data_off > file_size || off > data_avail || sz > data_avail - off) { - set_last_error(std::string("Qwen3-0.6B drafter GGUF is truncated or corrupt: tensor '") + set_last_error(std::string("Qwen3-0.6B GGUF is truncated or corrupt: tensor '") + gguf_get_tensor_name(gctx, i) + "' data ends at " + std::to_string(data_off + off + sz) + " but file is only " + std::to_string(file_size) - + " bytes. Re-download the drafter model (" + path + ")."); + + " bytes. Re-download the model (" + path + ")."); gguf_free(gctx); ggml_backend_buffer_free(out.buf); ggml_free(out.ctx); @@ -387,31 +302,13 @@ bool load_qwen3_drafter_model(const std::string & path, out.ctx = nullptr; return false; } - if (const char * head_path = std::getenv("PFLASH_SCORING_HEAD_GGUF")) { - constexpr const char * expected_drafter_sha256 = - "f9c9f1d3c1e21755b82d4e165f88dbbbd4355646d632fb5d6cef7c66ed4ee04e"; - const auto drafter_identity = read_gguf_metadata(path, true); - if (!*head_path || out.n_layer < 14 || !drafter_identity.ok || - drafter_identity.sha256 != expected_drafter_sha256 || - !load_scoring_head(head_path, drafter_identity.sha256, out)) { - if (drafter_identity.sha256 != expected_drafter_sha256) { - set_last_error("scoring head requires the pinned Qwen3-0.6B drafter GGUF"); - } - ggml_backend_buffer_free(out.buf); - ggml_free(out.ctx); - out.buf = nullptr; - out.ctx = nullptr; - return false; - } - } return true; } -void free_qwen3_drafter_model(Qwen3DrafterWeights & w) { +void free_qwen3_model(Qwen3Weights & w) { if (w.buf) { ggml_backend_buffer_free(w.buf); w.buf = nullptr; } if (w.ctx) { ggml_free(w.ctx); w.ctx = nullptr; } w.layers.clear(); - w.scoring_head_loaded = false; w.tok_embd = w.out_norm = w.output = nullptr; w.backend = nullptr; } diff --git a/server/src/qwen3/qwen3_model.h b/server/src/qwen3/qwen3_model.h new file mode 100644 index 000000000..b38bd1683 --- /dev/null +++ b/server/src/qwen3/qwen3_model.h @@ -0,0 +1,74 @@ +// Qwen3-0.6B model weights, loaded in-process (no libllama). +// +// Used by Qwen3Backend for standalone inference; the pflash drafter moved to +// the Qwen3.5-0.8B scorer under src/pflash/. +// +// Public API: +// bool load_qwen3_model(path, backend, out) → load GGUF weights +// void free_qwen3_model(weights) +// +#pragma once + +#include "ggml.h" + +#include +#include +#include +#include +#include + +struct ggml_context; +struct ggml_tensor; +struct ggml_backend; +typedef struct ggml_backend * ggml_backend_t; +struct ggml_backend_buffer; +typedef struct ggml_backend_buffer * ggml_backend_buffer_t; + +namespace dflash::common { + +struct Qwen3Layer { + ggml_tensor * attn_norm = nullptr; // [hidden] + ggml_tensor * wq = nullptr; // [hidden, q_dim] = [1024, 2048] + ggml_tensor * wk = nullptr; // [hidden, kv_dim] = [1024, 1024] + ggml_tensor * wv = nullptr; // [hidden, kv_dim] + ggml_tensor * wo = nullptr; // [q_dim, hidden] = [2048, 1024] + ggml_tensor * q_norm = nullptr; // [head_dim] = [128] + ggml_tensor * k_norm = nullptr; // [head_dim] + ggml_tensor * ffn_norm = nullptr; // [hidden] + ggml_tensor * ffn_gate = nullptr; // [hidden, ffn] + ggml_tensor * ffn_up = nullptr; // [hidden, ffn] + ggml_tensor * ffn_down = nullptr; // [ffn, hidden] +}; + +struct Qwen3Weights { + ggml_context * ctx = nullptr; + ggml_backend_t backend = nullptr; + ggml_backend_buffer_t buf = nullptr; + ggml_type weight_type = GGML_TYPE_BF16; + ggml_type compute_type = GGML_TYPE_BF16; + + ggml_tensor * tok_embd = nullptr; // [hidden, vocab] + ggml_tensor * out_norm = nullptr; // [hidden] + ggml_tensor * output = nullptr; // [hidden, vocab] (lm_head) + + std::vector layers; // size = n_layer = 28 + + // Architecture metadata. + int n_layer = 28; + int n_head = 16; + int n_head_kv = 8; + int n_embd = 1024; + int n_ff = 3072; + int head_dim = 128; + int n_vocab = 151936; + int n_ctx_max = 40960; + float rope_theta = 1000000.0f; +}; + +bool load_qwen3_model(const std::string & gguf_path, + ggml_backend_t backend, + Qwen3Weights & out); + +void free_qwen3_model(Qwen3Weights & w); + +} // namespace dflash::common diff --git a/server/src/qwen35/qwen35_backend.cpp b/server/src/qwen35/qwen35_backend.cpp index 4e1150bed..e7a0ad09c 100644 --- a/server/src/qwen35/qwen35_backend.cpp +++ b/server/src/qwen35/qwen35_backend.cpp @@ -25,8 +25,8 @@ #include "common/restore_delta.h" #include "common/specla_mode.h" #include "qwen35_tensor_parallel.h" -#include "qwen3/qwen3_drafter.h" -#include "qwen3/qwen3_kvflash_scorer.h" +#include "pflash/pflash_drafter.h" +#include "pflash/kvflash_drafter_scorer.h" #include "ggml-cuda.h" #include "ggml-backend-impl.h" @@ -649,7 +649,7 @@ bool Qwen35Backend::init() { kvflash_qk_policy_ ? "qk (target pooled-K vs decode query)" : !kvflash_drafter_path_.empty() ? "drafter (attaches on first reselect)" - : "lru (recency-only: no Qwen3-0.6B drafter found " + : "lru (recency-only: no Qwen3.5-0.8B drafter found " "next to the model or in --prefill-drafter)"); std::fflush(stdout); } @@ -1194,10 +1194,14 @@ std::vector Qwen35Backend::compress_batch( if (request.input_ids.empty() || request.drafter_path.empty()) continue; auto & result = results[index]; + // score_query_end < 0 is the legacy "tail window" request value; + // the qwen35 scorer requires an explicit end. + const int score_query_end = request.score_query_end >= 0 + ? request.score_query_end : (int)request.input_ids.size(); result.compressed_ids = drafter_score_and_compress( drafter_ctx_, request.input_ids, request.keep_ratio, /*chunk_size=*/32, request.score_query_tokens, /*pool_kernel=*/13, - request.score_query_end, request.required_instruction_spans); + score_query_end, request.required_instruction_spans); result.ok = !result.compressed_ids.empty(); if (result.ok) { std::fprintf(stderr, "[compress] %zu -> %zu tokens\n", @@ -1240,7 +1244,7 @@ bool Qwen35Backend::handle_compress(const std::string & line, const DaemonIO & i req.keep_ratio = (float)keep_x1000 / 1000.0f; req.drafter_path = (n >= 3 && drafter_path[0]) ? drafter_path - : "/opt/lucebox/models/drafter/Qwen3-0.6B-BF16.gguf"; + : "/opt/lucebox/models/drafter/Qwen3.5-0.8B-BF16.gguf"; { size_t total_vram = 0; int dev = 0; diff --git a/server/src/qwen35/qwen35_backend.h b/server/src/qwen35/qwen35_backend.h index e55671b9e..b42146609 100644 --- a/server/src/qwen35/qwen35_backend.h +++ b/server/src/qwen35/qwen35_backend.h @@ -26,7 +26,7 @@ #include "common/concurrency/paged_kv_pool.h" #include "concurrency/qwen35_seq_engine.h" #include "internal.h" // TargetWeights, TargetCache, DraftWeights, PrefixSnapshot -#include "qwen3/qwen3_drafter.h" // DrafterContext, load_drafter, free_drafter, drafter_score_and_compress +#include "pflash/pflash_drafter.h" // DrafterContext, load_drafter, free_drafter, drafter_score_and_compress #include "kvflash_pager.h" // bounded KV residency pool #include "kvflash_scorer.h" // chunk-relevance policy interface #include "kvflash_qk.h" // target-QK scorer (pooled keys + query) diff --git a/server/src/qwen35/qwen35_layer_split_adapter.cpp b/server/src/qwen35/qwen35_layer_split_adapter.cpp index 3189647ec..7acfcd41e 100644 --- a/server/src/qwen35/qwen35_layer_split_adapter.cpp +++ b/server/src/qwen35/qwen35_layer_split_adapter.cpp @@ -14,8 +14,8 @@ #include "qwen35/layer_split_forward.h" #include "qwen35/qwen35_layer_split_dflash_target.h" #include "qwen35/prefill_helpers.h" -#include "qwen3/qwen3_drafter.h" -#include "qwen3/qwen3_kvflash_scorer.h" +#include "pflash/pflash_drafter.h" +#include "pflash/kvflash_drafter_scorer.h" #include "kv_quant.h" #include "ggml-cuda.h" @@ -210,7 +210,7 @@ bool Qwen35LayerSplitAdapter::kvflash_attach() { kvflash_tau_, !kvflash_drafter_path_.empty() ? "drafter (attaches on first reselect)" - : "lru (recency-only: no Qwen3-0.6B drafter found)"); + : "lru (recency-only: no Qwen3.5-0.8B drafter found)"); std::fflush(stdout); return true; } @@ -1360,7 +1360,7 @@ bool Qwen35LayerSplitAdapter::decode_dflash( } const char * Qwen35LayerSplitAdapter::default_compress_drafter_path() const { - return "/opt/lucebox/models/drafter/Qwen3-0.6B-BF16.gguf"; + return "/opt/lucebox/models/drafter/Qwen3.5-0.8B-BF16.gguf"; } ModelBackend::CompressResult @@ -1385,10 +1385,14 @@ Qwen35LayerSplitAdapter::compress(const ModelBackend::CompressRequest & req) { std::fprintf(stderr, "[target-split][compress] drafter ready\n"); } + // score_query_end < 0 is the legacy "tail window" request value; the + // qwen35 scorer requires an explicit end. + const int score_query_end = req.score_query_end >= 0 + ? req.score_query_end : (int)req.input_ids.size(); result.compressed_ids = drafter_score_and_compress( pflash_drafter_, req.input_ids, req.keep_ratio, /*chunk_size=*/32, req.score_query_tokens, /*pool_kernel=*/13, - req.score_query_end, req.required_instruction_spans); + score_query_end, req.required_instruction_spans); result.ok = !result.compressed_ids.empty(); if (result.ok) { std::fprintf(stderr, "[target-split][compress] %zu -> %zu tokens\n", diff --git a/server/src/qwen35/qwen35_layer_split_adapter.h b/server/src/qwen35/qwen35_layer_split_adapter.h index 06fb11c40..7d62d343f 100644 --- a/server/src/qwen35/qwen35_layer_split_adapter.h +++ b/server/src/qwen35/qwen35_layer_split_adapter.h @@ -12,7 +12,7 @@ #include "placement/placement_config.h" #include "placement/remote_draft_config.h" #include "placement/remote_target_shard_config.h" -#include "qwen3/qwen3_drafter.h" +#include "pflash/pflash_drafter.h" #include "qwen35_target_shard_ipc.h" #include "step_graph.h" #include "internal.h" diff --git a/server/src/server/http_server.cpp b/server/src/server/http_server.cpp index 933a19ad7..37199888e 100644 --- a/server/src/server/http_server.cpp +++ b/server/src/server/http_server.cpp @@ -27,7 +27,7 @@ #include "pin_friendly_prompt.h" #include "common/kv_rotation.h" #include "common/sha1.h" -#include "qwen3/pflash_selection.h" +#include "pflash/pflash_selection.h" #include "freeze_history.h" #ifdef DFLASH_HAS_CURL @@ -3326,7 +3326,7 @@ void HttpServer::apply_flowkv_compression( std::string HttpServer::apply_pflash_compression( const ParsedRequest & req, PreparedPrompt & prepared) { const bool selection_environment = - dflash::qwen3::has_pflash_selection_environment(); + dflash::pflash::has_pflash_selection_environment(); auto [full_slot, full_len] = prefix_cache_.lookup_full(req.prompt_tokens); if (http_detail::pflash_full_cache_restore_allowed( selection_environment) && full_slot >= 0) { @@ -3350,9 +3350,9 @@ std::string HttpServer::apply_pflash_compression( return "PFlash drafter tokenizer produced an empty prompt"; } - dflash::qwen3::PFlashSelectionConfig experiment; + dflash::pflash::PFlashSelectionConfig experiment; std::string experiment_error; - if (!dflash::qwen3::resolve_pflash_selection( + if (!dflash::pflash::resolve_pflash_selection( (int) drafter_ids.size(), 32, experiment, experiment_error)) { return "invalid PFlash strict selection config: " + experiment_error; } @@ -3399,12 +3399,12 @@ std::string HttpServer::apply_pflash_compression( int boundary_index = (int) messages.size() - 1; if (!raw_text_input && experiment.query_parser == - dflash::qwen3::PFlashQueryParser::SemanticUser) { + dflash::pflash::PFlashQueryParser::SemanticUser) { boundary_index = last_user_index; } if (boundary_index < 0 || (experiment.query_parser == - dflash::qwen3::PFlashQueryParser::SemanticUser && + dflash::pflash::PFlashQueryParser::SemanticUser && !raw_text_input && last_user_text.empty())) { return "PFlash strict selection latest-user boundary is unavailable"; } @@ -3566,7 +3566,7 @@ std::string HttpServer::apply_pflash_compression( // merges like " What" inside the span. if (!req.pflash_query.empty() && experiment.query_parser == - dflash::qwen3::PFlashQueryParser::SemanticUser) { + dflash::pflash::PFlashQueryParser::SemanticUser) { explicit_query_span = http_detail::pflash_decoded_text_span( *drafter_tokenizer_, drafter_ids, @@ -3583,7 +3583,7 @@ std::string HttpServer::apply_pflash_compression( http_detail::canonicalize_pflash_token_spans( std::move(required_instruction_spans)); std::string instruction_error; - if (!dflash::qwen3::validate_pflash_instruction_spans( + if (!dflash::pflash::validate_pflash_instruction_spans( required_instruction_spans, (int) drafter_ids.size(), instruction_error)) { return "PFlash strict selection instruction mapping failed: " + @@ -3634,7 +3634,7 @@ std::string HttpServer::apply_pflash_compression( query_content_end, query_content_begin); } else if (experiment.configured && experiment.query_parser == - dflash::qwen3::PFlashQueryParser::ArbitraryTail) { + dflash::pflash::PFlashQueryParser::ArbitraryTail) { parser_selection_rule = "prompt_tail"; query_window = http_detail::pflash_tail_query_window( drafter_ids, experiment.query_tokens, query_content_end); @@ -3687,7 +3687,7 @@ std::string HttpServer::apply_pflash_compression( {"input_kind", parser_input_kind}, {"selection_rule", parser_selection_rule}, {"query_parser", - dflash::qwen3::pflash_query_parser_name( + dflash::pflash::pflash_query_parser_name( experiment.query_parser)}, {"input_tokens", (int) compress_request.input_ids.size()}, {"input_fingerprint_fnv1a64", @@ -3867,11 +3867,11 @@ HttpServer::PreparedPrompt HttpServer::prepare_prompt( const bool continuation = should_compress && is_continuation_request(req.messages); const bool selection_environment = - dflash::qwen3::has_pflash_selection_environment(); + dflash::pflash::has_pflash_selection_environment(); if (should_compress && selection_environment) { - dflash::qwen3::PFlashSelectionConfig experiment; + dflash::pflash::PFlashSelectionConfig experiment; std::string experiment_error; - if (!dflash::qwen3::resolve_pflash_selection( + if (!dflash::pflash::resolve_pflash_selection( 0, 32, experiment, experiment_error)) { prepared.error_status = 500; prepared.error = "invalid PFlash strict selection config: " + diff --git a/server/src/server/http_server.h b/server/src/server/http_server.h index c7fad0165..2a628c132 100644 --- a/server/src/server/http_server.h +++ b/server/src/server/http_server.h @@ -216,7 +216,7 @@ struct ServerConfig { PflashMode pflash_mode = PflashMode::OFF; int pflash_threshold = 32000; // token count threshold for AUTO mode float pflash_keep_ratio = 0.05f; // fraction of tokens to keep - std::string pflash_drafter_path; // path to drafter GGUF (Qwen3-0.6B) + std::string pflash_drafter_path; // path to drafter GGUF (Qwen3.5-0.8B) int pflash_drafter_gpu = 0; // backend-local GPU for PFlash drafter bool pflash_remote_drafter = false; // use IPC drafter for mixed backends RemoteDraftConfig pflash_remote; // IPC binary/work-dir for remote PFlash drafter diff --git a/server/src/server/server_main.cpp b/server/src/server/server_main.cpp index 8382a7c0c..bdb419353 100644 --- a/server/src/server/server_main.cpp +++ b/server/src/server/server_main.cpp @@ -201,7 +201,7 @@ static void print_usage(const char * prog) { " (token,ratio) breakpoints; linear interp.\n" " Overrides --prefill-keep-ratio. Example:\n" " 10000:0.5 40000:0.2 100000:0.1\n" - " --prefill-drafter Drafter GGUF for compression (Qwen3-0.6B)\n" + " --prefill-drafter Drafter GGUF for compression (Qwen3.5-0.8B)\n" " --prefill-skip-park Skip park/unpark (for >=32GB GPUs)\n" " --draft-residency auto|persistent|request-scoped\n" " Drafter lifetime policy (default: auto)\n" diff --git a/server/test/bench_laguna_pflash.cpp b/server/test/bench_laguna_pflash.cpp index 46b2311e0..c02e1b54c 100644 --- a/server/test/bench_laguna_pflash.cpp +++ b/server/test/bench_laguna_pflash.cpp @@ -1,7 +1,7 @@ // End-to-end PFlash + Laguna TTFT bench. Mirrors the qwen3.6-27B PFlash flow: // // 1. Tokenize input (synthetic in DRAFTER vocab for the bench) -// 2. Drafter (Qwen3-0.6B BF16) score_and_compress -> surviving Qwen3 IDs +// 2. Drafter (Qwen3.5-0.8B BF16) score_and_compress -> surviving Qwen3.5 IDs // 3. Cross-tokenizer mapping Qwen3 IDs -> Laguna IDs (NOT plumbed yet; we // use a fake target token for compute-time-only measurement) // 4. Laguna build_laguna_graph dense prefill on the COMPRESSED sequence @@ -13,7 +13,8 @@ #include "laguna_internal.h" #include "internal.h" -#include "qwen3_drafter.h" +#include "pflash/pflash_drafter.h" +#include "pflash/qwen35_drafter.h" #include "dflash27b.h" #include @@ -102,12 +103,13 @@ int main(int argc, char ** argv) { } auto td1 = std::chrono::steady_clock::now(); std::printf("[pflash] drafter loaded in %.2fs vocab=%d\n", - std::chrono::duration(td1 - td0).count(), drafter.weights.n_vocab); + std::chrono::duration(td1 - td0).count(), drafter.state->weights.n_vocab); std::vector input(N, fake_q); auto tc0 = std::chrono::steady_clock::now(); std::vector compressed = drafter_score_and_compress( - drafter, input, keep_r, /*chunk_size=*/32, /*n_lookahead=*/8, /*pool_kernel=*/13); + drafter, input, keep_r, /*chunk_size=*/32, /*n_lookahead=*/8, /*pool_kernel=*/13, + /*score_query_end=*/(int)input.size()); auto tc1 = std::chrono::steady_clock::now(); if (compressed.empty()) { std::fprintf(stderr, "drafter compress failed: %s\n", dflash27b_last_error()); diff --git a/server/test/pflash_daemon.cpp b/server/test/pflash_daemon.cpp index 38e291e8a..ecdc4e47c 100644 --- a/server/test/pflash_daemon.cpp +++ b/server/test/pflash_daemon.cpp @@ -1,6 +1,6 @@ // Persistent PFlash compressor daemon. // -// Loads the Qwen3-0.6B PFlash drafter once, then accepts stdin commands: +// Loads the Qwen3.5-0.8B PFlash drafter once, then accepts stdin commands: // // compress // quit @@ -10,7 +10,8 @@ // values to --stream-fd=, terminated by -1. Logs go to stdout/stderr. #include "dflash27b.h" -#include "qwen3_drafter.h" +#include "pflash/pflash_drafter.h" +#include "pflash/qwen35_drafter.h" #include #include @@ -68,7 +69,7 @@ static void stream_ids(int stream_fd, const std::vector & ids) { int main(int argc, char ** argv) { if (argc < 2) { - std::fprintf(stderr, "usage: %s [--stream-fd=N]\n", argv[0]); + std::fprintf(stderr, "usage: %s [--stream-fd=N]\n", argv[0]); return 2; } @@ -89,7 +90,7 @@ int main(int argc, char ** argv) { auto t_load1 = std::chrono::steady_clock::now(); std::printf("[pflash-daemon] ready load=%.3fs vocab=%d\n", std::chrono::duration(t_load1 - t_load0).count(), - ctx.weights.n_vocab); + ctx.state->weights.n_vocab); std::fflush(stdout); std::string line; @@ -139,7 +140,7 @@ int main(int argc, char ** argv) { std::fflush(stdout); auto t0 = std::chrono::steady_clock::now(); - std::vector out = drafter_score_and_compress(ctx, ids, keep_ratio, chunk, lookahead, pool); + std::vector out = drafter_score_and_compress(ctx, ids, keep_ratio, chunk, lookahead, pool, (int)ids.size()); auto t1 = std::chrono::steady_clock::now(); const double secs = std::chrono::duration(t1 - t0).count(); diff --git a/server/test/smoke_qwen3_forward.cpp b/server/test/smoke_qwen3_forward.cpp index 4efa41781..df552230d 100644 --- a/server/test/smoke_qwen3_forward.cpp +++ b/server/test/smoke_qwen3_forward.cpp @@ -1,4 +1,4 @@ -// Smoke test for the custom Qwen3-0.6B drafter forward path. +// Smoke test for the Qwen3.5-0.8B PFlash drafter forward path. // // Loads the BF16 GGUF, generates a synthetic token sequence at the requested // length, runs drafter_score_and_compress end-to-end, and prints timing + @@ -8,12 +8,13 @@ // Usage: // smoke_qwen3_forward [keep_ratio] // Examples: -// smoke_qwen3_forward .../Qwen3-0.6B-BF16.gguf 140000 0.02 -// smoke_qwen3_forward .../Qwen3-0.6B-BF16.gguf FILE:/tmp/niah_32k.bin 0.05 +// smoke_qwen3_forward .../Qwen3.5-0.8B-BF16.gguf 140000 0.02 +// smoke_qwen3_forward .../Qwen3.5-0.8B-BF16.gguf FILE:/tmp/niah_32k.bin 0.05 // // Token file format: little-endian u32 count, then count int32 token IDs. -#include "qwen3_drafter.h" +#include "pflash/pflash_drafter.h" +#include "pflash/qwen35_drafter.h" #include "dflash27b.h" #include @@ -71,7 +72,7 @@ int main(int argc, char ** argv) { auto t_load1 = std::chrono::steady_clock::now(); std::printf("[smoke] load_drafter %.2fs vocab=%d\n", std::chrono::duration(t_load1 - t_load0).count(), - ctx.weights.n_vocab); + ctx.state->weights.n_vocab); std::vector ids; if (from_file) { @@ -79,7 +80,7 @@ int main(int argc, char ** argv) { } else { ids.resize((size_t)S); std::mt19937 rng(42); - std::uniform_int_distribution dist(0, ctx.weights.n_vocab - 1); + std::uniform_int_distribution dist(0, ctx.state->weights.n_vocab - 1); for (int i = 0; i < S; ++i) ids[i] = dist(rng); } @@ -90,7 +91,8 @@ int main(int argc, char ** argv) { auto t0 = std::chrono::steady_clock::now(); std::vector out = drafter_score_and_compress( ctx, ids, keep_ratio, - /*chunk_size=*/32, /*n_lookahead=*/8, /*pool_kernel=*/13); + /*chunk_size=*/32, /*n_lookahead=*/8, /*pool_kernel=*/13, + /*score_query_end=*/(int)ids.size()); auto t1 = std::chrono::steady_clock::now(); if (out.empty()) { diff --git a/server/test/test_anchor_params.cpp b/server/test/test_anchor_params.cpp index 285d945a0..fd0bf952c 100644 --- a/server/test/test_anchor_params.cpp +++ b/server/test/test_anchor_params.cpp @@ -1,7 +1,7 @@ // Unit tests for resolve_anchor_params() — no GPU, no model files. #include "CppUnitTestFramework.hpp" -#include "qwen3/anchor_params.h" +#include "pflash/anchor_params.h" using namespace dflash::common; diff --git a/server/test/test_anchor_transitive.cpp b/server/test/test_anchor_transitive.cpp index c2ba4005f..c94750f43 100644 --- a/server/test/test_anchor_transitive.cpp +++ b/server/test/test_anchor_transitive.cpp @@ -2,7 +2,7 @@ // T1: single-pass match; T2: single-pass misses hops; T3: transitive rescues all hops. #include "CppUnitTestFramework.hpp" -#include "../src/qwen3/anchor_scan.h" +#include "../src/pflash/anchor_scan.h" #include #include @@ -52,9 +52,9 @@ static void t1_single_pass_match() { const int n_chunks = (N + CHUNK - 1) / CHUNK; std::vector forced((size_t)n_chunks, 0); - dflash::qwen3::AnchorScanCfg cfg{CHUNK, /*anchor_radius=*/0, + dflash::pflash::AnchorScanCfg cfg{CHUNK, /*anchor_radius=*/0, /*max_anchor_hits=*/8, /*ngram=*/4}; - dflash::qwen3::scan_and_force(ids, q0, query_pool, cfg, forced); + dflash::pflash::scan_and_force(ids, q0, query_pool, cfg, forced); // Chunk containing pos 100 must be forced. const int target_chunk = 100 / CHUNK; // chunk 1 @@ -88,9 +88,9 @@ static void t2_single_pass_misses_hops() { const int n_chunks = (N + CHUNK - 1) / CHUNK; std::vector forced((size_t)n_chunks, 0); - dflash::qwen3::AnchorScanCfg cfg{CHUNK, /*anchor_radius=*/0, + dflash::pflash::AnchorScanCfg cfg{CHUNK, /*anchor_radius=*/0, /*max_anchor_hits=*/8, /*ngram=*/4}; - dflash::qwen3::scan_and_force(ids, q0, query_pool, cfg, forced); + dflash::pflash::scan_and_force(ids, q0, query_pool, cfg, forced); const int chunk_hop3 = 1200 / CHUNK; // 18 const int chunk_hop2 = 600 / CHUNK; // 9 @@ -126,9 +126,9 @@ static void t3_transitive_rescues_all() { const int n_chunks = (N + CHUNK - 1) / CHUNK; std::vector forced((size_t)n_chunks, 0); - dflash::qwen3::AnchorScanCfg cfg{CHUNK, /*anchor_radius=*/0, + dflash::pflash::AnchorScanCfg cfg{CHUNK, /*anchor_radius=*/0, /*max_anchor_hits=*/8, /*ngram=*/4}; - dflash::qwen3::scan_and_force_transitive(ids, q0, initial_query_pool, + dflash::pflash::scan_and_force_transitive(ids, q0, initial_query_pool, cfg, /*max_iters=*/3, forced); const int chunk_hop3 = 1200 / CHUNK; @@ -187,10 +187,10 @@ static void t4_rare_token_bridges_different_context() { const int n_chunks = (N + CHUNK - 1) / CHUNK; std::vector forced((size_t)n_chunks, 0); - dflash::qwen3::AnchorScanCfg cfg{CHUNK, /*anchor_radius=*/0, + dflash::pflash::AnchorScanCfg cfg{CHUNK, /*anchor_radius=*/0, /*max_anchor_hits=*/8, /*ngram=*/4, /*rare_token_max_freq=*/8}; - dflash::qwen3::scan_and_force_transitive(ids, q0, initial_query_pool, + dflash::pflash::scan_and_force_transitive(ids, q0, initial_query_pool, cfg, /*max_iters=*/3, forced); const int chunk_hop3 = 1200 / CHUNK; // 18 @@ -250,12 +250,12 @@ static void t5_gate_closes_when_pass1_finds_many() { // --- Test A: gate CLOSED (cascade_min_anchor_count=5) --- { std::vector forced_a((size_t)n_chunks, 0); - dflash::qwen3::AnchorScanCfg cfg{CHUNK, /*anchor_radius=*/0, + dflash::pflash::AnchorScanCfg cfg{CHUNK, /*anchor_radius=*/0, /*max_anchor_hits=*/64, /*ngram=*/4, /*rare_token_max_freq=*/2, /*cascade_min_anchor_count=*/5, /*max_forced_count=*/INT_MAX}; - dflash::qwen3::scan_and_force_transitive(ids, q0, query_pool, + dflash::pflash::scan_and_force_transitive(ids, q0, query_pool, cfg, /*max_iters=*/3, forced_a); // Pass-1 forces chunks 0..49 (50 chunks); gate closes → cascade skipped. @@ -272,12 +272,12 @@ static void t5_gate_closes_when_pass1_finds_many() { // --- Test B: gate OPEN (cascade_min_anchor_count=0) → cascade forces chunk 60 --- { std::vector forced_b((size_t)n_chunks, 0); - dflash::qwen3::AnchorScanCfg cfg{CHUNK, /*anchor_radius=*/0, + dflash::pflash::AnchorScanCfg cfg{CHUNK, /*anchor_radius=*/0, /*max_anchor_hits=*/64, /*ngram=*/4, /*rare_token_max_freq=*/2, /*cascade_min_anchor_count=*/0, /*max_forced_count=*/INT_MAX}; - dflash::qwen3::scan_and_force_transitive(ids, q0, query_pool, + dflash::pflash::scan_and_force_transitive(ids, q0, query_pool, cfg, /*max_iters=*/3, forced_b); // Cascade runs; chunk 5 is forced by pass-1 and contains RT; @@ -330,12 +330,12 @@ static void t6_hard_cap_prevents_runaway() { // Without cap: cascade forces chunks 0..20 (21 chunks total). // With cap=5: stops at 5. std::vector forced((size_t)n_chunks, 0); - dflash::qwen3::AnchorScanCfg cfg{CHUNK, /*anchor_radius=*/0, + dflash::pflash::AnchorScanCfg cfg{CHUNK, /*anchor_radius=*/0, /*max_anchor_hits=*/8, /*ngram=*/4, /*rare_token_max_freq=*/2, /*cascade_min_anchor_count=*/0, /*max_forced_count=*/5}; - dflash::qwen3::scan_and_force_transitive(ids, q0, query_pool, + dflash::pflash::scan_and_force_transitive(ids, q0, query_pool, cfg, /*max_iters=*/25, forced); int total_forced = 0; diff --git a/server/test/test_dflash.cpp b/server/test/test_dflash.cpp index d33399d97..5fcc2bc4c 100644 --- a/server/test/test_dflash.cpp +++ b/server/test/test_dflash.cpp @@ -25,7 +25,7 @@ #include "specla_commit_cuda.h" #include "specla_mode.h" #include "draft_graph.h" -#include "qwen3_drafter.h" +#include "pflash/pflash_drafter.h" #include "gpu_runtime_compat.h" #include "chain_rollback_policy.h" #include "draft_swa.h" @@ -2379,8 +2379,8 @@ int main(int argc, char ** argv) { // Format: "compress [drafter_arch]" // src_bin_path: int32 token IDs file (drafter vocab) // keep_ratio_x1000: integer keep ratio × 1000 (e.g. 20 → 0.020) - // drafter_gguf: path to drafter GGUF (loaded lazily once) - // drafter_arch: qwen3-0.6b (default) or qwen35-0.8b + // drafter_gguf: path to the Qwen3.5-0.8B drafter GGUF (loaded lazily once) + // drafter_arch: accepted for compatibility, ignored (Qwen3.5-0.8B only) // Output: stream of int32 compressed token IDs, terminated by -1. // Drafter coexists with target+draft via libllama in the same // ggml allocator — no park/unpark needed for compression itself. @@ -2388,7 +2388,7 @@ int main(int argc, char ** argv) { char ppath[1024]; int keep_x1000 = 0; char drafter_path[1024]; - char arch_name[64] = "qwen3-0.6b"; + char arch_name[64] = ""; int n = std::sscanf(line.c_str() + 9, "%1023s %d %1023s %63s", ppath, &keep_x1000, drafter_path, arch_name); if (n < 3) { @@ -2396,11 +2396,6 @@ int main(int argc, char ** argv) { "[compress] bad args, need: [drafter_arch]\n"); stream_emit(-1); continue; } - dflash::common::DrafterArch drafter_arch; - if (!dflash::common::parse_drafter_arch(arch_name, drafter_arch)) { - std::fprintf(stderr, "[compress] bad drafter_arch: %s\n", arch_name); - stream_emit(-1); continue; - } auto src_ids = read_int32_file(ppath); if (src_ids.empty()) { std::fprintf(stderr, "[compress] empty input\n"); @@ -2429,31 +2424,21 @@ int main(int argc, char ** argv) { } if (!drafter_loaded) { - if (!dflash::common::load_drafter(drafter_path, /*gpu_layers=*/999, drafter_arch, drafter_ctx)) { + if (!dflash::common::load_drafter(drafter_path, /*gpu_layers=*/999, drafter_ctx)) { std::fprintf(stderr, "[compress] load_drafter failed: %s\n", dflash27b_last_error()); stream_emit(-1); continue; } drafter_loaded = true; - if (drafter_arch == dflash::common::DrafterArch::Qwen3_0p6b) { - std::printf("[drafter] loaded %s arch=%s (n_layer=%d n_head=%d n_head_kv=%d)\n", - drafter_path, dflash::common::drafter_arch_name(drafter_arch), drafter_ctx.weights.n_layer, - drafter_ctx.weights.n_head, drafter_ctx.weights.n_head_kv); - } else { - std::printf("[drafter] loaded %s arch=%s\n", - drafter_path, dflash::common::drafter_arch_name(drafter_arch)); - } + std::printf("[drafter] loaded %s\n", drafter_path); std::fflush(stdout); - } else if (drafter_ctx.arch != drafter_arch) { - std::fprintf(stderr, "[compress] requested arch=%s but loaded arch=%s\n", - dflash::common::drafter_arch_name(drafter_arch), - dflash::common::drafter_arch_name(drafter_ctx.arch)); - stream_emit(-1); continue; } float keep = (float)keep_x1000 / 1000.0f; auto compressed = dflash::common::drafter_score_and_compress( - drafter_ctx, src_ids, keep); + drafter_ctx, src_ids, keep, + /*chunk_size=*/32, /*n_lookahead=*/8, /*pool_kernel=*/13, + (int)src_ids.size()); std::printf("[compress] %zu -> %zu tokens (keep_ratio=%.3f)\n", src_ids.size(), compressed.size(), keep); std::fflush(stdout); diff --git a/server/test/test_drafter_early_exit_score_range.cpp b/server/test/test_drafter_early_exit_score_range.cpp deleted file mode 100644 index 047e7e604..000000000 --- a/server/test/test_drafter_early_exit_score_range.cpp +++ /dev/null @@ -1,90 +0,0 @@ -// Unit tests for dflash::common::compute_score_range(). -// SCORE_LAYERS is relative to fwd_layer_limit: ee7+sl7 → [0,7), not phantom-empty [7,7). - -#include "CppUnitTestFramework.hpp" -#include "score_range.h" - -#include -#include - -using dflash::common::ScoreRange; -using dflash::common::compute_score_range; - -namespace { -struct DrafterEarlyExitScoreRangeFixture : CppUnitTestFramework::CommonFixture { - using CppUnitTestFramework::CommonFixture::CommonFixture; - - void t1_bug_scenario() { - ScoreRange r = compute_score_range(/*n_layer=*/28, - /*score_layers=*/7, - /*fwd_layer_limit=*/7); - REQUIRE(r.start == 0 && "score_layer_start must be 0"); - REQUIRE(r.end == 7 && "score_layer_end must equal fwd_layer_limit"); - REQUIRE(!r.empty() && "range must be non-empty"); - REQUIRE(r.count() == 7); - printf("T1 pass: early_exit_n=7 score_layers=7 n_layer=28 -> [%d,%d)\n", - r.start, r.end); - } - - void t2_no_early_exit() { - ScoreRange r = compute_score_range(28, 7, 28); - REQUIRE(r.start == 21); - REQUIRE(r.end == 28); - REQUIRE(!r.empty()); - REQUIRE(r.count() == 7); - printf("T2 pass: no early exit score_layers=7 -> [%d,%d)\n", r.start, r.end); - } - - void t3_all_layers_no_exit() { - ScoreRange r = compute_score_range(28, -1, 28); - REQUIRE(r.start == 0); - REQUIRE(r.end == 28); - REQUIRE(!r.empty()); - printf("T3 pass: score_layers=-1 no exit -> [%d,%d)\n", r.start, r.end); - } - - void t4_all_layers_with_exit() { - ScoreRange r = compute_score_range(28, -1, 14); - REQUIRE(r.start == 0); - REQUIRE(r.end == 14); - REQUIRE(!r.empty()); - printf("T4 pass: score_layers=-1 early_exit=14 -> [%d,%d)\n", r.start, r.end); - } - - void t5_score_layers_exceeds_exit() { - ScoreRange r = compute_score_range(28, 14, 7); - REQUIRE(r.start == 0); - REQUIRE(r.end == 7); - REQUIRE(!r.empty()); - printf("T5 pass: score_layers=14 early_exit=7 -> [%d,%d)\n", r.start, r.end); - } - - void t6_score_layers_equals_n_layer() { - ScoreRange r = compute_score_range(28, 28, 28); - REQUIRE(r.start == 0); - REQUIRE(r.end == 28); - REQUIRE(!r.empty()); - printf("T6 pass: score_layers=n_layer=28 -> [%d,%d)\n", r.start, r.end); - } - - void t7_partial_exit_partial_score() { - ScoreRange r = compute_score_range(28, 7, 14); - REQUIRE(r.start == 7); - REQUIRE(r.end == 14); - REQUIRE(!r.empty()); - REQUIRE(r.count() == 7); - printf("T7 pass: early_exit=14 score_layers=7 -> [%d,%d)\n", r.start, r.end); - } -}; -} - -TEST_CASE(DrafterEarlyExitScoreRangeFixture, score_range_suite) { - t1_bug_scenario(); - t2_no_early_exit(); - t3_all_layers_no_exit(); - t4_all_layers_with_exit(); - t5_score_layers_exceeds_exit(); - t6_score_layers_equals_n_layer(); - t7_partial_exit_partial_score(); - printf("\nAll score_range tests passed.\n"); -} diff --git a/server/test/test_drafter_tail_capture_guard.cpp b/server/test/test_drafter_tail_capture_guard.cpp deleted file mode 100644 index 1ce9d176f..000000000 --- a/server/test/test_drafter_tail_capture_guard.cpp +++ /dev/null @@ -1,118 +0,0 @@ -// Unit tests for the tail-capture chunk-boundary guard in qwen3_graph.cpp. -// Reproduces Bug #42: ggml_view_3d overrun when S % chunk_size ∈ {1..7} -// and n_lookahead == 8. -// -// Pure integer arithmetic — no ggml, no GPU, no server deps. -// -// Root cause (codex's diagnosis, confirmed by momus's data audit): -// tail_lo = S - n_lookahead -// When chunk 0 contains S = chunk_size + r tokens (r ∈ {1..7}), a second -// chunk was dispatched but we still evaluate the first chunk's guard with -// cs=0, cl=chunk_size. tail_lo = chunk_size + r - n_lookahead = 4088 + r. -// -// OLD guard: tail_lo >= cs && tail_lo < cs + cl -// r=1..7: (4088+r) >= 0 && (4088+r) < 4096 → TRUE ← BUG: tail overruns -// -// NEW guard: tail_lo >= cs && tail_lo + n_lookahead <= cs + cl -// r=1..7: (4088+r) + 8 <= 4096 → 4096+r <= 4096 → FALSE ← correct: skip -// -// TDD RED/GREEN: -// RED (before patch): TAIL_GUARD_USE_NEW_FORMULA undefined → old guard inline → test FAILS. -// GREEN (after patch): TAIL_GUARD_USE_NEW_FORMULA defined via compiler flag → test PASSES. -// The patch to qwen3_graph.cpp changes the same 2 lines as this toggle. - -#include "CppUnitTestFramework.hpp" - -#include -#include - -static bool tail_fits(int tail_lo, int cs, int cl, int n_lookahead); - -namespace { -struct DrafterTailCaptureGuardFixture : CppUnitTestFramework::CommonFixture { - using CppUnitTestFramework::CommonFixture::CommonFixture; - - void t1_straddling_tail_must_be_skipped() { - const int chunk_size = 4096, n_lookahead = 8; - const int cs = 0, cl = chunk_size; - - for (int r = 1; r <= 7; r++) { - const int S = chunk_size + r; - const int tail_lo = S - n_lookahead; - - const bool result = tail_fits(tail_lo, cs, cl, n_lookahead); - std::printf("T1 r=%d S=%d tail_lo=%d tail_hi=%d chunk=[%d,%d): fits=%d (expect 0)\n", - r, S, tail_lo, tail_lo + n_lookahead, cs, cs + cl, (int)result); - REQUIRE(!result && "tail overruns chunk boundary — guard must return false"); - } - } - - void t2_tail_fits_exactly_at_chunk_end() { - const int chunk_size = 4096, n_lookahead = 8; - const int cs = 0, cl = chunk_size; - const int S = chunk_size; - const int tail_lo = S - n_lookahead; - - const bool result = tail_fits(tail_lo, cs, cl, n_lookahead); - std::printf("T2 r=0 S=%d tail_lo=%d: fits=%d (expect 1)\n", S, tail_lo, (int)result); - REQUIRE(result && "tail fits exactly at chunk end — must return true"); - } - - void t3_tail_starts_outside_chunk() { - const int chunk_size = 4096, n_lookahead = 8; - const int cs = 0, cl = chunk_size; - const int S = chunk_size + 8; - const int tail_lo = S - n_lookahead; - - const bool result = tail_fits(tail_lo, cs, cl, n_lookahead); - std::printf("T3 r=8 S=%d tail_lo=%d: fits=%d (expect 0)\n", S, tail_lo, (int)result); - REQUIRE(!result && "tail starts at next chunk — must return false"); - } - - void t4_second_chunk_tail_fits_exactly() { - const int chunk_size = 4096, n_lookahead = 8; - const int cs = chunk_size, cl = chunk_size; - const int S = 2 * chunk_size; - const int tail_lo = S - n_lookahead; - - const bool result = tail_fits(tail_lo, cs, cl, n_lookahead); - std::printf("T4 second chunk S=%d tail_lo=%d cs=%d: fits=%d (expect 1)\n", - S, tail_lo, cs, (int)result); - REQUIRE(result && "tail fits exactly in second chunk — must return true"); - } - - void t5_second_chunk_straddling_tail_skipped() { - const int chunk_size = 4096, n_lookahead = 8; - const int cs = chunk_size, cl = chunk_size; - const int r = 3; - const int S = 2 * chunk_size + r; - const int tail_lo = S - n_lookahead; - - const bool result = tail_fits(tail_lo, cs, cl, n_lookahead); - std::printf("T5 second chunk r=%d S=%d tail_lo=%d: fits=%d (expect 0)\n", - r, S, tail_lo, (int)result); - REQUIRE(!result && "tail straddles end of second chunk — must return false"); - } -}; -} - -// The guard being tested — toggled by compile-time flag to reproduce RED/GREEN. -#ifdef TAIL_GUARD_USE_NEW_FORMULA -static bool tail_fits(int tail_lo, int cs, int cl, int n_lookahead) { - return tail_lo >= cs && tail_lo + n_lookahead <= cs + cl; // NEW (fix) -} -#else -static bool tail_fits(int tail_lo, int cs, int cl, int n_lookahead) { - (void)n_lookahead; - return tail_lo >= cs && tail_lo < cs + cl; // OLD (Bug #42) -} -#endif - -TEST_CASE(DrafterTailCaptureGuardFixture, tail_capture_guard_suite) { - t1_straddling_tail_must_be_skipped(); - t2_tail_fits_exactly_at_chunk_end(); - t3_tail_starts_outside_chunk(); - t4_second_chunk_tail_fits_exactly(); - t5_second_chunk_straddling_tail_skipped(); - std::printf("All tail_capture guard tests passed.\n"); -} diff --git a/server/test/test_drafter_warm_path_regression.cpp b/server/test/test_drafter_warm_path_regression.cpp deleted file mode 100644 index 7d16683e6..000000000 --- a/server/test/test_drafter_warm_path_regression.cpp +++ /dev/null @@ -1,170 +0,0 @@ -// Regression test: K_norope_v/Q_norope_v sized to n_score_layers, not n_layer. -// Old code allocated 28 entries (~5.6 GB wasted at 128K); fix uses score_range.count(). - -#include "CppUnitTestFramework.hpp" -#include "score_range.h" - -#include -#include -#include -#include - -using dflash::common::ScoreRange; -using dflash::common::compute_score_range; - -#define TEST_ASSERT(cond) do { \ - if (!(cond)) { \ - throw std::runtime_error(std::string(__FILE__) + ":" + \ - std::to_string(__LINE__) + ": " + #cond); \ - } \ -} while (0) -#undef assert -#define assert(cond) TEST_ASSERT(cond) - -namespace { -struct DrafterWarmPathRegressionFixture {}; -} - -// Helper: compute n_score_layers as the fixed allocator does. -static int score_layer_count(int n_layer, int score_layers_env, int early_exit_env) { - const int fwd_limit = (early_exit_env > 0 && early_exit_env < n_layer) - ? early_exit_env : n_layer; - ScoreRange r = compute_score_range(n_layer, score_layers_env, fwd_limit); - return r.count(); -} - -// T1: baseline case — SCORE_LAYERS unset (-1), no early exit. -// K_norope_v should have n_layer entries. -static void t1_baseline_full_alloc() { - int n = score_layer_count(28, -1, -1); - assert(n == 28 && "baseline: all 28 layers must be allocated"); - printf("T1 pass: baseline n_score_layers=%d\n", n); -} - -// T2: L7 case — SCORE_LAYERS=7, no early exit. -// OLD: allocated 28 entries (5.6 GB wasted). NEW: 7 entries. -static void t2_l7_trimmed_alloc() { - int n = score_layer_count(28, 7, -1); - assert(n == 7 && "L7: only 7 K_norope entries must be allocated"); - printf("T2 pass: L7 n_score_layers=%d (was 28 before fix)\n", n); -} - -// T3: early-exit=14, SCORE_LAYERS=7. Scoring range [7,14), 7 layers. -static void t3_early_exit_with_score_layers() { - int n = score_layer_count(28, 7, 14); - assert(n == 7); - printf("T3 pass: early_exit=14 score_layers=7 -> n_score_layers=%d\n", n); -} - -// T4: early-exit=7, SCORE_LAYERS=7 (the classic double-7 composition). -// Range [0,7), 7 layers. -static void t4_ee7_score7_composition() { - int n = score_layer_count(28, 7, 7); - assert(n == 7); - printf("T4 pass: ee7+score7 n_score_layers=%d\n", n); -} - -// T5: SCORE_LAYERS not set (all layers), early-exit=14. -// Scoring range [0,14), 14 layers needed. -static void t5_all_score_with_early_exit() { - int n = score_layer_count(28, -1, 14); - assert(n == 14); - printf("T5 pass: score_all early_exit=14 n_score_layers=%d\n", n); -} - -// T6: validate that score_layer_start_pre matches score_layer_start used -// in the scoring loop (must be identical for correct buffer indexing). -static void t6_start_pre_matches_loop_start() { - // Replicate the pre-alloc computation. - const int n_layer = 28, score_layers_env = 7, early_exit_env = -1; - const int fwd_limit = (early_exit_env > 0 && early_exit_env < n_layer) - ? early_exit_env : n_layer; - ScoreRange pre = compute_score_range(n_layer, score_layers_env, fwd_limit); - // Scoring loop uses the same fwd_layer_limit (== fwd_limit) and same env. - ScoreRange loop = compute_score_range(n_layer, score_layers_env, fwd_limit); - assert(pre.start == loop.start && "score_layer_start_pre must equal score_layer_start"); - assert(pre.end == loop.end); - printf("T6 pass: pre_start=%d loop_start=%d (match)\n", pre.start, loop.start); -} - -// T7: alloc loop boundary check — the alloc loop iterates 0..n_layer but must only -// fill K_norope_v for layers in [score_layer_start_pre, fwd_layer_limit_pre). -// This replicates the guard added to the alloc loop: il >= start AND il < fwd_limit. -// Before the fix: il was only bounded below (il >= start), causing K_norope_v[si] -// out-of-bounds when n_score_layers < n_layer (e.g. ee14: si 0..27 but vec size 14). -static void t7_alloc_loop_upper_bound() { - struct FakeVec { - int capacity; - int max_si_written = -1; - void write(int si) { - assert(si >= 0 && si < capacity && "si out of bounds"); - if (si > max_si_written) max_si_written = si; - } - }; - - // Simulate ee14 (no SCORE_LAYERS, early_exit=14, n_layer=28). - { - const int n_layer = 28, score_layers = -1, early_exit = 14; - const int fwd_limit = early_exit; - ScoreRange r = compute_score_range(n_layer, score_layers, fwd_limit); - const int n_score = r.count(); // 14 - FakeVec v{n_score}; - int writes = 0; - for (int il = 0; il < n_layer; ++il) { - // Correct guard: il >= start AND il < fwd_limit (the fix) - if (il >= r.start && il < fwd_limit) { - v.write(il - r.start); - writes++; - } - } - assert(writes == n_score && "ee14: must write exactly n_score_layers entries"); - printf("T7a pass: ee14 alloc writes=%d capacity=%d (no overflow)\n", writes, n_score); - } - - // Simulate ee7 (SCORE_LAYERS=7, early_exit=7, n_layer=28). - { - const int n_layer = 28, score_layers = 7, early_exit = 7; - const int fwd_limit = early_exit; - ScoreRange r = compute_score_range(n_layer, score_layers, fwd_limit); - const int n_score = r.count(); // 7 - FakeVec v{n_score}; - int writes = 0; - for (int il = 0; il < n_layer; ++il) { - if (il >= r.start && il < fwd_limit) { - v.write(il - r.start); - writes++; - } - } - assert(writes == n_score && "ee7: must write exactly 7 entries"); - printf("T7b pass: ee7 alloc writes=%d capacity=%d (no overflow)\n", writes, n_score); - } - - // Simulate baseline (no ee, no score_layers). - { - const int n_layer = 28, score_layers = -1, early_exit = -1; - const int fwd_limit = n_layer; - ScoreRange r = compute_score_range(n_layer, score_layers, fwd_limit); - const int n_score = r.count(); // 28 - FakeVec v{n_score}; - int writes = 0; - for (int il = 0; il < n_layer; ++il) { - if (il >= r.start && il < fwd_limit) { - v.write(il - r.start); - writes++; - } - } - assert(writes == n_score && "baseline: must write 28 entries"); - printf("T7c pass: baseline alloc writes=%d capacity=%d (no overflow)\n", writes, n_score); - } -} - -TEST_CASE(DrafterWarmPathRegressionFixture, warm_path_regression_suite) { - t1_baseline_full_alloc(); - t2_l7_trimmed_alloc(); - t3_early_exit_with_score_layers(); - t4_ee7_score7_composition(); - t5_all_score_with_early_exit(); - t6_start_pre_matches_loop_start(); - t7_alloc_loop_upper_bound(); - printf("\nAll warm-path regression tests passed.\n"); -} diff --git a/server/test/test_kvflash.cpp b/server/test/test_kvflash.cpp index 9cfbdd114..03044c900 100644 --- a/server/test/test_kvflash.cpp +++ b/server/test/test_kvflash.cpp @@ -27,8 +27,8 @@ #include "kvflash_qk.h" #include "attn_masks.h" #include "prefill_helpers.h" -#include "qwen3_drafter.h" -#include "qwen3_kvflash_scorer.h" +#include "pflash/pflash_drafter.h" +#include "pflash/kvflash_drafter_scorer.h" #include "ggml.h" #include "ggml-alloc.h" @@ -212,7 +212,7 @@ struct Stepper { std::vector make_prompt(int n, int vocab) { std::vector p(n); uint64_t s = 0x9E3779B97F4A7C15ull; - // Cap below the drafter vocab too (Qwen3-0.6B ~151K) so the same ids + // Cap below the drafter vocab too (Qwen3.5-0.8B) so the same ids // are scoreable by the indexer in run F. const int cap = std::min(vocab, 100000); for (int i = 0; i < n; i++) { @@ -502,7 +502,7 @@ int main(int argc, char ** argv) { KvFlashDrafterScorer dscorer(&dctx); if (is_drafter) { const char * dpath = arg_str(argc, argv, "--qk-drafter", - "/opt/lucebox/models/drafter/Qwen3-0.6B-BF16.gguf"); + "/opt/lucebox/models/drafter/Qwen3.5-0.8B-BF16.gguf"); if (!load_drafter(dpath, 0, dctx)) { std::fprintf(stderr, "drafter load failed\n"); return 1; @@ -720,7 +720,7 @@ int main(int argc, char ** argv) { for (int mode = 0; mode < 2; mode++) { // 0=baseline 1=pool if (only_mode >= 0 && mode != only_mode) continue; if (mode == 1 && !dctx.loaded && - !load_drafter("/opt/lucebox/models/drafter/Qwen3-0.6B-BF16.gguf", 0, dctx)) { + !load_drafter("/opt/lucebox/models/drafter/Qwen3.5-0.8B-BF16.gguf", 0, dctx)) { std::fprintf(stderr, "drafter load failed\n"); return 1; } @@ -808,7 +808,7 @@ int main(int argc, char ** argv) { // inside the recency window is the induction control (distance-free). if (arg_flag(argc, argv, "--niah256")) { DrafterContext dctx; - if (!load_drafter("/opt/lucebox/models/drafter/Qwen3-0.6B-BF16.gguf", 0, dctx)) { + if (!load_drafter("/opt/lucebox/models/drafter/Qwen3.5-0.8B-BF16.gguf", 0, dctx)) { std::fprintf(stderr, "drafter load failed\n"); return 1; } @@ -895,7 +895,7 @@ int main(int argc, char ** argv) { if (arg_flag(argc, argv, "--niah")) { DrafterContext dctx; const bool have_drafter = - load_drafter("/opt/lucebox/models/drafter/Qwen3-0.6B-BF16.gguf", 0, dctx); + load_drafter("/opt/lucebox/models/drafter/Qwen3.5-0.8B-BF16.gguf", 0, dctx); if (!have_drafter) std::printf("[niah] drafter unavailable, skipping drafter policy\n"); KvFlashDrafterScorer scorer(&dctx); if (have_drafter) { @@ -1198,7 +1198,7 @@ int main(int argc, char ** argv) { // reselect() repages the pool. PASS requires at least one genuine // drafter-driven recall of a chunk evicted earlier. { - const char * drafter_path = "/opt/lucebox/models/drafter/Qwen3-0.6B-BF16.gguf"; + const char * drafter_path = "/opt/lucebox/models/drafter/Qwen3.5-0.8B-BF16.gguf"; DrafterContext dctx; if (!load_drafter(drafter_path, 0, dctx)) { std::printf("FAIL indexer run: drafter load failed (%s)\n", dflash27b_last_error()); diff --git a/server/test/test_pflash_drafter_ipc.cpp b/server/test/test_pflash_drafter_ipc.cpp index e1d45518d..8ff5959bd 100644 --- a/server/test/test_pflash_drafter_ipc.cpp +++ b/server/test/test_pflash_drafter_ipc.cpp @@ -2,7 +2,7 @@ #include "common/pflash_drafter_ipc.h" #include "common/model_backend.h" -#include "qwen3/pflash_selection.h" +#include "pflash/pflash_selection.h" #include #include @@ -81,20 +81,20 @@ TEST_CASE(PFlashDrafterIpcFixture, compress3_matches_the_local_selector_contract REQUIRE(remote.required_instruction_spans == local.required_instruction_spans); const auto select = [&] (const std::vector & spans) { - std::vector candidates; + std::vector candidates; constexpr double scores[]{0.0, 9.0, 1.0, 0.0, 2.0, 3.0, 4.0, 0.0}; for (int chunk = 0; chunk < 8; ++chunk) { const int begin = chunk * 4; const int end = begin + 4; candidates.push_back({ (size_t) chunk, begin, end, scores[chunk], - dflash::qwen3::pflash_chunk_is_structurally_required( + dflash::pflash::pflash_chunk_is_structurally_required( begin, end, 28, 32, 32, spans), }); } - return dflash::qwen3::select_pflash_candidates( + return dflash::pflash::select_pflash_candidates( candidates, {16, 0.95}, - dflash::qwen3::PFlashSelectionMode::BudgetOnly); + dflash::pflash::PFlashSelectionMode::BudgetOnly); }; const auto local_result = select(local.required_instruction_spans); const auto remote_result = select(remote.required_instruction_spans); @@ -107,13 +107,13 @@ TEST_CASE(PFlashDrafterIpcFixture, compress3_matches_the_local_selector_contract const std::vector out_of_range{{0, 33}}; std::string local_error; std::string remote_error; - REQUIRE(!dflash::qwen3::validate_pflash_instruction_spans( + REQUIRE(!dflash::pflash::validate_pflash_instruction_spans( out_of_range, (int) local.input_ids.size(), local_error)); REQUIRE(format_pflash_drafter_ipc_compress_command( local.keep_ratio, local.score_query_end, local.score_query_tokens, out_of_range, "/tmp/ids.bin", line, error)); REQUIRE(parse_pflash_drafter_ipc_compress_command(line, remote, error)); - REQUIRE(!dflash::qwen3::validate_pflash_instruction_spans( + REQUIRE(!dflash::pflash::validate_pflash_instruction_spans( remote.required_instruction_spans, (int) local.input_ids.size(), remote_error)); REQUIRE(local_error == remote_error); diff --git a/server/test/test_pflash_selection.cpp b/server/test/test_pflash_selection.cpp index 5539bc640..d56fe8a90 100644 --- a/server/test/test_pflash_selection.cpp +++ b/server/test/test_pflash_selection.cpp @@ -1,7 +1,7 @@ #include "CppUnitTestFramework.hpp" -#include "qwen3/pflash_selection.h" -#include "qwen3/qwen3_drafter_model.h" +#include "pflash/pflash_selection.h" +#include "pflash/pflash_compress.h" #include "scoped_env.h" #include @@ -12,7 +12,7 @@ #include #include -using namespace dflash::qwen3; +using namespace dflash::pflash; namespace { diff --git a/server/test/test_qwen3_buffer_plan.cpp b/server/test/test_qwen3_buffer_plan.cpp deleted file mode 100644 index d26809875..000000000 --- a/server/test/test_qwen3_buffer_plan.cpp +++ /dev/null @@ -1,71 +0,0 @@ -#include "CppUnitTestFramework.hpp" - -#include "qwen3/qwen3_buffer_plan.h" - -#include -#include - -using dflash::common::qwen3_drafter_buffer_plan; - -namespace { -struct Qwen3BufferPlanFixture : CppUnitTestFramework::CommonFixture { - using CppUnitTestFramework::CommonFixture::CommonFixture; - - void nope_tail_reuses_current_layer_kv() { - const auto plan = qwen3_drafter_buffer_plan(true, 28); - REQUIRE(plan.rope_k_buffers == (size_t)1); - REQUIRE(plan.value_buffers == (size_t)1); - REQUIRE(plan.rope_q_tail_buffers == (size_t)0); - REQUIRE(plan.layer_cache_index(0) == (size_t)0); - REQUIRE(plan.layer_cache_index(1) == (size_t)0); - REQUIRE(plan.layer_cache_index(27) == (size_t)0); - } - - void legacy_rope_scoring_retains_per_layer_state() { - const auto plan = qwen3_drafter_buffer_plan(false, 28); - REQUIRE(plan.rope_k_buffers == (size_t)28); - REQUIRE(plan.value_buffers == (size_t)1); - REQUIRE(plan.rope_q_tail_buffers == (size_t)28); - REQUIRE(plan.layer_cache_index(0) == (size_t)0); - REQUIRE(plan.layer_cache_index(1) == (size_t)1); - REQUIRE(plan.layer_cache_index(27) == (size_t)27); - } - - void empty_model_has_no_layer_buffers() { - const auto plan = qwen3_drafter_buffer_plan(true, 0); - REQUIRE(plan.rope_k_buffers == (size_t)0); - REQUIRE(plan.value_buffers == (size_t)0); - REQUIRE(plan.rope_q_tail_buffers == (size_t)0); - } - - void single_layer_mapping_is_in_bounds() { - const auto nope_plan = qwen3_drafter_buffer_plan(true, 1); - const auto legacy_plan = qwen3_drafter_buffer_plan(false, 1); - REQUIRE(nope_plan.layer_cache_index(0) == (size_t)0); - REQUIRE(legacy_plan.layer_cache_index(0) == (size_t)0); - } - - void nope_tail_removes_reported_per_layer_allocation_growth() { - constexpr size_t heads_kv = 8; - constexpr size_t head_dim = 128; - constexpr size_t bf16_bytes = 2; - const auto bytes_per_kv = [](size_t seq_len) { - return seq_len * heads_kv * head_dim * bf16_bytes; - }; - - REQUIRE(bytes_per_kv(179262) == (size_t)367128576); - REQUIRE(bytes_per_kv(199530) == (size_t)408637440); - - const auto plan = qwen3_drafter_buffer_plan(true, 28); - REQUIRE(plan.rope_k_buffers + plan.value_buffers == (size_t)2); - } -}; -} - -TEST_CASE(Qwen3BufferPlanFixture, allocation_policy) { - nope_tail_reuses_current_layer_kv(); - legacy_rope_scoring_retains_per_layer_state(); - empty_model_has_no_layer_buffers(); - single_layer_mapping_is_in_bounds(); - nope_tail_removes_reported_per_layer_allocation_growth(); -} diff --git a/server/test/test_server_unit.cpp b/server/test/test_server_unit.cpp index 5743f4712..a5342b6a9 100644 --- a/server/test/test_server_unit.cpp +++ b/server/test/test_server_unit.cpp @@ -45,8 +45,9 @@ #include "qwen35moe/qwen35moe_ffn.h" #include "ggml-cpu.h" #include "server/prompt_normalize.h" -#include "qwen3_drafter.h" -#include "qwen3_drafter_model.h" +#include "pflash/pflash_drafter.h" +#include "qwen3_model.h" +#include "pflash/pflash_compress.h" #include "dflash27b.h" #include "gguf.h" #include @@ -855,7 +856,6 @@ TEST_CASE(ServerUnitFixture, test_pflash_score_validation_counts_nan_and_inf) { TEST_CASE(ServerUnitFixture, test_qwen35_pflash_rejects_missing_query_window) { DrafterContext ctx; ctx.loaded = true; - ctx.arch = DrafterArch::Qwen35_0p8b; const std::vector ids(16, 1); const auto compressed = drafter_score_and_compress( @@ -8672,20 +8672,20 @@ TEST_CASE(ServerUnitFixture, test_flowkv_session_keep_ratio_override) { } // ═══════════════════════════════════════════════════════════════════════ -// Qwen3-0.6B drafter loader: truncated GGUF guard (bug #438) +// Qwen3-0.6B model loader: truncated GGUF guard (bug #438) // ═══════════════════════════════════════════════════════════════════════ // // Builds a minimal but structurally valid Qwen3-0.6B-style GGUF on disk, then -// verifies that load_qwen3_drafter_model: +// verifies that load_qwen3_model: // (1) loads the full, untruncated file successfully (positive control), and // (2) fails cleanly with a "truncated or corrupt" error when the tensor-data // section is truncated — instead of letting the H2D copy read past the // end of the mmap and SIGSEGV inside the device copy. -// Write a tiny valid drafter GGUF and return its path. The loader fixes -// n_vocab at 151936 (Qwen3DrafterWeights default), so token_embd stays the +// Write a tiny valid model GGUF and return its path. The loader fixes +// n_vocab at 151936 (Qwen3Weights default), so token_embd stays the // largest tensor (~2.4 MB BF16) while every other tensor is minimal. -static std::string write_qwen3_drafter_fixture_gguf() { +static std::string write_qwen3_model_fixture_gguf() { const int n_embd = 8; const int n_head = 2; const int head_dim = 4; @@ -8741,7 +8741,7 @@ static std::string write_qwen3_drafter_fixture_gguf() { add_tensor("blk.0.ffn_down.weight", GGML_TYPE_BF16, 2, n_ff, n_embd); const std::string path = test_tmp_path( - "dflash_test_qwen3_drafter_438.gguf").string(); + "dflash_test_qwen3_model_438.gguf").string(); gguf_write_to_file(g, path.c_str(), /*only_meta=*/false); gguf_free(g); @@ -8749,18 +8749,18 @@ static std::string write_qwen3_drafter_fixture_gguf() { return path; } -TEST_CASE(ServerUnitFixture, test_qwen3_drafter_rejects_truncated_gguf) { - const std::string path = write_qwen3_drafter_fixture_gguf(); +TEST_CASE(ServerUnitFixture, test_qwen3_model_rejects_truncated_gguf) { + const std::string path = write_qwen3_model_fixture_gguf(); ggml_backend_t backend = ggml_backend_cpu_init(); TEST_ASSERT(backend != nullptr); // Positive control: the full, untruncated file loads cleanly. { - Qwen3DrafterWeights w; - bool ok = load_qwen3_drafter_model(path, backend, w); + Qwen3Weights w; + bool ok = load_qwen3_model(path, backend, w); TEST_ASSERT_MSG(ok, dflash27b_last_error()); - free_qwen3_drafter_model(w); + free_qwen3_model(w); } // Truncate inside the tensor-data section. The header, kv block, and tensor @@ -8776,13 +8776,13 @@ TEST_CASE(ServerUnitFixture, test_qwen3_drafter_rejects_truncated_gguf) { // The loader must fail cleanly (no SIGSEGV) with a descriptive error. { - Qwen3DrafterWeights w; - bool ok = load_qwen3_drafter_model(path, backend, w); + Qwen3Weights w; + bool ok = load_qwen3_model(path, backend, w); TEST_ASSERT(!ok); const std::string err = dflash27b_last_error(); TEST_ASSERT_MSG(err.find("truncated or corrupt") != std::string::npos, err.c_str()); - free_qwen3_drafter_model(w); + free_qwen3_model(w); } ggml_backend_free(backend); From ed2fd89706685ac0708fa96206f139764517b5c2 Mon Sep 17 00:00:00 2001 From: Graffioh Date: Sun, 20 Sep 2026 10:48:26 +0000 Subject: [PATCH 05/26] docs(pflash): clarify pflash vs pflash_drafter boundary MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit PFlash is the whole compression concept; pflash_drafter is the scorer model behind it — Qwen3.5-0.8B for now, not permanently. Make the seam explicit in the file headers instead of implying the coupling is fixed. Generated with [Devin](https://devin.ai) Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- server/src/pflash/kvflash_drafter_scorer.h | 4 ++-- server/src/pflash/pflash_drafter.cpp | 6 ++++-- server/src/pflash/pflash_drafter.h | 18 ++++++++++++------ server/src/pflash/qwen35_drafter.h | 2 +- 4 files changed, 19 insertions(+), 11 deletions(-) diff --git a/server/src/pflash/kvflash_drafter_scorer.h b/server/src/pflash/kvflash_drafter_scorer.h index 2451716b9..3883efbda 100644 --- a/server/src/pflash/kvflash_drafter_scorer.h +++ b/server/src/pflash/kvflash_drafter_scorer.h @@ -1,8 +1,8 @@ // KvFlashDrafterScorer — pflash drafter as the KV pager's Memory Indexer. // // Scores 64-token chunks with the same tail-attention scoring that pflash -// compression uses (the Qwen3.5-0.8B drafter), but returns the per-chunk -// relevance scores instead of a compressed token list. The DrafterContext +// compression uses (the pflash drafter — Qwen3.5-0.8B for now), but returns +// the per-chunk relevance scores instead of a compressed token list. The DrafterContext // is borrowed: the daemon shares its pflash drafter; the pager itself never // depends on this file (see common/kvflash_scorer.h). diff --git a/server/src/pflash/pflash_drafter.cpp b/server/src/pflash/pflash_drafter.cpp index 484ede122..51a074c21 100644 --- a/server/src/pflash/pflash_drafter.cpp +++ b/server/src/pflash/pflash_drafter.cpp @@ -1,5 +1,7 @@ -// PFlash drafter entry points: load/free the Qwen3.5-0.8B scorer and run -// drafter_score_and_compress. +// PFlash drafter entry points: load/free the scorer and run +// drafter_score_and_compress. The pflash drafter is Qwen3.5-0.8B for now — +// this file is the dispatch seam where a different drafter model would +// slot in. // // Wires three pieces: // - qwen35_loader.cpp : mmap GGUF + populate ggml tensors on backend, diff --git a/server/src/pflash/pflash_drafter.h b/server/src/pflash/pflash_drafter.h index 7749e437e..0ab3a6235 100644 --- a/server/src/pflash/pflash_drafter.h +++ b/server/src/pflash/pflash_drafter.h @@ -1,10 +1,13 @@ // In-process PFlash drafter for speculative prefill. // -// The drafter is the Qwen3.5-0.8B scorer (qwen35_drafter.cpp + -// qwen35_loader.cpp): it runs the model's first fifteen blocks and scores -// the context with block 15's NoPE Q/K attention-mass head, with the -// all-layer running-max scorer kept as an opt-in alternative -// (PFLASH_QWEN35_LEGACY_SCORER=1 or the PFLASH scorer config). +// "PFlash" is the whole compression concept (score -> select -> emit); the +// pflash drafter is the scorer model behind it — Qwen3.5-0.8B for now +// (qwen35_drafter.cpp + qwen35_loader.cpp): it runs the model's first +// fifteen blocks and scores the context with block 15's NoPE Q/K +// attention-mass head, with the all-layer running-max scorer kept as an +// opt-in alternative (PFLASH_QWEN35_LEGACY_SCORER=1 or the PFLASH scorer +// config). This header is the model-agnostic API surface; a future drafter +// slots in behind load_drafter / drafter_score_and_compress. // // Hosted in the SAME process / SAME ggml allocator as the dflash target, so // we never pay the cross-process VRAM contention that broke the Python @@ -32,12 +35,15 @@ struct Qwen35DrafterState; struct DrafterContext { ggml_backend_t backend = nullptr; // owned (created in load_drafter) + // Scorer state for the current drafter (Qwen3.5-0.8B). The public API + // below never exposes it; backends that need internals include + // qwen35_drafter.h explicitly. Qwen35DrafterState * state = nullptr; // owned scorer weights + heads int gpu = -1; bool loaded = false; }; -// Load the drafter GGUF (a Qwen3.5-0.8B GGUF). +// Load the drafter GGUF (a Qwen3.5-0.8B GGUF today). // Creates a fresh GPU backend if `backend` is null. Otherwise uses the // caller-provided backend (so the drafter shares the daemon's allocator). // diff --git a/server/src/pflash/qwen35_drafter.h b/server/src/pflash/qwen35_drafter.h index 434353612..c68d46e3f 100644 --- a/server/src/pflash/qwen35_drafter.h +++ b/server/src/pflash/qwen35_drafter.h @@ -1,4 +1,4 @@ -// Internal interface of the Qwen3.5-0.8B drafter — the only pflash scorer. +// Internal interface of the Qwen3.5-0.8B drafter — the current pflash scorer. // // The scorer runs on the Qwen3.5 target architecture (TargetWeights, // build_qwen35_layer): qwen35_loader.cpp loads the GGUF, the optional From 1ee9cc16ff8f8ef122cfe2ab9b9beab576a1cc3e Mon Sep 17 00:00:00 2001 From: Graffioh Date: Tue, 22 Sep 2026 06:18:02 +0000 Subject: [PATCH 06/26] feat(pflash): add a rank-based top_k selection mode The budget selector fills a fixed fraction of the prompt by descending segment score, but on compact-evidence traffic the evidence sits in the first few ranks, so a rank rule reaches it for far fewer tokens. TopK keeps the mandatory candidates and then the K highest-scoring optional ones in score order, with the token budget still a hard ceiling and oversized candidates skipped exactly as budget_only does, and stops with top_k_reached when K binds before the budget. PFLASH_SELECT_TOPK is required and validated in that mode, the config line and the compression trace both carry the K that applied, and split selection rejects the mode rather than quietly keeping up to 2K segments. Co-Authored-By: Claude Opus 5 (1M context) --- server/README.md | 1 + server/src/pflash/pflash_compress.cpp | 10 +- server/src/pflash/pflash_compress.h | 2 + server/src/pflash/pflash_drafter.cpp | 5 +- server/src/pflash/pflash_selection.cpp | 42 +++++- server/src/pflash/pflash_selection.h | 8 ++ server/test/test_pflash_selection.cpp | 185 +++++++++++++++++++++++++ 7 files changed, 245 insertions(+), 8 deletions(-) diff --git a/server/README.md b/server/README.md index 7e11d8780..e2a032642 100644 --- a/server/README.md +++ b/server/README.md @@ -376,6 +376,7 @@ the whole request's device footprint. `/status/json` reports | `--prefill-upstream-base ` | none | Enable compression-proxy mode. | | `--prefill-upstream-key ` | none | Bearer token for the upstream. | | `--prefill-upstream-model ` | none | Model name forwarded upstream. | +| `PFLASH_SELECT_MODE=top_k` + `PFLASH_SELECT_TOPK ` | budget-only fill | Rank rule: keep the K highest-scoring optional segments in score order instead of filling the keep ratio, with the keep-ratio budget still a hard ceiling (min(K segments, the budget)). Use it where the evidence is compact and sits in the first few ranks -- needle retrieval, passage QA, code -- so a small K reaches it for a fraction of the budget's tokens. Do not use it where the answer needs a whole document identified, since the evidence there spans many segments and K cuts it off. | With a Qwen3.5-0.8B drafter and strict budget selection (`PFLASH_SELECT_MODE=budget_only`, `PFLASH_SELECT_CHUNK_SIZE`, diff --git a/server/src/pflash/pflash_compress.cpp b/server/src/pflash/pflash_compress.cpp index 07b7410db..1a14d061d 100644 --- a/server/src/pflash/pflash_compress.cpp +++ b/server/src/pflash/pflash_compress.cpp @@ -87,13 +87,14 @@ void write_compression_trace( std::fprintf(file, "],\"query_begin\":%d,\"query_end\":%d," "\"selector_mode\":\"%s\",\"query_parser\":\"%s\"," - "\"token_budget\":%d," + "\"token_budget\":%d,\"top_k\":%d," "\"retained_tokens\":%d", trace_fields->query_begin, trace_fields->query_end, dflash::pflash::pflash_selection_mode_name( trace_fields->selector_mode), dflash::pflash::pflash_query_parser_name(trace_fields->query_parser), - trace_fields->token_budget, trace_fields->retained_tokens); + trace_fields->token_budget, trace_fields->top_k, + trace_fields->retained_tokens); std::fputs(",\"required_instruction_spans\":[", file); if (trace_fields->required_instruction_spans) { for (size_t index = 0; @@ -247,7 +248,8 @@ std::vector select_pflash_chunks( } const dflash::pflash::PFlashSelectionPolicy policy{selector_budget, config.top_p, - /*skip_oversized=*/ segments != nullptr}; + /*skip_oversized=*/ segments != nullptr, + config.top_k}; const auto selected = split ? dflash::pflash::select_pflash_split(candidates, other_candidates, policy, split_fraction, config.mode) : dflash::pflash::select_pflash_candidates(candidates, policy, config.mode); @@ -313,6 +315,8 @@ std::vector select_pflash_chunks( strict_fields.scorer = split ? "split" : dflash::pflash::pflash_scorer_name(config.scorer); strict_fields.split_fraction = split ? split_fraction : 0.0; strict_fields.other_chunk_scores = split ? &other_scores : nullptr; + strict_fields.top_k = + config.mode == dflash::pflash::PFlashSelectionMode::TopK ? config.top_k : 0; write_compression_trace( input_tokens, keep_ratio, trace_chunk, query_tokens, pool_kernel, n_keep_approx, chunk_means, selected_mask, diff --git a/server/src/pflash/pflash_compress.h b/server/src/pflash/pflash_compress.h index d97fcc024..4d17a3b39 100644 --- a/server/src/pflash/pflash_compress.h +++ b/server/src/pflash/pflash_compress.h @@ -103,6 +103,8 @@ struct PFlashTraceFields { const char * scorer = "head"; double split_fraction = 0.0; const std::vector * other_chunk_scores = nullptr; + // Rank-mode ceiling: the K that applied, 0 outside top_k mode. + int top_k = 0; }; void write_compression_trace( diff --git a/server/src/pflash/pflash_drafter.cpp b/server/src/pflash/pflash_drafter.cpp index 51a074c21..2a65f5128 100644 --- a/server/src/pflash/pflash_drafter.cpp +++ b/server/src/pflash/pflash_drafter.cpp @@ -145,11 +145,12 @@ std::vector drafter_score_and_compress( std::fprintf(stderr, "[pflash-select] config mode=%s active=%d chunk=%d " "query_parser=%s query_cap=%d query_actual=%d top_p=%.9g " - "input=%zu\n", + "top_k=%d input=%zu\n", dflash::pflash::pflash_selection_mode_name(experiment.mode), (int) experiment.selection_active, experiment.chunk_size, dflash::pflash::pflash_query_parser_name(experiment.query_parser), - experiment.query_tokens, n_lookahead, experiment.top_p, ids.size()); + experiment.query_tokens, n_lookahead, experiment.top_p, + experiment.top_k, ids.size()); std::fflush(stderr); } if (score_query_end < 0) { diff --git a/server/src/pflash/pflash_selection.cpp b/server/src/pflash/pflash_selection.cpp index f6d3fe6bf..a331d3e52 100644 --- a/server/src/pflash/pflash_selection.cpp +++ b/server/src/pflash/pflash_selection.cpp @@ -20,6 +20,7 @@ constexpr const char * kChunkEnv = "PFLASH_SELECT_CHUNK_SIZE"; constexpr const char * kQueryEnv = "PFLASH_SELECT_QUERY_TOKENS"; constexpr const char * kQueryParserEnv = "PFLASH_SELECT_QUERY_PARSER"; constexpr const char * kTopPEnv = "PFLASH_SELECT_TOP_P"; +constexpr const char * kTopKEnv = "PFLASH_SELECT_TOPK"; constexpr const char * kSegmentsEnv = "PFLASH_SELECT_SEGMENTS"; constexpr const char * kSelectEnv = "PFLASH_SELECT_SCORE"; constexpr const char * kScorerEnv = "PFLASH_SELECT_SCORER"; @@ -71,6 +72,7 @@ bool has_pflash_selection_environment() noexcept { std::getenv(kQueryEnv) != nullptr || std::getenv(kQueryParserEnv) != nullptr || std::getenv(kTopPEnv) != nullptr || + std::getenv(kTopKEnv) != nullptr || std::getenv(kSegmentsEnv) != nullptr || std::getenv(kSelectEnv) != nullptr || std::getenv(kScorerEnv) != nullptr || @@ -142,6 +144,9 @@ PFlashSelectionResult select_pflash_candidates( if (!std::isfinite(policy.top_p) || policy.top_p <= 0.0 || policy.top_p > 1.0) { return invalid_result("PFlash top_p must be finite and in (0, 1]"); } + if (mode == PFlashSelectionMode::TopK && policy.top_k <= 0) { + return invalid_result("PFlash top_k must be positive"); + } std::vector source_ranges; source_ranges.reserve(candidates.size()); @@ -216,12 +221,19 @@ PFlashSelectionResult select_pflash_candidates( } } + int kept_optional = 0; for (const auto * candidate : optional) { if (mode == PFlashSelectionMode::CumulativeTopP && result.retained_mass >= policy.top_p) { result.stop = PFlashSelectionStop::TopPReached; break; } + // Rank rule: K optional candidates in score order, the budget below + // still a ceiling. K binding here means the budget never was. + if (mode == PFlashSelectionMode::TopK && kept_optional >= policy.top_k) { + result.stop = PFlashSelectionStop::TopKReached; + break; + } const int length = candidate->end - candidate->begin; if (length > policy.token_budget - result.retained_tokens) { @@ -232,6 +244,7 @@ PFlashSelectionResult select_pflash_candidates( selected_candidates.push_back(candidate); result.retained_tokens += length; + ++kept_optional; if (!optional.empty()) { result.retained_mass += max_score > 0.0 ? (std::max(0.0, candidate->score) / max_score) / scaled_total @@ -256,6 +269,7 @@ const char * pflash_selection_mode_name(PFlashSelectionMode mode) noexcept { case PFlashSelectionMode::Legacy: return "legacy"; case PFlashSelectionMode::BudgetOnly: return "budget_only"; case PFlashSelectionMode::CumulativeTopP: return "top_p"; + case PFlashSelectionMode::TopK: return "top_k"; } return "unknown"; } @@ -263,6 +277,7 @@ const char * pflash_selection_mode_name(PFlashSelectionMode mode) noexcept { const char * pflash_selection_stop_name(PFlashSelectionStop stop) noexcept { switch (stop) { case PFlashSelectionStop::TopPReached: return "top_p_reached"; + case PFlashSelectionStop::TopKReached: return "top_k_reached"; case PFlashSelectionStop::BudgetReached: return "budget_reached"; case PFlashSelectionStop::CandidatesExhausted: return "candidates_exhausted"; case PFlashSelectionStop::InvalidInput: return "invalid_input"; @@ -300,6 +315,7 @@ bool resolve_pflash_selection( const char * query_raw = std::getenv(kQueryEnv); const char * query_parser_raw = std::getenv(kQueryParserEnv); const char * top_p_raw = std::getenv(kTopPEnv); + const char * top_k_raw = std::getenv(kTopKEnv); const char * segments_raw = std::getenv(kSegmentsEnv); const char * select_raw = std::getenv(kSelectEnv); const char * scorer_raw = std::getenv(kScorerEnv); @@ -307,8 +323,8 @@ bool resolve_pflash_selection( PFlashSelectionConfig config; config.configured = mode_raw || chunk_raw || query_raw || - query_parser_raw || top_p_raw || segments_raw || select_raw || - scorer_raw || split_raw; + query_parser_raw || top_p_raw || top_k_raw || segments_raw || + select_raw || scorer_raw || split_raw; if (scorer_raw) { if (std::strcmp(scorer_raw, "head") == 0) { config.scorer = PFlashScorer::Head; @@ -358,8 +374,11 @@ bool resolve_pflash_selection( config.mode = PFlashSelectionMode::BudgetOnly; } else if (std::strcmp(mode_raw, "top_p") == 0) { config.mode = PFlashSelectionMode::CumulativeTopP; + } else if (std::strcmp(mode_raw, "top_k") == 0) { + config.mode = PFlashSelectionMode::TopK; } else { - error = std::string(kModeEnv) + " must be budget_only or top_p"; + error = std::string(kModeEnv) + + " must be budget_only, top_p or top_k"; return false; } } @@ -400,6 +419,16 @@ bool resolve_pflash_selection( return false; } + if (top_k_raw && (!parse_int(top_k_raw, config.top_k) || config.top_k <= 0)) { + error = std::string(kTopKEnv) + " must be a positive integer"; + return false; + } + if (config.mode == PFlashSelectionMode::TopK && config.top_k <= 0) { + error = std::string(kTopKEnv) + " is required when " + + std::string(kModeEnv) + " is top_k"; + return false; + } + out = config; return true; } @@ -472,6 +501,13 @@ PFlashSelectionResult select_pflash_split( double head_fraction, PFlashSelectionMode mode) { PFlashSelectionResult result; + if (mode == PFlashSelectionMode::TopK) { + // A per-pass K would keep up to 2K segments, which is not the rule the + // mode names; fail closed rather than quietly double it. + result.stop = PFlashSelectionStop::InvalidInput; + result.error = "split selection does not support top_k"; + return result; + } if (head.size() != other.size() || !(head_fraction > 0.0 && head_fraction < 1.0)) { result.stop = PFlashSelectionStop::InvalidInput; result.error = "split selection needs matching candidate lists and a fraction in (0, 1)"; diff --git a/server/src/pflash/pflash_selection.h b/server/src/pflash/pflash_selection.h index a88e4c6f0..7df7df7f6 100644 --- a/server/src/pflash/pflash_selection.h +++ b/server/src/pflash/pflash_selection.h @@ -12,6 +12,9 @@ enum class PFlashSelectionMode { Legacy, BudgetOnly, CumulativeTopP, + // Rank rule: keep the K highest-scoring optional candidates, the token + // budget still a hard ceiling -- min(K segments, the budget). + TopK, }; enum class PFlashQueryParser { @@ -21,6 +24,7 @@ enum class PFlashQueryParser { enum class PFlashSelectionStop { TopPReached, + TopKReached, BudgetReached, CandidatesExhausted, InvalidInput, @@ -42,6 +46,9 @@ struct PFlashSelectionPolicy { // budget is skipped instead of ending the fill, so smaller segments // ranked below it can still be kept. bool skip_oversized = false; + // TopK mode only: how many optional candidates to keep. Must be positive + // in that mode and is ignored in the others. + int top_k = 0; }; struct PFlashSelectionResult { @@ -92,6 +99,7 @@ struct PFlashSelectionConfig { int chunk_size = 0; int query_tokens = 8; double top_p = 0.95; + int top_k = 0; PFlashSegmentation segmentation = PFlashSegmentation::Auto; PFlashCandidateScore candidate_score = PFlashCandidateScore::Auto; PFlashScorer scorer = PFlashScorer::Head; diff --git a/server/test/test_pflash_selection.cpp b/server/test/test_pflash_selection.cpp index d56fe8a90..2dfb7c19f 100644 --- a/server/test/test_pflash_selection.cpp +++ b/server/test/test_pflash_selection.cpp @@ -21,6 +21,7 @@ constexpr const char * kChunkEnv = "PFLASH_SELECT_CHUNK_SIZE"; constexpr const char * kQueryEnv = "PFLASH_SELECT_QUERY_TOKENS"; constexpr const char * kQueryParserEnv = "PFLASH_SELECT_QUERY_PARSER"; constexpr const char * kTopPEnv = "PFLASH_SELECT_TOP_P"; +constexpr const char * kTopKEnv = "PFLASH_SELECT_TOPK"; struct CleanPFlashEnv { luce_test::ScopedEnvVar mode{kModeEnv, nullptr}; @@ -28,6 +29,7 @@ struct CleanPFlashEnv { luce_test::ScopedEnvVar query{kQueryEnv, nullptr}; luce_test::ScopedEnvVar query_parser{kQueryParserEnv, nullptr}; luce_test::ScopedEnvVar top_p{kTopPEnv, nullptr}; + luce_test::ScopedEnvVar top_k{kTopKEnv, nullptr}; }; void set_env(const char * name, const char * value) { @@ -673,3 +675,186 @@ TEST_CASE(PFlashSelectionFixture, scorer_and_split_environment_resolve_or_fail) REQUIRE(!resolve_pflash_selection(4096, 1024, invalid, error)); REQUIRE(error.find("PFLASH_SELECT_SPLIT") != std::string::npos); } + +namespace { + +// Six equal-length optional candidates, scores descending with the ordinal. +std::vector ranked_candidates() { + std::vector candidates; + for (size_t index = 0; index < 6; ++index) { + const int begin = (int) index * 100; + candidates.push_back( + candidate(index, begin, begin + 100, 6.0 - (double) index)); + } + return candidates; +} + +} // namespace + +TEST_CASE(PFlashSelectionFixture, top_k_keeps_the_k_highest_scoring_optional_candidates) { + const auto candidates = ranked_candidates(); + // Budget far above the six candidates: K alone decides. + const auto result = select_pflash_candidates( + candidates, PFlashSelectionPolicy{100000, 0.95, false, 2}, + PFlashSelectionMode::TopK); + + REQUIRE(result.ok); + REQUIRE(result.stop == PFlashSelectionStop::TopKReached); + REQUIRE(result.retained_tokens == 200); + require_ordinals(result, {0, 1}); +} + +TEST_CASE(PFlashSelectionFixture, top_k_above_the_candidate_count_falls_back_to_budget_behaviour) { + const auto candidates = ranked_candidates(); + // K larger than the candidate list: every candidate fits, so the run ends + // exactly where budget_only would. + const auto roomy = select_pflash_candidates( + candidates, PFlashSelectionPolicy{100000, 0.95, false, 50}, + PFlashSelectionMode::TopK); + const auto budget_only = select_pflash_candidates( + candidates, PFlashSelectionPolicy{100000, 0.95, false}, + PFlashSelectionMode::BudgetOnly); + + REQUIRE(roomy.ok); + REQUIRE(roomy.stop == PFlashSelectionStop::CandidatesExhausted); + REQUIRE(roomy.stop == budget_only.stop); + REQUIRE(roomy.retained_tokens == budget_only.retained_tokens); + require_ordinals(roomy, {0, 1, 2, 3, 4, 5}); + + // The same K against a budget that bites: the budget stops the fill. + const auto tight = select_pflash_candidates( + candidates, PFlashSelectionPolicy{250, 0.95, false, 50}, + PFlashSelectionMode::TopK); + REQUIRE(tight.ok); + REQUIRE(tight.stop == PFlashSelectionStop::BudgetReached); + REQUIRE(tight.retained_tokens == 200); + require_ordinals(tight, {0, 1}); +} + +TEST_CASE(PFlashSelectionFixture, top_k_never_exceeds_the_token_budget) { + const auto candidates = ranked_candidates(); + // K=5 wants 500 tokens; the budget caps the run at two candidates. + const auto result = select_pflash_candidates( + candidates, PFlashSelectionPolicy{250, 0.95, false, 5}, + PFlashSelectionMode::TopK); + + REQUIRE(result.ok); + REQUIRE(result.stop == PFlashSelectionStop::BudgetReached); + REQUIRE(result.retained_tokens <= 250); + REQUIRE(result.retained_tokens == 200); + require_ordinals(result, {0, 1}); +} + +TEST_CASE(PFlashSelectionFixture, top_k_skips_an_oversized_candidate_and_keeps_ranking) { + // Variable-length segments: the 300-token second-ranked candidate does not + // fit, so the fill continues below it instead of ending, exactly as + // budget_only does, and K counts only the candidates actually kept. + const std::vector candidates{ + candidate(0, 0, 100, 10.0), + candidate(1, 100, 400, 9.0), + candidate(2, 400, 500, 8.0), + candidate(3, 500, 600, 7.0), + }; + const auto result = select_pflash_candidates( + candidates, PFlashSelectionPolicy{250, 0.95, true, 3}, + PFlashSelectionMode::TopK); + + REQUIRE(result.ok); + REQUIRE(result.stop == PFlashSelectionStop::BudgetReached); + REQUIRE(result.retained_tokens == 200); + require_ordinals(result, {0, 2}); +} + +TEST_CASE(PFlashSelectionFixture, top_k_keeps_mandatory_candidates_outside_the_rank) { + // Ordinal 3 is mandatory and scores lowest: it is kept and charged, and K + // still buys one optional candidate on top of it. + std::vector candidates{ + candidate(0, 0, 100, 5.0), + candidate(1, 100, 200, 4.0), + candidate(2, 200, 300, 3.0), + candidate(3, 300, 400, 0.0, true), + }; + const auto result = select_pflash_candidates( + candidates, PFlashSelectionPolicy{100000, 0.95, false, 1}, + PFlashSelectionMode::TopK); + + REQUIRE(result.ok); + REQUIRE(result.stop == PFlashSelectionStop::TopKReached); + REQUIRE(result.retained_tokens == 200); + require_ordinals(result, {0, 3}); + + // A mandatory span that cannot fit still fails closed under top_k. + const auto overflow = select_pflash_candidates( + candidates, PFlashSelectionPolicy{50, 0.95, false, 1}, + PFlashSelectionMode::TopK); + REQUIRE(!overflow.ok); + REQUIRE(overflow.stop == PFlashSelectionStop::MandatoryQueryExceedsBudget); +} + +TEST_CASE(PFlashSelectionFixture, top_k_requires_a_positive_k) { + const auto candidates = ranked_candidates(); + REQUIRE(!select_pflash_candidates( + candidates, PFlashSelectionPolicy{100000, 0.95, false, 0}, + PFlashSelectionMode::TopK).ok); + REQUIRE(!select_pflash_candidates( + candidates, PFlashSelectionPolicy{100000, 0.95, false, -1}, + PFlashSelectionMode::TopK).ok); + // The other modes ignore the field. + REQUIRE(select_pflash_candidates( + candidates, PFlashSelectionPolicy{100000, 0.95, false, 0}, + PFlashSelectionMode::BudgetOnly).ok); +} + +TEST_CASE(PFlashSelectionFixture, top_k_environment_resolves_or_fails) { + CleanPFlashEnv clean; + set_env(kModeEnv, "top_k"); + + // Missing K in top_k mode is rejected by name. + PFlashSelectionConfig invalid; + std::string error; + REQUIRE(!resolve_pflash_selection(32768, 1024, invalid, error)); + REQUIRE(error.find(kTopKEnv) != std::string::npos); + + for (const char * bad : {"0", "-3", "abc", "20.5", ""}) { + set_env(kTopKEnv, bad); + REQUIRE(!resolve_pflash_selection(32768, 1024, invalid, error)); + REQUIRE(error.find(kTopKEnv) != std::string::npos); + } + + set_env(kTopKEnv, "20"); + const auto config = resolve_or_fail(32768, 1024); + REQUIRE(config.configured); + REQUIRE(config.selection_active); + REQUIRE(config.mode == PFlashSelectionMode::TopK); + REQUIRE(config.top_k == 20); + + // K alone, without the mode, parses but leaves the mode untouched. + set_env(kModeEnv, "budget_only"); + const auto budget = resolve_or_fail(32768, 1024); + REQUIRE(budget.mode == PFlashSelectionMode::BudgetOnly); + REQUIRE(budget.top_k == 20); + + // An unknown mode still names the legal set. + set_env(kModeEnv, "top_q"); + REQUIRE(!resolve_pflash_selection(32768, 1024, invalid, error)); + REQUIRE(error.find("top_k") != std::string::npos); + + REQUIRE(std::string(pflash_selection_mode_name(PFlashSelectionMode::TopK)) == "top_k"); + REQUIRE(std::string(pflash_selection_stop_name(PFlashSelectionStop::TopKReached)) == + "top_k_reached"); +} + +TEST_CASE(PFlashSelectionFixture, split_selection_rejects_top_k) { + std::vector head; + std::vector other; + for (size_t i = 0; i < 4; ++i) { + const int begin = (int) i * 100; + head.push_back(candidate(i, begin, begin + 100, 4.0 - (double) i)); + other.push_back(candidate(i, begin, begin + 100, (double) i)); + } + const auto result = select_pflash_split( + head, other, PFlashSelectionPolicy{400, 0.95, false, 2}, 0.5, + PFlashSelectionMode::TopK); + REQUIRE(!result.ok); + REQUIRE(result.stop == PFlashSelectionStop::InvalidInput); +} From 2455efb8c500b74c7c3e1f42e9203b1bf790e38a Mon Sep 17 00:00:00 2001 From: Graffioh Date: Tue, 22 Sep 2026 08:39:36 +0000 Subject: [PATCH 07/26] feat(pflash): rank segments with a document mass prior Per-segment density ranking throws away what the head already knows at document level: on BRIGHT the gold document ranks first by aggregated mass in 16 of 28 prompts while its median gold segment ranks 68 of 215. PFLASH_SELECT_DOC_PRIOR scales an optional candidate's ranking score by its document's share of the prompt's total mass raised to the exponent, so a document the head likes as a whole lifts its own segments. Zero, the default, leaves the ranking untouched, and the prior applies in every mode, so it composes with top_k. Documents come from a new pflash_documents request field or, absent it, from the served prompt's own "Document :" and [DOC-] markers; fewer than three documents makes the prior a no-op, which the trace records with the exponent and the document count. The drafter IPC compress protocol carries no document spans, so the prior is inert on that path until the wire format gains them. Co-Authored-By: Claude Opus 5 (1M context) --- server/README.md | 1 + server/src/common/model_backend.h | 3 + .../src/common/pflash_drafter_ipc_daemon.cpp | 3 + server/src/deepseek4/deepseek4_backend.cpp | 3 +- server/src/pflash/pflash_compress.cpp | 33 +++- server/src/pflash/pflash_compress.h | 12 +- server/src/pflash/pflash_drafter.cpp | 11 +- server/src/pflash/pflash_drafter.h | 5 +- server/src/pflash/pflash_selection.cpp | 58 +++++- server/src/pflash/pflash_selection.h | 21 +++ server/src/pflash/qwen35_drafter.cpp | 32 ++-- server/src/pflash/qwen35_drafter.h | 9 +- server/src/qwen35/qwen35_backend.cpp | 3 +- .../src/qwen35/qwen35_layer_split_adapter.cpp | 3 +- server/src/server/http_server.cpp | 146 +++++++++++++++ server/src/server/http_server.h | 55 ++++++ server/test/test_pflash_selection.cpp | 167 ++++++++++++++++++ 17 files changed, 532 insertions(+), 33 deletions(-) diff --git a/server/README.md b/server/README.md index e2a032642..3f87240d5 100644 --- a/server/README.md +++ b/server/README.md @@ -377,6 +377,7 @@ the whole request's device footprint. `/status/json` reports | `--prefill-upstream-key ` | none | Bearer token for the upstream. | | `--prefill-upstream-model ` | none | Model name forwarded upstream. | | `PFLASH_SELECT_MODE=top_k` + `PFLASH_SELECT_TOPK ` | budget-only fill | Rank rule: keep the K highest-scoring optional segments in score order instead of filling the keep ratio, with the keep-ratio budget still a hard ceiling (min(K segments, the budget)). Use it where the evidence is compact and sits in the first few ranks -- needle retrieval, passage QA, code -- so a small K reaches it for a fraction of the budget's tokens. Do not use it where the answer needs a whole document identified, since the evidence there spans many segments and K cuts it off. | +| `PFLASH_SELECT_DOC_PRIOR ` | `0` (off) | Document prior: rank an optional segment by `max(0, score) * document_mass_share ^ E` instead of its score alone, where a document's mass is the sum over its segments of `max(0, score) * tokens` normalised by the prompt total. The head identifies the relevant document when its mass is aggregated per document, which per-segment density ranking throws away, so this helps where the answer needs a whole document found among many. `0` leaves the ranking unchanged and it applies in every mode, so it composes with `top_k`. Documents come from the request's `pflash_documents` ranges, or from the prompt's own `Document :` / `[DOC-]` markers; fewer than three documents makes it a no-op, which the compression trace records alongside the exponent. | With a Qwen3.5-0.8B drafter and strict budget selection (`PFLASH_SELECT_MODE=budget_only`, `PFLASH_SELECT_CHUNK_SIZE`, diff --git a/server/src/common/model_backend.h b/server/src/common/model_backend.h index 0f98bc9a9..f601376c3 100644 --- a/server/src/common/model_backend.h +++ b/server/src/common/model_backend.h @@ -271,6 +271,9 @@ struct ModelBackend { // Role-derived instruction structure in drafter-token coordinates. // Empty is a valid instruction-free or legacy request. std::vector required_instruction_spans; + // Document starts in drafter-token coordinates, for the document + // prior. Empty is a valid single-document or prior-free request. + std::vector document_spans; std::string drafter_path; // GGUF path (for lazy-load) int drafter_gpu = 0; // backend-local GPU for PFlash drafter bool skip_park = false; // true on >=32GB GPUs diff --git a/server/src/common/pflash_drafter_ipc_daemon.cpp b/server/src/common/pflash_drafter_ipc_daemon.cpp index 972ac4a8b..abfd8e357 100644 --- a/server/src/common/pflash_drafter_ipc_daemon.cpp +++ b/server/src/common/pflash_drafter_ipc_daemon.cpp @@ -70,6 +70,9 @@ int run_pflash_drafter_ipc_daemon(const char * drafter_path, request.score_query_tokens, /*pool_kernel=*/13, score_query_end, request.required_instruction_spans); + // The IPC compress protocol carries no document spans, so the + // document prior (PFLASH_SELECT_DOC_PRIOR) is inert on this path + // until the wire format gains them. if (compressed.empty()) { std::fprintf(stderr, "[pflash-ipc-daemon] compress returned empty\n"); stream_status(stream_fd, -1); diff --git a/server/src/deepseek4/deepseek4_backend.cpp b/server/src/deepseek4/deepseek4_backend.cpp index f0fbcac16..2abe9b077 100644 --- a/server/src/deepseek4/deepseek4_backend.cpp +++ b/server/src/deepseek4/deepseek4_backend.cpp @@ -3181,7 +3181,8 @@ std::vector DeepSeek4Backend::compress_batch( result.compressed_ids = drafter_score_and_compress( pflash_drafter_ctx_, request.input_ids, request.keep_ratio, /*chunk_size=*/32, request.score_query_tokens, /*pool_kernel=*/13, - score_query_end, request.required_instruction_spans); + score_query_end, request.required_instruction_spans, + request.document_spans); result.ok = !result.compressed_ids.empty(); } diff --git a/server/src/pflash/pflash_compress.cpp b/server/src/pflash/pflash_compress.cpp index 1a14d061d..3e2e9c2c8 100644 --- a/server/src/pflash/pflash_compress.cpp +++ b/server/src/pflash/pflash_compress.cpp @@ -88,12 +88,16 @@ void write_compression_trace( "],\"query_begin\":%d,\"query_end\":%d," "\"selector_mode\":\"%s\",\"query_parser\":\"%s\"," "\"token_budget\":%d,\"top_k\":%d," + "\"doc_prior_exponent\":%.9g,\"documents\":%zu," + "\"doc_prior_applied\":%s," "\"retained_tokens\":%d", trace_fields->query_begin, trace_fields->query_end, dflash::pflash::pflash_selection_mode_name( trace_fields->selector_mode), dflash::pflash::pflash_query_parser_name(trace_fields->query_parser), trace_fields->token_budget, trace_fields->top_k, + trace_fields->doc_prior_exponent, trace_fields->documents, + trace_fields->doc_prior_applied ? "true" : "false", trace_fields->retained_tokens); std::fputs(",\"required_instruction_spans\":[", file); if (trace_fields->required_instruction_spans) { @@ -193,7 +197,8 @@ std::vector select_pflash_chunks( const std::vector * segments, bool density, const std::vector * other_token_scores, - double split_fraction) { + double split_fraction, + const std::vector * documents) { const int input_tokens = (int) ids.size(); const int query_end = score_query_end < 0 ? input_tokens : score_query_end; const int query_tokens = std::min(n_lookahead, query_end); @@ -226,7 +231,16 @@ std::vector select_pflash_chunks( dflash::pflash::pflash_chunk_is_structurally_required( begin, end, query_begin, query_end, input_tokens, required_instruction_spans); - candidates.push_back({(size_t) chunk, begin, end, score, mandatory}); + // The last document starting at or before this candidate; everything + // ahead of the first document start belongs to the first document. + size_t document = 0; + if (documents) { + for (size_t index = 0; index < documents->size(); ++index) { + if ((*documents)[index].begin <= begin) document = index; + else break; + } + } + candidates.push_back({(size_t) chunk, begin, end, score, mandatory, document}); chunk_means.push_back({(float) score, chunk}); exact_chunk_scores.push_back(score); } @@ -242,14 +256,16 @@ std::vector select_pflash_chunks( score += (*other_token_scores)[(size_t) token]; } score /= (double) std::max(1, candidate.end - candidate.begin); - other_candidates.push_back({candidate.ordinal, candidate.begin, candidate.end, score, candidate.mandatory}); + other_candidates.push_back({candidate.ordinal, candidate.begin, candidate.end, + score, candidate.mandatory, candidate.document}); other_scores.push_back(score); } } const dflash::pflash::PFlashSelectionPolicy policy{selector_budget, config.top_p, /*skip_oversized=*/ segments != nullptr, - config.top_k}; + config.top_k, + config.doc_prior_exponent}; const auto selected = split ? dflash::pflash::select_pflash_split(candidates, other_candidates, policy, split_fraction, config.mode) : dflash::pflash::select_pflash_candidates(candidates, policy, config.mode); @@ -289,14 +305,16 @@ std::vector select_pflash_chunks( std::fprintf(stderr, "[pflash-select] selected mode=%s scorer=%s segments=%s score=%s chunk=%d query=%d " - "budget=%d selected_tokens=%zu chunks=%zu/%d stop=%s mass=%.9g\n", + "budget=%d selected_tokens=%zu chunks=%zu/%d stop=%s mass=%.9g " + "docs=%zu doc_prior=%.9g applied=%d\n", dflash::pflash::pflash_selection_mode_name(config.mode), split ? "split" : "single", segments ? "probe" : "fixed", density ? "density" : "sum", segments ? 0 : config.chunk_size, query_tokens, selector_budget, output.size(), selected.ordinals.size(), n_chunks, dflash::pflash::pflash_selection_stop_name(selected.stop), - selected.retained_mass); + selected.retained_mass, selected.documents, config.doc_prior_exponent, + (int) selected.doc_prior_applied); std::fflush(stderr); if (write_trace) { @@ -317,6 +335,9 @@ std::vector select_pflash_chunks( strict_fields.other_chunk_scores = split ? &other_scores : nullptr; strict_fields.top_k = config.mode == dflash::pflash::PFlashSelectionMode::TopK ? config.top_k : 0; + strict_fields.doc_prior_exponent = config.doc_prior_exponent; + strict_fields.documents = selected.documents; + strict_fields.doc_prior_applied = selected.doc_prior_applied; write_compression_trace( input_tokens, keep_ratio, trace_chunk, query_tokens, pool_kernel, n_keep_approx, chunk_means, selected_mask, diff --git a/server/src/pflash/pflash_compress.h b/server/src/pflash/pflash_compress.h index 4d17a3b39..240a3580a 100644 --- a/server/src/pflash/pflash_compress.h +++ b/server/src/pflash/pflash_compress.h @@ -105,6 +105,11 @@ struct PFlashTraceFields { const std::vector * other_chunk_scores = nullptr; // Rank-mode ceiling: the K that applied, 0 outside top_k mode. int top_k = 0; + // Document prior: the configured exponent, the documents the selector + // saw, and whether the prior actually reweighted the ranking. + double doc_prior_exponent = 0.0; + size_t documents = 0; + bool doc_prior_applied = false; }; void write_compression_trace( @@ -134,6 +139,11 @@ std::vector select_pflash_chunks( const std::vector * segments = nullptr, bool density = false, const std::vector * other_token_scores = nullptr, - double split_fraction = 0.0); + double split_fraction = 0.0, + // Document starts in prompt-token coordinates, for the document + // prior. A candidate belongs to the last document starting at or + // before its first token. Null or shorter than + // ``kPFlashMinPriorDocuments`` leaves the prior a no-op. + const std::vector * documents = nullptr); } // namespace dflash::common diff --git a/server/src/pflash/pflash_drafter.cpp b/server/src/pflash/pflash_drafter.cpp index 2a65f5128..86db52149 100644 --- a/server/src/pflash/pflash_drafter.cpp +++ b/server/src/pflash/pflash_drafter.cpp @@ -104,7 +104,8 @@ std::vector drafter_score_and_compress( int n_lookahead, int pool_kernel, int score_query_end, - const std::vector & required_instruction_spans) { + const std::vector & required_instruction_spans, + const std::vector & document_spans) { if (!ctx.loaded) { set_last_error("drafter not loaded"); return {}; @@ -145,12 +146,13 @@ std::vector drafter_score_and_compress( std::fprintf(stderr, "[pflash-select] config mode=%s active=%d chunk=%d " "query_parser=%s query_cap=%d query_actual=%d top_p=%.9g " - "top_k=%d input=%zu\n", + "top_k=%d doc_prior=%.9g doc_spans=%zu input=%zu\n", dflash::pflash::pflash_selection_mode_name(experiment.mode), (int) experiment.selection_active, experiment.chunk_size, dflash::pflash::pflash_query_parser_name(experiment.query_parser), experiment.query_tokens, n_lookahead, experiment.top_p, - experiment.top_k, ids.size()); + experiment.top_k, experiment.doc_prior_exponent, + document_spans.size(), ids.size()); std::fflush(stderr); } if (score_query_end < 0) { @@ -159,7 +161,8 @@ std::vector drafter_score_and_compress( } return qwen35_drafter_score_and_compress( ctx, ids, keep_ratio, chunk_size, n_lookahead, pool_kernel, - score_query_end, experiment, required_instruction_spans); + score_query_end, experiment, required_instruction_spans, + document_spans); } } // namespace dflash::common diff --git a/server/src/pflash/pflash_drafter.h b/server/src/pflash/pflash_drafter.h index 0ab3a6235..b01c4f083 100644 --- a/server/src/pflash/pflash_drafter.h +++ b/server/src/pflash/pflash_drafter.h @@ -82,6 +82,9 @@ std::vector drafter_score_and_compress( int pool_kernel = 13, int score_query_end = -1, const std::vector & - required_instruction_spans = {}); + required_instruction_spans = {}, + // Document starts in prompt-token coordinates, for the document prior + // (PFLASH_SELECT_DOC_PRIOR). Empty leaves the prior a no-op. + const std::vector & document_spans = {}); } // namespace dflash::common diff --git a/server/src/pflash/pflash_selection.cpp b/server/src/pflash/pflash_selection.cpp index a331d3e52..d411a36ff 100644 --- a/server/src/pflash/pflash_selection.cpp +++ b/server/src/pflash/pflash_selection.cpp @@ -21,6 +21,7 @@ constexpr const char * kQueryEnv = "PFLASH_SELECT_QUERY_TOKENS"; constexpr const char * kQueryParserEnv = "PFLASH_SELECT_QUERY_PARSER"; constexpr const char * kTopPEnv = "PFLASH_SELECT_TOP_P"; constexpr const char * kTopKEnv = "PFLASH_SELECT_TOPK"; +constexpr const char * kDocPriorEnv = "PFLASH_SELECT_DOC_PRIOR"; constexpr const char * kSegmentsEnv = "PFLASH_SELECT_SEGMENTS"; constexpr const char * kSelectEnv = "PFLASH_SELECT_SCORE"; constexpr const char * kScorerEnv = "PFLASH_SELECT_SCORER"; @@ -73,6 +74,7 @@ bool has_pflash_selection_environment() noexcept { std::getenv(kQueryParserEnv) != nullptr || std::getenv(kTopPEnv) != nullptr || std::getenv(kTopKEnv) != nullptr || + std::getenv(kDocPriorEnv) != nullptr || std::getenv(kSegmentsEnv) != nullptr || std::getenv(kSelectEnv) != nullptr || std::getenv(kScorerEnv) != nullptr || @@ -147,6 +149,11 @@ PFlashSelectionResult select_pflash_candidates( if (mode == PFlashSelectionMode::TopK && policy.top_k <= 0) { return invalid_result("PFlash top_k must be positive"); } + if (!std::isfinite(policy.doc_prior_exponent) || + policy.doc_prior_exponent < 0.0) { + return invalid_result( + "PFlash document prior exponent must be finite and non-negative"); + } std::vector source_ranges; source_ranges.reserve(candidates.size()); @@ -202,10 +209,43 @@ PFlashSelectionResult select_pflash_candidates( } } + // Document prior. A document's mass is the sum over its candidates -- + // mandatory ones included, since they are part of the prompt -- of + // max(0, score) * tokens, normalised by the prompt total. Ranking by + // score * share^exponent lets a document the head likes as a whole lift + // its own segments, which per-segment density ranking throws away. + std::vector> document_mass; + for (const auto & candidate : candidates) { + const double mass = std::max(0.0, candidate.score) * + (double) (candidate.end - candidate.begin); + auto it = std::find_if(document_mass.begin(), document_mass.end(), + [&](const auto & entry) { return entry.first == candidate.document; }); + if (it == document_mass.end()) { + document_mass.push_back({candidate.document, mass}); + } else { + it->second += mass; + } + } + double total_mass = 0.0; + for (const auto & entry : document_mass) total_mass += entry.second; + result.documents = document_mass.size(); + result.doc_prior_applied = policy.doc_prior_exponent > 0.0 && + document_mass.size() >= kPFlashMinPriorDocuments && total_mass > 0.0; + + const auto rank_score = [&](const PFlashSelectionCandidate * candidate) { + const double base = std::max(0.0, candidate->score); + if (!result.doc_prior_applied) return base; + double share = 0.0; + for (const auto & entry : document_mass) { + if (entry.first == candidate->document) { share = entry.second; break; } + } + return base * std::pow(share / total_mass, policy.doc_prior_exponent); + }; + std::sort(optional.begin(), optional.end(), - [](const auto * left, const auto * right) { - const double left_score = std::max(0.0, left->score); - const double right_score = std::max(0.0, right->score); + [&](const auto * left, const auto * right) { + const double left_score = rank_score(left); + const double right_score = rank_score(right); if (left_score != right_score) return left_score > right_score; return left->ordinal < right->ordinal; }); @@ -316,6 +356,7 @@ bool resolve_pflash_selection( const char * query_parser_raw = std::getenv(kQueryParserEnv); const char * top_p_raw = std::getenv(kTopPEnv); const char * top_k_raw = std::getenv(kTopKEnv); + const char * doc_prior_raw = std::getenv(kDocPriorEnv); const char * segments_raw = std::getenv(kSegmentsEnv); const char * select_raw = std::getenv(kSelectEnv); const char * scorer_raw = std::getenv(kScorerEnv); @@ -323,8 +364,8 @@ bool resolve_pflash_selection( PFlashSelectionConfig config; config.configured = mode_raw || chunk_raw || query_raw || - query_parser_raw || top_p_raw || top_k_raw || segments_raw || - select_raw || scorer_raw || split_raw; + query_parser_raw || top_p_raw || top_k_raw || doc_prior_raw || + segments_raw || select_raw || scorer_raw || split_raw; if (scorer_raw) { if (std::strcmp(scorer_raw, "head") == 0) { config.scorer = PFlashScorer::Head; @@ -423,6 +464,13 @@ bool resolve_pflash_selection( error = std::string(kTopKEnv) + " must be a positive integer"; return false; } + if (doc_prior_raw && + (!parse_double(doc_prior_raw, config.doc_prior_exponent) || + config.doc_prior_exponent < 0.0)) { + error = std::string(kDocPriorEnv) + + " must be a non-negative number"; + return false; + } if (config.mode == PFlashSelectionMode::TopK && config.top_k <= 0) { error = std::string(kTopKEnv) + " is required when " + std::string(kModeEnv) + " is top_k"; diff --git a/server/src/pflash/pflash_selection.h b/server/src/pflash/pflash_selection.h index 7df7df7f6..9ef9f278d 100644 --- a/server/src/pflash/pflash_selection.h +++ b/server/src/pflash/pflash_selection.h @@ -37,6 +37,9 @@ struct PFlashSelectionCandidate { int end = 0; double score = 0.0; bool mandatory = false; + // Which document this candidate falls in, for the document prior below. + // All-zero (one document) leaves the prior a no-op. + size_t document = 0; }; struct PFlashSelectionPolicy { @@ -49,8 +52,20 @@ struct PFlashSelectionPolicy { // TopK mode only: how many optional candidates to keep. Must be positive // in that mode and is ignored in the others. int top_k = 0; + // Document prior: an optional candidate is ranked by + // ``max(0, score) * document_mass_share ^ doc_prior_exponent``, where a + // document's mass is the sum over its candidates of + // ``max(0, score) * tokens`` normalised by the prompt total. 0 disables + // it and leaves the ranking byte-identical. It applies in every mode, so + // it composes with TopK. Fewer than ``kPFlashMinPriorDocuments`` distinct + // documents makes it a no-op: with one or two documents the shares carry + // no ranking information worth a reweight. + double doc_prior_exponent = 0.0; }; +// Below this many distinct documents the document prior is a no-op. +constexpr size_t kPFlashMinPriorDocuments = 3; + struct PFlashSelectionResult { bool ok = false; std::vector ordinals; @@ -58,6 +73,11 @@ struct PFlashSelectionResult { double retained_mass = 0.0; PFlashSelectionStop stop = PFlashSelectionStop::InvalidInput; std::string error; + // Distinct documents seen among the candidates, and whether the document + // prior actually reweighted the ranking (exponent > 0 and enough + // documents). Both are recorded in the compression trace. + size_t documents = 0; + bool doc_prior_applied = false; }; bool pflash_chunk_is_structurally_required( @@ -100,6 +120,7 @@ struct PFlashSelectionConfig { int query_tokens = 8; double top_p = 0.95; int top_k = 0; + double doc_prior_exponent = 0.0; PFlashSegmentation segmentation = PFlashSegmentation::Auto; PFlashCandidateScore candidate_score = PFlashCandidateScore::Auto; PFlashScorer scorer = PFlashScorer::Head; diff --git a/server/src/pflash/qwen35_drafter.cpp b/server/src/pflash/qwen35_drafter.cpp index 7f34d8b4c..63aa89008 100644 --- a/server/src/pflash/qwen35_drafter.cpp +++ b/server/src/pflash/qwen35_drafter.cpp @@ -100,7 +100,8 @@ std::vector qwen35_score_and_compress( int score_query_end, const dflash::pflash::PFlashSelectionConfig & experiment, const std::vector & required_instruction_spans, - std::vector * token_scores_out) { + std::vector * token_scores_out, + const std::vector * document_spans) { const int S = (int)ids.size(); const int hidden = w.n_embd; @@ -379,7 +380,10 @@ std::vector qwen35_score_and_compress( if (experiment.selection_active) { return select_pflash_chunks( ids, smooth_score, keep_ratio, n_lookahead, score_query_end, - pk, experiment, required_instruction_spans, false, true); + pk, experiment, required_instruction_spans, false, true, + /*segments=*/nullptr, /*density=*/false, + /*other_token_scores=*/nullptr, /*split_fraction=*/0.0, + document_spans); } std::vector> chunk_means; @@ -533,7 +537,8 @@ std::vector qwen35_strict_score_and_compress( const std::vector & required_instruction_spans, std::vector * token_mass_out, std::vector * segments_out, - bool * density_out) { + bool * density_out, + const std::vector * document_spans) { TargetWeights & w = st.weights; const int S = (int)ids.size(); @@ -919,7 +924,9 @@ std::vector qwen35_strict_score_and_compress( ids, token_mass, keep_ratio, n_lookahead, score_query_end, /*pool_kernel=*/1, experiment, required_instruction_spans, /*direct_mass=*/true, /*write_trace=*/true, - segments.empty() ? nullptr : &segments, density); + segments.empty() ? nullptr : &segments, density, + /*other_token_scores=*/nullptr, /*split_fraction=*/0.0, + document_spans); } std::vector qwen35_drafter_score_and_compress( @@ -931,7 +938,8 @@ std::vector qwen35_drafter_score_and_compress( int pool_kernel, int score_query_end, const dflash::pflash::PFlashSelectionConfig & experiment, - const std::vector & required_instruction_spans) { + const std::vector & required_instruction_spans, + const std::vector & document_spans) { if (!ctx.state) { set_last_error("qwen35 drafter state missing"); return {}; @@ -953,14 +961,14 @@ std::vector qwen35_drafter_score_and_compress( if (qwen35_strict_score_and_compress( *st, ids, keep_ratio, n_lookahead, score_query_end, experiment, required_instruction_spans, &head_mass, &head_segments, - &head_density).empty()) { + &head_density, &document_spans).empty()) { return {}; } std::vector other_scores; if (qwen35_score_and_compress(st->weights, ids, keep_ratio, chunk_size, n_lookahead, pool_kernel, score_query_end, experiment, required_instruction_spans, - &other_scores).empty()) { + &other_scores, &document_spans).empty()) { return {}; } if (other_scores.size() != head_mass.size()) { @@ -977,12 +985,14 @@ std::vector qwen35_drafter_score_and_compress( /*pool_kernel=*/1, experiment, required_instruction_spans, /*direct_mass=*/true, /*write_trace=*/true, head_segments.empty() ? nullptr : &head_segments, head_density, - &other_scores, experiment.split_fraction); + &other_scores, experiment.split_fraction, &document_spans); } if (experiment.selection_active && !force_legacy) { return qwen35_strict_score_and_compress( *st, ids, keep_ratio, n_lookahead, score_query_end, experiment, - required_instruction_spans); + required_instruction_spans, /*token_mass_out=*/nullptr, + /*segments_out=*/nullptr, /*density_out=*/nullptr, + &document_spans); } if (st->head_loaded && !experiment.selection_active) { set_last_error("Qwen3.5 scoring head requires strict selection"); @@ -991,7 +1001,9 @@ std::vector qwen35_drafter_score_and_compress( return qwen35_score_and_compress(st->weights, ids, keep_ratio, chunk_size, n_lookahead, pool_kernel, score_query_end, experiment, - required_instruction_spans); + required_instruction_spans, + /*token_scores_out=*/nullptr, + &document_spans); } } // namespace dflash::common diff --git a/server/src/pflash/qwen35_drafter.h b/server/src/pflash/qwen35_drafter.h index c68d46e3f..675552d7c 100644 --- a/server/src/pflash/qwen35_drafter.h +++ b/server/src/pflash/qwen35_drafter.h @@ -76,7 +76,8 @@ std::vector qwen35_score_and_compress( int score_query_end, const dflash::pflash::PFlashSelectionConfig & experiment, const std::vector & required_instruction_spans, - std::vector * token_scores_out = nullptr); + std::vector * token_scores_out = nullptr, + const std::vector * document_spans = nullptr); // The block-15 scoring head under strict budget selection. std::vector qwen35_strict_score_and_compress( @@ -89,7 +90,8 @@ std::vector qwen35_strict_score_and_compress( const std::vector & required_instruction_spans, std::vector * token_mass_out = nullptr, std::vector * segments_out = nullptr, - bool * density_out = nullptr); + bool * density_out = nullptr, + const std::vector * document_spans = nullptr); // Arch dispatch target of drafter_score_and_compress. std::vector qwen35_drafter_score_and_compress( @@ -101,6 +103,7 @@ std::vector qwen35_drafter_score_and_compress( int pool_kernel, int score_query_end, const dflash::pflash::PFlashSelectionConfig & experiment, - const std::vector & required_instruction_spans); + const std::vector & required_instruction_spans, + const std::vector & document_spans = {}); } // namespace dflash::common diff --git a/server/src/qwen35/qwen35_backend.cpp b/server/src/qwen35/qwen35_backend.cpp index e7a0ad09c..9d784ce1b 100644 --- a/server/src/qwen35/qwen35_backend.cpp +++ b/server/src/qwen35/qwen35_backend.cpp @@ -1201,7 +1201,8 @@ std::vector Qwen35Backend::compress_batch( result.compressed_ids = drafter_score_and_compress( drafter_ctx_, request.input_ids, request.keep_ratio, /*chunk_size=*/32, request.score_query_tokens, /*pool_kernel=*/13, - score_query_end, request.required_instruction_spans); + score_query_end, request.required_instruction_spans, + request.document_spans); result.ok = !result.compressed_ids.empty(); if (result.ok) { std::fprintf(stderr, "[compress] %zu -> %zu tokens\n", diff --git a/server/src/qwen35/qwen35_layer_split_adapter.cpp b/server/src/qwen35/qwen35_layer_split_adapter.cpp index 7acfcd41e..9ef900d4b 100644 --- a/server/src/qwen35/qwen35_layer_split_adapter.cpp +++ b/server/src/qwen35/qwen35_layer_split_adapter.cpp @@ -1392,7 +1392,8 @@ Qwen35LayerSplitAdapter::compress(const ModelBackend::CompressRequest & req) { result.compressed_ids = drafter_score_and_compress( pflash_drafter_, req.input_ids, req.keep_ratio, /*chunk_size=*/32, req.score_query_tokens, /*pool_kernel=*/13, - score_query_end, req.required_instruction_spans); + score_query_end, req.required_instruction_spans, + req.document_spans); result.ok = !result.compressed_ids.empty(); if (result.ok) { std::fprintf(stderr, "[target-split][compress] %zu -> %zu tokens\n", diff --git a/server/src/server/http_server.cpp b/server/src/server/http_server.cpp index 37199888e..6ff18473b 100644 --- a/server/src/server/http_server.cpp +++ b/server/src/server/http_server.cpp @@ -371,6 +371,133 @@ PFlashTokenSpan pflash_changed_token_span( : PFlashTokenSpan{-1, -1}; } +namespace { + +// Decoded prompt text plus the character offset each token starts at. +// Tokenizer::decode concatenates per-token text, so the offsets are exact. +struct DecodedPrompt { + std::string text; + std::vector token_begin; // size == prompt.size() +}; + +DecodedPrompt decode_prompt_with_offsets( + const Tokenizer & tokenizer, + const std::vector & prompt) { + DecodedPrompt out; + out.token_begin.reserve(prompt.size()); + std::vector one(1, 0); + for (int32_t id : prompt) { + out.token_begin.push_back(out.text.size()); + one[0] = id; + out.text += tokenizer.decode(one); + } + return out; +} + +// The token containing character offset `at`, clamped into the prompt. +int token_at_offset(const DecodedPrompt & decoded, size_t at) { + const auto upper = std::upper_bound( + decoded.token_begin.begin(), decoded.token_begin.end(), at); + if (upper == decoded.token_begin.begin()) return 0; + return (int) (upper - decoded.token_begin.begin() - 1); +} + +bool is_digit_run(const std::string & text, size_t at, size_t & after) { + after = at; + while (after < text.size() && text[after] >= '0' && text[after] <= '9') ++after; + return after > at; +} + +// Character offsets where a document marker opens: a line beginning +// "Document :" or a "[DOC-]" tag anywhere in the line. +std::vector document_marker_offsets(const std::string & text) { + static const std::string kLabel = "Document "; + static const std::string kTag = "[DOC-"; + std::vector offsets; + for (size_t at = text.find(kLabel); at != std::string::npos; + at = text.find(kLabel, at + 1)) { + if (at != 0 && text[at - 1] != '\n') continue; + size_t after = 0; + if (!is_digit_run(text, at + kLabel.size(), after)) continue; + if (after >= text.size() || text[after] != ':') continue; + offsets.push_back(at); + } + for (size_t at = text.find(kTag); at != std::string::npos; + at = text.find(kTag, at + 1)) { + size_t after = 0; + if (!is_digit_run(text, at + kTag.size(), after)) continue; + if (after >= text.size() || text[after] != ']') continue; + offsets.push_back(at); + } + std::sort(offsets.begin(), offsets.end()); + offsets.erase(std::unique(offsets.begin(), offsets.end()), offsets.end()); + return offsets; +} + +// Tile [0, tokens) from ascending document start tokens. Starts before the +// first marker stay with the first document, so every token has a document. +std::vector tile_document_starts( + std::vector starts, int tokens) { + std::vector spans; + std::sort(starts.begin(), starts.end()); + starts.erase(std::unique(starts.begin(), starts.end()), starts.end()); + while (!starts.empty() && starts.back() >= tokens) starts.pop_back(); + if ((size_t) starts.size() < dflash::pflash::kPFlashMinPriorDocuments) { + return spans; + } + if (starts.front() != 0) starts.insert(starts.begin(), 0); + for (size_t index = 0; index < starts.size(); ++index) { + const int end = index + 1 < starts.size() ? starts[index + 1] : tokens; + spans.push_back({starts[index], end}); + } + return spans; +} + +} // namespace + +std::vector pflash_detect_document_spans( + const Tokenizer & tokenizer, + const std::vector & prompt) { + if (prompt.empty()) return {}; + const DecodedPrompt decoded = decode_prompt_with_offsets(tokenizer, prompt); + const auto offsets = document_marker_offsets(decoded.text); + if (offsets.size() < dflash::pflash::kPFlashMinPriorDocuments) return {}; + std::vector starts; + starts.reserve(offsets.size()); + for (size_t offset : offsets) { + starts.push_back(token_at_offset(decoded, offset)); + } + return tile_document_starts(std::move(starts), (int) prompt.size()); +} + +std::vector pflash_document_spans_from_ranges( + const Tokenizer & tokenizer, + const std::vector & prompt, + const std::vector> & ranges) { + if (prompt.empty() || + ranges.size() < dflash::pflash::kPFlashMinPriorDocuments) { + return {}; + } + const int tokens = (int) prompt.size(); + bool token_coordinates = true; + for (const auto & range : ranges) { + if (range.second > tokens) { token_coordinates = false; break; } + } + std::vector starts; + starts.reserve(ranges.size()); + if (token_coordinates) { + for (const auto & range : ranges) starts.push_back(range.first); + } else { + const DecodedPrompt decoded = + decode_prompt_with_offsets(tokenizer, prompt); + for (const auto & range : ranges) { + if ((size_t) range.first >= decoded.text.size()) continue; + starts.push_back(token_at_offset(decoded, (size_t) range.first)); + } + } + return tile_document_starts(std::move(starts), tokens); +} + std::vector canonicalize_pflash_token_spans( std::vector spans) { std::sort(spans.begin(), spans.end(), [] ( @@ -2649,6 +2776,7 @@ bool HttpServer::handle_model_request(SocketHandle fd, ParsedRequest & req, req.session_id = parse_session_id_from_body(body); req.pflash_query = parse_pflash_query_from_body(body); req.pflash_required = parse_pflash_required_from_body(body); + req.pflash_documents = parse_pflash_documents_from_body(body); // PPP rearrange (optional): peel ephemeral system banners into a // following system message so the first chat boundary is stable. @@ -3666,10 +3794,28 @@ std::string HttpServer::apply_pflash_compression( } } + // Document prior input: client-declared ranges when the request carries + // them, otherwise the served prompt's own document markers. Either way + // fewer than three documents yields no spans and the prior stays off. + std::vector document_spans; + if (experiment.doc_prior_exponent > 0.0 && drafter_tokenizer_) { + document_spans = req.pflash_documents.empty() + ? http_detail::pflash_detect_document_spans( + *drafter_tokenizer_, drafter_ids) + : http_detail::pflash_document_spans_from_ranges( + *drafter_tokenizer_, drafter_ids, req.pflash_documents); + std::fprintf(stderr, + "[pflash-docs] source=%s documents=%zu exponent=%.9g\n", + req.pflash_documents.empty() ? "detected" : "request", + document_spans.size(), experiment.doc_prior_exponent); + std::fflush(stderr); + } + ModelBackend::CompressRequest compress_request; compress_request.input_ids = std::move(drafter_ids); compress_request.required_instruction_spans = std::move(required_instruction_spans); + compress_request.document_spans = std::move(document_spans); compress_request.keep_ratio = http_detail::resolve_pflash_keep_ratio( pflash_keep_ratio(config_, prompt_tokens), req.session_id, sessions_); if (query_window.valid()) { diff --git a/server/src/server/http_server.h b/server/src/server/http_server.h index 2a628c132..febcdb771 100644 --- a/server/src/server/http_server.h +++ b/server/src/server/http_server.h @@ -350,6 +350,27 @@ int pflash_query_search_begin_from_sentinel( std::string pflash_token_fingerprint( const std::vector & ids); +// Document starts detected in the served prompt, in prompt-token +// coordinates. A line matching ``Document :`` or a ``[DOC-]`` tag opens +// a document; the returned spans tile [0, prompt.size()) so every token +// belongs to exactly one document. Fewer than +// ``dflash::pflash::kPFlashMinPriorDocuments`` starts returns empty, which +// leaves the document prior a no-op -- with one or two documents the mass +// shares carry no ranking information worth a reweight. +std::vector pflash_detect_document_spans( + const Tokenizer & tokenizer, + const std::vector & prompt); + +// Turn client-declared ``pflash_documents`` ranges into document spans. +// Ranges are token indices when every value fits inside the prompt's token +// count, and character offsets into the decoded prompt text otherwise. The +// result tiles [0, prompt.size()) from the range starts, exactly like the +// detected form. Returns empty when the ranges are unusable or too few. +std::vector pflash_document_spans_from_ranges( + const Tokenizer & tokenizer, + const std::vector & prompt, + const std::vector> & ranges); + bool pflash_full_cache_restore_allowed( bool selection_environment_present) noexcept; bool pflash_continuation_must_fail_closed( @@ -407,6 +428,9 @@ struct ParsedRequest { // compression (e.g. an answer-format directive embedded in a user // message). Each occurrence is mapped and retained as a mandatory span. std::vector pflash_required; + // Client-declared document ranges for the document prior. Empty falls + // back to detecting the served prompt's document markers. + std::vector> pflash_documents; DiskPrefixCachePolicy disk_cache_policy; // PPP: stable pin cut for tool-heavy requests (0 = use default boundary). int pin_end_token = 0; @@ -808,6 +832,37 @@ inline std::vector parse_pflash_required_from_body(const json & bod return result; } +// PFlash: client-declared document ranges for the document prior, as a list +// of two-element [begin, end] arrays. Accepted at the top level or under +// extra_body, like pflash_query. Absent means the runtime detects documents +// from the served prompt's markers instead. +inline std::vector> parse_pflash_documents_from_body( + const json & body) { + const json * field = nullptr; + if (body.contains("extra_body")) { + const auto & eb = body["extra_body"]; + if (eb.is_object() && eb.contains("pflash_documents") && + eb["pflash_documents"].is_array()) { + field = &eb["pflash_documents"]; + } + } + if (!field && body.contains("pflash_documents") && + body["pflash_documents"].is_array()) { + field = &body["pflash_documents"]; + } + std::vector> result; + if (!field) return result; + for (const auto & entry : *field) { + if (!entry.is_array() || entry.size() != 2) continue; + if (!entry[0].is_number_integer() || !entry[1].is_number_integer()) continue; + const int begin = entry[0].get(); + const int end = entry[1].get(); + if (begin < 0 || end <= begin) continue; + result.push_back({begin, end}); + } + return result; +} + inline std::string parse_session_id_from_body(const json & body) { if (body.contains("extra_body")) { const auto & eb = body["extra_body"]; diff --git a/server/test/test_pflash_selection.cpp b/server/test/test_pflash_selection.cpp index 2dfb7c19f..3814246c5 100644 --- a/server/test/test_pflash_selection.cpp +++ b/server/test/test_pflash_selection.cpp @@ -858,3 +858,170 @@ TEST_CASE(PFlashSelectionFixture, split_selection_rejects_top_k) { REQUIRE(!result.ok); REQUIRE(result.stop == PFlashSelectionStop::InvalidInput); } + +namespace { + +// Three documents of two candidates each. Document 2 (ordinals 4,5) has the +// lowest per-segment scores but the most mass; document 0 holds the single +// highest-scoring segment. Per-segment ranking prefers ordinal 0, the +// document prior prefers document 2. +std::vector document_candidates() { + const double scores[6] = {9.0, 1.0, 2.0, 2.0, 5.0, 5.0}; + const size_t documents[6] = {0, 0, 1, 1, 2, 2}; + std::vector candidates; + for (size_t index = 0; index < 6; ++index) { + const int begin = (int) index * 100; + PFlashSelectionCandidate c = candidate(index, begin, begin + 100, scores[index]); + c.document = documents[index]; + candidates.push_back(c); + } + return candidates; +} + +} // namespace + +TEST_CASE(PFlashSelectionFixture, document_prior_reweights_by_document_mass_share) { + const auto candidates = document_candidates(); + // Masses: doc0 = (9+1)*100 = 1000, doc1 = 400, doc2 = 1000; total 2400. + // Shares: 0.41667, 0.16667, 0.41667. + const auto plain = select_pflash_candidates( + candidates, PFlashSelectionPolicy{200, 0.95, false}, + PFlashSelectionMode::BudgetOnly); + REQUIRE(plain.ok); + REQUIRE(!plain.doc_prior_applied); + REQUIRE(plain.documents == 3); + // Without the prior the two highest raw scores win: ordinal 0 (9) and + // one of the 5s. + require_ordinals(plain, {0, 4}); + + const auto prior = select_pflash_candidates( + candidates, PFlashSelectionPolicy{200, 0.95, false, 0, 1.0}, + PFlashSelectionMode::BudgetOnly); + REQUIRE(prior.ok); + REQUIRE(prior.doc_prior_applied); + REQUIRE(prior.documents == 3); + // 9 * 0.41667 = 3.75 still beats 5 * 0.41667 = 2.083? No: 2.083 < 3.75, + // so ordinal 0 stays first; ordinal 4 (2.083) beats ordinal 1 + // (1 * 0.41667 = 0.4167) and ordinal 2 (2 * 0.16667 = 0.333). + require_ordinals(prior, {0, 4}); + + // A larger exponent sharpens the shares; the weak document falls further + // behind, which is the whole point of the knob. + const auto sharp = select_pflash_candidates( + candidates, PFlashSelectionPolicy{400, 0.95, false, 0, 3.0}, + PFlashSelectionMode::BudgetOnly); + REQUIRE(sharp.ok); + REQUIRE(sharp.doc_prior_applied); + // doc1's segments (share 0.1667^3 = 0.00463) rank below every segment of + // doc0 and doc2 (share 0.41667^3 = 0.0723). + require_ordinals(sharp, {0, 1, 4, 5}); +} + +TEST_CASE(PFlashSelectionFixture, document_prior_exponent_zero_is_todays_ranking) { + const auto candidates = document_candidates(); + const auto off = select_pflash_candidates( + candidates, PFlashSelectionPolicy{300, 0.95, false, 0, 0.0}, + PFlashSelectionMode::BudgetOnly); + const auto shipped = select_pflash_candidates( + candidates, PFlashSelectionPolicy{300, 0.95, false}, + PFlashSelectionMode::BudgetOnly); + REQUIRE(off.ok); + REQUIRE(!off.doc_prior_applied); + REQUIRE(off.ordinals == shipped.ordinals); + REQUIRE(off.retained_tokens == shipped.retained_tokens); + REQUIRE(off.stop == shipped.stop); +} + +TEST_CASE(PFlashSelectionFixture, document_prior_is_a_no_op_below_three_documents) { + // Two documents: the shares are real but the prior stays off by rule. + std::vector candidates = document_candidates(); + for (auto & c : candidates) if (c.document == 2) c.document = 1; + const auto two = select_pflash_candidates( + candidates, PFlashSelectionPolicy{200, 0.95, false, 0, 1.0}, + PFlashSelectionMode::BudgetOnly); + REQUIRE(two.ok); + REQUIRE(two.documents == 2); + REQUIRE(!two.doc_prior_applied); + + // One document (a single book, as NoLiMa serves) is a no-op twice over: + // by the rule and because the only share is 1. + for (auto & c : candidates) c.document = 0; + const auto one = select_pflash_candidates( + candidates, PFlashSelectionPolicy{200, 0.95, false, 0, 1.0}, + PFlashSelectionMode::BudgetOnly); + REQUIRE(one.ok); + REQUIRE(one.documents == 1); + REQUIRE(!one.doc_prior_applied); + require_ordinals(one, {0, 4}); +} + +TEST_CASE(PFlashSelectionFixture, document_prior_composes_with_top_k) { + const auto candidates = document_candidates(); + // K=2 with the prior: the rank order is the prior's, the count is K's. + const auto composed = select_pflash_candidates( + candidates, PFlashSelectionPolicy{100000, 0.95, false, 2, 3.0}, + PFlashSelectionMode::TopK); + REQUIRE(composed.ok); + REQUIRE(composed.stop == PFlashSelectionStop::TopKReached); + REQUIRE(composed.doc_prior_applied); + REQUIRE(composed.retained_tokens == 200); + // Sharpened shares put doc0's 9 first and doc2's 5 second; doc1 is out. + require_ordinals(composed, {0, 4}); + + // The same K without the prior keeps the raw top two, which here is the + // same pair -- so also check a K that exposes the reordering below them. + const auto plain_three = select_pflash_candidates( + candidates, PFlashSelectionPolicy{100000, 0.95, false, 3}, + PFlashSelectionMode::TopK); + const auto prior_three = select_pflash_candidates( + candidates, PFlashSelectionPolicy{100000, 0.95, false, 3, 3.0}, + PFlashSelectionMode::TopK); + REQUIRE(plain_three.ok); + REQUIRE(prior_three.ok); + require_ordinals(plain_three, {0, 4, 5}); + require_ordinals(prior_three, {0, 4, 5}); + // At rank four the orders diverge: raw picks doc1's 2, the prior picks + // doc0's 1 because doc1's share is cubed away. + const auto plain_four = select_pflash_candidates( + candidates, PFlashSelectionPolicy{100000, 0.95, false, 4}, + PFlashSelectionMode::TopK); + const auto prior_four = select_pflash_candidates( + candidates, PFlashSelectionPolicy{100000, 0.95, false, 4, 3.0}, + PFlashSelectionMode::TopK); + require_ordinals(plain_four, {0, 2, 4, 5}); + require_ordinals(prior_four, {0, 1, 4, 5}); +} + +TEST_CASE(PFlashSelectionFixture, document_prior_rejects_a_negative_exponent) { + const auto candidates = document_candidates(); + REQUIRE(!select_pflash_candidates( + candidates, PFlashSelectionPolicy{200, 0.95, false, 0, -1.0}, + PFlashSelectionMode::BudgetOnly).ok); +} + +TEST_CASE(PFlashSelectionFixture, document_prior_environment_resolves_or_fails) { + CleanPFlashEnv clean; + luce_test::ScopedEnvVar doc_prior{"PFLASH_SELECT_DOC_PRIOR", nullptr}; + set_env(kModeEnv, "budget_only"); + REQUIRE(resolve_or_fail(32768, 1024).doc_prior_exponent == 0.0); + + set_env("PFLASH_SELECT_DOC_PRIOR", "1.0"); + auto config = resolve_or_fail(32768, 1024); + REQUIRE(std::fabs(config.doc_prior_exponent - 1.0) < 1e-12); + + // It composes with top_k in the resolved config too. + set_env(kModeEnv, "top_k"); + set_env(kTopKEnv, "20"); + config = resolve_or_fail(32768, 1024); + REQUIRE(config.mode == PFlashSelectionMode::TopK); + REQUIRE(config.top_k == 20); + REQUIRE(std::fabs(config.doc_prior_exponent - 1.0) < 1e-12); + + PFlashSelectionConfig invalid; + std::string error; + for (const char * bad : {"-1", "abc", ""}) { + set_env("PFLASH_SELECT_DOC_PRIOR", bad); + REQUIRE(!resolve_pflash_selection(32768, 1024, invalid, error)); + REQUIRE(error.find("PFLASH_SELECT_DOC_PRIOR") != std::string::npos); + } +} From 739a8c3a6dcd136cf88cdecafb000896ca5a0338 Mon Sep 17 00:00:00 2001 From: Graffioh Date: Tue, 22 Sep 2026 08:55:19 +0000 Subject: [PATCH 08/26] feat(pflash): force the headers of the best documents Keeping a document's body while dropping the header that names it leaves the model unable to cite what it is quoting, and the score ranking drops headers first because they are short and low density. PFLASH_SELECT_FORCE_DOC_HEADS keeps the first segment of each of the D highest-mass documents, charged after the structurally required spans and before the fill, so a header can never starve a mandatory span and a header that no longer fits is dropped rather than fatal. It reuses the document detection the prior added and composes with the prior and with top_k. The trace and the selector log record the configured D and the headers actually forced. Co-Authored-By: Claude Opus 5 (1M context) --- server/README.md | 1 + server/src/pflash/pflash_compress.cpp | 12 ++- server/src/pflash/pflash_compress.h | 3 + server/src/pflash/pflash_drafter.cpp | 4 +- server/src/pflash/pflash_selection.cpp | 111 ++++++++++++++++----- server/src/pflash/pflash_selection.h | 9 ++ server/src/server/http_server.cpp | 8 +- server/test/test_pflash_selection.cpp | 132 +++++++++++++++++++++++++ 8 files changed, 247 insertions(+), 33 deletions(-) diff --git a/server/README.md b/server/README.md index 3f87240d5..3ab82a8a7 100644 --- a/server/README.md +++ b/server/README.md @@ -378,6 +378,7 @@ the whole request's device footprint. `/status/json` reports | `--prefill-upstream-model ` | none | Model name forwarded upstream. | | `PFLASH_SELECT_MODE=top_k` + `PFLASH_SELECT_TOPK ` | budget-only fill | Rank rule: keep the K highest-scoring optional segments in score order instead of filling the keep ratio, with the keep-ratio budget still a hard ceiling (min(K segments, the budget)). Use it where the evidence is compact and sits in the first few ranks -- needle retrieval, passage QA, code -- so a small K reaches it for a fraction of the budget's tokens. Do not use it where the answer needs a whole document identified, since the evidence there spans many segments and K cuts it off. | | `PFLASH_SELECT_DOC_PRIOR ` | `0` (off) | Document prior: rank an optional segment by `max(0, score) * document_mass_share ^ E` instead of its score alone, where a document's mass is the sum over its segments of `max(0, score) * tokens` normalised by the prompt total. The head identifies the relevant document when its mass is aggregated per document, which per-segment density ranking throws away, so this helps where the answer needs a whole document found among many. `0` leaves the ranking unchanged and it applies in every mode, so it composes with `top_k`. Documents come from the request's `pflash_documents` ranges, or from the prompt's own `Document :` / `[DOC-]` markers; fewer than three documents makes it a no-op, which the compression trace records alongside the exponent. | +| `PFLASH_SELECT_FORCE_DOC_HEADS ` | `0` (off) | Attribution: keep the first segment of each of the D highest-mass documents whatever its own score, charged with the mandatory spans before the fill. A compressed context that keeps a document's text but drops the header identifying it leaves the model unable to cite its source, and a header is short, so it costs far less than the body it names. Headers that no longer fit the budget are dropped rather than fatal, and a structurally required span is always charged first. Uses the same documents as the prior and composes with it and with `top_k`. Note when reading BRIGHT numbers: that benchmark as adapted here is scored by naming a document tag, so part of any gain there is an artefact of the adaptation rather than better evidence -- judge it on body-evidence retention and on sets whose answers live in the text. | With a Qwen3.5-0.8B drafter and strict budget selection (`PFLASH_SELECT_MODE=budget_only`, `PFLASH_SELECT_CHUNK_SIZE`, diff --git a/server/src/pflash/pflash_compress.cpp b/server/src/pflash/pflash_compress.cpp index 3e2e9c2c8..112a5b0ba 100644 --- a/server/src/pflash/pflash_compress.cpp +++ b/server/src/pflash/pflash_compress.cpp @@ -90,6 +90,7 @@ void write_compression_trace( "\"token_budget\":%d,\"top_k\":%d," "\"doc_prior_exponent\":%.9g,\"documents\":%zu," "\"doc_prior_applied\":%s," + "\"force_doc_heads\":%d,\"forced_doc_heads\":%d," "\"retained_tokens\":%d", trace_fields->query_begin, trace_fields->query_end, dflash::pflash::pflash_selection_mode_name( @@ -98,6 +99,7 @@ void write_compression_trace( trace_fields->token_budget, trace_fields->top_k, trace_fields->doc_prior_exponent, trace_fields->documents, trace_fields->doc_prior_applied ? "true" : "false", + trace_fields->force_doc_heads, trace_fields->forced_doc_heads, trace_fields->retained_tokens); std::fputs(",\"required_instruction_spans\":[", file); if (trace_fields->required_instruction_spans) { @@ -265,7 +267,8 @@ std::vector select_pflash_chunks( const dflash::pflash::PFlashSelectionPolicy policy{selector_budget, config.top_p, /*skip_oversized=*/ segments != nullptr, config.top_k, - config.doc_prior_exponent}; + config.doc_prior_exponent, + config.force_doc_heads}; const auto selected = split ? dflash::pflash::select_pflash_split(candidates, other_candidates, policy, split_fraction, config.mode) : dflash::pflash::select_pflash_candidates(candidates, policy, config.mode); @@ -306,7 +309,7 @@ std::vector select_pflash_chunks( std::fprintf(stderr, "[pflash-select] selected mode=%s scorer=%s segments=%s score=%s chunk=%d query=%d " "budget=%d selected_tokens=%zu chunks=%zu/%d stop=%s mass=%.9g " - "docs=%zu doc_prior=%.9g applied=%d\n", + "docs=%zu doc_prior=%.9g applied=%d heads=%d/%d\n", dflash::pflash::pflash_selection_mode_name(config.mode), split ? "split" : "single", segments ? "probe" : "fixed", density ? "density" : "sum", @@ -314,7 +317,8 @@ std::vector select_pflash_chunks( selected.ordinals.size(), n_chunks, dflash::pflash::pflash_selection_stop_name(selected.stop), selected.retained_mass, selected.documents, config.doc_prior_exponent, - (int) selected.doc_prior_applied); + (int) selected.doc_prior_applied, selected.forced_doc_heads, + config.force_doc_heads); std::fflush(stderr); if (write_trace) { @@ -338,6 +342,8 @@ std::vector select_pflash_chunks( strict_fields.doc_prior_exponent = config.doc_prior_exponent; strict_fields.documents = selected.documents; strict_fields.doc_prior_applied = selected.doc_prior_applied; + strict_fields.force_doc_heads = config.force_doc_heads; + strict_fields.forced_doc_heads = selected.forced_doc_heads; write_compression_trace( input_tokens, keep_ratio, trace_chunk, query_tokens, pool_kernel, n_keep_approx, chunk_means, selected_mask, diff --git a/server/src/pflash/pflash_compress.h b/server/src/pflash/pflash_compress.h index 240a3580a..5bbedfe15 100644 --- a/server/src/pflash/pflash_compress.h +++ b/server/src/pflash/pflash_compress.h @@ -110,6 +110,9 @@ struct PFlashTraceFields { double doc_prior_exponent = 0.0; size_t documents = 0; bool doc_prior_applied = false; + // Attribution: the configured D and the headers actually forced. + int force_doc_heads = 0; + int forced_doc_heads = 0; }; void write_compression_trace( diff --git a/server/src/pflash/pflash_drafter.cpp b/server/src/pflash/pflash_drafter.cpp index 86db52149..e0e90415e 100644 --- a/server/src/pflash/pflash_drafter.cpp +++ b/server/src/pflash/pflash_drafter.cpp @@ -146,13 +146,13 @@ std::vector drafter_score_and_compress( std::fprintf(stderr, "[pflash-select] config mode=%s active=%d chunk=%d " "query_parser=%s query_cap=%d query_actual=%d top_p=%.9g " - "top_k=%d doc_prior=%.9g doc_spans=%zu input=%zu\n", + "top_k=%d doc_prior=%.9g doc_heads=%d doc_spans=%zu input=%zu\n", dflash::pflash::pflash_selection_mode_name(experiment.mode), (int) experiment.selection_active, experiment.chunk_size, dflash::pflash::pflash_query_parser_name(experiment.query_parser), experiment.query_tokens, n_lookahead, experiment.top_p, experiment.top_k, experiment.doc_prior_exponent, - document_spans.size(), ids.size()); + experiment.force_doc_heads, document_spans.size(), ids.size()); std::fflush(stderr); } if (score_query_end < 0) { diff --git a/server/src/pflash/pflash_selection.cpp b/server/src/pflash/pflash_selection.cpp index d411a36ff..2da03ce32 100644 --- a/server/src/pflash/pflash_selection.cpp +++ b/server/src/pflash/pflash_selection.cpp @@ -22,6 +22,7 @@ constexpr const char * kQueryParserEnv = "PFLASH_SELECT_QUERY_PARSER"; constexpr const char * kTopPEnv = "PFLASH_SELECT_TOP_P"; constexpr const char * kTopKEnv = "PFLASH_SELECT_TOPK"; constexpr const char * kDocPriorEnv = "PFLASH_SELECT_DOC_PRIOR"; +constexpr const char * kDocHeadsEnv = "PFLASH_SELECT_FORCE_DOC_HEADS"; constexpr const char * kSegmentsEnv = "PFLASH_SELECT_SEGMENTS"; constexpr const char * kSelectEnv = "PFLASH_SELECT_SCORE"; constexpr const char * kScorerEnv = "PFLASH_SELECT_SCORER"; @@ -75,6 +76,7 @@ bool has_pflash_selection_environment() noexcept { std::getenv(kTopPEnv) != nullptr || std::getenv(kTopKEnv) != nullptr || std::getenv(kDocPriorEnv) != nullptr || + std::getenv(kDocHeadsEnv) != nullptr || std::getenv(kSegmentsEnv) != nullptr || std::getenv(kSelectEnv) != nullptr || std::getenv(kScorerEnv) != nullptr || @@ -149,6 +151,9 @@ PFlashSelectionResult select_pflash_candidates( if (mode == PFlashSelectionMode::TopK && policy.top_k <= 0) { return invalid_result("PFlash top_k must be positive"); } + if (policy.force_doc_heads < 0) { + return invalid_result("PFlash forced document head count must not be negative"); + } if (!std::isfinite(policy.doc_prior_exponent) || policy.doc_prior_exponent < 0.0) { return invalid_result( @@ -191,29 +196,9 @@ PFlashSelectionResult select_pflash_candidates( std::vector selected_candidates; selected_candidates.reserve(candidates.size()); - std::vector optional; - optional.reserve(candidates.size()); - for (const auto & candidate : candidates) { - if (candidate.mandatory) { - const int length = candidate.end - candidate.begin; - if (length > policy.token_budget - result.retained_tokens) { - result = {}; - result.stop = PFlashSelectionStop::MandatoryQueryExceedsBudget; - result.error = "mandatory PFlash retention tokens exceed the token budget"; - return result; - } - selected_candidates.push_back(&candidate); - result.retained_tokens += length; - } else { - optional.push_back(&candidate); - } - } - - // Document prior. A document's mass is the sum over its candidates -- - // mandatory ones included, since they are part of the prompt -- of - // max(0, score) * tokens, normalised by the prompt total. Ranking by - // score * share^exponent lets a document the head likes as a whole lift - // its own segments, which per-segment density ranking throws away. + // Per-document mass: the sum over a document's candidates -- mandatory + // ones included, since they are part of the prompt -- of + // max(0, score) * tokens. It feeds both document rules below. std::vector> document_mass; for (const auto & candidate : candidates) { const double mass = std::max(0.0, candidate.score) * @@ -229,8 +214,77 @@ PFlashSelectionResult select_pflash_candidates( double total_mass = 0.0; for (const auto & entry : document_mass) total_mass += entry.second; result.documents = document_mass.size(); - result.doc_prior_applied = policy.doc_prior_exponent > 0.0 && + const bool enough_documents = document_mass.size() >= kPFlashMinPriorDocuments && total_mass > 0.0; + // Document prior: ranking by score * share^exponent lets a document the + // head likes as a whole lift its own segments, which per-segment density + // ranking throws away. + result.doc_prior_applied = policy.doc_prior_exponent > 0.0 && enough_documents; + + // Forced document headers: the first candidate of each of the D + // highest-mass documents joins the mandatory set. Attribution, not + // evidence -- a compressed context that drops the header identifying a + // document leaves the model unable to cite what it is quoting, and the + // header is short, so it costs far less than the body it names. + std::vector forced_heads; + if (policy.force_doc_heads > 0 && enough_documents) { + std::vector> ranked = document_mass; + std::sort(ranked.begin(), ranked.end(), + [](const auto & left, const auto & right) { + if (left.second != right.second) return left.second > right.second; + return left.first < right.first; + }); + const size_t wanted = std::min( + (size_t) policy.force_doc_heads, ranked.size()); + for (size_t index = 0; index < wanted; ++index) { + const PFlashSelectionCandidate * head = nullptr; + for (const auto & candidate : candidates) { + if (candidate.document != ranked[index].first) continue; + if (!head || candidate.begin < head->begin) head = &candidate; + } + // Already-mandatory headers are kept by the rule above anyway. + if (head && !head->mandatory) forced_heads.push_back(head); + } + } + // True mandatory candidates are charged first, so a forced header can + // never push a structurally required span out of the budget. + for (const auto & candidate : candidates) { + if (!candidate.mandatory) continue; + const int length = candidate.end - candidate.begin; + if (length > policy.token_budget - result.retained_tokens) { + const size_t documents_seen = result.documents; + result = {}; + result.documents = documents_seen; + result.stop = PFlashSelectionStop::MandatoryQueryExceedsBudget; + result.error = "mandatory PFlash retention tokens exceed the token budget"; + return result; + } + selected_candidates.push_back(&candidate); + result.retained_tokens += length; + } + // Then the forced headers, best document first. One that no longer fits + // is dropped rather than fatal: attribution yields to the budget, and it + // stays eligible for the ordinary fill below. + std::vector kept_heads; + for (const auto * head : forced_heads) { + const int length = head->end - head->begin; + if (length > policy.token_budget - result.retained_tokens) continue; + selected_candidates.push_back(head); + result.retained_tokens += length; + kept_heads.push_back(head->ordinal); + } + result.forced_doc_heads = (int) kept_heads.size(); + const auto head_was_kept = [&](size_t ordinal) { + return std::find(kept_heads.begin(), kept_heads.end(), ordinal) != + kept_heads.end(); + }; + + std::vector optional; + optional.reserve(candidates.size()); + for (const auto & candidate : candidates) { + if (candidate.mandatory || head_was_kept(candidate.ordinal)) continue; + optional.push_back(&candidate); + } const auto rank_score = [&](const PFlashSelectionCandidate * candidate) { const double base = std::max(0.0, candidate->score); @@ -357,6 +411,7 @@ bool resolve_pflash_selection( const char * top_p_raw = std::getenv(kTopPEnv); const char * top_k_raw = std::getenv(kTopKEnv); const char * doc_prior_raw = std::getenv(kDocPriorEnv); + const char * doc_heads_raw = std::getenv(kDocHeadsEnv); const char * segments_raw = std::getenv(kSegmentsEnv); const char * select_raw = std::getenv(kSelectEnv); const char * scorer_raw = std::getenv(kScorerEnv); @@ -365,7 +420,7 @@ bool resolve_pflash_selection( PFlashSelectionConfig config; config.configured = mode_raw || chunk_raw || query_raw || query_parser_raw || top_p_raw || top_k_raw || doc_prior_raw || - segments_raw || select_raw || scorer_raw || split_raw; + doc_heads_raw || segments_raw || select_raw || scorer_raw || split_raw; if (scorer_raw) { if (std::strcmp(scorer_raw, "head") == 0) { config.scorer = PFlashScorer::Head; @@ -471,6 +526,12 @@ bool resolve_pflash_selection( " must be a non-negative number"; return false; } + if (doc_heads_raw && + (!parse_int(doc_heads_raw, config.force_doc_heads) || + config.force_doc_heads < 0)) { + error = std::string(kDocHeadsEnv) + " must be a non-negative integer"; + return false; + } if (config.mode == PFlashSelectionMode::TopK && config.top_k <= 0) { error = std::string(kTopKEnv) + " is required when " + std::string(kModeEnv) + " is top_k"; diff --git a/server/src/pflash/pflash_selection.h b/server/src/pflash/pflash_selection.h index 9ef9f278d..2738bf778 100644 --- a/server/src/pflash/pflash_selection.h +++ b/server/src/pflash/pflash_selection.h @@ -61,6 +61,12 @@ struct PFlashSelectionPolicy { // documents makes it a no-op: with one or two documents the shares carry // no ranking information worth a reweight. double doc_prior_exponent = 0.0; + // Attribution: force the first candidate of each of the D highest-mass + // documents into the mandatory set before the fill, so a compressed + // context never drops the header that identifies a document it quotes. + // 0 disables it. It composes with the prior and with TopK, and the forced + // headers are charged against the budget like any mandatory candidate. + int force_doc_heads = 0; }; // Below this many distinct documents the document prior is a no-op. @@ -78,6 +84,8 @@ struct PFlashSelectionResult { // documents). Both are recorded in the compression trace. size_t documents = 0; bool doc_prior_applied = false; + // Document headers promoted to mandatory by ``force_doc_heads``. + int forced_doc_heads = 0; }; bool pflash_chunk_is_structurally_required( @@ -121,6 +129,7 @@ struct PFlashSelectionConfig { double top_p = 0.95; int top_k = 0; double doc_prior_exponent = 0.0; + int force_doc_heads = 0; PFlashSegmentation segmentation = PFlashSegmentation::Auto; PFlashCandidateScore candidate_score = PFlashCandidateScore::Auto; PFlashScorer scorer = PFlashScorer::Head; diff --git a/server/src/server/http_server.cpp b/server/src/server/http_server.cpp index 6ff18473b..fe0534cf9 100644 --- a/server/src/server/http_server.cpp +++ b/server/src/server/http_server.cpp @@ -3798,16 +3798,18 @@ std::string HttpServer::apply_pflash_compression( // them, otherwise the served prompt's own document markers. Either way // fewer than three documents yields no spans and the prior stays off. std::vector document_spans; - if (experiment.doc_prior_exponent > 0.0 && drafter_tokenizer_) { + if ((experiment.doc_prior_exponent > 0.0 || + experiment.force_doc_heads > 0) && drafter_tokenizer_) { document_spans = req.pflash_documents.empty() ? http_detail::pflash_detect_document_spans( *drafter_tokenizer_, drafter_ids) : http_detail::pflash_document_spans_from_ranges( *drafter_tokenizer_, drafter_ids, req.pflash_documents); std::fprintf(stderr, - "[pflash-docs] source=%s documents=%zu exponent=%.9g\n", + "[pflash-docs] source=%s documents=%zu exponent=%.9g heads=%d\n", req.pflash_documents.empty() ? "detected" : "request", - document_spans.size(), experiment.doc_prior_exponent); + document_spans.size(), experiment.doc_prior_exponent, + experiment.force_doc_heads); std::fflush(stderr); } diff --git a/server/test/test_pflash_selection.cpp b/server/test/test_pflash_selection.cpp index 3814246c5..44f191359 100644 --- a/server/test/test_pflash_selection.cpp +++ b/server/test/test_pflash_selection.cpp @@ -1025,3 +1025,135 @@ TEST_CASE(PFlashSelectionFixture, document_prior_environment_resolves_or_fails) REQUIRE(error.find("PFLASH_SELECT_DOC_PRIOR") != std::string::npos); } } + +TEST_CASE(PFlashSelectionFixture, forced_document_heads_keep_the_best_documents_identifiers) { + // Six candidates, three documents, two each. The first candidate of a + // document is its header: short and low-scoring, exactly what a score + // ranking drops. Masses: doc0 = 10*10 + 9*100 = 1000, + // doc1 = 1*10 + 3*100 = 310, doc2 = 10*10 + 5*100 = 600. + std::vector candidates{ + candidate(0, 0, 10, 10.0), // doc0 header + candidate(1, 10, 110, 9.0), // doc0 body + candidate(2, 110, 120, 1.0), // doc1 header + candidate(3, 120, 220, 3.0), // doc1 body + candidate(4, 220, 230, 10.0), // doc2 header + candidate(5, 230, 330, 5.0), // doc2 body + }; + const size_t documents[6] = {0, 0, 1, 1, 2, 2}; + for (size_t index = 0; index < candidates.size(); ++index) { + candidates[index].document = documents[index]; + } + + // Without forcing, a tight budget spends itself on the best body and + // leaves the weakest document with nothing -- not even its identifier. + const auto plain = select_pflash_candidates( + candidates, PFlashSelectionPolicy{130, 0.95, false}, + PFlashSelectionMode::BudgetOnly); + REQUIRE(plain.ok); + REQUIRE(plain.forced_doc_heads == 0); + require_ordinals(plain, {0, 1, 4}); + + // Forcing the headers of the top two documents by mass (doc0, doc2) + // charges them first; they are short, so the fill still gets the bodies. + const auto heads = select_pflash_candidates( + candidates, PFlashSelectionPolicy{230, 0.95, false, 0, 0.0, 2}, + PFlashSelectionMode::BudgetOnly); + REQUIRE(heads.ok); + REQUIRE(heads.forced_doc_heads == 2); + require_ordinals(heads, {0, 1, 4, 5}); + + // D above the document count forces every header that fits. At the same + // budget that kept no identifier for doc1 above, its header now survives. + const auto all_heads = select_pflash_candidates( + candidates, PFlashSelectionPolicy{130, 0.95, false, 0, 0.0, 9}, + PFlashSelectionMode::BudgetOnly); + REQUIRE(all_heads.ok); + REQUIRE(all_heads.forced_doc_heads == 3); + REQUIRE(all_heads.retained_tokens <= 130); + REQUIRE(all_heads.retained_tokens == 130); + // Three headers (30 tokens) plus the single body that still fits. + require_ordinals(all_heads, {0, 1, 2, 4}); +} + +TEST_CASE(PFlashSelectionFixture, forced_document_heads_compose_with_top_k_and_respect_the_budget) { + std::vector candidates{ + candidate(0, 0, 10, 1.0), + candidate(1, 10, 110, 9.0), + candidate(2, 110, 120, 1.0), + candidate(3, 120, 220, 3.0), + candidate(4, 220, 230, 1.0), + candidate(5, 230, 330, 5.0), + }; + const size_t documents[6] = {0, 0, 1, 1, 2, 2}; + for (size_t index = 0; index < candidates.size(); ++index) { + candidates[index].document = documents[index]; + } + // K counts only the ordinary fill; the forced headers are mandatory and + // sit outside it, so top-1 plus heads-of-3 keeps four candidates. + const auto composed = select_pflash_candidates( + candidates, PFlashSelectionPolicy{100000, 0.95, true, 1, 0.0, 3}, + PFlashSelectionMode::TopK); + REQUIRE(composed.ok); + REQUIRE(composed.stop == PFlashSelectionStop::TopKReached); + REQUIRE(composed.forced_doc_heads == 3); + require_ordinals(composed, {0, 1, 2, 4}); + REQUIRE(composed.retained_tokens == 130); + + // A budget too small for every header drops the ones that no longer fit + // rather than failing, and never exceeds the ceiling. + const auto tight = select_pflash_candidates( + candidates, PFlashSelectionPolicy{25, 0.95, true, 1, 0.0, 3}, + PFlashSelectionMode::TopK); + REQUIRE(tight.ok); + REQUIRE(tight.retained_tokens <= 25); + REQUIRE(tight.forced_doc_heads == 2); + + // A structurally required span is charged before any header, so forcing + // headers can never starve it. + std::vector with_mandatory = candidates; + with_mandatory[5].mandatory = true; + const auto safe = select_pflash_candidates( + with_mandatory, PFlashSelectionPolicy{110, 0.95, true, 1, 0.0, 3}, + PFlashSelectionMode::TopK); + REQUIRE(safe.ok); + REQUIRE(safe.retained_tokens <= 110); + // Ordinal 5 (100 tokens, mandatory) plus the headers that still fit. + REQUIRE(std::find(safe.ordinals.begin(), safe.ordinals.end(), (size_t) 5) != + safe.ordinals.end()); +} + +TEST_CASE(PFlashSelectionFixture, forced_document_heads_need_three_documents_and_a_valid_count) { + auto candidates = document_candidates(); + for (auto & c : candidates) c.document = 0; + const auto single = select_pflash_candidates( + candidates, PFlashSelectionPolicy{100000, 0.95, false, 0, 0.0, 5}, + PFlashSelectionMode::BudgetOnly); + REQUIRE(single.ok); + REQUIRE(single.forced_doc_heads == 0); + + REQUIRE(!select_pflash_candidates( + document_candidates(), PFlashSelectionPolicy{200, 0.95, false, 0, 0.0, -1}, + PFlashSelectionMode::BudgetOnly).ok); +} + +TEST_CASE(PFlashSelectionFixture, forced_document_heads_environment_resolves_or_fails) { + CleanPFlashEnv clean; + luce_test::ScopedEnvVar heads{"PFLASH_SELECT_FORCE_DOC_HEADS", nullptr}; + set_env(kModeEnv, "top_k"); + set_env(kTopKEnv, "20"); + REQUIRE(resolve_or_fail(32768, 1024).force_doc_heads == 0); + + set_env("PFLASH_SELECT_FORCE_DOC_HEADS", "5"); + const auto config = resolve_or_fail(32768, 1024); + REQUIRE(config.force_doc_heads == 5); + REQUIRE(config.mode == PFlashSelectionMode::TopK); + REQUIRE(config.top_k == 20); + + PFlashSelectionConfig invalid; + std::string error; + for (const char * bad : {"-1", "abc", ""}) { + set_env("PFLASH_SELECT_FORCE_DOC_HEADS", bad); + REQUIRE(!resolve_pflash_selection(32768, 1024, invalid, error)); + REQUIRE(error.find("PFLASH_SELECT_FORCE_DOC_HEADS") != std::string::npos); + } +} From 3e65d3cf720f33f407e9efe245ef8b9776c00336 Mon Sep 17 00:00:00 2001 From: Graffioh Date: Tue, 22 Sep 2026 15:01:39 +0000 Subject: [PATCH 09/26] pflash: derive the scorer query from the rendered chat tail Production PFlash is now chat-only: the scorer query is the last N tokens of the final message's content, located by the rendered prompt's own control markers (ChatMarkers) instead of benchmark message bookkeeping. Template machinery ("<|im_start|>", role names, "<|im_end|>", generation prompts) is never scored, and the last turn's role header is pinned mandatory so compression keeps the role envelope. Marker-less prompts and the latest_user parser keep the sentinel-render path; pflash_query stays a benchmark-only override, mapped against decoded text and pinned in full while the scorer consumes its bounded tail. Strict selection now owns multi-turn chit-chat: continuations run whole-prompt PFlash on the full history plus current turn instead of failing closed or routing to FlowKV. Request-scoped FlowKV disk compression remains the only rejected combination; unconfigured requests keep the legacy FlowKV continuation path. Generated with [Devin](https://devin.ai) Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- server/src/pflash/pflash_selection.h | 4 +- server/src/server/http_server.cpp | 282 +++++++++++++++++---- server/src/server/http_server.h | 28 ++- server/test/test_pflash_selection.cpp | 5 + server/test/test_server_unit.cpp | 341 +++++++++++++++++++++++--- 5 files changed, 582 insertions(+), 78 deletions(-) diff --git a/server/src/pflash/pflash_selection.h b/server/src/pflash/pflash_selection.h index 83b017ebb..084b6268a 100644 --- a/server/src/pflash/pflash_selection.h +++ b/server/src/pflash/pflash_selection.h @@ -123,7 +123,9 @@ enum class PFlashScorer { Head, Legacy, Split }; struct PFlashSelectionConfig { PFlashSelectionMode mode = PFlashSelectionMode::Legacy; - PFlashQueryParser query_parser = PFlashQueryParser::SemanticUser; + // Chat-first default: the scorer query is the tail of the last message's + // content. latest_user stays selectable for benchmark experiments. + PFlashQueryParser query_parser = PFlashQueryParser::ArbitraryTail; int chunk_size = 0; int query_tokens = 8; double top_p = 0.95; diff --git a/server/src/server/http_server.cpp b/server/src/server/http_server.cpp index f3f2b9f32..1bedce34c 100644 --- a/server/src/server/http_server.cpp +++ b/server/src/server/http_server.cpp @@ -534,15 +534,149 @@ std::string pflash_token_fingerprint( return encoded; } +PflashChatTailSpan pflash_last_message_content_span( + const Tokenizer & marker_tokenizer, + const ChatMarkers & markers, + const Tokenizer & tokenizer, + const std::vector & prompt) { + PflashChatTailSpan tail; + if (prompt.empty()) return tail; -bool pflash_full_cache_restore_allowed( - bool selection_environment_present) noexcept { - return !selection_environment_present; + const auto seq_text = [&marker_tokenizer]( + const std::vector & seq) { + std::string text; + for (const int32_t id : seq) text += marker_tokenizer.token_text(id); + return text; + }; + std::vector role_marks; + for (const auto & seq : markers.next_role_starts) { + std::string text = seq_text(seq); + if (!text.empty()) role_marks.push_back(std::move(text)); + } + std::vector end_marks; + for (const auto & seq : markers.end_msg_seqs) { + std::string text = seq_text(seq); + if (!text.empty()) end_marks.push_back(std::move(text)); + } + if (role_marks.empty() && end_marks.empty()) return tail; + + const DecodedPrompt decoded = + decode_prompt_with_offsets(tokenizer, prompt); + const std::string & text = decoded.text; + + // Last occurrence strictly before `before` (npos: anywhere), as + // (offset, length). Marker strings never overlap themselves, so the + // rightmost start across all needles is the answer. + const auto last_before = [&text]( + const std::vector & needles, size_t before) + -> std::pair { + size_t best = std::string::npos, len = 0; + for (const auto & needle : needles) { + const size_t at = before == std::string::npos + ? text.rfind(needle) + : before == 0 + ? std::string::npos + : text.rfind(needle, before - 1); + if (at != std::string::npos && + (best == std::string::npos || at > best)) { + best = at; + len = needle.size(); + } + } + return {best, len}; + }; + + const auto last_role = last_before(role_marks, std::string::npos); + const auto last_end = last_before(end_marks, std::string::npos); + if (last_role.first == std::string::npos && + last_end.first == std::string::npos) { + return tail; + } + + // Families whose markers already name the role (DeepSeek "<｜User｜>", + // Laguna "") delimit content at the marker itself. Generic markers + // (Qwen "<|im_start|>", Gemma "<|turn>") are followed by a "name\n" + // header line. + const bool marker_carries_role = + markers.role_starts_delimit || markers.family == "laguna"; + + // The character offset where the final message's content stops. The + // prompt's last marker is either the message's own end marker, the + // assistant generation marker (a role marker followed only by a + // role-name header), or a role marker that opened an unterminated + // message — in which case content runs to the prompt end. + size_t content_end_text; + if (last_role.first != std::string::npos && + (last_end.first == std::string::npos || + last_role.first > last_end.first)) { + const size_t after = last_role.first + last_role.second; + bool generation; + if (marker_carries_role) { + generation = text.find_first_not_of(" \t\n\r", after) == + std::string::npos; + } else { + size_t i = after; + while (i < text.size() && i - after < 16 && + std::isalpha((unsigned char) text[i])) ++i; + generation = i > after && i < text.size() && text[i] == '\n' && + text.find_first_not_of(" \t\n\r", i + 1) == + std::string::npos; + } + if (generation) { + if (marker_carries_role) { + content_end_text = last_role.first; + } else { + // Content ends at this message's end marker — the last one + // after the role marker that opened it, not an earlier + // turn's. + const auto prev_role = last_before(role_marks, last_role.first); + content_end_text = + (last_end.first != std::string::npos && + (prev_role.first == std::string::npos || + last_end.first > prev_role.first)) + ? last_end.first : last_role.first; + } + } else { + content_end_text = text.size(); + } + } else { + content_end_text = last_end.first; + } + + // The marker that opened the message containing content_end. + const auto role = last_before(role_marks, content_end_text); + size_t content_begin_text = 0; + if (role.first != std::string::npos) { + tail.role_begin = token_at_offset(decoded, role.first); + size_t begin = role.first + role.second; + if (!marker_carries_role) { + size_t i = begin; + while (i < content_end_text && i - begin < 16 && + std::isalpha((unsigned char) text[i])) ++i; + if (i > begin && i < content_end_text && text[i] == '\n') { + begin = i + 1; + } + } + content_begin_text = begin; + } + + tail.content_end = content_end_text >= text.size() + ? (int) prompt.size() + : token_at_offset(decoded, content_end_text); + // First token whose text starts at-or-after the content offset — a token + // merged across the header/content boundary stays with the header. + tail.content_begin = (int) (std::lower_bound( + decoded.token_begin.begin(), decoded.token_begin.end(), + content_begin_text) - decoded.token_begin.begin()); + if (tail.role_begin > tail.content_begin) { + tail.role_begin = tail.content_begin; + } + return tail; } -bool pflash_continuation_must_fail_closed( +bool pflash_full_cache_restore_allowed( bool selection_environment_present) noexcept { - return selection_environment_present; + return !selection_environment_present; } int pflash_target_token_ceiling( @@ -3493,6 +3627,8 @@ std::string HttpServer::apply_pflash_compression( const bool raw_text_input = req.messages.is_string(); const char * parser_input_kind = messages_input ? "messages" : (raw_text_input ? "raw_text" : "unsupported"); + const bool tail_parser = experiment.query_parser == + luce::pflash::PFlashQueryParser::ArbitraryTail; std::string parser_selection_rule; std::string last_user_text; int query_content_begin = -1; @@ -3501,7 +3637,23 @@ std::string HttpServer::apply_pflash_compression( // content, when it was mapped against the decoded token text. The strict // selector keeps the whole span mandatory; the scorer window is its tail. PFlashTokenSpan explicit_query_span{-1, -1}; + // Header ("<|im_start|>user\n") opening the last message, when the chat + // markers resolved it — pinned mandatory so a compressed prompt keeps + // the current turn's role envelope. + PFlashTokenSpan last_role_header{-1, -1}; std::vector required_instruction_spans; + // Chat-first scorer query: the last message's content span, located by + // the rendered prompt's own control markers. Feeds the strict tail + // parser and the legacy window; unused when a benchmark parser + // (latest_user) or a marker-less prompt needs the sentinel mapping. + http_detail::PflashChatTailSpan chat_tail; + if (!experiment.configured || tail_parser) { + ChatMarkers chat_markers; + if (resolve_chat_markers(tokenizer_, chat_markers)) { + chat_tail = http_detail::pflash_last_message_content_span( + tokenizer_, chat_markers, *drafter_tokenizer_, drafter_ids); + } + } if (experiment.configured) { if (!messages_input && !raw_text_input) { return "PFlash strict selection input has no parseable text"; @@ -3524,18 +3676,12 @@ std::string HttpServer::apply_pflash_compression( last_user_text = messages[(size_t) last_user_index].content; } + const bool semantic_parser = experiment.query_parser == + luce::pflash::PFlashQueryParser::SemanticUser; int boundary_index = (int) messages.size() - 1; - if (!raw_text_input && - experiment.query_parser == - luce::pflash::PFlashQueryParser::SemanticUser) { + if (!raw_text_input && semantic_parser) { boundary_index = last_user_index; } - if (boundary_index < 0 || - (experiment.query_parser == - luce::pflash::PFlashQueryParser::SemanticUser && - !raw_text_input && last_user_text.empty())) { - return "PFlash strict selection latest-user boundary is unavailable"; - } static constexpr const char * kContentBegin = "__LUCE_PFLASH_CONTENT_BEGIN_02C47F91__"; @@ -3617,15 +3763,36 @@ std::string HttpServer::apply_pflash_compression( }; std::string boundary_error; - if (!map_message_content( - (size_t) boundary_index, - query_content_begin, query_content_end, - boundary_error)) { - return "PFlash strict selection " + boundary_error; + // Chat default: the last message's content bounds come from the + // rendered prompt's control markers — no sentinel re-renders. + if (tail_parser && chat_tail.valid()) { + query_content_begin = chat_tail.content_begin; + query_content_end = chat_tail.content_end; + if (chat_tail.role_begin >= 0 && + chat_tail.role_begin < chat_tail.content_begin) { + last_role_header = {chat_tail.role_begin, + chat_tail.content_begin}; + } + } + if (query_content_begin < 0) { + // Marker-less prompts (and the benchmark's latest_user + // parser) still locate the boundary through sentinel + // renders of the boundary message. + if (boundary_index < 0 || + (semantic_parser && !raw_text_input && + last_user_text.empty())) { + return "PFlash strict selection latest-user boundary is unavailable"; + } + if (!map_message_content( + (size_t) boundary_index, + query_content_begin, query_content_end, + boundary_error)) { + return "PFlash strict selection " + boundary_error; + } } if (query_content_begin < 0 || query_content_end <= query_content_begin || - query_content_end >= (int) drafter_ids.size()) { + query_content_end > (int) drafter_ids.size()) { return "PFlash strict selection content boundary mapping failed"; } @@ -3691,22 +3858,36 @@ std::string HttpServer::apply_pflash_compression( // whole question is mandatory even though the scorer only // consumes its bounded tail. Mapping against the decoded // content text (not a standalone encoding) keeps BPE boundary - // merges like " What" inside the span. - if (!req.pflash_query.empty() && - experiment.query_parser == - luce::pflash::PFlashQueryParser::SemanticUser) { + // merges like " What" inside the span. Benchmark-only: under + // the chat tail parser the query may sit anywhere before the + // closing markers; latest_user still scopes it to the user + // message. + if (!req.pflash_query.empty()) { + const int query_search_begin = + semantic_parser ? query_content_begin : 0; + const int query_search_end = semantic_parser + ? query_content_end + : (query_content_end > 0 + ? query_content_end : (int) drafter_ids.size()); explicit_query_span = http_detail::pflash_decoded_text_span( *drafter_tokenizer_, drafter_ids, - query_content_begin, query_content_end, + query_search_begin, query_search_end, req.pflash_query); if (explicit_query_span.begin < 0) { - return "PFlash strict selection explicit query mapping " - "failed: pflash_query does not occur in the " - "latest user content"; + return semantic_parser + ? "PFlash strict selection explicit query mapping " + "failed: pflash_query does not occur in the " + "latest user content" + : "PFlash strict selection explicit query mapping " + "failed: pflash_query does not occur in the " + "prompt"; } required_instruction_spans.push_back(explicit_query_span); } + if (last_role_header.begin >= 0) { + required_instruction_spans.push_back(last_role_header); + } required_instruction_spans = http_detail::canonicalize_pflash_token_spans( std::move(required_instruction_spans)); @@ -3755,18 +3936,7 @@ std::string HttpServer::apply_pflash_compression( if (!last_user_text.empty()) { semantic_query_ids = drafter_tokenizer_->encode(last_user_text); } - if (experiment.configured && raw_text_input) { - parser_selection_rule = "content_tail"; - query_window = http_detail::pflash_tail_query_window( - drafter_ids, experiment.query_tokens, - query_content_end, query_content_begin); - } else if (experiment.configured && - experiment.query_parser == - luce::pflash::PFlashQueryParser::ArbitraryTail) { - parser_selection_rule = "prompt_tail"; - query_window = http_detail::pflash_tail_query_window( - drafter_ids, experiment.query_tokens, query_content_end); - } else if (explicit_query_span.begin >= 0) { + if (explicit_query_span.begin >= 0) { // The explicit query was already mapped against the decoded content // text and pinned as a mandatory span. The scorer consumes the span's // bounded tail window; the complete span stays in the target prompt. @@ -3778,6 +3948,18 @@ std::string HttpServer::apply_pflash_compression( expected_query_ids.assign( drafter_ids.begin() + (query_window.end - query_window.tokens), drafter_ids.begin() + query_window.end); + } else if (experiment.configured && (raw_text_input || tail_parser)) { + // Chat default: the scorer query is the tail of the last message's + // content — the prompt region before the closing/generation markers, + // clamped so it never swallows the role header. + parser_selection_rule = raw_text_input ? "content_tail" : "prompt_tail"; + query_window = http_detail::pflash_tail_query_window( + drafter_ids, experiment.query_tokens, + query_content_end, query_content_begin); + } else if (!experiment.configured && chat_tail.valid()) { + // Legacy (unconfigured) chat mode uses the same marker-derived tail. + query_window = http_detail::pflash_tail_query_window( + drafter_ids, 8, chat_tail.content_end, chat_tail.content_begin); } else if (!semantic_query_ids.empty()) { if (experiment.configured) parser_selection_rule = "semantic_suffix"; query_window = http_detail::find_pflash_query_window( @@ -4016,7 +4198,13 @@ HttpServer::PreparedPrompt HttpServer::prepare_prompt( is_continuation_request(req.messages); const bool selection_environment = luce::pflash::has_pflash_selection_environment(); - if (should_compress && selection_environment) { + // With a strict-selection environment PFlash owns every turn — + // multi-turn chit-chat compresses the whole rendered history plus + // the current user turn instead of routing to FlowKV. Only the + // per-request FlowKV disk-compression mode still conflicts. + const bool selection_owns_compression = + should_compress && selection_environment; + if (selection_owns_compression) { luce::pflash::PFlashSelectionConfig experiment; std::string experiment_error; if (!luce::pflash::resolve_pflash_selection( @@ -4026,24 +4214,24 @@ HttpServer::PreparedPrompt HttpServer::prepare_prompt( experiment_error; return prepared; } - if (http_detail::pflash_continuation_must_fail_closed( - selection_environment) && - (continuation || req.disk_cache_policy.compress)) { + if (req.disk_cache_policy.compress) { prepared.error_status = 500; prepared.error = - "PFlash strict selection does not support continuation or FlowKV compression"; + "PFlash strict selection does not support FlowKV compression"; return prepared; } } - if (should_compress && continuation && req.messages.is_array()) { + if (should_compress && continuation && req.messages.is_array() && + !selection_owns_compression) { // FlowKV owns continuation compression automatically. Falling // back to whole-prompt compression would destroy the reusable // system/tool prefix anchor, and requiring a separate disk-cache // flag made --prefill-compression auto silently do nothing. apply_flowkv_compression(req, prepared); should_compress = false; - } else if (should_compress && continuation) { + } else if (should_compress && continuation && + !selection_owns_compression) { should_compress = false; std::fprintf(stderr, "[pflash] skip-compress (continuation without messages array)\n"); diff --git a/server/src/server/http_server.h b/server/src/server/http_server.h index 3a86ccd25..35a7127af 100644 --- a/server/src/server/http_server.h +++ b/server/src/server/http_server.h @@ -333,6 +333,32 @@ PFlashTokenSpan pflash_decoded_text_span( int end, const std::string & needle); +// The last chat message's content span inside a rendered prompt, located by +// the model's own chat control markers rather than message bookkeeping. +// ``role_begin`` is the marker opening that message (the header to pin); +// ``content_begin`` skips the role-name line ("<|im_start|>user\n") when the +// family uses generic role markers; ``content_end`` sits before the closing +// or generation marker. Offsets are token indices in ``prompt``'s own +// vocabulary. ``markers`` were resolved on ``marker_tokenizer`` (the target +// model's); its marker strings are searched in the decoded prompt text, so +// a drafter whose vocabulary lacks the control tokens still maps correctly. +// Invalid when the prompt carries no chat markers. +struct PflashChatTailSpan { + int role_begin = -1; + int content_begin = -1; + int content_end = -1; + + bool valid() const { + return content_begin >= 0 && content_end > content_begin; + } +}; + +PflashChatTailSpan pflash_last_message_content_span( + const Tokenizer & marker_tokenizer, + const ChatMarkers & markers, + const Tokenizer & tokenizer, + const std::vector & prompt); + // Return the original prompt offset immediately before the stable trailing // suffix shared with a version whose latest user message carries a sentinel. // Invalid when no such bounded suffix can establish the semantic boundary. @@ -373,8 +399,6 @@ std::vector pflash_document_spans_from_ranges( bool pflash_full_cache_restore_allowed( bool selection_environment_present) noexcept; -bool pflash_continuation_must_fail_closed( - bool selection_environment_present) noexcept; int pflash_target_token_ceiling( int original_target_tokens, double keep_ratio) noexcept; diff --git a/server/test/test_pflash_selection.cpp b/server/test/test_pflash_selection.cpp index 56dbd1a51..07b89a5e9 100644 --- a/server/test/test_pflash_selection.cpp +++ b/server/test/test_pflash_selection.cpp @@ -397,6 +397,11 @@ TEST_CASE(PFlashSelectionFixture, any_selection_environment_is_observable_before TEST_CASE(PFlashSelectionFixture, resolver_selects_explicit_query_parser) { CleanPFlashEnv env; + // Chat-first default: no override selects the rendered tail parser. + const auto fallback = resolve_or_fail(120000, 32); + REQUIRE(fallback.configured == false); + REQUIRE(fallback.query_parser == PFlashQueryParser::ArbitraryTail); + set_env(kQueryParserEnv, "arbitrary_tail"); const auto arbitrary = resolve_or_fail(120000, 32); REQUIRE(arbitrary.configured); diff --git a/server/test/test_server_unit.cpp b/server/test/test_server_unit.cpp index 88cef1ccf..3b0c5aa9d 100644 --- a/server/test/test_server_unit.cpp +++ b/server/test/test_server_unit.cpp @@ -161,6 +161,10 @@ struct HttpServerTestAccess { HttpServer::PreparedPrompt prepared; return server.apply_pflash_compression(req, prepared); } + static HttpServer::PreparedPrompt prepare_prompt( + HttpServer & server, const ParsedRequest & req) { + return server.prepare_prompt(req); + } }; } @@ -459,6 +463,31 @@ static std::string write_pflash_bpe_tokenizer_fixture( return path; } +static std::string write_deepseek_marker_tokenizer_fixture() { + gguf_context * g = gguf_init_empty(); + const char * tokens[] = { + "x", + "<｜begin▁of▁sentence｜>", + "<｜end▁of▁sentence｜>", + "<｜User｜>", + "<｜Assistant｜>", + }; + const uint32_t token_types[] = {1, 3, 3, 3, 3}; + gguf_set_arr_str(g, "tokenizer.ggml.tokens", tokens, + sizeof(tokens) / sizeof(tokens[0])); + gguf_set_arr_data(g, "tokenizer.ggml.token_type", GGUF_TYPE_UINT32, + token_types, sizeof(token_types) / sizeof(token_types[0])); + gguf_set_val_str(g, "tokenizer.ggml.model", "gpt2"); + gguf_set_val_str(g, "tokenizer.ggml.pre", "qwen35"); + gguf_set_val_u32(g, "tokenizer.ggml.bos_token_id", 1); + gguf_set_val_u32(g, "tokenizer.ggml.eos_token_id", 2); + + const std::string path = test_tmp_path("luce_test_deepseek_markers.gguf").string(); + gguf_write_to_file(g, path.c_str(), /*only_meta=*/false); + gguf_free(g); + return path; +} + TEST_CASE(ServerUnitFixture, test_pflash_decoded_span_covers_bpe_merged_first_token) { const std::string content = @@ -789,6 +818,114 @@ TEST_CASE(ServerUnitFixture, test_pflash_tail_query_window) { TEST_ASSERT(!http_detail::pflash_tail_query_window(long_prompt, 128, 201).valid()); } +TEST_CASE(ServerUnitFixture, + test_pflash_last_content_span_stops_before_chat_markers) { + const std::string rendered = + "<|im_start|>system\nYou are helpful.<|im_end|>\n" + "<|im_start|>user\nfirst turn<|im_end|>\n" + "<|im_start|>assistant\nSure.<|im_end|>\n" + "<|im_start|>user\nWhat is the answer?<|im_end|>\n" + "<|im_start|>assistant\n"; + const std::string path = write_pflash_bpe_tokenizer_fixture( + {"What", " is", " the", " answer", "?", "user", "assistant", + "system", "\n", "You", " are", " helpful", ".", "first", " turn", + "Sure"}, + rendered); + Tokenizer tok; + TEST_ASSERT(tok.load_from_gguf(path.c_str())); + + const auto prompt = tok.encode(rendered); + ChatMarkers markers; + TEST_ASSERT(resolve_chat_markers(tok, markers)); + TEST_ASSERT(markers.family == "qwen"); + + const auto span = http_detail::pflash_last_message_content_span( + tok, markers, tok, prompt); + TEST_ASSERT(span.valid()); + TEST_ASSERT(tok.decode({prompt.begin() + span.content_begin, + prompt.begin() + span.content_end}) + == "What is the answer?"); + // The pinned header is exactly the role envelope of the last turn. + TEST_ASSERT(tok.decode({prompt.begin() + span.role_begin, + prompt.begin() + span.content_begin}) + == "<|im_start|>user\n"); + // Everything after the content is template machinery — never part of + // the scorer query. + TEST_ASSERT(tok.decode({prompt.begin() + span.content_end, + prompt.end()}) + == "<|im_end|>\n<|im_start|>assistant\n"); + + const auto window = http_detail::pflash_tail_query_window( + prompt, 8, span.content_end, span.content_begin); + TEST_ASSERT(window.valid()); + TEST_ASSERT(window.end == span.content_end); + TEST_ASSERT(window.tokens == + std::min(8, span.content_end - span.content_begin)); + unlink(path.c_str()); +} + +TEST_CASE(ServerUnitFixture, + test_pflash_last_content_span_deepseek_delimited_roles) { + const std::string path = write_deepseek_marker_tokenizer_fixture(); + Tokenizer tok; + TEST_ASSERT(tok.load_from_gguf(path.c_str())); + + ChatMarkers markers; + TEST_ASSERT(resolve_chat_markers(tok, markers)); + TEST_ASSERT(markers.family == "deepseek"); + TEST_ASSERT(markers.role_starts_delimit); + + // bos + system text + <|User|> + content + <|Assistant|> generation + const std::vector prompt = {1, 0, 3, 0, 0, 4}; + const auto span = http_detail::pflash_last_message_content_span( + tok, markers, tok, prompt); + TEST_ASSERT(span.valid()); + TEST_ASSERT(span.role_begin == 2); + TEST_ASSERT(span.content_begin == 3); + TEST_ASSERT(span.content_end == 5); + TEST_ASSERT(tok.decode({prompt.begin() + span.content_begin, + prompt.begin() + span.content_end}) == "xx"); + remove_test_path(path); +} + +TEST_CASE(ServerUnitFixture, + test_pflash_last_content_span_open_tail_runs_to_prompt_end) { + const std::string rendered = "<|im_start|>user\nhello there"; + const std::string path = write_pflash_bpe_tokenizer_fixture( + {"hello", " there", "user", "\n"}, rendered); + Tokenizer tok; + TEST_ASSERT(tok.load_from_gguf(path.c_str())); + + const auto prompt = tok.encode(rendered); + ChatMarkers markers; + TEST_ASSERT(resolve_chat_markers(tok, markers)); + const auto span = http_detail::pflash_last_message_content_span( + tok, markers, tok, prompt); + TEST_ASSERT(span.valid()); + TEST_ASSERT(span.content_end == (int) prompt.size()); + TEST_ASSERT(tok.decode({prompt.begin() + span.content_begin, + prompt.begin() + span.content_end}) + == "hello there"); + unlink(path.c_str()); +} + +TEST_CASE(ServerUnitFixture, + test_pflash_last_content_span_rejects_markerless_text) { + const std::string rendered = "just some raw text, no chat markers"; + const std::string path = write_pflash_bpe_tokenizer_fixture( + {"just", " some", " raw", " text"}, rendered); + Tokenizer tok; + TEST_ASSERT(tok.load_from_gguf(path.c_str())); + + const auto prompt = tok.encode(rendered); + ChatMarkers markers; + TEST_ASSERT(resolve_chat_markers(tok, markers)); + const auto span = http_detail::pflash_last_message_content_span( + tok, markers, tok, prompt); + TEST_ASSERT(!span.valid()); + unlink(path.c_str()); +} + TEST_CASE(ServerUnitFixture, test_pflash_normalizes_multipart_latest_user_for_reverse_lookup) { ToolMemory tool_memory; const json messages = json::array({ @@ -817,8 +954,6 @@ TEST_CASE(ServerUnitFixture, test_pflash_normalizes_multipart_latest_user_for_re TEST_CASE(ServerUnitFixture, test_pflash_selection_cache_and_continuation_policy) { TEST_ASSERT(http_detail::pflash_full_cache_restore_allowed(false)); TEST_ASSERT(!http_detail::pflash_full_cache_restore_allowed(true)); - TEST_ASSERT(!http_detail::pflash_continuation_must_fail_closed(false)); - TEST_ASSERT(http_detail::pflash_continuation_must_fail_closed(true)); } TEST_CASE(ServerUnitFixture, test_pflash_target_token_ceiling_floors) { @@ -3683,32 +3818,6 @@ TEST_CASE(ServerUnitFixture, test_stop_sequence_holdback_extends) { // Prefix cache hash tests (model-free) // ═══════════════════════════════════════════════════════════════════════ -static std::string write_deepseek_marker_tokenizer_fixture() { - gguf_context * g = gguf_init_empty(); - const char * tokens[] = { - "x", - "<｜begin▁of▁sentence｜>", - "<｜end▁of▁sentence｜>", - "<｜User｜>", - "<｜Assistant｜>", - }; - const uint32_t token_types[] = {1, 3, 3, 3, 3}; - gguf_set_arr_str(g, "tokenizer.ggml.tokens", tokens, - sizeof(tokens) / sizeof(tokens[0])); - gguf_set_arr_data(g, "tokenizer.ggml.token_type", GGUF_TYPE_UINT32, - token_types, - sizeof(token_types) / sizeof(token_types[0])); - gguf_set_val_str(g, "tokenizer.ggml.model", "gpt2"); - gguf_set_val_str(g, "tokenizer.ggml.pre", "qwen35"); - gguf_set_val_u32(g, "tokenizer.ggml.bos_token_id", 1); - gguf_set_val_u32(g, "tokenizer.ggml.eos_token_id", 2); - - const std::string path = test_tmp_path("luce_test_deepseek_markers.gguf").string(); - gguf_write_to_file(g, path.c_str(), /*only_meta=*/false); - gguf_free(g); - return path; -} - TEST_CASE(ServerUnitFixture, test_resolve_deepseek_chat_markers) { const std::string path = write_deepseek_marker_tokenizer_fixture(); Tokenizer tokenizer; @@ -6641,6 +6750,182 @@ TEST_CASE(ServerUnitFixture, test_pflash_default_raw_text_maps_user_query) { unlink(tokenizer_path.c_str()); } +TEST_CASE(ServerUnitFixture, + test_pflash_strict_chat_tail_query_uses_last_message_content) { + luce_test::ScopedEnvVar mode{"PFLASH_SELECT_MODE", "top_p"}; + + const std::string rendered = + "<|im_start|>system\nYou are helpful.<|im_end|>\n" + "<|im_start|>user\nWhat is the answer?<|im_end|>\n" + "<|im_start|>assistant\n"; + const std::string path = write_pflash_bpe_tokenizer_fixture( + {"What", " is", " the", " answer", "?", "user", "assistant", + "system", "\n", "You", " are", " helpful", "."}, + 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; + request.messages = json::array({ + {{"role", "user"}, {"content", "What is the answer?"}}, + }); + 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); + const auto & ids = backend.last_request.input_ids; + const int im_end = tokenizer.token_to_id("<|im_end|>"); + int last_im_end = -1; + for (int i = 0; i < (int) ids.size(); ++i) { + if (ids[i] == im_end) last_im_end = i; + } + TEST_ASSERT(last_im_end > 0); + // The scorer window ends where the user content does — the generation + // markers ("<|im_end|>\n<|im_start|>assistant\n") are never scored. + TEST_ASSERT(backend.last_request.score_query_end == last_im_end); + const int query_begin = backend.last_request.score_query_end - + backend.last_request.score_query_tokens; + TEST_ASSERT(query_begin >= 0); + TEST_ASSERT(tokenizer.decode({ids.begin() + query_begin, + ids.begin() + last_im_end}) + == "What is the answer?"); + // The last turn's role header is pinned mandatory. + bool header_pinned = false; + for (const auto & span : backend.last_request.required_instruction_spans) { + if (tokenizer.decode({ids.begin() + span.begin, + ids.begin() + span.end}) + == "<|im_start|>user\n") { + header_pinned = true; + } + } + TEST_ASSERT(header_pinned); + unlink(path.c_str()); +} + +TEST_CASE(ServerUnitFixture, + test_pflash_strict_selection_owns_chat_continuations) { + luce_test::ScopedEnvVar mode{"PFLASH_SELECT_MODE", "top_p"}; + + const std::string rendered = + "<|im_start|>user\nfirst<|im_end|>\n" + "<|im_start|>assistant\nSure.<|im_end|>\n" + "<|im_start|>user\nsecond question<|im_end|>\n" + "<|im_start|>assistant\n"; + const std::string path = write_pflash_bpe_tokenizer_fixture( + {"first", "second", " question", "user", "assistant", "\n", "Sure", + "."}, + 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_mode = ServerConfig::PflashMode::ALWAYS; + config.pflash_keep_ratio = 1.0f; + config.max_ctx = 8192; + 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; + request.messages = json::array({ + {{"role", "user"}, {"content", "first"}}, + {{"role", "assistant"}, {"content", "Sure."}}, + {{"role", "user"}, {"content", "second question"}}, + }); + request.prompt_tokens = tokenizer.encode(rendered); + + const auto prepared = + HttpServerTestAccess::prepare_prompt(server, request); + TEST_ASSERT_MSG(prepared.error.empty(), prepared.error); + TEST_ASSERT(prepared.compressed); + TEST_ASSERT(!prepared.flowkv); + } + + // Whole-prompt PFlash ran on the multi-turn prompt and scored against + // the last user turn's content tail. + TEST_ASSERT(backend.compress_calls == 1); + const auto & ids = backend.last_request.input_ids; + const int im_end = tokenizer.token_to_id("<|im_end|>"); + int last_im_end = -1; + for (int i = 0; i < (int) ids.size(); ++i) { + if (ids[i] == im_end) last_im_end = i; + } + TEST_ASSERT(backend.last_request.score_query_end == last_im_end); + const int query_begin = backend.last_request.score_query_end - + backend.last_request.score_query_tokens; + TEST_ASSERT(tokenizer.decode({ids.begin() + query_begin, + ids.begin() + last_im_end}) + == "second question"); + unlink(path.c_str()); +} + +TEST_CASE(ServerUnitFixture, + test_pflash_default_continuation_stays_on_flowkv) { + const std::string rendered = + "<|im_start|>user\nfirst<|im_end|>\n" + "<|im_start|>assistant\nSure.<|im_end|>\n" + "<|im_start|>user\nsecond<|im_end|>\n" + "<|im_start|>assistant\n"; + const std::string path = write_pflash_bpe_tokenizer_fixture( + {"first", "second", "user", "assistant", "\n", "Sure", "."}, + 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_mode = ServerConfig::PflashMode::ALWAYS; + config.max_ctx = 8192; + 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; + request.messages = json::array({ + {{"role", "user"}, {"content", "first"}}, + {{"role", "assistant"}, {"content", "Sure."}}, + {{"role", "user"}, {"content", "second"}}, + }); + request.prompt_tokens = tokenizer.encode(rendered); + + const auto prepared = + HttpServerTestAccess::prepare_prompt(server, request); + TEST_ASSERT_MSG(prepared.error.empty(), prepared.error); + } + + // Without a strict-selection environment FlowKV keeps owning + // continuations; whole-prompt PFlash never ran. + TEST_ASSERT(backend.compress_calls == 0); + unlink(path.c_str()); +} + struct MockBatchCompressBackend : MockBackend { int compress_calls = 0; From 5418882a28096d66e68e9a9064c052f8ad3316a1 Mon Sep 17 00:00:00 2001 From: Graffioh Date: Tue, 22 Sep 2026 17:49:41 +0000 Subject: [PATCH 10/26] revert(pflash): drop the document prior and forced document heads PFlash scores the input prompt against the query and nothing else: no document detection, no client pflash_documents ranges, no mass prior and no forced document headers. Reverts 2455efb8c and 739a8c3a6; the decoded-offset helpers they introduced stay because the chat-tail query mapping uses them. Co-Authored-By: Claude Opus 5.5 (1M context) --- server/README.md | 2 - server/src/common/model_backend.h | 3 - .../src/common/pflash_drafter_ipc_daemon.cpp | 3 - server/src/deepseek4/deepseek4_backend.cpp | 3 +- server/src/pflash/pflash_compress.cpp | 39 +-- server/src/pflash/pflash_compress.h | 15 +- server/src/pflash/pflash_drafter.cpp | 11 +- server/src/pflash/pflash_drafter.h | 5 +- server/src/pflash/pflash_selection.cpp | 145 ++------- server/src/pflash/pflash_selection.h | 30 -- server/src/pflash/qwen35_drafter.cpp | 32 +- server/src/pflash/qwen35_drafter.h | 9 +- server/src/qwen35/qwen35_backend.cpp | 3 +- .../src/qwen35/qwen35_layer_split_adapter.cpp | 3 +- server/src/server/http_server.cpp | 115 ------- server/src/server/http_server.h | 55 ---- server/test/test_pflash_selection.cpp | 299 ------------------ 17 files changed, 46 insertions(+), 726 deletions(-) diff --git a/server/README.md b/server/README.md index c428b0aa8..b939762bc 100644 --- a/server/README.md +++ b/server/README.md @@ -377,8 +377,6 @@ the whole request's device footprint. `/status/json` reports | `--prefill-upstream-key ` | none | Bearer token for the upstream. | | `--prefill-upstream-model ` | none | Model name forwarded upstream. | | `PFLASH_SELECT_MODE=top_k` + `PFLASH_SELECT_TOPK ` | budget-only fill | Rank rule: keep the K highest-scoring optional segments in score order instead of filling the keep ratio, with the keep-ratio budget still a hard ceiling (min(K segments, the budget)). Use it where the evidence is compact and sits in the first few ranks -- needle retrieval, passage QA, code -- so a small K reaches it for a fraction of the budget's tokens. Do not use it where the answer needs a whole document identified, since the evidence there spans many segments and K cuts it off. | -| `PFLASH_SELECT_DOC_PRIOR ` | `0` (off) | Document prior: rank an optional segment by `max(0, score) * document_mass_share ^ E` instead of its score alone, where a document's mass is the sum over its segments of `max(0, score) * tokens` normalised by the prompt total. The head identifies the relevant document when its mass is aggregated per document, which per-segment density ranking throws away, so this helps where the answer needs a whole document found among many. `0` leaves the ranking unchanged and it applies in every mode, so it composes with `top_k`. Documents come from the request's `pflash_documents` ranges, or from the prompt's own `Document :` / `[DOC-]` markers; fewer than three documents makes it a no-op, which the compression trace records alongside the exponent. | -| `PFLASH_SELECT_FORCE_DOC_HEADS ` | `0` (off) | Attribution: keep the first segment of each of the D highest-mass documents whatever its own score, charged with the mandatory spans before the fill. A compressed context that keeps a document's text but drops the header identifying it leaves the model unable to cite its source, and a header is short, so it costs far less than the body it names. Headers that no longer fit the budget are dropped rather than fatal, and a structurally required span is always charged first. Uses the same documents as the prior and composes with it and with `top_k`. Note when reading BRIGHT numbers: that benchmark as adapted here is scored by naming a document tag, so part of any gain there is an artefact of the adaptation rather than better evidence -- judge it on body-evidence retention and on sets whose answers live in the text. | With a Qwen3.5-0.8B drafter and strict budget selection (`PFLASH_SELECT_MODE=budget_only`, `PFLASH_SELECT_CHUNK_SIZE`, diff --git a/server/src/common/model_backend.h b/server/src/common/model_backend.h index fcb9de720..53f4ac5a1 100644 --- a/server/src/common/model_backend.h +++ b/server/src/common/model_backend.h @@ -271,9 +271,6 @@ struct ModelBackend { // Role-derived instruction structure in drafter-token coordinates. // Empty is a valid instruction-free or legacy request. std::vector required_instruction_spans; - // Document starts in drafter-token coordinates, for the document - // prior. Empty is a valid single-document or prior-free request. - std::vector document_spans; std::string drafter_path; // GGUF path (for lazy-load) int drafter_gpu = 0; // backend-local GPU for PFlash drafter bool skip_park = false; // true on >=32GB GPUs diff --git a/server/src/common/pflash_drafter_ipc_daemon.cpp b/server/src/common/pflash_drafter_ipc_daemon.cpp index d917366fa..061d262cc 100644 --- a/server/src/common/pflash_drafter_ipc_daemon.cpp +++ b/server/src/common/pflash_drafter_ipc_daemon.cpp @@ -70,9 +70,6 @@ int run_pflash_drafter_ipc_daemon(const char * drafter_path, request.score_query_tokens, /*pool_kernel=*/13, score_query_end, request.required_instruction_spans); - // The IPC compress protocol carries no document spans, so the - // document prior (PFLASH_SELECT_DOC_PRIOR) is inert on this path - // until the wire format gains them. if (compressed.empty()) { std::fprintf(stderr, "[pflash-ipc-daemon] compress returned empty\n"); stream_status(stream_fd, -1); diff --git a/server/src/deepseek4/deepseek4_backend.cpp b/server/src/deepseek4/deepseek4_backend.cpp index 979919be4..473cedc75 100644 --- a/server/src/deepseek4/deepseek4_backend.cpp +++ b/server/src/deepseek4/deepseek4_backend.cpp @@ -3181,8 +3181,7 @@ std::vector DeepSeek4Backend::compress_batch( result.compressed_ids = drafter_score_and_compress( pflash_drafter_ctx_, request.input_ids, request.keep_ratio, /*chunk_size=*/32, request.score_query_tokens, /*pool_kernel=*/13, - score_query_end, request.required_instruction_spans, - request.document_spans); + score_query_end, request.required_instruction_spans); result.ok = !result.compressed_ids.empty(); } diff --git a/server/src/pflash/pflash_compress.cpp b/server/src/pflash/pflash_compress.cpp index 9f815bf9a..30949bb76 100644 --- a/server/src/pflash/pflash_compress.cpp +++ b/server/src/pflash/pflash_compress.cpp @@ -88,18 +88,12 @@ void write_compression_trace( "],\"query_begin\":%d,\"query_end\":%d," "\"selector_mode\":\"%s\",\"query_parser\":\"%s\"," "\"token_budget\":%d,\"top_k\":%d," - "\"doc_prior_exponent\":%.9g,\"documents\":%zu," - "\"doc_prior_applied\":%s," - "\"force_doc_heads\":%d,\"forced_doc_heads\":%d," "\"retained_tokens\":%d", trace_fields->query_begin, trace_fields->query_end, luce::pflash::pflash_selection_mode_name( trace_fields->selector_mode), luce::pflash::pflash_query_parser_name(trace_fields->query_parser), trace_fields->token_budget, trace_fields->top_k, - trace_fields->doc_prior_exponent, trace_fields->documents, - trace_fields->doc_prior_applied ? "true" : "false", - trace_fields->force_doc_heads, trace_fields->forced_doc_heads, trace_fields->retained_tokens); std::fputs(",\"required_instruction_spans\":[", file); if (trace_fields->required_instruction_spans) { @@ -199,8 +193,7 @@ std::vector select_pflash_chunks( const std::vector * segments, bool density, const std::vector * other_token_scores, - double split_fraction, - const std::vector * documents) { + double split_fraction) { const int input_tokens = (int) ids.size(); const int query_end = score_query_end < 0 ? input_tokens : score_query_end; const int query_tokens = std::min(n_lookahead, query_end); @@ -233,16 +226,7 @@ std::vector select_pflash_chunks( luce::pflash::pflash_chunk_is_structurally_required( begin, end, query_begin, query_end, input_tokens, required_instruction_spans); - // The last document starting at or before this candidate; everything - // ahead of the first document start belongs to the first document. - size_t document = 0; - if (documents) { - for (size_t index = 0; index < documents->size(); ++index) { - if ((*documents)[index].begin <= begin) document = index; - else break; - } - } - candidates.push_back({(size_t) chunk, begin, end, score, mandatory, document}); + candidates.push_back({(size_t) chunk, begin, end, score, mandatory}); chunk_means.push_back({(float) score, chunk}); exact_chunk_scores.push_back(score); } @@ -258,17 +242,14 @@ std::vector select_pflash_chunks( score += (*other_token_scores)[(size_t) token]; } score /= (double) std::max(1, candidate.end - candidate.begin); - other_candidates.push_back({candidate.ordinal, candidate.begin, candidate.end, - score, candidate.mandatory, candidate.document}); + other_candidates.push_back({candidate.ordinal, candidate.begin, candidate.end, score, candidate.mandatory}); other_scores.push_back(score); } } const luce::pflash::PFlashSelectionPolicy policy{selector_budget, config.top_p, /*skip_oversized=*/ segments != nullptr, - config.top_k, - config.doc_prior_exponent, - config.force_doc_heads}; + config.top_k}; const auto selected = split ? luce::pflash::select_pflash_split(candidates, other_candidates, policy, split_fraction, config.mode) : luce::pflash::select_pflash_candidates(candidates, policy, config.mode); @@ -308,17 +289,14 @@ std::vector select_pflash_chunks( std::fprintf(stderr, "[pflash-select] selected mode=%s scorer=%s segments=%s score=%s chunk=%d query=%d " - "budget=%d selected_tokens=%zu chunks=%zu/%d stop=%s mass=%.9g " - "docs=%zu doc_prior=%.9g applied=%d heads=%d/%d\n", + "budget=%d selected_tokens=%zu chunks=%zu/%d stop=%s mass=%.9g\n", luce::pflash::pflash_selection_mode_name(config.mode), split ? "split" : "single", segments ? "probe" : "fixed", density ? "density" : "sum", segments ? 0 : config.chunk_size, query_tokens, selector_budget, output.size(), selected.ordinals.size(), n_chunks, luce::pflash::pflash_selection_stop_name(selected.stop), - selected.retained_mass, selected.documents, config.doc_prior_exponent, - (int) selected.doc_prior_applied, selected.forced_doc_heads, - config.force_doc_heads); + selected.retained_mass); std::fflush(stderr); if (write_trace) { @@ -339,11 +317,6 @@ std::vector select_pflash_chunks( strict_fields.other_chunk_scores = split ? &other_scores : nullptr; strict_fields.top_k = config.mode == luce::pflash::PFlashSelectionMode::TopK ? config.top_k : 0; - strict_fields.doc_prior_exponent = config.doc_prior_exponent; - strict_fields.documents = selected.documents; - strict_fields.doc_prior_applied = selected.doc_prior_applied; - strict_fields.force_doc_heads = config.force_doc_heads; - strict_fields.forced_doc_heads = selected.forced_doc_heads; write_compression_trace( input_tokens, keep_ratio, trace_chunk, query_tokens, pool_kernel, n_keep_approx, chunk_means, selected_mask, diff --git a/server/src/pflash/pflash_compress.h b/server/src/pflash/pflash_compress.h index ea10cae00..9fa333852 100644 --- a/server/src/pflash/pflash_compress.h +++ b/server/src/pflash/pflash_compress.h @@ -105,14 +105,6 @@ struct PFlashTraceFields { const std::vector * other_chunk_scores = nullptr; // Rank-mode ceiling: the K that applied, 0 outside top_k mode. int top_k = 0; - // Document prior: the configured exponent, the documents the selector - // saw, and whether the prior actually reweighted the ranking. - double doc_prior_exponent = 0.0; - size_t documents = 0; - bool doc_prior_applied = false; - // Attribution: the configured D and the headers actually forced. - int force_doc_heads = 0; - int forced_doc_heads = 0; }; void write_compression_trace( @@ -142,11 +134,6 @@ std::vector select_pflash_chunks( const std::vector * segments = nullptr, bool density = false, const std::vector * other_token_scores = nullptr, - double split_fraction = 0.0, - // Document starts in prompt-token coordinates, for the document - // prior. A candidate belongs to the last document starting at or - // before its first token. Null or shorter than - // ``kPFlashMinPriorDocuments`` leaves the prior a no-op. - const std::vector * documents = nullptr); + double split_fraction = 0.0); } // namespace luce::common diff --git a/server/src/pflash/pflash_drafter.cpp b/server/src/pflash/pflash_drafter.cpp index d3af756a0..5fd2f2d5f 100644 --- a/server/src/pflash/pflash_drafter.cpp +++ b/server/src/pflash/pflash_drafter.cpp @@ -104,8 +104,7 @@ std::vector drafter_score_and_compress( int n_lookahead, int pool_kernel, int score_query_end, - const std::vector & required_instruction_spans, - const std::vector & document_spans) { + const std::vector & required_instruction_spans) { if (!ctx.loaded) { set_last_error("drafter not loaded"); return {}; @@ -146,13 +145,12 @@ std::vector drafter_score_and_compress( std::fprintf(stderr, "[pflash-select] config mode=%s active=%d chunk=%d " "query_parser=%s query_cap=%d query_actual=%d top_p=%.9g " - "top_k=%d doc_prior=%.9g doc_heads=%d doc_spans=%zu input=%zu\n", + "top_k=%d input=%zu\n", luce::pflash::pflash_selection_mode_name(experiment.mode), (int) experiment.selection_active, experiment.chunk_size, luce::pflash::pflash_query_parser_name(experiment.query_parser), experiment.query_tokens, n_lookahead, experiment.top_p, - experiment.top_k, experiment.doc_prior_exponent, - experiment.force_doc_heads, document_spans.size(), ids.size()); + experiment.top_k, ids.size()); std::fflush(stderr); } if (score_query_end < 0) { @@ -161,8 +159,7 @@ std::vector drafter_score_and_compress( } return qwen35_drafter_score_and_compress( ctx, ids, keep_ratio, chunk_size, n_lookahead, pool_kernel, - score_query_end, experiment, required_instruction_spans, - document_spans); + score_query_end, experiment, required_instruction_spans); } } // namespace luce::common diff --git a/server/src/pflash/pflash_drafter.h b/server/src/pflash/pflash_drafter.h index 2ace1b745..8340114b3 100644 --- a/server/src/pflash/pflash_drafter.h +++ b/server/src/pflash/pflash_drafter.h @@ -82,9 +82,6 @@ std::vector drafter_score_and_compress( int pool_kernel = 13, int score_query_end = -1, const std::vector & - required_instruction_spans = {}, - // Document starts in prompt-token coordinates, for the document prior - // (PFLASH_SELECT_DOC_PRIOR). Empty leaves the prior a no-op. - const std::vector & document_spans = {}); + required_instruction_spans = {}); } // namespace luce::common diff --git a/server/src/pflash/pflash_selection.cpp b/server/src/pflash/pflash_selection.cpp index 5007992b8..03a20ce9d 100644 --- a/server/src/pflash/pflash_selection.cpp +++ b/server/src/pflash/pflash_selection.cpp @@ -21,8 +21,6 @@ constexpr const char * kQueryEnv = "PFLASH_SELECT_QUERY_TOKENS"; constexpr const char * kQueryParserEnv = "PFLASH_SELECT_QUERY_PARSER"; constexpr const char * kTopPEnv = "PFLASH_SELECT_TOP_P"; constexpr const char * kTopKEnv = "PFLASH_SELECT_TOPK"; -constexpr const char * kDocPriorEnv = "PFLASH_SELECT_DOC_PRIOR"; -constexpr const char * kDocHeadsEnv = "PFLASH_SELECT_FORCE_DOC_HEADS"; constexpr const char * kSegmentsEnv = "PFLASH_SELECT_SEGMENTS"; constexpr const char * kSelectEnv = "PFLASH_SELECT_SCORE"; constexpr const char * kScorerEnv = "PFLASH_SELECT_SCORER"; @@ -75,8 +73,6 @@ bool has_pflash_selection_environment() noexcept { std::getenv(kQueryParserEnv) != nullptr || std::getenv(kTopPEnv) != nullptr || std::getenv(kTopKEnv) != nullptr || - std::getenv(kDocPriorEnv) != nullptr || - std::getenv(kDocHeadsEnv) != nullptr || std::getenv(kSegmentsEnv) != nullptr || std::getenv(kSelectEnv) != nullptr || std::getenv(kScorerEnv) != nullptr || @@ -151,14 +147,6 @@ PFlashSelectionResult select_pflash_candidates( if (mode == PFlashSelectionMode::TopK && policy.top_k <= 0) { return invalid_result("PFlash top_k must be positive"); } - if (policy.force_doc_heads < 0) { - return invalid_result("PFlash forced document head count must not be negative"); - } - if (!std::isfinite(policy.doc_prior_exponent) || - policy.doc_prior_exponent < 0.0) { - return invalid_result( - "PFlash document prior exponent must be finite and non-negative"); - } std::vector source_ranges; source_ranges.reserve(candidates.size()); @@ -196,110 +184,28 @@ PFlashSelectionResult select_pflash_candidates( std::vector selected_candidates; selected_candidates.reserve(candidates.size()); - // Per-document mass: the sum over a document's candidates -- mandatory - // ones included, since they are part of the prompt -- of - // max(0, score) * tokens. It feeds both document rules below. - std::vector> document_mass; - for (const auto & candidate : candidates) { - const double mass = std::max(0.0, candidate.score) * - (double) (candidate.end - candidate.begin); - auto it = std::find_if(document_mass.begin(), document_mass.end(), - [&](const auto & entry) { return entry.first == candidate.document; }); - if (it == document_mass.end()) { - document_mass.push_back({candidate.document, mass}); - } else { - it->second += mass; - } - } - double total_mass = 0.0; - for (const auto & entry : document_mass) total_mass += entry.second; - result.documents = document_mass.size(); - const bool enough_documents = - document_mass.size() >= kPFlashMinPriorDocuments && total_mass > 0.0; - // Document prior: ranking by score * share^exponent lets a document the - // head likes as a whole lift its own segments, which per-segment density - // ranking throws away. - result.doc_prior_applied = policy.doc_prior_exponent > 0.0 && enough_documents; - - // Forced document headers: the first candidate of each of the D - // highest-mass documents joins the mandatory set. Attribution, not - // evidence -- a compressed context that drops the header identifying a - // document leaves the model unable to cite what it is quoting, and the - // header is short, so it costs far less than the body it names. - std::vector forced_heads; - if (policy.force_doc_heads > 0 && enough_documents) { - std::vector> ranked = document_mass; - std::sort(ranked.begin(), ranked.end(), - [](const auto & left, const auto & right) { - if (left.second != right.second) return left.second > right.second; - return left.first < right.first; - }); - const size_t wanted = std::min( - (size_t) policy.force_doc_heads, ranked.size()); - for (size_t index = 0; index < wanted; ++index) { - const PFlashSelectionCandidate * head = nullptr; - for (const auto & candidate : candidates) { - if (candidate.document != ranked[index].first) continue; - if (!head || candidate.begin < head->begin) head = &candidate; - } - // Already-mandatory headers are kept by the rule above anyway. - if (head && !head->mandatory) forced_heads.push_back(head); - } - } - // True mandatory candidates are charged first, so a forced header can - // never push a structurally required span out of the budget. - for (const auto & candidate : candidates) { - if (!candidate.mandatory) continue; - const int length = candidate.end - candidate.begin; - if (length > policy.token_budget - result.retained_tokens) { - const size_t documents_seen = result.documents; - result = {}; - result.documents = documents_seen; - result.stop = PFlashSelectionStop::MandatoryQueryExceedsBudget; - result.error = "mandatory PFlash retention tokens exceed the token budget"; - return result; - } - selected_candidates.push_back(&candidate); - result.retained_tokens += length; - } - // Then the forced headers, best document first. One that no longer fits - // is dropped rather than fatal: attribution yields to the budget, and it - // stays eligible for the ordinary fill below. - std::vector kept_heads; - for (const auto * head : forced_heads) { - const int length = head->end - head->begin; - if (length > policy.token_budget - result.retained_tokens) continue; - selected_candidates.push_back(head); - result.retained_tokens += length; - kept_heads.push_back(head->ordinal); - } - result.forced_doc_heads = (int) kept_heads.size(); - const auto head_was_kept = [&](size_t ordinal) { - return std::find(kept_heads.begin(), kept_heads.end(), ordinal) != - kept_heads.end(); - }; - std::vector optional; optional.reserve(candidates.size()); for (const auto & candidate : candidates) { - if (candidate.mandatory || head_was_kept(candidate.ordinal)) continue; - optional.push_back(&candidate); - } - - const auto rank_score = [&](const PFlashSelectionCandidate * candidate) { - const double base = std::max(0.0, candidate->score); - if (!result.doc_prior_applied) return base; - double share = 0.0; - for (const auto & entry : document_mass) { - if (entry.first == candidate->document) { share = entry.second; break; } + if (candidate.mandatory) { + const int length = candidate.end - candidate.begin; + if (length > policy.token_budget - result.retained_tokens) { + result = {}; + result.stop = PFlashSelectionStop::MandatoryQueryExceedsBudget; + result.error = "mandatory PFlash retention tokens exceed the token budget"; + return result; + } + selected_candidates.push_back(&candidate); + result.retained_tokens += length; + } else { + optional.push_back(&candidate); } - return base * std::pow(share / total_mass, policy.doc_prior_exponent); - }; + } std::sort(optional.begin(), optional.end(), - [&](const auto * left, const auto * right) { - const double left_score = rank_score(left); - const double right_score = rank_score(right); + [](const auto * left, const auto * right) { + const double left_score = std::max(0.0, left->score); + const double right_score = std::max(0.0, right->score); if (left_score != right_score) return left_score > right_score; return left->ordinal < right->ordinal; }); @@ -410,8 +316,6 @@ bool resolve_pflash_selection( const char * query_parser_raw = std::getenv(kQueryParserEnv); const char * top_p_raw = std::getenv(kTopPEnv); const char * top_k_raw = std::getenv(kTopKEnv); - const char * doc_prior_raw = std::getenv(kDocPriorEnv); - const char * doc_heads_raw = std::getenv(kDocHeadsEnv); const char * segments_raw = std::getenv(kSegmentsEnv); const char * select_raw = std::getenv(kSelectEnv); const char * scorer_raw = std::getenv(kScorerEnv); @@ -419,8 +323,8 @@ bool resolve_pflash_selection( PFlashSelectionConfig config; config.configured = mode_raw || chunk_raw || query_raw || - query_parser_raw || top_p_raw || top_k_raw || doc_prior_raw || - doc_heads_raw || segments_raw || select_raw || scorer_raw || split_raw; + query_parser_raw || top_p_raw || top_k_raw || segments_raw || + select_raw || scorer_raw || split_raw; if (scorer_raw) { if (std::strcmp(scorer_raw, "head") == 0) { config.scorer = PFlashScorer::Head; @@ -519,19 +423,6 @@ bool resolve_pflash_selection( error = std::string(kTopKEnv) + " must be a positive integer"; return false; } - if (doc_prior_raw && - (!parse_double(doc_prior_raw, config.doc_prior_exponent) || - config.doc_prior_exponent < 0.0)) { - error = std::string(kDocPriorEnv) + - " must be a non-negative number"; - return false; - } - if (doc_heads_raw && - (!parse_int(doc_heads_raw, config.force_doc_heads) || - config.force_doc_heads < 0)) { - error = std::string(kDocHeadsEnv) + " must be a non-negative integer"; - return false; - } if (config.mode == PFlashSelectionMode::TopK && config.top_k <= 0) { error = std::string(kTopKEnv) + " is required when " + std::string(kModeEnv) + " is top_k"; diff --git a/server/src/pflash/pflash_selection.h b/server/src/pflash/pflash_selection.h index 084b6268a..67ea85a56 100644 --- a/server/src/pflash/pflash_selection.h +++ b/server/src/pflash/pflash_selection.h @@ -37,9 +37,6 @@ struct PFlashSelectionCandidate { int end = 0; double score = 0.0; bool mandatory = false; - // Which document this candidate falls in, for the document prior below. - // All-zero (one document) leaves the prior a no-op. - size_t document = 0; }; struct PFlashSelectionPolicy { @@ -52,26 +49,8 @@ struct PFlashSelectionPolicy { // TopK mode only: how many optional candidates to keep. Must be positive // in that mode and is ignored in the others. int top_k = 0; - // Document prior: an optional candidate is ranked by - // ``max(0, score) * document_mass_share ^ doc_prior_exponent``, where a - // document's mass is the sum over its candidates of - // ``max(0, score) * tokens`` normalised by the prompt total. 0 disables - // it and leaves the ranking byte-identical. It applies in every mode, so - // it composes with TopK. Fewer than ``kPFlashMinPriorDocuments`` distinct - // documents makes it a no-op: with one or two documents the shares carry - // no ranking information worth a reweight. - double doc_prior_exponent = 0.0; - // Attribution: force the first candidate of each of the D highest-mass - // documents into the mandatory set before the fill, so a compressed - // context never drops the header that identifies a document it quotes. - // 0 disables it. It composes with the prior and with TopK, and the forced - // headers are charged against the budget like any mandatory candidate. - int force_doc_heads = 0; }; -// Below this many distinct documents the document prior is a no-op. -constexpr size_t kPFlashMinPriorDocuments = 3; - struct PFlashSelectionResult { bool ok = false; std::vector ordinals; @@ -79,13 +58,6 @@ struct PFlashSelectionResult { double retained_mass = 0.0; PFlashSelectionStop stop = PFlashSelectionStop::InvalidInput; std::string error; - // Distinct documents seen among the candidates, and whether the document - // prior actually reweighted the ranking (exponent > 0 and enough - // documents). Both are recorded in the compression trace. - size_t documents = 0; - bool doc_prior_applied = false; - // Document headers promoted to mandatory by ``force_doc_heads``. - int forced_doc_heads = 0; }; bool pflash_chunk_is_structurally_required( @@ -130,8 +102,6 @@ struct PFlashSelectionConfig { int query_tokens = 8; double top_p = 0.95; int top_k = 0; - double doc_prior_exponent = 0.0; - int force_doc_heads = 0; PFlashSegmentation segmentation = PFlashSegmentation::Auto; PFlashCandidateScore candidate_score = PFlashCandidateScore::Auto; PFlashScorer scorer = PFlashScorer::Head; diff --git a/server/src/pflash/qwen35_drafter.cpp b/server/src/pflash/qwen35_drafter.cpp index f8d8c4fdf..1dbd62979 100644 --- a/server/src/pflash/qwen35_drafter.cpp +++ b/server/src/pflash/qwen35_drafter.cpp @@ -100,8 +100,7 @@ std::vector qwen35_score_and_compress( int score_query_end, const luce::pflash::PFlashSelectionConfig & experiment, const std::vector & required_instruction_spans, - std::vector * token_scores_out, - const std::vector * document_spans) { + std::vector * token_scores_out) { const int S = (int)ids.size(); const int hidden = w.n_embd; @@ -380,10 +379,7 @@ std::vector qwen35_score_and_compress( if (experiment.selection_active) { return select_pflash_chunks( ids, smooth_score, keep_ratio, n_lookahead, score_query_end, - pk, experiment, required_instruction_spans, false, true, - /*segments=*/nullptr, /*density=*/false, - /*other_token_scores=*/nullptr, /*split_fraction=*/0.0, - document_spans); + pk, experiment, required_instruction_spans, false, true); } std::vector> chunk_means; @@ -537,8 +533,7 @@ std::vector qwen35_strict_score_and_compress( const std::vector & required_instruction_spans, std::vector * token_mass_out, std::vector * segments_out, - bool * density_out, - const std::vector * document_spans) { + bool * density_out) { TargetWeights & w = st.weights; const int S = (int)ids.size(); @@ -924,9 +919,7 @@ std::vector qwen35_strict_score_and_compress( ids, token_mass, keep_ratio, n_lookahead, score_query_end, /*pool_kernel=*/1, experiment, required_instruction_spans, /*direct_mass=*/true, /*write_trace=*/true, - segments.empty() ? nullptr : &segments, density, - /*other_token_scores=*/nullptr, /*split_fraction=*/0.0, - document_spans); + segments.empty() ? nullptr : &segments, density); } std::vector qwen35_drafter_score_and_compress( @@ -938,8 +931,7 @@ std::vector qwen35_drafter_score_and_compress( int pool_kernel, int score_query_end, const luce::pflash::PFlashSelectionConfig & experiment, - const std::vector & required_instruction_spans, - const std::vector & document_spans) { + const std::vector & required_instruction_spans) { if (!ctx.state) { set_last_error("qwen35 drafter state missing"); return {}; @@ -961,14 +953,14 @@ std::vector qwen35_drafter_score_and_compress( if (qwen35_strict_score_and_compress( *st, ids, keep_ratio, n_lookahead, score_query_end, experiment, required_instruction_spans, &head_mass, &head_segments, - &head_density, &document_spans).empty()) { + &head_density).empty()) { return {}; } std::vector other_scores; if (qwen35_score_and_compress(st->weights, ids, keep_ratio, chunk_size, n_lookahead, pool_kernel, score_query_end, experiment, required_instruction_spans, - &other_scores, &document_spans).empty()) { + &other_scores).empty()) { return {}; } if (other_scores.size() != head_mass.size()) { @@ -985,14 +977,12 @@ std::vector qwen35_drafter_score_and_compress( /*pool_kernel=*/1, experiment, required_instruction_spans, /*direct_mass=*/true, /*write_trace=*/true, head_segments.empty() ? nullptr : &head_segments, head_density, - &other_scores, experiment.split_fraction, &document_spans); + &other_scores, experiment.split_fraction); } if (experiment.selection_active && !force_legacy) { return qwen35_strict_score_and_compress( *st, ids, keep_ratio, n_lookahead, score_query_end, experiment, - required_instruction_spans, /*token_mass_out=*/nullptr, - /*segments_out=*/nullptr, /*density_out=*/nullptr, - &document_spans); + required_instruction_spans); } if (st->head_loaded && !experiment.selection_active) { set_last_error("Qwen3.5 scoring head requires strict selection"); @@ -1001,9 +991,7 @@ std::vector qwen35_drafter_score_and_compress( return qwen35_score_and_compress(st->weights, ids, keep_ratio, chunk_size, n_lookahead, pool_kernel, score_query_end, experiment, - required_instruction_spans, - /*token_scores_out=*/nullptr, - &document_spans); + required_instruction_spans); } } // namespace luce::common diff --git a/server/src/pflash/qwen35_drafter.h b/server/src/pflash/qwen35_drafter.h index 3e6f25979..c39e730f6 100644 --- a/server/src/pflash/qwen35_drafter.h +++ b/server/src/pflash/qwen35_drafter.h @@ -76,8 +76,7 @@ std::vector qwen35_score_and_compress( int score_query_end, const luce::pflash::PFlashSelectionConfig & experiment, const std::vector & required_instruction_spans, - std::vector * token_scores_out = nullptr, - const std::vector * document_spans = nullptr); + std::vector * token_scores_out = nullptr); // The block-15 scoring head under strict budget selection. std::vector qwen35_strict_score_and_compress( @@ -90,8 +89,7 @@ std::vector qwen35_strict_score_and_compress( const std::vector & required_instruction_spans, std::vector * token_mass_out = nullptr, std::vector * segments_out = nullptr, - bool * density_out = nullptr, - const std::vector * document_spans = nullptr); + bool * density_out = nullptr); // Arch dispatch target of drafter_score_and_compress. std::vector qwen35_drafter_score_and_compress( @@ -103,7 +101,6 @@ std::vector qwen35_drafter_score_and_compress( int pool_kernel, int score_query_end, const luce::pflash::PFlashSelectionConfig & experiment, - const std::vector & required_instruction_spans, - const std::vector & document_spans = {}); + const std::vector & required_instruction_spans); } // namespace luce::common diff --git a/server/src/qwen35/qwen35_backend.cpp b/server/src/qwen35/qwen35_backend.cpp index cd2a05699..c7f4b8731 100644 --- a/server/src/qwen35/qwen35_backend.cpp +++ b/server/src/qwen35/qwen35_backend.cpp @@ -1201,8 +1201,7 @@ std::vector Qwen35Backend::compress_batch( result.compressed_ids = drafter_score_and_compress( drafter_ctx_, request.input_ids, request.keep_ratio, /*chunk_size=*/32, request.score_query_tokens, /*pool_kernel=*/13, - score_query_end, request.required_instruction_spans, - request.document_spans); + score_query_end, request.required_instruction_spans); result.ok = !result.compressed_ids.empty(); if (result.ok) { std::fprintf(stderr, "[compress] %zu -> %zu tokens\n", diff --git a/server/src/qwen35/qwen35_layer_split_adapter.cpp b/server/src/qwen35/qwen35_layer_split_adapter.cpp index a18a5c04b..7a4a04aa0 100644 --- a/server/src/qwen35/qwen35_layer_split_adapter.cpp +++ b/server/src/qwen35/qwen35_layer_split_adapter.cpp @@ -1392,8 +1392,7 @@ Qwen35LayerSplitAdapter::compress(const ModelBackend::CompressRequest & req) { result.compressed_ids = drafter_score_and_compress( pflash_drafter_, req.input_ids, req.keep_ratio, /*chunk_size=*/32, req.score_query_tokens, /*pool_kernel=*/13, - score_query_end, req.required_instruction_spans, - req.document_spans); + score_query_end, req.required_instruction_spans); result.ok = !result.compressed_ids.empty(); if (result.ok) { std::fprintf(stderr, "[target-split][compress] %zu -> %zu tokens\n", diff --git a/server/src/server/http_server.cpp b/server/src/server/http_server.cpp index 1bedce34c..554c55fbb 100644 --- a/server/src/server/http_server.cpp +++ b/server/src/server/http_server.cpp @@ -402,102 +402,8 @@ int token_at_offset(const DecodedPrompt & decoded, size_t at) { return (int) (upper - decoded.token_begin.begin() - 1); } -bool is_digit_run(const std::string & text, size_t at, size_t & after) { - after = at; - while (after < text.size() && text[after] >= '0' && text[after] <= '9') ++after; - return after > at; -} - -// Character offsets where a document marker opens: a line beginning -// "Document :" or a "[DOC-]" tag anywhere in the line. -std::vector document_marker_offsets(const std::string & text) { - static const std::string kLabel = "Document "; - static const std::string kTag = "[DOC-"; - std::vector offsets; - for (size_t at = text.find(kLabel); at != std::string::npos; - at = text.find(kLabel, at + 1)) { - if (at != 0 && text[at - 1] != '\n') continue; - size_t after = 0; - if (!is_digit_run(text, at + kLabel.size(), after)) continue; - if (after >= text.size() || text[after] != ':') continue; - offsets.push_back(at); - } - for (size_t at = text.find(kTag); at != std::string::npos; - at = text.find(kTag, at + 1)) { - size_t after = 0; - if (!is_digit_run(text, at + kTag.size(), after)) continue; - if (after >= text.size() || text[after] != ']') continue; - offsets.push_back(at); - } - std::sort(offsets.begin(), offsets.end()); - offsets.erase(std::unique(offsets.begin(), offsets.end()), offsets.end()); - return offsets; -} - -// Tile [0, tokens) from ascending document start tokens. Starts before the -// first marker stay with the first document, so every token has a document. -std::vector tile_document_starts( - std::vector starts, int tokens) { - std::vector spans; - std::sort(starts.begin(), starts.end()); - starts.erase(std::unique(starts.begin(), starts.end()), starts.end()); - while (!starts.empty() && starts.back() >= tokens) starts.pop_back(); - if ((size_t) starts.size() < luce::pflash::kPFlashMinPriorDocuments) { - return spans; - } - if (starts.front() != 0) starts.insert(starts.begin(), 0); - for (size_t index = 0; index < starts.size(); ++index) { - const int end = index + 1 < starts.size() ? starts[index + 1] : tokens; - spans.push_back({starts[index], end}); - } - return spans; -} - } // namespace -std::vector pflash_detect_document_spans( - const Tokenizer & tokenizer, - const std::vector & prompt) { - if (prompt.empty()) return {}; - const DecodedPrompt decoded = decode_prompt_with_offsets(tokenizer, prompt); - const auto offsets = document_marker_offsets(decoded.text); - if (offsets.size() < luce::pflash::kPFlashMinPriorDocuments) return {}; - std::vector starts; - starts.reserve(offsets.size()); - for (size_t offset : offsets) { - starts.push_back(token_at_offset(decoded, offset)); - } - return tile_document_starts(std::move(starts), (int) prompt.size()); -} - -std::vector pflash_document_spans_from_ranges( - const Tokenizer & tokenizer, - const std::vector & prompt, - const std::vector> & ranges) { - if (prompt.empty() || - ranges.size() < luce::pflash::kPFlashMinPriorDocuments) { - return {}; - } - const int tokens = (int) prompt.size(); - bool token_coordinates = true; - for (const auto & range : ranges) { - if (range.second > tokens) { token_coordinates = false; break; } - } - std::vector starts; - starts.reserve(ranges.size()); - if (token_coordinates) { - for (const auto & range : ranges) starts.push_back(range.first); - } else { - const DecodedPrompt decoded = - decode_prompt_with_offsets(tokenizer, prompt); - for (const auto & range : ranges) { - if ((size_t) range.first >= decoded.text.size()) continue; - starts.push_back(token_at_offset(decoded, (size_t) range.first)); - } - } - return tile_document_starts(std::move(starts), tokens); -} - std::vector canonicalize_pflash_token_spans( std::vector spans) { std::sort(spans.begin(), spans.end(), [] ( @@ -2910,7 +2816,6 @@ bool HttpServer::handle_model_request(SocketHandle fd, ParsedRequest & req, req.session_id = parse_session_id_from_body(body); req.pflash_query = parse_pflash_query_from_body(body); req.pflash_required = parse_pflash_required_from_body(body); - req.pflash_documents = parse_pflash_documents_from_body(body); // PPP rearrange (optional): peel ephemeral system banners into a // following system message so the first chat boundary is stable. @@ -3976,30 +3881,10 @@ std::string HttpServer::apply_pflash_compression( } } - // Document prior input: client-declared ranges when the request carries - // them, otherwise the served prompt's own document markers. Either way - // fewer than three documents yields no spans and the prior stays off. - std::vector document_spans; - if ((experiment.doc_prior_exponent > 0.0 || - experiment.force_doc_heads > 0) && drafter_tokenizer_) { - document_spans = req.pflash_documents.empty() - ? http_detail::pflash_detect_document_spans( - *drafter_tokenizer_, drafter_ids) - : http_detail::pflash_document_spans_from_ranges( - *drafter_tokenizer_, drafter_ids, req.pflash_documents); - std::fprintf(stderr, - "[pflash-docs] source=%s documents=%zu exponent=%.9g heads=%d\n", - req.pflash_documents.empty() ? "detected" : "request", - document_spans.size(), experiment.doc_prior_exponent, - experiment.force_doc_heads); - std::fflush(stderr); - } - ModelBackend::CompressRequest compress_request; compress_request.input_ids = std::move(drafter_ids); compress_request.required_instruction_spans = std::move(required_instruction_spans); - compress_request.document_spans = std::move(document_spans); compress_request.keep_ratio = http_detail::resolve_pflash_keep_ratio( pflash_keep_ratio(config_, prompt_tokens), req.session_id, sessions_); if (query_window.valid()) { diff --git a/server/src/server/http_server.h b/server/src/server/http_server.h index 35a7127af..0e96f1dd5 100644 --- a/server/src/server/http_server.h +++ b/server/src/server/http_server.h @@ -376,27 +376,6 @@ int pflash_query_search_begin_from_sentinel( std::string pflash_token_fingerprint( const std::vector & ids); -// Document starts detected in the served prompt, in prompt-token -// coordinates. A line matching ``Document :`` or a ``[DOC-]`` tag opens -// a document; the returned spans tile [0, prompt.size()) so every token -// belongs to exactly one document. Fewer than -// ``luce::pflash::kPFlashMinPriorDocuments`` starts returns empty, which -// leaves the document prior a no-op -- with one or two documents the mass -// shares carry no ranking information worth a reweight. -std::vector pflash_detect_document_spans( - const Tokenizer & tokenizer, - const std::vector & prompt); - -// Turn client-declared ``pflash_documents`` ranges into document spans. -// Ranges are token indices when every value fits inside the prompt's token -// count, and character offsets into the decoded prompt text otherwise. The -// result tiles [0, prompt.size()) from the range starts, exactly like the -// detected form. Returns empty when the ranges are unusable or too few. -std::vector pflash_document_spans_from_ranges( - const Tokenizer & tokenizer, - const std::vector & prompt, - const std::vector> & ranges); - bool pflash_full_cache_restore_allowed( bool selection_environment_present) noexcept; int pflash_target_token_ceiling( @@ -452,9 +431,6 @@ struct ParsedRequest { // compression (e.g. an answer-format directive embedded in a user // message). Each occurrence is mapped and retained as a mandatory span. std::vector pflash_required; - // Client-declared document ranges for the document prior. Empty falls - // back to detecting the served prompt's document markers. - std::vector> pflash_documents; DiskPrefixCachePolicy disk_cache_policy; // PPP: stable pin cut for tool-heavy requests (0 = use default boundary). int pin_end_token = 0; @@ -856,37 +832,6 @@ inline std::vector parse_pflash_required_from_body(const json & bod return result; } -// PFlash: client-declared document ranges for the document prior, as a list -// of two-element [begin, end] arrays. Accepted at the top level or under -// extra_body, like pflash_query. Absent means the runtime detects documents -// from the served prompt's markers instead. -inline std::vector> parse_pflash_documents_from_body( - const json & body) { - const json * field = nullptr; - if (body.contains("extra_body")) { - const auto & eb = body["extra_body"]; - if (eb.is_object() && eb.contains("pflash_documents") && - eb["pflash_documents"].is_array()) { - field = &eb["pflash_documents"]; - } - } - if (!field && body.contains("pflash_documents") && - body["pflash_documents"].is_array()) { - field = &body["pflash_documents"]; - } - std::vector> result; - if (!field) return result; - for (const auto & entry : *field) { - if (!entry.is_array() || entry.size() != 2) continue; - if (!entry[0].is_number_integer() || !entry[1].is_number_integer()) continue; - const int begin = entry[0].get(); - const int end = entry[1].get(); - if (begin < 0 || end <= begin) continue; - result.push_back({begin, end}); - } - return result; -} - inline std::string parse_session_id_from_body(const json & body) { if (body.contains("extra_body")) { const auto & eb = body["extra_body"]; diff --git a/server/test/test_pflash_selection.cpp b/server/test/test_pflash_selection.cpp index 07b89a5e9..ccdb2189f 100644 --- a/server/test/test_pflash_selection.cpp +++ b/server/test/test_pflash_selection.cpp @@ -863,302 +863,3 @@ TEST_CASE(PFlashSelectionFixture, split_selection_rejects_top_k) { REQUIRE(!result.ok); REQUIRE(result.stop == PFlashSelectionStop::InvalidInput); } - -namespace { - -// Three documents of two candidates each. Document 2 (ordinals 4,5) has the -// lowest per-segment scores but the most mass; document 0 holds the single -// highest-scoring segment. Per-segment ranking prefers ordinal 0, the -// document prior prefers document 2. -std::vector document_candidates() { - const double scores[6] = {9.0, 1.0, 2.0, 2.0, 5.0, 5.0}; - const size_t documents[6] = {0, 0, 1, 1, 2, 2}; - std::vector candidates; - for (size_t index = 0; index < 6; ++index) { - const int begin = (int) index * 100; - PFlashSelectionCandidate c = candidate(index, begin, begin + 100, scores[index]); - c.document = documents[index]; - candidates.push_back(c); - } - return candidates; -} - -} // namespace - -TEST_CASE(PFlashSelectionFixture, document_prior_reweights_by_document_mass_share) { - const auto candidates = document_candidates(); - // Masses: doc0 = (9+1)*100 = 1000, doc1 = 400, doc2 = 1000; total 2400. - // Shares: 0.41667, 0.16667, 0.41667. - const auto plain = select_pflash_candidates( - candidates, PFlashSelectionPolicy{200, 0.95, false}, - PFlashSelectionMode::BudgetOnly); - REQUIRE(plain.ok); - REQUIRE(!plain.doc_prior_applied); - REQUIRE(plain.documents == 3); - // Without the prior the two highest raw scores win: ordinal 0 (9) and - // one of the 5s. - require_ordinals(plain, {0, 4}); - - const auto prior = select_pflash_candidates( - candidates, PFlashSelectionPolicy{200, 0.95, false, 0, 1.0}, - PFlashSelectionMode::BudgetOnly); - REQUIRE(prior.ok); - REQUIRE(prior.doc_prior_applied); - REQUIRE(prior.documents == 3); - // 9 * 0.41667 = 3.75 still beats 5 * 0.41667 = 2.083? No: 2.083 < 3.75, - // so ordinal 0 stays first; ordinal 4 (2.083) beats ordinal 1 - // (1 * 0.41667 = 0.4167) and ordinal 2 (2 * 0.16667 = 0.333). - require_ordinals(prior, {0, 4}); - - // A larger exponent sharpens the shares; the weak document falls further - // behind, which is the whole point of the knob. - const auto sharp = select_pflash_candidates( - candidates, PFlashSelectionPolicy{400, 0.95, false, 0, 3.0}, - PFlashSelectionMode::BudgetOnly); - REQUIRE(sharp.ok); - REQUIRE(sharp.doc_prior_applied); - // doc1's segments (share 0.1667^3 = 0.00463) rank below every segment of - // doc0 and doc2 (share 0.41667^3 = 0.0723). - require_ordinals(sharp, {0, 1, 4, 5}); -} - -TEST_CASE(PFlashSelectionFixture, document_prior_exponent_zero_is_todays_ranking) { - const auto candidates = document_candidates(); - const auto off = select_pflash_candidates( - candidates, PFlashSelectionPolicy{300, 0.95, false, 0, 0.0}, - PFlashSelectionMode::BudgetOnly); - const auto shipped = select_pflash_candidates( - candidates, PFlashSelectionPolicy{300, 0.95, false}, - PFlashSelectionMode::BudgetOnly); - REQUIRE(off.ok); - REQUIRE(!off.doc_prior_applied); - REQUIRE(off.ordinals == shipped.ordinals); - REQUIRE(off.retained_tokens == shipped.retained_tokens); - REQUIRE(off.stop == shipped.stop); -} - -TEST_CASE(PFlashSelectionFixture, document_prior_is_a_no_op_below_three_documents) { - // Two documents: the shares are real but the prior stays off by rule. - std::vector candidates = document_candidates(); - for (auto & c : candidates) if (c.document == 2) c.document = 1; - const auto two = select_pflash_candidates( - candidates, PFlashSelectionPolicy{200, 0.95, false, 0, 1.0}, - PFlashSelectionMode::BudgetOnly); - REQUIRE(two.ok); - REQUIRE(two.documents == 2); - REQUIRE(!two.doc_prior_applied); - - // One document (a single book, as NoLiMa serves) is a no-op twice over: - // by the rule and because the only share is 1. - for (auto & c : candidates) c.document = 0; - const auto one = select_pflash_candidates( - candidates, PFlashSelectionPolicy{200, 0.95, false, 0, 1.0}, - PFlashSelectionMode::BudgetOnly); - REQUIRE(one.ok); - REQUIRE(one.documents == 1); - REQUIRE(!one.doc_prior_applied); - require_ordinals(one, {0, 4}); -} - -TEST_CASE(PFlashSelectionFixture, document_prior_composes_with_top_k) { - const auto candidates = document_candidates(); - // K=2 with the prior: the rank order is the prior's, the count is K's. - const auto composed = select_pflash_candidates( - candidates, PFlashSelectionPolicy{100000, 0.95, false, 2, 3.0}, - PFlashSelectionMode::TopK); - REQUIRE(composed.ok); - REQUIRE(composed.stop == PFlashSelectionStop::TopKReached); - REQUIRE(composed.doc_prior_applied); - REQUIRE(composed.retained_tokens == 200); - // Sharpened shares put doc0's 9 first and doc2's 5 second; doc1 is out. - require_ordinals(composed, {0, 4}); - - // The same K without the prior keeps the raw top two, which here is the - // same pair -- so also check a K that exposes the reordering below them. - const auto plain_three = select_pflash_candidates( - candidates, PFlashSelectionPolicy{100000, 0.95, false, 3}, - PFlashSelectionMode::TopK); - const auto prior_three = select_pflash_candidates( - candidates, PFlashSelectionPolicy{100000, 0.95, false, 3, 3.0}, - PFlashSelectionMode::TopK); - REQUIRE(plain_three.ok); - REQUIRE(prior_three.ok); - require_ordinals(plain_three, {0, 4, 5}); - require_ordinals(prior_three, {0, 4, 5}); - // At rank four the orders diverge: raw picks doc1's 2, the prior picks - // doc0's 1 because doc1's share is cubed away. - const auto plain_four = select_pflash_candidates( - candidates, PFlashSelectionPolicy{100000, 0.95, false, 4}, - PFlashSelectionMode::TopK); - const auto prior_four = select_pflash_candidates( - candidates, PFlashSelectionPolicy{100000, 0.95, false, 4, 3.0}, - PFlashSelectionMode::TopK); - require_ordinals(plain_four, {0, 2, 4, 5}); - require_ordinals(prior_four, {0, 1, 4, 5}); -} - -TEST_CASE(PFlashSelectionFixture, document_prior_rejects_a_negative_exponent) { - const auto candidates = document_candidates(); - REQUIRE(!select_pflash_candidates( - candidates, PFlashSelectionPolicy{200, 0.95, false, 0, -1.0}, - PFlashSelectionMode::BudgetOnly).ok); -} - -TEST_CASE(PFlashSelectionFixture, document_prior_environment_resolves_or_fails) { - CleanPFlashEnv clean; - luce_test::ScopedEnvVar doc_prior{"PFLASH_SELECT_DOC_PRIOR", nullptr}; - set_env(kModeEnv, "budget_only"); - REQUIRE(resolve_or_fail(32768, 1024).doc_prior_exponent == 0.0); - - set_env("PFLASH_SELECT_DOC_PRIOR", "1.0"); - auto config = resolve_or_fail(32768, 1024); - REQUIRE(std::fabs(config.doc_prior_exponent - 1.0) < 1e-12); - - // It composes with top_k in the resolved config too. - set_env(kModeEnv, "top_k"); - set_env(kTopKEnv, "20"); - config = resolve_or_fail(32768, 1024); - REQUIRE(config.mode == PFlashSelectionMode::TopK); - REQUIRE(config.top_k == 20); - REQUIRE(std::fabs(config.doc_prior_exponent - 1.0) < 1e-12); - - PFlashSelectionConfig invalid; - std::string error; - for (const char * bad : {"-1", "abc", ""}) { - set_env("PFLASH_SELECT_DOC_PRIOR", bad); - REQUIRE(!resolve_pflash_selection(32768, 1024, invalid, error)); - REQUIRE(error.find("PFLASH_SELECT_DOC_PRIOR") != std::string::npos); - } -} - -TEST_CASE(PFlashSelectionFixture, forced_document_heads_keep_the_best_documents_identifiers) { - // Six candidates, three documents, two each. The first candidate of a - // document is its header: short and low-scoring, exactly what a score - // ranking drops. Masses: doc0 = 10*10 + 9*100 = 1000, - // doc1 = 1*10 + 3*100 = 310, doc2 = 10*10 + 5*100 = 600. - std::vector candidates{ - candidate(0, 0, 10, 10.0), // doc0 header - candidate(1, 10, 110, 9.0), // doc0 body - candidate(2, 110, 120, 1.0), // doc1 header - candidate(3, 120, 220, 3.0), // doc1 body - candidate(4, 220, 230, 10.0), // doc2 header - candidate(5, 230, 330, 5.0), // doc2 body - }; - const size_t documents[6] = {0, 0, 1, 1, 2, 2}; - for (size_t index = 0; index < candidates.size(); ++index) { - candidates[index].document = documents[index]; - } - - // Without forcing, a tight budget spends itself on the best body and - // leaves the weakest document with nothing -- not even its identifier. - const auto plain = select_pflash_candidates( - candidates, PFlashSelectionPolicy{130, 0.95, false}, - PFlashSelectionMode::BudgetOnly); - REQUIRE(plain.ok); - REQUIRE(plain.forced_doc_heads == 0); - require_ordinals(plain, {0, 1, 4}); - - // Forcing the headers of the top two documents by mass (doc0, doc2) - // charges them first; they are short, so the fill still gets the bodies. - const auto heads = select_pflash_candidates( - candidates, PFlashSelectionPolicy{230, 0.95, false, 0, 0.0, 2}, - PFlashSelectionMode::BudgetOnly); - REQUIRE(heads.ok); - REQUIRE(heads.forced_doc_heads == 2); - require_ordinals(heads, {0, 1, 4, 5}); - - // D above the document count forces every header that fits. At the same - // budget that kept no identifier for doc1 above, its header now survives. - const auto all_heads = select_pflash_candidates( - candidates, PFlashSelectionPolicy{130, 0.95, false, 0, 0.0, 9}, - PFlashSelectionMode::BudgetOnly); - REQUIRE(all_heads.ok); - REQUIRE(all_heads.forced_doc_heads == 3); - REQUIRE(all_heads.retained_tokens <= 130); - REQUIRE(all_heads.retained_tokens == 130); - // Three headers (30 tokens) plus the single body that still fits. - require_ordinals(all_heads, {0, 1, 2, 4}); -} - -TEST_CASE(PFlashSelectionFixture, forced_document_heads_compose_with_top_k_and_respect_the_budget) { - std::vector candidates{ - candidate(0, 0, 10, 1.0), - candidate(1, 10, 110, 9.0), - candidate(2, 110, 120, 1.0), - candidate(3, 120, 220, 3.0), - candidate(4, 220, 230, 1.0), - candidate(5, 230, 330, 5.0), - }; - const size_t documents[6] = {0, 0, 1, 1, 2, 2}; - for (size_t index = 0; index < candidates.size(); ++index) { - candidates[index].document = documents[index]; - } - // K counts only the ordinary fill; the forced headers are mandatory and - // sit outside it, so top-1 plus heads-of-3 keeps four candidates. - const auto composed = select_pflash_candidates( - candidates, PFlashSelectionPolicy{100000, 0.95, true, 1, 0.0, 3}, - PFlashSelectionMode::TopK); - REQUIRE(composed.ok); - REQUIRE(composed.stop == PFlashSelectionStop::TopKReached); - REQUIRE(composed.forced_doc_heads == 3); - require_ordinals(composed, {0, 1, 2, 4}); - REQUIRE(composed.retained_tokens == 130); - - // A budget too small for every header drops the ones that no longer fit - // rather than failing, and never exceeds the ceiling. - const auto tight = select_pflash_candidates( - candidates, PFlashSelectionPolicy{25, 0.95, true, 1, 0.0, 3}, - PFlashSelectionMode::TopK); - REQUIRE(tight.ok); - REQUIRE(tight.retained_tokens <= 25); - REQUIRE(tight.forced_doc_heads == 2); - - // A structurally required span is charged before any header, so forcing - // headers can never starve it. - std::vector with_mandatory = candidates; - with_mandatory[5].mandatory = true; - const auto safe = select_pflash_candidates( - with_mandatory, PFlashSelectionPolicy{110, 0.95, true, 1, 0.0, 3}, - PFlashSelectionMode::TopK); - REQUIRE(safe.ok); - REQUIRE(safe.retained_tokens <= 110); - // Ordinal 5 (100 tokens, mandatory) plus the headers that still fit. - REQUIRE(std::find(safe.ordinals.begin(), safe.ordinals.end(), (size_t) 5) != - safe.ordinals.end()); -} - -TEST_CASE(PFlashSelectionFixture, forced_document_heads_need_three_documents_and_a_valid_count) { - auto candidates = document_candidates(); - for (auto & c : candidates) c.document = 0; - const auto single = select_pflash_candidates( - candidates, PFlashSelectionPolicy{100000, 0.95, false, 0, 0.0, 5}, - PFlashSelectionMode::BudgetOnly); - REQUIRE(single.ok); - REQUIRE(single.forced_doc_heads == 0); - - REQUIRE(!select_pflash_candidates( - document_candidates(), PFlashSelectionPolicy{200, 0.95, false, 0, 0.0, -1}, - PFlashSelectionMode::BudgetOnly).ok); -} - -TEST_CASE(PFlashSelectionFixture, forced_document_heads_environment_resolves_or_fails) { - CleanPFlashEnv clean; - luce_test::ScopedEnvVar heads{"PFLASH_SELECT_FORCE_DOC_HEADS", nullptr}; - set_env(kModeEnv, "top_k"); - set_env(kTopKEnv, "20"); - REQUIRE(resolve_or_fail(32768, 1024).force_doc_heads == 0); - - set_env("PFLASH_SELECT_FORCE_DOC_HEADS", "5"); - const auto config = resolve_or_fail(32768, 1024); - REQUIRE(config.force_doc_heads == 5); - REQUIRE(config.mode == PFlashSelectionMode::TopK); - REQUIRE(config.top_k == 20); - - PFlashSelectionConfig invalid; - std::string error; - for (const char * bad : {"-1", "abc", ""}) { - set_env("PFLASH_SELECT_FORCE_DOC_HEADS", bad); - REQUIRE(!resolve_pflash_selection(32768, 1024, invalid, error)); - REQUIRE(error.find("PFLASH_SELECT_FORCE_DOC_HEADS") != std::string::npos); - } -} From 903b714eda057c30ff7bc358c7adafb68b3b5cb6 Mon Sep 17 00:00:00 2001 From: Graffioh Date: Tue, 22 Sep 2026 18:02:58 +0000 Subject: [PATCH 11/26] fix(pflash): take the chat scorer query from the latest user turn MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Every chat template appends a think or channel prefix after the assistant generation marker ("<|im_start|>assistant\n\n", "<｜Assistant｜>", "<|turn>model\n<|channel>thought..."). The marker-based tail extractor required only whitespace there, so on real prompts it took the generation prompt as the last message: the scorer query became "", the assistant header was pinned instead of the user's, and pflash_required failed to map. The legacy unconfigured path used the same span and regressed with it. The extractor now scans the rendered prompt's turns: an assistant turn left open at the end is the generation prompt whatever follows its marker, tool output wrapped in a user turn does not count, and the query comes from the latest user turn (else the latest turn with content). The derived query then takes the explicit pflash_query path: its span is the scorer window, pinned under strict selection together with the turn's role header; an explicit pflash_query still replaces it. Tests render prompts through render_chat_template for Qwen, Gemma, DeepSeek and Laguna with thinking on and off, instead of hand-written strings that omitted the think prefix. The stale default-parser assertion is updated, and the BPE tokenizer fixture now writes per-process paths so parallel ctest runs stop overwriting each other. Co-Authored-By: Claude Opus 5.5 (1M context) --- server/README.md | 10 + server/src/server/http_server.cpp | 367 ++++++++++++++------------ server/src/server/http_server.h | 27 +- server/test/test_pflash_selection.cpp | 7 +- server/test/test_server_unit.cpp | 306 +++++++++++++++------ 5 files changed, 454 insertions(+), 263 deletions(-) diff --git a/server/README.md b/server/README.md index b939762bc..1ab4b5b8f 100644 --- a/server/README.md +++ b/server/README.md @@ -388,6 +388,16 @@ restores the previous all-layer running-max scorer. The Qwen3.5 attention runs dense (`ggml_flash_attn_ext`); the block-sparse FlashPrefill kernels still dispatch head dimension 128 only. +The scorer query is the tail (`PFLASH_SELECT_QUERY_TOKENS`, default 8) of the +latest user turn, located by the model's own chat markers in the rendered +prompt. Tool output wrapped in a user turn and the generation prompt, with +its think prefix, never count as that turn. Strict selection keeps the query +and its turn's role header, and it runs on every turn of a multi-turn chat. +A request's `pflash_query` string replaces the derived query and keeps its +whole span; it is meant for benchmarks. `PFLASH_SELECT_QUERY_PARSER=latest_user` +selects the benchmark parser, which finds the latest user message through +sentinel renders. + `PFLASH_SEGMENT_PROBE_GGUF` loads a segment probe (schema `qwen3_5_0_8b_segment_probe_v1`): a 264K-parameter network on the same block-14 tap that scores every token for "a new unit of text starts here". With it diff --git a/server/src/server/http_server.cpp b/server/src/server/http_server.cpp index 554c55fbb..07b356a42 100644 --- a/server/src/server/http_server.cpp +++ b/server/src/server/http_server.cpp @@ -440,144 +440,161 @@ std::string pflash_token_fingerprint( return encoded; } -PflashChatTailSpan pflash_last_message_content_span( +PflashChatTurnSpan pflash_chat_query_turn( const Tokenizer & marker_tokenizer, const ChatMarkers & markers, const Tokenizer & tokenizer, const std::vector & prompt) { - PflashChatTailSpan tail; - if (prompt.empty()) return tail; - + PflashChatTurnSpan chosen; + if (prompt.empty()) return chosen; + + // Marker strings, searched in the decoded prompt text so a drafter whose + // vocabulary spells the control tokens differently still maps. + struct Mark { + size_t at = 0; + size_t len = 0; + bool role = false; + std::string text; + }; const auto seq_text = [&marker_tokenizer]( const std::vector & seq) { std::string text; for (const int32_t id : seq) text += marker_tokenizer.token_text(id); return text; }; - std::vector role_marks; + std::vector> needles; for (const auto & seq : markers.next_role_starts) { std::string text = seq_text(seq); - if (!text.empty()) role_marks.push_back(std::move(text)); + if (!text.empty()) needles.emplace_back(std::move(text), true); } - std::vector end_marks; for (const auto & seq : markers.end_msg_seqs) { std::string text = seq_text(seq); - if (!text.empty()) end_marks.push_back(std::move(text)); + if (!text.empty()) needles.emplace_back(std::move(text), false); } - if (role_marks.empty() && end_marks.empty()) return tail; + if (needles.empty()) return chosen; const DecodedPrompt decoded = decode_prompt_with_offsets(tokenizer, prompt); const std::string & text = decoded.text; - - // Last occurrence strictly before `before` (npos: anywhere), as - // (offset, length). Marker strings never overlap themselves, so the - // rightmost start across all needles is the answer. - const auto last_before = [&text]( - const std::vector & needles, size_t before) - -> std::pair { - size_t best = std::string::npos, len = 0; - for (const auto & needle : needles) { - const size_t at = before == std::string::npos - ? text.rfind(needle) - : before == 0 - ? std::string::npos - : text.rfind(needle, before - 1); - if (at != std::string::npos && - (best == std::string::npos || at > best)) { - best = at; - len = needle.size(); - } + std::vector marks; + for (const auto & [needle, role] : needles) { + for (size_t at = text.find(needle); at != std::string::npos; + at = text.find(needle, at + needle.size())) { + marks.push_back({at, needle.size(), role, needle}); } - return {best, len}; - }; - - const auto last_role = last_before(role_marks, std::string::npos); - const auto last_end = last_before(end_marks, std::string::npos); - if (last_role.first == std::string::npos && - last_end.first == std::string::npos) { - return tail; } + if (marks.empty()) return chosen; + std::sort(marks.begin(), marks.end(), + [] (const Mark & a, const Mark & b) { return a.at < b.at; }); - // Families whose markers already name the role (DeepSeek "<｜User｜>", - // Laguna "") delimit content at the marker itself. Generic markers - // (Qwen "<|im_start|>", Gemma "<|turn>") are followed by a "name\n" - // header line. + // Families whose markers name the role (DeepSeek "<｜User｜>", Laguna + // "") start content at the marker; generic markers (Qwen + // "<|im_start|>", Gemma "<|turn>") are followed by a "name\n" line. const bool marker_carries_role = markers.role_starts_delimit || markers.family == "laguna"; + const auto role_from_marker = [] (const std::string & marker) { + std::string role; + for (const char c : marker) { + if (std::isalpha((unsigned char) c)) { + role += (char) std::tolower((unsigned char) c); + } + } + return role; + }; + const auto is_space = [] (char c) { + return c == ' ' || c == '\t' || c == '\n' || c == '\r'; + }; + const auto starts_with = [&text] (size_t at, const char * prefix) { + return text.compare(at, std::strlen(prefix), prefix) == 0; + }; - // The character offset where the final message's content stops. The - // prompt's last marker is either the message's own end marker, the - // assistant generation marker (a role marker followed only by a - // role-name header), or a role marker that opened an unterminated - // message — in which case content runs to the prompt end. - size_t content_end_text; - if (last_role.first != std::string::npos && - (last_end.first == std::string::npos || - last_role.first > last_end.first)) { - const size_t after = last_role.first + last_role.second; - bool generation; + struct Turn { + size_t role_at = 0; + size_t content_at = 0; + size_t content_end = 0; + std::string role; + bool closed = false; + }; + std::vector turns; + for (size_t index = 0; index < marks.size(); ++index) { + const Mark & mark = marks[index]; + if (!mark.role) continue; + Turn turn; + turn.role_at = mark.at; + size_t content_at = mark.at + mark.len; if (marker_carries_role) { - generation = text.find_first_not_of(" \t\n\r", after) == - std::string::npos; + turn.role = role_from_marker(mark.text); } else { - size_t i = after; - while (i < text.size() && i - after < 16 && - std::isalpha((unsigned char) text[i])) ++i; - generation = i > after && i < text.size() && text[i] == '\n' && - text.find_first_not_of(" \t\n\r", i + 1) == - std::string::npos; - } - if (generation) { - if (marker_carries_role) { - content_end_text = last_role.first; - } else { - // Content ends at this message's end marker — the last one - // after the role marker that opened it, not an earlier - // turn's. - const auto prev_role = last_before(role_marks, last_role.first); - content_end_text = - (last_end.first != std::string::npos && - (prev_role.first == std::string::npos || - last_end.first > prev_role.first)) - ? last_end.first : last_role.first; + size_t name_end = content_at; + while (name_end < text.size() && name_end - content_at < 16 && + std::isalpha((unsigned char) text[name_end])) { + ++name_end; } - } else { - content_end_text = text.size(); - } - } else { - content_end_text = last_end.first; - } - - // The marker that opened the message containing content_end. - const auto role = last_before(role_marks, content_end_text); - size_t content_begin_text = 0; - if (role.first != std::string::npos) { - tail.role_begin = token_at_offset(decoded, role.first); - size_t begin = role.first + role.second; - if (!marker_carries_role) { - size_t i = begin; - while (i < content_end_text && i - begin < 16 && - std::isalpha((unsigned char) text[i])) ++i; - if (i > begin && i < content_end_text && text[i] == '\n') { - begin = i + 1; + turn.role = text.substr(content_at, name_end - content_at); + if (name_end < text.size() && text[name_end] == '\n') { + content_at = name_end + 1; } } - content_begin_text = begin; + // A turn ends at its end marker, or -- when role markers delimit + // (DeepSeek user turns) -- at the next role marker. Nothing after + // it leaves the turn open to the prompt end. + size_t content_end = text.size(); + if (index + 1 < marks.size()) { + const Mark & next = marks[index + 1]; + content_end = next.at; + turn.closed = !next.role || markers.role_starts_delimit; + } + // Content ignores the whitespace the template wraps it in. + while (content_at < content_end && is_space(text[content_at])) { + ++content_at; + } + while (content_end > content_at && is_space(text[content_end - 1])) { + --content_end; + } + turn.content_at = content_at; + turn.content_end = content_end; + // Tool output travels in user turns on some templates. + if (starts_with(content_at, "") || + starts_with(content_at, "")) { + turn.role = "tool"; + } + turns.push_back(std::move(turn)); + } + if (turns.empty()) return chosen; + + // An assistant turn left open at the prompt end is the generation prompt + // ("<|im_start|>assistant\n\n"): template machinery, never query. + size_t usable = turns.size(); + const Turn & last = turns.back(); + if (!last.closed && (last.role == "assistant" || last.role == "model")) { + --usable; + } + // The query comes from the latest user turn; a conversation without one + // falls back to its latest turn with content. + const Turn * query = nullptr; + for (size_t index = usable; index-- > 0;) { + const Turn & turn = turns[index]; + if (turn.content_end <= turn.content_at) continue; + if (turn.role == "user") { query = &turn; break; } + if (!query) query = &turn; + } + if (!query) return chosen; + + // Content bounds in tokens: the first token starting at-or-after each + // character offset, so a token merged across a boundary stays with the + // content it ends. + const auto token_from = [&decoded] (size_t at) { + return (int) (std::lower_bound( + decoded.token_begin.begin(), decoded.token_begin.end(), at) - + decoded.token_begin.begin()); + }; + chosen.role_begin = token_at_offset(decoded, query->role_at); + chosen.content_begin = token_from(query->content_at); + chosen.content_end = token_from(query->content_end); + if (chosen.role_begin > chosen.content_begin) { + chosen.role_begin = chosen.content_begin; } - - tail.content_end = content_end_text >= text.size() - ? (int) prompt.size() - : token_at_offset(decoded, content_end_text); - // First token whose text starts at-or-after the content offset — a token - // merged across the header/content boundary stays with the header. - tail.content_begin = (int) (std::lower_bound( - decoded.token_begin.begin(), decoded.token_begin.end(), - content_begin_text) - decoded.token_begin.begin()); - if (tail.role_begin > tail.content_begin) { - tail.role_begin = tail.content_begin; - } - return tail; + return chosen; } bool pflash_full_cache_restore_allowed( @@ -3538,24 +3555,26 @@ std::string HttpServer::apply_pflash_compression( std::string last_user_text; int query_content_begin = -1; int query_content_end = -1; - // Complete token span of an explicit pflash_query inside the boundary - // content, when it was mapped against the decoded token text. The strict - // selector keeps the whole span mandatory; the scorer window is its tail. - PFlashTokenSpan explicit_query_span{-1, -1}; - // Header ("<|im_start|>user\n") opening the last message, when the chat + // Complete token span of the scorer query: an explicit pflash_query + // (benchmark override) mapped against the decoded token text, or the + // tail of the latest user turn (chat default). The strict selector keeps + // the whole span mandatory; the scorer window is its tail. + PFlashTokenSpan query_span{-1, -1}; + std::string query_span_rule; + // Header ("<|im_start|>user\n") opening the query's turn, when the chat // markers resolved it — pinned mandatory so a compressed prompt keeps // the current turn's role envelope. - PFlashTokenSpan last_role_header{-1, -1}; + PFlashTokenSpan query_role_header{-1, -1}; std::vector required_instruction_spans; - // Chat-first scorer query: the last message's content span, located by - // the rendered prompt's own control markers. Feeds the strict tail + // Chat-first scorer query: the latest user turn's content span, located + // by the rendered prompt's own control markers. Feeds the strict tail // parser and the legacy window; unused when a benchmark parser // (latest_user) or a marker-less prompt needs the sentinel mapping. - http_detail::PflashChatTailSpan chat_tail; + http_detail::PflashChatTurnSpan chat_turn; if (!experiment.configured || tail_parser) { ChatMarkers chat_markers; if (resolve_chat_markers(tokenizer_, chat_markers)) { - chat_tail = http_detail::pflash_last_message_content_span( + chat_turn = http_detail::pflash_chat_query_turn( tokenizer_, chat_markers, *drafter_tokenizer_, drafter_ids); } } @@ -3583,8 +3602,10 @@ std::string HttpServer::apply_pflash_compression( const bool semantic_parser = experiment.query_parser == luce::pflash::PFlashQueryParser::SemanticUser; + // The latest user message bounds the query; the chat parser + // falls back to the last message when there is none. int boundary_index = (int) messages.size() - 1; - if (!raw_text_input && semantic_parser) { + if (!raw_text_input && (semantic_parser || last_user_index >= 0)) { boundary_index = last_user_index; } @@ -3668,15 +3689,15 @@ std::string HttpServer::apply_pflash_compression( }; std::string boundary_error; - // Chat default: the last message's content bounds come from the - // rendered prompt's control markers — no sentinel re-renders. - if (tail_parser && chat_tail.valid()) { - query_content_begin = chat_tail.content_begin; - query_content_end = chat_tail.content_end; - if (chat_tail.role_begin >= 0 && - chat_tail.role_begin < chat_tail.content_begin) { - last_role_header = {chat_tail.role_begin, - chat_tail.content_begin}; + // Chat default: the latest user turn's content bounds come from + // the rendered prompt's control markers — no sentinel re-renders. + if (tail_parser && chat_turn.valid()) { + query_content_begin = chat_turn.content_begin; + query_content_end = chat_turn.content_end; + if (chat_turn.role_begin >= 0 && + chat_turn.role_begin < chat_turn.content_begin) { + query_role_header = {chat_turn.role_begin, + chat_turn.content_begin}; } } if (query_content_begin < 0) { @@ -3700,6 +3721,19 @@ std::string HttpServer::apply_pflash_compression( query_content_end > (int) drafter_ids.size()) { return "PFlash strict selection content boundary mapping failed"; } + // Chat default: without an explicit pflash_query the scorer + // query is the tail of the boundary content, and it takes the + // explicit query's path from here on. + if (tail_parser && req.pflash_query.empty()) { + const auto window = http_detail::pflash_tail_query_window( + drafter_ids, experiment.query_tokens, + query_content_end, query_content_begin); + if (window.valid()) { + query_span = {window.end - window.tokens, window.end}; + query_span_rule = chat_turn.valid() ? "chat_user_tail" + : (raw_text_input ? "content_tail" : "prompt_tail"); + } + } if (experiment.selection_active) { const auto instruction_plan = @@ -3759,27 +3793,23 @@ std::string HttpServer::apply_pflash_compression( } required_instruction_spans.push_back(required_span); } - // An explicit scorer query also pins its complete span: the - // whole question is mandatory even though the scorer only - // consumes its bounded tail. Mapping against the decoded - // content text (not a standalone encoding) keeps BPE boundary - // merges like " What" inside the span. Benchmark-only: under - // the chat tail parser the query may sit anywhere before the - // closing markers; latest_user still scopes it to the user - // message. + // An explicit scorer query replaces the chat-derived one and + // pins its complete span: the whole question is mandatory + // even though the scorer only consumes its bounded tail. + // Mapping against the decoded content text (not a standalone + // encoding) keeps BPE boundary merges like " What" inside the + // span. Benchmark-only: under the chat tail parser the query + // may sit anywhere before the user turn's closing marker; + // latest_user still scopes it to the user message. if (!req.pflash_query.empty()) { const int query_search_begin = semantic_parser ? query_content_begin : 0; - const int query_search_end = semantic_parser - ? query_content_end - : (query_content_end > 0 - ? query_content_end : (int) drafter_ids.size()); - explicit_query_span = - http_detail::pflash_decoded_text_span( - *drafter_tokenizer_, drafter_ids, - query_search_begin, query_search_end, - req.pflash_query); - if (explicit_query_span.begin < 0) { + query_span = http_detail::pflash_decoded_text_span( + *drafter_tokenizer_, drafter_ids, + query_search_begin, query_content_end, + req.pflash_query); + query_span_rule = "explicit_query_span"; + if (query_span.begin < 0) { return semantic_parser ? "PFlash strict selection explicit query mapping " "failed: pflash_query does not occur in the " @@ -3788,10 +3818,12 @@ std::string HttpServer::apply_pflash_compression( "failed: pflash_query does not occur in the " "prompt"; } - required_instruction_spans.push_back(explicit_query_span); } - if (last_role_header.begin >= 0) { - required_instruction_spans.push_back(last_role_header); + if (query_span.begin >= 0) { + required_instruction_spans.push_back(query_span); + } + if (query_role_header.begin >= 0) { + required_instruction_spans.push_back(query_role_header); } required_instruction_spans = http_detail::canonicalize_pflash_token_spans( @@ -3841,30 +3873,37 @@ std::string HttpServer::apply_pflash_compression( if (!last_user_text.empty()) { semantic_query_ids = drafter_tokenizer_->encode(last_user_text); } - if (explicit_query_span.begin >= 0) { - // The explicit query was already mapped against the decoded content - // text and pinned as a mandatory span. The scorer consumes the span's - // bounded tail window; the complete span stays in the target prompt. - parser_selection_rule = "explicit_query_span"; - query_window.end = explicit_query_span.end; + // Unconfigured (legacy) chat mode derives the query the same way, from + // the latest user turn's tail. + if (!experiment.configured && req.pflash_query.empty() && + chat_turn.valid()) { + const auto window = http_detail::pflash_tail_query_window( + drafter_ids, experiment.query_tokens, + chat_turn.content_end, chat_turn.content_begin); + if (window.valid()) { + query_span = {window.end - window.tokens, window.end}; + query_span_rule = "chat_user_tail"; + } + } + if (query_span.begin >= 0) { + // The query span — explicit or chat-derived — was mapped onto the + // prompt's own tokens and, under strict selection, pinned mandatory. + // The scorer consumes the span's bounded tail window; the complete + // span stays in the target prompt. + parser_selection_rule = query_span_rule; + query_window.end = query_span.end; query_window.tokens = (std::min)( - experiment.query_tokens, - explicit_query_span.end - explicit_query_span.begin); + experiment.query_tokens, query_span.end - query_span.begin); expected_query_ids.assign( drafter_ids.begin() + (query_window.end - query_window.tokens), drafter_ids.begin() + query_window.end); } else if (experiment.configured && (raw_text_input || tail_parser)) { - // Chat default: the scorer query is the tail of the last message's - // content — the prompt region before the closing/generation markers, - // clamped so it never swallows the role header. + // Content located, but no query span (a parser without a derived + // query, or selection inactive): score the content's tail. parser_selection_rule = raw_text_input ? "content_tail" : "prompt_tail"; query_window = http_detail::pflash_tail_query_window( drafter_ids, experiment.query_tokens, query_content_end, query_content_begin); - } else if (!experiment.configured && chat_tail.valid()) { - // Legacy (unconfigured) chat mode uses the same marker-derived tail. - query_window = http_detail::pflash_tail_query_window( - drafter_ids, 8, chat_tail.content_end, chat_tail.content_begin); } else if (!semantic_query_ids.empty()) { if (experiment.configured) parser_selection_rule = "semantic_suffix"; query_window = http_detail::find_pflash_query_window( @@ -3912,8 +3951,8 @@ std::string HttpServer::apply_pflash_compression( {"content_end", query_content_end}, {"query_begin", query_begin}, {"query_end", query_window.end}, - {"query_span_begin", explicit_query_span.begin}, - {"query_span_end", explicit_query_span.end}, + {"query_span_begin", query_span.begin}, + {"query_span_end", query_span.end}, {"requested_query_tokens", experiment.query_tokens}, {"required_text_count", req.pflash_required.size()}, {"expected_query_ids", expected_query_ids}, diff --git a/server/src/server/http_server.h b/server/src/server/http_server.h index 0e96f1dd5..adb1eae94 100644 --- a/server/src/server/http_server.h +++ b/server/src/server/http_server.h @@ -333,17 +333,22 @@ PFlashTokenSpan pflash_decoded_text_span( int end, const std::string & needle); -// The last chat message's content span inside a rendered prompt, located by -// the model's own chat control markers rather than message bookkeeping. -// ``role_begin`` is the marker opening that message (the header to pin); +// The chat turn the scorer query comes from, located by the model's own +// chat control markers in the rendered prompt rather than message +// bookkeeping: the latest user turn (tool output wrapped in a user turn does +// not count), else the latest turn with content. An assistant turn left open +// at the prompt end is the generation prompt -- whatever think or channel +// prefix the template adds to it -- and never a candidate. +// ``role_begin`` is the marker opening the turn (the header to pin); // ``content_begin`` skips the role-name line ("<|im_start|>user\n") when the -// family uses generic role markers; ``content_end`` sits before the closing -// or generation marker. Offsets are token indices in ``prompt``'s own -// vocabulary. ``markers`` were resolved on ``marker_tokenizer`` (the target -// model's); its marker strings are searched in the decoded prompt text, so -// a drafter whose vocabulary lacks the control tokens still maps correctly. -// Invalid when the prompt carries no chat markers. -struct PflashChatTailSpan { +// family uses generic role markers; ``content_end`` sits before the turn's +// closing marker. Both trim the whitespace the template wraps content in. +// Offsets are token indices in ``prompt``'s own vocabulary. ``markers`` were +// resolved on ``marker_tokenizer`` (the target model's); its marker strings +// are searched in the decoded prompt text, so a drafter whose vocabulary +// lacks the control tokens still maps correctly. Invalid when the prompt +// carries no chat markers. +struct PflashChatTurnSpan { int role_begin = -1; int content_begin = -1; int content_end = -1; @@ -353,7 +358,7 @@ struct PflashChatTailSpan { } }; -PflashChatTailSpan pflash_last_message_content_span( +PflashChatTurnSpan pflash_chat_query_turn( const Tokenizer & marker_tokenizer, const ChatMarkers & markers, const Tokenizer & tokenizer, diff --git a/server/test/test_pflash_selection.cpp b/server/test/test_pflash_selection.cpp index ccdb2189f..035d92b9f 100644 --- a/server/test/test_pflash_selection.cpp +++ b/server/test/test_pflash_selection.cpp @@ -348,7 +348,7 @@ TEST_CASE(PFlashSelectionFixture, resolver_defaults_to_legacy_arguments) { REQUIRE(!config.configured); REQUIRE(!config.selection_active); REQUIRE(config.mode == PFlashSelectionMode::Legacy); - REQUIRE(config.query_parser == PFlashQueryParser::SemanticUser); + REQUIRE(config.query_parser == PFlashQueryParser::ArbitraryTail); REQUIRE(config.chunk_size == 32); REQUIRE(config.query_tokens == 8); REQUIRE(std::abs(config.top_p - 0.95) < 1e-12); @@ -397,11 +397,6 @@ TEST_CASE(PFlashSelectionFixture, any_selection_environment_is_observable_before TEST_CASE(PFlashSelectionFixture, resolver_selects_explicit_query_parser) { CleanPFlashEnv env; - // Chat-first default: no override selects the rendered tail parser. - const auto fallback = resolve_or_fail(120000, 32); - REQUIRE(fallback.configured == false); - REQUIRE(fallback.query_parser == PFlashQueryParser::ArbitraryTail); - set_env(kQueryParserEnv, "arbitrary_tail"); const auto arbitrary = resolve_or_fail(120000, 32); REQUIRE(arbitrary.configured); diff --git a/server/test/test_server_unit.cpp b/server/test/test_server_unit.cpp index 3b0c5aa9d..5681d18a1 100644 --- a/server/test/test_server_unit.cpp +++ b/server/test/test_server_unit.cpp @@ -73,6 +73,7 @@ #include #else #include +#include #endif #if defined(_WIN32) @@ -432,9 +433,11 @@ static std::string test_gpt2_encode(const std::string & text) { static std::string write_pflash_bpe_tokenizer_fixture( const std::vector & raw_tokens, - const std::string & byte_cover) { - std::vector tokens{"<|im_start|>", "<|im_end|>"}; - std::vector types{3, 3}; + const std::string & byte_cover, + const std::vector & control_tokens = + {"<|im_start|>", "<|im_end|>"}) { + std::vector tokens = control_tokens; + std::vector types(tokens.size(), 3); const auto add = [&](const std::string & encoded, uint32_t type) { if (std::find(tokens.begin(), tokens.end(), encoded) == tokens.end()) { tokens.push_back(encoded); @@ -455,9 +458,17 @@ static std::string write_pflash_bpe_tokenizer_fixture( gguf_set_val_str(g, "tokenizer.ggml.pre", "qwen35"); gguf_set_val_u32(g, "tokenizer.ggml.bos_token_id", 0); gguf_set_val_u32(g, "tokenizer.ggml.eos_token_id", 1); + // Per-process names: ctest runs each case in its own process, in + // parallel, and every process starts the serial at zero. static int fixture_serial = 0; - const std::string path = "/tmp/dflash_test_pflash_bpe_" + - std::to_string(++fixture_serial) + ".gguf"; +#if defined(_WIN32) + const long long pid = (long long) _getpid(); +#else + const long long pid = (long long) getpid(); +#endif + const std::string path = test_tmp_path(( + "luce_test_pflash_bpe_" + std::to_string(pid) + "_" + + std::to_string(++fixture_serial) + ".gguf").c_str()).string(); gguf_write_to_file(g, path.c_str(), /*only_meta=*/false); gguf_free(g); return path; @@ -818,78 +829,139 @@ TEST_CASE(ServerUnitFixture, test_pflash_tail_query_window) { TEST_ASSERT(!http_detail::pflash_tail_query_window(long_prompt, 128, 201).valid()); } -TEST_CASE(ServerUnitFixture, - test_pflash_last_content_span_stops_before_chat_markers) { - const std::string rendered = - "<|im_start|>system\nYou are helpful.<|im_end|>\n" - "<|im_start|>user\nfirst turn<|im_end|>\n" - "<|im_start|>assistant\nSure.<|im_end|>\n" - "<|im_start|>user\nWhat is the answer?<|im_end|>\n" - "<|im_start|>assistant\n"; +// Renders ``messages`` with the server's own chat template and returns the +// chat-query turn the scorer would use, decoded as {header, content}. +struct PflashRenderedQueryTurn { + bool valid = false; + std::string family; + std::string header; + std::string content; + std::string after; +}; + +static PflashRenderedQueryTurn pflash_rendered_query_turn( + const std::vector & messages, + ChatFormat format, + bool thinking, + const std::vector & control_tokens) { + const std::string rendered = render_chat_template( + messages, format, /*add_generation_prompt=*/true, thinking); const std::string path = write_pflash_bpe_tokenizer_fixture( - {"What", " is", " the", " answer", "?", "user", "assistant", - "system", "\n", "You", " are", " helpful", ".", "first", " turn", - "Sure"}, - rendered); + {"What", " is", " the", " answer", "?", "user", "assistant", "model", + "system", "\n", "Sure", "."}, + rendered, control_tokens); Tokenizer tok; - TEST_ASSERT(tok.load_from_gguf(path.c_str())); - + PflashRenderedQueryTurn out; + if (!tok.load_from_gguf(path.c_str())) { + unlink(path.c_str()); + return out; + } const auto prompt = tok.encode(rendered); ChatMarkers markers; - TEST_ASSERT(resolve_chat_markers(tok, markers)); - TEST_ASSERT(markers.family == "qwen"); - - const auto span = http_detail::pflash_last_message_content_span( - tok, markers, tok, prompt); - TEST_ASSERT(span.valid()); - TEST_ASSERT(tok.decode({prompt.begin() + span.content_begin, - prompt.begin() + span.content_end}) - == "What is the answer?"); - // The pinned header is exactly the role envelope of the last turn. - TEST_ASSERT(tok.decode({prompt.begin() + span.role_begin, - prompt.begin() + span.content_begin}) - == "<|im_start|>user\n"); - // Everything after the content is template machinery — never part of - // the scorer query. - TEST_ASSERT(tok.decode({prompt.begin() + span.content_end, - prompt.end()}) - == "<|im_end|>\n<|im_start|>assistant\n"); - - const auto window = http_detail::pflash_tail_query_window( - prompt, 8, span.content_end, span.content_begin); - TEST_ASSERT(window.valid()); - TEST_ASSERT(window.end == span.content_end); - TEST_ASSERT(window.tokens == - std::min(8, span.content_end - span.content_begin)); + if (resolve_chat_markers(tok, markers)) { + out.family = markers.family; + const auto turn = http_detail::pflash_chat_query_turn( + tok, markers, tok, prompt); + out.valid = turn.valid(); + if (out.valid) { + out.header = tok.decode({prompt.begin() + turn.role_begin, + prompt.begin() + turn.content_begin}); + out.content = tok.decode({prompt.begin() + turn.content_begin, + prompt.begin() + turn.content_end}); + out.after = tok.decode({prompt.begin() + turn.content_end, + prompt.end()}); + } + } unlink(path.c_str()); + return out; } TEST_CASE(ServerUnitFixture, - test_pflash_last_content_span_deepseek_delimited_roles) { - const std::string path = write_deepseek_marker_tokenizer_fixture(); - Tokenizer tok; - TEST_ASSERT(tok.load_from_gguf(path.c_str())); + test_pflash_chat_query_turn_skips_rendered_generation_prompt) { + // Every family's generation prompt carries a think/channel prefix after + // the assistant marker; the query must still be the latest user turn. + const std::vector messages{ + {"system", "You are helpful.", ""}, + {"user", "first turn", ""}, + {"assistant", "Sure.", ""}, + {"user", "What is the answer?", ""}, + }; + struct Family { + ChatFormat format; + const char * name; + std::vector control_tokens; + const char * header; + }; + const std::vector families{ + {ChatFormat::QWEN3, "qwen", {"<|im_start|>", "<|im_end|>"}, + "<|im_start|>user\n"}, + {ChatFormat::GEMMA4, "gemma", {"<|turn>", ""}, + "<|turn>user\n"}, + {ChatFormat::DEEPSEEK4, "deepseek", + {"<｜begin▁of▁sentence｜>", "<｜end▁of▁sentence｜>", "<｜User｜>", + "<｜Assistant｜>"}, + "<｜User｜>"}, + {ChatFormat::LAGUNA, "laguna", + {"", "", "", "", "", + ""}, + "\n"}, + }; + for (const auto & family : families) { + for (const bool thinking : {false, true}) { + const auto turn = pflash_rendered_query_turn( + messages, family.format, thinking, family.control_tokens); + TEST_ASSERT_MSG(turn.family == family.name, family.name); + TEST_ASSERT_MSG(turn.valid, family.name); + TEST_ASSERT_MSG(turn.content == "What is the answer?", + std::string(family.name) + " thinking=" + + (thinking ? "1" : "0") + " content=[" + + turn.content + "]"); + TEST_ASSERT_MSG(turn.header == family.header, + std::string(family.name) + " header=[" + + turn.header + "]"); + // The rendered tail after the content is template machinery. + TEST_ASSERT_MSG(turn.after.find("answer") == std::string::npos, + family.name); + } + } +} - ChatMarkers markers; - TEST_ASSERT(resolve_chat_markers(tok, markers)); - TEST_ASSERT(markers.family == "deepseek"); - TEST_ASSERT(markers.role_starts_delimit); +TEST_CASE(ServerUnitFixture, + test_pflash_chat_query_turn_skips_tool_output_turns) { + // Agent loop: tool results render inside user turns on Qwen and + // DeepSeek. The query stays on the user's own latest turn. + const std::vector messages{ + {"user", "What is the answer?", ""}, + {"assistant", "Sure.", ""}, + {"tool", "tool output here", "call-1"}, + }; + const auto qwen = pflash_rendered_query_turn( + messages, ChatFormat::QWEN3, false, {"<|im_start|>", "<|im_end|>"}); + TEST_ASSERT(qwen.valid); + TEST_ASSERT_MSG(qwen.content == "What is the answer?", qwen.content); + TEST_ASSERT(qwen.header == "<|im_start|>user\n"); - // bos + system text + <|User|> + content + <|Assistant|> generation - const std::vector prompt = {1, 0, 3, 0, 0, 4}; - const auto span = http_detail::pflash_last_message_content_span( - tok, markers, tok, prompt); - TEST_ASSERT(span.valid()); - TEST_ASSERT(span.role_begin == 2); - TEST_ASSERT(span.content_begin == 3); - TEST_ASSERT(span.content_end == 5); - TEST_ASSERT(tok.decode({prompt.begin() + span.content_begin, - prompt.begin() + span.content_end}) == "xx"); - remove_test_path(path); + const auto deepseek = pflash_rendered_query_turn( + messages, ChatFormat::DEEPSEEK4, true, + {"<｜begin▁of▁sentence｜>", "<｜end▁of▁sentence｜>", "<｜User｜>", + "<｜Assistant｜>"}); + TEST_ASSERT(deepseek.valid); + TEST_ASSERT_MSG(deepseek.content == "What is the answer?", + deepseek.content); } TEST_CASE(ServerUnitFixture, - test_pflash_last_content_span_open_tail_runs_to_prompt_end) { + test_pflash_chat_query_turn_falls_back_without_user_turn) { + const auto turn = pflash_rendered_query_turn( + {{"system", "You are helpful.", ""}}, ChatFormat::QWEN3, true, + {"<|im_start|>", "<|im_end|>"}); + TEST_ASSERT(turn.valid); + TEST_ASSERT(turn.content == "You are helpful."); + TEST_ASSERT(turn.header == "<|im_start|>system\n"); +} + +TEST_CASE(ServerUnitFixture, + test_pflash_chat_query_turn_open_tail_runs_to_prompt_end) { const std::string rendered = "<|im_start|>user\nhello there"; const std::string path = write_pflash_bpe_tokenizer_fixture( {"hello", " there", "user", "\n"}, rendered); @@ -899,7 +971,7 @@ TEST_CASE(ServerUnitFixture, const auto prompt = tok.encode(rendered); ChatMarkers markers; TEST_ASSERT(resolve_chat_markers(tok, markers)); - const auto span = http_detail::pflash_last_message_content_span( + const auto span = http_detail::pflash_chat_query_turn( tok, markers, tok, prompt); TEST_ASSERT(span.valid()); TEST_ASSERT(span.content_end == (int) prompt.size()); @@ -910,7 +982,7 @@ TEST_CASE(ServerUnitFixture, } TEST_CASE(ServerUnitFixture, - test_pflash_last_content_span_rejects_markerless_text) { + test_pflash_chat_query_turn_rejects_markerless_text) { const std::string rendered = "just some raw text, no chat markers"; const std::string path = write_pflash_bpe_tokenizer_fixture( {"just", " some", " raw", " text"}, rendered); @@ -920,7 +992,7 @@ TEST_CASE(ServerUnitFixture, const auto prompt = tok.encode(rendered); ChatMarkers markers; TEST_ASSERT(resolve_chat_markers(tok, markers)); - const auto span = http_detail::pflash_last_message_content_span( + const auto span = http_detail::pflash_chat_query_turn( tok, markers, tok, prompt); TEST_ASSERT(!span.valid()); unlink(path.c_str()); @@ -6751,13 +6823,16 @@ TEST_CASE(ServerUnitFixture, test_pflash_default_raw_text_maps_user_query) { } TEST_CASE(ServerUnitFixture, - test_pflash_strict_chat_tail_query_uses_last_message_content) { + test_pflash_strict_chat_tail_query_uses_last_user_content) { luce_test::ScopedEnvVar mode{"PFLASH_SELECT_MODE", "top_p"}; - const std::string rendered = - "<|im_start|>system\nYou are helpful.<|im_end|>\n" - "<|im_start|>user\nWhat is the answer?<|im_end|>\n" - "<|im_start|>assistant\n"; + // The server's own Qwen rendering, generation prompt and its think + // prefix included. + const std::string rendered = render_chat_template( + {{"system", "You are helpful.", ""}, + {"user", "What is the answer?", ""}}, + ChatFormat::QWEN3, /*add_generation_prompt=*/true, + /*enable_thinking=*/true); const std::string path = write_pflash_bpe_tokenizer_fixture( {"What", " is", " the", " answer", "?", "user", "assistant", "system", "\n", "You", " are", " helpful", "."}, @@ -6779,6 +6854,7 @@ TEST_CASE(ServerUnitFixture, ParsedRequest request; request.format = ApiFormat::OPENAI_CHAT; request.messages = json::array({ + {{"role", "system"}, {"content", "You are helpful."}}, {{"role", "user"}, {"content", "What is the answer?"}}, }); request.prompt_tokens = tokenizer.encode(rendered); @@ -6797,7 +6873,8 @@ TEST_CASE(ServerUnitFixture, } TEST_ASSERT(last_im_end > 0); // The scorer window ends where the user content does — the generation - // markers ("<|im_end|>\n<|im_start|>assistant\n") are never scored. + // prompt ("<|im_end|>\n<|im_start|>assistant\n\n") is never + // scored. TEST_ASSERT(backend.last_request.score_query_end == last_im_end); const int query_begin = backend.last_request.score_query_end - backend.last_request.score_query_tokens; @@ -6805,16 +6882,80 @@ TEST_CASE(ServerUnitFixture, TEST_ASSERT(tokenizer.decode({ids.begin() + query_begin, ids.begin() + last_im_end}) == "What is the answer?"); - // The last turn's role header is pinned mandatory. + // The query span and its turn's role header are pinned mandatory. bool header_pinned = false; + bool query_pinned = false; for (const auto & span : backend.last_request.required_instruction_spans) { - if (tokenizer.decode({ids.begin() + span.begin, - ids.begin() + span.end}) - == "<|im_start|>user\n") { + const std::string text = tokenizer.decode( + {ids.begin() + span.begin, ids.begin() + span.end}); + if (text.find("<|im_start|>user\n") != std::string::npos) { header_pinned = true; } + if (span.begin <= query_begin && span.end >= last_im_end) { + query_pinned = true; + } } TEST_ASSERT(header_pinned); + TEST_ASSERT(query_pinned); + unlink(path.c_str()); +} + +TEST_CASE(ServerUnitFixture, + test_pflash_legacy_chat_query_uses_last_user_turn) { + // No strict-selection environment: the legacy selector derives the same + // query, and a trailing tool result does not replace the user's turn. + const std::vector messages{ + {"user", "What is the answer?", ""}, + {"assistant", "Sure.", ""}, + {"tool", "tool output here", "call-1"}, + }; + const std::string rendered = render_chat_template( + messages, ChatFormat::QWEN3, /*add_generation_prompt=*/true, + /*enable_thinking=*/false); + const std::string path = write_pflash_bpe_tokenizer_fixture( + {"What", " is", " the", " answer", "?", "user", "assistant", "\n", + "Sure", "."}, + 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.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; + request.messages = json::array({ + {{"role", "user"}, {"content", "What is the answer?"}}, + {{"role", "assistant"}, {"content", "Sure."}}, + {{"role", "tool"}, {"content", "tool output here"}, + {"tool_call_id", "call-1"}}, + }); + 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); + const auto & ids = backend.last_request.input_ids; + const int query_end = backend.last_request.score_query_end; + const int query_begin = query_end - backend.last_request.score_query_tokens; + TEST_ASSERT(query_begin >= 0); + TEST_ASSERT(backend.last_request.score_query_tokens <= 8); + const std::string query = tokenizer.decode( + {ids.begin() + query_begin, ids.begin() + query_end}); + TEST_ASSERT_MSG(std::string("What is the answer?").size() >= query.size() && + std::string("What is the answer?").compare( + 19 - query.size(), query.size(), query) == 0, + query); unlink(path.c_str()); } @@ -6822,11 +6963,12 @@ TEST_CASE(ServerUnitFixture, test_pflash_strict_selection_owns_chat_continuations) { luce_test::ScopedEnvVar mode{"PFLASH_SELECT_MODE", "top_p"}; - const std::string rendered = - "<|im_start|>user\nfirst<|im_end|>\n" - "<|im_start|>assistant\nSure.<|im_end|>\n" - "<|im_start|>user\nsecond question<|im_end|>\n" - "<|im_start|>assistant\n"; + const std::string rendered = render_chat_template( + {{"user", "first", ""}, + {"assistant", "Sure.", ""}, + {"user", "second question", ""}}, + ChatFormat::QWEN3, /*add_generation_prompt=*/true, + /*enable_thinking=*/false); const std::string path = write_pflash_bpe_tokenizer_fixture( {"first", "second", " question", "user", "assistant", "\n", "Sure", "."}, From 8e7fc2cc7a21fd1b2f8a128b7dd85c326cd8f6ac Mon Sep 17 00:00:00 2001 From: Graffioh Date: Tue, 22 Sep 2026 18:20:10 +0000 Subject: [PATCH 12/26] feat(pflash): score the turns after the query in agent loops In an agent loop the latest user turn is followed by assistant and tool turns. Selection treated every token after the query window as a kept suffix and the head scorer masked those keys, so the tool output was pinned whole: a 2K-token tool result against a 0.3 keep ratio failed with mandatory_query_exceeds_budget. The prompt minus the query is now the candidate pool whatever side of the query it sits on. When turns follow the query's turn, the server pins only the rest of that turn through its closing marker and the generation prompt, and sets CompressRequest::query_suffix_candidates. The block-15 head then scores keys after the query window too (NoPE: no position term) and the selector stops treating that suffix as structural. A query in the final user turn keeps the old contract, so single-turn and plain multi-turn requests select exactly as before. The running-max scorer, alone or in the split, and the remote drafter IPC keep the kept-suffix contract. Co-Authored-By: Claude Opus 5.5 (1M context) --- server/README.md | 3 + server/src/common/model_backend.h | 4 + server/src/deepseek4/deepseek4_backend.cpp | 3 +- server/src/pflash/pflash_compress.cpp | 3 +- server/src/pflash/pflash_drafter.cpp | 10 ++- server/src/pflash/pflash_drafter.h | 5 +- server/src/pflash/pflash_selection.cpp | 5 +- server/src/pflash/pflash_selection.h | 15 +++- server/src/pflash/qwen35_drafter.cpp | 25 ++++-- server/src/qwen3/qwen3_backend.cpp | 3 +- server/src/qwen35/qwen35_backend.cpp | 3 +- .../src/qwen35/qwen35_layer_split_adapter.cpp | 3 +- server/src/server/http_server.cpp | 40 ++++++++- server/src/server/http_server.h | 6 ++ server/test/test_pflash_selection.cpp | 17 ++++ server/test/test_server_unit.cpp | 84 +++++++++++++++++++ 16 files changed, 206 insertions(+), 23 deletions(-) diff --git a/server/README.md b/server/README.md index 1ab4b5b8f..75270e758 100644 --- a/server/README.md +++ b/server/README.md @@ -393,6 +393,9 @@ latest user turn, located by the model's own chat markers in the rendered prompt. Tool output wrapped in a user turn and the generation prompt, with its think prefix, never count as that turn. Strict selection keeps the query and its turn's role header, and it runs on every turn of a multi-turn chat. +In an agent loop the assistant and tool turns after the user's turn are +scored against the query like the context before it; only the generation +prompt is kept with them. A request's `pflash_query` string replaces the derived query and keeps its whole span; it is meant for benchmarks. `PFLASH_SELECT_QUERY_PARSER=latest_user` selects the benchmark parser, which finds the latest user message through diff --git a/server/src/common/model_backend.h b/server/src/common/model_backend.h index 53f4ac5a1..3acb8ef24 100644 --- a/server/src/common/model_backend.h +++ b/server/src/common/model_backend.h @@ -271,6 +271,10 @@ struct ModelBackend { // Role-derived instruction structure in drafter-token coordinates. // Empty is a valid instruction-free or legacy request. std::vector required_instruction_spans; + // Strict selection with the block-15 head: the tokens after the query + // window are scored candidates rather than a kept suffix. The caller + // pins what of that suffix must stay (the generation prompt). + bool query_suffix_candidates = false; std::string drafter_path; // GGUF path (for lazy-load) int drafter_gpu = 0; // backend-local GPU for PFlash drafter bool skip_park = false; // true on >=32GB GPUs diff --git a/server/src/deepseek4/deepseek4_backend.cpp b/server/src/deepseek4/deepseek4_backend.cpp index 473cedc75..c4ebb489e 100644 --- a/server/src/deepseek4/deepseek4_backend.cpp +++ b/server/src/deepseek4/deepseek4_backend.cpp @@ -3181,7 +3181,8 @@ std::vector DeepSeek4Backend::compress_batch( result.compressed_ids = drafter_score_and_compress( pflash_drafter_ctx_, request.input_ids, request.keep_ratio, /*chunk_size=*/32, request.score_query_tokens, /*pool_kernel=*/13, - score_query_end, request.required_instruction_spans); + score_query_end, request.required_instruction_spans, + request.query_suffix_candidates); result.ok = !result.compressed_ids.empty(); } diff --git a/server/src/pflash/pflash_compress.cpp b/server/src/pflash/pflash_compress.cpp index 30949bb76..8ed55593c 100644 --- a/server/src/pflash/pflash_compress.cpp +++ b/server/src/pflash/pflash_compress.cpp @@ -225,7 +225,8 @@ std::vector select_pflash_chunks( const bool mandatory = luce::pflash::pflash_chunk_is_structurally_required( begin, end, query_begin, query_end, input_tokens, - required_instruction_spans); + required_instruction_spans, + /*query_suffix_structural=*/ !config.query_suffix_candidates); candidates.push_back({(size_t) chunk, begin, end, score, mandatory}); chunk_means.push_back({(float) score, chunk}); exact_chunk_scores.push_back(score); diff --git a/server/src/pflash/pflash_drafter.cpp b/server/src/pflash/pflash_drafter.cpp index 5fd2f2d5f..719659fb1 100644 --- a/server/src/pflash/pflash_drafter.cpp +++ b/server/src/pflash/pflash_drafter.cpp @@ -104,7 +104,8 @@ std::vector drafter_score_and_compress( int n_lookahead, int pool_kernel, int score_query_end, - const std::vector & required_instruction_spans) { + const std::vector & required_instruction_spans, + bool query_suffix_candidates) { if (!ctx.loaded) { set_last_error("drafter not loaded"); return {}; @@ -121,6 +122,8 @@ std::vector drafter_score_and_compress( return {}; } chunk_size = experiment.chunk_size; + experiment.query_suffix_candidates = + query_suffix_candidates && experiment.selection_active; if (!experiment.selection_active && !required_instruction_spans.empty()) { set_last_error( "PFlash instruction spans require strict budget selection"); @@ -145,12 +148,13 @@ std::vector drafter_score_and_compress( std::fprintf(stderr, "[pflash-select] config mode=%s active=%d chunk=%d " "query_parser=%s query_cap=%d query_actual=%d top_p=%.9g " - "top_k=%d input=%zu\n", + "top_k=%d suffix_candidates=%d input=%zu\n", luce::pflash::pflash_selection_mode_name(experiment.mode), (int) experiment.selection_active, experiment.chunk_size, luce::pflash::pflash_query_parser_name(experiment.query_parser), experiment.query_tokens, n_lookahead, experiment.top_p, - experiment.top_k, ids.size()); + experiment.top_k, (int) experiment.query_suffix_candidates, + ids.size()); std::fflush(stderr); } if (score_query_end < 0) { diff --git a/server/src/pflash/pflash_drafter.h b/server/src/pflash/pflash_drafter.h index 8340114b3..364d2d477 100644 --- a/server/src/pflash/pflash_drafter.h +++ b/server/src/pflash/pflash_drafter.h @@ -71,6 +71,8 @@ void free_drafter_weights(DrafterContext & ctx); // pool_kernel AvgPool kernel for score smoothing (default 13) // score_query_end exclusive end of the scorer query window in ids; // required (negative values are rejected) +// query_suffix_candidates strict selection only: tokens after the query +// window are scored candidates, not a kept suffix // // On failure returns empty vector + sets last_error. std::vector drafter_score_and_compress( @@ -82,6 +84,7 @@ std::vector drafter_score_and_compress( int pool_kernel = 13, int score_query_end = -1, const std::vector & - required_instruction_spans = {}); + required_instruction_spans = {}, + bool query_suffix_candidates = false); } // namespace luce::common diff --git a/server/src/pflash/pflash_selection.cpp b/server/src/pflash/pflash_selection.cpp index 03a20ce9d..5b975b6ce 100644 --- a/server/src/pflash/pflash_selection.cpp +++ b/server/src/pflash/pflash_selection.cpp @@ -86,14 +86,15 @@ bool pflash_chunk_is_structurally_required( int query_end, int input_tokens, const std::vector & - required_instruction_spans) noexcept { + required_instruction_spans, + bool query_suffix_structural) noexcept { if (begin < 0 || end <= begin || query_begin < 0 || query_end < query_begin || input_tokens < query_end || end > input_tokens) { return false; } const bool query_chunk = begin < query_end && end > query_begin; - const bool structural_suffix_chunk = + const bool structural_suffix_chunk = query_suffix_structural && begin < input_tokens && end > query_end; if (query_chunk || structural_suffix_chunk) return true; for (const auto & span : required_instruction_spans) { diff --git a/server/src/pflash/pflash_selection.h b/server/src/pflash/pflash_selection.h index 67ea85a56..dede1a51a 100644 --- a/server/src/pflash/pflash_selection.h +++ b/server/src/pflash/pflash_selection.h @@ -60,6 +60,9 @@ struct PFlashSelectionResult { std::string error; }; +// A chunk is kept whatever its score when it overlaps the query window or a +// required instruction span, or -- with ``query_suffix_structural`` -- any +// token after the query window. bool pflash_chunk_is_structurally_required( int begin, int end, @@ -67,7 +70,8 @@ bool pflash_chunk_is_structurally_required( int query_end, int input_tokens, const std::vector & - required_instruction_spans = {}) noexcept; + required_instruction_spans = {}, + bool query_suffix_structural = true) noexcept; bool validate_pflash_instruction_spans( const std::vector & spans, @@ -95,8 +99,8 @@ enum class PFlashScorer { Head, Legacy, Split }; struct PFlashSelectionConfig { PFlashSelectionMode mode = PFlashSelectionMode::Legacy; - // Chat-first default: the scorer query is the tail of the last message's - // content. latest_user stays selectable for benchmark experiments. + // Chat-first default: the scorer query is the tail of the latest user + // turn. latest_user stays selectable for benchmark experiments. PFlashQueryParser query_parser = PFlashQueryParser::ArbitraryTail; int chunk_size = 0; int query_tokens = 8; @@ -108,6 +112,11 @@ struct PFlashSelectionConfig { double split_fraction = 0.5; bool configured = false; bool selection_active = false; + // Per request, never from the environment: the tokens after the query + // window are candidates scored against it instead of a kept suffix (a + // chat whose latest user turn is followed by assistant and tool turns). + // The caller pins whatever of that suffix must stay. + bool query_suffix_candidates = false; }; // Segment probe: cut the context before every token whose boundary score is diff --git a/server/src/pflash/qwen35_drafter.cpp b/server/src/pflash/qwen35_drafter.cpp index 1dbd62979..04b50575c 100644 --- a/server/src/pflash/qwen35_drafter.cpp +++ b/server/src/pflash/qwen35_drafter.cpp @@ -707,9 +707,16 @@ std::vector qwen35_strict_score_and_compress( return {}; } { + // Keys are the context before the query window and, when the tokens + // after it are candidates too, the context after it. NoPE scoring has + // no position term, so a later key scores like an earlier one. std::vector m((size_t)n_lookahead * S, -INFINITY); for (int t = 0; t < n_lookahead; ++t) { std::fill_n(m.begin() + (size_t)t * S, (size_t)query_start, 0.0f); + if (experiment.query_suffix_candidates) { + std::fill_n(m.begin() + (size_t)t * S + query_end, + (size_t)(S - query_end), 0.0f); + } } ggml_backend_tensor_set(mask, m.data(), 0, m.size() * sizeof(float)); } @@ -877,7 +884,11 @@ std::vector qwen35_strict_score_and_compress( forced.push_back(span.end); } int boundaries_in_context = 0; - for (int t = 1; t < query_begin; ++t) { + for (int t = 1; t < S; ++t) { + if (t >= query_begin && + (t < query_end || !experiment.query_suffix_candidates)) { + continue; + } if (boundary[(size_t) t] > st.probe_threshold) ++boundaries_in_context; } const bool forced_probe = @@ -943,6 +954,10 @@ std::vector qwen35_drafter_score_and_compress( const char * legacy_scorer = std::getenv("PFLASH_QWEN35_LEGACY_SCORER"); const bool force_legacy = (legacy_scorer && std::string(legacy_scorer) == "1") || experiment.scorer == luce::pflash::PFlashScorer::Legacy; + // Only the block-15 head scores keys after the query window; the + // running-max scorer, alone or in the split, keeps the suffix. + luce::pflash::PFlashSelectionConfig suffix_kept = experiment; + suffix_kept.query_suffix_candidates = false; if (experiment.selection_active && experiment.scorer == luce::pflash::PFlashScorer::Split) { // Two scorers, one budget: the block-15 head ranks (and segments) @@ -951,7 +966,7 @@ std::vector qwen35_drafter_score_and_compress( std::vector head_segments; bool head_density = false; if (qwen35_strict_score_and_compress( - *st, ids, keep_ratio, n_lookahead, score_query_end, experiment, + *st, ids, keep_ratio, n_lookahead, score_query_end, suffix_kept, required_instruction_spans, &head_mass, &head_segments, &head_density).empty()) { return {}; @@ -959,7 +974,7 @@ std::vector qwen35_drafter_score_and_compress( std::vector other_scores; if (qwen35_score_and_compress(st->weights, ids, keep_ratio, chunk_size, n_lookahead, pool_kernel, score_query_end, - experiment, required_instruction_spans, + suffix_kept, required_instruction_spans, &other_scores).empty()) { return {}; } @@ -974,7 +989,7 @@ std::vector qwen35_drafter_score_and_compress( std::fflush(stderr); return select_pflash_chunks( ids, head_mass, keep_ratio, n_lookahead, score_query_end, - /*pool_kernel=*/1, experiment, required_instruction_spans, + /*pool_kernel=*/1, suffix_kept, required_instruction_spans, /*direct_mass=*/true, /*write_trace=*/true, head_segments.empty() ? nullptr : &head_segments, head_density, &other_scores, experiment.split_fraction); @@ -990,7 +1005,7 @@ std::vector qwen35_drafter_score_and_compress( } return qwen35_score_and_compress(st->weights, ids, keep_ratio, chunk_size, n_lookahead, pool_kernel, score_query_end, - experiment, + suffix_kept, required_instruction_spans); } diff --git a/server/src/qwen3/qwen3_backend.cpp b/server/src/qwen3/qwen3_backend.cpp index b7cdf1d6d..021ce83e4 100644 --- a/server/src/qwen3/qwen3_backend.cpp +++ b/server/src/qwen3/qwen3_backend.cpp @@ -976,7 +976,8 @@ ModelBackend::CompressResult Qwen3Backend::compress(const CompressRequest & req) result = CompressResult::from_compressed_ids(drafter_score_and_compress( drafter_ctx_, req.input_ids, req.keep_ratio, /*chunk_size=*/32, req.score_query_tokens, /*pool_kernel=*/13, - score_query_end, req.required_instruction_spans)); + score_query_end, req.required_instruction_spans, + req.query_suffix_candidates)); if (req.residency_action == DraftResidencyAction::ReleaseAfterUse) { free_drafter(); diff --git a/server/src/qwen35/qwen35_backend.cpp b/server/src/qwen35/qwen35_backend.cpp index c7f4b8731..cbb297393 100644 --- a/server/src/qwen35/qwen35_backend.cpp +++ b/server/src/qwen35/qwen35_backend.cpp @@ -1201,7 +1201,8 @@ std::vector Qwen35Backend::compress_batch( result.compressed_ids = drafter_score_and_compress( drafter_ctx_, request.input_ids, request.keep_ratio, /*chunk_size=*/32, request.score_query_tokens, /*pool_kernel=*/13, - score_query_end, request.required_instruction_spans); + score_query_end, request.required_instruction_spans, + request.query_suffix_candidates); result.ok = !result.compressed_ids.empty(); if (result.ok) { std::fprintf(stderr, "[compress] %zu -> %zu tokens\n", diff --git a/server/src/qwen35/qwen35_layer_split_adapter.cpp b/server/src/qwen35/qwen35_layer_split_adapter.cpp index 7a4a04aa0..895602473 100644 --- a/server/src/qwen35/qwen35_layer_split_adapter.cpp +++ b/server/src/qwen35/qwen35_layer_split_adapter.cpp @@ -1392,7 +1392,8 @@ Qwen35LayerSplitAdapter::compress(const ModelBackend::CompressRequest & req) { result.compressed_ids = drafter_score_and_compress( pflash_drafter_, req.input_ids, req.keep_ratio, /*chunk_size=*/32, req.score_query_tokens, /*pool_kernel=*/13, - score_query_end, req.required_instruction_spans); + score_query_end, req.required_instruction_spans, + req.query_suffix_candidates); result.ok = !result.compressed_ids.empty(); if (result.ok) { std::fprintf(stderr, "[target-split][compress] %zu -> %zu tokens\n", diff --git a/server/src/server/http_server.cpp b/server/src/server/http_server.cpp index 07b356a42..bfa610b7a 100644 --- a/server/src/server/http_server.cpp +++ b/server/src/server/http_server.cpp @@ -512,6 +512,7 @@ PflashChatTurnSpan pflash_chat_query_turn( size_t role_at = 0; size_t content_at = 0; size_t content_end = 0; + size_t close_end = 0; // past the closing marker, if any std::string role; bool closed = false; }; @@ -539,10 +540,12 @@ PflashChatTurnSpan pflash_chat_query_turn( // (DeepSeek user turns) -- at the next role marker. Nothing after // it leaves the turn open to the prompt end. size_t content_end = text.size(); + turn.close_end = text.size(); if (index + 1 < marks.size()) { const Mark & next = marks[index + 1]; content_end = next.at; turn.closed = !next.role || markers.role_starts_delimit; + turn.close_end = next.role ? next.at : next.at + next.len; } // Content ignores the whitespace the template wraps it in. while (content_at < content_end && is_space(text[content_at])) { @@ -571,14 +574,15 @@ PflashChatTurnSpan pflash_chat_query_turn( } // The query comes from the latest user turn; a conversation without one // falls back to its latest turn with content. - const Turn * query = nullptr; + size_t query_index = turns.size(); for (size_t index = usable; index-- > 0;) { const Turn & turn = turns[index]; if (turn.content_end <= turn.content_at) continue; - if (turn.role == "user") { query = &turn; break; } - if (!query) query = &turn; + if (turn.role == "user") { query_index = index; break; } + if (query_index == turns.size()) query_index = index; } - if (!query) return chosen; + if (query_index == turns.size()) return chosen; + const Turn * query = &turns[query_index]; // Content bounds in tokens: the first token starting at-or-after each // character offset, so a token merged across a boundary stays with the @@ -591,6 +595,11 @@ PflashChatTurnSpan pflash_chat_query_turn( chosen.role_begin = token_at_offset(decoded, query->role_at); chosen.content_begin = token_from(query->content_at); chosen.content_end = token_from(query->content_end); + chosen.turn_end = token_from(query->close_end); + chosen.generation_begin = usable < turns.size() + ? token_at_offset(decoded, turns.back().role_at) + : (int) prompt.size(); + chosen.later_turns = query_index + 1 < usable; if (chosen.role_begin > chosen.content_begin) { chosen.role_begin = chosen.content_begin; } @@ -3561,6 +3570,9 @@ std::string HttpServer::apply_pflash_compression( // the whole span mandatory; the scorer window is its tail. PFlashTokenSpan query_span{-1, -1}; std::string query_span_rule; + // Assistant and tool turns after the query's turn are context the query + // scores, not a kept suffix (strict selection only). + bool query_suffix_candidates = false; // Header ("<|im_start|>user\n") opening the query's turn, when the chat // markers resolved it — pinned mandatory so a compressed prompt keeps // the current turn's role envelope. @@ -3825,6 +3837,24 @@ std::string HttpServer::apply_pflash_compression( if (query_role_header.begin >= 0) { required_instruction_spans.push_back(query_role_header); } + // An agent loop puts assistant and tool turns after the + // user's: they compete for the budget like the context before + // the query. What stays is the rest of the query's turn + // through its closing marker, and the generation prompt. + if (chat_turn.valid() && chat_turn.later_turns && + query_span.begin >= chat_turn.content_begin && + query_span.end <= chat_turn.content_end) { + query_suffix_candidates = true; + if (chat_turn.turn_end > query_span.end) { + required_instruction_spans.push_back( + {query_span.end, chat_turn.turn_end}); + } + if (chat_turn.generation_begin < (int) drafter_ids.size()) { + required_instruction_spans.push_back( + {chat_turn.generation_begin, + (int) drafter_ids.size()}); + } + } required_instruction_spans = http_detail::canonicalize_pflash_token_spans( std::move(required_instruction_spans)); @@ -3924,6 +3954,7 @@ std::string HttpServer::apply_pflash_compression( compress_request.input_ids = std::move(drafter_ids); compress_request.required_instruction_spans = std::move(required_instruction_spans); + compress_request.query_suffix_candidates = query_suffix_candidates; compress_request.keep_ratio = http_detail::resolve_pflash_keep_ratio( pflash_keep_ratio(config_, prompt_tokens), req.session_id, sessions_); if (query_window.valid()) { @@ -3953,6 +3984,7 @@ std::string HttpServer::apply_pflash_compression( {"query_end", query_window.end}, {"query_span_begin", query_span.begin}, {"query_span_end", query_span.end}, + {"query_suffix_candidates", query_suffix_candidates}, {"requested_query_tokens", experiment.query_tokens}, {"required_text_count", req.pflash_required.size()}, {"expected_query_ids", expected_query_ids}, diff --git a/server/src/server/http_server.h b/server/src/server/http_server.h index adb1eae94..bfc00eaf3 100644 --- a/server/src/server/http_server.h +++ b/server/src/server/http_server.h @@ -343,6 +343,9 @@ PFlashTokenSpan pflash_decoded_text_span( // ``content_begin`` skips the role-name line ("<|im_start|>user\n") when the // family uses generic role markers; ``content_end`` sits before the turn's // closing marker. Both trim the whitespace the template wraps content in. +// ``turn_end`` sits past that closing marker; ``generation_begin`` is the +// generation prompt's marker (the prompt end when there is none); +// ``later_turns`` says assistant or tool turns sit between the two. // Offsets are token indices in ``prompt``'s own vocabulary. ``markers`` were // resolved on ``marker_tokenizer`` (the target model's); its marker strings // are searched in the decoded prompt text, so a drafter whose vocabulary @@ -352,6 +355,9 @@ struct PflashChatTurnSpan { int role_begin = -1; int content_begin = -1; int content_end = -1; + int turn_end = -1; + int generation_begin = -1; + bool later_turns = false; bool valid() const { return content_begin >= 0 && content_end > content_begin; diff --git a/server/test/test_pflash_selection.cpp b/server/test/test_pflash_selection.cpp index 035d92b9f..a08e35ded 100644 --- a/server/test/test_pflash_selection.cpp +++ b/server/test/test_pflash_selection.cpp @@ -112,6 +112,23 @@ TEST_CASE(PFlashSelectionFixture, structural_suffix_only_chunk_is_mandatory_and_ require_ordinals(result, {1, 2}); } +TEST_CASE(PFlashSelectionFixture, query_suffix_candidates_are_optional_unless_pinned) { + // Agent loop: turns after the query are scored context. Only the query + // window and pinned spans (the generation prompt) stay mandatory. + constexpr int input_tokens = 200; + constexpr int query_begin = 20; + constexpr int query_end = 30; + const std::vector pinned{{190, 200}}; + REQUIRE(pflash_chunk_is_structurally_required( + 0, 32, query_begin, query_end, input_tokens, pinned, false)); + REQUIRE(!pflash_chunk_is_structurally_required( + 32, 64, query_begin, query_end, input_tokens, pinned, false)); + REQUIRE(pflash_chunk_is_structurally_required( + 32, 64, query_begin, query_end, input_tokens, pinned, true)); + REQUIRE(pflash_chunk_is_structurally_required( + 160, 200, query_begin, query_end, input_tokens, pinned, false)); +} + TEST_CASE(PFlashSelectionFixture, instruction_overlap_is_mandatory_without_changing_optional_ranking) { constexpr int input_tokens = 120000; constexpr int query_begin = 119872; diff --git a/server/test/test_server_unit.cpp b/server/test/test_server_unit.cpp index 5681d18a1..305ce8c73 100644 --- a/server/test/test_server_unit.cpp +++ b/server/test/test_server_unit.cpp @@ -833,10 +833,13 @@ TEST_CASE(ServerUnitFixture, test_pflash_tail_query_window) { // chat-query turn the scorer would use, decoded as {header, content}. struct PflashRenderedQueryTurn { bool valid = false; + bool later_turns = false; std::string family; std::string header; std::string content; std::string after; + std::string closing; // content end .. turn end + std::string generation; // generation prompt .. prompt end }; static PflashRenderedQueryTurn pflash_rendered_query_turn( @@ -870,6 +873,11 @@ static PflashRenderedQueryTurn pflash_rendered_query_turn( prompt.begin() + turn.content_end}); out.after = tok.decode({prompt.begin() + turn.content_end, prompt.end()}); + out.closing = tok.decode({prompt.begin() + turn.content_end, + prompt.begin() + turn.turn_end}); + out.generation = tok.decode( + {prompt.begin() + turn.generation_begin, prompt.end()}); + out.later_turns = turn.later_turns; } } unlink(path.c_str()); @@ -922,6 +930,7 @@ TEST_CASE(ServerUnitFixture, // The rendered tail after the content is template machinery. TEST_ASSERT_MSG(turn.after.find("answer") == std::string::npos, family.name); + TEST_ASSERT_MSG(!turn.later_turns, family.name); } } } @@ -940,6 +949,10 @@ TEST_CASE(ServerUnitFixture, TEST_ASSERT(qwen.valid); TEST_ASSERT_MSG(qwen.content == "What is the answer?", qwen.content); TEST_ASSERT(qwen.header == "<|im_start|>user\n"); + TEST_ASSERT(qwen.later_turns); + TEST_ASSERT_MSG(qwen.closing == "<|im_end|>", qwen.closing); + TEST_ASSERT_MSG(qwen.generation == "<|im_start|>assistant\n\n\n\n\n", + qwen.generation); const auto deepseek = pflash_rendered_query_turn( messages, ChatFormat::DEEPSEEK4, true, @@ -6900,6 +6913,77 @@ TEST_CASE(ServerUnitFixture, unlink(path.c_str()); } +TEST_CASE(ServerUnitFixture, + test_pflash_strict_agent_turns_after_query_are_candidates) { + luce_test::ScopedEnvVar mode{"PFLASH_SELECT_MODE", "top_p"}; + + const std::vector messages{ + {"user", "What is the answer?", ""}, + {"assistant", "Sure.", ""}, + {"tool", "tool output here", "call-1"}, + }; + const std::string rendered = render_chat_template( + messages, ChatFormat::QWEN3, /*add_generation_prompt=*/true, + /*enable_thinking=*/true); + const std::string path = write_pflash_bpe_tokenizer_fixture( + {"What", " is", " the", " answer", "?", "user", "assistant", "\n", + "Sure", "."}, + 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; + request.messages = json::array({ + {{"role", "user"}, {"content", "What is the answer?"}}, + {{"role", "assistant"}, {"content", "Sure."}}, + {{"role", "tool"}, {"content", "tool output here"}, + {"tool_call_id", "call-1"}}, + }); + 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); + const auto & request = backend.last_request; + const auto & ids = request.input_ids; + TEST_ASSERT(request.query_suffix_candidates); + const int query_end = request.score_query_end; + const int query_begin = query_end - request.score_query_tokens; + TEST_ASSERT(tokenizer.decode({ids.begin() + query_begin, + ids.begin() + query_end}) + == "What is the answer?"); + // The generation prompt is pinned; the assistant and tool turns between + // the query and it are not. + bool generation_pinned = false; + for (const auto & span : request.required_instruction_spans) { + const std::string text = tokenizer.decode( + {ids.begin() + span.begin, ids.begin() + span.end}); + TEST_ASSERT_MSG(text.find("tool output") == std::string::npos, text); + TEST_ASSERT_MSG(text.find("Sure") == std::string::npos, text); + if (span.end == (int) ids.size() && + text.find("<|im_start|>assistant\n\n") != std::string::npos) { + generation_pinned = true; + } + } + TEST_ASSERT(generation_pinned); + unlink(path.c_str()); +} + TEST_CASE(ServerUnitFixture, test_pflash_legacy_chat_query_uses_last_user_turn) { // No strict-selection environment: the legacy selector derives the same From 7e57eab9ed44b3466e1434e09ebe5f52385ec9bd Mon Sep 17 00:00:00 2001 From: Graffioh Date: Tue, 22 Sep 2026 19:05:18 +0000 Subject: [PATCH 13/26] feat(pflash): spend the keep ratio on the droppable tokens Strict selection keeps system and developer messages, tool definitions, the query with its turn's envelope and the generation prompt, and it charged them against keep_ratio x the whole prompt. A 15K system prompt under the default 5% ratio exhausted the budget of anything shorter than 300K tokens and the request failed with mandatory_query_exceeds_budget. The server now counts the kept tokens with the selector's own mandatory-chunk rule and passes an effective ratio, (kept + keep_ratio x (input - kept) + 1) / input, so the selector budget and the target-token ceiling both add the kept part on top of the ratio's share of the rest. Auto mode compares the threshold with the droppable tokens, so a large system prompt plus a short chat is served as is. Instructions that alone would not fit the context (a document pasted into the system prompt) lose their pin and compete for the budget against the query. Live on the R9700 (27B target, 0.8B drafter, keep 0.3, 16K context): a 2K system-prompt document answers instead of failing; a 19K one compresses to 5.8K and answers correctly; the four chat shapes still answer, keeping their kept tokens on top (2237 -> 701 instead of 637). Co-Authored-By: Claude Opus 5.5 (1M context) --- server/README.md | 8 ++ server/src/server/http_server.cpp | 111 +++++++++++++++- server/src/server/http_server.h | 15 +++ server/test/test_server_unit.cpp | 210 ++++++++++++++++++++++++++++++ 4 files changed, 342 insertions(+), 2 deletions(-) diff --git a/server/README.md b/server/README.md index 75270e758..b14fceb30 100644 --- a/server/README.md +++ b/server/README.md @@ -396,6 +396,14 @@ and its turn's role header, and it runs on every turn of a multi-turn chat. In an agent loop the assistant and tool turns after the user's turn are scored against the query like the context before it; only the generation prompt is kept with them. + +The keep ratio applies to the droppable tokens only: what strict selection +keeps anyway (system and developer messages, tool definitions, the query and +its turn's envelope, the generation prompt) is added on top, so a long +system prompt no longer exhausts the budget. Auto mode compares +`--prefill-threshold` with the droppable tokens too. Instructions that alone +would not fit the context lose their pin and are scored like any other +context. A request's `pflash_query` string replaces the derived query and keeps its whole span; it is meant for benchmarks. `PFLASH_SELECT_QUERY_PARSER=latest_user` selects the benchmark parser, which finds the latest user message through diff --git a/server/src/server/http_server.cpp b/server/src/server/http_server.cpp index bfa610b7a..75fc584e6 100644 --- a/server/src/server/http_server.cpp +++ b/server/src/server/http_server.cpp @@ -611,6 +611,38 @@ bool pflash_full_cache_restore_allowed( return !selection_environment_present; } +int pflash_kept_tokens( + int input_tokens, + int chunk_size, + int query_begin, + int query_end, + const std::vector & kept_spans, + bool query_suffix_structural) noexcept { + if (input_tokens <= 0 || chunk_size <= 0) return 0; + int kept = 0; + for (int begin = 0; begin < input_tokens; begin += chunk_size) { + const int end = (std::min)(input_tokens, begin + chunk_size); + if (luce::pflash::pflash_chunk_is_structurally_required( + begin, end, query_begin, query_end, input_tokens, + kept_spans, query_suffix_structural)) { + kept += end - begin; + } + } + return kept; +} + +double pflash_effective_keep_ratio( + int input_tokens, int kept_tokens, double keep_ratio) noexcept { + if (input_tokens <= 0 || !std::isfinite(keep_ratio) || keep_ratio <= 0.0) { + return keep_ratio; + } + const int kept = (std::max)(0, (std::min)(kept_tokens, input_tokens)); + // One token of slack so flooring the budget never lands below `kept`. + const double budget = + (double) kept + keep_ratio * (double) (input_tokens - kept) + 1.0; + return (std::min)(1.0, budget / (double) input_tokens); +} + int pflash_target_token_ceiling( int original_target_tokens, double keep_ratio) noexcept { if (original_target_tokens < 0 || !std::isfinite(keep_ratio) || @@ -3578,6 +3610,9 @@ std::string HttpServer::apply_pflash_compression( // the current turn's role envelope. PFlashTokenSpan query_role_header{-1, -1}; std::vector required_instruction_spans; + // System, developer and tool-definition spans: kept verbatim like the + // required spans, except when they alone would not fit the context. + std::vector instruction_role_spans; // Chat-first scorer query: the latest user turn's content span, located // by the rendered prompt's own control markers. Feeds the strict tail // parser and the legacy window; unused when a benchmark parser @@ -3760,7 +3795,7 @@ std::string HttpServer::apply_pflash_compression( return "PFlash strict selection instruction mapping failed: " + boundary_error; } - required_instruction_spans.push_back(instruction_span); + instruction_role_spans.push_back(instruction_span); } if (!req.tools.is_null() && !req.tools.empty()) { @@ -3784,7 +3819,7 @@ std::string HttpServer::apply_pflash_compression( return "PFlash strict selection tool mapping failed: " "tools did not produce a retained prompt span"; } - required_instruction_spans.push_back(tool_span); + instruction_role_spans.push_back(tool_span); } // Client-declared literal text that must survive compression @@ -3858,9 +3893,15 @@ std::string HttpServer::apply_pflash_compression( required_instruction_spans = http_detail::canonicalize_pflash_token_spans( std::move(required_instruction_spans)); + instruction_role_spans = + http_detail::canonicalize_pflash_token_spans( + std::move(instruction_role_spans)); std::string instruction_error; if (!luce::pflash::validate_pflash_instruction_spans( required_instruction_spans, + (int) drafter_ids.size(), instruction_error) || + !luce::pflash::validate_pflash_instruction_spans( + instruction_role_spans, (int) drafter_ids.size(), instruction_error)) { return "PFlash strict selection instruction mapping failed: " + instruction_error; @@ -3950,6 +3991,60 @@ std::string HttpServer::apply_pflash_compression( } } + // Strict selection spends the keep ratio on the droppable tokens only: + // what it keeps anyway (instructions, tools, the query and its turn's + // envelope, the generation prompt) is already cheap -- a stable system + // prefix hits the prefix cache from the second turn on -- and must not + // exhaust the budget of the history it rides with. + int kept_tokens = 0; + if (experiment.selection_active && query_window.valid()) { + const int input_tokens = (int) drafter_ids.size(); + const int query_begin = query_window.end - query_window.tokens; + const auto kept_with = [&] ( + const std::vector & spans) { + return http_detail::pflash_kept_tokens( + input_tokens, experiment.chunk_size, query_begin, + query_window.end, spans, !query_suffix_candidates); + }; + const auto target_estimate = [&] (int drafter_tokens) { + return (int) std::ceil((double) prompt_tokens * + (double) drafter_tokens / (double) input_tokens); + }; + auto kept_spans = required_instruction_spans; + kept_spans.insert(kept_spans.end(), instruction_role_spans.begin(), + instruction_role_spans.end()); + kept_spans = http_detail::canonicalize_pflash_token_spans( + std::move(kept_spans)); + kept_tokens = kept_with(kept_spans); + // Instructions that alone overflow the context are data (a document + // pasted into the system prompt), not a preamble: they compete for + // the budget against the query like any other context. + if (!instruction_role_spans.empty() && config_.max_ctx > 0 && + target_estimate(kept_tokens) + req.max_output > config_.max_ctx) { + std::fprintf(stderr, + "[pflash-select] kept instructions do not fit the context " + "(~%d + %d > %d target tokens); scoring them as context\n", + target_estimate(kept_tokens), req.max_output, config_.max_ctx); + kept_spans = required_instruction_spans; + kept_tokens = kept_with(kept_spans); + } + required_instruction_spans = std::move(kept_spans); + // Auto mode compresses when the droppable part is long enough, not + // the whole prompt: a large system prompt plus a short chat has + // nothing worth selecting. + const int droppable_target = + prompt_tokens - target_estimate(kept_tokens); + if (config_.pflash_mode == ServerConfig::PflashMode::AUTO && + droppable_target < config_.pflash_threshold) { + std::fprintf(stderr, + "[pflash] skip-compress (droppable ~%d < threshold %d; " + "kept %d of %d drafter tokens)\n", + droppable_target, config_.pflash_threshold, kept_tokens, + input_tokens); + return {}; + } + } + ModelBackend::CompressRequest compress_request; compress_request.input_ids = std::move(drafter_ids); compress_request.required_instruction_spans = @@ -3957,6 +4052,18 @@ std::string HttpServer::apply_pflash_compression( compress_request.query_suffix_candidates = query_suffix_candidates; compress_request.keep_ratio = http_detail::resolve_pflash_keep_ratio( pflash_keep_ratio(config_, prompt_tokens), req.session_id, sessions_); + if (experiment.selection_active && query_window.valid()) { + const double effective = http_detail::pflash_effective_keep_ratio( + (int) compress_request.input_ids.size(), kept_tokens, + compress_request.keep_ratio); + std::fprintf(stderr, + "[pflash-select] kept=%d droppable=%d keep_ratio=%.6f " + "effective=%.6f\n", + kept_tokens, + (int) compress_request.input_ids.size() - kept_tokens, + (double) compress_request.keep_ratio, effective); + compress_request.keep_ratio = (float) effective; + } if (query_window.valid()) { compress_request.score_query_end = query_window.end; compress_request.score_query_tokens = query_window.tokens; diff --git a/server/src/server/http_server.h b/server/src/server/http_server.h index bfc00eaf3..1e4827613 100644 --- a/server/src/server/http_server.h +++ b/server/src/server/http_server.h @@ -389,6 +389,21 @@ std::string pflash_token_fingerprint( bool pflash_full_cache_restore_allowed( bool selection_environment_present) noexcept; +// Tokens the strict selector keeps whatever their score, counted the way it +// charges them: every fixed chunk overlapping the query window or a kept span +// and, with ``query_suffix_structural``, every chunk after the query. Probe +// segments cut exactly at those edges, so this is an upper bound for them. +int pflash_kept_tokens( + int input_tokens, + int chunk_size, + int query_begin, + int query_end, + const std::vector & kept_spans, + bool query_suffix_structural) noexcept; +// The keep ratio that spends ``keep_ratio`` on the droppable tokens only: +// (kept + keep_ratio * (input - kept) + 1) / input, capped at 1. +double pflash_effective_keep_ratio( + int input_tokens, int kept_tokens, double keep_ratio) noexcept; int pflash_target_token_ceiling( int original_target_tokens, double keep_ratio) noexcept; diff --git a/server/test/test_server_unit.cpp b/server/test/test_server_unit.cpp index 305ce8c73..6d8889e8b 100644 --- a/server/test/test_server_unit.cpp +++ b/server/test/test_server_unit.cpp @@ -1041,6 +1041,34 @@ TEST_CASE(ServerUnitFixture, test_pflash_selection_cache_and_continuation_policy TEST_ASSERT(!http_detail::pflash_full_cache_restore_allowed(true)); } +TEST_CASE(ServerUnitFixture, test_pflash_kept_tokens_follow_selector_chunks) { + // 100 tokens in chunks of 10; query [80, 85); instruction span [3, 12) + // touches chunks 0 and 1. + const std::vector kept{{3, 12}}; + // Suffix structural: chunks 0, 1, 8, 9. + TEST_ASSERT(http_detail::pflash_kept_tokens(100, 10, 80, 85, kept, true) == 40); + // Suffix scored: only the query's chunk after it. + TEST_ASSERT(http_detail::pflash_kept_tokens(100, 10, 80, 85, kept, false) == 30); + TEST_ASSERT(http_detail::pflash_kept_tokens(0, 10, 0, 0, kept, true) == 0); +} + +TEST_CASE(ServerUnitFixture, test_pflash_effective_keep_ratio_spends_on_droppable) { + // 1000 tokens, 400 kept, 5 %: budget 400 + 30 + 1 slack. + TEST_ASSERT(std::abs(http_detail::pflash_effective_keep_ratio( + 1000, 400, 0.05) - 0.431) < 1e-12); + // Nothing kept: the plain ratio plus the slack token. + TEST_ASSERT(std::abs(http_detail::pflash_effective_keep_ratio( + 1000, 0, 0.05) - 0.051) < 1e-12); + // Everything kept caps at 1. + TEST_ASSERT(http_detail::pflash_effective_keep_ratio(1000, 1000, 0.05) == 1.0); + // The floored budget never lands below the kept tokens. + for (int kept = 0; kept <= 997; kept += 7) { + const double ratio = + http_detail::pflash_effective_keep_ratio(997, kept, 0.013); + TEST_ASSERT((int) std::floor(997.0 * ratio) >= kept); + } +} + TEST_CASE(ServerUnitFixture, test_pflash_target_token_ceiling_floors) { TEST_ASSERT(http_detail::pflash_target_token_ceiling(7, 0.5) == 3); TEST_ASSERT(http_detail::pflash_target_token_ceiling(120000, 16384.0 / 120000.0) == 16384); @@ -6984,6 +7012,188 @@ TEST_CASE(ServerUnitFixture, unlink(path.c_str()); } +// Keeps a prefix of the input that fits the requested ratio, so the target +// ceiling check sees a real compression. +struct MockPflashBudgetBackend : MockBackend { + int compress_calls = 0; + CompressRequest last_request; + + CompressResult compress(const CompressRequest & request) override { + ++compress_calls; + last_request = request; + const size_t keep = (size_t) std::max(1.0, std::floor( + (double) request.input_ids.size() * request.keep_ratio) - 2.0); + return CompressResult::from_compressed_ids(std::vector( + request.input_ids.begin(), + request.input_ids.begin() + (long) std::min(keep, request.input_ids.size()))); + } +}; + +struct PflashSystemPromptCase { + std::string rendered; + json messages; + std::vector vocab; +}; + +static PflashSystemPromptCase pflash_long_system_prompt_case() { + std::string system; + for (int i = 0; i < 30; ++i) system += "You are helpful. "; + std::string history; + for (int i = 0; i < 30; ++i) history += "Sure. "; + PflashSystemPromptCase out; + out.rendered = render_chat_template( + {{"system", system, ""}, + {"user", history, ""}, + {"assistant", "Sure.", ""}, + {"user", "What is the answer?", ""}}, + // As the server renders it: ParsedRequest defaults to thinking on, + // and the instruction spans come from re-renders of this prompt. + ChatFormat::QWEN3, /*add_generation_prompt=*/true, + /*enable_thinking=*/true); + out.messages = json::array({ + {{"role", "system"}, {"content", system}}, + {{"role", "user"}, {"content", history}}, + {{"role", "assistant"}, {"content", "Sure."}}, + {{"role", "user"}, {"content", "What is the answer?"}}, + }); + out.vocab = {"What", " is", " the", " answer", "?", "user", "assistant", + "system", "\n", "You", " are", " helpful", ".", " ", "Sure"}; + return out; +} + +TEST_CASE(ServerUnitFixture, + test_pflash_strict_budget_spends_keep_ratio_on_droppable_tokens) { + // A system prompt larger than keep_ratio x prompt used to exhaust the + // budget (mandatory_query_exceeds_budget). It is kept and the ratio now + // applies to the rest. + luce_test::ScopedEnvVar mode{"PFLASH_SELECT_MODE", "budget_only"}; + luce_test::ScopedEnvVar chunk{"PFLASH_SELECT_CHUNK_SIZE", "4"}; + const auto prompt = pflash_long_system_prompt_case(); + const std::string path = + write_pflash_bpe_tokenizer_fixture(prompt.vocab, prompt.rendered); + Tokenizer tokenizer; + TEST_ASSERT(tokenizer.load_from_gguf(path.c_str())); + + auto backend_owner = std::make_unique(); + MockPflashBudgetBackend & backend = *backend_owner; + LuceEngine engine(std::move(backend_owner)); + ServerConfig config; + config.pflash_keep_ratio = 0.05f; + config.prefix_cache_cap = 0; + config.prefill_cache_cap = 0; + std::string error; + { + HttpServer server(engine, tokenizer, config); + server.set_drafter_tokenizer(&tokenizer); + ParsedRequest request; + request.format = ApiFormat::OPENAI_CHAT; + request.messages = prompt.messages; + request.prompt_tokens = tokenizer.encode(prompt.rendered); + error = HttpServerTestAccess::apply_pflash_compression(server, request); + } + TEST_ASSERT_MSG(error.empty(), error); + TEST_ASSERT(backend.compress_calls == 1); + const auto & request = backend.last_request; + const int input = (int) request.input_ids.size(); + int system_end = -1; + for (const auto & span : request.required_instruction_spans) { + const std::string text = tokenizer.decode( + {request.input_ids.begin() + span.begin, + request.input_ids.begin() + span.end}); + if (text.find("You are helpful") != std::string::npos) { + system_end = span.end; + } + } + TEST_ASSERT(system_end > 0); // the system prompt is still kept + // Its tokens are charged on top of 5 % of the droppable rest. + TEST_ASSERT(request.keep_ratio > (double) system_end / input); + TEST_ASSERT(request.keep_ratio < 1.0f); + unlink(path.c_str()); +} + +TEST_CASE(ServerUnitFixture, + test_pflash_strict_scores_instructions_that_overflow_context) { + // Instructions that alone do not fit the context are data: they lose + // their pin and compete for the budget. + luce_test::ScopedEnvVar mode{"PFLASH_SELECT_MODE", "budget_only"}; + luce_test::ScopedEnvVar chunk{"PFLASH_SELECT_CHUNK_SIZE", "4"}; + const auto prompt = pflash_long_system_prompt_case(); + const std::string path = + write_pflash_bpe_tokenizer_fixture(prompt.vocab, prompt.rendered); + Tokenizer tokenizer; + TEST_ASSERT(tokenizer.load_from_gguf(path.c_str())); + + auto backend_owner = std::make_unique(); + MockPflashBudgetBackend & backend = *backend_owner; + LuceEngine engine(std::move(backend_owner)); + ServerConfig config; + config.pflash_keep_ratio = 0.3f; + config.max_ctx = 200; // smaller than the system prompt + max_output + config.prefix_cache_cap = 0; + config.prefill_cache_cap = 0; + std::string error; + { + HttpServer server(engine, tokenizer, config); + server.set_drafter_tokenizer(&tokenizer); + ParsedRequest request; + request.format = ApiFormat::OPENAI_CHAT; + request.messages = prompt.messages; + request.prompt_tokens = tokenizer.encode(prompt.rendered); + request.max_output = 16; + error = HttpServerTestAccess::apply_pflash_compression(server, request); + } + TEST_ASSERT_MSG(error.empty(), error); + TEST_ASSERT(backend.compress_calls == 1); + const auto & request = backend.last_request; + for (const auto & span : request.required_instruction_spans) { + const std::string text = tokenizer.decode( + {request.input_ids.begin() + span.begin, + request.input_ids.begin() + span.end}); + TEST_ASSERT_MSG(text.find("You are helpful") == std::string::npos, text); + } + unlink(path.c_str()); +} + +TEST_CASE(ServerUnitFixture, + test_pflash_auto_threshold_counts_droppable_tokens) { + // Auto mode: the whole prompt clears the threshold, the droppable part + // does not -- nothing worth selecting, so the prompt goes through as is. + luce_test::ScopedEnvVar mode{"PFLASH_SELECT_MODE", "budget_only"}; + luce_test::ScopedEnvVar chunk{"PFLASH_SELECT_CHUNK_SIZE", "4"}; + const auto prompt = pflash_long_system_prompt_case(); + const std::string path = + write_pflash_bpe_tokenizer_fixture(prompt.vocab, prompt.rendered); + Tokenizer tokenizer; + TEST_ASSERT(tokenizer.load_from_gguf(path.c_str())); + const auto ids = tokenizer.encode(prompt.rendered); + + auto backend_owner = std::make_unique(); + MockPflashBudgetBackend & backend = *backend_owner; + LuceEngine engine(std::move(backend_owner)); + ServerConfig config; + config.pflash_mode = ServerConfig::PflashMode::AUTO; + config.pflash_threshold = (int) ids.size() - 20; + config.pflash_keep_ratio = 0.3f; + config.max_ctx = 8192; + 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; + request.messages = prompt.messages; + request.prompt_tokens = ids; + const auto prepared = + HttpServerTestAccess::prepare_prompt(server, request); + TEST_ASSERT_MSG(prepared.error.empty(), prepared.error); + TEST_ASSERT(!prepared.compressed); + TEST_ASSERT(prepared.tokens == ids); + } + TEST_ASSERT(backend.compress_calls == 0); + unlink(path.c_str()); +} + TEST_CASE(ServerUnitFixture, test_pflash_legacy_chat_query_uses_last_user_turn) { // No strict-selection environment: the legacy selector derives the same From 36da94420d9f9b5a23a6399539a2d6865867763c Mon Sep 17 00:00:00 2001 From: Graffioh Date: Tue, 22 Sep 2026 19:19:27 +0000 Subject: [PATCH 14/26] fix(pflash): never compress the system prompt A system prompt that alone does not fit the context now fails the request instead of losing its pin: PFlash does not compress system prompts. Developer messages and tool definitions that would not fit still lose their pin and are scored like any other context. Co-Authored-By: Claude Opus 5.5 (1M context) --- server/README.md | 7 ++-- server/src/server/http_server.cpp | 61 ++++++++++++++++++++++-------- server/test/test_server_unit.cpp | 62 +++++++++++++++++++++++-------- 3 files changed, 96 insertions(+), 34 deletions(-) diff --git a/server/README.md b/server/README.md index b14fceb30..5d0e374b3 100644 --- a/server/README.md +++ b/server/README.md @@ -401,9 +401,10 @@ The keep ratio applies to the droppable tokens only: what strict selection keeps anyway (system and developer messages, tool definitions, the query and its turn's envelope, the generation prompt) is added on top, so a long system prompt no longer exhausts the budget. Auto mode compares -`--prefill-threshold` with the droppable tokens too. Instructions that alone -would not fit the context lose their pin and are scored like any other -context. +`--prefill-threshold` with the droppable tokens too. PFlash never compresses +the system prompt: one that alone would not fit the context fails the +request. Developer messages and tool definitions that would not fit lose +their pin and are scored like any other context. A request's `pflash_query` string replaces the derived query and keeps its whole span; it is meant for benchmarks. `PFLASH_SELECT_QUERY_PARSER=latest_user` selects the benchmark parser, which finds the latest user message through diff --git a/server/src/server/http_server.cpp b/server/src/server/http_server.cpp index 75fc584e6..61879dbb4 100644 --- a/server/src/server/http_server.cpp +++ b/server/src/server/http_server.cpp @@ -3610,8 +3610,11 @@ std::string HttpServer::apply_pflash_compression( // the current turn's role envelope. PFlashTokenSpan query_role_header{-1, -1}; std::vector required_instruction_spans; - // System, developer and tool-definition spans: kept verbatim like the - // required spans, except when they alone would not fit the context. + // Kept verbatim like the required spans. The system prompt never loses + // its pin: when it alone does not fit the context the request fails. + // Developer messages and tool definitions that do not fit are scored + // like any other context. + std::vector system_spans; std::vector instruction_role_spans; // Chat-first scorer query: the latest user turn's content span, located // by the rendered prompt's own control markers. Feeds the strict tail @@ -3795,7 +3798,9 @@ std::string HttpServer::apply_pflash_compression( return "PFlash strict selection instruction mapping failed: " + boundary_error; } - instruction_role_spans.push_back(instruction_span); + (messages[instruction_index].role == "system" + ? system_spans : instruction_role_spans) + .push_back(instruction_span); } if (!req.tools.is_null() && !req.tools.empty()) { @@ -3896,12 +3901,18 @@ std::string HttpServer::apply_pflash_compression( instruction_role_spans = http_detail::canonicalize_pflash_token_spans( std::move(instruction_role_spans)); + system_spans = + http_detail::canonicalize_pflash_token_spans( + std::move(system_spans)); std::string instruction_error; if (!luce::pflash::validate_pflash_instruction_spans( required_instruction_spans, (int) drafter_ids.size(), instruction_error) || !luce::pflash::validate_pflash_instruction_spans( instruction_role_spans, + (int) drafter_ids.size(), instruction_error) || + !luce::pflash::validate_pflash_instruction_spans( + system_spans, (int) drafter_ids.size(), instruction_error)) { return "PFlash strict selection instruction mapping failed: " + instruction_error; @@ -4010,24 +4021,44 @@ std::string HttpServer::apply_pflash_compression( return (int) std::ceil((double) prompt_tokens * (double) drafter_tokens / (double) input_tokens); }; - auto kept_spans = required_instruction_spans; - kept_spans.insert(kept_spans.end(), instruction_role_spans.begin(), - instruction_role_spans.end()); - kept_spans = http_detail::canonicalize_pflash_token_spans( - std::move(kept_spans)); + const auto merged = [&] (bool with_instructions) { + auto spans = required_instruction_spans; + spans.insert(spans.end(), system_spans.begin(), system_spans.end()); + if (with_instructions) { + spans.insert(spans.end(), instruction_role_spans.begin(), + instruction_role_spans.end()); + } + return http_detail::canonicalize_pflash_token_spans( + std::move(spans)); + }; + const auto fits = [&] (int drafter_tokens) { + return config_.max_ctx <= 0 || + target_estimate(drafter_tokens) + req.max_output <= + config_.max_ctx; + }; + auto kept_spans = merged(/*with_instructions=*/true); kept_tokens = kept_with(kept_spans); - // Instructions that alone overflow the context are data (a document - // pasted into the system prompt), not a preamble: they compete for - // the budget against the query like any other context. - if (!instruction_role_spans.empty() && config_.max_ctx > 0 && - target_estimate(kept_tokens) + req.max_output > config_.max_ctx) { + // Developer messages and tool definitions that alone overflow the + // context are data (a document pasted into them), not a preamble: + // they compete for the budget against the query like any other + // context. The system prompt keeps its pin whatever its size. + if (!fits(kept_tokens) && !instruction_role_spans.empty()) { std::fprintf(stderr, "[pflash-select] kept instructions do not fit the context " - "(~%d + %d > %d target tokens); scoring them as context\n", + "(~%d + %d > %d target tokens); scoring developer and tool " + "spans as context\n", target_estimate(kept_tokens), req.max_output, config_.max_ctx); - kept_spans = required_instruction_spans; + kept_spans = merged(/*with_instructions=*/false); kept_tokens = kept_with(kept_spans); } + if (!fits(kept_with(system_spans))) { + return "PFlash strict selection: the system prompt alone does not " + "fit the context (~" + + std::to_string(target_estimate(kept_with(system_spans))) + + " + " + std::to_string(req.max_output) + " > " + + std::to_string(config_.max_ctx) + + " target tokens); PFlash does not compress system prompts"; + } required_instruction_spans = std::move(kept_spans); // Auto mode compresses when the droppable part is long enough, not // the whole prompt: a large system prompt plus a short chat has diff --git a/server/test/test_server_unit.cpp b/server/test/test_server_unit.cpp index 6d8889e8b..6ff9eadc5 100644 --- a/server/test/test_server_unit.cpp +++ b/server/test/test_server_unit.cpp @@ -7035,14 +7035,15 @@ struct PflashSystemPromptCase { std::vector vocab; }; -static PflashSystemPromptCase pflash_long_system_prompt_case() { +static PflashSystemPromptCase pflash_long_system_prompt_case( + const std::string & instruction_role = "system") { std::string system; for (int i = 0; i < 30; ++i) system += "You are helpful. "; std::string history; for (int i = 0; i < 30; ++i) history += "Sure. "; PflashSystemPromptCase out; out.rendered = render_chat_template( - {{"system", system, ""}, + {{instruction_role, system, ""}, {"user", history, ""}, {"assistant", "Sure.", ""}, {"user", "What is the answer?", ""}}, @@ -7051,7 +7052,7 @@ static PflashSystemPromptCase pflash_long_system_prompt_case() { ChatFormat::QWEN3, /*add_generation_prompt=*/true, /*enable_thinking=*/true); out.messages = json::array({ - {{"role", "system"}, {"content", system}}, + {{"role", instruction_role}, {"content", system}}, {{"role", "user"}, {"content", history}}, {{"role", "assistant"}, {"content", "Sure."}}, {{"role", "user"}, {"content", "What is the answer?"}}, @@ -7111,13 +7112,12 @@ TEST_CASE(ServerUnitFixture, unlink(path.c_str()); } -TEST_CASE(ServerUnitFixture, - test_pflash_strict_scores_instructions_that_overflow_context) { - // Instructions that alone do not fit the context are data: they lose - // their pin and compete for the budget. +static std::string pflash_overflowing_instruction_error( + const std::string & role, bool & compressed, + std::vector & kept_texts) { luce_test::ScopedEnvVar mode{"PFLASH_SELECT_MODE", "budget_only"}; luce_test::ScopedEnvVar chunk{"PFLASH_SELECT_CHUNK_SIZE", "4"}; - const auto prompt = pflash_long_system_prompt_case(); + const auto prompt = pflash_long_system_prompt_case(role); const std::string path = write_pflash_bpe_tokenizer_fixture(prompt.vocab, prompt.rendered); Tokenizer tokenizer; @@ -7128,7 +7128,7 @@ TEST_CASE(ServerUnitFixture, LuceEngine engine(std::move(backend_owner)); ServerConfig config; config.pflash_keep_ratio = 0.3f; - config.max_ctx = 200; // smaller than the system prompt + max_output + config.max_ctx = 200; // smaller than the instructions + max_output config.prefix_cache_cap = 0; config.prefill_cache_cap = 0; std::string error; @@ -7142,16 +7142,46 @@ TEST_CASE(ServerUnitFixture, request.max_output = 16; error = HttpServerTestAccess::apply_pflash_compression(server, request); } + compressed = backend.compress_calls == 1; + kept_texts.clear(); + if (compressed) { + const auto & request = backend.last_request; + for (const auto & span : request.required_instruction_spans) { + kept_texts.push_back(tokenizer.decode( + {request.input_ids.begin() + span.begin, + request.input_ids.begin() + span.end})); + } + } + unlink(path.c_str()); + return error; +} + +TEST_CASE(ServerUnitFixture, + test_pflash_strict_refuses_system_prompt_that_overflows_context) { + // The system prompt is never compressed: when it alone does not fit the + // context the request fails before the drafter runs. + bool compressed = true; + std::vector kept; + const std::string error = + pflash_overflowing_instruction_error("system", compressed, kept); + TEST_ASSERT_MSG(error.find("system prompt alone does not fit") != + std::string::npos, error); + TEST_ASSERT(!compressed); +} + +TEST_CASE(ServerUnitFixture, + test_pflash_strict_scores_developer_text_that_overflows_context) { + // A developer message too large for the context is data: it loses its + // pin and competes for the budget. + bool compressed = false; + std::vector kept; + const std::string error = + pflash_overflowing_instruction_error("developer", compressed, kept); TEST_ASSERT_MSG(error.empty(), error); - TEST_ASSERT(backend.compress_calls == 1); - const auto & request = backend.last_request; - for (const auto & span : request.required_instruction_spans) { - const std::string text = tokenizer.decode( - {request.input_ids.begin() + span.begin, - request.input_ids.begin() + span.end}); + TEST_ASSERT(compressed); + for (const auto & text : kept) { TEST_ASSERT_MSG(text.find("You are helpful") == std::string::npos, text); } - unlink(path.c_str()); } TEST_CASE(ServerUnitFixture, From 613c71f029eec92d184927e36ece320720ae4390 Mon Sep 17 00:00:00 2001 From: Graffioh Date: Tue, 22 Sep 2026 19:32:53 +0000 Subject: [PATCH 15/26] feat(pflash): multi-turn views that append, recall and rebuild Strict selection re-selected the whole history on every turn, so each turn's compressed prompt diverged from the last early on and the target prefilled it from scratch; evidence the model had just used could also vanish between turns. The server now remembers the prompt it served for each conversation, matched by raw prompt prefix up to the old generation prompt (most clients send no session id). The next turn serves that view plus the new turns. The fresh compression still runs: whatever it keeps for the new question that the view lacks is recalled as excerpts at the start of the new user turn, after everything the target cached, so the view always holds the fresh selection. The view is rebuilt from the fresh prompt when it outgrows twice the fresh prompt or the context. The selector reports its kept spans (CompressResult::kept_spans) for this, and the view asks for its prefix-cache snapshot at the start of its generation prompt, where the next turn branches off; the default second-to-last boundary lands on the system prompt of a first turn. PFLASH_CHAT_VIEW=0 turns views off. Live on the R9700 (27B target, 0.8B drafter, keep 0.1, 20.6K-token conversation): three turns answer correctly, turn 2 recalls the segment turn 1 dropped, and target prefill drops from 7.7 s to 2.8 s (turn 2, 2048 tokens restored) and from 2.7 s to 0.9 s (turn 3, 3584 restored). The drafter still reads the whole history each turn. Co-Authored-By: Claude Opus 5.5 (1M context) --- server/README.md | 11 + server/src/common/model_backend.h | 3 + server/src/deepseek4/deepseek4_backend.cpp | 2 + server/src/pflash/pflash_compress.cpp | 19 ++ server/src/pflash/pflash_compress.h | 7 + server/src/pflash/pflash_drafter.cpp | 2 + server/src/qwen35/qwen35_backend.cpp | 2 + .../src/qwen35/qwen35_layer_split_adapter.cpp | 2 + server/src/server/http_server.cpp | 285 ++++++++++++++++++ server/src/server/http_server.h | 63 ++++ server/test/test_server_unit.cpp | 187 ++++++++++++ 11 files changed, 583 insertions(+) diff --git a/server/README.md b/server/README.md index 5d0e374b3..3c4af655d 100644 --- a/server/README.md +++ b/server/README.md @@ -405,6 +405,17 @@ system prompt no longer exhausts the budget. Auto mode compares the system prompt: one that alone would not fit the context fails the request. Developer messages and tool definitions that would not fit lose their pin and are scored like any other context. + +Multi-turn chats keep a view: the prompt served for a turn is remembered, +and when the next request's prompt continues it (same tokens up to the old +generation prompt), PFlash serves that view plus the new turns instead of a +fresh compression, so the target restores its prefix-cache snapshot of the +view (taken at the start of its generation prompt) and prefills only what is +new. Segments the fresh selection keeps for the new question that the view +lacks are recalled as excerpts at the start of the new user turn. When the +view grows past twice the fresh prompt, or past the context, the fresh prompt +starts a new view. `PFLASH_CHAT_VIEW=0` serves the fresh compression every +turn. A request's `pflash_query` string replaces the derived query and keeps its whole span; it is meant for benchmarks. `PFLASH_SELECT_QUERY_PARSER=latest_user` selects the benchmark parser, which finds the latest user message through diff --git a/server/src/common/model_backend.h b/server/src/common/model_backend.h index 3acb8ef24..1174f81dd 100644 --- a/server/src/common/model_backend.h +++ b/server/src/common/model_backend.h @@ -284,6 +284,9 @@ struct ModelBackend { struct CompressResult { bool ok = false; std::vector compressed_ids; // surviving token IDs + // Strict selection: the input spans behind compressed_ids, ascending. + // Empty when the backend does not report them (remote drafter). + std::vector kept_spans; static CompressResult from_compressed_ids( std::vector ids) { diff --git a/server/src/deepseek4/deepseek4_backend.cpp b/server/src/deepseek4/deepseek4_backend.cpp index c4ebb489e..981d47d0c 100644 --- a/server/src/deepseek4/deepseek4_backend.cpp +++ b/server/src/deepseek4/deepseek4_backend.cpp @@ -12,6 +12,7 @@ #include "common/peer_access.h" #include "common/platform_env.h" #include "common/sampler.h" +#include "pflash/pflash_compress.h" #if defined(LUCE_BACKEND_HIP) || defined(GGML_USE_HIP) #include "common/gpu_runtime_compat.h" @@ -3184,6 +3185,7 @@ std::vector DeepSeek4Backend::compress_batch( score_query_end, request.required_instruction_spans, request.query_suffix_candidates); result.ok = !result.compressed_ids.empty(); + if (result.ok) result.kept_spans = pflash_last_kept_spans(); } if (load_request->residency_action == diff --git a/server/src/pflash/pflash_compress.cpp b/server/src/pflash/pflash_compress.cpp index 8ed55593c..387854084 100644 --- a/server/src/pflash/pflash_compress.cpp +++ b/server/src/pflash/pflash_compress.cpp @@ -179,6 +179,18 @@ void write_compression_trace( std::fclose(file); } +namespace { +thread_local std::vector g_last_kept_spans; +} // namespace + +const std::vector & pflash_last_kept_spans() { + return g_last_kept_spans; +} + +void pflash_clear_kept_spans() { + g_last_kept_spans.clear(); +} + std::vector select_pflash_chunks( const std::vector & ids, const std::vector & token_scores, @@ -281,11 +293,18 @@ std::vector select_pflash_chunks( std::vector output; output.reserve((size_t) selected.retained_tokens); + g_last_kept_spans.clear(); for (const auto & candidate : candidates) { if (!selected_mask[candidate.ordinal]) continue; output.insert(output.end(), ids.begin() + candidate.begin, ids.begin() + candidate.end); + if (!g_last_kept_spans.empty() && + g_last_kept_spans.back().end == candidate.begin) { + g_last_kept_spans.back().end = candidate.end; + } else { + g_last_kept_spans.push_back({candidate.begin, candidate.end}); + } } std::fprintf(stderr, diff --git a/server/src/pflash/pflash_compress.h b/server/src/pflash/pflash_compress.h index 9fa333852..efcdf080f 100644 --- a/server/src/pflash/pflash_compress.h +++ b/server/src/pflash/pflash_compress.h @@ -120,6 +120,13 @@ void write_compression_trace( const std::vector & compressed_ids, const PFlashTraceFields * trace_fields = nullptr); +// The spans the last strict selection on this thread kept, in input +// coordinates, ascending and merged. Cleared at the start of every +// drafter_score_and_compress call; empty when the call did not reach a +// strict selection (legacy selection, errors). +const std::vector & pflash_last_kept_spans(); +void pflash_clear_kept_spans(); + std::vector select_pflash_chunks( const std::vector & ids, const std::vector & token_scores, diff --git a/server/src/pflash/pflash_drafter.cpp b/server/src/pflash/pflash_drafter.cpp index 719659fb1..ab372215a 100644 --- a/server/src/pflash/pflash_drafter.cpp +++ b/server/src/pflash/pflash_drafter.cpp @@ -18,6 +18,7 @@ #include "qwen35_drafter.h" #include "pflash_selection.h" +#include "pflash_compress.h" #include "common/dspark_head.h" #include "internal.h" @@ -106,6 +107,7 @@ std::vector drafter_score_and_compress( int score_query_end, const std::vector & required_instruction_spans, bool query_suffix_candidates) { + pflash_clear_kept_spans(); if (!ctx.loaded) { set_last_error("drafter not loaded"); return {}; diff --git a/server/src/qwen35/qwen35_backend.cpp b/server/src/qwen35/qwen35_backend.cpp index cbb297393..4c85cbabe 100644 --- a/server/src/qwen35/qwen35_backend.cpp +++ b/server/src/qwen35/qwen35_backend.cpp @@ -26,6 +26,7 @@ #include "common/specla_mode.h" #include "qwen35_tensor_parallel.h" #include "pflash/pflash_drafter.h" +#include "pflash/pflash_compress.h" #include "pflash/kvflash_drafter_scorer.h" #include "ggml-cuda.h" @@ -1204,6 +1205,7 @@ std::vector Qwen35Backend::compress_batch( score_query_end, request.required_instruction_spans, request.query_suffix_candidates); result.ok = !result.compressed_ids.empty(); + if (result.ok) result.kept_spans = pflash_last_kept_spans(); if (result.ok) { std::fprintf(stderr, "[compress] %zu -> %zu tokens\n", request.input_ids.size(), result.compressed_ids.size()); diff --git a/server/src/qwen35/qwen35_layer_split_adapter.cpp b/server/src/qwen35/qwen35_layer_split_adapter.cpp index 895602473..4aae63430 100644 --- a/server/src/qwen35/qwen35_layer_split_adapter.cpp +++ b/server/src/qwen35/qwen35_layer_split_adapter.cpp @@ -15,6 +15,7 @@ #include "qwen35/qwen35_layer_split_dflash_target.h" #include "qwen35/prefill_helpers.h" #include "pflash/pflash_drafter.h" +#include "pflash/pflash_compress.h" #include "pflash/kvflash_drafter_scorer.h" #include "kv_quant.h" @@ -1395,6 +1396,7 @@ Qwen35LayerSplitAdapter::compress(const ModelBackend::CompressRequest & req) { score_query_end, req.required_instruction_spans, req.query_suffix_candidates); result.ok = !result.compressed_ids.empty(); + if (result.ok) result.kept_spans = pflash_last_kept_spans(); if (result.ok) { std::fprintf(stderr, "[target-split][compress] %zu -> %zu tokens\n", req.input_ids.size(), result.compressed_ids.size()); diff --git a/server/src/server/http_server.cpp b/server/src/server/http_server.cpp index 61879dbb4..293e29271 100644 --- a/server/src/server/http_server.cpp +++ b/server/src/server/http_server.cpp @@ -606,6 +606,122 @@ PflashChatTurnSpan pflash_chat_query_turn( return chosen; } +std::vector pflash_subtract_token_spans( + const std::vector & spans, + const std::vector & minus) { + std::vector out; + size_t cut = 0; + for (const auto & span : spans) { + int begin = span.begin; + while (cut < minus.size() && minus[cut].end <= begin) ++cut; + for (size_t index = cut; + index < minus.size() && minus[index].begin < span.end; ++index) { + if (minus[index].begin > begin) { + out.push_back({begin, minus[index].begin}); + } + begin = (std::max)(begin, minus[index].end); + } + if (begin < span.end) out.push_back({begin, span.end}); + } + return out; +} + +std::string pflash_recall_excerpt( + const std::string & text, + const std::vector & role_markers, + const std::vector & end_markers, + bool generic_role_lines) { + std::string out; + out.reserve(text.size()); + size_t at = 0; + while (at < text.size()) { + bool matched = false; + for (const auto & marker : role_markers) { + if (marker.empty() || text.compare(at, marker.size(), marker) != 0) { + continue; + } + at += marker.size(); + if (generic_role_lines) { + size_t name_end = at; + while (name_end < text.size() && name_end - at < 16 && + std::isalpha((unsigned char) text[name_end])) { + ++name_end; + } + if (name_end < text.size() && text[name_end] == '\n') { + at = name_end + 1; + } + } + out += '\n'; + matched = true; + break; + } + if (matched) continue; + for (const auto & marker : end_markers) { + if (marker.empty() || text.compare(at, marker.size(), marker) != 0) { + continue; + } + at += marker.size(); + out += '\n'; + matched = true; + break; + } + if (matched) continue; + out += text[at++]; + } + const size_t first = out.find_first_not_of(" \t\r\n"); + if (first == std::string::npos) return {}; + const size_t last = out.find_last_not_of(" \t\r\n"); + return out.substr(first, last - first + 1); +} + +bool PflashChatViewStore::find( + const std::vector & raw_tokens, + const std::vector & drafter_ids, + PflashChatView & out) const { + std::lock_guard lock(mutex_); + const PflashChatView * best = nullptr; + for (const auto & view : views_) { + if (view.raw_gen_begin <= 0 || view.drafter_gen_begin <= 0 || + (size_t) view.raw_gen_begin > raw_tokens.size() || + (size_t) view.drafter_gen_begin > drafter_ids.size()) { + continue; + } + if (!std::equal(view.raw_tokens.begin(), + view.raw_tokens.begin() + view.raw_gen_begin, + raw_tokens.begin()) || + !std::equal(view.drafter_ids.begin(), + view.drafter_ids.begin() + view.drafter_gen_begin, + drafter_ids.begin())) { + continue; + } + if (!best || view.raw_gen_begin > best->raw_gen_begin) best = &view; + } + if (!best) return false; + out = *best; + return true; +} + +void PflashChatViewStore::remember(PflashChatView view) { + std::lock_guard lock(mutex_); + // Drop the views this one continues: same conversation, older turn. + views_.erase(std::remove_if(views_.begin(), views_.end(), + [&view] (const PflashChatView & old) { + return old.raw_gen_begin > 0 && + old.raw_gen_begin <= view.raw_gen_begin && + (size_t) old.raw_gen_begin <= view.raw_tokens.size() && + std::equal(old.raw_tokens.begin(), + old.raw_tokens.begin() + old.raw_gen_begin, + view.raw_tokens.begin()); + }), views_.end()); + views_.push_back(std::move(view)); + while (views_.size() > capacity_) views_.erase(views_.begin()); +} + +size_t PflashChatViewStore::size() const { + std::lock_guard lock(mutex_); + return views_.size(); +} + bool pflash_full_cache_restore_allowed( bool selection_environment_present) noexcept { return !selection_environment_present; @@ -4266,6 +4382,12 @@ std::string HttpServer::apply_pflash_compression( std::to_string(target_ceiling) + ")"; } } + if (experiment.selection_active && messages_input && chat_turn.valid() && + !result.kept_spans.empty()) { + final_tokens = continue_pflash_chat_view( + req, compress_request.input_ids, chat_turn, result.kept_spans, + std::move(final_tokens), prepared.snapshot_cut); + } prepared.tokens = std::move(final_tokens); prepared.compressed = true; std::fprintf(stderr, @@ -4276,6 +4398,168 @@ std::string HttpServer::apply_pflash_compression( return {}; } +std::vector HttpServer::continue_pflash_chat_view( + const ParsedRequest & req, + const std::vector & drafter_ids, + const http_detail::PflashChatTurnSpan & turn, + const std::vector & kept_spans, + std::vector fresh, + int & snapshot_cut) { + snapshot_cut = -1; + const char * disabled = std::getenv("PFLASH_CHAT_VIEW"); + if (disabled && std::string(disabled) == "0") return fresh; + const int input = (int) drafter_ids.size(); + if (turn.generation_begin <= 0 || turn.generation_begin >= input) { + return fresh; + } + // The generation prompt, in target tokens: the raw prompt and every + // served prompt end with it (strict selection keeps it verbatim). + const auto generation = tokenizer_.encode(drafter_tokenizer_->decode( + std::vector(drafter_ids.begin() + turn.generation_begin, + drafter_ids.end()))); + const auto ends_with_generation = [&generation] ( + const std::vector & tokens) { + return !generation.empty() && tokens.size() > generation.size() && + std::equal(generation.begin(), generation.end(), + tokens.end() - (long) generation.size()); + }; + if (!ends_with_generation(req.prompt_tokens) || + !ends_with_generation(fresh)) { + return fresh; + } + const int raw_gen_begin = + (int) (req.prompt_tokens.size() - generation.size()); + + http_detail::PflashChatView next; + next.raw_tokens = req.prompt_tokens; + next.raw_gen_begin = raw_gen_begin; + next.drafter_ids = drafter_ids; + next.drafter_gen_begin = turn.generation_begin; + + http_detail::PflashChatView view; + const bool continues = + pflash_views_.find(req.prompt_tokens, drafter_ids, view); + const auto serve_fresh = [&] (const char * why, int turns) { + next.view_tokens = fresh; + next.view_gen_begin = (int) (fresh.size() - generation.size()); + next.spans = kept_spans; + next.turns = turns; + snapshot_cut = next.view_gen_begin; + std::fprintf(stderr, + "[pflash-view] %s turn=%d served=%zu\n", why, turns, fresh.size()); + std::fflush(stderr); + pflash_views_.remember(std::move(next)); + return std::move(fresh); + }; + if (!continues) return serve_fresh("fresh", 1); + if (view.raw_tokens == req.prompt_tokens) { + // The same prompt again (a retry): serve what was served. + std::fprintf(stderr, "[pflash-view] repeat turn=%d served=%zu\n", + view.turns, view.view_tokens.size()); + std::fflush(stderr); + snapshot_cut = view.view_gen_begin; + return view.view_tokens; + } + if (view.drafter_gen_begin >= turn.generation_begin || + view.view_gen_begin <= 0 || + (size_t) view.view_gen_begin > view.view_tokens.size()) { + return serve_fresh("fresh", 1); + } + + // Recall: what the fresh selection keeps for the new query that the view + // does not hold. Only a new user turn brings a new query; an agent step + // (assistant call plus tool output) appends without recalling. + std::vector recalled; + if (turn.role_begin >= view.drafter_gen_begin) { + auto in_view = view.spans; + in_view.push_back({view.drafter_gen_begin, input}); + recalled = http_detail::pflash_subtract_token_spans( + kept_spans, + http_detail::canonicalize_pflash_token_spans(std::move(in_view))); + } + std::string recall_block; + int recalled_tokens = 0; + if (!recalled.empty()) { + ChatMarkers markers; + std::vector role_markers; + std::vector end_markers; + bool generic_roles = false; + if (resolve_chat_markers(tokenizer_, markers)) { + const auto seq_text = [this] (const std::vector & seq) { + std::string text; + for (const int32_t id : seq) text += tokenizer_.token_text(id); + return text; + }; + for (const auto & seq : markers.next_role_starts) { + role_markers.push_back(seq_text(seq)); + } + for (const auto & seq : markers.end_msg_seqs) { + end_markers.push_back(seq_text(seq)); + } + generic_roles = !markers.role_starts_delimit && + markers.family != "laguna"; + } + std::string excerpts; + for (const auto & span : recalled) { + const std::string excerpt = http_detail::pflash_recall_excerpt( + drafter_tokenizer_->decode(std::vector( + drafter_ids.begin() + span.begin, + drafter_ids.begin() + span.end)), + role_markers, end_markers, generic_roles); + if (excerpt.empty()) continue; + if (!excerpts.empty()) excerpts += "\n\n"; + excerpts += excerpt; + recalled_tokens += span.end - span.begin; + } + if (!excerpts.empty()) { + recall_block = "[Earlier in this conversation]\n" + excerpts + + "\n[End of earlier excerpts]\n\n"; + } + } + + // The previous view without its generation prompt, then this turn's new + // tokens from where that generation prompt started; recalled excerpts + // open the new user turn's content, after everything the target cached. + const auto decode_range = [&] (int begin, int end) { + return drafter_tokenizer_->decode(std::vector( + drafter_ids.begin() + begin, drafter_ids.begin() + end)); + }; + const std::string delta = recall_block.empty() + ? decode_range(view.drafter_gen_begin, input) + : decode_range(view.drafter_gen_begin, turn.content_begin) + + recall_block + decode_range(turn.content_begin, input); + std::vector served(view.view_tokens.begin(), + view.view_tokens.begin() + view.view_gen_begin); + const auto delta_tokens = tokenizer_.encode(delta); + served.insert(served.end(), delta_tokens.begin(), delta_tokens.end()); + if (!ends_with_generation(served)) { + return serve_fresh("fresh", 1); + } + // Rebuild when the view outgrew what a fresh selection keeps, or the + // context: the fresh prompt starts a new view, prefilled from scratch. + const bool outgrown = served.size() > 2 * fresh.size() || + (config_.max_ctx > 0 && + (int) served.size() + req.max_output > config_.max_ctx); + if (outgrown) return serve_fresh("rebuild", view.turns + 1); + + auto spans = view.spans; + spans.push_back({view.drafter_gen_begin, input}); + spans.insert(spans.end(), recalled.begin(), recalled.end()); + next.view_tokens = served; + next.view_gen_begin = (int) (served.size() - generation.size()); + next.spans = http_detail::canonicalize_pflash_token_spans(std::move(spans)); + next.turns = view.turns + 1; + snapshot_cut = next.view_gen_begin; + std::fprintf(stderr, + "[pflash-view] continue turn=%d served=%zu reused=%d delta=%zu " + "recalled=%d fresh=%zu\n", + next.turns, served.size(), view.view_gen_begin, delta_tokens.size(), + recalled_tokens, fresh.size()); + std::fflush(stderr); + pflash_views_.remember(std::move(next)); + return served; +} + HttpServer::PreparedPrompt HttpServer::prepare_prompt( const ParsedRequest & req) { PreparedPrompt prepared; @@ -4444,6 +4728,7 @@ HttpServer::GenerationCacheState HttpServer::prepare_generation_cache( const bool prefer_tools_boundary = ppp_prefers_tools_boundary(config_.ppp_enabled, prefer_inline_snap); int forced_cut = req.pin_end_token; + if (forced_cut <= 0) forced_cut = prepared.snapshot_cut; // PPP runs *before* lookup. Default (rearrange=0): annotate a sticky // pin_end only — never mutate tokens. Token-level DiffPin rewrite diff --git a/server/src/server/http_server.h b/server/src/server/http_server.h index 1e4827613..8248ac36b 100644 --- a/server/src/server/http_server.h +++ b/server/src/server/http_server.h @@ -304,6 +304,52 @@ PFlashTokenSpan pflash_changed_token_span( std::vector canonicalize_pflash_token_spans( std::vector spans); +// The parts of ``spans`` that ``minus`` does not cover. Both canonical. +std::vector pflash_subtract_token_spans( + const std::vector & spans, + const std::vector & minus); + +// Text of a recalled segment made safe to quote inside a user turn: chat +// control markers are removed, with the role-name line that follows a +// generic role marker ("<|im_start|>assistant\n"), and the result is trimmed. +std::string pflash_recall_excerpt( + const std::string & text, + const std::vector & role_markers, + const std::vector & end_markers, + bool generic_role_lines); + +// A multi-turn PFlash view: what was served for one turn of a conversation, +// kept so the next turn can append to it instead of recompressing. The next +// request continues the view when its raw prompt starts with this one's up +// to the generation prompt, in target and drafter tokens alike. +struct PflashChatView { + std::vector raw_tokens; // target tokens, raw prompt + int raw_gen_begin = -1; // its generation prompt + std::vector drafter_ids; // drafter tokens, raw prompt + int drafter_gen_begin = -1; + std::vector view_tokens; // target tokens served + int view_gen_begin = -1; + std::vector spans; // raw drafter spans in the view + int turns = 0; +}; + +class PflashChatViewStore { +public: + explicit PflashChatViewStore(size_t capacity = 8) : capacity_(capacity) {} + // The stored view the prompt continues (longest match), if any. + bool find(const std::vector & raw_tokens, + const std::vector & drafter_ids, + PflashChatView & out) const; + // Store a view, replacing the one it continues. + void remember(PflashChatView view); + size_t size() const; + +private: + mutable std::mutex mutex_; + size_t capacity_; + std::vector views_; // most recent last +}; + // Find the last sufficiently-specific suffix of the user query inside the // rendered drafter-tokenized prompt. Public for model-free regression tests. PflashQueryWindow find_pflash_query_window( @@ -554,6 +600,10 @@ class HttpServer { int full_cache_served_tokens = -1; int full_cache_hit_slot = -1; int full_cache_hit_len = 0; + // Where to take this request's prefix-cache snapshot when nothing + // else asks for one: a multi-turn PFlash view sets the start of its + // generation prompt, where the next turn's prompt branches off. + int snapshot_cut = -1; int error_status = 0; std::string error; }; @@ -565,6 +615,16 @@ class HttpServer { PreparedPrompt & prepared); std::string apply_pflash_compression(const ParsedRequest & req, PreparedPrompt & prepared); + // Multi-turn: serve the conversation's previous view plus this turn's + // new tokens (and the segments the fresh selection wants that the view + // lacks) when one continues into this prompt; else the fresh prompt. + std::vector continue_pflash_chat_view( + const ParsedRequest & req, + const std::vector & drafter_ids, + const http_detail::PflashChatTurnSpan & turn, + const std::vector & kept_spans, + std::vector fresh, + int & snapshot_cut); bool forward_upstream(ServerJob * job, const ParsedRequest & req, const PreparedPrompt & prepared); @@ -713,6 +773,9 @@ class HttpServer { // Per-session adaptive keep_ratio bandit state. HttpServerSessions sessions_; + // Multi-turn PFlash views, matched by raw prompt prefix. + http_detail::PflashChatViewStore pflash_views_; + // Live status tracker (read by /status/json, written by worker thread). ServerStatus status_; diff --git a/server/test/test_server_unit.cpp b/server/test/test_server_unit.cpp index 6ff9eadc5..26cb594f8 100644 --- a/server/test/test_server_unit.cpp +++ b/server/test/test_server_unit.cpp @@ -1041,6 +1041,56 @@ TEST_CASE(ServerUnitFixture, test_pflash_selection_cache_and_continuation_policy TEST_ASSERT(!http_detail::pflash_full_cache_restore_allowed(true)); } +TEST_CASE(ServerUnitFixture, test_pflash_subtract_token_spans) { + const std::vector spans{{0, 10}, {20, 30}, {40, 50}}; + const std::vector minus{{5, 22}, {25, 26}, {40, 50}}; + const auto out = http_detail::pflash_subtract_token_spans(spans, minus); + TEST_ASSERT(out.size() == 3); + TEST_ASSERT(out[0].begin == 0 && out[0].end == 5); + TEST_ASSERT(out[1].begin == 22 && out[1].end == 25); + TEST_ASSERT(out[2].begin == 26 && out[2].end == 30); + TEST_ASSERT(http_detail::pflash_subtract_token_spans(spans, {}).size() == 3); + TEST_ASSERT(http_detail::pflash_subtract_token_spans(spans, {{0, 60}}).empty()); +} + +TEST_CASE(ServerUnitFixture, test_pflash_recall_excerpt_strips_chat_markers) { + const std::string text = + "tail of a fact<|im_end|>\n<|im_start|>assistant\nSure, noted." + "<|im_end|>\n<|im_start|>user\n"; + const std::string excerpt = http_detail::pflash_recall_excerpt( + text, {"<|im_start|>"}, {"<|im_end|>"}, /*generic_role_lines=*/true); + TEST_ASSERT_MSG(excerpt == "tail of a fact\n\n\nSure, noted.", excerpt); + TEST_ASSERT(http_detail::pflash_recall_excerpt( + "<|im_end|>\n", {"<|im_start|>"}, {"<|im_end|>"}, true).empty()); +} + +TEST_CASE(ServerUnitFixture, test_pflash_chat_view_store_matches_prompt_prefix) { + http_detail::PflashChatViewStore store(2); + http_detail::PflashChatView first; + first.raw_tokens = {1, 2, 3, 9, 9}; + first.raw_gen_begin = 3; + first.drafter_ids = {1, 2, 3, 9, 9}; + first.drafter_gen_begin = 3; + store.remember(first); + + http_detail::PflashChatView found; + // The next turn keeps the prefix before the old generation prompt. + TEST_ASSERT(store.find({1, 2, 3, 4, 5, 9, 9}, {1, 2, 3, 4, 5, 9, 9}, found)); + TEST_ASSERT(found.raw_gen_begin == 3); + // A different conversation does not match. + TEST_ASSERT(!store.find({1, 7, 3, 4}, {1, 7, 3, 4}, found)); + // A continuation replaces the view it continues. + http_detail::PflashChatView second = first; + second.raw_tokens = {1, 2, 3, 4, 5, 9, 9}; + second.raw_gen_begin = 5; + second.drafter_ids = second.raw_tokens; + second.drafter_gen_begin = 5; + store.remember(second); + TEST_ASSERT(store.size() == 1); + TEST_ASSERT(store.find({1, 2, 3, 4, 5, 6, 9}, {1, 2, 3, 4, 5, 6, 9}, found)); + TEST_ASSERT(found.raw_gen_begin == 5); +} + TEST_CASE(ServerUnitFixture, test_pflash_kept_tokens_follow_selector_chunks) { // 100 tokens in chunks of 10; query [80, 85); instruction span [3, 12) // touches chunks 0 and 1. @@ -7224,6 +7274,143 @@ TEST_CASE(ServerUnitFixture, unlink(path.c_str()); } +// Keeps the required spans, the query window through the end, and whatever +// ``pick`` adds; reports the kept spans like the in-process drafter does. +struct MockPflashSpanBackend : MockBackend { + int compress_calls = 0; + CompressRequest last_request; + std::function(const CompressRequest &)> pick; + + CompressResult compress(const CompressRequest & request) override { + ++compress_calls; + last_request = request; + auto spans = request.required_instruction_spans; + spans.push_back({request.score_query_end - request.score_query_tokens, + (int) request.input_ids.size()}); + if (pick) { + const auto extra = pick(request); + spans.insert(spans.end(), extra.begin(), extra.end()); + } + spans = http_detail::canonicalize_pflash_token_spans(std::move(spans)); + CompressResult result; + for (const auto & span : spans) { + result.compressed_ids.insert(result.compressed_ids.end(), + request.input_ids.begin() + span.begin, + request.input_ids.begin() + span.end); + } + result.kept_spans = spans; + result.ok = !result.compressed_ids.empty(); + return result; + } +}; + +TEST_CASE(ServerUnitFixture, + test_pflash_chat_view_appends_turns_and_recalls_missing_segments) { + luce_test::ScopedEnvVar mode{"PFLASH_SELECT_MODE", "budget_only"}; + luce_test::ScopedEnvVar view_env{"PFLASH_CHAT_VIEW", nullptr}; + + std::string system; + for (int i = 0; i < 20; ++i) system += "You are helpful. "; + const std::string document = + "alpha facts live here. filler filler filler. beta facts live here."; + const std::vector turn1{ + {"system", system, ""}, + {"user", document + " Question one?", ""}, + }; + auto turn2 = turn1; + turn2.push_back({"assistant", "Answer one.", ""}); + turn2.push_back({"user", "Question two?", ""}); + const auto render = [] (const std::vector & messages) { + return render_chat_template(messages, ChatFormat::QWEN3, + /*add_generation_prompt=*/true, + /*enable_thinking=*/true); + }; + const auto to_json = [] (const std::vector & messages) { + json out = json::array(); + for (const auto & message : messages) { + out.push_back({{"role", message.role}, {"content", message.content}}); + } + return out; + }; + const std::string path = write_pflash_bpe_tokenizer_fixture( + {"alpha", " facts", "beta", " live", " here", ".", " filler", + "Question", " one", " two", "?", "Answer", "user", "assistant", + "system", "\n", "You", " are", " helpful"}, + render(turn2) + + "[Earlier in this conversation]\n[End of earlier excerpts]\n"); + Tokenizer tokenizer; + TEST_ASSERT(tokenizer.load_from_gguf(path.c_str())); + + auto backend_owner = std::make_unique(); + MockPflashSpanBackend & backend = *backend_owner; + const char * wanted = "alpha facts"; + backend.pick = [&] (const ModelBackend::CompressRequest & request) { + const auto span = http_detail::pflash_decoded_text_span( + tokenizer, request.input_ids, 0, (int) request.input_ids.size(), + wanted); + return span.begin < 0 ? std::vector{} + : std::vector{span}; + }; + LuceEngine engine(std::move(backend_owner)); + ServerConfig config; + config.pflash_mode = ServerConfig::PflashMode::ALWAYS; + config.pflash_keep_ratio = 1.0f; + config.max_ctx = 8192; + config.prefix_cache_cap = 0; + config.prefill_cache_cap = 0; + std::vector served1; + std::vector served2; + std::vector served3; + { + HttpServer server(engine, tokenizer, config); + server.set_drafter_tokenizer(&tokenizer); + const auto run = [&] (const std::vector & messages, + std::vector & served) { + ParsedRequest request; + request.format = ApiFormat::OPENAI_CHAT; + request.messages = to_json(messages); + request.prompt_tokens = tokenizer.encode(render(messages)); + const auto prepared = + HttpServerTestAccess::prepare_prompt(server, request); + TEST_ASSERT_MSG(prepared.error.empty(), prepared.error); + TEST_ASSERT(prepared.compressed); + served = prepared.tokens; + // The snapshot lands where the next turn's prompt branches off. + const auto generation = + tokenizer.encode("<|im_start|>assistant\n\n"); + TEST_ASSERT(prepared.snapshot_cut == + (int) (served.size() - generation.size())); + }; + run(turn1, served1); + wanted = "beta facts"; // the new query wants what turn 1 dropped + run(turn2, served2); + run(turn2, served3); // a retry serves the same view + } + const std::string text1 = tokenizer.decode(served1); + const std::string text2 = tokenizer.decode(served2); + TEST_ASSERT(text1.find("alpha facts") != std::string::npos); + TEST_ASSERT(text1.find("beta facts") == std::string::npos); + + // Turn 2 extends turn 1's served prompt: everything before its + // generation prompt is reused token for token. + const auto generation = tokenizer.encode("<|im_start|>assistant\n\n"); + TEST_ASSERT(served1.size() > generation.size()); + const size_t reused = served1.size() - generation.size(); + TEST_ASSERT(served2.size() > reused); + TEST_ASSERT(std::equal(served1.begin(), served1.begin() + (long) reused, + served2.begin())); + // The recalled segment opens the new user turn, after the answer. + const size_t answer = text2.find("Answer one."); + const size_t recall = text2.find("[Earlier in this conversation]"); + const size_t question = text2.find("Question two?"); + TEST_ASSERT(answer != std::string::npos); + TEST_ASSERT(recall != std::string::npos && recall > answer); + TEST_ASSERT(text2.find("beta facts", recall) != std::string::npos); + TEST_ASSERT(question != std::string::npos && question > recall); + TEST_ASSERT(served3 == served2); + unlink(path.c_str()); +} + TEST_CASE(ServerUnitFixture, test_pflash_legacy_chat_query_uses_last_user_turn) { // No strict-selection environment: the legacy selector derives the same From 55475e6fba479e268fadaa1598629b056d003f55 Mon Sep 17 00:00:00 2001 From: Graffioh Date: Tue, 22 Sep 2026 19:51:26 +0000 Subject: [PATCH 16/26] feat(pflash): per-conversation drafter scoring sessions The Qwen3.5 drafter re-read the whole prompt on every call, so every turn of a conversation paid a forward over its entire history (about 4 s at 20K tokens on the R9700) even when only the last few hundred tokens were new. Blocks 0..14 read left to right, so the cache state and the block-15 keys of a shared prefix never change, and NoPE keys carry no position. The strict scorer now keeps up to PFLASH_DRAFTER_SESSIONS (default 2) scoring sessions: a blocks-0..14 cache with headroom, the block-15 keys, the probe's raw logits, the last query window's block-14 rows, and a recurrent-state checkpoint 64 tokens before the prompt's end. A prompt that shares a session's prefix resumes from the session's end or, when it diverged before it (the previous turn's generation prompt), from the checkpoint, runs only its new tokens, and scores the query against every stored key. The query rows are reused while the query stays put (an agent step appends tool output after the same user turn). Anything else starts the least recently used session over; any failure forgets the session's prompt. Sessions live with the loaded drafter. Live on the R9700 (27B target, 0.8B drafter resident, keep 0.1, 20.6K tokens, three turns): drafter forward 3.9 s -> 0.04 s on turn 2 (resumed at 20571, 90 new) and 3.3 s -> 0.13 s on turn 3, with the same selections, served prompts and answers as scoring from scratch; turns take 3.9 s and 3.3 s instead of 8.8 s and 6.2 s. An agent step resumes at 10993 with the stored query rows and runs only the 8435-token tool output (1.9 s). Co-Authored-By: Claude Opus 5.5 (1M context) --- server/README.md | 12 + server/src/pflash/qwen35_drafter.cpp | 463 +++++++++++++++++++++------ server/src/pflash/qwen35_drafter.h | 31 ++ server/src/pflash/qwen35_loader.cpp | 4 + 4 files changed, 418 insertions(+), 92 deletions(-) diff --git a/server/README.md b/server/README.md index 3c4af655d..b31efe96b 100644 --- a/server/README.md +++ b/server/README.md @@ -416,6 +416,18 @@ lacks are recalled as excerpts at the start of the new user turn. When the view grows past twice the fresh prompt, or past the context, the fresh prompt starts a new view. `PFLASH_CHAT_VIEW=0` serves the fresh compression every turn. + +The drafter keeps a scoring session per conversation +(`PFLASH_DRAFTER_SESSIONS`, default 2, least recently used evicted; 0 scores +every prompt from scratch): the cache of blocks 0-14, the block-15 keys and +the probe logits of the prompt it last scored, with the recurrent state +checkpointed 64 tokens before its end. A prompt that shares that prefix runs +only its new tokens through the drafter, from the end or from the +checkpoint (the previous turn's generation prompt is replaced), and the new +query scores against every stored key. Sessions live with the loaded drafter, +so they pay off with `--draft-residency persistent` (and `--prefill-skip-park` +where the target and drafter fit together); the default releases the drafter +after each compression. A request's `pflash_query` string replaces the derived query and keeps its whole span; it is meant for benchmarks. `PFLASH_SELECT_QUERY_PARSER=latest_user` selects the benchmark parser, which finds the latest user message through diff --git a/server/src/pflash/qwen35_drafter.cpp b/server/src/pflash/qwen35_drafter.cpp index 04b50575c..56a8e7ac0 100644 --- a/server/src/pflash/qwen35_drafter.cpp +++ b/server/src/pflash/qwen35_drafter.cpp @@ -519,6 +519,184 @@ std::vector qwen35_score_and_compress( return out_ids; } +void free_qwen35_scoring_session(Qwen35ScoringSession & session) { + free_target_cache(session.cache); + session.cache = TargetCache{}; + if (session.key_buf) ggml_backend_buffer_free(session.key_buf); + if (session.key_ctx) ggml_free(session.key_ctx); + session.key_buf = nullptr; + session.key_ctx = nullptr; + session.keys = nullptr; + session.capacity = 0; + session.ids.clear(); + session.checkpoint = 0; + session.probe_raw.clear(); + session.subunit_raw.clear(); + session.query_begin = session.query_end = -1; + session.query_rows.clear(); +} + +namespace { + +int scoring_session_limit() { + const char * raw = std::getenv("PFLASH_DRAFTER_SESSIONS"); + if (!raw || !*raw) return 2; + char * end = nullptr; + const long value = std::strtol(raw, &end, 10); + if (end == raw || *end != '\0' || value < 0) return 2; + return (int) std::min(value, 16); +} + +bool allocate_scoring_session(TargetWeights & w, int capacity, + Qwen35ScoringSession & session) { + { + ScopedKvTq3Off tq3_off; + if (!create_target_cache_partial(w, capacity, 0, w.backend, session.cache, + /*prefill_only=*/true, 0, kQwen35HeadBlock, + /*allocate_target_feat=*/false)) { + return false; + } + } + if (!ensure_ssm_snapshot(session.cache, w.backend)) { + free_qwen35_scoring_session(session); + return false; + } + ggml_init_params kp{}; + kp.mem_size = ggml_tensor_overhead() + 1024; + kp.no_alloc = true; + session.key_ctx = ggml_init(kp); + if (session.key_ctx) { + session.keys = ggml_new_tensor_3d(session.key_ctx, GGML_TYPE_F32, + w.n_embd_head_k, w.n_head_kv, capacity); + session.key_buf = ggml_backend_alloc_ctx_tensors(session.key_ctx, w.backend); + } + if (!session.key_buf) { + free_qwen35_scoring_session(session); + return false; + } + session.capacity = capacity; + return true; +} + +void zero_recurrent_state(TargetCache & cache) { + for (size_t i = 0; i < cache.ssm_state.size(); ++i) { + if (cache.ssm_state[i]) { + ggml_backend_tensor_memset(cache.ssm_state[i], 0, 0, + ggml_nbytes(cache.ssm_state[i])); + } + if (i < cache.conv_state.size() && cache.conv_state[i]) { + ggml_backend_tensor_memset(cache.conv_state[i], 0, 0, + ggml_nbytes(cache.conv_state[i])); + } + } +} + +// The session to score ``ids`` with, and the token it resumes from: the +// longest prefix a stored session already covers -- its live end, or its +// checkpoint when the prompt diverged before the end (the previous turn's +// generation prompt) -- provided the query rows and probe logits the +// scoring needs are covered too. Otherwise the least recently used session +// (or ``scratch`` with sessions off) starts over from token 0. +Qwen35ScoringSession * acquire_scoring_session( + Qwen35DrafterState & st, + const std::vector & ids, + int query_start, + int query_end, + bool need_probe, + bool need_subunit, + int & resume, + std::unique_ptr & scratch) { + TargetWeights & w = st.weights; + const int S = (int) ids.size(); + const int limit = scoring_session_limit(); + const size_t row_floats = (size_t) w.n_embd * (size_t) (query_end - query_start); + resume = 0; + + Qwen35ScoringSession * best = nullptr; + bool best_restore = false; + for (auto & owned : st.sessions) { + Qwen35ScoringSession * session = owned.get(); + if (!session || session->capacity < S || session->ids.empty() || + session->keys_trained != st.head_loaded) { + continue; + } + const size_t n = std::min(session->ids.size(), ids.size()); + const int shared = (int) (std::mismatch(session->ids.begin(), + session->ids.begin() + (long) n, ids.begin()).first - + session->ids.begin()); + int r = 0; + bool restore = false; + if (shared == (int) session->ids.size()) { + r = shared; + } else if (session->checkpoint > 0 && session->checkpoint <= shared) { + r = session->checkpoint; + restore = true; + } + if (r > query_start) { + const bool rows = session->query_begin == query_start && + session->query_end == query_end && query_end <= shared && + session->query_rows.size() == row_floats; + if (!rows) { + r = session->checkpoint > 0 && session->checkpoint <= query_start && + session->checkpoint <= shared + ? session->checkpoint : 0; + restore = r > 0; + } + } + if ((need_probe && (int) session->probe_raw.size() < r) || + (need_subunit && (int) session->subunit_raw.size() < r)) { + r = 0; + } + if (r > resume) { + resume = r; + best = session; + best_restore = restore; + } + } + if (best) { + if (best_restore && !restore_ssm_state(best->cache, w.backend)) { + resume = 0; + } else { + return best; + } + } + + Qwen35ScoringSession * target = best; + if (!target) { + if (limit == 0) { + scratch = std::make_unique(); + target = scratch.get(); + } else if ((int) st.sessions.size() < limit) { + st.sessions.push_back(std::make_unique()); + target = st.sessions.back().get(); + } else { + target = std::min_element(st.sessions.begin(), st.sessions.end(), + [] (const auto & a, const auto & b) { + return a->last_used < b->last_used; + })->get(); + } + } + if (target->capacity < S) { + free_qwen35_scoring_session(*target); + // Headroom so the next turns append without reallocating. + const int capacity = limit == 0 ? S : S + S / 2 + 4096; + if (!allocate_scoring_session(w, capacity, *target)) { + set_last_error("qwen35 scoring session allocation failed"); + return nullptr; + } + } + zero_recurrent_state(target->cache); + target->ids.clear(); + target->checkpoint = 0; + target->probe_raw.clear(); + target->subunit_raw.clear(); + target->query_begin = target->query_end = -1; + target->query_rows.clear(); + return target; +} + +} // namespace + // Scoring-head selection for the Qwen3.5-0.8B drafter: run blocks 0..14, then // score every context token against the query window with block 15's NoPE // Q/K (or a trained replacement) and select chunks by attention mass. This @@ -557,75 +735,112 @@ std::vector qwen35_strict_score_and_compress( } const int query_start = query_end - n_lookahead; const TargetLayer & L = w.layers[(size_t)kQwen35HeadBlock]; + const bool use_probe = st.probe_loaded && + experiment.segmentation != luce::pflash::PFlashSegmentation::Fixed; auto t0 = std::chrono::steady_clock::now(); - TargetCache cache; - { - ScopedKvTq3Off tq3_off; - if (!create_target_cache(w, S, 0, w.backend, cache, true)) { - return {}; + int resume = 0; + std::unique_ptr scratch; + Qwen35ScoringSession * session = acquire_scoring_session( + st, ids, query_start, query_end, use_probe, + use_probe && st.probe_sub_fc2_w != nullptr, resume, scratch); + if (!session) return {}; + // A session is released (freed or kept) on every exit below. + struct SessionExit { + std::unique_ptr & scratch; + ~SessionExit() { + if (scratch) free_qwen35_scoring_session(*scratch); } - } + } session_exit{scratch}; + TargetCache & cache = session->cache; + const int n_new = S - resume; + // The next turn replaces this prompt's generation prompt; checkpoint the + // recurrent state a little before the end so it can resume there. The + // same prompt again (a retry) keeps the checkpoint it has. + const int checkpoint = n_new == 0 && session->checkpoint > 0 + ? session->checkpoint : std::max(resume, S - 64); ggml_init_params act_ip{}; act_ip.mem_size = (size_t)8 * ggml_tensor_overhead() + 4096; act_ip.no_alloc = true; ggml_context * act_ctx = ggml_init(act_ip); if (!act_ctx) { - free_target_cache(cache); + session->ids.clear(); set_last_error("qwen35 drafter activation ctx init failed"); return {}; } - ggml_tensor * act_in = ggml_new_tensor_2d(act_ctx, GGML_TYPE_F32, hidden, S); - ggml_tensor * act_out = ggml_new_tensor_2d(act_ctx, GGML_TYPE_F32, hidden, S); + ggml_tensor * act_in = ggml_new_tensor_2d(act_ctx, GGML_TYPE_F32, hidden, std::max(1, n_new)); + ggml_tensor * act_out = ggml_new_tensor_2d(act_ctx, GGML_TYPE_F32, hidden, std::max(1, n_new)); ggml_backend_buffer_t act_buf = ggml_backend_alloc_ctx_tensors(act_ctx, w.backend); if (!act_buf) { ggml_free(act_ctx); - free_target_cache(cache); + session->ids.clear(); set_last_error("qwen35 drafter activation allocation failed"); return {}; } + // Any failure below leaves the session's state half-written: forget its + // prompt so the next call starts it over. auto cleanup = [&]() { ggml_backend_buffer_free(act_buf); ggml_free(act_ctx); - free_target_cache(cache); + }; + auto fail = [&](const char * message) -> std::vector { + cleanup(); + session->ids.clear(); + set_last_error(message); + return {}; }; { const int batch = 2048; std::vector emb((size_t)hidden * batch); - for (int i = 0; i < S; i += batch) { - const int n = std::min(batch, S - i); - if (!w.embedder.embed(ids.data() + i, n, emb.data())) { - cleanup(); - set_last_error("qwen35 drafter embedding failed"); - return {}; + for (int i = 0; i < n_new; i += batch) { + const int n = std::min(batch, n_new - i); + if (!w.embedder.embed(ids.data() + resume + i, n, emb.data())) { + return fail("qwen35 drafter embedding failed"); } ggml_backend_tensor_set(act_in, emb.data(), (size_t)i * act_in->nb[1], (size_t)hidden * n * sizeof(float)); } } + // Blocks 0..14 over the new tokens only, layer by layer. Each + // DeltaNet layer's recurrent state is copied once it reaches the + // checkpoint. + const auto snapshot_layer = [&](int il) { + int dn = 0; + for (int l = 0; l < il; ++l) { + if (((l + 1) % w.full_attention_interval) != 0) ++dn; + } + if (dn < (int) cache.ssm_state.size() && cache.ssm_state[(size_t) dn] && + cache.ssm_state_snap[(size_t) dn]) { + ggml_backend_tensor_copy(cache.ssm_state[(size_t) dn], + cache.ssm_state_snap[(size_t) dn]); + ggml_backend_tensor_copy(cache.conv_state[(size_t) dn], + cache.conv_state_snap[(size_t) dn]); + } + }; ggml_gallocr_t alloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(w.backend)); const int ubatch = 1024; std::vector mask_bits; for (int il = 0; il < kQwen35HeadBlock; ++il) { const bool is_attn = (((il + 1) % w.full_attention_interval) == 0); - for (int start = 0; start < S; start += ubatch) { - const int n = std::min(ubatch, S - start); + if (!is_attn && checkpoint == resume) snapshot_layer(il); + for (int start = resume; start < S;) { + const int stop = start < checkpoint ? checkpoint : S; + const int n = std::min(ubatch, stop - start); const int kv_len = start + n; ggml_init_params ip{}; ip.mem_size = 512 * 1024 * 1024; ip.no_alloc = true; ggml_context * ctx = ggml_init(ip); if (!ctx) { - ggml_gallocr_free(alloc); cleanup(); - set_last_error("qwen35 drafter layer graph ctx init failed"); - return {}; + ggml_gallocr_free(alloc); + return fail("qwen35 drafter layer graph ctx init failed"); } ggml_cgraph * gf = ggml_new_graph_custom(ctx, 16384, false); ggml_tensor * inp = ggml_view_2d(ctx, act_in, hidden, n, act_in->nb[1], - (size_t)start * act_in->nb[1]); + (size_t)(start - resume) * act_in->nb[1]); ggml_tensor * pos = nullptr; ggml_tensor * mask = nullptr; if (is_attn) { @@ -638,17 +853,15 @@ std::vector qwen35_strict_score_and_compress( ggml_tensor * out = build_qwen35_layer(ctx, gf, w, cache, il, inp, pos, mask, start, n, false, 0); ggml_tensor * dst = ggml_view_2d(ctx, act_out, hidden, n, act_out->nb[1], - (size_t)start * act_out->nb[1]); + (size_t)(start - resume) * act_out->nb[1]); if (ggml_nelements(out) != ggml_nelements(dst)) { - ggml_free(ctx); ggml_gallocr_free(alloc); cleanup(); - set_last_error("qwen35 layer output shape mismatch"); - return {}; + ggml_free(ctx); ggml_gallocr_free(alloc); + return fail("qwen35 layer output shape mismatch"); } ggml_build_forward_expand(gf, ggml_cpy(ctx, out, dst)); if (!ggml_gallocr_alloc_graph(alloc, gf)) { - ggml_free(ctx); ggml_gallocr_free(alloc); cleanup(); - set_last_error("qwen35 drafter graph allocation failed"); - return {}; + ggml_free(ctx); ggml_gallocr_free(alloc); + return fail("qwen35 drafter graph allocation failed"); } if (is_attn) { std::vector p4((size_t)4 * n, 0); @@ -666,9 +879,12 @@ std::vector qwen35_strict_score_and_compress( const auto status = ggml_backend_graph_compute(w.backend, gf); ggml_free(ctx); if (status != GGML_STATUS_SUCCESS) { - ggml_gallocr_free(alloc); cleanup(); - set_last_error("qwen35 drafter graph compute failed"); - return {}; + ggml_gallocr_free(alloc); + return fail("qwen35 drafter graph compute failed"); + } + start += n; + if (!is_attn && start == checkpoint && checkpoint > resume) { + snapshot_layer(il); } } std::swap(act_in, act_out); @@ -676,36 +892,133 @@ std::vector qwen35_strict_score_and_compress( ggml_gallocr_free(alloc); auto t1 = std::chrono::steady_clock::now(); - // Block-15 NoPE Q/K scoring: softmax over keys before the query window, - // then mean over heads and query tokens. The query never scores itself. - // Keys are projected in chunks so no intermediate tensor puts the - // sequence length into a HIP grid y/z dimension (65,535 limit); the - // logits land in one [S, n_lookahead, H] buffer for a single softmax. + // The query window's block-14 rows: this call's when it computed them, + // else the session's from the turn that did. + std::vector query_rows((size_t)hidden * n_lookahead); + if (query_start >= resume) { + ggml_backend_tensor_get(act_in, query_rows.data(), + (size_t)(query_start - resume) * act_in->nb[1], + query_rows.size() * sizeof(float)); + } else { + query_rows = session->query_rows; + } + + ggml_tensor * wk_src = st.head_loaded ? st.head_wk : L.wk; + session->keys_trained = st.head_loaded; + session->probe_raw.resize(use_probe ? (size_t) resume : 0); + session->subunit_raw.resize( + use_probe && st.probe_sub_fc2_w ? (size_t) resume : 0); + + // Keys (and probe logits) of the new tokens, into the session. Keys are + // projected in chunks so no intermediate tensor puts the sequence length + // into a HIP grid y/z dimension (65,535 limit). const int key_chunk = 8192; + if (n_new > 0) { + ggml_init_params nip{}; + nip.mem_size = (size_t)4 * ggml_tensor_overhead() + 4096; + nip.no_alloc = true; + ggml_context * nctx = ggml_init(nip); + ggml_tensor * probe_new = use_probe ? ggml_new_tensor_1d(nctx, GGML_TYPE_F32, n_new) : nullptr; + ggml_tensor * subunit_new = use_probe && st.probe_sub_fc2_w + ? ggml_new_tensor_1d(nctx, GGML_TYPE_F32, n_new) : nullptr; + ggml_backend_buffer_t nbuf = use_probe + ? ggml_backend_alloc_ctx_tensors(nctx, w.backend) : nullptr; + if (use_probe && !nbuf) { + ggml_free(nctx); + return fail("qwen35 probe buffer allocation failed"); + } + const int n_key_chunks = (n_new + key_chunk - 1) / key_chunk; + ggml_init_params kip{}; + kip.mem_size = ggml_tensor_overhead() * (size_t)(64 + 24 * n_key_chunks) + + ggml_graph_overhead_custom(4096, false) + 64 * 1024; + kip.no_alloc = true; + ggml_context * kctx = ggml_init(kip); + ggml_cgraph * kgf = ggml_new_graph_custom(kctx, 4096, false); + for (int b = 0; b < n_new; b += key_chunk) { + const int n = std::min(key_chunk, n_new - b); + ggml_tensor * x_c = ggml_view_2d(kctx, act_in, hidden, n, act_in->nb[1], + (size_t)b * act_in->nb[1]); + ggml_tensor * x_norm = ggml_mul(kctx, ggml_rms_norm(kctx, x_c, w.rms_eps), L.attn_norm); + ggml_tensor * K = ggml_reshape_3d(kctx, ggml_mul_mat(kctx, wk_src, x_norm), D, Hk, n); + K = ggml_mul(kctx, ggml_rms_norm(kctx, K, w.rms_eps), L.k_norm); + ggml_tensor * k_dst = ggml_view_3d(kctx, session->keys, D, Hk, n, + session->keys->nb[1], session->keys->nb[2], + (size_t)(resume + b) * session->keys->nb[2]); + ggml_build_forward_expand(kgf, ggml_cpy(kctx, K, k_dst)); + if (use_probe) { + // Segment probe on the same tap: LayerNorm -> fc1 -> GELU trunk, + // then one fc2 row per head (unit always; sub-unit when shipped). + ggml_tensor * p = ggml_norm(kctx, x_c, st.probe_norm_eps); + p = ggml_add(kctx, ggml_mul(kctx, p, st.probe_norm_w), st.probe_norm_b); + p = ggml_gelu(kctx, ggml_add(kctx, ggml_mul_mat(kctx, st.probe_fc1_w, p), + st.probe_fc1_b)); // [width, n] + ggml_tensor * unit = ggml_add(kctx, ggml_mul_mat(kctx, st.probe_fc2_w, p), + st.probe_fc2_b); // [1, n] + ggml_tensor * p_dst = ggml_view_1d(kctx, probe_new, n, + (size_t)b * ggml_element_size(probe_new)); + ggml_build_forward_expand(kgf, ggml_cpy(kctx, ggml_reshape_1d(kctx, unit, n), p_dst)); + if (subunit_new) { + ggml_tensor * sub = ggml_add(kctx, + ggml_mul_mat(kctx, st.probe_sub_fc2_w, p), st.probe_sub_fc2_b); + ggml_tensor * s_dst = ggml_view_1d(kctx, subunit_new, n, + (size_t)b * ggml_element_size(subunit_new)); + ggml_build_forward_expand(kgf, + ggml_cpy(kctx, ggml_reshape_1d(kctx, sub, n), s_dst)); + } + } + } + ggml_gallocr_t kalloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(w.backend)); + const bool key_ok = ggml_gallocr_alloc_graph(kalloc, kgf) && + ggml_backend_graph_compute(w.backend, kgf) == GGML_STATUS_SUCCESS; + ggml_gallocr_free(kalloc); + ggml_free(kctx); + if (key_ok && use_probe) { + session->probe_raw.resize((size_t) S); + ggml_backend_tensor_get(probe_new, session->probe_raw.data() + resume, 0, + (size_t) n_new * sizeof(float)); + if (subunit_new) { + session->subunit_raw.resize((size_t) S); + ggml_backend_tensor_get(subunit_new, session->subunit_raw.data() + resume, 0, + (size_t) n_new * sizeof(float)); + } + } + if (nbuf) ggml_backend_buffer_free(nbuf); + ggml_free(nctx); + if (!key_ok) return fail("qwen35 key graph compute failed"); + } + cleanup(); + // The session now covers this prompt, resumable at the checkpoint and, + // while the query stays put, reusing its rows. + session->ids = ids; + session->checkpoint = checkpoint; + session->query_begin = query_start; + session->query_end = query_end; + session->query_rows = query_rows; + session->last_used = ++st.session_clock; + + // Block-15 NoPE Q/K scoring: softmax over the keys outside the query + // window (before it, and after it when those tokens are candidates), + // then mean over heads and query tokens. The query never scores itself. + // The logits land in one [S, n_lookahead, H] buffer for a single softmax. const int n_key_chunks = (S + key_chunk - 1) / key_chunk; ggml_init_params lip{}; lip.mem_size = (size_t)8 * ggml_tensor_overhead() + 4096; lip.no_alloc = true; ggml_context * lctx = ggml_init(lip); if (!lctx) { - cleanup(); set_last_error("qwen35 score buffer ctx allocation failed"); return {}; } ggml_tensor * logits = ggml_new_tensor_3d(lctx, GGML_TYPE_F32, S, n_lookahead, H); ggml_tensor * mask = ggml_new_tensor_2d(lctx, GGML_TYPE_F32, S, n_lookahead); - const bool use_probe = st.probe_loaded && - experiment.segmentation != luce::pflash::PFlashSegmentation::Fixed; - ggml_tensor * probe_logits = use_probe - ? ggml_new_tensor_1d(lctx, GGML_TYPE_F32, S) : nullptr; - ggml_tensor * subunit_logits = use_probe && st.probe_sub_fc2_w - ? ggml_new_tensor_1d(lctx, GGML_TYPE_F32, S) : nullptr; + ggml_tensor * x_q = ggml_new_tensor_2d(lctx, GGML_TYPE_F32, hidden, n_lookahead); ggml_backend_buffer_t lbuf = ggml_backend_alloc_ctx_tensors(lctx, w.backend); if (!lbuf) { - ggml_free(lctx); cleanup(); + ggml_free(lctx); set_last_error("qwen35 score buffer allocation failed"); return {}; } + ggml_backend_tensor_set(x_q, query_rows.data(), 0, query_rows.size() * sizeof(float)); { // Keys are the context before the query window and, when the tokens // after it are candidates too, the context after it. NoPE scoring has @@ -726,14 +1039,11 @@ std::vector qwen35_strict_score_and_compress( sip.no_alloc = true; ggml_context * sctx = ggml_init(sip); if (!sctx) { - ggml_backend_buffer_free(lbuf); ggml_free(lctx); cleanup(); + ggml_backend_buffer_free(lbuf); ggml_free(lctx); set_last_error("qwen35 score graph ctx allocation failed"); return {}; } ggml_cgraph * sgf = ggml_new_graph_custom(sctx, 4096, false); - ggml_tensor * wk_src = st.head_loaded ? st.head_wk : L.wk; - ggml_tensor * x_q = ggml_view_2d(sctx, act_in, hidden, n_lookahead, act_in->nb[1], - (size_t)query_start * act_in->nb[1]); ggml_tensor * q_in = ggml_mul(sctx, ggml_rms_norm(sctx, x_q, w.rms_eps), L.attn_norm); ggml_tensor * Q = nullptr; if (st.head_loaded) { @@ -750,11 +1060,9 @@ std::vector qwen35_strict_score_and_compress( ggml_tensor * Q_perm = ggml_cont(sctx, ggml_permute(sctx, Q, 0, 2, 1, 3)); // [D, n_lookahead, H] for (int b = 0; b < S; b += key_chunk) { const int n = std::min(key_chunk, S - b); - ggml_tensor * x_c = ggml_view_2d(sctx, act_in, hidden, n, act_in->nb[1], - (size_t)b * act_in->nb[1]); - ggml_tensor * x_norm = ggml_mul(sctx, ggml_rms_norm(sctx, x_c, w.rms_eps), L.attn_norm); - ggml_tensor * K = ggml_reshape_3d(sctx, ggml_mul_mat(sctx, wk_src, x_norm), D, Hk, n); - K = ggml_mul(sctx, ggml_rms_norm(sctx, K, w.rms_eps), L.k_norm); + ggml_tensor * K = ggml_view_3d(sctx, session->keys, D, Hk, n, + session->keys->nb[1], session->keys->nb[2], + (size_t)b * session->keys->nb[2]); K = ggml_cont(sctx, ggml_permute(sctx, K, 0, 2, 1, 3)); // [D, n, Hk] ggml_tensor * K_score = K; if (H != Hk) { @@ -768,27 +1076,6 @@ std::vector qwen35_strict_score_and_compress( logits->nb[1], logits->nb[2], (size_t)b * logits->nb[0]); ggml_build_forward_expand(sgf, ggml_cpy(sctx, part, dst)); - if (use_probe) { - // Segment probe on the same tap: LayerNorm -> fc1 -> GELU trunk, - // then one fc2 row per head (unit always; sub-unit when shipped). - ggml_tensor * p = ggml_norm(sctx, x_c, st.probe_norm_eps); - p = ggml_add(sctx, ggml_mul(sctx, p, st.probe_norm_w), st.probe_norm_b); - p = ggml_gelu(sctx, ggml_add(sctx, ggml_mul_mat(sctx, st.probe_fc1_w, p), - st.probe_fc1_b)); // [width, n] - ggml_tensor * unit = ggml_add(sctx, ggml_mul_mat(sctx, st.probe_fc2_w, p), - st.probe_fc2_b); // [1, n] - ggml_tensor * p_dst = ggml_view_1d(sctx, probe_logits, n, - (size_t)b * ggml_element_size(probe_logits)); - ggml_build_forward_expand(sgf, ggml_cpy(sctx, ggml_reshape_1d(sctx, unit, n), p_dst)); - if (subunit_logits) { - ggml_tensor * sub = ggml_add(sctx, - ggml_mul_mat(sctx, st.probe_sub_fc2_w, p), st.probe_sub_fc2_b); - ggml_tensor * s_dst = ggml_view_1d(sctx, subunit_logits, n, - (size_t)b * ggml_element_size(subunit_logits)); - ggml_build_forward_expand(sgf, - ggml_cpy(sctx, ggml_reshape_1d(sctx, sub, n), s_dst)); - } - } } ggml_tensor * probs = ggml_soft_max_ext(sctx, logits, mask, 1.0f / std::sqrt((float)D), 0.0f); @@ -797,35 +1084,25 @@ std::vector qwen35_strict_score_and_compress( ggml_gallocr_t salloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(w.backend)); if (!ggml_gallocr_alloc_graph(salloc, sgf)) { ggml_gallocr_free(salloc); ggml_free(sctx); - ggml_backend_buffer_free(lbuf); ggml_free(lctx); cleanup(); + ggml_backend_buffer_free(lbuf); ggml_free(lctx); set_last_error("qwen35 score graph allocation failed"); return {}; } const auto score_status = ggml_backend_graph_compute(w.backend, sgf); if (score_status != GGML_STATUS_SUCCESS) { ggml_gallocr_free(salloc); ggml_free(sctx); - ggml_backend_buffer_free(lbuf); ggml_free(lctx); cleanup(); + ggml_backend_buffer_free(lbuf); ggml_free(lctx); set_last_error("qwen35 score graph compute failed"); return {}; } std::vector probs_h((size_t)S * n_lookahead * H); ggml_backend_tensor_get(probs, probs_h.data(), 0, probs_h.size() * sizeof(float)); - std::vector probe_raw; - std::vector subunit_raw; - if (use_probe) { - probe_raw.resize((size_t) S); - ggml_backend_tensor_get(probe_logits, probe_raw.data(), 0, probe_raw.size() * sizeof(float)); - if (subunit_logits) { - subunit_raw.resize((size_t) S); - ggml_backend_tensor_get(subunit_logits, subunit_raw.data(), 0, - subunit_raw.size() * sizeof(float)); - } - } + const std::vector & probe_raw = session->probe_raw; + const std::vector & subunit_raw = session->subunit_raw; ggml_gallocr_free(salloc); ggml_free(sctx); ggml_backend_buffer_free(lbuf); ggml_free(lctx); - cleanup(); const size_t nonfinite = count_nonfinite_scores(probs_h.data(), probs_h.size()); if (nonfinite != 0) { const std::string message = @@ -833,6 +1110,7 @@ std::vector qwen35_strict_score_and_compress( "/" + std::to_string(probs_h.size()); std::fprintf(stderr, "[pflash] ERROR: %s\n", message.c_str()); std::fflush(stderr); + session->ids.clear(); set_last_error(message); return {}; } @@ -840,9 +1118,10 @@ std::vector qwen35_strict_score_and_compress( scoring_head_mean_token_mass(probs_h.data(), S, n_lookahead, H, token_mass); auto t2 = std::chrono::steady_clock::now(); std::fprintf(stderr, - "[qwen35-scorer] forward %.2fs (blocks 0-%d, S=%d) score %.2fs " - "total %.2fs head=%s\n", + "[qwen35-scorer] forward %.2fs (blocks 0-%d, S=%d, resumed at %d, " + "%d new) score %.2fs total %.2fs head=%s\n", std::chrono::duration(t1 - t0).count(), kQwen35HeadBlock - 1, S, + resume, n_new, std::chrono::duration(t2 - t1).count(), std::chrono::duration(t2 - t0).count(), st.head_loaded ? "trained" : "native-block15"); diff --git a/server/src/pflash/qwen35_drafter.h b/server/src/pflash/qwen35_drafter.h index c39e730f6..8ef7cf16d 100644 --- a/server/src/pflash/qwen35_drafter.h +++ b/server/src/pflash/qwen35_drafter.h @@ -15,6 +15,8 @@ #include "ggml.h" #include "ggml-backend.h" +#include +#include #include #include @@ -27,6 +29,31 @@ namespace luce::common { // projections. static constexpr int kQwen35HeadBlock = 15; +// What the scorer computed for one conversation's prompt, kept so the next +// turn only runs the new tokens. Blocks 0..14 read left to right, so the +// cache state and block-15 keys of a shared prefix never change; the keys +// carry no position (NoPE), and the probe's raw logits are per token. +struct Qwen35ScoringSession { + std::vector ids; // prompt tokens the state covers + int checkpoint = 0; // recurrent-state snapshot position + int capacity = 0; // tokens the cache and keys can hold + TargetCache cache; // blocks 0..14 only + ggml_context * key_ctx = nullptr; + ggml_backend_buffer_t key_buf = nullptr; + ggml_tensor * keys = nullptr; // [head_dim, n_head_kv, capacity] f32 + bool keys_trained = false; + std::vector probe_raw; // per token, unit logit + std::vector subunit_raw; // per token, when the probe has one + // Block-14 output of the last query window, reused while the query stays + // put (an agent step appends tool output after the same user turn). + int query_begin = -1; + int query_end = -1; + std::vector query_rows; // [hidden, query_end - query_begin] + uint64_t last_used = 0; +}; + +void free_qwen35_scoring_session(Qwen35ScoringSession & session); + struct Qwen35DrafterState { TargetWeights weights; std::string gguf_sha256; @@ -56,6 +83,10 @@ struct Qwen35DrafterState { int probe_max_segment = 2048; int probe_width = 0; bool probe_loaded = false; + // Strict scorer sessions, least recently used evicted + // (PFLASH_DRAFTER_SESSIONS, default 2; 0 scores every prompt from scratch). + std::vector> sessions; + uint64_t session_clock = 0; }; // Defined in qwen35_loader.cpp. diff --git a/server/src/pflash/qwen35_loader.cpp b/server/src/pflash/qwen35_loader.cpp index d737cea4f..4a2ca718f 100644 --- a/server/src/pflash/qwen35_loader.cpp +++ b/server/src/pflash/qwen35_loader.cpp @@ -362,6 +362,10 @@ bool load_qwen35_drafter(const std::string & gguf_path, void free_qwen35_drafter_state(DrafterContext & ctx) { auto * st = static_cast(ctx.state); + for (auto & session : st->sessions) { + if (session) free_qwen35_scoring_session(*session); + } + st->sessions.clear(); free_qwen35_head(*st); free_qwen35_segment_probe(*st); free_target_weights(st->weights); From 55d7f9c0602680f4750fed8df773dfc23e6e3bc5 Mon Sep 17 00:00:00 2001 From: Graffioh Date: Tue, 22 Sep 2026 20:06:09 +0000 Subject: [PATCH 17/26] feat(pflash): multi-turn skeleton and history-query scoring A rebuild (and every turn's fresh selection) scored the conversation against the latest question alone and could drop the dialogue itself: earlier questions, the model's own short answers, role headers. The chat-turn scan now reports every turn. Under strict selection every other turn keeps its role header, and user turns and assistant answers up to PFLASH_CHAT_SKELETON_TOKENS (default 256; 0 keeps headers) stay whole; like developer text and tool definitions, the skeleton is scored as context when it alone would not fit. The tails of the last PFLASH_CHAT_HISTORY_QUERIES (default 3) earlier user turns score the context alongside the current query (CompressRequest::history_query_spans): the head scores each window against the query's key set and mixes the masses at weights 1, 1/2, 1/4, 1/8. The drafter session keeps the block-14 rows of up to 8 recent windows, so a new turn's history windows come from the turns that computed them. Live on the R9700 (20.6K-token conversation, sessions on): turn 2 recalls 64 tokens instead of 1472 and completes in 1.0 s instead of 3.9 s; turn 3 in 3.0 s; all three answers correct. Mixing in old questions trades the new question's share of the budget, so both knobs are ablations for the multi-turn evaluation. Co-Authored-By: Claude Opus 5.5 (1M context) --- server/README.md | 9 + server/src/common/model_backend.h | 3 + server/src/deepseek4/deepseek4_backend.cpp | 2 +- server/src/pflash/pflash_drafter.cpp | 11 +- server/src/pflash/pflash_drafter.h | 5 +- server/src/pflash/pflash_selection.h | 4 + server/src/pflash/qwen35_drafter.cpp | 345 +++++++++++------- server/src/pflash/qwen35_drafter.h | 15 +- server/src/qwen3/qwen3_backend.cpp | 2 +- server/src/qwen35/qwen35_backend.cpp | 2 +- .../src/qwen35/qwen35_layer_split_adapter.cpp | 2 +- server/src/server/http_server.cpp | 79 ++++ server/src/server/http_server.h | 24 ++ server/test/test_server_unit.cpp | 83 ++++- 14 files changed, 436 insertions(+), 150 deletions(-) diff --git a/server/README.md b/server/README.md index b31efe96b..e91155c15 100644 --- a/server/README.md +++ b/server/README.md @@ -417,6 +417,15 @@ view grows past twice the fresh prompt, or past the context, the fresh prompt starts a new view. `PFLASH_CHAT_VIEW=0` serves the fresh compression every turn. +Every other turn of a multi-turn chat keeps its role header, and user +turns and assistant answers up to `PFLASH_CHAT_SKELETON_TOKENS` (default 256 +drafter tokens; 0 keeps headers only) stay whole: the conversation's +skeleton, as opposed to the material it quotes. Like instructions, the +skeleton is scored as context when it alone would not fit. The last +`PFLASH_CHAT_HISTORY_QUERIES` (default 3) earlier user questions score the +context alongside the current one, their masses mixed in at weights 1/2, +1/4, 1/8, so what the conversation keeps coming back to stays selected. + The drafter keeps a scoring session per conversation (`PFLASH_DRAFTER_SESSIONS`, default 2, least recently used evicted; 0 scores every prompt from scratch): the cache of blocks 0-14, the block-15 keys and diff --git a/server/src/common/model_backend.h b/server/src/common/model_backend.h index 1174f81dd..cc6ef062d 100644 --- a/server/src/common/model_backend.h +++ b/server/src/common/model_backend.h @@ -275,6 +275,9 @@ struct ModelBackend { // window are scored candidates rather than a kept suffix. The caller // pins what of that suffix must stay (the generation prompt). bool query_suffix_candidates = false; + // Earlier user questions (their scorer windows), most recent first: + // they score the context alongside the query at halving weights. + std::vector history_query_spans; std::string drafter_path; // GGUF path (for lazy-load) int drafter_gpu = 0; // backend-local GPU for PFlash drafter bool skip_park = false; // true on >=32GB GPUs diff --git a/server/src/deepseek4/deepseek4_backend.cpp b/server/src/deepseek4/deepseek4_backend.cpp index 981d47d0c..0d3d82c68 100644 --- a/server/src/deepseek4/deepseek4_backend.cpp +++ b/server/src/deepseek4/deepseek4_backend.cpp @@ -3183,7 +3183,7 @@ std::vector DeepSeek4Backend::compress_batch( pflash_drafter_ctx_, request.input_ids, request.keep_ratio, /*chunk_size=*/32, request.score_query_tokens, /*pool_kernel=*/13, score_query_end, request.required_instruction_spans, - request.query_suffix_candidates); + request.query_suffix_candidates, request.history_query_spans); result.ok = !result.compressed_ids.empty(); if (result.ok) result.kept_spans = pflash_last_kept_spans(); } diff --git a/server/src/pflash/pflash_drafter.cpp b/server/src/pflash/pflash_drafter.cpp index ab372215a..fad60ccd9 100644 --- a/server/src/pflash/pflash_drafter.cpp +++ b/server/src/pflash/pflash_drafter.cpp @@ -106,7 +106,8 @@ std::vector drafter_score_and_compress( int pool_kernel, int score_query_end, const std::vector & required_instruction_spans, - bool query_suffix_candidates) { + bool query_suffix_candidates, + const std::vector & history_queries) { pflash_clear_kept_spans(); if (!ctx.loaded) { set_last_error("drafter not loaded"); @@ -126,6 +127,14 @@ std::vector drafter_score_and_compress( chunk_size = experiment.chunk_size; experiment.query_suffix_candidates = query_suffix_candidates && experiment.selection_active; + if (experiment.selection_active) { + for (const auto & window : history_queries) { + if (window.begin >= 0 && window.end > window.begin && + window.end <= (int) ids.size()) { + experiment.history_queries.push_back(window); + } + } + } if (!experiment.selection_active && !required_instruction_spans.empty()) { set_last_error( "PFlash instruction spans require strict budget selection"); diff --git a/server/src/pflash/pflash_drafter.h b/server/src/pflash/pflash_drafter.h index 364d2d477..bd89d2746 100644 --- a/server/src/pflash/pflash_drafter.h +++ b/server/src/pflash/pflash_drafter.h @@ -73,6 +73,8 @@ void free_drafter_weights(DrafterContext & ctx); // required (negative values are rejected) // query_suffix_candidates strict selection only: tokens after the query // window are scored candidates, not a kept suffix +// history_queries strict selection only: earlier questions' windows, most +// recent first, mixed into the scores at halving weights // // On failure returns empty vector + sets last_error. std::vector drafter_score_and_compress( @@ -85,6 +87,7 @@ std::vector drafter_score_and_compress( int score_query_end = -1, const std::vector & required_instruction_spans = {}, - bool query_suffix_candidates = false); + bool query_suffix_candidates = false, + const std::vector & history_queries = {}); } // namespace luce::common diff --git a/server/src/pflash/pflash_selection.h b/server/src/pflash/pflash_selection.h index dede1a51a..fb1efef8d 100644 --- a/server/src/pflash/pflash_selection.h +++ b/server/src/pflash/pflash_selection.h @@ -117,6 +117,10 @@ struct PFlashSelectionConfig { // chat whose latest user turn is followed by assistant and tool turns). // The caller pins whatever of that suffix must stay. bool query_suffix_candidates = false; + // Per request: earlier user questions' scorer windows, most recent + // 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; }; // Segment probe: cut the context before every token whose boundary score is diff --git a/server/src/pflash/qwen35_drafter.cpp b/server/src/pflash/qwen35_drafter.cpp index 56a8e7ac0..25c1aac4f 100644 --- a/server/src/pflash/qwen35_drafter.cpp +++ b/server/src/pflash/qwen35_drafter.cpp @@ -532,8 +532,7 @@ void free_qwen35_scoring_session(Qwen35ScoringSession & session) { session.checkpoint = 0; session.probe_raw.clear(); session.subunit_raw.clear(); - session.query_begin = session.query_end = -1; - session.query_rows.clear(); + session.query_windows.clear(); } namespace { @@ -605,15 +604,18 @@ Qwen35ScoringSession * acquire_scoring_session( bool need_probe, bool need_subunit, int & resume, + int & shared_prefix, std::unique_ptr & scratch) { TargetWeights & w = st.weights; const int S = (int) ids.size(); const int limit = scoring_session_limit(); const size_t row_floats = (size_t) w.n_embd * (size_t) (query_end - query_start); resume = 0; + shared_prefix = 0; Qwen35ScoringSession * best = nullptr; bool best_restore = false; + int best_shared = 0; for (auto & owned : st.sessions) { Qwen35ScoringSession * session = owned.get(); if (!session || session->capacity < S || session->ids.empty() || @@ -633,9 +635,12 @@ Qwen35ScoringSession * acquire_scoring_session( restore = true; } if (r > query_start) { - const bool rows = session->query_begin == query_start && - session->query_end == query_end && query_end <= shared && - session->query_rows.size() == row_floats; + bool rows = false; + for (const auto & window : session->query_windows) { + rows = rows || (window.begin == query_start && + window.end == query_end && query_end <= shared && + window.rows.size() == row_floats); + } if (!rows) { r = session->checkpoint > 0 && session->checkpoint <= query_start && session->checkpoint <= shared @@ -651,12 +656,14 @@ Qwen35ScoringSession * acquire_scoring_session( resume = r; best = session; best_restore = restore; + best_shared = shared; } } if (best) { if (best_restore && !restore_ssm_state(best->cache, w.backend)) { resume = 0; } else { + shared_prefix = best_shared; return best; } } @@ -690,8 +697,7 @@ Qwen35ScoringSession * acquire_scoring_session( target->checkpoint = 0; target->probe_raw.clear(); target->subunit_raw.clear(); - target->query_begin = target->query_end = -1; - target->query_rows.clear(); + target->query_windows.clear(); return target; } @@ -740,10 +746,12 @@ std::vector qwen35_strict_score_and_compress( auto t0 = std::chrono::steady_clock::now(); int resume = 0; + int shared_prefix = 0; std::unique_ptr scratch; Qwen35ScoringSession * session = acquire_scoring_session( st, ids, query_start, query_end, use_probe, - use_probe && st.probe_sub_fc2_w != nullptr, resume, scratch); + use_probe && st.probe_sub_fc2_w != nullptr, resume, shared_prefix, + scratch); if (!session) return {}; // A session is released (freed or kept) on every exit below. struct SessionExit { @@ -892,15 +900,55 @@ std::vector qwen35_strict_score_and_compress( ggml_gallocr_free(alloc); auto t1 = std::chrono::steady_clock::now(); - // The query window's block-14 rows: this call's when it computed them, - // else the session's from the turn that did. - std::vector query_rows((size_t)hidden * n_lookahead); - if (query_start >= resume) { - ggml_backend_tensor_get(act_in, query_rows.data(), - (size_t)(query_start - resume) * act_in->nb[1], - query_rows.size() * sizeof(float)); - } else { - query_rows = session->query_rows; + // Block-14 rows of each query window -- the query, then earlier + // questions -- from this call when it computed them, else from the + // session while they sit in the shared prefix. The query's are always + // available (the session was chosen for them); a history window whose + // rows are gone is skipped. + struct ScoredWindow { + int begin = 0; + int end = 0; + double weight = 1.0; + std::vector rows; + }; + std::vector windows; + const auto rows_for = [&](int begin, int end, std::vector & rows) { + rows.assign((size_t)hidden * (size_t)(end - begin), 0.0f); + if (begin >= resume) { + ggml_backend_tensor_get(act_in, rows.data(), + (size_t)(begin - resume) * act_in->nb[1], + 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; + }; + { + ScoredWindow query; + query.begin = query_start; + query.end = query_end; + if (!rows_for(query_start, query_end, query.rows)) { + return fail("qwen35 scorer query rows unavailable"); + } + windows.push_back(std::move(query)); + double weight = 1.0; + for (const auto & span : experiment.history_queries) { + weight *= 0.5; + if (span.end > query_start || span.end - span.begin < 1) continue; + ScoredWindow history; + history.begin = span.begin; + history.end = span.end; + history.weight = weight; + if (rows_for(span.begin, span.end, history.rows)) { + windows.push_back(std::move(history)); + } + } } ggml_tensor * wk_src = st.head_loaded ? st.head_wk : L.wk; @@ -991,138 +1039,165 @@ std::vector qwen35_strict_score_and_compress( // while the query stays put, reusing its rows. session->ids = ids; session->checkpoint = checkpoint; - session->query_begin = query_start; - session->query_end = query_end; - session->query_rows = query_rows; + for (const auto & window : windows) { + 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()); + } session->last_used = ++st.session_clock; - // Block-15 NoPE Q/K scoring: softmax over the keys outside the query - // window (before it, and after it when those tokens are candidates), - // then mean over heads and query tokens. The query never scores itself. - // The logits land in one [S, n_lookahead, H] buffer for a single softmax. + // Block-15 NoPE Q/K scoring, once per query window: softmax over the + // keys outside the query window (before it, and after it when those + // tokens are candidates), then mean over heads and window tokens. No + // window scores itself. The logits land in one [S, rows, H] buffer for + // a single softmax. History windows share the query's key set and mix + // into its mass at their weights. const int n_key_chunks = (S + key_chunk - 1) / key_chunk; - ggml_init_params lip{}; - lip.mem_size = (size_t)8 * ggml_tensor_overhead() + 4096; - lip.no_alloc = true; - ggml_context * lctx = ggml_init(lip); - if (!lctx) { - set_last_error("qwen35 score buffer ctx allocation failed"); - return {}; - } - ggml_tensor * logits = ggml_new_tensor_3d(lctx, GGML_TYPE_F32, S, n_lookahead, H); - ggml_tensor * mask = ggml_new_tensor_2d(lctx, GGML_TYPE_F32, S, n_lookahead); - ggml_tensor * x_q = ggml_new_tensor_2d(lctx, GGML_TYPE_F32, hidden, n_lookahead); - ggml_backend_buffer_t lbuf = ggml_backend_alloc_ctx_tensors(lctx, w.backend); - if (!lbuf) { - ggml_free(lctx); - set_last_error("qwen35 score buffer allocation failed"); - return {}; - } - ggml_backend_tensor_set(x_q, query_rows.data(), 0, query_rows.size() * sizeof(float)); - { - // Keys are the context before the query window and, when the tokens - // after it are candidates too, the context after it. NoPE scoring has - // no position term, so a later key scores like an earlier one. - std::vector m((size_t)n_lookahead * S, -INFINITY); - for (int t = 0; t < n_lookahead; ++t) { - std::fill_n(m.begin() + (size_t)t * S, (size_t)query_start, 0.0f); - if (experiment.query_suffix_candidates) { - std::fill_n(m.begin() + (size_t)t * S + query_end, - (size_t)(S - query_end), 0.0f); + const auto score_window = [&](const ScoredWindow & window, + std::vector & mass) -> bool { + const int nq = window.end - window.begin; + ggml_init_params lip{}; + lip.mem_size = (size_t)8 * ggml_tensor_overhead() + 4096; + lip.no_alloc = true; + ggml_context * lctx = ggml_init(lip); + if (!lctx) { + set_last_error("qwen35 score buffer ctx allocation failed"); + return false; + } + ggml_tensor * logits = ggml_new_tensor_3d(lctx, GGML_TYPE_F32, S, nq, H); + ggml_tensor * mask = ggml_new_tensor_2d(lctx, GGML_TYPE_F32, S, nq); + ggml_tensor * x_q = ggml_new_tensor_2d(lctx, GGML_TYPE_F32, hidden, nq); + ggml_backend_buffer_t lbuf = ggml_backend_alloc_ctx_tensors(lctx, w.backend); + if (!lbuf) { + ggml_free(lctx); + set_last_error("qwen35 score buffer allocation failed"); + return false; + } + ggml_backend_tensor_set(x_q, window.rows.data(), 0, + window.rows.size() * sizeof(float)); + { + // Keys are the context before the query window and, when the + // tokens after it are candidates too, the context after it. NoPE + // scoring has no position term, so a later key scores like an + // earlier one. + std::vector m((size_t)nq * S, -INFINITY); + for (int t = 0; t < nq; ++t) { + float * row = m.data() + (size_t)t * S; + std::fill_n(row, (size_t)query_start, 0.0f); + if (experiment.query_suffix_candidates) { + std::fill_n(row + query_end, (size_t)(S - query_end), 0.0f); + } + if (window.end <= query_start) { + std::fill_n(row + window.begin, (size_t)nq, -INFINITY); + } } + ggml_backend_tensor_set(mask, m.data(), 0, m.size() * sizeof(float)); } - ggml_backend_tensor_set(mask, m.data(), 0, m.size() * sizeof(float)); - } - ggml_init_params sip{}; - sip.mem_size = ggml_tensor_overhead() * (size_t)(64 + 24 * n_key_chunks) + - ggml_graph_overhead_custom(4096, false) + 64 * 1024; - sip.no_alloc = true; - ggml_context * sctx = ggml_init(sip); - if (!sctx) { - ggml_backend_buffer_free(lbuf); ggml_free(lctx); - set_last_error("qwen35 score graph ctx allocation failed"); - return {}; - } - ggml_cgraph * sgf = ggml_new_graph_custom(sctx, 4096, false); - ggml_tensor * q_in = ggml_mul(sctx, ggml_rms_norm(sctx, x_q, w.rms_eps), L.attn_norm); - ggml_tensor * Q = nullptr; - if (st.head_loaded) { - Q = ggml_reshape_3d(sctx, ggml_mul_mat(sctx, st.head_wq, q_in), D, H, n_lookahead); - } else { - // Native block 15 packs query and gate rows per head; keep the query half. - ggml_tensor * QG = ggml_reshape_3d(sctx, ggml_mul_mat(sctx, L.wq, q_in), - D * 2, H, n_lookahead); - Q = ggml_view_3d(sctx, QG, D, H, n_lookahead, - ggml_element_size(QG) * D * 2, - ggml_element_size(QG) * D * 2 * H, 0); - } - Q = ggml_mul(sctx, ggml_rms_norm(sctx, Q, w.rms_eps), L.q_norm); - ggml_tensor * Q_perm = ggml_cont(sctx, ggml_permute(sctx, Q, 0, 2, 1, 3)); // [D, n_lookahead, H] - for (int b = 0; b < S; b += key_chunk) { - const int n = std::min(key_chunk, S - b); - ggml_tensor * K = ggml_view_3d(sctx, session->keys, D, Hk, n, - session->keys->nb[1], session->keys->nb[2], - (size_t)b * session->keys->nb[2]); - K = ggml_cont(sctx, ggml_permute(sctx, K, 0, 2, 1, 3)); // [D, n, Hk] - ggml_tensor * K_score = K; - if (H != Hk) { - const int gqa = H / Hk; - ggml_tensor * K_4d = ggml_reshape_4d(sctx, K, D, n, 1, Hk); - ggml_tensor * K_tpl = ggml_new_tensor_4d(sctx, GGML_TYPE_F32, D, n, gqa, Hk); - K_score = ggml_reshape_3d(sctx, ggml_repeat(sctx, K_4d, K_tpl), D, n, H); + ggml_init_params sip{}; + sip.mem_size = ggml_tensor_overhead() * (size_t)(64 + 24 * n_key_chunks) + + ggml_graph_overhead_custom(4096, false) + 64 * 1024; + sip.no_alloc = true; + ggml_context * sctx = ggml_init(sip); + if (!sctx) { + ggml_backend_buffer_free(lbuf); ggml_free(lctx); + set_last_error("qwen35 score graph ctx allocation failed"); + return false; } - ggml_tensor * part = ggml_mul_mat(sctx, K_score, Q_perm); // [n, n_lookahead, H] - ggml_tensor * dst = ggml_view_3d(sctx, logits, n, n_lookahead, H, - logits->nb[1], logits->nb[2], - (size_t)b * logits->nb[0]); - ggml_build_forward_expand(sgf, ggml_cpy(sctx, part, dst)); - } - ggml_tensor * probs = ggml_soft_max_ext(sctx, logits, mask, - 1.0f / std::sqrt((float)D), 0.0f); - ggml_set_output(probs); - ggml_build_forward_expand(sgf, probs); - ggml_gallocr_t salloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(w.backend)); - if (!ggml_gallocr_alloc_graph(salloc, sgf)) { - ggml_gallocr_free(salloc); ggml_free(sctx); - ggml_backend_buffer_free(lbuf); ggml_free(lctx); - set_last_error("qwen35 score graph allocation failed"); - return {}; - } - const auto score_status = ggml_backend_graph_compute(w.backend, sgf); - if (score_status != GGML_STATUS_SUCCESS) { - ggml_gallocr_free(salloc); ggml_free(sctx); - ggml_backend_buffer_free(lbuf); ggml_free(lctx); - set_last_error("qwen35 score graph compute failed"); - return {}; + ggml_cgraph * sgf = ggml_new_graph_custom(sctx, 4096, false); + ggml_tensor * q_in = ggml_mul(sctx, ggml_rms_norm(sctx, x_q, w.rms_eps), L.attn_norm); + ggml_tensor * Q = nullptr; + if (st.head_loaded) { + Q = ggml_reshape_3d(sctx, ggml_mul_mat(sctx, st.head_wq, q_in), D, H, nq); + } else { + // Native block 15 packs query and gate rows per head; keep the query half. + ggml_tensor * QG = ggml_reshape_3d(sctx, ggml_mul_mat(sctx, L.wq, q_in), + D * 2, H, nq); + Q = ggml_view_3d(sctx, QG, D, H, nq, + ggml_element_size(QG) * D * 2, + ggml_element_size(QG) * D * 2 * H, 0); + } + Q = ggml_mul(sctx, ggml_rms_norm(sctx, Q, w.rms_eps), L.q_norm); + ggml_tensor * Q_perm = ggml_cont(sctx, ggml_permute(sctx, Q, 0, 2, 1, 3)); // [D, nq, H] + for (int b = 0; b < S; b += key_chunk) { + const int n = std::min(key_chunk, S - b); + ggml_tensor * K = ggml_view_3d(sctx, session->keys, D, Hk, n, + session->keys->nb[1], session->keys->nb[2], + (size_t)b * session->keys->nb[2]); + K = ggml_cont(sctx, ggml_permute(sctx, K, 0, 2, 1, 3)); // [D, n, Hk] + ggml_tensor * K_score = K; + if (H != Hk) { + const int gqa = H / Hk; + ggml_tensor * K_4d = ggml_reshape_4d(sctx, K, D, n, 1, Hk); + ggml_tensor * K_tpl = ggml_new_tensor_4d(sctx, GGML_TYPE_F32, D, n, gqa, Hk); + K_score = ggml_reshape_3d(sctx, ggml_repeat(sctx, K_4d, K_tpl), D, n, H); + } + ggml_tensor * part = ggml_mul_mat(sctx, K_score, Q_perm); // [n, nq, H] + ggml_tensor * dst = ggml_view_3d(sctx, logits, n, nq, H, + logits->nb[1], logits->nb[2], + (size_t)b * logits->nb[0]); + ggml_build_forward_expand(sgf, ggml_cpy(sctx, part, dst)); + } + ggml_tensor * probs = ggml_soft_max_ext(sctx, logits, mask, + 1.0f / std::sqrt((float)D), 0.0f); + ggml_set_output(probs); + ggml_build_forward_expand(sgf, probs); + ggml_gallocr_t salloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(w.backend)); + const bool ok = ggml_gallocr_alloc_graph(salloc, sgf) && + ggml_backend_graph_compute(w.backend, sgf) == GGML_STATUS_SUCCESS; + std::vector probs_h; + if (ok) { + probs_h.resize((size_t)S * nq * H); + ggml_backend_tensor_get(probs, probs_h.data(), 0, probs_h.size() * sizeof(float)); + } + ggml_gallocr_free(salloc); + ggml_free(sctx); + ggml_backend_buffer_free(lbuf); + ggml_free(lctx); + if (!ok) { + set_last_error("qwen35 score graph compute failed"); + return false; + } + const size_t nonfinite = count_nonfinite_scores(probs_h.data(), probs_h.size()); + if (nonfinite != 0) { + const std::string message = + "non-finite Qwen3.5 scoring-head scores: " + std::to_string(nonfinite) + + "/" + std::to_string(probs_h.size()); + std::fprintf(stderr, "[pflash] ERROR: %s\n", message.c_str()); + std::fflush(stderr); + set_last_error(message); + return false; + } + scoring_head_mean_token_mass(probs_h.data(), S, nq, H, mass); + return true; + }; + std::vector token_mass; + double total_weight = 0.0; + for (const auto & window : windows) { + std::vector mass; + if (!score_window(window, mass)) { + session->ids.clear(); + return {}; + } + if (token_mass.empty()) token_mass.assign(mass.size(), 0.0f); + for (size_t i = 0; i < mass.size(); ++i) { + token_mass[i] += (float) window.weight * mass[i]; + } + total_weight += window.weight; } - std::vector probs_h((size_t)S * n_lookahead * H); - ggml_backend_tensor_get(probs, probs_h.data(), 0, probs_h.size() * sizeof(float)); + for (auto & value : token_mass) value = (float) (value / total_weight); const std::vector & probe_raw = session->probe_raw; const std::vector & subunit_raw = session->subunit_raw; - ggml_gallocr_free(salloc); - ggml_free(sctx); - ggml_backend_buffer_free(lbuf); - ggml_free(lctx); - const size_t nonfinite = count_nonfinite_scores(probs_h.data(), probs_h.size()); - if (nonfinite != 0) { - const std::string message = - "non-finite Qwen3.5 scoring-head scores: " + std::to_string(nonfinite) + - "/" + std::to_string(probs_h.size()); - std::fprintf(stderr, "[pflash] ERROR: %s\n", message.c_str()); - std::fflush(stderr); - session->ids.clear(); - set_last_error(message); - return {}; - } - std::vector token_mass; - scoring_head_mean_token_mass(probs_h.data(), S, n_lookahead, H, token_mass); auto t2 = std::chrono::steady_clock::now(); std::fprintf(stderr, "[qwen35-scorer] forward %.2fs (blocks 0-%d, S=%d, resumed at %d, " - "%d new) score %.2fs total %.2fs head=%s\n", + "%d new) score %.2fs (%zu query windows) total %.2fs head=%s\n", std::chrono::duration(t1 - t0).count(), kQwen35HeadBlock - 1, S, resume, n_new, - std::chrono::duration(t2 - t1).count(), + std::chrono::duration(t2 - t1).count(), windows.size(), std::chrono::duration(t2 - t0).count(), st.head_loaded ? "trained" : "native-block15"); std::fflush(stderr); diff --git a/server/src/pflash/qwen35_drafter.h b/server/src/pflash/qwen35_drafter.h index 8ef7cf16d..8272fa656 100644 --- a/server/src/pflash/qwen35_drafter.h +++ b/server/src/pflash/qwen35_drafter.h @@ -44,11 +44,16 @@ struct Qwen35ScoringSession { bool keys_trained = false; std::vector probe_raw; // per token, unit logit std::vector subunit_raw; // per token, when the probe has one - // Block-14 output of the last query window, reused while the query stays - // put (an agent step appends tool output after the same user turn). - int query_begin = -1; - int query_end = -1; - std::vector query_rows; // [hidden, query_end - query_begin] + // Block-14 output of recent query windows (the query and earlier + // 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; uint64_t last_used = 0; }; diff --git a/server/src/qwen3/qwen3_backend.cpp b/server/src/qwen3/qwen3_backend.cpp index 021ce83e4..3b0a04b94 100644 --- a/server/src/qwen3/qwen3_backend.cpp +++ b/server/src/qwen3/qwen3_backend.cpp @@ -977,7 +977,7 @@ ModelBackend::CompressResult Qwen3Backend::compress(const CompressRequest & req) drafter_ctx_, req.input_ids, req.keep_ratio, /*chunk_size=*/32, req.score_query_tokens, /*pool_kernel=*/13, score_query_end, req.required_instruction_spans, - req.query_suffix_candidates)); + req.query_suffix_candidates, req.history_query_spans)); if (req.residency_action == DraftResidencyAction::ReleaseAfterUse) { free_drafter(); diff --git a/server/src/qwen35/qwen35_backend.cpp b/server/src/qwen35/qwen35_backend.cpp index 4c85cbabe..c0e8b3cf4 100644 --- a/server/src/qwen35/qwen35_backend.cpp +++ b/server/src/qwen35/qwen35_backend.cpp @@ -1203,7 +1203,7 @@ std::vector Qwen35Backend::compress_batch( drafter_ctx_, request.input_ids, request.keep_ratio, /*chunk_size=*/32, request.score_query_tokens, /*pool_kernel=*/13, score_query_end, request.required_instruction_spans, - request.query_suffix_candidates); + request.query_suffix_candidates, request.history_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 4aae63430..4a3005ad5 100644 --- a/server/src/qwen35/qwen35_layer_split_adapter.cpp +++ b/server/src/qwen35/qwen35_layer_split_adapter.cpp @@ -1394,7 +1394,7 @@ Qwen35LayerSplitAdapter::compress(const ModelBackend::CompressRequest & req) { pflash_drafter_, req.input_ids, req.keep_ratio, /*chunk_size=*/32, req.score_query_tokens, /*pool_kernel=*/13, score_query_end, req.required_instruction_spans, - req.query_suffix_candidates); + req.query_suffix_candidates, req.history_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 293e29271..bb02f3be6 100644 --- a/server/src/server/http_server.cpp +++ b/server/src/server/http_server.cpp @@ -603,9 +603,39 @@ PflashChatTurnSpan pflash_chat_query_turn( if (chosen.role_begin > chosen.content_begin) { chosen.role_begin = chosen.content_begin; } + for (size_t index = 0; index < usable; ++index) { + const Turn & turn = turns[index]; + PflashChatTurn out; + out.role_begin = token_at_offset(decoded, turn.role_at); + out.content_begin = token_from(turn.content_at); + out.content_end = token_from(turn.content_end); + out.turn_end = token_from(turn.close_end); + out.role_begin = (std::min)(out.role_begin, out.content_begin); + out.role = turn.role; + chosen.turns.push_back(std::move(out)); + } + chosen.query_turn = (int) query_index; return chosen; } +int pflash_chat_skeleton_tokens() noexcept { + const char * raw = std::getenv("PFLASH_CHAT_SKELETON_TOKENS"); + if (!raw || !*raw) return 256; + char * end = nullptr; + const long value = std::strtol(raw, &end, 10); + if (end == raw || *end != '\0' || value < 0) return 256; + return (int) (std::min)(value, 1L << 20); +} + +int pflash_chat_history_queries() noexcept { + const char * raw = std::getenv("PFLASH_CHAT_HISTORY_QUERIES"); + if (!raw || !*raw) return 3; + char * end = nullptr; + const long value = std::strtol(raw, &end, 10); + if (end == raw || *end != '\0' || value < 0) return 3; + return (int) (std::min)(value, 8L); +} + std::vector pflash_subtract_token_spans( const std::vector & spans, const std::vector & minus) { @@ -3721,6 +3751,9 @@ std::string HttpServer::apply_pflash_compression( // Assistant and tool turns after the query's turn are context the query // scores, not a kept suffix (strict selection only). bool query_suffix_candidates = false; + // 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; // Header ("<|im_start|>user\n") opening the query's turn, when the chat // markers resolved it — pinned mandatory so a compressed prompt keeps // the current turn's role envelope. @@ -4011,6 +4044,50 @@ std::string HttpServer::apply_pflash_compression( (int) drafter_ids.size()}); } } + // Multi-turn skeleton: every other turn keeps its role + // header, and short user turns and assistant answers stay + // whole -- what the conversation said rather than the + // material it quoted. Like instructions, they are scored as + // context when they alone would not fit. + if (chat_turn.valid()) { + const int skeleton_tokens = + http_detail::pflash_chat_skeleton_tokens(); + for (size_t index = 0; index < chat_turn.turns.size(); + ++index) { + if ((int) index == chat_turn.query_turn) continue; + const auto & turn = chat_turn.turns[index]; + if (turn.role == "system") continue; + if (turn.content_begin > turn.role_begin) { + instruction_role_spans.push_back( + {turn.role_begin, turn.content_begin}); + } + const bool conversational = turn.role == "user" || + turn.role == "assistant" || turn.role == "model"; + if (conversational && skeleton_tokens > 0 && + turn.content_end - turn.content_begin <= + skeleton_tokens && + turn.turn_end > turn.role_begin) { + instruction_role_spans.push_back( + {turn.role_begin, turn.turn_end}); + } + } + const size_t history_queries = + (size_t) http_detail::pflash_chat_history_queries(); + for (int index = chat_turn.query_turn - 1; + index >= 0 && + history_query_spans.size() < history_queries; + --index) { + const auto & turn = chat_turn.turns[(size_t) index]; + if (turn.role != "user") continue; + const auto window = http_detail::pflash_tail_query_window( + drafter_ids, experiment.query_tokens, + turn.content_end, turn.content_begin); + if (window.valid()) { + history_query_spans.push_back( + {window.end - window.tokens, window.end}); + } + } + } required_instruction_spans = http_detail::canonicalize_pflash_token_spans( std::move(required_instruction_spans)); @@ -4197,6 +4274,7 @@ std::string HttpServer::apply_pflash_compression( compress_request.required_instruction_spans = std::move(required_instruction_spans); compress_request.query_suffix_candidates = query_suffix_candidates; + compress_request.history_query_spans = history_query_spans; compress_request.keep_ratio = http_detail::resolve_pflash_keep_ratio( pflash_keep_ratio(config_, prompt_tokens), req.session_id, sessions_); if (experiment.selection_active && query_window.valid()) { @@ -4239,6 +4317,7 @@ std::string HttpServer::apply_pflash_compression( {"query_span_begin", query_span.begin}, {"query_span_end", query_span.end}, {"query_suffix_candidates", query_suffix_candidates}, + {"history_queries", history_query_spans.size()}, {"requested_query_tokens", experiment.query_tokens}, {"required_text_count", req.pflash_required.size()}, {"expected_query_ids", expected_query_ids}, diff --git a/server/src/server/http_server.h b/server/src/server/http_server.h index 8248ac36b..f59d3336e 100644 --- a/server/src/server/http_server.h +++ b/server/src/server/http_server.h @@ -304,6 +304,17 @@ PFlashTokenSpan pflash_changed_token_span( std::vector canonicalize_pflash_token_spans( std::vector spans); +// Content length (drafter tokens) up to which a user turn or assistant +// answer of a multi-turn chat is kept whole: PFLASH_CHAT_SKELETON_TOKENS, +// default 256; 0 keeps only role headers. +int pflash_chat_skeleton_tokens() noexcept; + +// Earlier user questions of a multi-turn chat that score alongside the +// current one, most recent first at weights 1/2, 1/4, ...: +// PFLASH_CHAT_HISTORY_QUERIES, default 3, at most 8; 0 scores the current +// question alone. +int pflash_chat_history_queries() noexcept; + // The parts of ``spans`` that ``minus`` does not cover. Both canonical. std::vector pflash_subtract_token_spans( const std::vector & spans, @@ -397,6 +408,14 @@ PFlashTokenSpan pflash_decoded_text_span( // are searched in the decoded prompt text, so a drafter whose vocabulary // lacks the control tokens still maps correctly. Invalid when the prompt // carries no chat markers. +struct PflashChatTurn { + int role_begin = -1; + int content_begin = -1; + int content_end = -1; + int turn_end = -1; + std::string role; // "user", "assistant", "system", "tool", ... +}; + struct PflashChatTurnSpan { int role_begin = -1; int content_begin = -1; @@ -404,6 +423,11 @@ struct PflashChatTurnSpan { int turn_end = -1; int generation_begin = -1; bool later_turns = false; + // Every turn before the generation prompt, in order; ``query_turn`` + // indexes the one above. Tool output wrapped in a user turn has role + // "tool". + std::vector turns; + int query_turn = -1; bool valid() const { return content_begin >= 0 && content_end > content_begin; diff --git a/server/test/test_server_unit.cpp b/server/test/test_server_unit.cpp index 26cb594f8..8d08f365e 100644 --- a/server/test/test_server_unit.cpp +++ b/server/test/test_server_unit.cpp @@ -840,6 +840,7 @@ struct PflashRenderedQueryTurn { std::string after; std::string closing; // content end .. turn end std::string generation; // generation prompt .. prompt end + std::vector roles; }; static PflashRenderedQueryTurn pflash_rendered_query_turn( @@ -878,6 +879,7 @@ static PflashRenderedQueryTurn pflash_rendered_query_turn( out.generation = tok.decode( {prompt.begin() + turn.generation_begin, prompt.end()}); out.later_turns = turn.later_turns; + for (const auto & each : turn.turns) out.roles.push_back(each.role); } } unlink(path.c_str()); @@ -950,6 +952,7 @@ TEST_CASE(ServerUnitFixture, TEST_ASSERT_MSG(qwen.content == "What is the answer?", qwen.content); TEST_ASSERT(qwen.header == "<|im_start|>user\n"); TEST_ASSERT(qwen.later_turns); + TEST_ASSERT(qwen.roles == std::vector({"user", "assistant", "tool"})); TEST_ASSERT_MSG(qwen.closing == "<|im_end|>", qwen.closing); TEST_ASSERT_MSG(qwen.generation == "<|im_start|>assistant\n\n\n\n\n", qwen.generation); @@ -7045,14 +7048,14 @@ TEST_CASE(ServerUnitFixture, TEST_ASSERT(tokenizer.decode({ids.begin() + query_begin, ids.begin() + query_end}) == "What is the answer?"); - // The generation prompt is pinned; the assistant and tool turns between - // the query and it are not. + // The generation prompt is pinned, and the short assistant turn stays as + // part of the conversation's skeleton; the tool output between the query + // and the generation prompt is scored, not pinned. bool generation_pinned = false; for (const auto & span : request.required_instruction_spans) { const std::string text = tokenizer.decode( {ids.begin() + span.begin, ids.begin() + span.end}); TEST_ASSERT_MSG(text.find("tool output") == std::string::npos, text); - TEST_ASSERT_MSG(text.find("Sure") == std::string::npos, text); if (span.end == (int) ids.size() && text.find("<|im_start|>assistant\n\n") != std::string::npos) { generation_pinned = true; @@ -7089,8 +7092,9 @@ static PflashSystemPromptCase pflash_long_system_prompt_case( const std::string & instruction_role = "system") { std::string system; for (int i = 0; i < 30; ++i) system += "You are helpful. "; + // Longer than the multi-turn skeleton keeps whole: droppable material. std::string history; - for (int i = 0; i < 30; ++i) history += "Sure. "; + for (int i = 0; i < 120; ++i) history += "Sure. "; PflashSystemPromptCase out; out.rendered = render_chat_template( {{instruction_role, system, ""}, @@ -7411,6 +7415,77 @@ TEST_CASE(ServerUnitFixture, unlink(path.c_str()); } +TEST_CASE(ServerUnitFixture, + test_pflash_strict_multi_turn_keeps_skeleton_and_history_queries) { + 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. "; + const std::vector messages{ + {"system", "You are helpful.", ""}, + {"user", "Here is text: " + material + "What is the first answer?", ""}, + {"assistant", "Sure.", ""}, + {"user", "What is the answer?", ""}, + }; + const std::string rendered = render_chat_template( + messages, 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", "\n", "Sure", ".", " filler", " words", " here"}, + 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 : messages) { + 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); + const auto & request = backend.last_request; + const auto & ids = request.input_ids; + const auto text_of = [&] (const PFlashTokenSpan & span) { + return tokenizer.decode({ids.begin() + span.begin, ids.begin() + span.end}); + }; + // The short assistant answer stays whole; the long first user turn keeps + // only its header, its material competes for the budget. + bool answer_kept = false; + bool material_kept = false; + for (const auto & span : request.required_instruction_spans) { + const std::string text = text_of(span); + answer_kept = answer_kept || text.find("Sure.") != std::string::npos; + material_kept = material_kept || + text.find("filler words") != std::string::npos; + } + TEST_ASSERT(answer_kept); + TEST_ASSERT(!material_kept); + // The earlier question scores alongside the current one. + TEST_ASSERT(request.history_query_spans.size() == 1); + TEST_ASSERT_MSG(text_of(request.history_query_spans[0]).find("first answer?") != + std::string::npos, + text_of(request.history_query_spans[0])); + unlink(path.c_str()); +} + TEST_CASE(ServerUnitFixture, test_pflash_legacy_chat_query_uses_last_user_turn) { // No strict-selection environment: the legacy selector derives the same From 2ce924afbea541ffdd817a3105b6a8e4cdd60932 Mon Sep 17 00:00:00 2001 From: Graffioh Date: Tue, 22 Sep 2026 20:09:17 +0000 Subject: [PATCH 18/26] feat(pflash): report compression and view details in usage timings Benchmarks could only see PFlash's multi-turn behaviour by scraping server logs. usage.timings now carries a pflash object on compressed requests: compress time, drafter input, kept and compressed tokens, the effective keep ratio, the query rule, the history-query count, the drafter session's resume point, new tokens and forward time, and the view outcome (fresh, continue, rebuild or repeat, with served, reused, delta, recalled and fresh token counts). Co-Authored-By: Claude Opus 5.5 (1M context) --- server/src/common/model_backend.h | 5 +++ server/src/deepseek4/deepseek4_backend.cpp | 6 ++++ server/src/pflash/pflash_compress.cpp | 10 ++++++ server/src/pflash/pflash_compress.h | 12 +++++++ server/src/pflash/qwen35_drafter.cpp | 2 ++ server/src/qwen35/qwen35_backend.cpp | 6 ++++ .../src/qwen35/qwen35_layer_split_adapter.cpp | 6 ++++ server/src/server/http_server.cpp | 31 +++++++++++++++++-- server/src/server/http_server.h | 5 ++- server/src/server/sse_emitter.cpp | 4 ++- server/src/server/sse_emitter.h | 3 ++ server/test/test_server_unit.cpp | 15 +++++++++ 12 files changed, 101 insertions(+), 4 deletions(-) diff --git a/server/src/common/model_backend.h b/server/src/common/model_backend.h index cc6ef062d..2e4ec6099 100644 --- a/server/src/common/model_backend.h +++ b/server/src/common/model_backend.h @@ -290,6 +290,11 @@ struct ModelBackend { // Strict selection: the input spans behind compressed_ids, ascending. // Empty when the backend does not report them (remote drafter). std::vector kept_spans; + // Drafter session reuse: the token scoring resumed from and the + // tokens it ran (-1 when unknown), and its forward time. + int scorer_resume = -1; + int scorer_new_tokens = -1; + double scorer_forward_s = 0.0; static CompressResult from_compressed_ids( std::vector ids) { diff --git a/server/src/deepseek4/deepseek4_backend.cpp b/server/src/deepseek4/deepseek4_backend.cpp index 0d3d82c68..cb0989688 100644 --- a/server/src/deepseek4/deepseek4_backend.cpp +++ b/server/src/deepseek4/deepseek4_backend.cpp @@ -3186,6 +3186,12 @@ std::vector DeepSeek4Backend::compress_batch( request.query_suffix_candidates, request.history_query_spans); result.ok = !result.compressed_ids.empty(); if (result.ok) result.kept_spans = pflash_last_kept_spans(); + if (result.ok) { + const auto & scoring = pflash_last_scoring_stats(); + result.scorer_resume = scoring.resume; + result.scorer_new_tokens = scoring.new_tokens; + result.scorer_forward_s = scoring.forward_s; + } } if (load_request->residency_action == diff --git a/server/src/pflash/pflash_compress.cpp b/server/src/pflash/pflash_compress.cpp index 387854084..f50bef0e5 100644 --- a/server/src/pflash/pflash_compress.cpp +++ b/server/src/pflash/pflash_compress.cpp @@ -181,14 +181,24 @@ void write_compression_trace( namespace { thread_local std::vector g_last_kept_spans; +thread_local PFlashScoringStats g_last_scoring_stats; } // namespace +const PFlashScoringStats & pflash_last_scoring_stats() { + return g_last_scoring_stats; +} + +void pflash_set_scoring_stats(const PFlashScoringStats & stats) { + g_last_scoring_stats = stats; +} + const std::vector & pflash_last_kept_spans() { return g_last_kept_spans; } void pflash_clear_kept_spans() { g_last_kept_spans.clear(); + g_last_scoring_stats = {}; } std::vector select_pflash_chunks( diff --git a/server/src/pflash/pflash_compress.h b/server/src/pflash/pflash_compress.h index efcdf080f..e0911a9c3 100644 --- a/server/src/pflash/pflash_compress.h +++ b/server/src/pflash/pflash_compress.h @@ -127,6 +127,18 @@ void write_compression_trace( const std::vector & pflash_last_kept_spans(); void pflash_clear_kept_spans(); +// What the last strict scoring on this thread reused: the token its drafter +// session resumed from, how many tokens it ran, how many query windows it +// scored. Cleared with the kept spans. +struct PFlashScoringStats { + int resume = -1; + int new_tokens = -1; + int query_windows = 0; + double forward_s = 0.0; +}; +const PFlashScoringStats & pflash_last_scoring_stats(); +void pflash_set_scoring_stats(const PFlashScoringStats & stats); + std::vector select_pflash_chunks( const std::vector & ids, const std::vector & token_scores, diff --git a/server/src/pflash/qwen35_drafter.cpp b/server/src/pflash/qwen35_drafter.cpp index 25c1aac4f..0be2be6a1 100644 --- a/server/src/pflash/qwen35_drafter.cpp +++ b/server/src/pflash/qwen35_drafter.cpp @@ -1201,6 +1201,8 @@ std::vector qwen35_strict_score_and_compress( std::chrono::duration(t2 - t0).count(), st.head_loaded ? "trained" : "native-block15"); std::fflush(stderr); + pflash_set_scoring_stats({resume, n_new, (int) windows.size(), + std::chrono::duration(t1 - t0).count()}); std::vector segments; bool density = experiment.candidate_score == luce::pflash::PFlashCandidateScore::Density; diff --git a/server/src/qwen35/qwen35_backend.cpp b/server/src/qwen35/qwen35_backend.cpp index c0e8b3cf4..a7e061593 100644 --- a/server/src/qwen35/qwen35_backend.cpp +++ b/server/src/qwen35/qwen35_backend.cpp @@ -1206,6 +1206,12 @@ std::vector Qwen35Backend::compress_batch( request.query_suffix_candidates, request.history_query_spans); result.ok = !result.compressed_ids.empty(); if (result.ok) result.kept_spans = pflash_last_kept_spans(); + if (result.ok) { + const auto & scoring = pflash_last_scoring_stats(); + result.scorer_resume = scoring.resume; + result.scorer_new_tokens = scoring.new_tokens; + result.scorer_forward_s = scoring.forward_s; + } if (result.ok) { std::fprintf(stderr, "[compress] %zu -> %zu tokens\n", request.input_ids.size(), result.compressed_ids.size()); diff --git a/server/src/qwen35/qwen35_layer_split_adapter.cpp b/server/src/qwen35/qwen35_layer_split_adapter.cpp index 4a3005ad5..312f4cfa7 100644 --- a/server/src/qwen35/qwen35_layer_split_adapter.cpp +++ b/server/src/qwen35/qwen35_layer_split_adapter.cpp @@ -1397,6 +1397,12 @@ Qwen35LayerSplitAdapter::compress(const ModelBackend::CompressRequest & req) { req.query_suffix_candidates, req.history_query_spans); result.ok = !result.compressed_ids.empty(); if (result.ok) result.kept_spans = pflash_last_kept_spans(); + if (result.ok) { + const auto & scoring = pflash_last_scoring_stats(); + result.scorer_resume = scoring.resume; + result.scorer_new_tokens = scoring.new_tokens; + result.scorer_forward_s = scoring.forward_s; + } if (result.ok) { std::fprintf(stderr, "[target-split][compress] %zu -> %zu tokens\n", req.input_ids.size(), result.compressed_ids.size()); diff --git a/server/src/server/http_server.cpp b/server/src/server/http_server.cpp index bb02f3be6..8cdbbaa0d 100644 --- a/server/src/server/http_server.cpp +++ b/server/src/server/http_server.cpp @@ -4364,6 +4364,7 @@ std::string HttpServer::apply_pflash_compression( ModelBackend::CompressResult result; std::vector final_tokens; + const auto compress_started = std::chrono::steady_clock::now(); for (int attempt = 0; ; ++attempt) { result = {}; if (config_.pflash_remote_drafter) { @@ -4461,11 +4462,25 @@ std::string HttpServer::apply_pflash_compression( std::to_string(target_ceiling) + ")"; } } + prepared.pflash_stats = { + {"compress_ms", std::round(std::chrono::duration( + std::chrono::steady_clock::now() - compress_started).count() * 10.0) / 10.0}, + {"drafter_input_tokens", compress_request.input_ids.size()}, + {"kept_tokens", kept_tokens}, + {"keep_ratio", compress_request.keep_ratio}, + {"compressed_tokens", final_tokens.size()}, + {"query_rule", parser_selection_rule}, + {"history_queries", history_query_spans.size()}, + {"scorer_resume", result.scorer_resume}, + {"scorer_new_tokens", result.scorer_new_tokens}, + {"scorer_forward_ms", std::round(result.scorer_forward_s * 10000.0) / 10.0}, + }; if (experiment.selection_active && messages_input && chat_turn.valid() && !result.kept_spans.empty()) { final_tokens = continue_pflash_chat_view( req, compress_request.input_ids, chat_turn, result.kept_spans, - std::move(final_tokens), prepared.snapshot_cut); + std::move(final_tokens), prepared.snapshot_cut, + prepared.pflash_stats["view"]); } prepared.tokens = std::move(final_tokens); prepared.compressed = true; @@ -4483,8 +4498,10 @@ std::vector HttpServer::continue_pflash_chat_view( const http_detail::PflashChatTurnSpan & turn, const std::vector & kept_spans, std::vector fresh, - int & snapshot_cut) { + int & snapshot_cut, + json & stats) { snapshot_cut = -1; + stats = nullptr; const char * disabled = std::getenv("PFLASH_CHAT_VIEW"); if (disabled && std::string(disabled) == "0") return fresh; const int input = (int) drafter_ids.size(); @@ -4527,6 +4544,8 @@ std::vector HttpServer::continue_pflash_chat_view( std::fprintf(stderr, "[pflash-view] %s turn=%d served=%zu\n", why, turns, fresh.size()); std::fflush(stderr); + stats = {{"mode", why}, {"turn", turns}, {"served_tokens", fresh.size()}, + {"fresh_tokens", fresh.size()}}; pflash_views_.remember(std::move(next)); return std::move(fresh); }; @@ -4536,6 +4555,9 @@ std::vector HttpServer::continue_pflash_chat_view( std::fprintf(stderr, "[pflash-view] repeat turn=%d served=%zu\n", view.turns, view.view_tokens.size()); std::fflush(stderr); + stats = {{"mode", "repeat"}, {"turn", view.turns}, + {"served_tokens", view.view_tokens.size()}, + {"fresh_tokens", fresh.size()}}; snapshot_cut = view.view_gen_begin; return view.view_tokens; } @@ -4635,6 +4657,10 @@ std::vector HttpServer::continue_pflash_chat_view( next.turns, served.size(), view.view_gen_begin, delta_tokens.size(), recalled_tokens, fresh.size()); std::fflush(stderr); + stats = {{"mode", "continue"}, {"turn", next.turns}, + {"served_tokens", served.size()}, {"reused_tokens", view.view_gen_begin}, + {"delta_tokens", delta_tokens.size()}, + {"recalled_tokens", recalled_tokens}, {"fresh_tokens", fresh.size()}}; pflash_views_.remember(std::move(next)); return served; } @@ -5924,6 +5950,7 @@ void HttpServer::process_job(ServerJob * job) { effective_prompt_tokens - cached_prefix_tokens, effective_prompt_tokens, agent_turn_cache_hit, + prepared.pflash_stats, }; // Record performance for /status page. diff --git a/server/src/server/http_server.h b/server/src/server/http_server.h index f59d3336e..6efba1f8d 100644 --- a/server/src/server/http_server.h +++ b/server/src/server/http_server.h @@ -628,6 +628,8 @@ class HttpServer { // else asks for one: a multi-turn PFlash view sets the start of its // generation prompt, where the next turn's prompt branches off. int snapshot_cut = -1; + // PFlash details for usage.timings.pflash (see build_timings_json). + nlohmann::json pflash_stats; int error_status = 0; std::string error; }; @@ -648,7 +650,8 @@ class HttpServer { const http_detail::PflashChatTurnSpan & turn, const std::vector & kept_spans, std::vector fresh, - int & snapshot_cut); + int & snapshot_cut, + nlohmann::json & stats); bool forward_upstream(ServerJob * job, const ParsedRequest & req, const PreparedPrompt & prepared); diff --git a/server/src/server/sse_emitter.cpp b/server/src/server/sse_emitter.cpp index db86a1438..f50595a6d 100644 --- a/server/src/server/sse_emitter.cpp +++ b/server/src/server/sse_emitter.cpp @@ -76,7 +76,7 @@ json build_timings_json(const GenTimings & t, int completion_tokens) { const double decode_ms = round1(t.decode_s * 1000.0); const double tps = t.decode_s > 0.0 ? round1((double)completion_tokens / t.decode_s) : 0.0; - return json{ + json out{ {"prefill_ms", prefill_ms}, {"decode_ms", decode_ms}, {"decode_tokens_per_sec", tps}, @@ -86,6 +86,8 @@ json build_timings_json(const GenTimings & t, int completion_tokens) { {"effective_prompt_tokens", t.effective_prompt_tokens}, {"agent_turn_cache_hit", t.agent_turn_cache_hit} }; + if (!t.pflash.is_null()) out["pflash"] = t.pflash; + return out; } // ─── Constructor ──────────────────────────────────────────────────────── diff --git a/server/src/server/sse_emitter.h b/server/src/server/sse_emitter.h index bfc08082a..ae61c116e 100644 --- a/server/src/server/sse_emitter.h +++ b/server/src/server/sse_emitter.h @@ -48,6 +48,9 @@ struct GenTimings { int prefilled_tokens = 0; int effective_prompt_tokens = 0; bool agent_turn_cache_hit = false; + // PFlash compression details (compress time, kept/served tokens, the + // multi-turn view and drafter-session outcome); null when not compressed. + nlohmann::json pflash; }; // Build the `timings` sub-object emitted under `usage`. diff --git a/server/test/test_server_unit.cpp b/server/test/test_server_unit.cpp index 8d08f365e..ec8d26736 100644 --- a/server/test/test_server_unit.cpp +++ b/server/test/test_server_unit.cpp @@ -1044,6 +1044,14 @@ TEST_CASE(ServerUnitFixture, test_pflash_selection_cache_and_continuation_policy TEST_ASSERT(!http_detail::pflash_full_cache_restore_allowed(true)); } +TEST_CASE(ServerUnitFixture, test_timings_json_carries_pflash_details) { + GenTimings timings; + TEST_ASSERT(!build_timings_json(timings, 0).contains("pflash")); + timings.pflash = {{"compress_ms", 12.5}, {"view", {{"mode", "continue"}}}}; + const auto out = build_timings_json(timings, 0); + TEST_ASSERT(out["pflash"]["view"]["mode"] == "continue"); +} + TEST_CASE(ServerUnitFixture, test_pflash_subtract_token_spans) { const std::vector spans{{0, 10}, {20, 30}, {40, 50}}; const std::vector minus{{5, 22}, {25, 26}, {40, 50}}; @@ -7365,6 +7373,7 @@ TEST_CASE(ServerUnitFixture, std::vector served1; std::vector served2; std::vector served3; + std::vector modes; { HttpServer server(engine, tokenizer, config); server.set_drafter_tokenizer(&tokenizer); @@ -7379,6 +7388,11 @@ TEST_CASE(ServerUnitFixture, TEST_ASSERT_MSG(prepared.error.empty(), prepared.error); TEST_ASSERT(prepared.compressed); served = prepared.tokens; + // usage.timings.pflash reports the view outcome. + TEST_ASSERT(prepared.pflash_stats.contains("view")); + modes.push_back(prepared.pflash_stats["view"].value("mode", "")); + TEST_ASSERT(prepared.pflash_stats["view"].value("served_tokens", 0) == + (int) served.size()); // The snapshot lands where the next turn's prompt branches off. const auto generation = tokenizer.encode("<|im_start|>assistant\n\n"); @@ -7412,6 +7426,7 @@ TEST_CASE(ServerUnitFixture, TEST_ASSERT(text2.find("beta facts", recall) != std::string::npos); TEST_ASSERT(question != std::string::npos && question > recall); TEST_ASSERT(served3 == served2); + TEST_ASSERT(modes == std::vector({"fresh", "continue", "repeat"})); unlink(path.c_str()); } From 71bee68d073bcbb13295a2a9674d501a6509ef8e Mon Sep 17 00:00:00 2001 From: Graffioh Date: Wed, 23 Sep 2026 05:55:51 +0000 Subject: [PATCH 19/26] feat(pflash): compress only large follow-ups; recall switch; served trace A continuing turn appended all its new material verbatim, whatever its size, and a large paste or tool output only got compressed indirectly, by outgrowing the view and forcing a rebuild of everything. What a turn adds is now appended verbatim below PFLASH_CHAT_COMPRESS_NEW_TOKENS (default 16384), the way full prefill appends a follow-up; from there on only the parts the fresh selection keeps of the new material are appended, and the view before it stays cached (view mode "continue-compressed"). PFLASH_CHAT_RECALL=0 turns recall off, and then a small follow-up (like a repeated prompt) is served before any compression runs: no drafter, no scoring, exactly full prefill's cost on top of the cached view. The view logic is one function the server calls before compressing (for the turns that need no scoring) and after. PFLASH_VIEW_TRACE_PATH appends every compressed request's served prompt text with its PFlash details as JSONL, so an evaluation can tell evidence the selection dropped from answers the model got wrong. Co-Authored-By: Claude Opus 5.5 (1M context) --- server/README.md | 9 ++ server/src/server/http_server.cpp | 203 ++++++++++++++++++++++-------- server/src/server/http_server.h | 27 +++- server/test/test_server_unit.cpp | 113 +++++++++++++++++ 4 files changed, 292 insertions(+), 60 deletions(-) diff --git a/server/README.md b/server/README.md index e91155c15..f45b8f73f 100644 --- a/server/README.md +++ b/server/README.md @@ -417,6 +417,15 @@ view grows past twice the fresh prompt, or past the context, the fresh prompt starts a new view. `PFLASH_CHAT_VIEW=0` serves the fresh compression every turn. +What a turn adds is appended verbatim while it is small, the way full +prefill appends a follow-up; from `PFLASH_CHAT_COMPRESS_NEW_TOKENS` (default +16384) tokens of new material (a pasted document, a large tool output) only +what the fresh selection keeps of it is appended, and the view before it +stays cached. `PFLASH_CHAT_RECALL=0` turns recall off: a small follow-up is +then served without running the drafter at all. `PFLASH_VIEW_TRACE_PATH` +appends each compressed request's served prompt as JSONL, for evidence +checks in evaluations. + Every other turn of a multi-turn chat keeps its role header, and user turns and assistant answers up to `PFLASH_CHAT_SKELETON_TOKENS` (default 256 drafter tokens; 0 keeps headers only) stay whole: the conversation's diff --git a/server/src/server/http_server.cpp b/server/src/server/http_server.cpp index 8cdbbaa0d..2fe9ea82f 100644 --- a/server/src/server/http_server.cpp +++ b/server/src/server/http_server.cpp @@ -627,6 +627,20 @@ int pflash_chat_skeleton_tokens() noexcept { return (int) (std::min)(value, 1L << 20); } +bool pflash_chat_recall() noexcept { + const char * raw = std::getenv("PFLASH_CHAT_RECALL"); + return !(raw && std::string(raw) == "0"); +} + +int pflash_chat_compress_new_tokens() noexcept { + const char * raw = std::getenv("PFLASH_CHAT_COMPRESS_NEW_TOKENS"); + if (!raw || !*raw) return 16384; + char * end = nullptr; + const long value = std::strtol(raw, &end, 10); + if (end == raw || *end != '\0' || value < 1) return 16384; + return (int) (std::min)(value, 1L << 30); +} + int pflash_chat_history_queries() noexcept { const char * raw = std::getenv("PFLASH_CHAT_HISTORY_QUERIES"); if (!raw || !*raw) return 3; @@ -4362,6 +4376,29 @@ std::string HttpServer::apply_pflash_compression( } const float requested_keep_ratio = compress_request.keep_ratio; + // A turn the conversation's view serves without scoring -- the same + // prompt again, or a small follow-up appended verbatim with recall off + // -- skips the drafter altogether. + if (experiment.selection_active && messages_input && chat_turn.valid() && + !config_.pflash_remote_drafter) { + std::vector served; + json view_stats; + if (serve_pflash_chat_view( + req, compress_request.input_ids, chat_turn, nullptr, nullptr, + served, prepared.snapshot_cut, view_stats)) { + prepared.tokens = std::move(served); + prepared.compressed = true; + prepared.pflash_stats = { + {"compress_ms", 0.0}, + {"drafter_input_tokens", compress_request.input_ids.size()}, + {"query_rule", parser_selection_rule}, + {"view", view_stats}, + }; + trace_pflash_served(req, prepared); + return {}; + } + } + ModelBackend::CompressResult result; std::vector final_tokens; const auto compress_started = std::chrono::steady_clock::now(); @@ -4477,13 +4514,17 @@ std::string HttpServer::apply_pflash_compression( }; if (experiment.selection_active && messages_input && chat_turn.valid() && !result.kept_spans.empty()) { - final_tokens = continue_pflash_chat_view( - req, compress_request.input_ids, chat_turn, result.kept_spans, - std::move(final_tokens), prepared.snapshot_cut, - prepared.pflash_stats["view"]); + std::vector served; + if (serve_pflash_chat_view( + req, compress_request.input_ids, chat_turn, &final_tokens, + &result.kept_spans, served, prepared.snapshot_cut, + prepared.pflash_stats["view"])) { + final_tokens = std::move(served); + } } prepared.tokens = std::move(final_tokens); prepared.compressed = true; + trace_pflash_served(req, prepared); std::fprintf(stderr, "[pflash] %d -> %d -> %d tokens (%.1f%% kept)\n", prompt_tokens, (int) result.compressed_ids.size(), @@ -4492,21 +4533,38 @@ std::string HttpServer::apply_pflash_compression( return {}; } -std::vector HttpServer::continue_pflash_chat_view( +void HttpServer::trace_pflash_served( + const ParsedRequest & req, const PreparedPrompt & prepared) { + const char * path = std::getenv("PFLASH_VIEW_TRACE_PATH"); + if (!path || !*path) return; + const json record = { + {"schema_version", 1}, + {"prompt_tokens", req.prompt_tokens.size()}, + {"served_tokens", prepared.tokens.size()}, + {"pflash", prepared.pflash_stats}, + {"served_text", tokenizer_.decode(prepared.tokens)}, + }; + std::ofstream out(path, std::ios::app); + if (out) out << record.dump(-1, ' ', false, json::error_handler_t::replace) << "\n"; +} + +bool HttpServer::serve_pflash_chat_view( const ParsedRequest & req, const std::vector & drafter_ids, const http_detail::PflashChatTurnSpan & turn, - const std::vector & kept_spans, - std::vector fresh, + const std::vector * fresh, + const std::vector * kept_spans, + std::vector & served, int & snapshot_cut, json & stats) { snapshot_cut = -1; stats = nullptr; + const bool compressed = fresh != nullptr && kept_spans != nullptr; const char * disabled = std::getenv("PFLASH_CHAT_VIEW"); - if (disabled && std::string(disabled) == "0") return fresh; + if (disabled && std::string(disabled) == "0") return false; const int input = (int) drafter_ids.size(); if (turn.generation_begin <= 0 || turn.generation_begin >= input) { - return fresh; + return false; } // The generation prompt, in target tokens: the raw prompt and every // served prompt end with it (strict selection keeps it verbatim). @@ -4520,8 +4578,8 @@ std::vector HttpServer::continue_pflash_chat_view( tokens.end() - (long) generation.size()); }; if (!ends_with_generation(req.prompt_tokens) || - !ends_with_generation(fresh)) { - return fresh; + (compressed && !ends_with_generation(*fresh))) { + return false; } const int raw_gen_begin = (int) (req.prompt_tokens.size() - generation.size()); @@ -4536,46 +4594,60 @@ std::vector HttpServer::continue_pflash_chat_view( const bool continues = pflash_views_.find(req.prompt_tokens, drafter_ids, view); const auto serve_fresh = [&] (const char * why, int turns) { - next.view_tokens = fresh; - next.view_gen_begin = (int) (fresh.size() - generation.size()); - next.spans = kept_spans; + next.view_tokens = *fresh; + next.view_gen_begin = (int) (fresh->size() - generation.size()); + next.spans = *kept_spans; next.turns = turns; snapshot_cut = next.view_gen_begin; std::fprintf(stderr, - "[pflash-view] %s turn=%d served=%zu\n", why, turns, fresh.size()); + "[pflash-view] %s turn=%d served=%zu\n", why, turns, fresh->size()); std::fflush(stderr); - stats = {{"mode", why}, {"turn", turns}, {"served_tokens", fresh.size()}, - {"fresh_tokens", fresh.size()}}; + stats = {{"mode", why}, {"turn", turns}, {"served_tokens", fresh->size()}, + {"fresh_tokens", fresh->size()}}; + served = *fresh; pflash_views_.remember(std::move(next)); - return std::move(fresh); + return true; }; - if (!continues) return serve_fresh("fresh", 1); - if (view.raw_tokens == req.prompt_tokens) { + if (continues && view.raw_tokens == req.prompt_tokens) { // The same prompt again (a retry): serve what was served. std::fprintf(stderr, "[pflash-view] repeat turn=%d served=%zu\n", view.turns, view.view_tokens.size()); std::fflush(stderr); stats = {{"mode", "repeat"}, {"turn", view.turns}, - {"served_tokens", view.view_tokens.size()}, - {"fresh_tokens", fresh.size()}}; + {"served_tokens", view.view_tokens.size()}}; + if (compressed) stats["fresh_tokens"] = fresh->size(); snapshot_cut = view.view_gen_begin; - return view.view_tokens; - } - if (view.drafter_gen_begin >= turn.generation_begin || - view.view_gen_begin <= 0 || - (size_t) view.view_gen_begin > view.view_tokens.size()) { - return serve_fresh("fresh", 1); + served = view.view_tokens; + return true; } + const bool usable = continues && + view.drafter_gen_begin < turn.generation_begin && + view.view_gen_begin > 0 && + (size_t) view.view_gen_begin <= view.view_tokens.size(); + if (!usable) return compressed && serve_fresh("fresh", 1); + + // What this turn adds to the conversation, in target tokens. A small + // follow-up is appended verbatim, the way full prefill appends it; one + // of PFLASH_CHAT_COMPRESS_NEW_TOKENS or more (a pasted document, a large + // tool output) goes through the compressor, and only the new material + // is compressed, so the view it extends stays cached. + const int new_tokens = raw_gen_begin + (int) generation.size() - + view.raw_gen_begin; + const bool compress_new = + new_tokens >= http_detail::pflash_chat_compress_new_tokens(); + const bool new_question = turn.role_begin >= view.drafter_gen_begin; + const bool recall = new_question && http_detail::pflash_chat_recall(); + if (!compressed && (compress_new || recall)) return false; // Recall: what the fresh selection keeps for the new query that the view // does not hold. Only a new user turn brings a new query; an agent step // (assistant call plus tool output) appends without recalling. std::vector recalled; - if (turn.role_begin >= view.drafter_gen_begin) { + if (compressed && recall) { auto in_view = view.spans; in_view.push_back({view.drafter_gen_begin, input}); recalled = http_detail::pflash_subtract_token_spans( - kept_spans, + *kept_spans, http_detail::canonicalize_pflash_token_spans(std::move(in_view))); } std::string recall_block; @@ -4618,51 +4690,72 @@ std::vector HttpServer::continue_pflash_chat_view( } } - // The previous view without its generation prompt, then this turn's new - // tokens from where that generation prompt started; recalled excerpts - // open the new user turn's content, after everything the target cached. + // This turn's new tokens from where the previous generation prompt + // started: all of them, or the parts the fresh selection keeps when they + // are compressed. Recalled excerpts open the new user turn's content, + // after everything the target cached. const auto decode_range = [&] (int begin, int end) { - return drafter_tokenizer_->decode(std::vector( - drafter_ids.begin() + begin, drafter_ids.begin() + end)); + return end > begin + ? drafter_tokenizer_->decode(std::vector( + drafter_ids.begin() + begin, drafter_ids.begin() + end)) + : std::string(); }; - const std::string delta = recall_block.empty() - ? decode_range(view.drafter_gen_begin, input) - : decode_range(view.drafter_gen_begin, turn.content_begin) + - recall_block + decode_range(turn.content_begin, input); - std::vector served(view.view_tokens.begin(), - view.view_tokens.begin() + view.view_gen_begin); + std::vector delta_spans; + if (compress_new) { + delta_spans = http_detail::pflash_subtract_token_spans( + *kept_spans, {{0, view.drafter_gen_begin}}); + } else { + delta_spans.push_back({view.drafter_gen_begin, input}); + } + const int split = new_question && !recall_block.empty() + ? turn.content_begin : input; + std::string delta; + for (const auto & span : delta_spans) { + delta += decode_range(span.begin, (std::min)(span.end, split)); + } + delta += recall_block; + for (const auto & span : delta_spans) { + delta += decode_range((std::max)(span.begin, split), span.end); + } + served.assign(view.view_tokens.begin(), + view.view_tokens.begin() + view.view_gen_begin); const auto delta_tokens = tokenizer_.encode(delta); served.insert(served.end(), delta_tokens.begin(), delta_tokens.end()); if (!ends_with_generation(served)) { - return serve_fresh("fresh", 1); + return compressed && serve_fresh("fresh", 1); } // Rebuild when the view outgrew what a fresh selection keeps, or the // context: the fresh prompt starts a new view, prefilled from scratch. - const bool outgrown = served.size() > 2 * fresh.size() || - (config_.max_ctx > 0 && - (int) served.size() + req.max_output > config_.max_ctx); - if (outgrown) return serve_fresh("rebuild", view.turns + 1); + const bool too_long = config_.max_ctx > 0 && + (int) served.size() + req.max_output > config_.max_ctx; + const bool outgrown = compressed && served.size() > 2 * fresh->size(); + if (too_long || outgrown) { + return compressed && serve_fresh("rebuild", view.turns + 1); + } auto spans = view.spans; - spans.push_back({view.drafter_gen_begin, input}); + spans.insert(spans.end(), delta_spans.begin(), delta_spans.end()); spans.insert(spans.end(), recalled.begin(), recalled.end()); next.view_tokens = served; next.view_gen_begin = (int) (served.size() - generation.size()); next.spans = http_detail::canonicalize_pflash_token_spans(std::move(spans)); next.turns = view.turns + 1; snapshot_cut = next.view_gen_begin; + const char * mode = compress_new ? "continue-compressed" : "continue"; std::fprintf(stderr, - "[pflash-view] continue turn=%d served=%zu reused=%d delta=%zu " - "recalled=%d fresh=%zu\n", - next.turns, served.size(), view.view_gen_begin, delta_tokens.size(), - recalled_tokens, fresh.size()); + "[pflash-view] %s turn=%d served=%zu reused=%d new=%d delta=%zu " + "recalled=%d fresh=%d\n", + mode, next.turns, served.size(), view.view_gen_begin, new_tokens, + delta_tokens.size(), recalled_tokens, + compressed ? (int) fresh->size() : -1); std::fflush(stderr); - stats = {{"mode", "continue"}, {"turn", next.turns}, + stats = {{"mode", mode}, {"turn", next.turns}, {"served_tokens", served.size()}, {"reused_tokens", view.view_gen_begin}, - {"delta_tokens", delta_tokens.size()}, - {"recalled_tokens", recalled_tokens}, {"fresh_tokens", fresh.size()}}; + {"new_tokens", new_tokens}, {"delta_tokens", delta_tokens.size()}, + {"recalled_tokens", recalled_tokens}}; + if (compressed) stats["fresh_tokens"] = fresh->size(); pflash_views_.remember(std::move(next)); - return served; + return true; } HttpServer::PreparedPrompt HttpServer::prepare_prompt( diff --git a/server/src/server/http_server.h b/server/src/server/http_server.h index 6efba1f8d..6682035a0 100644 --- a/server/src/server/http_server.h +++ b/server/src/server/http_server.h @@ -315,6 +315,13 @@ int pflash_chat_skeleton_tokens() noexcept; // question alone. int pflash_chat_history_queries() noexcept; +// Multi-turn follow-ups: PFLASH_CHAT_RECALL=0 turns recall off (a small +// follow-up is then appended exactly as full prefill appends it, without +// scoring); PFLASH_CHAT_COMPRESS_NEW_TOKENS (default 16384) is the size of +// new material in one turn from which it is compressed instead of appended. +bool pflash_chat_recall() noexcept; +int pflash_chat_compress_new_tokens() noexcept; + // The parts of ``spans`` that ``minus`` does not cover. Both canonical. std::vector pflash_subtract_token_spans( const std::vector & spans, @@ -642,16 +649,26 @@ class HttpServer { std::string apply_pflash_compression(const ParsedRequest & req, PreparedPrompt & prepared); // Multi-turn: serve the conversation's previous view plus this turn's - // new tokens (and the segments the fresh selection wants that the view - // lacks) when one continues into this prompt; else the fresh prompt. - std::vector continue_pflash_chat_view( + // new material -- verbatim below PFLASH_CHAT_COMPRESS_NEW_TOKENS, else + // the parts the fresh selection keeps -- with the segments the fresh + // selection wants that the view lacks recalled at the new user turn. + // Without a fresh compression (``fresh`` null) it serves only what needs + // no scoring: a repeated prompt, or a small follow-up with recall off; + // it returns false for the caller to compress. With one it always + // serves: a continued view, or the fresh prompt as a new view. + bool serve_pflash_chat_view( const ParsedRequest & req, const std::vector & drafter_ids, const http_detail::PflashChatTurnSpan & turn, - const std::vector & kept_spans, - std::vector fresh, + const std::vector * fresh, + const std::vector * kept_spans, + std::vector & served, int & snapshot_cut, nlohmann::json & stats); + // PFLASH_VIEW_TRACE_PATH: one JSONL record per compressed request with + // the served prompt's text, for evidence checks in evaluations. + void trace_pflash_served(const ParsedRequest & req, + const PreparedPrompt & prepared); bool forward_upstream(ServerJob * job, const ParsedRequest & req, const PreparedPrompt & prepared); diff --git a/server/test/test_server_unit.cpp b/server/test/test_server_unit.cpp index ec8d26736..32d54c372 100644 --- a/server/test/test_server_unit.cpp +++ b/server/test/test_server_unit.cpp @@ -7501,6 +7501,119 @@ TEST_CASE(ServerUnitFixture, unlink(path.c_str()); } +// Two turns over one document through prepare_prompt; returns the served +// prompts, the view modes and how often the compressor ran. +struct PflashTwoTurnRun { + std::vector> served; + std::vector modes; + int compress_calls = 0; + std::string text2; +}; + +static PflashTwoTurnRun pflash_two_turn_run(const std::string & follow_up) { + std::string system; + for (int i = 0; i < 20; ++i) system += "You are helpful. "; + const std::string document = + "alpha facts live here. filler filler filler. beta facts live here."; + const std::vector turn1{ + {"system", system, ""}, + {"user", document + " Question one?", ""}, + }; + auto turn2 = turn1; + turn2.push_back({"assistant", "Answer one.", ""}); + turn2.push_back({"user", follow_up + " Question two?", ""}); + const auto render = [] (const std::vector & messages) { + return render_chat_template(messages, ChatFormat::QWEN3, + /*add_generation_prompt=*/true, + /*enable_thinking=*/true); + }; + const std::string path = write_pflash_bpe_tokenizer_fixture( + {"alpha", " facts", "beta", " live", " here", ".", " filler", + "Question", " one", " two", "?", "Answer", "user", "assistant", + "system", "\n", "You", " are", " helpful", " pasted", " notes"}, + render(turn2) + + "[Earlier in this conversation]\n[End of earlier excerpts]\n"); + Tokenizer tokenizer; + TEST_ASSERT(tokenizer.load_from_gguf(path.c_str())); + auto backend_owner = std::make_unique(); + MockPflashSpanBackend & backend = *backend_owner; + backend.pick = [&] (const ModelBackend::CompressRequest & request) { + const auto span = http_detail::pflash_decoded_text_span( + tokenizer, request.input_ids, 0, (int) request.input_ids.size(), + "alpha facts"); + return span.begin < 0 ? std::vector{} + : std::vector{span}; + }; + LuceEngine engine(std::move(backend_owner)); + ServerConfig config; + config.pflash_mode = ServerConfig::PflashMode::ALWAYS; + config.pflash_keep_ratio = 1.0f; + config.max_ctx = 8192; + config.prefix_cache_cap = 0; + config.prefill_cache_cap = 0; + PflashTwoTurnRun run; + { + HttpServer server(engine, tokenizer, config); + server.set_drafter_tokenizer(&tokenizer); + const std::vector *> turns{&turn1, &turn2}; + for (const auto * messages : turns) { + ParsedRequest request; + request.format = ApiFormat::OPENAI_CHAT; + json wire = json::array(); + for (const auto & message : *messages) { + wire.push_back({{"role", message.role}, {"content", message.content}}); + } + request.messages = wire; + request.prompt_tokens = tokenizer.encode(render(*messages)); + const auto prepared = HttpServerTestAccess::prepare_prompt(server, request); + TEST_ASSERT_MSG(prepared.error.empty(), prepared.error); + run.served.push_back(prepared.tokens); + run.modes.push_back(prepared.pflash_stats["view"].value("mode", "")); + } + } + run.compress_calls = backend.compress_calls; + run.text2 = tokenizer.decode(run.served[1]); + unlink(path.c_str()); + return run; +} + +TEST_CASE(ServerUnitFixture, + test_pflash_chat_view_small_follow_up_without_recall_skips_the_drafter) { + luce_test::ScopedEnvVar mode{"PFLASH_SELECT_MODE", "budget_only"}; + luce_test::ScopedEnvVar recall{"PFLASH_CHAT_RECALL", "0"}; + const auto run = pflash_two_turn_run(""); + // Turn 2 is appended exactly as full prefill appends it: no scoring. + TEST_ASSERT(run.compress_calls == 1); + TEST_ASSERT(run.modes == std::vector({"fresh", "continue"})); + TEST_ASSERT(run.text2.find("[Earlier in this conversation]") == std::string::npos); + TEST_ASSERT(run.text2.find("Answer one.") != std::string::npos); + TEST_ASSERT(run.text2.find("Question two?") != std::string::npos); +} + +TEST_CASE(ServerUnitFixture, + test_pflash_chat_view_compresses_large_follow_ups_only) { + luce_test::ScopedEnvVar mode{"PFLASH_SELECT_MODE", "budget_only"}; + luce_test::ScopedEnvVar threshold{"PFLASH_CHAT_COMPRESS_NEW_TOKENS", "40"}; + // Wide enough that the pinned query covers the whole question. + luce_test::ScopedEnvVar query{"PFLASH_SELECT_QUERY_TOKENS", "16"}; + std::string pasted; + for (int i = 0; i < 12; ++i) pasted += " pasted notes filler."; + const auto run = pflash_two_turn_run(pasted); + TEST_ASSERT(run.compress_calls == 2); + TEST_ASSERT(run.modes == std::vector({"fresh", "continue-compressed"})); + // The view before the new material is reused token for token... + const auto & first = run.served[0]; + const auto & second = run.served[1]; + TEST_ASSERT(second.size() > 0 && first.size() > 16); + const size_t prefix = first.size() - 8; + TEST_ASSERT(std::equal(first.begin(), first.begin() + (long) (prefix - 8), + second.begin())); + // ...and the pasted material is compressed: only the pinned question and + // what the selection keeps survive, not the whole paste. + TEST_ASSERT_MSG(run.text2.find("Question two?") != std::string::npos, run.text2); + TEST_ASSERT(run.text2.find(pasted) == std::string::npos); +} + TEST_CASE(ServerUnitFixture, test_pflash_legacy_chat_query_uses_last_user_turn) { // No strict-selection environment: the legacy selector derives the same From d90c27ef95f74f682b103af298ae0c0bc6ff651e Mon Sep 17 00:00:00 2001 From: Graffioh Date: Wed, 23 Sep 2026 06:40:29 +0000 Subject: [PATCH 20/26] feat(pflash): recall by attention lift; optional paragraph join On the multi-turn development chats every PFlash failure still had its gold evidence in the served prompt; the losses came from what surrounded it. For a content-free follow-up ("which documents support that answer?") recall filled its budget with 20-odd low-relevance documents right beside the question. Segment lifts from the head (mass per token relative to uniform attention) separate evidence (typically 20-160x) from background (90th percentile about 1x). The selector now reports every candidate's lift (CompressResult::candidate_lifts), and recall takes only segments the view lacks with a lift of at least PFLASH_CHAT_RECALL_MIN_LIFT (default 8), strongest first, up to PFLASH_CHAT_RECALL_TOKENS (default 2048 drafter tokens); without lifts it keeps the previous rule. PFLASH_SELECT_PARAGRAPH_JOIN=1 (off by default, under test) rebuilds the compressed text from the kept spans with a paragraph break between pieces that were not adjacent ("...other bands.Document 1:" otherwise), for fresh selections and compressed follow-ups; the breaks do not count against the target-token ceiling. Co-Authored-By: Claude Opus 5.5 (1M context) --- server/README.md | 9 +- server/src/common/model_backend.h | 3 + server/src/deepseek4/deepseek4_backend.cpp | 3 + server/src/pflash/pflash_compress.cpp | 21 +++ server/src/pflash/pflash_compress.h | 10 ++ server/src/qwen35/qwen35_backend.cpp | 3 + .../src/qwen35/qwen35_layer_split_adapter.cpp | 3 + server/src/server/http_server.cpp | 146 ++++++++++++++++-- server/src/server/http_server.h | 22 +++ server/test/test_server_unit.cpp | 48 ++++++ 10 files changed, 252 insertions(+), 16 deletions(-) diff --git a/server/README.md b/server/README.md index f45b8f73f..b44f8c7a2 100644 --- a/server/README.md +++ b/server/README.md @@ -422,7 +422,14 @@ prefill appends a follow-up; from `PFLASH_CHAT_COMPRESS_NEW_TOKENS` (default 16384) tokens of new material (a pasted document, a large tool output) only what the fresh selection keeps of it is appended, and the view before it stays cached. `PFLASH_CHAT_RECALL=0` turns recall off: a small follow-up is -then served without running the drafter at all. `PFLASH_VIEW_TRACE_PATH` +then served without running the drafter at all. Recall takes only segments +the new question clearly attends to: attention lift (mass per token relative +to uniform attention) of at least `PFLASH_CHAT_RECALL_MIN_LIFT` (default 8), +strongest first, up to `PFLASH_CHAT_RECALL_TOKENS` (default 2048 drafter +tokens), so a content-free follow-up ("which documents support that?") +recalls next to nothing instead of filling the budget with noise beside the +question. `PFLASH_SELECT_PARAGRAPH_JOIN=1` joins kept pieces that were not +adjacent with a paragraph break (off by default). `PFLASH_VIEW_TRACE_PATH` appends each compressed request's served prompt as JSONL, for evidence checks in evaluations. diff --git a/server/src/common/model_backend.h b/server/src/common/model_backend.h index 2e4ec6099..ce2e8c017 100644 --- a/server/src/common/model_backend.h +++ b/server/src/common/model_backend.h @@ -292,6 +292,9 @@ struct ModelBackend { std::vector kept_spans; // Drafter session reuse: the token scoring resumed from and the // tokens it ran (-1 when unknown), and its forward time. + // Strict selection with the head: every candidate's attention lift + // (mass per token relative to uniform); empty when unknown. + std::vector> candidate_lifts; int scorer_resume = -1; int scorer_new_tokens = -1; double scorer_forward_s = 0.0; diff --git a/server/src/deepseek4/deepseek4_backend.cpp b/server/src/deepseek4/deepseek4_backend.cpp index cb0989688..44126e8f4 100644 --- a/server/src/deepseek4/deepseek4_backend.cpp +++ b/server/src/deepseek4/deepseek4_backend.cpp @@ -3191,6 +3191,9 @@ std::vector DeepSeek4Backend::compress_batch( result.scorer_resume = scoring.resume; result.scorer_new_tokens = scoring.new_tokens; result.scorer_forward_s = scoring.forward_s; + for (const auto & candidate : pflash_last_candidate_lifts()) { + result.candidate_lifts.push_back({candidate.span, candidate.lift}); + } } } diff --git a/server/src/pflash/pflash_compress.cpp b/server/src/pflash/pflash_compress.cpp index f50bef0e5..f2a5fb603 100644 --- a/server/src/pflash/pflash_compress.cpp +++ b/server/src/pflash/pflash_compress.cpp @@ -182,8 +182,13 @@ void write_compression_trace( namespace { thread_local std::vector g_last_kept_spans; thread_local PFlashScoringStats g_last_scoring_stats; +thread_local std::vector g_last_candidate_lifts; } // namespace +const std::vector & pflash_last_candidate_lifts() { + return g_last_candidate_lifts; +} + const PFlashScoringStats & pflash_last_scoring_stats() { return g_last_scoring_stats; } @@ -199,6 +204,7 @@ const std::vector & pflash_last_kept_spans() { void pflash_clear_kept_spans() { g_last_kept_spans.clear(); g_last_scoring_stats = {}; + g_last_candidate_lifts.clear(); } std::vector select_pflash_chunks( @@ -304,6 +310,21 @@ std::vector select_pflash_chunks( std::vector output; output.reserve((size_t) selected.retained_tokens); g_last_kept_spans.clear(); + g_last_candidate_lifts.clear(); + if (direct_mass) { + // Head mass sums to one over the keys, so uniform attention gives + // each token 1/input of it. + for (const auto & candidate : candidates) { + double mass = 0.0; + for (int token = candidate.begin; token < candidate.end; ++token) { + mass += token_scores[(size_t) token]; + } + const int length = std::max(1, candidate.end - candidate.begin); + g_last_candidate_lifts.push_back( + {{candidate.begin, candidate.end}, + mass / (double) length * (double) input_tokens}); + } + } for (const auto & candidate : candidates) { if (!selected_mask[candidate.ordinal]) continue; output.insert(output.end(), diff --git a/server/src/pflash/pflash_compress.h b/server/src/pflash/pflash_compress.h index e0911a9c3..7bd37e5cf 100644 --- a/server/src/pflash/pflash_compress.h +++ b/server/src/pflash/pflash_compress.h @@ -139,6 +139,16 @@ struct PFlashScoringStats { const PFlashScoringStats & pflash_last_scoring_stats(); void pflash_set_scoring_stats(const PFlashScoringStats & stats); +// Every candidate of the last strict selection on this thread with its +// attention lift: mean per-token mass relative to uniform attention over +// the input (1 = average, 20 = twenty times average). Cleared with the +// kept spans. +struct PFlashCandidateLift { + PFlashTokenSpan span; + double lift = 0.0; +}; +const std::vector & pflash_last_candidate_lifts(); + std::vector select_pflash_chunks( const std::vector & ids, const std::vector & token_scores, diff --git a/server/src/qwen35/qwen35_backend.cpp b/server/src/qwen35/qwen35_backend.cpp index a7e061593..477fb2a75 100644 --- a/server/src/qwen35/qwen35_backend.cpp +++ b/server/src/qwen35/qwen35_backend.cpp @@ -1211,6 +1211,9 @@ std::vector Qwen35Backend::compress_batch( result.scorer_resume = scoring.resume; result.scorer_new_tokens = scoring.new_tokens; result.scorer_forward_s = scoring.forward_s; + for (const auto & candidate : pflash_last_candidate_lifts()) { + result.candidate_lifts.push_back({candidate.span, candidate.lift}); + } } if (result.ok) { std::fprintf(stderr, "[compress] %zu -> %zu tokens\n", diff --git a/server/src/qwen35/qwen35_layer_split_adapter.cpp b/server/src/qwen35/qwen35_layer_split_adapter.cpp index 312f4cfa7..79e8f642f 100644 --- a/server/src/qwen35/qwen35_layer_split_adapter.cpp +++ b/server/src/qwen35/qwen35_layer_split_adapter.cpp @@ -1402,6 +1402,9 @@ Qwen35LayerSplitAdapter::compress(const ModelBackend::CompressRequest & req) { result.scorer_resume = scoring.resume; result.scorer_new_tokens = scoring.new_tokens; result.scorer_forward_s = scoring.forward_s; + for (const auto & candidate : pflash_last_candidate_lifts()) { + result.candidate_lifts.push_back({candidate.span, candidate.lift}); + } } if (result.ok) { std::fprintf(stderr, "[target-split][compress] %zu -> %zu tokens\n", diff --git a/server/src/server/http_server.cpp b/server/src/server/http_server.cpp index 2fe9ea82f..aa597d073 100644 --- a/server/src/server/http_server.cpp +++ b/server/src/server/http_server.cpp @@ -627,6 +627,41 @@ int pflash_chat_skeleton_tokens() noexcept { return (int) (std::min)(value, 1L << 20); } +bool pflash_paragraph_join() noexcept { + const char * raw = std::getenv("PFLASH_SELECT_PARAGRAPH_JOIN"); + return raw && std::string(raw) == "1"; +} + +std::string pflash_join_kept_spans( + const Tokenizer & tokenizer, + const std::vector & ids, + const std::vector & spans) { + std::string out; + int previous_end = -1; + for (const auto & span : spans) { + if (span.begin < 0 || span.end > (int) ids.size() || span.end <= span.begin) { + continue; + } + std::string piece = tokenizer.decode(std::vector( + ids.begin() + span.begin, ids.begin() + span.end)); + if (previous_end >= 0 && span.begin > previous_end && !out.empty() && + !piece.empty()) { + const bool left_break = out.back() == '\n'; + const bool right_break = piece.front() == '\n'; + if (!left_break && !right_break) { + out += "\n\n"; + } else if (left_break != right_break && + !(out.size() >= 2 && out[out.size() - 2] == '\n') && + !(piece.size() >= 2 && piece[1] == '\n')) { + out += "\n"; + } + } + out += piece; + previous_end = span.end; + } + return out; +} + bool pflash_chat_recall() noexcept { const char * raw = std::getenv("PFLASH_CHAT_RECALL"); return !(raw && std::string(raw) == "0"); @@ -641,6 +676,50 @@ int pflash_chat_compress_new_tokens() noexcept { return (int) (std::min)(value, 1L << 30); } +double pflash_chat_recall_min_lift() noexcept { + const char * raw = std::getenv("PFLASH_CHAT_RECALL_MIN_LIFT"); + if (!raw || !*raw) return 8.0; + char * end = nullptr; + const double value = std::strtod(raw, &end); + if (end == raw || *end != '\0' || !std::isfinite(value) || value < 0.0) return 8.0; + return value; +} + +int pflash_chat_recall_tokens() noexcept { + const char * raw = std::getenv("PFLASH_CHAT_RECALL_TOKENS"); + if (!raw || !*raw) return 2048; + char * end = nullptr; + const long value = std::strtol(raw, &end, 10); + if (end == raw || *end != '\0' || value < 0) return 2048; + return (int) (std::min)(value, 1L << 24); +} + +std::vector pflash_recall_by_lift( + const std::vector> & lifts, + const std::vector & in_view, + double min_lift, + int max_tokens) { + std::vector> picks; + for (const auto & [span, lift] : lifts) { + if (!(lift >= min_lift)) continue; + for (const auto & part : pflash_subtract_token_spans({span}, in_view)) { + picks.push_back({lift, part}); + } + } + std::stable_sort(picks.begin(), picks.end(), [] (const auto & a, const auto & b) { + return a.first > b.first; + }); + std::vector chosen; + int used = 0; + for (const auto & [lift, part] : picks) { + const int length = part.end - part.begin; + if (used + length > max_tokens) continue; + chosen.push_back(part); + used += length; + } + return canonicalize_pflash_token_spans(std::move(chosen)); +} + int pflash_chat_history_queries() noexcept { const char * raw = std::getenv("PFLASH_CHAT_HISTORY_QUERIES"); if (!raw || !*raw) return 3; @@ -4385,7 +4464,7 @@ std::string HttpServer::apply_pflash_compression( json view_stats; if (serve_pflash_chat_view( req, compress_request.input_ids, chat_turn, nullptr, nullptr, - served, prepared.snapshot_cut, view_stats)) { + nullptr, served, prepared.snapshot_cut, view_stats)) { prepared.tokens = std::move(served); prepared.compressed = true; prepared.pflash_stats = { @@ -4402,6 +4481,7 @@ std::string HttpServer::apply_pflash_compression( ModelBackend::CompressResult result; std::vector final_tokens; const auto compress_started = std::chrono::steady_clock::now(); + int join_overhead = 0; for (int attempt = 0; ; ++attempt) { result = {}; if (config_.pflash_remote_drafter) { @@ -4433,6 +4513,21 @@ std::string HttpServer::apply_pflash_compression( std::string compressed_text = drafter_tokenizer_->decode(result.compressed_ids); + join_overhead = 0; + // PFLASH_SELECT_PARAGRAPH_JOIN=1: kept pieces that were not adjacent + // in the prompt are joined by a paragraph break when neither side + // already has one, so a cut does not glue two passages into one + // run-on line ("...other bands.Document 1:"). + if (http_detail::pflash_paragraph_join() && !result.kept_spans.empty()) { + const int plain = (int) tokenizer_.encode(compressed_text).size(); + compressed_text = http_detail::pflash_join_kept_spans( + *drafter_tokenizer_, compress_request.input_ids, + result.kept_spans); + // The breaks are layout, not retained context: the ceiling + // bounds what the selection kept. + join_overhead = (std::max)( + 0, (int) tokenizer_.encode(compressed_text).size() - plain); + } // Compression is allowed to be lossy, but the active user query must // survive. Re-append short queries when fewer than 80% of their tokens do. @@ -4467,10 +4562,11 @@ std::string HttpServer::apply_pflash_compression( final_tokens = tokenizer_.encode(compressed_text); if (!experiment.selection_active || - (int) final_tokens.size() <= target_ceiling) { + (int) final_tokens.size() - join_overhead <= target_ceiling) { break; } - const int overflow = (int) final_tokens.size() - target_ceiling; + const int overflow = + (int) final_tokens.size() - join_overhead - target_ceiling; const int tightened = target_ceiling - overflow - 1; if (attempt >= 2 || tightened <= 0) { break; @@ -4493,7 +4589,7 @@ std::string HttpServer::apply_pflash_compression( "[pflash-select] final target tokens=%zu ceiling=%d\n", final_tokens.size(), target_ceiling); std::fflush(stderr); - if ((int) final_tokens.size() > target_ceiling) { + if ((int) final_tokens.size() - join_overhead > target_ceiling) { return "PFlash strict selection final prompt exceeds target-token ceiling " "(" + std::to_string(final_tokens.size()) + " > " + std::to_string(target_ceiling) + ")"; @@ -4517,8 +4613,8 @@ std::string HttpServer::apply_pflash_compression( std::vector served; if (serve_pflash_chat_view( req, compress_request.input_ids, chat_turn, &final_tokens, - &result.kept_spans, served, prepared.snapshot_cut, - prepared.pflash_stats["view"])) { + &result.kept_spans, &result.candidate_lifts, served, + prepared.snapshot_cut, prepared.pflash_stats["view"])) { final_tokens = std::move(served); } } @@ -4554,6 +4650,7 @@ bool HttpServer::serve_pflash_chat_view( const http_detail::PflashChatTurnSpan & turn, const std::vector * fresh, const std::vector * kept_spans, + const std::vector> * lifts, std::vector & served, int & snapshot_cut, json & stats) { @@ -4642,13 +4739,20 @@ bool HttpServer::serve_pflash_chat_view( // Recall: what the fresh selection keeps for the new query that the view // does not hold. Only a new user turn brings a new query; an agent step // (assistant call plus tool output) appends without recalling. + // With the head's per-candidate lifts, recall only what the new + // question clearly attends to: a content-free question ("which + // documents support that?") recalls next to nothing instead of filling + // the budget with low-relevance segments beside the question. std::vector recalled; if (compressed && recall) { auto in_view = view.spans; in_view.push_back({view.drafter_gen_begin, input}); - recalled = http_detail::pflash_subtract_token_spans( - *kept_spans, - http_detail::canonicalize_pflash_token_spans(std::move(in_view))); + in_view = http_detail::canonicalize_pflash_token_spans(std::move(in_view)); + recalled = lifts && !lifts->empty() + ? http_detail::pflash_recall_by_lift( + *lifts, in_view, http_detail::pflash_chat_recall_min_lift(), + http_detail::pflash_chat_recall_tokens()) + : http_detail::pflash_subtract_token_spans(*kept_spans, in_view); } std::string recall_block; int recalled_tokens = 0; @@ -4709,14 +4813,26 @@ bool HttpServer::serve_pflash_chat_view( } const int split = new_question && !recall_block.empty() ? turn.content_begin : input; - std::string delta; + std::vector before_split; + std::vector after_split; for (const auto & span : delta_spans) { - delta += decode_range(span.begin, (std::min)(span.end, split)); - } - delta += recall_block; - for (const auto & span : delta_spans) { - delta += decode_range((std::max)(span.begin, split), span.end); + if (span.begin < split) { + before_split.push_back({span.begin, (std::min)(span.end, split)}); + } + if (span.end > split) { + after_split.push_back({(std::max)(span.begin, split), span.end}); + } } + const auto join = [&] (const std::vector & spans) { + if (compress_new && http_detail::pflash_paragraph_join()) { + return http_detail::pflash_join_kept_spans( + *drafter_tokenizer_, drafter_ids, spans); + } + std::string text; + for (const auto & span : spans) text += decode_range(span.begin, span.end); + return text; + }; + const std::string delta = join(before_split) + recall_block + join(after_split); served.assign(view.view_tokens.begin(), view.view_tokens.begin() + view.view_gen_begin); const auto delta_tokens = tokenizer_.encode(delta); diff --git a/server/src/server/http_server.h b/server/src/server/http_server.h index 6682035a0..4dbe984ad 100644 --- a/server/src/server/http_server.h +++ b/server/src/server/http_server.h @@ -320,8 +320,29 @@ int pflash_chat_history_queries() noexcept; // scoring); PFLASH_CHAT_COMPRESS_NEW_TOKENS (default 16384) is the size of // new material in one turn from which it is compressed instead of appended. bool pflash_chat_recall() noexcept; + +// PFLASH_SELECT_PARAGRAPH_JOIN=1: the compressed text is rebuilt from the +// kept spans with a paragraph break between pieces that were not adjacent +// in the prompt, unless one side already ends or starts a paragraph. +bool pflash_paragraph_join() noexcept; +std::string pflash_join_kept_spans( + const Tokenizer & tokenizer, + const std::vector & ids, + const std::vector & spans); int pflash_chat_compress_new_tokens() noexcept; +// Recall takes only segments the new question clearly attends to: attention +// lift (mass per token relative to uniform) of at least +// PFLASH_CHAT_RECALL_MIN_LIFT (default 8), strongest first, up to +// PFLASH_CHAT_RECALL_TOKENS (default 2048) drafter tokens. +double pflash_chat_recall_min_lift() noexcept; +int pflash_chat_recall_tokens() noexcept; +std::vector pflash_recall_by_lift( + const std::vector> & lifts, + const std::vector & in_view, + double min_lift, + int max_tokens); + // The parts of ``spans`` that ``minus`` does not cover. Both canonical. std::vector pflash_subtract_token_spans( const std::vector & spans, @@ -662,6 +683,7 @@ class HttpServer { const http_detail::PflashChatTurnSpan & turn, const std::vector * fresh, const std::vector * kept_spans, + const std::vector> * lifts, std::vector & served, int & snapshot_cut, nlohmann::json & stats); diff --git a/server/test/test_server_unit.cpp b/server/test/test_server_unit.cpp index 32d54c372..7873be39c 100644 --- a/server/test/test_server_unit.cpp +++ b/server/test/test_server_unit.cpp @@ -1052,6 +1052,54 @@ TEST_CASE(ServerUnitFixture, test_timings_json_carries_pflash_details) { TEST_ASSERT(out["pflash"]["view"]["mode"] == "continue"); } +TEST_CASE(ServerUnitFixture, test_pflash_join_kept_spans_breaks_between_pieces) { + const std::string text = "one fact.\ngap text.Two starts.\n\nThree."; + const std::string path = write_pflash_bpe_tokenizer_fixture({}, text); + Tokenizer tok; + TEST_ASSERT(tok.load_from_gguf(path.c_str())); + const auto ids = tok.encode(text); + const auto token_at = [&] (const std::string & needle) { + const auto span = http_detail::pflash_decoded_text_span( + tok, ids, 0, (int) ids.size(), needle); + return span; + }; + const auto one = token_at("one fact."); + const auto two = token_at("Two starts."); + const auto three = token_at("\n\nThree."); + // Non-adjacent pieces without a break get a paragraph break... + TEST_ASSERT_MSG(http_detail::pflash_join_kept_spans(tok, ids, {one, two}) == + "one fact.\n\nTwo starts.", + http_detail::pflash_join_kept_spans(tok, ids, {one, two})); + // ...a piece that already opens a paragraph is left alone... + TEST_ASSERT(http_detail::pflash_join_kept_spans(tok, ids, {one, three}) == + "one fact.\n\nThree."); + // ...and adjacent pieces are joined as they were. + TEST_ASSERT(http_detail::pflash_join_kept_spans(tok, ids, {{0, 3}, {3, 5}}) == + tok.decode({ids.begin(), ids.begin() + 5})); + unlink(path.c_str()); +} + +TEST_CASE(ServerUnitFixture, test_pflash_recall_by_lift_takes_clear_attention_only) { + const std::vector> lifts{ + {{0, 10}, 40.0}, // clearly attended, not in view + {{10, 20}, 2.0}, // background + {{20, 30}, 90.0}, // clearly attended, already in view + {{30, 60}, 12.0}, // attended, half in view + {{60, 70}, 25.0}, + }; + const std::vector in_view{{20, 45}}; + auto recalled = http_detail::pflash_recall_by_lift(lifts, in_view, 8.0, 1000); + TEST_ASSERT(recalled.size() == 2); // [0,10) and [45,70) merged + TEST_ASSERT(recalled[0].begin == 0 && recalled[0].end == 10); + TEST_ASSERT(recalled[1].begin == 45 && recalled[1].end == 70); + // A tight cap keeps the strongest: 40, then 25. + recalled = http_detail::pflash_recall_by_lift(lifts, in_view, 8.0, 20); + TEST_ASSERT(recalled.size() == 2); + TEST_ASSERT(recalled[0].begin == 0 && recalled[1].begin == 60); + // Nothing clears a high bar. + TEST_ASSERT(http_detail::pflash_recall_by_lift(lifts, in_view, 100.0, 1000).empty()); +} + TEST_CASE(ServerUnitFixture, test_pflash_subtract_token_spans) { const std::vector spans{{0, 10}, {20, 30}, {40, 50}}; const std::vector minus{{5, 22}, {25, 26}, {40, 50}}; From 554953cda42731c8196e4d780e47b6fbb2628a58 Mon Sep 17 00:00:00 2001 From: Graffioh Date: Wed, 23 Sep 2026 06:49:04 +0000 Subject: [PATCH 21/26] feat(pflash): join kept pieces with paragraph breaks by default On the multi-turn development chats (7 non-held-out 32K sessions, 42 turns) the paragraph join recovered every citation turn (6/6, like full prefill, against 5/6 without it) and was never worse elsewhere; PFLASH_SELECT_PARAGRAPH_JOIN=0 turns it off. Single-turn compressed prompts change accordingly: the article's single-turn numbers were measured without it. Co-Authored-By: Claude Opus 5.5 (1M context) --- server/README.md | 6 ++++-- server/src/server/http_server.cpp | 4 ++-- server/src/server/http_server.h | 7 ++++--- 3 files changed, 10 insertions(+), 7 deletions(-) diff --git a/server/README.md b/server/README.md index b44f8c7a2..bbcf70cea 100644 --- a/server/README.md +++ b/server/README.md @@ -428,8 +428,10 @@ to uniform attention) of at least `PFLASH_CHAT_RECALL_MIN_LIFT` (default 8), strongest first, up to `PFLASH_CHAT_RECALL_TOKENS` (default 2048 drafter tokens), so a content-free follow-up ("which documents support that?") recalls next to nothing instead of filling the budget with noise beside the -question. `PFLASH_SELECT_PARAGRAPH_JOIN=1` joins kept pieces that were not -adjacent with a paragraph break (off by default). `PFLASH_VIEW_TRACE_PATH` +question. Kept pieces that were not adjacent in the prompt are joined by a +paragraph break when neither side has one (`PFLASH_SELECT_PARAGRAPH_JOIN=0` +turns it off; the breaks do not count against the token ceiling). +`PFLASH_VIEW_TRACE_PATH` appends each compressed request's served prompt as JSONL, for evidence checks in evaluations. diff --git a/server/src/server/http_server.cpp b/server/src/server/http_server.cpp index d293fcdec..bdd2fb58c 100644 --- a/server/src/server/http_server.cpp +++ b/server/src/server/http_server.cpp @@ -630,7 +630,7 @@ int pflash_chat_skeleton_tokens() noexcept { bool pflash_paragraph_join() noexcept { const char * raw = std::getenv("PFLASH_SELECT_PARAGRAPH_JOIN"); - return raw && std::string(raw) == "1"; + return !(raw && std::string(raw) == "0"); } std::string pflash_join_kept_spans( @@ -4554,7 +4554,7 @@ std::string HttpServer::apply_pflash_compression( std::string compressed_text = drafter_tokenizer_->decode(result.compressed_ids); join_overhead = 0; - // PFLASH_SELECT_PARAGRAPH_JOIN=1: kept pieces that were not adjacent + // Kept pieces that were not adjacent // in the prompt are joined by a paragraph break when neither side // already has one, so a cut does not glue two passages into one // run-on line ("...other bands.Document 1:"). diff --git a/server/src/server/http_server.h b/server/src/server/http_server.h index c7d77b140..0ddb6fce8 100644 --- a/server/src/server/http_server.h +++ b/server/src/server/http_server.h @@ -322,9 +322,10 @@ int pflash_chat_history_queries() noexcept; // new material in one turn from which it is compressed instead of appended. bool pflash_chat_recall() noexcept; -// PFLASH_SELECT_PARAGRAPH_JOIN=1: the compressed text is rebuilt from the -// kept spans with a paragraph break between pieces that were not adjacent -// in the prompt, unless one side already ends or starts a paragraph. +// The compressed text is rebuilt from the kept spans with a paragraph break +// between pieces that were not adjacent in the prompt, unless one side +// already ends or starts a paragraph (PFLASH_SELECT_PARAGRAPH_JOIN=0 turns +// it off). bool pflash_paragraph_join() noexcept; std::string pflash_join_kept_spans( const Tokenizer & tokenizer, From b1c8fe5d44c693f209adcf9b24aa12d4287c3a07 Mon Sep 17 00:00:00 2001 From: Graffioh Date: Wed, 23 Sep 2026 10:45:18 +0000 Subject: [PATCH 22/26] feat(pflash): score chats from the prompt's last token The chat scorer query was the tail of the latest user turn, so a question before or inside a pasted document was missed. The query is now the prompt's last token -- where the model starts answering, having read the whole request -- and nothing is parsed out of the user's text: the latest turn follows the skeleton rule like the others, the generation prompt is pinned, and history queries are the last token of the header of the reply that followed each earlier user turn. Agent tool turns after the user's become ordinary context. Marker-less prompts keep the content tail; pflash_query keeps its explicit span. Co-Authored-By: Claude Opus 5.5 (1M context) --- server/README.md | 37 ++++++------ server/src/server/http_server.cpp | 58 +++++++++++++++---- server/test/test_server_unit.cpp | 95 ++++++++++++++++--------------- 3 files changed, 117 insertions(+), 73 deletions(-) diff --git a/server/README.md b/server/README.md index bbcf70cea..66d6d297a 100644 --- a/server/README.md +++ b/server/README.md @@ -388,14 +388,17 @@ restores the previous all-layer running-max scorer. The Qwen3.5 attention runs dense (`ggml_flash_attn_ext`); the block-sparse FlashPrefill kernels still dispatch head dimension 128 only. -The scorer query is the tail (`PFLASH_SELECT_QUERY_TOKENS`, default 8) of the -latest user turn, located by the model's own chat markers in the rendered -prompt. Tool output wrapped in a user turn and the generation prompt, with -its think prefix, never count as that turn. Strict selection keeps the query -and its turn's role header, and it runs on every turn of a multi-turn chat. -In an agent loop the assistant and tool turns after the user's turn are -scored against the query like the context before it; only the generation -prompt is kept with them. +The scorer query of a chat is the prompt's last token: the end of the +generation prompt, where the model starts answering, having read the whole +request. Nothing is parsed out of the user's text, so the question can sit +anywhere in the message -- before a pasted document, in the middle of it, or +among the user's own sentences -- and the user's own words score far above +the material they paste. Strict selection keeps the generation prompt and +the latest user turn's role header, and it runs on every turn of a +multi-turn chat. In an agent loop the assistant and tool turns after the +user's turn are scored like the rest of the conversation. A prompt without +chat markers scores the tail (`PFLASH_SELECT_QUERY_TOKENS`, default 8) of +its content. The keep ratio applies to the droppable tokens only: what strict selection keeps anyway (system and developer messages, tool definitions, the query and @@ -435,14 +438,16 @@ turns it off; the breaks do not count against the token ceiling). appends each compressed request's served prompt as JSONL, for evidence checks in evaluations. -Every other turn of a multi-turn chat keeps its role header, and user -turns and assistant answers up to `PFLASH_CHAT_SKELETON_TOKENS` (default 256 -drafter tokens; 0 keeps headers only) stay whole: the conversation's -skeleton, as opposed to the material it quotes. Like instructions, the -skeleton is scored as context when it alone would not fit. The last -`PFLASH_CHAT_HISTORY_QUERIES` (default 3) earlier user questions score the -context alongside the current one, their masses mixed in at weights 1/2, -1/4, 1/8, so what the conversation keeps coming back to stays selected. +Every turn of a multi-turn chat keeps its role header, and user turns (the +latest included) and assistant answers up to `PFLASH_CHAT_SKELETON_TOKENS` +(default 256 drafter tokens; 0 keeps headers only) stay whole: the +conversation's skeleton, as opposed to the material it quotes. Like +instructions, the skeleton is scored as context when it alone would not fit. +The last `PFLASH_CHAT_HISTORY_QUERIES` (default 3) earlier user turns score +the context alongside the current query, each through the last token of the +header of the reply that followed it (that turn's own prompt end), their +masses mixed in at weights 1/2, 1/4, 1/8, so what the conversation keeps +coming back to stays selected. The drafter keeps a scoring session per conversation (`PFLASH_DRAFTER_SESSIONS`, default 2, least recently used evicted; 0 scores diff --git a/server/src/server/http_server.cpp b/server/src/server/http_server.cpp index bdd2fb58c..c9f8645c1 100644 --- a/server/src/server/http_server.cpp +++ b/server/src/server/http_server.cpp @@ -4054,18 +4054,30 @@ std::string HttpServer::apply_pflash_compression( return "PFlash strict selection content boundary mapping failed"; } // Chat default: without an explicit pflash_query the scorer - // query is the tail of the boundary content, and it takes the - // explicit query's path from here on. + // query is the prompt's last token -- where the model starts + // answering, having read the whole request wherever the question + // sits in it. Nothing is parsed out of the user's text: the + // latest turn is scored like the rest of the conversation. A + // prompt without chat markers falls back to its content's tail. if (tail_parser && req.pflash_query.empty()) { - const auto window = http_detail::pflash_tail_query_window( - drafter_ids, experiment.query_tokens, - query_content_end, query_content_begin); - if (window.valid()) { - query_span = {window.end - window.tokens, window.end}; - query_span_rule = chat_turn.valid() ? "chat_user_tail" - : (raw_text_input ? "content_tail" : "prompt_tail"); + const int prompt_end = (int) drafter_ids.size(); + if (chat_turn.valid() && + chat_turn.generation_begin > 0 && + chat_turn.generation_begin < prompt_end) { + query_span = {prompt_end - 1, prompt_end}; + query_span_rule = "prompt_end"; + } else { + const auto window = http_detail::pflash_tail_query_window( + drafter_ids, experiment.query_tokens, + query_content_end, query_content_begin); + if (window.valid()) { + query_span = {window.end - window.tokens, window.end}; + query_span_rule = + raw_text_input ? "content_tail" : "prompt_tail"; + } } } + const bool prompt_end_query = query_span_rule == "prompt_end"; if (experiment.selection_active) { const auto instruction_plan = @@ -4156,6 +4168,12 @@ std::string HttpServer::apply_pflash_compression( if (query_span.begin >= 0) { required_instruction_spans.push_back(query_span); } + if (prompt_end_query) { + // The whole generation prompt stays verbatim; the query + // is its last token. + required_instruction_spans.push_back( + {chat_turn.generation_begin, (int) drafter_ids.size()}); + } if (query_role_header.begin >= 0) { required_instruction_spans.push_back(query_role_header); } @@ -4187,7 +4205,13 @@ std::string HttpServer::apply_pflash_compression( http_detail::pflash_chat_skeleton_tokens(); for (size_t index = 0; index < chat_turn.turns.size(); ++index) { - if ((int) index == chat_turn.query_turn) continue; + // With a prompt-end query the latest user turn is + // no longer pinned as the query: it follows the same + // rule as every other turn. + if ((int) index == chat_turn.query_turn && + !prompt_end_query) { + continue; + } const auto & turn = chat_turn.turns[index]; if (turn.role == "system") continue; if (turn.content_begin > turn.role_begin) { @@ -4212,6 +4236,20 @@ std::string HttpServer::apply_pflash_compression( --index) { const auto & turn = chat_turn.turns[(size_t) index]; if (turn.role != "user") continue; + if (prompt_end_query) { + // The earlier question's counterpart of the + // prompt's last token: the last token of the + // header of the reply that followed it. + const size_t reply = (size_t) index + 1; + if (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; + history_query_spans.push_back({end - 1, end}); + } + continue; + } const auto window = http_detail::pflash_tail_query_window( drafter_ids, experiment.query_tokens, turn.content_end, turn.content_begin); diff --git a/server/test/test_server_unit.cpp b/server/test/test_server_unit.cpp index 399eb7d28..928223d90 100644 --- a/server/test/test_server_unit.cpp +++ b/server/test/test_server_unit.cpp @@ -6975,7 +6975,7 @@ TEST_CASE(ServerUnitFixture, test_pflash_default_raw_text_maps_user_query) { } TEST_CASE(ServerUnitFixture, - test_pflash_strict_chat_tail_query_uses_last_user_content) { + test_pflash_strict_chat_query_is_the_prompt_end) { luce_test::ScopedEnvVar mode{"PFLASH_SELECT_MODE", "top_p"}; // The server's own Qwen rendering, generation prompt and its think @@ -7024,31 +7024,31 @@ TEST_CASE(ServerUnitFixture, if (ids[i] == im_end) last_im_end = i; } TEST_ASSERT(last_im_end > 0); - // The scorer window ends where the user content does — the generation - // prompt ("<|im_end|>\n<|im_start|>assistant\n\n") is never - // scored. - TEST_ASSERT(backend.last_request.score_query_end == last_im_end); - const int query_begin = backend.last_request.score_query_end - - backend.last_request.score_query_tokens; - TEST_ASSERT(query_begin >= 0); - TEST_ASSERT(tokenizer.decode({ids.begin() + query_begin, - ids.begin() + last_im_end}) - == "What is the answer?"); - // The query span and its turn's role header are pinned mandatory. + // The scorer query is the prompt's last token, where the model starts + // answering; nothing is parsed out of the user's text. + TEST_ASSERT(backend.last_request.score_query_end == (int) ids.size()); + TEST_ASSERT(backend.last_request.score_query_tokens == 1); + // The generation prompt, the turn's role header and -- a short turn -- + // the whole question stay. bool header_pinned = false; - bool query_pinned = false; + bool question_pinned = false; + bool generation_pinned = false; for (const auto & span : backend.last_request.required_instruction_spans) { const std::string text = tokenizer.decode( {ids.begin() + span.begin, ids.begin() + span.end}); if (text.find("<|im_start|>user\n") != std::string::npos) { header_pinned = true; } - if (span.begin <= query_begin && span.end >= last_im_end) { - query_pinned = true; + if (text.find("What is the answer?") != std::string::npos) { + question_pinned = true; + } + if (span.begin <= last_im_end + 2 && span.end == (int) ids.size()) { + generation_pinned = true; } } TEST_ASSERT(header_pinned); - TEST_ASSERT(query_pinned); + TEST_ASSERT(question_pinned); + TEST_ASSERT(generation_pinned); unlink(path.c_str()); } @@ -7100,12 +7100,11 @@ TEST_CASE(ServerUnitFixture, TEST_ASSERT(backend.compress_calls == 1); const auto & request = backend.last_request; const auto & ids = request.input_ids; - TEST_ASSERT(request.query_suffix_candidates); - const int query_end = request.score_query_end; - const int query_begin = query_end - request.score_query_tokens; - TEST_ASSERT(tokenizer.decode({ids.begin() + query_begin, - ids.begin() + query_end}) - == "What is the answer?"); + // The query is the prompt's end, so nothing follows it: the assistant + // and tool turns are ordinary context. + TEST_ASSERT(!request.query_suffix_candidates); + TEST_ASSERT(request.score_query_end == (int) ids.size()); + TEST_ASSERT(request.score_query_tokens == 1); // The generation prompt is pinned, and the short assistant turn stays as // part of the conversation's skeleton; the tool output between the query // and the generation prompt is scored, not pinned. @@ -7370,6 +7369,9 @@ TEST_CASE(ServerUnitFixture, test_pflash_chat_view_appends_turns_and_recalls_missing_segments) { luce_test::ScopedEnvVar mode{"PFLASH_SELECT_MODE", "budget_only"}; luce_test::ScopedEnvVar view_env{"PFLASH_CHAT_VIEW", nullptr}; + // Only turns of a few tokens stay whole, so the document turn is + // material the selection picks from. + luce_test::ScopedEnvVar skeleton{"PFLASH_CHAT_SKELETON_TOKENS", "4"}; std::string system; for (int i = 0; i < 20; ++i) system += "You are helpful. "; @@ -7543,11 +7545,14 @@ TEST_CASE(ServerUnitFixture, } TEST_ASSERT(answer_kept); TEST_ASSERT(!material_kept); - // The earlier question scores alongside the current one. + // The earlier question scores alongside the current one, through the + // last token of the header of the reply that followed it. TEST_ASSERT(request.history_query_spans.size() == 1); - TEST_ASSERT_MSG(text_of(request.history_query_spans[0]).find("first answer?") != - std::string::npos, - text_of(request.history_query_spans[0])); + const auto history = request.history_query_spans[0]; + TEST_ASSERT(history.end - history.begin == 1); + TEST_ASSERT_MSG(tokenizer.decode({ids.begin() + history.end, + ids.begin() + history.end + 2}) == "Sure.", + text_of(history)); unlink(path.c_str()); } @@ -7587,12 +7592,17 @@ static PflashTwoTurnRun pflash_two_turn_run(const std::string & follow_up) { TEST_ASSERT(tokenizer.load_from_gguf(path.c_str())); auto backend_owner = std::make_unique(); MockPflashSpanBackend & backend = *backend_owner; + // The fact and the questions: a prompt-end query ranks the user's own + // sentences first, so the scorer keeps them. backend.pick = [&] (const ModelBackend::CompressRequest & request) { - const auto span = http_detail::pflash_decoded_text_span( - tokenizer, request.input_ids, 0, (int) request.input_ids.size(), - "alpha facts"); - return span.begin < 0 ? std::vector{} - : std::vector{span}; + std::vector spans; + for (const char * text : {"alpha facts", "Question one?", "Question two?"}) { + const auto span = http_detail::pflash_decoded_text_span( + tokenizer, request.input_ids, 0, (int) request.input_ids.size(), + text); + if (span.begin >= 0) spans.push_back(span); + } + return spans; }; LuceEngine engine(std::move(backend_owner)); ServerConfig config; @@ -7644,8 +7654,8 @@ TEST_CASE(ServerUnitFixture, test_pflash_chat_view_compresses_large_follow_ups_only) { luce_test::ScopedEnvVar mode{"PFLASH_SELECT_MODE", "budget_only"}; luce_test::ScopedEnvVar threshold{"PFLASH_CHAT_COMPRESS_NEW_TOKENS", "40"}; - // Wide enough that the pinned query covers the whole question. - luce_test::ScopedEnvVar query{"PFLASH_SELECT_QUERY_TOKENS", "16"}; + // The pasted turn is material, not a short turn kept whole. + luce_test::ScopedEnvVar skeleton{"PFLASH_CHAT_SKELETON_TOKENS", "8"}; std::string pasted; for (int i = 0; i < 12; ++i) pasted += " pasted notes filler."; const auto run = pflash_two_turn_run(pasted); @@ -7658,8 +7668,8 @@ TEST_CASE(ServerUnitFixture, const size_t prefix = first.size() - 8; TEST_ASSERT(std::equal(first.begin(), first.begin() + (long) (prefix - 8), second.begin())); - // ...and the pasted material is compressed: only the pinned question and - // what the selection keeps survive, not the whole paste. + // ...and the pasted material is compressed: only what the selection + // keeps (the question) survives, not the whole paste. TEST_ASSERT_MSG(run.text2.find("Question two?") != std::string::npos, run.text2); TEST_ASSERT(run.text2.find(pasted) == std::string::npos); } @@ -7770,20 +7780,11 @@ TEST_CASE(ServerUnitFixture, } // Whole-prompt PFlash ran on the multi-turn prompt and scored against - // the last user turn's content tail. + // the prompt's last token. TEST_ASSERT(backend.compress_calls == 1); const auto & ids = backend.last_request.input_ids; - const int im_end = tokenizer.token_to_id("<|im_end|>"); - int last_im_end = -1; - for (int i = 0; i < (int) ids.size(); ++i) { - if (ids[i] == im_end) last_im_end = i; - } - TEST_ASSERT(backend.last_request.score_query_end == last_im_end); - const int query_begin = backend.last_request.score_query_end - - backend.last_request.score_query_tokens; - TEST_ASSERT(tokenizer.decode({ids.begin() + query_begin, - ids.begin() + last_im_end}) - == "second question"); + TEST_ASSERT(backend.last_request.score_query_end == (int) ids.size()); + TEST_ASSERT(backend.last_request.score_query_tokens == 1); unlink(path.c_str()); } From 63ccadd7486d0a3a95e39f365fe2e060416dfc96 Mon Sep 17 00:00:00 2001 From: Graffioh Date: Wed, 23 Sep 2026 10:51:46 +0000 Subject: [PATCH 23/26] feat(pflash): recall by one lift gate, rebuild past a third of fresh The 2048-token recall cap never bound once a lift gate was on (at most ~960 tokens on the dev chats). On the development chats (Jev-judged), a gate of 2 against the previous 8: 212 vs 197 of 252 turns on the query-layout chats (question first, in the middle, among the user's sentences, two questions, emails; full prefill 231), 37 vs 37 of 42 on the original chats, at ~190 recalled tokens per turn and 1.01 s vs 0.81 s median later-turn TTFT (full prefill 1.42 s). Recall now takes every out-of-view segment at lift >= 2 (PFLASH_CHAT_RECALL_MIN_LIFT); PFLASH_CHAT_RECALL_TOKENS is gone. A question that needs more than a third of a fresh selection rebuilds the view from it instead: past that, the fresh prompt costs about the same and keeps order. Co-Authored-By: Claude Opus 5.5 (1M context) --- server/README.md | 12 +++---- server/src/server/http_server.cpp | 57 ++++++++++++------------------- server/src/server/http_server.h | 9 ++--- server/test/test_server_unit.cpp | 13 +++---- 4 files changed, 38 insertions(+), 53 deletions(-) diff --git a/server/README.md b/server/README.md index 66d6d297a..37907fcdb 100644 --- a/server/README.md +++ b/server/README.md @@ -425,13 +425,13 @@ prefill appends a follow-up; from `PFLASH_CHAT_COMPRESS_NEW_TOKENS` (default 16384) tokens of new material (a pasted document, a large tool output) only what the fresh selection keeps of it is appended, and the view before it stays cached. `PFLASH_CHAT_RECALL=0` turns recall off: a small follow-up is -then served without running the drafter at all. Recall takes only segments +then served without running the drafter at all. Recall takes the segments the new question clearly attends to: attention lift (mass per token relative -to uniform attention) of at least `PFLASH_CHAT_RECALL_MIN_LIFT` (default 8), -strongest first, up to `PFLASH_CHAT_RECALL_TOKENS` (default 2048 drafter -tokens), so a content-free follow-up ("which documents support that?") -recalls next to nothing instead of filling the budget with noise beside the -question. Kept pieces that were not adjacent in the prompt are joined by a +to uniform attention) of at least `PFLASH_CHAT_RECALL_MIN_LIFT` (default 2), +so a content-free follow-up ("which documents support that?") recalls next +to nothing instead of filling the budget with noise beside the question. A +question that needs more than a third of a fresh selection starts a new view +from that selection instead. Kept pieces that were not adjacent in the prompt are joined by a paragraph break when neither side has one (`PFLASH_SELECT_PARAGRAPH_JOIN=0` turns it off; the breaks do not count against the token ceiling). `PFLASH_VIEW_TRACE_PATH` diff --git a/server/src/server/http_server.cpp b/server/src/server/http_server.cpp index c9f8645c1..26fb81678 100644 --- a/server/src/server/http_server.cpp +++ b/server/src/server/http_server.cpp @@ -679,45 +679,24 @@ int pflash_chat_compress_new_tokens() noexcept { double pflash_chat_recall_min_lift() noexcept { const char * raw = std::getenv("PFLASH_CHAT_RECALL_MIN_LIFT"); - if (!raw || !*raw) return 8.0; + if (!raw || !*raw) return 2.0; char * end = nullptr; const double value = std::strtod(raw, &end); - if (end == raw || *end != '\0' || !std::isfinite(value) || value < 0.0) return 8.0; + if (end == raw || *end != '\0' || !std::isfinite(value) || value < 0.0) return 2.0; return value; } -int pflash_chat_recall_tokens() noexcept { - const char * raw = std::getenv("PFLASH_CHAT_RECALL_TOKENS"); - if (!raw || !*raw) return 2048; - char * end = nullptr; - const long value = std::strtol(raw, &end, 10); - if (end == raw || *end != '\0' || value < 0) return 2048; - return (int) (std::min)(value, 1L << 24); -} - std::vector pflash_recall_by_lift( const std::vector> & lifts, const std::vector & in_view, - double min_lift, - int max_tokens) { - std::vector> picks; + double min_lift) { + std::vector chosen; for (const auto & [span, lift] : lifts) { if (!(lift >= min_lift)) continue; for (const auto & part : pflash_subtract_token_spans({span}, in_view)) { - picks.push_back({lift, part}); + chosen.push_back(part); } } - std::stable_sort(picks.begin(), picks.end(), [] (const auto & a, const auto & b) { - return a.first > b.first; - }); - std::vector chosen; - int used = 0; - for (const auto & [lift, part] : picks) { - const int length = part.end - part.begin; - if (used + length > max_tokens) continue; - chosen.push_back(part); - used += length; - } return canonicalize_pflash_token_spans(std::move(chosen)); } @@ -4814,13 +4793,15 @@ bool HttpServer::serve_pflash_chat_view( const bool recall = new_question && http_detail::pflash_chat_recall(); if (!compressed && (compress_new || recall)) return false; - // Recall: what the fresh selection keeps for the new query that the view - // does not hold. Only a new user turn brings a new query; an agent step - // (assistant call plus tool output) appends without recalling. - // With the head's per-candidate lifts, recall only what the new - // question clearly attends to: a content-free question ("which - // documents support that?") recalls next to nothing instead of filling - // the budget with low-relevance segments beside the question. + // Recall: what the new question clearly attends to that the view does + // not hold. Only a new user turn brings a new question; an agent step + // (assistant call plus tool output) appends without recalling. The head's + // per-candidate lifts decide it (the fresh selection minus the view when + // a scorer reports none), so a content-free question ("which documents + // support that?") recalls next to nothing. A question that needs more + // than a third of a fresh selection starts a new view from that + // selection instead: past that, prefilling the fresh prompt costs about + // the same and serves the material in order. std::vector recalled; if (compressed && recall) { auto in_view = view.spans; @@ -4828,9 +4809,15 @@ bool HttpServer::serve_pflash_chat_view( in_view = http_detail::canonicalize_pflash_token_spans(std::move(in_view)); recalled = lifts && !lifts->empty() ? http_detail::pflash_recall_by_lift( - *lifts, in_view, http_detail::pflash_chat_recall_min_lift(), - http_detail::pflash_chat_recall_tokens()) + *lifts, in_view, http_detail::pflash_chat_recall_min_lift()) : http_detail::pflash_subtract_token_spans(*kept_spans, in_view); + size_t recall_size = 0; + for (const auto & span : recalled) { + recall_size += (size_t) (span.end - span.begin); + } + if (3 * recall_size > fresh->size()) { + return serve_fresh("rebuild", view.turns + 1); + } } std::string recall_block; int recalled_tokens = 0; diff --git a/server/src/server/http_server.h b/server/src/server/http_server.h index 0ddb6fce8..b8b609d16 100644 --- a/server/src/server/http_server.h +++ b/server/src/server/http_server.h @@ -333,17 +333,14 @@ std::string pflash_join_kept_spans( const std::vector & spans); int pflash_chat_compress_new_tokens() noexcept; -// Recall takes only segments the new question clearly attends to: attention +// Recall takes the segments the new question clearly attends to: attention // lift (mass per token relative to uniform) of at least -// PFLASH_CHAT_RECALL_MIN_LIFT (default 8), strongest first, up to -// PFLASH_CHAT_RECALL_TOKENS (default 2048) drafter tokens. +// PFLASH_CHAT_RECALL_MIN_LIFT (default 2), minus what the view holds. double pflash_chat_recall_min_lift() noexcept; -int pflash_chat_recall_tokens() noexcept; std::vector pflash_recall_by_lift( const std::vector> & lifts, const std::vector & in_view, - double min_lift, - int max_tokens); + double min_lift); // The parts of ``spans`` that ``minus`` does not cover. Both canonical. std::vector pflash_subtract_token_spans( diff --git a/server/test/test_server_unit.cpp b/server/test/test_server_unit.cpp index 928223d90..72c639b70 100644 --- a/server/test/test_server_unit.cpp +++ b/server/test/test_server_unit.cpp @@ -1090,16 +1090,17 @@ TEST_CASE(ServerUnitFixture, test_pflash_recall_by_lift_takes_clear_attention_on {{60, 70}, 25.0}, }; const std::vector in_view{{20, 45}}; - auto recalled = http_detail::pflash_recall_by_lift(lifts, in_view, 8.0, 1000); + auto recalled = http_detail::pflash_recall_by_lift(lifts, in_view, 8.0); TEST_ASSERT(recalled.size() == 2); // [0,10) and [45,70) merged TEST_ASSERT(recalled[0].begin == 0 && recalled[0].end == 10); TEST_ASSERT(recalled[1].begin == 45 && recalled[1].end == 70); - // A tight cap keeps the strongest: 40, then 25. - recalled = http_detail::pflash_recall_by_lift(lifts, in_view, 8.0, 20); - TEST_ASSERT(recalled.size() == 2); - TEST_ASSERT(recalled[0].begin == 0 && recalled[1].begin == 60); + // A lower bar takes the background segment too. + recalled = http_detail::pflash_recall_by_lift(lifts, in_view, 1.0); + TEST_ASSERT(recalled.size() == 2); // [0,20) and [45,70) + TEST_ASSERT(recalled[0].begin == 0 && recalled[0].end == 20); + TEST_ASSERT(recalled[1].begin == 45 && recalled[1].end == 70); // Nothing clears a high bar. - TEST_ASSERT(http_detail::pflash_recall_by_lift(lifts, in_view, 100.0, 1000).empty()); + TEST_ASSERT(http_detail::pflash_recall_by_lift(lifts, in_view, 100.0).empty()); } TEST_CASE(ServerUnitFixture, test_pflash_subtract_token_spans) { From 5af2ed99af63f5a4498f6d396b0e5f43f19ff14d Mon Sep 17 00:00:00 2001 From: Graffioh Date: Wed, 23 Sep 2026 15:14:02 +0000 Subject: [PATCH 24/26] fix(pflash): rescore from scratch once on non-finite head scores One request in ~500 of the recall-gate development run failed with "non-finite Qwen3.5 scoring-head scores" on a fresh chat whose prompt the other arm scored cleanly; it did not reproduce. The failure already forgets the scoring session, so score the prompt once more from scratch (one drafter forward) before failing the request with a 500. Co-Authored-By: Claude Opus 5.5 (1M context) --- server/src/pflash/qwen35_drafter.cpp | 16 +++++++++++++++- 1 file changed, 15 insertions(+), 1 deletion(-) diff --git a/server/src/pflash/qwen35_drafter.cpp b/server/src/pflash/qwen35_drafter.cpp index 0be2be6a1..6d9cb98ea 100644 --- a/server/src/pflash/qwen35_drafter.cpp +++ b/server/src/pflash/qwen35_drafter.cpp @@ -1351,9 +1351,23 @@ std::vector qwen35_drafter_score_and_compress( &other_scores, experiment.split_fraction); } if (experiment.selection_active && !force_legacy) { - return qwen35_strict_score_and_compress( + auto kept = qwen35_strict_score_and_compress( *st, ids, keep_ratio, n_lookahead, score_query_end, experiment, required_instruction_spans); + // Non-finite head scores (once in ~500 development requests, not + // reproduced) leave the scoring session forgotten: score the prompt + // again from scratch, one drafter forward, instead of failing the + // request. + if (kept.empty() && + std::strncmp(luce_last_error(), "non-finite", 10) == 0) { + std::fprintf(stderr, + "[qwen35-scorer] non-finite scores; rescoring from scratch\n"); + std::fflush(stderr); + kept = qwen35_strict_score_and_compress( + *st, ids, keep_ratio, n_lookahead, score_query_end, experiment, + required_instruction_spans); + } + return kept; } if (st->head_loaded && !experiment.selection_active) { set_last_error("Qwen3.5 scoring head requires strict selection"); From 4d011e68f7429451b54c8a944d31eff859a84c85 Mon Sep 17 00:00:00 2001 From: Graffioh Date: Thu, 24 Sep 2026 06:36:23 +0000 Subject: [PATCH 25/26] feat(pflash): score the latest user turn's tail beside the prompt end On SCBench's multi-turn chats the prompt-end query answered 112 of 171 turns, the benchmark's explicit question 127: the gap is literal lookups (an identifier, a described function) the last token does not carry. Replaying those turns through the drafter, adding the latest user turn's tail (PFLASH_SELECT_QUERY_TOKENS) as a second query window at full weight serves the gold passage as often as the explicit query (kv follow-ups 31/32 recallable vs 5/32 from the last token alone; repoqa 40/48 vs 41/48 explicit, 28/48 last token). The last token still reads questions placed anywhere in the message. Co-Authored-By: Claude Opus 5.5 (1M context) --- server/README.md | 6 +++++- server/src/common/model_backend.h | 3 +++ server/src/deepseek4/deepseek4_backend.cpp | 3 ++- server/src/pflash/pflash_drafter.cpp | 7 ++++++- server/src/pflash/pflash_drafter.h | 5 ++++- server/src/pflash/pflash_selection.h | 6 ++++++ server/src/pflash/qwen35_drafter.cpp | 9 +++++++++ server/src/qwen3/qwen3_backend.cpp | 3 ++- server/src/qwen35/qwen35_backend.cpp | 3 ++- server/src/qwen35/qwen35_layer_split_adapter.cpp | 3 ++- server/src/server/http_server.cpp | 13 +++++++++++++ server/test/test_server_unit.cpp | 5 +++++ 12 files changed, 59 insertions(+), 7 deletions(-) diff --git a/server/README.md b/server/README.md index 37907fcdb..a6f8aaa46 100644 --- a/server/README.md +++ b/server/README.md @@ -393,7 +393,11 @@ generation prompt, where the model starts answering, having read the whole request. Nothing is parsed out of the user's text, so the question can sit anywhere in the message -- before a pasted document, in the middle of it, or among the user's own sentences -- and the user's own words score far above -the material they paste. Strict selection keeps the generation prompt and +the material they paste. The tail (`PFLASH_SELECT_QUERY_TOKENS`) of the +latest user turn scores alongside it as a second query window at the same +weight: the last token reads the whole request, the user's own tokens match +literal strings -- an identifier, a described function -- that it does not +carry. Strict selection keeps the generation prompt and the latest user turn's role header, and it runs on every turn of a multi-turn chat. In an agent loop the assistant and tool turns after the user's turn are scored like the rest of the conversation. A prompt without diff --git a/server/src/common/model_backend.h b/server/src/common/model_backend.h index f901433d1..89c2bb967 100644 --- a/server/src/common/model_backend.h +++ b/server/src/common/model_backend.h @@ -300,6 +300,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; + // 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}; std::string drafter_path; // GGUF path (for lazy-load) int drafter_gpu = 0; // backend-local GPU for PFlash drafter bool skip_park = false; // true on >=32GB GPUs diff --git a/server/src/deepseek4/deepseek4_backend.cpp b/server/src/deepseek4/deepseek4_backend.cpp index b12f927f7..74edf81a4 100644 --- a/server/src/deepseek4/deepseek4_backend.cpp +++ b/server/src/deepseek4/deepseek4_backend.cpp @@ -3643,7 +3643,8 @@ std::vector DeepSeek4Backend::compress_batch( pflash_drafter_ctx_, request.input_ids, request.keep_ratio, /*chunk_size=*/32, request.score_query_tokens, /*pool_kernel=*/13, score_query_end, request.required_instruction_spans, - request.query_suffix_candidates, request.history_query_spans); + request.query_suffix_candidates, request.history_query_spans, + request.turn_query_span); result.ok = !result.compressed_ids.empty(); if (result.ok) result.kept_spans = pflash_last_kept_spans(); if (result.ok) { diff --git a/server/src/pflash/pflash_drafter.cpp b/server/src/pflash/pflash_drafter.cpp index fad60ccd9..bae720e4f 100644 --- a/server/src/pflash/pflash_drafter.cpp +++ b/server/src/pflash/pflash_drafter.cpp @@ -107,7 +107,8 @@ std::vector drafter_score_and_compress( int score_query_end, const std::vector & required_instruction_spans, bool query_suffix_candidates, - const std::vector & history_queries) { + const std::vector & history_queries, + PFlashTokenSpan turn_query) { pflash_clear_kept_spans(); if (!ctx.loaded) { set_last_error("drafter not loaded"); @@ -134,6 +135,10 @@ std::vector drafter_score_and_compress( experiment.history_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; + } } if (!experiment.selection_active && !required_instruction_spans.empty()) { set_last_error( diff --git a/server/src/pflash/pflash_drafter.h b/server/src/pflash/pflash_drafter.h index bd89d2746..ceb2b847c 100644 --- a/server/src/pflash/pflash_drafter.h +++ b/server/src/pflash/pflash_drafter.h @@ -75,6 +75,8 @@ void free_drafter_weights(DrafterContext & ctx); // window are scored candidates, not a kept suffix // history_queries strict selection only: earlier questions' windows, most // recent first, mixed into the scores at halving weights +// turn_query strict selection only: the latest user turn's tail, mixed +// in at the query's own weight // // On failure returns empty vector + sets last_error. std::vector drafter_score_and_compress( @@ -88,6 +90,7 @@ std::vector drafter_score_and_compress( const std::vector & required_instruction_spans = {}, bool query_suffix_candidates = false, - const std::vector & history_queries = {}); + const std::vector & history_queries = {}, + PFlashTokenSpan turn_query = {-1, -1}); } // namespace luce::common diff --git a/server/src/pflash/pflash_selection.h b/server/src/pflash/pflash_selection.h index fb1efef8d..20cf95557 100644 --- a/server/src/pflash/pflash_selection.h +++ b/server/src/pflash/pflash_selection.h @@ -121,6 +121,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: 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 + // strings (an identifier, a function description) the last token does + // not carry. + luce::common::PFlashTokenSpan turn_query{-1, -1}; }; // Segment probe: cut the context before every token whose boundary score is diff --git a/server/src/pflash/qwen35_drafter.cpp b/server/src/pflash/qwen35_drafter.cpp index 6d9cb98ea..5313f22e0 100644 --- a/server/src/pflash/qwen35_drafter.cpp +++ b/server/src/pflash/qwen35_drafter.cpp @@ -937,6 +937,15 @@ std::vector qwen35_strict_score_and_compress( return fail("qwen35 scorer query rows unavailable"); } windows.push_back(std::move(query)); + const auto & turn = experiment.turn_query; + if (turn.begin >= 0 && turn.end <= query_start && turn.end > turn.begin) { + ScoredWindow tail; + tail.begin = turn.begin; + tail.end = turn.end; + if (rows_for(turn.begin, turn.end, tail.rows)) { + windows.push_back(std::move(tail)); + } + } double weight = 1.0; for (const auto & span : experiment.history_queries) { weight *= 0.5; diff --git a/server/src/qwen3/qwen3_backend.cpp b/server/src/qwen3/qwen3_backend.cpp index 3b0a04b94..4122b31a5 100644 --- a/server/src/qwen3/qwen3_backend.cpp +++ b/server/src/qwen3/qwen3_backend.cpp @@ -977,7 +977,8 @@ ModelBackend::CompressResult Qwen3Backend::compress(const CompressRequest & req) drafter_ctx_, req.input_ids, req.keep_ratio, /*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.query_suffix_candidates, req.history_query_spans, + req.turn_query_span)); if (req.residency_action == DraftResidencyAction::ReleaseAfterUse) { free_drafter(); diff --git a/server/src/qwen35/qwen35_backend.cpp b/server/src/qwen35/qwen35_backend.cpp index 488e86687..eab031219 100644 --- a/server/src/qwen35/qwen35_backend.cpp +++ b/server/src/qwen35/qwen35_backend.cpp @@ -1211,7 +1211,8 @@ std::vector Qwen35Backend::compress_batch( drafter_ctx_, request.input_ids, request.keep_ratio, /*chunk_size=*/32, request.score_query_tokens, /*pool_kernel=*/13, score_query_end, request.required_instruction_spans, - request.query_suffix_candidates, request.history_query_spans); + request.query_suffix_candidates, request.history_query_spans, + request.turn_query_span); 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 79e8f642f..07b28582e 100644 --- a/server/src/qwen35/qwen35_layer_split_adapter.cpp +++ b/server/src/qwen35/qwen35_layer_split_adapter.cpp @@ -1394,7 +1394,8 @@ Qwen35LayerSplitAdapter::compress(const ModelBackend::CompressRequest & req) { pflash_drafter_, req.input_ids, req.keep_ratio, /*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.query_suffix_candidates, req.history_query_spans, + req.turn_query_span); 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 26fb81678..75012e657 100644 --- a/server/src/server/http_server.cpp +++ b/server/src/server/http_server.cpp @@ -3866,6 +3866,10 @@ 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 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. + PFlashTokenSpan turn_query_span{-1, -1}; // Header ("<|im_start|>user\n") opening the query's turn, when the chat // markers resolved it — pinned mandatory so a compressed prompt keeps // the current turn's role envelope. @@ -4045,6 +4049,12 @@ std::string HttpServer::apply_pflash_compression( chat_turn.generation_begin < prompt_end) { query_span = {prompt_end - 1, prompt_end}; query_span_rule = "prompt_end"; + const auto tail = http_detail::pflash_tail_query_window( + drafter_ids, experiment.query_tokens, + chat_turn.content_end, chat_turn.content_begin); + if (tail.valid()) { + turn_query_span = {tail.end - tail.tokens, tail.end}; + } } else { const auto window = http_detail::pflash_tail_query_window( drafter_ids, experiment.query_tokens, @@ -4425,6 +4435,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.turn_query_span = turn_query_span; compress_request.keep_ratio = http_detail::resolve_pflash_keep_ratio( pflash_keep_ratio(config_, prompt_tokens), req.session_id, sessions_); if (experiment.selection_active && query_window.valid()) { @@ -4468,6 +4479,8 @@ std::string HttpServer::apply_pflash_compression( {"query_span_end", query_span.end}, {"query_suffix_candidates", query_suffix_candidates}, {"history_queries", history_query_spans.size()}, + {"turn_query_begin", turn_query_span.begin}, + {"turn_query_end", turn_query_span.end}, {"requested_query_tokens", experiment.query_tokens}, {"required_text_count", req.pflash_required.size()}, {"expected_query_ids", expected_query_ids}, diff --git a/server/test/test_server_unit.cpp b/server/test/test_server_unit.cpp index 72c639b70..c1984f495 100644 --- a/server/test/test_server_unit.cpp +++ b/server/test/test_server_unit.cpp @@ -7029,6 +7029,11 @@ TEST_CASE(ServerUnitFixture, // answering; nothing is parsed out of the user's text. TEST_ASSERT(backend.last_request.score_query_end == (int) ids.size()); TEST_ASSERT(backend.last_request.score_query_tokens == 1); + // The user turn's tail scores as a second query window. + const auto turn = backend.last_request.turn_query_span; + TEST_ASSERT(turn.begin >= 0 && turn.end == last_im_end); + TEST_ASSERT(tokenizer.decode({ids.begin() + turn.begin, ids.begin() + turn.end}) + == "What is the answer?"); // The generation prompt, the turn's role header and -- a short turn -- // the whole question stay. bool header_pinned = false; From d743a4ea0b1b15576f62edb285df3e7ac1ef120c Mon Sep 17 00:00:00 2001 From: Graffioh Date: Fri, 25 Sep 2026 19:27:41 +0000 Subject: [PATCH 26/26] feat(pflash): recall what a fresh selection keeps, as whole passages Recall decides what a follow-up gets beyond the cached compressed prompt. It used to take segments whose head lift cleared a gate (~300 tokens) and put them at the start of the new user turn. On MTRAG (IBM, 32 human conversations, 255 turns, a ~64K shared context per conversation, Jev-judged against its references) that served every gold passage of a follow-up on 17 of 207 turns, against 87 when each turn is compressed from scratch, and answered 102 turns against 126 from scratch and 120 for full prefill. Recall now takes what the fresh selection for the new question keeps that the view lacks, every kept piece the view does not fully hold, whole and in document order, before the new question (after it measured 91). Recalled passages are appended whatever their size; the view is rebuilt from the fresh selection only when it outgrows twice that selection or the context. PFLASH_CHAT_RECALL_MIN_LIFT is gone. MTRAG: 118 of 255 (dependent follow-ups 32 of 77, full prefill 30; standalone 68 of 146, full prefill 74), first turn 24.2 s against 98.1 s, follow-ups 7.6 s against 2.9 s, whole conversation 187 s against 273 s. With ~44K tokens pasted at turn 3: 120 against 110 correct, the paste answered in 33.5 s against 108.7 s. Co-Authored-By: Claude Opus 5.5 (1M context) --- server/README.md | 12 +++--- server/src/server/http_server.cpp | 65 ++++++++++--------------------- server/src/server/http_server.h | 9 ----- server/test/test_server_unit.cpp | 22 ----------- 4 files changed, 25 insertions(+), 83 deletions(-) diff --git a/server/README.md b/server/README.md index a6f8aaa46..81128f983 100644 --- a/server/README.md +++ b/server/README.md @@ -429,13 +429,11 @@ prefill appends a follow-up; from `PFLASH_CHAT_COMPRESS_NEW_TOKENS` (default 16384) tokens of new material (a pasted document, a large tool output) only what the fresh selection keeps of it is appended, and the view before it stays cached. `PFLASH_CHAT_RECALL=0` turns recall off: a small follow-up is -then served without running the drafter at all. Recall takes the segments -the new question clearly attends to: attention lift (mass per token relative -to uniform attention) of at least `PFLASH_CHAT_RECALL_MIN_LIFT` (default 2), -so a content-free follow-up ("which documents support that?") recalls next -to nothing instead of filling the budget with noise beside the question. A -question that needs more than a third of a fresh selection starts a new view -from that selection instead. Kept pieces that were not adjacent in the prompt are joined by a +then served without running the drafter at all. Recall takes what a fresh +selection for the new question keeps that the view lacks: a question on the +view's topic misses little and is served by appending it, and a question +that needs more than a third of a fresh selection -- the conversation moved +to other material -- starts a new view from that selection instead. Kept pieces that were not adjacent in the prompt are joined by a paragraph break when neither side has one (`PFLASH_SELECT_PARAGRAPH_JOIN=0` turns it off; the breaks do not count against the token ceiling). `PFLASH_VIEW_TRACE_PATH` diff --git a/server/src/server/http_server.cpp b/server/src/server/http_server.cpp index 75012e657..8ae125300 100644 --- a/server/src/server/http_server.cpp +++ b/server/src/server/http_server.cpp @@ -677,29 +677,6 @@ int pflash_chat_compress_new_tokens() noexcept { return (int) (std::min)(value, 1L << 30); } -double pflash_chat_recall_min_lift() noexcept { - const char * raw = std::getenv("PFLASH_CHAT_RECALL_MIN_LIFT"); - if (!raw || !*raw) return 2.0; - char * end = nullptr; - const double value = std::strtod(raw, &end); - if (end == raw || *end != '\0' || !std::isfinite(value) || value < 0.0) return 2.0; - return value; -} - -std::vector pflash_recall_by_lift( - const std::vector> & lifts, - const std::vector & in_view, - double min_lift) { - std::vector chosen; - for (const auto & [span, lift] : lifts) { - if (!(lift >= min_lift)) continue; - for (const auto & part : pflash_subtract_token_spans({span}, in_view)) { - chosen.push_back(part); - } - } - return canonicalize_pflash_token_spans(std::move(chosen)); -} - int pflash_chat_history_queries() noexcept { const char * raw = std::getenv("PFLASH_CHAT_HISTORY_QUERIES"); if (!raw || !*raw) return 3; @@ -4534,7 +4511,7 @@ std::string HttpServer::apply_pflash_compression( json view_stats; if (serve_pflash_chat_view( req, compress_request.input_ids, chat_turn, nullptr, nullptr, - nullptr, served, prepared.snapshot_cut, view_stats)) { + served, prepared.snapshot_cut, view_stats)) { prepared.tokens = std::move(served); prepared.compressed = true; prepared.pflash_stats = { @@ -4683,7 +4660,7 @@ std::string HttpServer::apply_pflash_compression( std::vector served; if (serve_pflash_chat_view( req, compress_request.input_ids, chat_turn, &final_tokens, - &result.kept_spans, &result.candidate_lifts, served, + &result.kept_spans, served, prepared.snapshot_cut, prepared.pflash_stats["view"])) { final_tokens = std::move(served); } @@ -4720,7 +4697,6 @@ bool HttpServer::serve_pflash_chat_view( const http_detail::PflashChatTurnSpan & turn, const std::vector * fresh, const std::vector * kept_spans, - const std::vector> * lifts, std::vector & served, int & snapshot_cut, json & stats) { @@ -4806,31 +4782,30 @@ bool HttpServer::serve_pflash_chat_view( const bool recall = new_question && http_detail::pflash_chat_recall(); if (!compressed && (compress_new || recall)) return false; - // Recall: what the new question clearly attends to that the view does - // not hold. Only a new user turn brings a new question; an agent step - // (assistant call plus tool output) appends without recalling. The head's - // per-candidate lifts decide it (the fresh selection minus the view when - // a scorer reports none), so a content-free question ("which documents - // support that?") recalls next to nothing. A question that needs more - // than a third of a fresh selection starts a new view from that - // selection instead: past that, prefilling the fresh prompt costs about - // the same and serves the material in order. + // Recall: what a fresh selection for the new question keeps that the + // view does not hold. Only a new user turn brings a new question; an + // agent step (assistant call plus tool output) appends without + // recalling. What is missing is appended, whatever its size, so the + // cached view is never thrown away for a new topic; the view starts over + // from the fresh selection only when it outgrows it (below). std::vector recalled; if (compressed && recall) { auto in_view = view.spans; in_view.push_back({view.drafter_gen_begin, input}); in_view = http_detail::canonicalize_pflash_token_spans(std::move(in_view)); - recalled = lifts && !lifts->empty() - ? http_detail::pflash_recall_by_lift( - *lifts, in_view, http_detail::pflash_chat_recall_min_lift()) - : http_detail::pflash_subtract_token_spans(*kept_spans, in_view); - size_t recall_size = 0; - for (const auto & span : recalled) { - recall_size += (size_t) (span.end - span.begin); - } - if (3 * recall_size > fresh->size()) { - return serve_fresh("rebuild", view.turns + 1); + const auto missing = + http_detail::pflash_subtract_token_spans(*kept_spans, in_view); + // Whole kept pieces: a passage the view holds only part of comes + // back in one piece, in order, not as a fragment far from the rest. + for (const auto & span : *kept_spans) { + for (const auto & part : missing) { + if (part.begin < span.end && part.end > span.begin) { + recalled.push_back(span); + break; + } + } } + recalled = http_detail::canonicalize_pflash_token_spans(std::move(recalled)); } std::string recall_block; int recalled_tokens = 0; diff --git a/server/src/server/http_server.h b/server/src/server/http_server.h index b8b609d16..3e46ce04a 100644 --- a/server/src/server/http_server.h +++ b/server/src/server/http_server.h @@ -333,14 +333,6 @@ std::string pflash_join_kept_spans( const std::vector & spans); int pflash_chat_compress_new_tokens() noexcept; -// Recall takes the segments the new question clearly attends to: attention -// lift (mass per token relative to uniform) of at least -// PFLASH_CHAT_RECALL_MIN_LIFT (default 2), minus what the view holds. -double pflash_chat_recall_min_lift() noexcept; -std::vector pflash_recall_by_lift( - const std::vector> & lifts, - const std::vector & in_view, - double min_lift); // The parts of ``spans`` that ``minus`` does not cover. Both canonical. std::vector pflash_subtract_token_spans( @@ -684,7 +676,6 @@ class HttpServer { const http_detail::PflashChatTurnSpan & turn, const std::vector * fresh, const std::vector * kept_spans, - const std::vector> * lifts, std::vector & served, int & snapshot_cut, nlohmann::json & stats); diff --git a/server/test/test_server_unit.cpp b/server/test/test_server_unit.cpp index c1984f495..cce52da89 100644 --- a/server/test/test_server_unit.cpp +++ b/server/test/test_server_unit.cpp @@ -1081,28 +1081,6 @@ TEST_CASE(ServerUnitFixture, test_pflash_join_kept_spans_breaks_between_pieces) unlink(path.c_str()); } -TEST_CASE(ServerUnitFixture, test_pflash_recall_by_lift_takes_clear_attention_only) { - const std::vector> lifts{ - {{0, 10}, 40.0}, // clearly attended, not in view - {{10, 20}, 2.0}, // background - {{20, 30}, 90.0}, // clearly attended, already in view - {{30, 60}, 12.0}, // attended, half in view - {{60, 70}, 25.0}, - }; - const std::vector in_view{{20, 45}}; - auto recalled = http_detail::pflash_recall_by_lift(lifts, in_view, 8.0); - TEST_ASSERT(recalled.size() == 2); // [0,10) and [45,70) merged - TEST_ASSERT(recalled[0].begin == 0 && recalled[0].end == 10); - TEST_ASSERT(recalled[1].begin == 45 && recalled[1].end == 70); - // A lower bar takes the background segment too. - recalled = http_detail::pflash_recall_by_lift(lifts, in_view, 1.0); - TEST_ASSERT(recalled.size() == 2); // [0,20) and [45,70) - TEST_ASSERT(recalled[0].begin == 0 && recalled[0].end == 20); - TEST_ASSERT(recalled[1].begin == 45 && recalled[1].end == 70); - // Nothing clears a high bar. - TEST_ASSERT(http_detail::pflash_recall_by_lift(lifts, in_view, 100.0).empty()); -} - TEST_CASE(ServerUnitFixture, test_pflash_subtract_token_spans) { const std::vector spans{{0, 10}, {20, 30}, {40, 50}}; const std::vector minus{{5, 22}, {25, 26}, {40, 50}};