fix(pflash): score the user query and reject non-finite scores

server/src/common/model_backend.h

@@ -423,6 +423,12 @@ struct ModelBackend {
     struct CompressRequest {
         std::vector<int32_t> 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

server/src/common/pflash_drafter_ipc.cpp

@@ -40,9 +40,12 @@ bool PFlashDrafterIpcClient::start(
 bool PFlashDrafterIpcClient::compress(
         const std::vector<int32_t> & input_ids,
         float keep_ratio,
-        std::vector<int32_t> & compressed_ids) {
+        std::vector<int32_t> & 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;

server/src/common/pflash_drafter_ipc.h

@@ -30,7 +30,9 @@ class PFlashDrafterIpcClient {
 
     bool compress(const std::vector<int32_t> & input_ids,
                   float keep_ratio,
-                  std::vector<int32_t> & compressed_ids);
+                  std::vector<int32_t> & compressed_ids,
+                  int score_query_end = -1,
+                  int score_query_tokens = 8);
 
     bool active() const { return active_; }
     void close();

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);

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) {

server/src/qwen3/qwen3_drafter.cpp

@@ -272,11 +272,18 @@ static std::vector<int32_t> 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<float> running_max((size_t)n_lookahead * S, -INFINITY);
@@ -441,7 +448,7 @@ static std::vector<int32_t> 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,15 +480,15 @@ static std::vector<int32_t> qwen35_score_and_compress(
             }
             std::vector<int32_t> 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;
             }
             ggml_backend_tensor_set(pos_tail, pos4.data(), 0, pos4.size() * sizeof(int32_t));
             std::vector<float> 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<int32_t> qwen35_score_and_compress(
             }
             std::vector<float> 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<int32_t> 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<int32_t> drafter_score_and_compress(
             return {};
         }
         auto * st = static_cast<Qwen35DrafterState *>(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<int32_t> drafter_score_and_compress(
     // ── 1. Custom forward + GPU tail-attention scoring ────────────────
     auto t0 = std::chrono::steady_clock::now();
     std::vector<float> 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();

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<int32_t> drafter_score_and_compress(
@@ -79,6 +80,7 @@ std::vector<int32_t> 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

server/src/qwen3/qwen3_drafter_model.h

@@ -13,6 +13,8 @@
 
 #include "ggml.h"
 
+#include <cmath>
+#include <cstddef>
 #include <cstdint>
 #include <string>
 #include <vector>
@@ -76,18 +78,52 @@ 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.
 bool forward_qwen3_drafter_model(
     const Qwen3DrafterWeights & w,
     const std::vector<int32_t> & ids,
     int n_lookahead,
-    std::vector<float> & running_max);
+    std::vector<float> & 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