diff --git a/server/src/common/model_backend.h b/server/src/common/model_backend.h index 221e54ab5..23e4a1b3c 100644 --- a/server/src/common/model_backend.h +++ b/server/src/common/model_backend.h @@ -423,6 +423,12 @@ struct ModelBackend { struct CompressRequest { std::vector input_ids; // drafter-tokenized prompt float keep_ratio; // fraction to keep (0.0–1.0) + // Exclusive end and width of the user-query token window inside + // input_ids. Negative end preserves the legacy trailing-token window. + // Keeping this separate from DFLASH_COMPRESS_QUERY_TOKENS matters: + // that knob controls lexical anchors, not neural scorer Q rows. + int score_query_end = -1; + int score_query_tokens = 8; 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.cpp b/server/src/common/pflash_drafter_ipc.cpp index 037a35f05..724717e49 100644 --- a/server/src/common/pflash_drafter_ipc.cpp +++ b/server/src/common/pflash_drafter_ipc.cpp @@ -40,9 +40,12 @@ bool PFlashDrafterIpcClient::start( bool PFlashDrafterIpcClient::compress( const std::vector & input_ids, float keep_ratio, - std::vector & compressed_ids) { + std::vector & compressed_ids, + int score_query_end, + int score_query_tokens) { #if defined(_WIN32) (void)input_ids; (void)keep_ratio; (void)compressed_ids; + (void)score_query_end; (void)score_query_tokens; return false; #else compressed_ids.clear(); @@ -58,7 +61,9 @@ bool PFlashDrafterIpcClient::compress( 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 %s\n", keep_x1000, path.c_str()); + std::fprintf(cmd, "compress %d %d %d %s\n", + keep_x1000, score_query_end, score_query_tokens, + path.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 1cc548eeb..e7ffecce4 100644 --- a/server/src/common/pflash_drafter_ipc.h +++ b/server/src/common/pflash_drafter_ipc.h @@ -30,7 +30,9 @@ class PFlashDrafterIpcClient { bool compress(const std::vector & input_ids, float keep_ratio, - std::vector & compressed_ids); + std::vector & compressed_ids, + int score_query_end = -1, + int score_query_tokens = 8); 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 8accfc547..33e114c7a 100644 --- a/server/src/common/pflash_drafter_ipc_daemon.cpp +++ b/server/src/common/pflash_drafter_ipc_daemon.cpp @@ -47,9 +47,12 @@ int run_pflash_drafter_ipc_daemon(const char * drafter_path, if (cmd == "quit" || cmd == "exit") break; if (cmd == "compress") { int keep_x1000 = 0; - iss >> keep_x1000; + 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 || path.empty()) { + if (keep_x1000 < 0 || keep_x1000 > 1000 || + score_query_tokens < 1 || path.empty()) { std::fprintf(stderr, "[pflash-ipc-daemon] bad compress: %s\n", line.c_str()); stream_status(stream_fd, -1); @@ -63,7 +66,9 @@ int run_pflash_drafter_ipc_daemon(const char * drafter_path, continue; } const float keep = (float)keep_x1000 / 1000.0f; - auto compressed = drafter_score_and_compress(ctx, input_ids, keep); + auto compressed = drafter_score_and_compress( + ctx, input_ids, keep, /*chunk_size=*/32, score_query_tokens, + /*pool_kernel=*/13, score_query_end); if (compressed.empty()) { std::fprintf(stderr, "[pflash-ipc-daemon] compress returned empty\n"); stream_status(stream_fd, -1); diff --git a/server/src/qwen3/qwen3_backend.cpp b/server/src/qwen3/qwen3_backend.cpp index 8748a67bc..633c40672 100644 --- a/server/src/qwen3/qwen3_backend.cpp +++ b/server/src/qwen3/qwen3_backend.cpp @@ -968,7 +968,9 @@ ModelBackend::CompressResult Qwen3Backend::compress(const CompressRequest & req) } result.compressed_ids = drafter_score_and_compress( - drafter_ctx_, req.input_ids, req.keep_ratio); + drafter_ctx_, req.input_ids, req.keep_ratio, + /*chunk_size=*/32, req.score_query_tokens, /*pool_kernel=*/13, + req.score_query_end); result.ok = true; if (req.residency_action == DraftResidencyAction::ReleaseAfterUse) { diff --git a/server/src/qwen3/qwen3_drafter.cpp b/server/src/qwen3/qwen3_drafter.cpp index 3acc1775a..35e4c89fa 100644 --- a/server/src/qwen3/qwen3_drafter.cpp +++ b/server/src/qwen3/qwen3_drafter.cpp @@ -272,11 +272,18 @@ static std::vector qwen35_score_and_compress( float keep_ratio, int chunk_size, int n_lookahead, - int pool_kernel) { + 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 < 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); @@ -441,7 +448,7 @@ static std::vector qwen35_score_and_compress( } 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)(S - n_lookahead) * act_in->nb[1]); + 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); @@ -473,7 +480,7 @@ static std::vector qwen35_score_and_compress( } std::vector pos4((size_t)4 * n_lookahead, 0); for (int i = 0; i < n_lookahead; ++i) { - int p = S - 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; @@ -481,7 +488,7 @@ static std::vector qwen35_score_and_compress( 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 = S - n_lookahead + t + 1; + 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; } @@ -495,6 +502,21 @@ static std::vector qwen35_score_and_compress( } 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) { @@ -684,7 +706,8 @@ std::vector drafter_score_and_compress( float keep_ratio, int chunk_size, int n_lookahead, - int pool_kernel) { + int pool_kernel, + int score_query_end) { if (!ctx.loaded) { set_last_error("drafter not loaded"); return {}; @@ -695,7 +718,8 @@ std::vector drafter_score_and_compress( return {}; } auto * st = static_cast(ctx.arch_state); - return qwen35_score_and_compress(st->weights, ids, keep_ratio, chunk_size, n_lookahead, pool_kernel); + return qwen35_score_and_compress(st->weights, ids, keep_ratio, chunk_size, + n_lookahead, pool_kernel, score_query_end); } const int S = (int)ids.size(); if (S < n_lookahead + 1) { @@ -706,7 +730,8 @@ std::vector drafter_score_and_compress( // ── 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)) { + if (!forward_qwen3_drafter_model( + ctx.weights, ids, n_lookahead, running_max, score_query_end)) { return {}; } auto t1 = std::chrono::steady_clock::now(); diff --git a/server/src/qwen3/qwen3_drafter.h b/server/src/qwen3/qwen3_drafter.h index e5424f9dd..bd006f6db 100644 --- a/server/src/qwen3/qwen3_drafter.h +++ b/server/src/qwen3/qwen3_drafter.h @@ -69,8 +69,9 @@ void free_drafter_weights(DrafterContext & ctx); // ids input token IDs of length S // keep_ratio fraction of `chunk_size`-token chunks to keep // chunk_size span granularity (default 32) -// n_lookahead trailing Q tokens used for tail attention (default 8) +// 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 // // On failure returns empty vector + sets last_error. std::vector drafter_score_and_compress( @@ -79,6 +80,7 @@ std::vector drafter_score_and_compress( float keep_ratio, int chunk_size = 32, int n_lookahead = 8, - int pool_kernel = 13); + int pool_kernel = 13, + int score_query_end = -1); } // namespace dflash::common diff --git a/server/src/qwen3/qwen3_drafter_model.h b/server/src/qwen3/qwen3_drafter_model.h index 8f9b9620b..c3f7c235d 100644 --- a/server/src/qwen3/qwen3_drafter_model.h +++ b/server/src/qwen3/qwen3_drafter_model.h @@ -13,6 +13,8 @@ #include "ggml.h" +#include +#include #include #include #include @@ -76,11 +78,12 @@ void free_qwen3_drafter_model(Qwen3DrafterWeights & w); // Inputs: // w — loaded weights (must be on the selected GPU backend) // ids — input token IDs of length S (drafter vocab) -// n_lookahead — number of trailing query tokens for tail attention (=8) +// 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_tail @ K^T / sqrt(D)) per (lookahead, key) pair. +// 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. @@ -88,6 +91,39 @@ bool forward_qwen3_drafter_model( const Qwen3DrafterWeights & w, const std::vector & ids, int n_lookahead, - std::vector & running_max); + 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; +} } // namespace dflash::common diff --git a/server/src/qwen3/qwen3_graph.cpp b/server/src/qwen3/qwen3_graph.cpp index 9fdc75139..3e40b08f5 100644 --- a/server/src/qwen3/qwen3_graph.cpp +++ b/server/src/qwen3/qwen3_graph.cpp @@ -224,7 +224,8 @@ bool forward_qwen3_drafter_model( const Qwen3DrafterWeights & w, const std::vector & ids, int n_lookahead, - std::vector & running_max) + std::vector & running_max, + int score_query_end) { if (!w.backend || !w.tok_embd) { set_last_error("forward_qwen3_drafter_model: weights not loaded"); @@ -250,10 +251,17 @@ bool forward_qwen3_drafter_model( return e == nullptr || std::string(e) != "0"; }(); - if (S < n_lookahead + 1) { + 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. @@ -353,7 +361,7 @@ bool forward_qwen3_drafter_model( { std::vector m((size_t)n_lookahead * S, 0.0f); for (int t = 0; t < n_lookahead; ++t) { - int visible_end = S - n_lookahead + t + 1; + const int visible_end = query_start + t + 1; for (int j = 0; j < S; ++j) { m[(size_t)t * S + j] = (j < visible_end) ? 0.0f : -INFINITY; } @@ -463,15 +471,21 @@ bool forward_qwen3_drafter_model( // 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 int tail_lo_nr = S - n_lookahead; - if (tail_lo_nr >= cs && tail_lo_nr + n_lookahead <= cs + cl) { - const int local_lo_nr = tail_lo_nr - cs; + 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, n_lookahead, + gA, Q, D, H, capture.tokens, Q->nb[1], Q->nb[2], - (size_t)local_lo_nr * 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_norope_v[si].t)); + ggml_cpy(gA, Q_prenrope_tail, Q_prenrope_dst)); } } Q = ggml_rope_ext(gA, Q, pos_chunk, nullptr, D, @@ -515,16 +529,22 @@ bool forward_qwen3_drafter_model( ggml_build_forward_expand(gfA, ggml_cpy(gA, K, K_dst)); ggml_build_forward_expand(gfA, ggml_cpy(gA, V, V_dst)); - // Copy Q tail to Q_last_v[il] in the chunk that contains the tail. - const int tail_lo = S - n_lookahead; - if (!nope_tail && tail_lo >= cs && tail_lo + n_lookahead <= cs + cl) { - int local_lo = tail_lo - cs; + // 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, n_lookahead, + gA, Q, D, H, capture.tokens, Q->nb[1], Q->nb[2], - (size_t)local_lo * 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_last_v[layer_cache_idx].t)); + ggml_cpy(gA, Q_tail_local, Q_tail_dst)); } auto tA_setup1 = std::chrono::steady_clock::now(); @@ -827,12 +847,33 @@ bool forward_qwen3_drafter_model( cleanup_all(); return false; } - ggml_backend_graph_compute(w.backend, gf); - ggml_backend_tensor_get(probs, probs_h.data(), 0, - probs_h.size() * sizeof(float)); + 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) { diff --git a/server/src/qwen35/qwen35_backend.cpp b/server/src/qwen35/qwen35_backend.cpp index 20a4b4f1a..a90c2d7fd 100644 --- a/server/src/qwen35/qwen35_backend.cpp +++ b/server/src/qwen35/qwen35_backend.cpp @@ -1187,7 +1187,9 @@ std::vector Qwen35Backend::compress_batch( auto & result = results[index]; result.compressed_ids = drafter_score_and_compress( - drafter_ctx_, request.input_ids, request.keep_ratio); + drafter_ctx_, request.input_ids, request.keep_ratio, + /*chunk_size=*/32, request.score_query_tokens, /*pool_kernel=*/13, + request.score_query_end); 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 9b63e86f2..f073f5925 100644 --- a/server/src/qwen35/qwen35_layer_split_adapter.cpp +++ b/server/src/qwen35/qwen35_layer_split_adapter.cpp @@ -1384,7 +1384,9 @@ Qwen35LayerSplitAdapter::compress(const ModelBackend::CompressRequest & req) { } result.compressed_ids = drafter_score_and_compress( - pflash_drafter_, req.input_ids, req.keep_ratio); + pflash_drafter_, 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(); 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 a909344e8..8d70d180e 100644 --- a/server/src/server/http_server.cpp +++ b/server/src/server/http_server.cpp @@ -186,6 +186,32 @@ HeartbeatSendResult try_send_sse_heartbeat( return HeartbeatSendResult::Complete; } +PflashQueryWindow find_pflash_query_window( + const std::vector & prompt, + const std::vector & query, + int max_tokens) { + PflashQueryWindow result; + if (prompt.empty() || query.empty() || max_tokens < 1) 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); + for (int width = widest; width >= narrowest; --width) { + const auto query_begin = query.end() - width; + for (int end = (int) prompt.size(); end >= width; --end) { + if (std::equal(query_begin, query.end(), + prompt.begin() + end - width)) { + result.end = end; + result.tokens = width; + return result; + } + } + } + return result; +} + } // namespace http_detail static std::string context_overflow_message(int max_ctx, int prompt_tokens, int max_output) { @@ -2819,6 +2845,31 @@ std::string HttpServer::apply_pflash_compression( const int prompt_tokens = (int) req.prompt_tokens.size(); const std::string prompt_text = tokenizer_.decode(req.prompt_tokens); auto drafter_ids = drafter_tokenizer_->encode(prompt_text); + + std::string last_user_text; + 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; + } + } + 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); if (drafter_ids.empty()) { return "PFlash drafter tokenizer produced an empty prompt"; } @@ -2827,6 +2878,19 @@ std::string HttpServer::apply_pflash_compression( compress_request.input_ids = std::move(drafter_ids); 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; + std::fprintf(stderr, + "[pflash] scorer query mapped to drafter tokens [%d,%d); " + "rendered suffix=%zu tokens\n", + query_window.end - query_window.tokens, query_window.end, + compress_request.input_ids.size() - (size_t) query_window.end); + } else { + std::fprintf(stderr, + "[pflash] ERROR: scorer query mapping failed; refusing compression\n"); + return "PFlash scorer query mapping failed"; + } compress_request.drafter_path = config_.pflash_drafter_path; compress_request.drafter_gpu = config_.pflash_drafter_gpu; compress_request.skip_park = config_.pflash_skip_park; @@ -2850,7 +2914,9 @@ std::string HttpServer::apply_pflash_compression( } result.ok = pflash_remote_.compress( compress_request.input_ids, compress_request.keep_ratio, - result.compressed_ids); + result.compressed_ids, + compress_request.score_query_end, + compress_request.score_query_tokens); if (residency == DraftResidencyAction::ReleaseAfterUse) { pflash_remote_.close(); } @@ -2869,28 +2935,7 @@ std::string HttpServer::apply_pflash_compression( // 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. - std::string last_user_text; - 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; - } - } - if (!last_user_text.empty()) { - const auto query_ids = drafter_tokenizer_->encode(last_user_text); int query_kept = 0; if (!query_ids.empty()) { int query_index = (int) query_ids.size() - 1; diff --git a/server/src/server/http_server.h b/server/src/server/http_server.h index 52c36473b..c2b4d4ea3 100644 --- a/server/src/server/http_server.h +++ b/server/src/server/http_server.h @@ -262,6 +262,20 @@ bool canonical_assistant_content( const std::string & generated_text, std::string & content); +struct PflashQueryWindow { + int end = -1; // exclusive token offset in the rendered prompt + int tokens = 0; // width of the matching query suffix + + bool valid() const { return end >= tokens && tokens > 0; } +}; + +// 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 max_tokens = 8); + } // namespace http_detail // ─── Parsed request ───────────────────────────────────────────────────── diff --git a/server/test/test_server_unit.cpp b/server/test/test_server_unit.cpp index 698b865be..b384b8153 100644 --- a/server/test/test_server_unit.cpp +++ b/server/test/test_server_unit.cpp @@ -108,6 +108,78 @@ TEST_CASE(ServerUnitFixture, test_api_format_names_are_total) { CHECK(std::string(api_format_name(ApiFormat::COMPLETIONS)) == "completions"); } +TEST_CASE(ServerUnitFixture, test_pflash_scorer_uses_user_query_before_chat_suffix) { + const std::vector query{ + 90, 91, 100, 101, 102, 103, 104, 105, 106, 107, + }; + const std::vector rendered{ + 1, 2, 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); + + TEST_ASSERT(window.valid()); + TEST_ASSERT(window.tokens == 8); + TEST_ASSERT(window.end == 10); + TEST_ASSERT((int)rendered.size() - window.end == 8); +} + +TEST_CASE(ServerUnitFixture, test_pflash_query_mapping_tolerates_one_bpe_boundary_token) { + 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, + }; + + 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_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).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); + TEST_ASSERT(short_window.valid()); + TEST_ASSERT(short_window.tokens == 3); + TEST_ASSERT(short_window.end == 4); +} + +TEST_CASE(ServerUnitFixture, test_pflash_score_validation_counts_nan_and_inf) { + const float values[]{ + 0.0f, + std::numeric_limits::quiet_NaN(), + std::numeric_limits::infinity(), + -std::numeric_limits::infinity(), + 1.0f, + }; + TEST_ASSERT(count_nonfinite_scores(values, 5) == 3); + TEST_ASSERT(count_nonfinite_scores(values, 1) == 0); +} + +TEST_CASE(ServerUnitFixture, test_pflash_query_capture_splits_across_chunks) { + const auto first = query_capture_slice(4093, 4101, 0, 4096); + TEST_ASSERT(first.valid()); + TEST_ASSERT(first.chunk_offset == 4093); + TEST_ASSERT(first.query_offset == 0); + TEST_ASSERT(first.tokens == 3); + + const auto second = query_capture_slice(4093, 4101, 4096, 4096); + TEST_ASSERT(second.valid()); + TEST_ASSERT(second.chunk_offset == 0); + TEST_ASSERT(second.query_offset == 3); + TEST_ASSERT(second.tokens == 5); + + TEST_ASSERT(!query_capture_slice(4093, 4101, 8192, 4096).valid()); +} + TEST_CASE(ServerUnitFixture, test_daemon_io_external_cancellation_latches) { bool cancel = false; DaemonIO io;