diff --git a/docs/specs/openapi-props.yaml b/docs/specs/openapi-props.yaml index f9550651e..b39c4bf61 100644 --- a/docs/specs/openapi-props.yaml +++ b/docs/specs/openapi-props.yaml @@ -136,6 +136,10 @@ paths: lm_head_fix: null mode: "off" skip_park: null + skip_park_mode: null + skip_park_estimate_bytes: null + skip_park_free_bytes: null + drafter_keep_loaded: null threshold: null prefix_cache: capacity: 0 @@ -638,6 +642,10 @@ components: - lm_head_fix - mode - skip_park + - skip_park_mode + - skip_park_estimate_bytes + - skip_park_free_bytes + - drafter_keep_loaded - threshold properties: bsa_alpha: @@ -673,7 +681,28 @@ components: example: "off" skip_park: type: ["boolean", "null"] - description: Whether to skip park/unpark (large-VRAM GPUs). + description: Resolved skip-park decision — whether compression + keeps target+draft resident while the drafter scores. + example: null + skip_park_mode: + type: ["string", "null"] + enum: ["auto", "on", "off", null] + description: Requested `--prefill-skip-park` policy. + example: null + skip_park_estimate_bytes: + type: ["integer", "null"] + description: Auto-mode drafter footprint estimate incl. margin. + example: null + skip_park_free_bytes: + type: ["integer", "null"] + description: Free VRAM measured on the drafter GPU at startup. + example: null + drafter_keep_loaded: + type: ["boolean", "null"] + description: Auto skip-park only — the estimate also fits the + drafter kept resident beside the target's compute reserve, so + draft-residency `auto` keeps it loaded between requests. False + for explicit `on`/`off`. example: null threshold: type: ["integer", "null"] diff --git a/docs/specs/props-endpoint.md b/docs/specs/props-endpoint.md index d8f2116db..ee2ac76cd 100644 --- a/docs/specs/props-endpoint.md +++ b/docs/specs/props-endpoint.md @@ -343,6 +343,10 @@ format and resolution order. "lm_head_fix": null, "mode": "off", "skip_park": null, + "skip_park_mode": null, + "skip_park_estimate_bytes": null, + "skip_park_free_bytes": null, + "drafter_keep_loaded": null, "threshold": null } ``` @@ -355,7 +359,13 @@ enabled, fields carry the runtime configuration: - `threshold` — token-count threshold for AUTO mode - `keep_ratio` — fraction of tokens retained after compression - `drafter_gguf` — path to the compression drafter GGUF -- `skip_park` — whether to skip park/unpark (large-VRAM GPUs) +- `skip_park` — whether compression actually skips park/unpark (resolved) +- `skip_park_mode` — `"auto" | "on" | "off"` requested policy +- `skip_park_estimate_bytes` / `skip_park_free_bytes` — the auto-mode + footprint estimate (incl. margin) and the free VRAM measured at startup +- `drafter_keep_loaded` — auto mode only: the estimate plus the target's + compute reserve also fits, so `--draft-residency auto` keeps the drafter + loaded between requests instead of reloading it per request - `bsa_enabled` / `bsa_alpha` / `lm_head_fix` — backend-specific PFlash tunables diff --git a/optimizations/kvflash/DESIGN.md b/optimizations/kvflash/DESIGN.md index 965dde709..856c53a75 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 b23cb6f0f..36b8b6bf5 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). luce_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 @@ luce_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: `LUCE_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 b53638c96..060f7da66 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 `LUCE_LM_HEAD_FIX=0` and `LUCE_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 `LUCE_LM_HEAD_FIX=0` a --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 `LUCE_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 luce 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 89c214892..942c6bba4 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 e289d2579..8f4dbb85a 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/server/models/draft/model.safetensors", help="draft model used for spec decoding (NOT drafter scorer)") - ap.add_argument("--drafter-gguf", default="/home/lucebox/lucebox-hub/server/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/server/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 d5d95d706..36b9e0dd9 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. @@ -476,6 +476,7 @@ set(LUCE_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 @@ -492,11 +493,14 @@ add_library(luce_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/qwen3_drafter.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 @@ -599,6 +603,7 @@ add_library(luce_common STATIC src/common/backend_precision.cpp src/common/daemon_loop.cpp src/common/gguf_inspect.cpp + src/placement/gpu_vmm_pool.cpp src/common/backend_plan.cpp src/common/backend_factory.cpp src/common/feature_gate.cpp @@ -695,7 +700,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 (LUCE_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(LUCE_GPU_BACKEND STREQUAL "hip") # rms_norm_hip.cu is needed by the HIP chunk-B graph path regardless of SM80_EQUIV. @@ -1379,12 +1384,6 @@ if(LUCE_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 luce_common) @@ -2077,10 +2076,8 @@ if(LUCE_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 test/test_gguf_mmap.cpp test/test_kv_quant.cpp @@ -2104,7 +2101,7 @@ if(LUCE_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 d95f32f4b..17f3e7d7f 100644 --- a/server/README.md +++ b/server/README.md @@ -194,7 +194,7 @@ server is byte-identical to local-inference mode. ```bash ./build/luce_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,11 +371,132 @@ 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-skip-park` | off | Keep target and decode draft resident while PFlash runs. | +| `--prefill-drafter ` | none | PFlash drafter GGUF (Qwen3.5-0.8B). | +| `--prefill-skip-park [auto\|on\|off]` | `auto` | Keep target and decode draft resident while PFlash runs. `auto` enables it when the drafter GGUF footprint fits measured free VRAM with a 25% margin, and also keeps the drafter loaded between requests (`--draft-residency auto`) when that footprint plus a 1.5 GiB target compute reserve fits; bare flag = `on` (explicit, no estimate, drafter residency unchanged). On CUDA builds with the VMM pool, cards under 32 GiB at max ctx > 64K always park (VMM fragmentation guard). A no-park window that runs out of device memory retries once parked, then parks 4 windows (doubling to 64 on repeats) before trying skip-park again; other failures do not retry. | | `--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`, +`PFLASH_SELECT_QUERY_TOKENS`), the drafter runs only its first fifteen +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 +still dispatch head dimension 128 only. + +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. 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 +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 +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. 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. + +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. + +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. 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 +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. + +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 +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, unless `--prefill-skip-park auto` found room to keep +it loaded (the estimate counts the kept sessions). +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 +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 @@ -520,9 +641,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' @@ -544,9 +665,9 @@ LUCE_KV_TYPE=q4_0 ./build/bench_laguna_ttft models/laguna-xs2-Q4_K_M.gguf '4096, # standalone test_laguna_daemon binary so it can run without luce_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 b4b2174a9..f0197cd86 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/luce_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 947064096..3bf8ef2fd 100644 --- a/server/docs/ENVIRONMENT.md +++ b/server/docs/ENVIRONMENT.md @@ -375,7 +375,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 171541778..a4cc3a38f 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, `LUCE_FP_USE_BSA=1`, +Qwen3.5-0.8B drafter, in-process daemon, `LUCE_FP_USE_BSA=1`, `LUCE_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 060f4a109..ee1979d6d 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 luce'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 luce'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 c936ab930..a075e735d 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 5e1f91006..f3cb88cf9 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("LUCE_TARGET", ROOT / "models" / "Qwen3.6-27B-Q4_K_M.gguf") DEFAULT_TARGET_DRAFT = env_path("LUCE_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 a8c4a459c..8649f7889 100644 --- a/server/scripts/quality_ab_simple.py +++ b/server/scripts/quality_ab_simple.py @@ -59,7 +59,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("LUCE_SERVER_BIN", "server/build/luce_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 07798a54f..6f448e56d 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("LUCE_SERVER_BIN", str(PROJECT_ROOT / "server/build/luce_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 68b2cbe7a..f579e8ea4 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 / "server/build/luce_server" for p, label in [ diff --git a/server/src/common/gguf_inspect.cpp b/server/src/common/gguf_inspect.cpp index fdc3298de..5a93787b6 100644 --- a/server/src/common/gguf_inspect.cpp +++ b/server/src/common/gguf_inspect.cpp @@ -1,5 +1,7 @@ #include "gguf_inspect.h" +#include "ggml.h" #include "gguf.h" +#include "kv_quant.h" #include #include @@ -421,4 +423,146 @@ GgufMetadata read_gguf_metadata(const std::string & path, return m; } +// ─── PFlash drafter footprint (skip-park estimator) ───────────────────── +// +// Worst-case per-token GPU state while the Qwen3.5 drafter scores a window +// (pflash/qwen35_drafter.cpp), on top of its weights: +// +// strict scorer (block-15 head, the default under strict selection): +// scoring sessions PFLASH_DRAFTER_SESSIONS (default 2) sessions kept +// for prefix reuse, each sized S + S/2 + 4096: +// KV of the full-attention layers in blocks 0..14 +// + f32 block-15 keys [head_dim, n_head_kv] +// act_in + act_out [n_embd, S] f32 → 2·n_embd·4 +// logits + mask [S, nq, n_head] + [S, nq] f32 → nq·(n_head+1)·4 +// (nq = the scorer query window, up to 512 tokens) +// probe raw scores [S] f32 ×2 +// legacy scorer (all-layer running max): full-depth KV + activations +// +// Both are bounded; the estimate takes the larger per-token cost. +// Fixed: SSM/conv state and snapshots, per-1024-token-ubatch transients, +// 8192-key score chunks, gallocr slack, and the sessions' 4096-token headroom. + +namespace { + +constexpr int64_t kLegacyLookahead = 8; +constexpr int kQwen35HeadBlocks = 15; // blocks 0..14 feed the head +// SSM/conv state + per-ubatch and key-chunk transients + gallocr slack. +constexpr int64_t kHybridFixedBytes = 384ll * 1024 * 1024; + +// The drafter cache wraps create_target_cache in ScopedKvTq3Off — TQ3 is never +// a drafter KV type, so suppress it while resolving the same env overrides. +struct ScopedKvTq3Suppress { + ScopedKvTq3Suppress() { + const char * raw = std::getenv("LUCE_KV_TQ3"); + had_ = raw != nullptr; + old_ = had_ ? raw : ""; +#if defined(_WIN32) + _putenv_s("LUCE_KV_TQ3", "0"); +#else + setenv("LUCE_KV_TQ3", "0", 1); +#endif + } + ~ScopedKvTq3Suppress() { +#if defined(_WIN32) + if (had_) _putenv_s("LUCE_KV_TQ3", old_.c_str()); + else _putenv_s("LUCE_KV_TQ3", ""); +#else + if (had_) setenv("LUCE_KV_TQ3", old_.c_str(), 1); + else unsetenv("LUCE_KV_TQ3"); +#endif + } + bool had_ = false; + std::string old_; +}; + +} // namespace + +bool inspect_drafter_footprint(const std::string & path, + int query_tokens, + int scoring_sessions, + SkipParkDrafterInfo & out) { + out = SkipParkDrafterInfo{}; + + struct stat st{}; + if (::stat(path.c_str(), &st) == 0) out.weights_bytes = int64_t(st.st_size); + + gguf_init_params gip{}; + gip.no_alloc = true; + gip.ctx = nullptr; + gguf_context * gctx = gguf_init_from_file(path.c_str(), gip); + if (!gctx) return false; + + std::string arch; + if (int64_t id = gguf_find_key(gctx, "general.architecture"); id >= 0) { + if (const char * v = gguf_get_val_str(gctx, id)) arch = v; + } + // The PFlash drafter is a Qwen3.5 hybrid (pflash/pflash_drafter.cpp); + // anything else is not a drafter the scorer can load. + if (arch != "qwen35") { gguf_free(gctx); return false; } + + auto get_i32 = [&](const char * suffix, int32_t & dst) -> bool { + const std::string key = arch + "." + suffix; + const int64_t id = gguf_find_key(gctx, key.c_str()); + if (id < 0) return false; + dst = int32_t(gguf_get_val_u32(gctx, id)); + return true; + }; + + // Qwen3.5-0.8B defaults, so a partial header still gets a bounded estimate. + int32_t n_layer = 24, n_embd = 1024, n_head = 8, n_head_kv = 2, + head_dim = 256, fai = 4, ctx_len = 262144; + get_i32("block_count", n_layer); + get_i32("embedding_length", n_embd); + get_i32("attention.head_count", n_head); + get_i32("attention.head_count_kv", n_head_kv); + get_i32("attention.key_length", head_dim); + get_i32("full_attention_interval", fai); + get_i32("context_length", ctx_len); + // Embedded NextN blocks inflate block_count on hybrid checkpoints. + int32_t nextn = 0; + get_i32("nextn_predict_layers", nextn); + gguf_free(gctx); + + if (n_layer <= 0 || n_embd <= 0 || n_head <= 0 || n_head_kv <= 0 || + head_dim <= 0 || fai <= 0 || ctx_len <= 0) { + return false; + } + if (nextn > 0 && nextn < n_layer) n_layer -= nextn; + + ggml_type kv_k = GGML_TYPE_Q4_0, kv_v = GGML_TYPE_Q4_0; + { + ScopedKvTq3Suppress tq3_off; + luce::resolve_kv_types(kv_k, kv_v); + } + + const int64_t nq = std::max(query_tokens, kLegacyLookahead); + const int64_t sessions = std::max(scoring_sessions, 1); + // PFLASH_DRAFTER_SESSIONS=0 scores from a scratch session sized S. + const int64_t capacity_x2 = scoring_sessions > 0 ? 3 : 2; // ×1.5 or ×1 + const int64_t session_per_token = + int64_t(kv_reservation_bytes_per_token( + std::min(n_layer, kQwen35HeadBlocks), fai, n_head_kv, + kv_k, head_dim, kv_v, head_dim)) + // KV + int64_t(head_dim) * n_head_kv * 4; // keys + const int64_t strict_per_token = + sessions * session_per_token * capacity_x2 / 2 + + int64_t(2) * n_embd * 4 + // act_in/out + nq * (n_head + 1) * 4 + // logits+mask + 2 * 4; // probe raw + const int64_t legacy_per_token = + int64_t(2) * n_embd * 4 + // act_in/out + int64_t(kv_reservation_bytes_per_token( + n_layer, fai, n_head_kv, kv_k, head_dim, kv_v, head_dim)) + + kLegacyLookahead * n_head * 4 + // logits + (kLegacyLookahead + 2) * 4; // mask+probe + + out.recognized = true; + out.runtime_bytes_per_token = std::max(strict_per_token, legacy_per_token); + out.fixed_bytes = kHybridFixedBytes + + (scoring_sessions > 0 ? sessions * 4096 * session_per_token : 0); + out.context_length = ctx_len; + return true; +} + } // namespace luce::common diff --git a/server/src/common/gguf_inspect.h b/server/src/common/gguf_inspect.h index 227075d46..0ab20b3a7 100644 --- a/server/src/common/gguf_inspect.h +++ b/server/src/common/gguf_inspect.h @@ -5,6 +5,8 @@ #pragma once +#include "placement/skip_park_guard.h" + #include #include @@ -73,4 +75,21 @@ struct GgufMetadata { GgufMetadata read_gguf_metadata(const std::string & path, bool compute_sha256); +// Read the PFlash drafter dims the skip-park estimator needs from the drafter +// GGUF header and derive its worst-case resident footprint. `weights_bytes` is +// the file size (upper bound on device weights); `runtime_bytes_per_token` and +// `fixed_bytes` mirror the buffers the Qwen3.5 scorer allocates in +// pflash/qwen35_drafter.cpp: `scoring_sessions` kept sessions (see +// pflash_scoring_sessions()) and a scorer query window of `query_tokens`. +// KV cache bytes honor the same LUCE_KV_* env resolution as the runtime +// cache, with TQ3 suppressed like the drafter's ScopedKvTq3Off. +// +// Returns false when the file can't be opened or is not a Qwen3.5 drafter — +// callers must then treat the footprint as unknown (auto skip-park resolves +// off). +bool inspect_drafter_footprint(const std::string & path, + int query_tokens, + int scoring_sessions, + SkipParkDrafterInfo & out); + } // namespace luce::common diff --git a/server/src/common/kvflash_pager.h b/server/src/common/kvflash_pager.h index 50e596359..71edc3f58 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 // (LUCE_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 cbce09dca..277cb29c9 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/model_backend.h b/server/src/common/model_backend.h index 1f6e6afc4..8af38b387 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 @@ -288,15 +290,51 @@ 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; + // 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; + // 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 + bool skip_park = false; // resolved --prefill-skip-park DraftResidencyAction residency_action = DraftResidencyAction::KeepLoaded; }; struct CompressResult { bool ok = false; + // Failed because a device allocation failed (drafter load, scorer + // buffers). Only this failure kind makes a skip-park window retry + // with parking; any other failure is returned as is. + bool out_of_memory = 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; + // 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; + + 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_drafter_ipc.cpp b/server/src/common/pflash_drafter_ipc.cpp index ab497f780..73b5c8ef8 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 luce::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 deb9adad3..f2dd82c26 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 luce::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 66e685aab..061d262cc 100644 --- a/server/src/common/pflash_drafter_ipc_daemon.cpp +++ b/server/src/common/pflash_drafter_ipc_daemon.cpp @@ -4,7 +4,7 @@ #include "luce.h" #include "dflash_draft_ipc.h" -#include "qwen3/qwen3_drafter.h" +#include "pflash/pflash_drafter.h" #include #include @@ -45,31 +45,31 @@ 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; + // 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, 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, + 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/common/pflash_types.h b/server/src/common/pflash_types.h new file mode 100644 index 000000000..711753056 --- /dev/null +++ b/server/src/common/pflash_types.h @@ -0,0 +1,27 @@ +#pragma once + +#include + +namespace luce::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 luce::common diff --git a/server/src/common/score_range.h b/server/src/common/score_range.h deleted file mode 100644 index 2d869c4f4..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 luce::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 luce::common diff --git a/server/src/deepseek4/deepseek4_backend.cpp b/server/src/deepseek4/deepseek4_backend.cpp index 38652ed66..67e15512d 100644 --- a/server/src/deepseek4/deepseek4_backend.cpp +++ b/server/src/deepseek4/deepseek4_backend.cpp @@ -17,6 +17,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" @@ -3600,52 +3601,109 @@ std::vector DeepSeek4Backend::compress_batch( } if (load_request == nullptr) return results; + const auto classify = [&requests, &valid_request]( + const std::vector & rs) { + auto outcome = SkipParkWindowOutcome::Ok; + for (size_t i = 0; i < requests.size(); ++i) { + if (!valid_request(requests[i])) continue; + if (i >= rs.size()) return SkipParkWindowOutcome::Failed; + if (rs[i].ok) continue; + if (rs[i].out_of_memory) return SkipParkWindowOutcome::OutOfMemory; + outcome = SkipParkWindowOutcome::Failed; + } + return outcome; + }; + return run_skip_park_window( + load_request->skip_park, skip_park_fallback_, + [&](bool park_window) { + return run_compress_window(requests, *load_request, park_window); + }, + classify, [this]() { release_pflash_drafter(); }, + "[deepseek4-pflash]"); +} + +std::vector DeepSeek4Backend::run_compress_window( + const std::vector & requests, + const CompressRequest & load_request, + bool park_window) { + std::vector results(requests.size()); + const auto valid_request = [](const CompressRequest & request) { + return !request.input_ids.empty() && !request.drafter_path.empty() && + std::isfinite(request.keep_ratio) && + request.keep_ratio >= 0.0f && request.keep_ratio <= 1.0f; + }; + // Parking releases target/cache buffers, including the expert backend. // Drain their queued work before releasing any of those dependencies. if (backend_) ggml_backend_synchronize(backend_); if (spec_backend_) ggml_backend_synchronize(spec_backend_); if (expert_backend_) ggml_backend_synchronize(expert_backend_); const bool was_parked = parked_; - if (!load_request->skip_park && !parked_ && + if (park_window && !parked_ && !park(ParkTarget::TargetModel)) { return results; } if (pflash_drafter_loaded_ && - (pflash_drafter_path_ != load_request->drafter_path || - pflash_drafter_gpu_ != load_request->drafter_gpu)) { + (pflash_drafter_path_ != load_request.drafter_path || + pflash_drafter_gpu_ != load_request.drafter_gpu)) { release_pflash_drafter(); } if (!pflash_drafter_loaded_) { - if (!load_drafter(load_request->drafter_path, 999, - load_request->drafter_gpu, + if (!load_drafter(load_request.drafter_path, 999, + load_request.drafter_gpu, pflash_drafter_ctx_)) { std::fprintf(stderr, "[deepseek4-pflash] load failed: %s\n", luce_last_error()); + const bool oom = luce::common::last_error_is_oom(); + for (size_t index = 0; index < requests.size(); ++index) { + if (valid_request(requests[index])) results[index].out_of_memory = oom; + } release_pflash_drafter(); - if (!load_request->skip_park && !was_parked) { + if (park_window && !was_parked) { unpark(ParkTarget::TargetModel); } return results; } pflash_drafter_loaded_ = true; - pflash_drafter_path_ = load_request->drafter_path; - pflash_drafter_gpu_ = load_request->drafter_gpu; + pflash_drafter_path_ = load_request.drafter_path; + pflash_drafter_gpu_ = load_request.drafter_gpu; } for (size_t index = 0; index < requests.size(); ++index) { 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, + request.query_suffix_candidates, request.history_query_spans, + request.turn_query_span); result.ok = !result.compressed_ids.empty(); + result.out_of_memory = !result.ok && luce::common::last_error_is_oom(); + 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; + for (const auto & candidate : pflash_last_candidate_lifts()) { + result.candidate_lifts.push_back({candidate.span, candidate.lift}); + } + } } - if (load_request->residency_action == - DraftResidencyAction::ReleaseAfterUse) { + // A recent out-of-memory window overrides KeepLoaded: VRAM is tight. + if (load_request.residency_action == + DraftResidencyAction::ReleaseAfterUse || + skip_park_fallback_.memory_tight()) { release_pflash_drafter(); } - if (!load_request->skip_park && !was_parked && + if (park_window && !was_parked && !unpark(ParkTarget::TargetModel)) { std::fill(results.begin(), results.end(), CompressResult{}); } diff --git a/server/src/deepseek4/deepseek4_backend.h b/server/src/deepseek4/deepseek4_backend.h index 4e45ca153..3c5dd84e7 100644 --- a/server/src/deepseek4/deepseek4_backend.h +++ b/server/src/deepseek4/deepseek4_backend.h @@ -15,11 +15,12 @@ #include "../common/moe_hybrid_stream.h" #include "deepseek4_internal.h" #include "deepseek4_dspark.h" +#include "pflash/pflash_drafter.h" +#include "placement/skip_park_guard.h" #include "deepseek4_vision.h" #include "deepseek4_image_prompt.h" #include "deepseek4_image_assembly.h" #include "deepseek4_image_admission.h" -#include "qwen3/qwen3_drafter.h" #include "deepseek4_seq_engine.h" #include "ggml.h" @@ -163,6 +164,9 @@ class DeepSeek4Backend : public ModelBackend { bool pflash_drafter_loaded_ = false; std::string pflash_drafter_path_; int pflash_drafter_gpu_ = -1; + // Skip-park fail-safe: parks a few windows after an out-of-memory + // no-park window recovered with parking (placement/skip_park_guard.h). + SkipParkFallback skip_park_fallback_; // Once a long prompt selects the fragmentation-safe prefill shape, retain // it for later requests so the HIP arenas never switch back under load. int hybrid_prefill_chunk_cap_ = 0; @@ -170,6 +174,14 @@ class DeepSeek4Backend : public ModelBackend { bool load_spec_drafter(); void release_spec_drafter(bool mark_parked); void release_pflash_drafter(); + // One compression window (sync → park → load drafter → score → restore) + // with the park step optional. compress_batch runs it through + // run_skip_park_window, which retries parked after an out-of-memory + // no-park attempt. + std::vector run_compress_window( + const std::vector & requests, + const CompressRequest & load_request, + bool park_window); void keep_spec_feature_tail(std::vector & features, size_t max_rows) const; // True when a wide prefill path returns per-token DSpark features and the diff --git a/server/src/errors.cpp b/server/src/errors.cpp index dca47dfc6..704dca8c2 100644 --- a/server/src/errors.cpp +++ b/server/src/errors.cpp @@ -12,11 +12,24 @@ namespace luce::common { namespace { std::mutex g_err_mu; std::string g_last_error; +bool g_last_error_oom = false; } void set_last_error(std::string msg) { std::lock_guard lk(g_err_mu); g_last_error = std::move(msg); + g_last_error_oom = false; +} + +void set_last_oom_error(std::string msg) { + std::lock_guard lk(g_err_mu); + g_last_error = std::move(msg); + g_last_error_oom = true; +} + +bool last_error_is_oom() { + std::lock_guard lk(g_err_mu); + return g_last_error_oom; } } // namespace luce::common diff --git a/server/src/flashprefill.h b/server/src/flashprefill.h index 13914fef9..e76f34b98 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/flashprefill_q8.cpp b/server/src/flashprefill_q8.cpp index 2a7e386e1..c3e974b73 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/gemma4/gemma4_backend.cpp b/server/src/gemma4/gemma4_backend.cpp index b1bafe922..f61526538 100644 --- a/server/src/gemma4/gemma4_backend.cpp +++ b/server/src/gemma4/gemma4_backend.cpp @@ -6,7 +6,7 @@ #include "gemma4_backend.h" #include "luce.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 3bf873dab..67ce888ea 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 6426fde23..06be90f92 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 "luce.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 41584325b..225b36c54 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/internal.h b/server/src/internal.h index 3d8387008..4be2fd6a1 100644 --- a/server/src/internal.h +++ b/server/src/internal.h @@ -34,6 +34,11 @@ struct MoeHybridStorage; // Single source of truth for error reporting. // All loaders / graph builders push into this via set_last_error(...). void set_last_error(std::string msg); +// Same, for a failed device allocation (buffer, graph allocator, cache). The +// PFlash skip-park fallback retries a window with parking only after this kind +// of failure; any later set_last_error() clears the flag. +void set_last_oom_error(std::string msg); +bool last_error_is_oom(); // ─── Target weights (Qwen3.5-27B, qwen35 hybrid, Q4_K_M in ggml context) ── // diff --git a/server/src/laguna/laguna_backend.cpp b/server/src/laguna/laguna_backend.cpp index c71bb779e..9cc97bc39 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 "luce.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 a4fba713e..addb2db99 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 bb1538a26..21c182ae4 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 "luce.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 5a3200820..c5bf40b85 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 99% rename from server/src/qwen3/anchor_scan.cpp rename to server/src/pflash/anchor_scan.cpp index 89c3e2ade..9265fe0c3 100644 --- a/server/src/qwen3/anchor_scan.cpp +++ b/server/src/pflash/anchor_scan.cpp @@ -5,7 +5,7 @@ #include #include -namespace luce::qwen3 { +namespace luce::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 luce::qwen3 +} // namespace luce::pflash diff --git a/server/src/qwen3/anchor_scan.h b/server/src/pflash/anchor_scan.h similarity index 96% rename from server/src/qwen3/anchor_scan.h rename to server/src/pflash/anchor_scan.h index ecace4395..bf7bddb45 100644 --- a/server/src/qwen3/anchor_scan.h +++ b/server/src/pflash/anchor_scan.h @@ -6,7 +6,7 @@ #include #include -namespace luce::qwen3 { +namespace luce::pflash { struct AnchorScanCfg { int chunk_size; @@ -39,4 +39,4 @@ void scan_and_force_transitive( std::vector& forced ); -} // namespace luce::qwen3 +} // namespace luce::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 f817e5929..6574ae812 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 luce::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 82% rename from server/src/qwen3/qwen3_kvflash_scorer.h rename to server/src/pflash/kvflash_drafter_scorer.h index 7c9737170..8de23df72 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 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). #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/pflash/pflash_compress.cpp b/server/src/pflash/pflash_compress.cpp new file mode 100644 index 000000000..f2a5fb603 --- /dev/null +++ b/server/src/pflash/pflash_compress.cpp @@ -0,0 +1,379 @@ +// PFlash scoring pipeline glue. See pflash_compress.h. + +#include "pflash_compress.h" + +#include "pflash_selection.h" +#include "internal.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace luce::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("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 != + luce::pflash::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,\"top_k\":%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->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 == + luce::pflash::PFlashSelectionMode::Legacy) { + std::fputs(",\"stop_reason\":null,\"retained_mass\":null", file); + } else { + std::fprintf(file, + ",\"stop_reason\":\"%s\",\"retained_mass\":%.17g", + luce::pflash::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); +} + +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; +} + +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 = {}; + g_last_candidate_lifts.clear(); +} + +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 luce::pflash::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 = + luce::pflash::pflash_chunk_is_structurally_required( + begin, end, query_begin, query_end, input_tokens, + 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); + } + // 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 luce::pflash::PFlashSelectionPolicy policy{selector_budget, config.top_p, + /*skip_oversized=*/ segments != nullptr, + 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); + 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", + luce::pflash::pflash_selection_mode_name(config.mode), + selector_budget, + luce::pflash::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); + 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(), + 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, + "[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", + 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); + 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" : luce::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 == luce::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, + mandatory_mask, output, &strict_fields); + } + return output; +} + +} // namespace luce::common diff --git a/server/src/pflash/pflash_compress.h b/server/src/pflash/pflash_compress.h new file mode 100644 index 000000000..7bd37e5cf --- /dev/null +++ b/server/src/pflash/pflash_compress.h @@ -0,0 +1,168 @@ +// 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 (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 luce::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; + luce::pflash::PFlashSelectionMode selector_mode = + luce::pflash::PFlashSelectionMode::Legacy; + luce::pflash::PFlashQueryParser query_parser = + luce::pflash::PFlashQueryParser::SemanticUser; + int token_budget = 0; + luce::pflash::PFlashSelectionStop stop = + luce::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; + // Rank-mode ceiling: the K that applied, 0 outside top_k mode. + int top_k = 0; +}; + +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); + +// 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(); + +// 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); + +// 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, + float keep_ratio, + int n_lookahead, + int score_query_end, + int pool_kernel, + const luce::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 luce::common diff --git a/server/src/pflash/pflash_drafter.cpp b/server/src/pflash/pflash_drafter.cpp new file mode 100644 index 000000000..bae720e4f --- /dev/null +++ b/server/src/pflash/pflash_drafter.cpp @@ -0,0 +1,185 @@ +// 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, +// 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 "pflash_compress.h" +#include "common/dspark_head.h" +#include "internal.h" + +#include "ggml.h" +#include "ggml-backend.h" + +#include +#include +#include + +namespace luce::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, + bool query_suffix_candidates, + const std::vector & history_queries, + PFlashTokenSpan turn_query) { + pflash_clear_kept_spans(); + if (!ctx.loaded) { + set_last_error("drafter not loaded"); + return {}; + } + + luce::pflash::PFlashSelectionConfig experiment; + std::string experiment_error; + if (!luce::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; + 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 (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( + "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 (!luce::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 " + "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, (int) experiment.query_suffix_candidates, + 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 luce::common diff --git a/server/src/pflash/pflash_drafter.h b/server/src/pflash/pflash_drafter.h new file mode 100644 index 000000000..3398cd1bc --- /dev/null +++ b/server/src/pflash/pflash_drafter.h @@ -0,0 +1,102 @@ +// In-process PFlash drafter for speculative prefill. +// +// "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 +// 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 +// token IDs (drafter vocab). + +#pragma once + +#include "common/pflash_types.h" + +#include +#include +#include +#include + +struct ggml_backend; +typedef struct ggml_backend * ggml_backend_t; + +namespace luce::common { + +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 today). +// Creates a fresh GPU backend if `backend` is null. Otherwise uses the +// caller-provided backend (so the drafter shares the daemon's allocator). +// +// `gpu_layers` is accepted for API compat but ignored — every layer goes on +// the GPU since the drafter weights are only ~1.5 GB. +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); + +void free_drafter(DrafterContext & ctx); + +// Scoring sessions the drafter keeps for prefix reuse while loaded +// (PFLASH_DRAFTER_SESSIONS, default 2; 0 scores every prompt from scratch). +// Each holds KV sized for its prompt plus headroom, so it counts toward the +// drafter's resident footprint (skip-park estimate, common/gguf_inspect.h). +int pflash_scoring_sessions(); + +// Free only model weights, keeping the backend alive for reuse. +// Avoids repeated ggml backend create/destroy during daemon reuse. +void free_drafter_weights(DrafterContext & ctx); + +// 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 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 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 +// 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( + DrafterContext & ctx, + const std::vector & ids, + float keep_ratio, + int chunk_size = 32, + int n_lookahead = 8, + int pool_kernel = 13, + int score_query_end = -1, + const std::vector & + required_instruction_spans = {}, + bool query_suffix_candidates = false, + const std::vector & history_queries = {}, + PFlashTokenSpan turn_query = {-1, -1}); + +} // namespace luce::common diff --git a/server/src/pflash/pflash_selection.cpp b/server/src/pflash/pflash_selection.cpp new file mode 100644 index 000000000..5b975b6ce --- /dev/null +++ b/server/src/pflash/pflash_selection.cpp @@ -0,0 +1,586 @@ +#include "pflash_selection.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace luce::pflash { + +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 * 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"; +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(kTopKEnv) != 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, + 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 = query_suffix_structural && + 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() > luce::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]"); + } + 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()); + 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; + } + } + + 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) { + result.stop = PFlashSelectionStop::BudgetReached; + if (policy.skip_oversized) continue; + break; + } + + 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 + : 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"; + case PFlashSelectionMode::TopK: return "top_k"; + } + return "unknown"; +} + +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"; + 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 * 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); + const char * split_raw = std::getenv(kSplitEnv); + + 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; + 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 if (std::strcmp(mode_raw, "top_k") == 0) { + config.mode = PFlashSelectionMode::TopK; + } else { + error = std::string(kModeEnv) + + " must be budget_only, top_p or top_k"; + 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; + } + + 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; +} + +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 luce::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 (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)"; + 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 luce::pflash diff --git a/server/src/pflash/pflash_selection.h b/server/src/pflash/pflash_selection.h new file mode 100644 index 000000000..20cf95557 --- /dev/null +++ b/server/src/pflash/pflash_selection.h @@ -0,0 +1,176 @@ +#pragma once + +#include "common/pflash_types.h" + +#include +#include +#include + +namespace luce::pflash { + +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 { + SemanticUser, + ArbitraryTail, +}; + +enum class PFlashSelectionStop { + TopPReached, + TopKReached, + 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; + // 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 { + bool ok = false; + std::vector ordinals; + int retained_tokens = 0; + double retained_mass = 0.0; + PFlashSelectionStop stop = PFlashSelectionStop::InvalidInput; + 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, + int query_begin, + int query_end, + int input_tokens, + const std::vector & + required_instruction_spans = {}, + bool query_suffix_structural = true) 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; + // 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; + double top_p = 0.95; + int top_k = 0; + 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; + // 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; + // 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; + // 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 +// 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 luce::pflash diff --git a/server/src/pflash/qwen35_drafter.cpp b/server/src/pflash/qwen35_drafter.cpp new file mode 100644 index 000000000..906152726 --- /dev/null +++ b/server/src/pflash/qwen35_drafter.cpp @@ -0,0 +1,1421 @@ +// 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; pflash_drafter.cpp dispatches into +// qwen35_drafter_score_and_compress. + +#include "qwen35_drafter.h" + +#include "pflash_drafter.h" +#include "pflash_compress.h" +#include "pflash_selection.h" +#include "common/gguf_inspect.h" +#include "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 luce::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 LUCE_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, "LUCE_KV_TQ3"); + had_ = raw != nullptr; + old_ = had_ ? raw : ""; + free(raw); + _putenv_s("LUCE_KV_TQ3", "0"); +#else + const char * raw = std::getenv("LUCE_KV_TQ3"); + had_ = raw != nullptr; + old_ = had_ ? raw : ""; + setenv("LUCE_KV_TQ3", "0", 1); +#endif + } + ~ScopedKvTq3Off() { +#if defined(_WIN32) + // _putenv_s with empty value removes the variable on MSVCRT. + _putenv_s("LUCE_KV_TQ3", had_ ? old_.c_str() : ""); +#else + if (had_) setenv("LUCE_KV_TQ3", old_.c_str(), 1); + else unsetenv("LUCE_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 luce::pflash::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_oom_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_oom_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); + if (st == GGML_STATUS_ALLOC_FAILED) { + set_last_oom_error("qwen35 drafter graph compute out of memory"); + } else { + 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_oom_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); + if (st == GGML_STATUS_ALLOC_FAILED) { + set_last_oom_error("qwen35 score graph compute out of memory"); + } else { + 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("LUCE_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("LUCE_COMPRESS_HEAD_CHUNKS", 8); + const int t_raw = env_int("LUCE_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("LUCE_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("LUCE_COMPRESS_ANCHOR_RADIUS", -1), + env_int("LUCE_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("LUCE_COMPRESS_REPEAT_MIN", 4); + const int repeat_max = env_int("LUCE_COMPRESS_REPEAT_MAX", 32); + const int repeat_limit = env_int("LUCE_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; +} + +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_windows.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, + 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() || + 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) { + 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 + ? 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; + best_shared = shared; + } + } + if (best) { + if (best_restore && !restore_ssm_state(best->cache, w.backend)) { + resume = 0; + } else { + shared_prefix = best_shared; + 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_oom_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_windows.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 +// 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 luce::pflash::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]; + const bool use_probe = st.probe_loaded && + experiment.segmentation != luce::pflash::PFlashSegmentation::Fixed; + + 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, shared_prefix, + 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) { + 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, 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); + session->ids.clear(); + set_last_oom_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); + }; + auto fail = [&](const char * message) -> std::vector { + cleanup(); + session->ids.clear(); + set_last_error(message); + return {}; + }; + auto fail_oom = [&](const char * message) -> std::vector { + cleanup(); + session->ids.clear(); + set_last_oom_error(message); + return {}; + }; + + { + const int batch = 2048; + std::vector emb((size_t)hidden * batch); + 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); + 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); + 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 - resume) * 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 - resume) * act_out->nb[1]); + if (ggml_nelements(out) != ggml_nelements(dst)) { + 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); + return fail_oom("qwen35 drafter graph allocation failed"); + } + 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); + return status == GGML_STATUS_ALLOC_FAILED + ? fail_oom("qwen35 drafter graph compute out of memory") + : fail("qwen35 drafter graph compute failed"); + } + start += n; + if (!is_attn && start == checkpoint && checkpoint > resume) { + snapshot_layer(il); + } + } + std::swap(act_in, act_out); + } + ggml_gallocr_free(alloc); + auto t1 = std::chrono::steady_clock::now(); + + // 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)); + 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; + 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; + 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_oom("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 ggml_status key_status = ggml_gallocr_alloc_graph(kalloc, kgf) + ? ggml_backend_graph_compute(w.backend, kgf) : GGML_STATUS_ALLOC_FAILED; + const bool key_ok = key_status == 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 key_status == GGML_STATUS_ALLOC_FAILED + ? fail_oom("qwen35 key graph allocation failed") + : 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; + 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, 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; + 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_oom_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_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_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 ggml_status score_status = ggml_gallocr_alloc_graph(salloc, sgf) + ? ggml_backend_graph_compute(w.backend, sgf) : GGML_STATUS_ALLOC_FAILED; + const bool ok = score_status == 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) { + if (score_status == GGML_STATUS_ALLOC_FAILED) { + set_last_oom_error("qwen35 score graph allocation failed"); + } else { + 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; + } + 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; + 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 (%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(), windows.size(), + 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; + 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 < 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 = + experiment.segmentation == luce::pflash::PFlashSegmentation::Probe; + if (boundaries_in_context >= 4 || forced_probe) { + segments = luce::pflash::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 == luce::pflash::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 luce::pflash::PFlashSelectionConfig & experiment, + const std::vector & required_instruction_spans) { + if (!ctx.state) { + set_last_error("qwen35 drafter state missing"); + return {}; + } + 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 == 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) + // 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, suffix_kept, + 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, + suffix_kept, 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, suffix_kept, 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) { + 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"); + return {}; + } + return qwen35_score_and_compress(st->weights, ids, keep_ratio, chunk_size, + n_lookahead, pool_kernel, score_query_end, + suffix_kept, + required_instruction_spans); +} + +int pflash_scoring_sessions() { + return scoring_session_limit(); +} + +} // namespace luce::common diff --git a/server/src/pflash/qwen35_drafter.h b/server/src/pflash/qwen35_drafter.h new file mode 100644 index 000000000..8272fa656 --- /dev/null +++ b/server/src/pflash/qwen35_drafter.h @@ -0,0 +1,142 @@ +// 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 +// scoring head and the segment probe; qwen35_drafter.cpp runs the two +// scorers. pflash_drafter.cpp owns the public entry points. + +#pragma once + +#include "pflash_drafter.h" +#include "pflash_selection.h" +#include "common/pflash_types.h" +#include "internal.h" + +#include "ggml.h" +#include "ggml-backend.h" + +#include +#include +#include +#include + +namespace luce::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; an optional trained head replaces those two +// 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 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; +}; + +void free_qwen35_scoring_session(Qwen35ScoringSession & session); + +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; + // 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. +bool qwen35_head_block_available(const TargetWeights & w, std::string & error); +bool load_qwen35_drafter(const std::string & gguf_path, 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 luce::pflash::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 luce::pflash::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 luce::pflash::PFlashSelectionConfig & experiment, + const std::vector & required_instruction_spans); + +} // namespace luce::common diff --git a/server/src/pflash/qwen35_loader.cpp b/server/src/pflash/qwen35_loader.cpp new file mode 100644 index 000000000..c150a8b16 --- /dev/null +++ b/server/src/pflash/qwen35_loader.cpp @@ -0,0 +1,384 @@ +// Qwen3.5-0.8B drafter loading: the drafter GGUF, the optional trained +// block-15 scoring head and the optional segment probe. +// +// The scorer is built on the Qwen3.5 target weights +// (load_target_gguf_partial). + +#include "qwen35_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 luce::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. +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) { + fail("scoring head buffer allocation failed"); + set_last_oom_error("scoring head buffer allocation failed"); + return false; + } + 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) { + fail("segment probe buffer allocation failed"); + set_last_oom_error("segment probe buffer allocation failed"); + return false; + } + 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, + 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.state = st; + out.loaded = true; + std::fprintf(stderr, + "[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", + 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.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); + delete st; + ctx.state = nullptr; +} + +} // namespace luce::common diff --git a/server/src/placement/draft_residency.h b/server/src/placement/draft_residency.h index fc22b33ab..1b4695a70 100644 --- a/server/src/placement/draft_residency.h +++ b/server/src/placement/draft_residency.h @@ -31,6 +31,12 @@ struct DraftResidencyContext { DraftResidencyUse use = DraftResidencyUse::PFlashCompress; bool low_vram_hint = false; bool has_decode_draft = false; + // True when the startup skip-park probe proved the target, decode draft + // and pflash drafter fit co-resident with margin. Under Auto it upgrades + // PFlashCompress to KeepLoaded — the drafter stays loaded between + // requests, removing its per-request reload. Ignored by other uses and + // by explicit Persistent/RequestScoped policies. + bool ample_vram = false; }; inline const char * draft_residency_policy_name(DraftResidencyPolicy policy) { @@ -71,8 +77,14 @@ inline DraftResidencyAction resolve_draft_residency_action( switch (ctx.use) { case DraftResidencyUse::PFlashCompress: - // Auto releases the pflash drafter after scoring: resident drafter starves target prefill on 24GB cards; lazy reload costs ~2s. - return DraftResidencyAction::ReleaseAfterUse; + // Auto releases the pflash drafter after scoring on constrained + // cards (resident drafter starves target prefill on 24GB; lazy + // reload costs ~2s). With proven ample VRAM the same startup probe + // that enabled skip-park also keeps the drafter resident — its + // footprint is already accounted and the reload is pure overhead. + return ctx.ample_vram + ? DraftResidencyAction::KeepLoaded + : DraftResidencyAction::ReleaseAfterUse; case DraftResidencyUse::DFlashDecode: // DFlash draft is latency-sensitive; keep it resident unless the // operator explicitly opted into the low-VRAM/request-scoped path. diff --git a/server/src/placement/gpu_vmm_pool.cpp b/server/src/placement/gpu_vmm_pool.cpp new file mode 100644 index 000000000..e2a2456c7 --- /dev/null +++ b/server/src/placement/gpu_vmm_pool.cpp @@ -0,0 +1,32 @@ +#include "placement/gpu_vmm_pool.h" + +#include "ggml-backend.h" + +#include + +namespace luce::common { + +bool gpu_backend_uses_vmm_pool() { + for (size_t i = 0; i < ggml_backend_reg_count(); ++i) { + ggml_backend_reg_t reg = ggml_backend_reg_get(i); + bool has_gpu = false; + for (size_t d = 0; d < ggml_backend_reg_dev_count(reg); ++d) { + if (ggml_backend_dev_type(ggml_backend_reg_dev_get(reg, d)) == + GGML_BACKEND_DEVICE_TYPE_GPU) { + has_gpu = true; + break; + } + } + if (!has_gpu) continue; + auto get_features = (ggml_backend_get_features_t) + ggml_backend_reg_get_proc_address(reg, "ggml_backend_get_features"); + if (!get_features) return true; + for (ggml_backend_feature * f = get_features(reg); f && f->name; ++f) { + if (std::strcmp(f->name, "NO_VMM") == 0) return false; + } + return true; + } + return true; +} + +} // namespace luce::common diff --git a/server/src/placement/gpu_vmm_pool.h b/server/src/placement/gpu_vmm_pool.h new file mode 100644 index 000000000..62893d96c --- /dev/null +++ b/server/src/placement/gpu_vmm_pool.h @@ -0,0 +1,15 @@ +// Whether the loaded ggml GPU backend allocates compute scratch from its VMM +// pool — the pool whose cuMemSetAccess failed under skip-park fragmentation +// (placement/skip_park_guard.h). + +#pragma once + +namespace luce::common { + +// Reads the GPU backend registry's feature list: ggml-cuda (also compiled as +// ggml-hip) reports NO_VMM when built without the VMM pool, which is the HIP +// default (GGML_HIP_NO_VMM=ON). Returns true when no GPU backend can be +// inspected, so an unknown build keeps the guard. +bool gpu_backend_uses_vmm_pool(); + +} // namespace luce::common diff --git a/server/src/placement/skip_park_guard.h b/server/src/placement/skip_park_guard.h index 00249f925..1d1678380 100644 --- a/server/src/placement/skip_park_guard.h +++ b/server/src/placement/skip_park_guard.h @@ -1,12 +1,264 @@ -// Footprint-aware guard: downgrade --prefill-skip-park on <32GB GPUs at max_ctx>65536. +// Skip-park guard + startup policy resolution for PFlash compression. +// +// During PFlash compression the server normally parks the resident target (and +// decode draft), loads the scoring drafter, compresses, frees the drafter and +// reloads everything. When VRAM is ample the park/unpark round-trip is pure +// overhead: the drafter can coexist with the resident models for the whole +// compression window ("skip park"). +// +// This header owns the policy, not the mechanics: +// - SkipParkMode: the requested policy (--prefill-skip-park auto|on|off) +// - skip_park_allowed: VMM crash guard (<32GiB with ctx>64K, VMM pool only) +// - resolve_skip_park: the single startup decision (auto estimate vs. free +// VRAM with margin, with explicit on/off precedence), plus whether auto +// may also keep the drafter loaded between requests +// - SkipParkFallback / run_skip_park_window: the runtime fail-safe that +// retries an out-of-memory no-park window with parking +// +// The per-drafter footprint inputs come from inspect_drafter_footprint() +// (common/gguf_inspect.h), which reads dims from the drafter GGUF header and +// mirrors the buffers the Qwen3.5 scorer allocates (pflash/qwen35_drafter.cpp). + #pragma once + +#include #include +#include +#include +#include +#include namespace luce::common { -// Returns false only when dual-residency is unsafe (VMM VA-fragmentation risk). -inline bool skip_park_allowed(bool requested, size_t total_vram_bytes, int max_ctx) { - return requested && (total_vram_bytes >= 32ull*1024*1024*1024 || max_ctx <= 65536); +// Requested policy for --prefill-skip-park. Bare `--prefill-skip-park` (no +// value) keeps the historical boolean meaning = On. Auto is the default: +// the startup estimate decides; On/Off are explicit operator overrides that +// bypass the estimate (the VMM guard still applies to On). +enum class SkipParkMode { Auto, On, Off }; + +inline const char * skip_park_mode_name(SkipParkMode mode) { + switch (mode) { + case SkipParkMode::Auto: return "auto"; + case SkipParkMode::On: return "on"; + case SkipParkMode::Off: return "off"; + } + return "auto"; +} + +inline bool parse_skip_park_mode(const std::string & value, SkipParkMode & out) { + if (value == "auto") { out = SkipParkMode::Auto; return true; } + if (value == "on") { out = SkipParkMode::On; return true; } + if (value == "off") { out = SkipParkMode::Off; return true; } + return false; +} + +// VMM crash guard (1c562eb4d): on a 24 GB CUDA card at max_ctx=131072, +// dual-resident target+drafter fragmented the VMM pool's virtual address space +// and its cuMemSetAccess failed. The failure lives in ggml's VMM pool, so the +// guard only applies when the GPU backend uses it (`vmm_pool`); HIP builds +// default to GGML_HIP_NO_VMM and allocate from the legacy pool, where the +// failure mode does not exist and the auto estimate alone decides. Applies to +// Auto and On alike — a forced-on that reliably crashes is worse than +// ignoring the request. +inline bool skip_park_allowed(bool requested, size_t total_vram_bytes, int max_ctx, + bool vmm_pool = true) { + return requested && + (!vmm_pool || total_vram_bytes >= 32ull*1024*1024*1024 || + max_ctx <= 65536); +} + +// Worst-case incremental VRAM while the PFlash drafter scores a window, on top +// of the resident target + decode draft it skips parking for. Populated by +// inspect_drafter_footprint() from the drafter GGUF header. +struct SkipParkDrafterInfo { + bool recognized = false; // false → auto resolves off + int64_t weights_bytes = 0; // GGUF file size (upper bound) + int64_t runtime_bytes_per_token = 0; // sessions + activations + score buffers + int64_t fixed_bytes = 0; // SSM/conv state, gallocr, chunk transients + int context_length = 0; // drafter native ctx — caps the window +}; + +// Headroom multiplier applied to the whole footprint. Covers fragmentation, +// gallocr slack and per-arch details the estimator intentionally ignores. +constexpr int64_t kSkipParkSafetyMarginPercent = 25; + +// Room the target's own prefill/decode compute buffers need beyond what is +// resident at startup — the same 1.5 GiB the KV pool budgets reserve for +// "runtime graph buffers" (qwen35/laguna/gemma4 make_kvflash_budget). Only +// the keep-loaded decision uses it: a drafter kept resident between requests +// sits beside the target's prefill, not just beside its weights. +constexpr int64_t kTargetComputeReserveBytes = 1536ll * 1024 * 1024; + +inline int64_t skip_park_with_margin(int64_t raw) { + return raw + raw * kSkipParkSafetyMarginPercent / 100; +} + +inline int64_t skip_park_raw_bytes(const SkipParkDrafterInfo & info, + int64_t window_tokens) { + return info.weights_bytes + info.fixed_bytes + + info.runtime_bytes_per_token * window_tokens; +} + +inline int64_t skip_park_required_bytes(const SkipParkDrafterInfo & info, + int64_t window_tokens) { + return skip_park_with_margin(skip_park_raw_bytes(info, window_tokens)); +} + +// Footprint to keep the drafter loaded between requests. Everything the +// drafter allocated for a window may still be resident when the target +// prefills the next prompt (its weights, the scoring sessions it keeps for +// prefix reuse, its backend's compute pool), so the whole window estimate is +// counted beside the target's compute reserve. +inline int64_t skip_park_keep_loaded_bytes(const SkipParkDrafterInfo & info, + int64_t window_tokens) { + return skip_park_with_margin(skip_park_raw_bytes(info, window_tokens) + + kTargetComputeReserveBytes); +} + +struct SkipParkDecision { + bool enabled = false; + // Auto only: the estimate also leaves room to keep the drafter loaded + // between requests (draft-residency auto → KeepLoaded). Never set by an + // explicit `on`, which runs no estimate. + bool keep_drafter_loaded = false; + std::string reason; // short human-readable log token + int64_t required_bytes = 0; // footprint incl. margin (auto only) + int64_t keep_loaded_bytes = 0; // keep-loaded footprint incl. margin + int64_t window_tokens = 0; // min(max_ctx, drafter ctx) +}; + +// The one authoritative decision, resolved once at startup: +// off → never skip park +// on → skip park unless the VMM guard blocks it; the drafter keeps the +// residency the draft-residency policy gives it +// auto → skip park iff the guard passes AND the estimated footprint fits +// measured free VRAM with margin; additionally keep the drafter +// loaded between requests iff the footprint plus the target's +// compute reserve fits too +inline SkipParkDecision resolve_skip_park( + SkipParkMode mode, bool drafter_configured, + const SkipParkDrafterInfo & info, + int64_t free_vram_bytes, int64_t total_vram_bytes, int64_t max_ctx, + bool vmm_pool = true) { + SkipParkDecision d; + if (mode == SkipParkMode::Off) { + d.reason = "off (explicit)"; + return d; + } + if (!drafter_configured) { + d.reason = "off (no local pflash drafter)"; + return d; + } + if (!skip_park_allowed(true, size_t(std::max(total_vram_bytes, 0)), + int(std::min(max_ctx, INT32_MAX)), + vmm_pool)) { + d.reason = "off (guard: VMM pool, <32GiB VRAM with ctx>64K)"; + return d; + } + if (mode == SkipParkMode::On) { + d.enabled = true; + d.reason = "on (explicit)"; + return d; + } + if (!info.recognized || info.context_length <= 0 || + info.runtime_bytes_per_token <= 0 || free_vram_bytes < 0) { + d.reason = "off (auto: drafter footprint unknown)"; + return d; + } + d.window_tokens = std::min(max_ctx, info.context_length); + d.required_bytes = skip_park_required_bytes(info, d.window_tokens); + d.keep_loaded_bytes = skip_park_keep_loaded_bytes(info, d.window_tokens); + d.enabled = d.required_bytes <= free_vram_bytes; + d.keep_drafter_loaded = d.enabled && d.keep_loaded_bytes <= free_vram_bytes; + d.reason = !d.enabled ? "off (auto: estimate exceeds free VRAM)" + : d.keep_drafter_loaded + ? "on (auto: estimate fits free VRAM, drafter kept loaded)" + : "on (auto: estimate fits free VRAM, drafter released)"; + return d; +} + +// ── Runtime fail-safe ────────────────────────────────────────────────────── +// +// The startup estimate can still be wrong (another process takes VRAM, a +// prompt shape the estimator does not model). A skip-park window whose +// drafter work fails for lack of device memory is retried once with the +// target parked. Any other failure (non-finite scores, invalid spans) is not +// a VRAM problem and is returned as is: parking would not fix it. +// +// After a parked retry recovers, the next `backoff` windows park outright, +// then skip-park is probed again. Each further out-of-memory recovery doubles +// the backoff (4 → 8 → … → 64 windows); a successful no-park window resets +// it. While any backoff is pending the drafter is released after each window +// even when the residency policy would keep it loaded. +enum class SkipParkWindowOutcome { Ok, OutOfMemory, Failed }; + +class SkipParkFallback { +public: + static constexpr int kInitialBackoffWindows = 4; + static constexpr int kMaxBackoffWindows = 64; + + // This window must park even though skip-park was requested. + bool parking_forced() const { return forced_windows_ > 0; } + // A recent out-of-memory window has not yet been followed by a clean + // no-park window: VRAM is tight, so do not keep the drafter resident. + bool memory_tight() const { return backoff_ > 0; } + int forced_windows() const { return forced_windows_; } + int backoff() const { return backoff_; } + + void on_forced_window() { + if (forced_windows_ > 0) --forced_windows_; + } + void on_oom_recovered() { + backoff_ = backoff_ == 0 + ? kInitialBackoffWindows + : std::min(backoff_ * 2, kMaxBackoffWindows); + forced_windows_ = backoff_; + } + void on_skip_park_ok() { + backoff_ = 0; + forced_windows_ = 0; + } + +private: + int forced_windows_ = 0; + int backoff_ = 0; +}; + +// One compression window under the skip-park policy and fail-safe. +// run_window(park) runs the window, parking the resident models iff park +// classify(results) → SkipParkWindowOutcome +// drop_drafter() frees whatever the failed attempt left allocated +// Returns the results of the attempt that counts. +template +auto run_skip_park_window(bool skip_park_requested, SkipParkFallback & fallback, + RunWindow && run_window, Classify && classify, + DropDrafter && drop_drafter, const char * tag) + -> decltype(run_window(true)) { + if (!skip_park_requested) return run_window(true); + if (fallback.parking_forced()) { + fallback.on_forced_window(); + return run_window(true); + } + auto results = run_window(false); + const SkipParkWindowOutcome outcome = classify(results); + if (outcome == SkipParkWindowOutcome::Ok) { + fallback.on_skip_park_ok(); + return results; + } + if (outcome != SkipParkWindowOutcome::OutOfMemory) return results; + + std::fprintf(stderr, + "%s skip-park window ran out of device memory — parking and " + "retrying once\n", tag); + drop_drafter(); + auto retry = run_window(true); + if (classify(retry) == SkipParkWindowOutcome::Ok) { + fallback.on_oom_recovered(); + std::fprintf(stderr, + "%s parked retry succeeded — parking the next %d windows before " + "probing skip-park again\n", tag, fallback.forced_windows()); + } + return retry; } } // namespace luce::common diff --git a/server/src/qwen3/qwen3_backend.cpp b/server/src/qwen3/qwen3_backend.cpp index 61d0a5f44..ccbcb755f 100644 --- a/server/src/qwen3/qwen3_backend.cpp +++ b/server/src/qwen3/qwen3_backend.cpp @@ -6,7 +6,8 @@ // After all layers, out_norm + lm_head produces logits for the last token. #include "qwen3_backend.h" -#include "qwen3_drafter.h" +#include "internal.h" +#include "pflash/pflash_drafter.h" #include "luce.h" #include "common/sampler.h" #include "common/io_utils.h" @@ -25,7 +26,7 @@ namespace luce::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 +92,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", luce_last_error()); return false; } @@ -145,8 +146,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; } @@ -957,29 +958,59 @@ ModelBackend::CompressResult Qwen3Backend::compress(const CompressRequest & req) CompressResult result; if (req.input_ids.empty()) return result; - const bool was_parked = parked_; - if (!req.skip_park && !parked_) park(ParkTarget::TargetModel); - - if (!drafter_loaded_) { - if (!load_drafter(req.drafter_path, 999, req.drafter_gpu, drafter_ctx_)) { - std::fprintf(stderr, "[compress] load failed: %s\n", luce_last_error()); - if (!req.skip_park && !was_parked) unpark(ParkTarget::TargetModel); - return result; + auto attempt = [&](bool park_window) { + CompressResult r; + const bool was_parked = parked_; + if (park_window && !parked_) park(ParkTarget::TargetModel); + + if (!drafter_loaded_) { + if (!load_drafter(req.drafter_path, 999, req.drafter_gpu, + drafter_ctx_)) { + std::fprintf(stderr, "[compress] load failed: %s\n", + luce_last_error()); + r.out_of_memory = luce::common::last_error_is_oom(); + if (park_window && !was_parked) unpark(ParkTarget::TargetModel); + return r; + } + drafter_loaded_ = true; } - drafter_loaded_ = true; - } - - result.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(); - if (req.residency_action == DraftResidencyAction::ReleaseAfterUse) { - free_drafter(); - } + // 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(); + r = 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, + req.query_suffix_candidates, req.history_query_spans, + req.turn_query_span)); + r.out_of_memory = !r.ok && luce::common::last_error_is_oom(); + + // A recent out-of-memory window overrides KeepLoaded: VRAM is tight. + if (req.residency_action == DraftResidencyAction::ReleaseAfterUse || + skip_park_fallback_.memory_tight()) { + free_drafter(); + } - if (!req.skip_park && !was_parked) unpark(ParkTarget::TargetModel); + if (park_window && !was_parked) unpark(ParkTarget::TargetModel); + return r; + }; + + const auto classify = [](const CompressResult & r) { + return r.ok ? SkipParkWindowOutcome::Ok + : r.out_of_memory ? SkipParkWindowOutcome::OutOfMemory + : SkipParkWindowOutcome::Failed; + }; + result = run_skip_park_window( + req.skip_park, skip_park_fallback_, attempt, classify, + [this]() { + // Unconditional free: a failed load_drafter can leave a live + // backend in the ctx even though drafter_loaded_ is still false. + luce::common::free_drafter(drafter_ctx_); + drafter_loaded_ = false; + }, + "[compress]"); return result; } @@ -1030,7 +1061,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); @@ -1038,7 +1071,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() { @@ -1064,7 +1097,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 b4e61397b..6e59fb241 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,9 @@ #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 "placement/skip_park_guard.h" #include "common/sampler.h" #include "ggml.h" @@ -48,7 +49,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,13 +111,16 @@ class Qwen3Backend : public ModelBackend { private: Qwen3BackendConfig cfg_; ggml_backend_t backend_ = nullptr; - Qwen3DrafterWeights w_; + Qwen3Weights w_; Qwen3Cache cache_; bool parked_ = false; // Pflash drafter (lazy-loaded, reuses the same model for compress) DrafterContext drafter_ctx_; bool drafter_loaded_ = false; + // Skip-park fail-safe: parks a few requests after an out-of-memory + // no-park compress recovered with parking (placement/skip_park_guard.h). + SkipParkFallback skip_park_fallback_; // Sampler SamplerCfg sampler_; diff --git a/server/src/qwen3/qwen3_buffer_plan.h b/server/src/qwen3/qwen3_buffer_plan.h deleted file mode 100644 index cc9409d98..000000000 --- a/server/src/qwen3/qwen3_buffer_plan.h +++ /dev/null @@ -1,29 +0,0 @@ -#pragma once - -#include - -namespace luce::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 luce::common diff --git a/server/src/qwen3/qwen3_drafter.cpp b/server/src/qwen3/qwen3_drafter.cpp deleted file mode 100644 index 27979e939..000000000 --- a/server/src/qwen3/qwen3_drafter.cpp +++ /dev/null @@ -1,887 +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) -// - chunk-top-K + span merge (this file) -// -// 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/anchor_params.h" -#include "common/backend_precision.h" -#include "internal.h" -#include "anchor_scan.h" - -#include "ggml.h" -#include "ggml-alloc.h" -#include "ggml-backend.h" - -#include -#include -#include -#include -#include -#include -#include -#include - -namespace luce::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(LUCE_BACKEND_HIP) -bool prewarm_drafter_once(const Qwen3DrafterWeights & w) { - static bool warmed = false; - if (warmed || std::getenv("LUCE_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) { - 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; - } - - 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(LUCE_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) { - auto * st = static_cast(ctx.arch_state); - free_target_weights(st->weights); - delete st; - ctx.arch_state = nullptr; - } - if (ctx.loaded) { - if (ctx.arch == DrafterArch::Qwen3_0p6b) { - free_qwen3_drafter_model(ctx.weights); - } - } - ctx.loaded = false; -} - -static 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 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, "LUCE_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("LUCE_KV_TQ3", "0"); - auto restore_tq3 = [&]() { - // _putenv_s with empty value removes the variable on MSVCRT. - _putenv_s("LUCE_KV_TQ3", had_old_tq3 ? old_tq3_s.c_str() : ""); - }; -#else - const char * old_tq3 = std::getenv("LUCE_KV_TQ3"); - std::string old_tq3_s = old_tq3 ? old_tq3 : ""; - const bool had_old_tq3 = (old_tq3 != nullptr); - setenv("LUCE_KV_TQ3", "0", 1); - auto restore_tq3 = [&]() { - if (had_old_tq3) setenv("LUCE_KV_TQ3", old_tq3_s.c_str(), 1); - else unsetenv("LUCE_KV_TQ3"); - }; -#endif - if (!create_target_cache(w, S, 0, w.backend, cache, true)) { - restore_tq3(); - 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"); - 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("LUCE_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("LUCE_COMPRESS_HEAD_CHUNKS", 8); - const int t_raw = env_int("LUCE_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("LUCE_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("LUCE_COMPRESS_ANCHOR_RADIUS", -1), - env_int("LUCE_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("LUCE_COMPRESS_REPEAT_MIN", 4); - const int repeat_max = env_int("LUCE_COMPRESS_REPEAT_MAX", 32); - const int repeat_limit = env_int("LUCE_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; -} - -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 (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); - } - 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; - } - - // ── 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("LUCE_COMPRESS_HEAD_CHUNKS", 8); - const int t_raw = env_int("LUCE_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("LUCE_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("LUCE_COMPRESS_ANCHOR_RADIUS", -1), - env_int("LUCE_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); - - return out; -} - -} // namespace luce::common diff --git a/server/src/qwen3/qwen3_drafter.h b/server/src/qwen3/qwen3_drafter.h deleted file mode 100644 index 75c31f802..000000000 --- a/server/src/qwen3/qwen3_drafter.h +++ /dev/null @@ -1,87 +0,0 @@ -// In-process Qwen3-0.6B drafter for pflash speculative prefill. -// -// 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. -// -// Public entry point: drafter_score_and_compress() takes raw input token IDs, -// runs the full pflash compression pipeline in C++, returns the surviving -// token IDs (drafter vocab). - -#pragma once - -#include -#include -#include -#include - -#include "qwen3_drafter_model.h" - -struct ggml_backend; -typedef struct ggml_backend * ggml_backend_t; - -namespace luce::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 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; - bool loaded = false; -}; - -// Load the drafter GGUF (e.g. /opt/lucebox/models/drafter/Qwen3-0.6B-BF16.gguf). -// Creates a fresh GPU backend if `backend` is null. Otherwise uses the -// caller-provided backend (so the drafter shares the daemon's allocator). -// -// `gpu_layers` is accepted for API compat but ignored — every layer goes on -// the GPU since the drafter weights are only ~1.5 GB. -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); - -// Free only model weights, keeping the backend alive for reuse. -// 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). -// -// 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 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 -// -// On failure returns empty vector + sets last_error. -std::vector drafter_score_and_compress( - DrafterContext & ctx, - const std::vector & ids, - float keep_ratio, - int chunk_size = 32, - int n_lookahead = 8, - int pool_kernel = 13, - int score_query_end = -1); - -} // namespace luce::common diff --git a/server/src/qwen3/qwen3_drafter_model.h b/server/src/qwen3/qwen3_drafter_model.h deleted file mode 100644 index a6b7b1240..000000000 --- a/server/src/qwen3/qwen3_drafter_model.h +++ /dev/null @@ -1,129 +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 luce::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 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; -} - -} // namespace luce::common diff --git a/server/src/qwen3/qwen3_graph.cpp b/server/src/qwen3/qwen3_graph.cpp deleted file mode 100644 index 29f21e262..000000000 --- a/server/src/qwen3/qwen3_graph.cpp +++ /dev/null @@ -1,906 +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 (LUCE_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 LUCE_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 luce::common { - -namespace { - -constexpr int FA_WINDOW = 512; - -int chunk_s_ff() { - if (const char * e = std::getenv("LUCE_FP_CHUNK_S")) { - int v = std::atoi(e); - if (v >= 1024) return v; - } -#if defined(LUCE_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(LUCE_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; -} - -void warm_hip_chunk_graph_b_once(ggml_backend_t backend, HipChunkGraphB & out) { - static bool warmed = false; - if (warmed) { - return; - } - - 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) { - return; - } - } - - ggml_backend_graph_compute(backend, out.gf_proj_add); - ggml_backend_graph_compute(backend, out.gf_ffn); - warmed = 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(LUCE_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 LUCE_FP_NOPE_TAIL=0 to disable (saves ~K_curr_v memory). - static const bool nope_tail = []() -> bool { - const char * e = std::getenv("LUCE_FP_NOPE_TAIL"); - return e == nullptr || std::string(e) != "0"; - }(); - - 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 = (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 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 = 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)); - ggml_backend_graph_compute(w.backend, gf); - 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)); - - flashprefill::FlashPrefillConfig fp_cfg; -#if defined(LUCE_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 LUCE_FP_ALPHA for quality/speed sweeps. - fp_cfg.alpha = 0.95f; -#endif - if (const char* a = std::getenv("LUCE_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 = (early_exit_n > 0 && early_exit_n < w.n_layer) - ? early_exit_n : w.n_layer; - - 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("LUCE_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(); - 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 (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); - } - - // ── 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; - } - cudaDeviceSynchronize(); - 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(LUCE_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(); - warm_hip_chunk_graph_b_once(w.backend, gb); - 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; - } - 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(); - ggml_backend_graph_compute(w.backend, gf); - auto ts1 = std::chrono::steady_clock::now(); - acc = std::chrono::duration(ts1 - ts0).count(); - }; - one(gb.gf_proj_add, proj_s); - - 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); - cudaDeviceSynchronize(); - 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); - 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) { - 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; - } - } - } - - 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 luce::common diff --git a/server/src/qwen3/qwen3_loader.cpp b/server/src/qwen3/qwen3_loader.cpp index 2bf58a62f..a620a1a74 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,8 +22,9 @@ // We mmap the GGUF file and copy each tensor's bytes to the backend buffer // (mirrors the luce 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" #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; @@ -110,9 +125,9 @@ float get_f32(gguf_context * g, const char * key, float def) { } // 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; @@ -194,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; @@ -236,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); @@ -290,7 +305,7 @@ bool load_qwen3_drafter_model(const std::string & path, 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(); diff --git a/server/src/qwen3/qwen3_model.h b/server/src/qwen3/qwen3_model.h new file mode 100644 index 000000000..45f6812cb --- /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 luce::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 luce::common diff --git a/server/src/qwen35/gguf_target_loader.cpp b/server/src/qwen35/gguf_target_loader.cpp index c3565c7aa..3964e98fb 100644 --- a/server/src/qwen35/gguf_target_loader.cpp +++ b/server/src/qwen35/gguf_target_loader.cpp @@ -576,7 +576,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; @@ -847,7 +849,7 @@ bool load_target_gguf_partial(const std::string & path, } else if (ggml_backend_buft_is_meta(buft)) { out.buf = ggml_backend_alloc_ctx_tensors(out.ctx, backend); if (!out.buf) { - set_last_error("ggml_backend_alloc_ctx_tensors failed (target TP)"); + set_last_oom_error("ggml_backend_alloc_ctx_tensors failed (target TP)"); gguf_free(gctx); return false; } @@ -855,7 +857,7 @@ bool load_target_gguf_partial(const std::string & path, } else { out.buf = ggml_backend_alloc_buffer(backend, alloc_total); if (!out.buf) { - set_last_error("ggml_backend_alloc_ctx_tensors failed (target)"); + set_last_oom_error("ggml_backend_alloc_ctx_tensors failed (target)"); gguf_free(gctx); return false; } diff --git a/server/src/qwen35/qwen35_backend.cpp b/server/src/qwen35/qwen35_backend.cpp index 2dbbaa4dd..9868065ae 100644 --- a/server/src/qwen35/qwen35_backend.cpp +++ b/server/src/qwen35/qwen35_backend.cpp @@ -6,6 +6,7 @@ #include "common/spec_acceptance.h" #include "common/draft_block_size.h" #include "common/draft_swa.h" +#include "placement/gpu_vmm_pool.h" #include "placement/skip_park_guard.h" #include "qwen35_dflash_target.h" #include "graph_builders.h" @@ -26,8 +27,9 @@ #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/pflash_compress.h" +#include "pflash/kvflash_drafter_scorer.h" #include "ggml-cuda.h" #include "ggml-backend-impl.h" @@ -651,7 +653,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); } @@ -1145,9 +1147,44 @@ std::vector Qwen35Backend::compress_batch( } } if (load_request == nullptr) return results; - const bool should_park = !load_request->skip_park; + + const auto classify = [&requests](const std::vector & rs) { + auto outcome = SkipParkWindowOutcome::Ok; + for (size_t i = 0; i < requests.size(); ++i) { + const auto & r = requests[i]; + if (r.input_ids.empty() || r.drafter_path.empty()) continue; + if (i >= rs.size()) return SkipParkWindowOutcome::Failed; + if (rs[i].ok) continue; + if (rs[i].out_of_memory) return SkipParkWindowOutcome::OutOfMemory; + outcome = SkipParkWindowOutcome::Failed; + } + return outcome; + }; + const auto drop_drafter = [this]() { + // free_drafter() handles the loaded case and the kvflash scorer + // borrow; the unconditional free clears a backend/weights + // half-initialized by a failed load_drafter. + free_drafter(); + luce::common::free_drafter(drafter_ctx_); + drafter_loaded_ = false; + }; + return run_skip_park_window( + load_request->skip_park, skip_park_fallback_, + [&](bool park_window) { + return run_compress_window(requests, *load_request, park_window); + }, + classify, drop_drafter, "[compress]"); +} + +std::vector Qwen35Backend::run_compress_window( + const std::vector & requests, + const CompressRequest & load_request, + bool park_window) { + std::vector results(requests.size()); + // A recent out-of-memory window overrides KeepLoaded: VRAM is tight. const bool release_after_use = - load_request->residency_action == DraftResidencyAction::ReleaseAfterUse; + load_request.residency_action == DraftResidencyAction::ReleaseAfterUse || + skip_park_fallback_.memory_tight(); // Park target+draft to free VRAM for the drafter (unless skip_park). // A FlowKV request may contain many aged messages. Keep this residency @@ -1155,7 +1192,7 @@ std::vector Qwen35Backend::compress_batch( // models independently. const bool was_target_parked = target_parked_; const bool was_draft_parked = draft_parked_; - if (should_park) { + if (park_window) { step_graph_destroy(sg_); if (!target_parked_) park(ParkTarget::TargetModel); if (!draft_parked_) park(ParkTarget::DraftModel); @@ -1175,12 +1212,14 @@ std::vector Qwen35Backend::compress_batch( if (!drafter_loaded_) { // drafter_ctx_.backend == nullptr → load_drafter creates its own std::fprintf(stderr, "[compress] loading drafter from %s ...\n", - load_request->drafter_path.c_str()); - if (!load_drafter(load_request->drafter_path, /*gpu_layers=*/999, - load_request->drafter_gpu, drafter_ctx_)) { + load_request.drafter_path.c_str()); + if (!load_drafter(load_request.drafter_path, /*gpu_layers=*/999, + load_request.drafter_gpu, drafter_ctx_)) { std::fprintf(stderr, "[compress] drafter init failed: %s\n", luce_last_error()); - if (should_park) { + const bool oom = luce::common::last_error_is_oom(); + for (auto & result : results) result.out_of_memory = oom; + if (park_window) { if (!was_target_parked) unpark(ParkTarget::TargetModel); if (!was_draft_parked) unpark(ParkTarget::DraftModel); } @@ -1202,11 +1241,28 @@ 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); + score_query_end, request.required_instruction_spans, + request.query_suffix_candidates, request.history_query_spans, + request.turn_query_span); result.ok = !result.compressed_ids.empty(); + result.out_of_memory = !result.ok && luce::common::last_error_is_oom(); + 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; + 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", request.input_ids.size(), result.compressed_ids.size()); @@ -1218,7 +1274,7 @@ std::vector Qwen35Backend::compress_batch( } // Restore park state - if (should_park) { + if (park_window) { if (!was_target_parked) unpark(ParkTarget::TargetModel); if (!was_draft_parked) unpark(ParkTarget::DraftModel); } @@ -1248,7 +1304,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; @@ -1257,7 +1313,8 @@ bool Qwen35Backend::handle_compress(const std::string & line, const DaemonIO & i if (cudaGetDeviceProperties(&prop, dev) == cudaSuccess) total_vram = prop.totalGlobalMem; const bool allowed = luce::common::skip_park_allowed( - skip_park, total_vram, cfg_.device.max_ctx); + skip_park, total_vram, cfg_.device.max_ctx, + luce::common::gpu_backend_uses_vmm_pool()); if (skip_park && !allowed) { std::fprintf(stderr, "[server] --prefill-skip-park downgraded: <32GB GPU with max_ctx>65536" diff --git a/server/src/qwen35/qwen35_backend.h b/server/src/qwen35/qwen35_backend.h index 423bb95d8..74e3410ed 100644 --- a/server/src/qwen35/qwen35_backend.h +++ b/server/src/qwen35/qwen35_backend.h @@ -26,8 +26,9 @@ #include "common/concurrency/paged_kv_pool.h" #include "concurrency/qwen35_seq_engine.h" #include "internal.h" // TargetWeights, TargetCache, DraftWeights, PrefixSnapshot +#include "placement/skip_park_guard.h" #include "qwen35_vision.h" -#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) @@ -275,6 +276,15 @@ class Qwen35Backend : public ModelBackend { void kvflash_ensure_scorer(); private: + // One compression window (park → load drafter → score → restore) with + // the park step optional. compress_batch runs it through + // run_skip_park_window, which retries parked after an out-of-memory + // no-park attempt. + std::vector run_compress_window( + const std::vector & requests, + const CompressRequest & load_request, + bool park_window); + // ── GPU backends ───────────────────────────────────────────────── ggml_backend_t target_backend_ = nullptr; ggml_backend_t draft_backend_ = nullptr; @@ -323,6 +333,9 @@ class Qwen35Backend : public ModelBackend { // ── Pflash drafter (lazy-loaded) ───────────────────────────────── DrafterContext drafter_ctx_; bool drafter_loaded_ = false; + // Skip-park fail-safe: parks a few windows after an out-of-memory + // no-park window recovered with parking (placement/skip_park_guard.h). + luce::common::SkipParkFallback skip_park_fallback_; // ── Sampler state ──────────────────────────────────────────────── SamplerCfg sampler_; diff --git a/server/src/qwen35/qwen35_layer_split_adapter.cpp b/server/src/qwen35/qwen35_layer_split_adapter.cpp index 3bb38c2cc..07b28582e 100644 --- a/server/src/qwen35/qwen35_layer_split_adapter.cpp +++ b/server/src/qwen35/qwen35_layer_split_adapter.cpp @@ -14,8 +14,9 @@ #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/pflash_compress.h" +#include "pflash/kvflash_drafter_scorer.h" #include "kv_quant.h" #include "ggml-cuda.h" @@ -210,7 +211,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 +1361,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,11 +1386,27 @@ 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); + score_query_end, req.required_instruction_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) { + 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; + 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", req.input_ids.size(), result.compressed_ids.size()); diff --git a/server/src/qwen35/qwen35_layer_split_adapter.h b/server/src/qwen35/qwen35_layer_split_adapter.h index 8b1b2dea5..36bd2ff99 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/qwen35/qwen35_target_graph.cpp b/server/src/qwen35/qwen35_target_graph.cpp index 57d4249e3..968a75e37 100644 --- a/server/src/qwen35/qwen35_target_graph.cpp +++ b/server/src/qwen35/qwen35_target_graph.cpp @@ -310,7 +310,7 @@ bool create_target_cache_partial(const TargetWeights & w, out.base_buf = ggml_backend_alloc_ctx_tensors(out.base_ctx, backend); if (!out.base_buf) { - set_last_error("ggml_backend_alloc_ctx_tensors failed for base cache"); + set_last_oom_error("ggml_backend_alloc_ctx_tensors failed for base cache"); ggml_free(out.base_ctx); out.base_ctx = nullptr; return false; @@ -388,7 +388,7 @@ bool create_target_cache_partial(const TargetWeights & w, ? f32_total - q8_total : 0); } if (!out.rollback_buf) { - set_last_error("ggml_backend_alloc_ctx_tensors failed for rollback cache"); + set_last_oom_error("ggml_backend_alloc_ctx_tensors failed for rollback cache"); ggml_free(out.rollback_ctx); out.rollback_ctx = nullptr; return false; @@ -721,7 +721,7 @@ bool migrate_prefill_cache(const TargetWeights & w, n_delta, max_verify_tokens, checkpoint_bytes); } if (!cache.rollback_buf) { - set_last_error("ggml_backend_alloc_ctx_tensors failed for rollback cache"); + set_last_oom_error("ggml_backend_alloc_ctx_tensors failed for rollback cache"); ggml_free(cache.rollback_ctx); cache.rollback_ctx = nullptr; return false; @@ -1077,7 +1077,7 @@ bool ensure_ssm_snapshot(TargetCache & c, ggml_backend_t backend) { c.rollback_buf = ggml_backend_alloc_ctx_tensors(c.rollback_ctx, backend); if (!c.rollback_buf) { - set_last_error("ensure_ssm_snapshot alloc_ctx_tensors failed"); + set_last_oom_error("ensure_ssm_snapshot alloc_ctx_tensors failed"); // Null the snap pointers so a later snapshot/restore_ssm_state (which // iterates ssm_state.size()) skips them instead of dereferencing // tensors from the freed rollback_ctx. diff --git a/server/src/server/adaptive_keep_ratio.h b/server/src/server/adaptive_keep_ratio.h index a83746c1f..1ac3aa9f9 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 18d8c7973..f2d8b66b1 100644 --- a/server/src/server/http_server.cpp +++ b/server/src/server/http_server.cpp @@ -28,6 +28,7 @@ #include "pin_friendly_prompt.h" #include "common/kv_rotation.h" #include "common/sha1.h" +#include "pflash/pflash_selection.h" #include "freeze_history.h" #ifdef LUCE_HAS_CURL @@ -38,6 +39,7 @@ #include #include #include +#include #include #include #include @@ -190,38 +192,688 @@ 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}; +} + +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); +} + +} // namespace + +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; +} + +PflashChatTurnSpan pflash_chat_query_turn( + const Tokenizer & marker_tokenizer, + const ChatMarkers & markers, + const Tokenizer & tokenizer, + const std::vector & prompt) { + 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> needles; + for (const auto & seq : markers.next_role_starts) { + std::string text = seq_text(seq); + if (!text.empty()) needles.emplace_back(std::move(text), true); + } + for (const auto & seq : markers.end_msg_seqs) { + std::string text = seq_text(seq); + if (!text.empty()) needles.emplace_back(std::move(text), false); + } + if (needles.empty()) return chosen; + + const DecodedPrompt decoded = + decode_prompt_with_offsets(tokenizer, prompt); + const std::string & text = decoded.text; + 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}); + } + } + 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 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; + }; + + struct 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; + }; + 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) { + turn.role = role_from_marker(mark.text); + } else { + 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; + } + turn.role = text.substr(content_at, name_end - content_at); + if (name_end < text.size() && text[name_end] == '\n') { + content_at = name_end + 1; + } + } + // 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(); + 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])) { + ++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. + 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_index = index; break; } + if (query_index == turns.size()) query_index = index; + } + 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 + // 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); + 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; + } + 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); +} + +bool pflash_paragraph_join() noexcept { + const char * raw = std::getenv("PFLASH_SELECT_PARAGRAPH_JOIN"); + return !(raw && std::string(raw) == "0"); +} + +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"); +} + +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); +} + +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; + 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) { + 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; +} + +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) || + 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 @@ -281,9 +933,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( @@ -791,6 +1445,10 @@ json build_props_body(const ServerConfig & config, {"keep_ratio", nullptr}, {"drafter_gguf", nullptr}, {"skip_park", nullptr}, + {"skip_park_mode", nullptr}, + {"skip_park_estimate_bytes", nullptr}, + {"skip_park_free_bytes", nullptr}, + {"drafter_keep_loaded", nullptr}, {"bsa_enabled", nullptr}, {"bsa_alpha", nullptr}, {"lm_head_fix", nullptr}, @@ -817,6 +1475,16 @@ json build_props_body(const ServerConfig & config, ? json(nullptr) : json(config.pflash_drafter_path)}, {"skip_park", config.pflash_skip_park}, + {"skip_park_mode", skip_park_mode_name(config.pflash_skip_park_mode)}, + {"drafter_keep_loaded", config.pflash_keep_drafter_loaded}, + {"skip_park_estimate_bytes", + config.pflash_skip_park_required_bytes > 0 + ? json(config.pflash_skip_park_required_bytes) + : json(nullptr)}, + {"skip_park_free_bytes", + config.pflash_skip_park_free_bytes > 0 + ? json(config.pflash_skip_park_free_bytes) + : json(nullptr)}, {"bsa_enabled", (bsa_env != nullptr && *bsa_env && std::strcmp(bsa_env, "0") != 0)}, {"bsa_alpha", bsa_alpha}, {"lm_head_fix", (lmfix_env != nullptr && *lmfix_env && std::strcmp(lmfix_env, "0") != 0)}, @@ -2465,6 +3133,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. @@ -3026,6 +3696,7 @@ void HttpServer::apply_flowkv_compression( DraftResidencyUse::PFlashCompress, config_.lazy_draft, !config_.draft_path.empty(), + config_.pflash_keep_drafter_loaded, // ample_vram: auto estimate }); std::vector compress_requests; @@ -3152,8 +3823,11 @@ void HttpServer::apply_flowkv_compression( std::string HttpServer::apply_pflash_compression( const ParsedRequest & req, PreparedPrompt & prepared) { + const bool selection_environment = + luce::pflash::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); @@ -3170,45 +3844,667 @@ 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"; } + luce::pflash::PFlashSelectionConfig experiment; + std::string experiment_error; + if (!luce::pflash::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"); + 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; + int query_content_end = -1; + // 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; + // 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; + // 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. + PFlashTokenSpan query_role_header{-1, -1}; + std::vector required_instruction_spans; + // 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 + // parser and the legacy window; unused when a benchmark parser + // (latest_user) or a marker-less prompt needs the sentinel mapping. + http_detail::PflashChatTurnSpan chat_turn; + if (!experiment.configured || tail_parser) { + ChatMarkers chat_markers; + if (resolve_chat_markers(tokenizer_, chat_markers)) { + chat_turn = http_detail::pflash_chat_query_turn( + 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"; + } + 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; + } + + 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 || last_user_index >= 0)) { + boundary_index = last_user_index; + } + + static constexpr const char * kContentBegin = + "__LUCE_PFLASH_CONTENT_BEGIN_02C47F91__"; + static constexpr const char * kContentEnd = + "__LUCE_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; + // 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) { + // 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()) { + return "PFlash strict selection content boundary mapping failed"; + } + // Chat default: without an explicit pflash_query the scorer + // 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 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"; + 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, + 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 = + 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; + } + (messages[instruction_index].role == "system" + ? system_spans : instruction_role_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"; + } + instruction_role_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 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; + 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 " + "latest user content" + : "PFlash strict selection explicit query mapping " + "failed: pflash_query does not occur in the " + "prompt"; + } + } + 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); + } + // 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()}); + } + } + // 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) { + // 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) { + 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; + 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); + 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)); + 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; + } + } + } 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); + } + // 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, 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)) { + // 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 (!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); + } + } + + // 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); + }; + 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); + // 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 developer and tool " + "spans as context\n", + target_estimate(kept_tokens), req.max_output, config_.max_ctx); + 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 + // 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 = + 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()) { + 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; + 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", + luce::pflash::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", query_span.begin}, + {"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}, + {"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", @@ -3228,72 +4524,189 @@ std::string HttpServer::apply_pflash_compression( DraftResidencyUse::PFlashCompress, config_.lazy_draft, !config_.draft_path.empty(), + config_.pflash_keep_drafter_loaded, // ample_vram: auto estimate }); compress_request.residency_action = residency; - 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(); + // 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; + + // 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, + 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 {}; } - } 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"; - } + 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) { + 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); + } - std::string compressed_text = - drafter_tokenizer_->decode(result.compressed_ids); + if (!result.ok || result.compressed_ids.empty()) { + return config_.pflash_remote_drafter + ? "remote PFlash drafter compression failed" + : "PFlash compression failed"; + } - // 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; + std::string compressed_text = + drafter_tokenizer_->decode(result.compressed_ids); + join_overhead = 0; + // 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. + 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() - join_overhead <= target_ceiling) { + break; + } + const int overflow = + (int) final_tokens.size() - join_overhead - 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() - 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) + ")"; + } + } + 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()) { + std::vector served; + if (serve_pflash_chat_view( + req, compress_request.input_ids, chat_turn, &final_tokens, + &result.kept_spans, &result.candidate_lifts, served, + prepared.snapshot_cut, prepared.pflash_stats["view"])) { + final_tokens = std::move(served); } } - - prepared.tokens = tokenizer_.encode(compressed_text); + 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(), @@ -3302,6 +4715,259 @@ std::string HttpServer::apply_pflash_compression( return {}; } +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 * fresh, + const std::vector * kept_spans, + const std::vector> * lifts, + 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 false; + const int input = (int) drafter_ids.size(); + if (turn.generation_begin <= 0 || turn.generation_begin >= input) { + return false; + } + // 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) || + (compressed && !ends_with_generation(*fresh))) { + return false; + } + 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); + stats = {{"mode", why}, {"turn", turns}, {"served_tokens", fresh->size()}, + {"fresh_tokens", fresh->size()}}; + served = *fresh; + pflash_views_.remember(std::move(next)); + return true; + }; + 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()}}; + if (compressed) stats["fresh_tokens"] = fresh->size(); + snapshot_cut = view.view_gen_begin; + 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 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; + 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); + } + } + 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"; + } + } + + // 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 end > begin + ? drafter_tokenizer_->decode(std::vector( + drafter_ids.begin() + begin, drafter_ids.begin() + end)) + : std::string(); + }; + 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::vector before_split; + std::vector after_split; + for (const auto & span : delta_spans) { + 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); + served.insert(served.end(), delta_tokens.begin(), delta_tokens.end()); + if (!ends_with_generation(served)) { + 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 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.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] %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", mode}, {"turn", next.turns}, + {"served_tokens", served.size()}, {"reused_tokens", view.view_gen_begin}, + {"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 true; +} + HttpServer::PreparedPrompt HttpServer::prepare_prompt( const ParsedRequest & req) { PreparedPrompt prepared; @@ -3325,15 +4991,42 @@ HttpServer::PreparedPrompt HttpServer::prepare_prompt( prompt_tokens >= config_.pflash_threshold); const bool continuation = should_compress && is_continuation_request(req.messages); + const bool selection_environment = + luce::pflash::has_pflash_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( + 0, 32, experiment, experiment_error)) { + prepared.error_status = 500; + prepared.error = "invalid PFlash strict selection config: " + + experiment_error; + return prepared; + } + if (req.disk_cache_policy.compress) { + prepared.error_status = 500; + prepared.error = + "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"); @@ -3454,6 +5147,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 @@ -4476,9 +6170,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, @@ -4573,6 +6269,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 c6a63b3bd..d558c9077 100644 --- a/server/src/server/http_server.h +++ b/server/src/server/http_server.h @@ -26,6 +26,7 @@ #include "api_types.h" #include "placement/draft_residency.h" #include "placement/remote_draft_config.h" +#include "placement/skip_park_guard.h" #include "common/pflash_drafter_ipc.h" #include "model_card.h" #include "adaptive_keep_ratio.h" @@ -217,11 +218,26 @@ 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 - bool pflash_skip_park = false; // skip park/unpark for >=32GB GPUs + // --prefill-skip-park auto|on|off. `pflash_skip_park` is the RESOLVED + // decision, fixed once at startup by resolve_skip_park(): under Auto it + // compares the drafter GGUF footprint against measured free VRAM with a + // margin; On/Off are explicit overrides (the <32GiB/ctx>64K hardware + // guard still applies to On). Remote drafter and upstream-proxy modes + // resolve off. See placement/skip_park_guard.h. + SkipParkMode pflash_skip_park_mode = SkipParkMode::Auto; + bool pflash_skip_park = false; + // Auto only: the estimate also fits the drafter kept resident beside the + // target's compute reserve, so draft-residency auto keeps the pflash + // drafter loaded between requests. False for explicit on/off. + bool pflash_keep_drafter_loaded = false; + // Auto estimate (bytes incl. margin) and free VRAM at resolve time — + // exposed in logs and /props for forensics. 0 when not estimated. + int64_t pflash_skip_park_required_bytes = 0; + int64_t pflash_skip_park_free_bytes = 0; // Passthrough proxy — forward to upstream OpenAI-compatible server std::string pflash_upstream_base; // e.g. "http://localhost:8080/v1" std::string pflash_upstream_key; // Bearer token for upstream @@ -282,22 +298,227 @@ 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); + +// 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; + +// 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; + +// 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, + const std::vector & ids, + 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( + 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( 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); + +// 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 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 +// 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; + int content_end = -1; + 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; + } +}; + +PflashChatTurnSpan pflash_chat_query_turn( + 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. +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; +// 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; } // namespace http_detail @@ -345,6 +566,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; @@ -418,6 +644,7 @@ class HttpServer { private: friend struct SchedulerTestHarness; + friend struct HttpServerTestAccess; // Client thread: read HTTP request, parse, enqueue job, wait. void handle_client(SocketHandle fd); @@ -442,6 +669,12 @@ 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; + // PFlash details for usage.timings.pflash (see build_timings_json). + nlohmann::json pflash_stats; int error_status = 0; std::string error; }; @@ -453,6 +686,28 @@ 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 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 * fresh, + const std::vector * kept_spans, + const std::vector> * lifts, + 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); @@ -601,6 +856,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_; @@ -706,6 +964,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/src/server/server_main.cpp b/server/src/server/server_main.cpp index a70e05081..6068c25a8 100644 --- a/server/src/server/server_main.cpp +++ b/server/src/server/server_main.cpp @@ -16,6 +16,8 @@ #include "model_card.h" #include "common/backend_factory.h" #include "common/chain_rollback_policy.h" +#include "common/gguf_inspect.h" +#include "common/gpu_runtime_compat.h" #include "common/layer_split_utils.h" #include "common/model_capabilities.h" #include "common/spark_corpus.h" @@ -27,6 +29,9 @@ #include "engine/luce_engine.h" #include "placement/pflash_placement.h" #include "placement/draft_residency.h" +#include "placement/gpu_vmm_pool.h" +#include "pflash/pflash_drafter.h" +#include "pflash/pflash_selection.h" #include "kvflash_pager.h" #include "kv_quant.h" @@ -202,8 +207,12 @@ 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-skip-park Skip park/unpark (for >=32GB GPUs)\n" + " --prefill-drafter Drafter GGUF for compression (Qwen3.5-0.8B)\n" + " --prefill-skip-park [auto|on|off]\n" + " Keep target+draft resident while the\n" + " drafter scores (default: auto — resolved\n" + " from drafter footprint vs. free VRAM;\n" + " bare flag = on)\n" " --draft-residency auto|persistent|request-scoped\n" " Drafter lifetime policy (default: auto)\n" " --lazy-draft Legacy alias for --draft-residency=request-scoped\n" @@ -706,7 +715,16 @@ static int parse_model_options(int argc, char ** argv, ModelOptions & model, } else if (std::strcmp(argv[i], "--prefill-drafter") == 0 && i + 1 < argc) { sconfig.pflash_drafter_path = argv[++i]; } else if (std::strcmp(argv[i], "--prefill-skip-park") == 0) { - sconfig.pflash_skip_park = true; + // Tri-state: bare flag = on (historical boolean). The value is + // optional — consume the next arg only when it parses as a mode + // so a stray positional isn't mistaken for a value. + if (i + 1 < argc && argv[i + 1][0] != '-' && + parse_skip_park_mode(argv[i + 1], + sconfig.pflash_skip_park_mode)) { + ++i; + } else { + sconfig.pflash_skip_park_mode = SkipParkMode::On; + } } else if (std::strcmp(argv[i], "--prefill-upstream-base") == 0 && i + 1 < argc) { sconfig.pflash_upstream_base = argv[++i]; // Strip trailing slash @@ -1115,11 +1133,11 @@ static int load_model(ModelOptions & model, LoadedModel & loaded, bool multi_mod std::fprintf(stderr, "[server] drafter tokenizer load failed\n"); return 1; } - std::fprintf(stderr, "[server] pflash: mode=%s threshold=%d keep=%.3f drafter_gpu=%d skip_park=%d\n", + std::fprintf(stderr, "[server] pflash: mode=%s threshold=%d keep=%.3f drafter_gpu=%d skip_park=%s\n", sconfig.pflash_mode == ServerConfig::PflashMode::AUTO ? "auto" : "always", sconfig.pflash_threshold, sconfig.pflash_keep_ratio, sconfig.pflash_drafter_gpu, - (int)sconfig.pflash_skip_park); + skip_park_mode_name(sconfig.pflash_skip_park_mode)); if (!sconfig.pflash_curve.empty()) { std::fprintf(stderr, "[server] pflash curve:"); for (const auto & p : sconfig.pflash_curve) @@ -1206,6 +1224,78 @@ static int load_model(ModelOptions & model, LoadedModel & loaded, bool multi_mod backend->shutdown(); return 2; } + + // ── Skip-park startup resolution ────────────────────────────────────── + // Resolve --prefill-skip-park once, here: after the backend has fully + // loaded (so free VRAM reflects the resident target + decode draft) and + // before sconfig is frozen into the HttpServer. Auto compares the + // drafter's GGUF-derived footprint against measured free VRAM on the + // drafter device; on/off are explicit. Remote-IPC and upstream-proxy + // drafters never park in-process, so they resolve off unconditionally. + if (pflash_enabled && !sconfig.pflash_drafter_path.empty() && + !sconfig.pflash_remote_drafter && + sconfig.pflash_upstream_base.empty()) { + // The scorer's query window sizes its logits; an invalid selection + // config fails every request later, so the widest window is a safe + // stand-in here. + luce::pflash::PFlashSelectionConfig selection; + std::string selection_error; + const int query_tokens = + luce::pflash::resolve_pflash_selection( + sconfig.max_ctx, /*legacy_chunk_size=*/32, selection, + selection_error) + ? selection.query_tokens : 512; + SkipParkDrafterInfo footprint; + const bool footprint_ok = inspect_drafter_footprint( + sconfig.pflash_drafter_path, query_tokens, + pflash_scoring_sessions(), footprint); + const bool vmm_pool = gpu_backend_uses_vmm_pool(); + int64_t total_vram = -1, free_vram = -1; + int prev_dev = -1; + if (cudaGetDevice(&prev_dev) == cudaSuccess) { + cudaDeviceProp prop{}; + if (cudaGetDeviceProperties(&prop, sconfig.pflash_drafter_gpu) == + cudaSuccess) { + total_vram = int64_t(prop.totalGlobalMem); + } + if (cudaSetDevice(sconfig.pflash_drafter_gpu) == cudaSuccess) { + size_t free_b = 0, total_b = 0; + if (cudaMemGetInfo(&free_b, &total_b) == cudaSuccess) { + free_vram = int64_t(free_b); + } + (void)cudaSetDevice(prev_dev); + } + } + const SkipParkDecision decision = resolve_skip_park( + sconfig.pflash_skip_park_mode, + /*drafter_configured=*/true, + footprint_ok ? footprint : SkipParkDrafterInfo{}, + free_vram, total_vram, sconfig.max_ctx, vmm_pool); + sconfig.pflash_skip_park = decision.enabled; + sconfig.pflash_keep_drafter_loaded = decision.keep_drafter_loaded; + sconfig.pflash_skip_park_required_bytes = decision.required_bytes; + sconfig.pflash_skip_park_free_bytes = + std::max(free_vram, 0); + std::fprintf(stderr, + "[server] pflash skip-park: mode=%s → %s " + "(need %.2f GiB incl. margin, keep-loaded %.2f GiB, free %.2f GiB, " + "window %lld tok, query %d tok, vmm_pool=%d)\n", + skip_park_mode_name(sconfig.pflash_skip_park_mode), + decision.reason.c_str(), + decision.required_bytes / double(1ll << 30), + decision.keep_loaded_bytes / double(1ll << 30), + std::max(free_vram, 0) / double(1ll << 30), + (long long)decision.window_tokens, query_tokens, (int)vmm_pool); + } else { + sconfig.pflash_skip_park = false; + sconfig.pflash_keep_drafter_loaded = false; + if (pflash_enabled) { + std::fprintf(stderr, + "[server] pflash skip-park: off (remote/upstream drafter or " + "no drafter)\n"); + } + } + // ── Thinking-budget v2: resolve model card and apply to ServerConfig ── // Reuse the metadata captured during factory preparation instead of // opening the GGUF header again. @@ -1499,7 +1589,9 @@ static int load_model(ModelOptions & model, LoadedModel & loaded, bool multi_mod std::fprintf(stderr, "[server] │ pflash_drafter_gpu= %d\n", sconfig.pflash_drafter_gpu); std::fprintf(stderr, "[server] │ pflash_drafter_exec= %s\n", sconfig.pflash_remote_drafter ? "remote-ipc" : "local"); - std::fprintf(stderr, "[server] │ pflash_skip_park= %s\n", sconfig.pflash_skip_park ? "ON" : "off"); + std::fprintf(stderr, "[server] │ pflash_skip_park= %s (mode=%s)\n", + sconfig.pflash_skip_park ? "ON" : "off", + skip_park_mode_name(sconfig.pflash_skip_park_mode)); std::fprintf(stderr, "[server] │ fp_use_bsa = %s\n", getenv("LUCE_FP_USE_BSA") ? "ON" : "off"); std::fprintf(stderr, "[server] │ fp_alpha = %s\n", getenv("LUCE_FP_ALPHA") ? getenv("LUCE_FP_ALPHA") : "0.12 (default)"); } 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/bench_laguna_pflash.cpp b/server/test/bench_laguna_pflash.cpp index 8502f08e7..b58dcc76e 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 "luce.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", luce_last_error()); diff --git a/server/test/pflash_daemon.cpp b/server/test/pflash_daemon.cpp index c16c3afee..209912afa 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 "luce.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 a14969a1a..623b15742 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 "luce.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 765703824..e05cc2e5a 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 luce::common; diff --git a/server/test/test_anchor_transitive.cpp b/server/test/test_anchor_transitive.cpp index b8c00ab1d..047ad516b 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); - luce::qwen3::AnchorScanCfg cfg{CHUNK, /*anchor_radius=*/0, + luce::pflash::AnchorScanCfg cfg{CHUNK, /*anchor_radius=*/0, /*max_anchor_hits=*/8, /*ngram=*/4}; - luce::qwen3::scan_and_force(ids, q0, query_pool, cfg, forced); + luce::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); - luce::qwen3::AnchorScanCfg cfg{CHUNK, /*anchor_radius=*/0, + luce::pflash::AnchorScanCfg cfg{CHUNK, /*anchor_radius=*/0, /*max_anchor_hits=*/8, /*ngram=*/4}; - luce::qwen3::scan_and_force(ids, q0, query_pool, cfg, forced); + luce::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); - luce::qwen3::AnchorScanCfg cfg{CHUNK, /*anchor_radius=*/0, + luce::pflash::AnchorScanCfg cfg{CHUNK, /*anchor_radius=*/0, /*max_anchor_hits=*/8, /*ngram=*/4}; - luce::qwen3::scan_and_force_transitive(ids, q0, initial_query_pool, + luce::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); - luce::qwen3::AnchorScanCfg cfg{CHUNK, /*anchor_radius=*/0, + luce::pflash::AnchorScanCfg cfg{CHUNK, /*anchor_radius=*/0, /*max_anchor_hits=*/8, /*ngram=*/4, /*rare_token_max_freq=*/8}; - luce::qwen3::scan_and_force_transitive(ids, q0, initial_query_pool, + luce::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); - luce::qwen3::AnchorScanCfg cfg{CHUNK, /*anchor_radius=*/0, + luce::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}; - luce::qwen3::scan_and_force_transitive(ids, q0, query_pool, + luce::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); - luce::qwen3::AnchorScanCfg cfg{CHUNK, /*anchor_radius=*/0, + luce::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}; - luce::qwen3::scan_and_force_transitive(ids, q0, query_pool, + luce::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); - luce::qwen3::AnchorScanCfg cfg{CHUNK, /*anchor_radius=*/0, + luce::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}; - luce::qwen3::scan_and_force_transitive(ids, q0, query_pool, + luce::pflash::scan_and_force_transitive(ids, q0, query_pool, cfg, /*max_iters=*/25, forced); int total_forced = 0; diff --git a/server/test/test_bandit_integration.cpp b/server/test/test_bandit_integration.cpp index 77da5b7c6..7b0041908 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_dflash.cpp b/server/test/test_dflash.cpp index 77cf23c0f..d47e659cb 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,7 @@ int main(int argc, char ** argv) { "[compress] bad args, need: [drafter_arch]\n"); stream_emit(-1); continue; } - luce::common::DrafterArch drafter_arch; - if (!luce::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 +2425,21 @@ int main(int argc, char ** argv) { } if (!drafter_loaded) { - if (!luce::common::load_drafter(drafter_path, /*gpu_layers=*/999, drafter_arch, drafter_ctx)) { + if (!luce::common::load_drafter(drafter_path, /*gpu_layers=*/999, drafter_ctx)) { std::fprintf(stderr, "[compress] load_drafter failed: %s\n", luce_last_error()); stream_emit(-1); continue; } drafter_loaded = true; - if (drafter_arch == luce::common::DrafterArch::Qwen3_0p6b) { - std::printf("[drafter] loaded %s arch=%s (n_layer=%d n_head=%d n_head_kv=%d)\n", - drafter_path, luce::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, luce::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", - luce::common::drafter_arch_name(drafter_arch), - luce::common::drafter_arch_name(drafter_ctx.arch)); - stream_emit(-1); continue; } float keep = (float)keep_x1000 / 1000.0f; auto compressed = luce::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 7aedf2d9d..000000000 --- a/server/test/test_drafter_early_exit_score_range.cpp +++ /dev/null @@ -1,90 +0,0 @@ -// Unit tests for luce::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 luce::common::ScoreRange; -using luce::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 734dcf44c..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 luce::common::ScoreRange; -using luce::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 70960ccf5..a778390a2 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", luce_last_error()); diff --git a/server/test/test_pflash_drafter_ipc.cpp b/server/test/test_pflash_drafter_ipc.cpp new file mode 100644 index 000000000..9aca1b2b9 --- /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 "pflash/pflash_selection.h" + +#include +#include + +using namespace luce::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], + luce::pflash::pflash_chunk_is_structurally_required( + begin, end, 28, 32, 32, spans), + }); + } + return luce::pflash::select_pflash_candidates( + candidates, {16, 0.95}, + luce::pflash::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(!luce::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(!luce::pflash::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..a08e35ded --- /dev/null +++ b/server/test/test_pflash_selection.cpp @@ -0,0 +1,877 @@ +#include "CppUnitTestFramework.hpp" + +#include "pflash/pflash_selection.h" +#include "pflash/pflash_compress.h" +#include "scoped_env.h" + +#include +#include +#include +#include +#include +#include +#include + +using namespace luce::pflash; + +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 * kTopKEnv = "PFLASH_SELECT_TOPK"; + +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}; + luce_test::ScopedEnvVar top_k{kTopKEnv, 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, 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; + 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::ArbitraryTail); + 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; + luce::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); + + luce::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); +} + +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); +} diff --git a/server/test/test_qwen3_buffer_plan.cpp b/server/test/test_qwen3_buffer_plan.cpp deleted file mode 100644 index 3081dddb9..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 luce::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 7f7348eb6..d6d01e834 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" @@ -45,13 +46,15 @@ #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 "luce.h" #include "gguf.h" #include #include +#include #include #include #include @@ -72,6 +75,7 @@ #include #else #include +#include #endif #if defined(_WIN32) @@ -153,6 +157,18 @@ 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); + } + static HttpServer::PreparedPrompt prepare_prompt( + HttpServer & server, const ParsedRequest & req) { + return server.prepare_prompt(req); + } +}; } namespace { @@ -201,26 +217,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) { @@ -229,25 +252,950 @@ 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, + 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); + 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); + // 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; +#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; +} + +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 = + "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()); +} + +// 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; + 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 + std::vector roles; +}; + +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", "model", + "system", "\n", "Sure", "."}, + rendered, control_tokens); + Tokenizer tok; + PflashRenderedQueryTurn out; + if (!tok.load_from_gguf(path.c_str())) { + unlink(path.c_str()); + return out; + } + const auto prompt = tok.encode(rendered); + ChatMarkers markers; + 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()}); + 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; + for (const auto & each : turn.turns) out.roles.push_back(each.role); + } + } + unlink(path.c_str()); + return out; +} + +TEST_CASE(ServerUnitFixture, + 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); + TEST_ASSERT_MSG(!turn.later_turns, family.name); + } + } +} + +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"); + 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); + + 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_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); + 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_chat_query_turn( + 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_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); + 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_chat_query_turn( + 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({ + {{"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_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_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); + 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}}; + 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. + 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); + 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); @@ -268,7 +1216,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( @@ -281,12 +1228,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) { @@ -1707,8 +2666,8 @@ TEST_CASE(ServerUnitFixture, test_parse_dsml_tool_calls_unclosed_final_invoke_be "<｜DSML｜tool_calls>\n" "<｜DSML｜invoke name=\"edit\">\n" "<｜DSML｜parameter name=\"file_path\" string=\"true\">/home/dpavlin/aimax/LUCEBOX_STRIX_HALO_GUIDE.md\n" - "<｜DSML｜parameter name=\"new_string\" string=\"true\">### B. Paged Attention WMMA (head-256, RDNA4) — `DFLASH27B_PAGED_WMMA` (default off, BURN-IN)\n" - "- Env `DFLASH27B_PAGED_WMMA=1` routes paged full-attention layers (head 256, F16/Q8_0/Q4_0 KV, non-tree) to the WMMA kernel. Default (unset/0) keeps the V_DOT2 decode kernel.\n" + "<｜DSML｜parameter name=\"new_string\" string=\"true\">### B. Paged Attention WMMA (head-256, RDNA4) — `LUCE_PAGED_WMMA` (default off, BURN-IN)\n" + "- Env `LUCE_PAGED_WMMA=1` routes paged full-attention layers (head 256, F16/Q8_0/Q4_0 KV, non-tree) to the WMMA kernel. Default (unset/0) keeps the V_DOT2 decode kernel.\n" "- Differential: single-prompt TTFT −21% @12K, −42% @44K; batched 8K-pool prefill slightly ahead. Kernel-level 20.6–22.4 TFLOP/s vs 6.4–6.8 (3.1–3.3×); end-to-end bounded by attention's prefill share (~8% @44K, ~4% @12K).\n" "- Two-mode CTest: `test_paged_attn_wmma` (V_DOT2, env=0) / `paged_attn_wmma_route` (env=1), diff via `server/test/compare_paged_attn.py --tol 6e-3`. Need `--reconfig` after the upstream pull for `CMakeLists.txt` to register targets.\n" "- Source of truth: `server/docs/PAGED_ATTN_WMMA_HANDOFF.md`.\n" @@ -1735,7 +2694,7 @@ TEST_CASE(ServerUnitFixture, test_parse_dsml_tool_calls_unclosed_final_invoke_be TEST_ASSERT(result.tool_calls[0].name == "edit"); auto args = json::parse(result.tool_calls[0].arguments); TEST_ASSERT(args["file_path"] == "/home/dpavlin/aimax/LUCEBOX_STRIX_HALO_GUIDE.md"); - TEST_ASSERT(args["new_string"].get().find("DFLASH27B_PAGED_WMMA") != std::string::npos); + TEST_ASSERT(args["new_string"].get().find("LUCE_PAGED_WMMA") != std::string::npos); TEST_ASSERT(args["old_string"].get().find("Prefill Mode") != std::string::npos); } } @@ -3084,32 +4043,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; @@ -4548,6 +5481,29 @@ TEST_CASE(ServerUnitFixture, test_draft_residency_pflash_auto) { /*has_decode_draft=*/true, }); TEST_ASSERT(action == DraftResidencyAction::ReleaseAfterUse); + + // Proven ample VRAM (resolved skip-park) upgrades auto to KeepLoaded — + // the drafter stays resident between requests. + action = resolve_draft_residency_action( + DraftResidencyPolicy::Auto, + DraftResidencyContext{ + DraftResidencyUse::PFlashCompress, + /*low_vram_hint=*/false, + /*has_decode_draft=*/false, + /*ample_vram=*/true, + }); + TEST_ASSERT(action == DraftResidencyAction::KeepLoaded); + + // Explicit policies still override the hint both ways. + action = resolve_draft_residency_action( + DraftResidencyPolicy::RequestScoped, + DraftResidencyContext{ + DraftResidencyUse::PFlashCompress, + /*low_vram_hint=*/false, + /*has_decode_draft=*/false, + /*ample_vram=*/true, + }); + TEST_ASSERT(action == DraftResidencyAction::ReleaseAfterUse); } TEST_CASE(ServerUnitFixture, test_draft_residency_dflash_auto_and_request_scoped) { @@ -5993,6 +6949,919 @@ 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()); +} + +TEST_CASE(ServerUnitFixture, + 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 + // 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", "."}, + 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", "system"}, {"content", "You are helpful."}}, + {{"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 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 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; + 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 (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(question_pinned); + TEST_ASSERT(generation_pinned); + 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; + // 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. + 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); + 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()); +} + +// 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( + 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 < 120; ++i) history += "Sure. "; + PflashSystemPromptCase out; + out.rendered = render_chat_template( + {{instruction_role, 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", instruction_role}, {"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()); +} + +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(role); + 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 instructions + 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); + } + 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(compressed); + for (const auto & text : kept) { + TEST_ASSERT_MSG(text.find("You are helpful") == std::string::npos, text); + } +} + +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()); +} + +// 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}; + // 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. "; + 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; + std::vector modes; + { + 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; + // 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"); + 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); + TEST_ASSERT(modes == std::vector({"fresh", "continue", "repeat"})); + 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, through the + // last token of the header of the reply that followed it. + TEST_ASSERT(request.history_query_spans.size() == 1); + 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()); +} + +// 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; + // 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) { + 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; + 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"}; + // 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); + 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 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); +} + +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()); +} + +TEST_CASE(ServerUnitFixture, + test_pflash_strict_selection_owns_chat_continuations) { + luce_test::ScopedEnvVar mode{"PFLASH_SELECT_MODE", "top_p"}; + + 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", + "."}, + 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 prompt's last token. + TEST_ASSERT(backend.compress_calls == 1); + const auto & ids = backend.last_request.input_ids; + TEST_ASSERT(backend.last_request.score_query_end == (int) ids.size()); + TEST_ASSERT(backend.last_request.score_query_tokens == 1); + 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; @@ -8387,23 +10256,36 @@ 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); } // ═══════════════════════════════════════════════════════════════════════ -// 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; @@ -8459,7 +10341,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( - "luce_test_qwen3_drafter_438.gguf").string(); + "luce_test_qwen3_model_438.gguf").string(); gguf_write_to_file(g, path.c_str(), /*only_meta=*/false); gguf_free(g); @@ -8467,18 +10349,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, luce_last_error()); - free_qwen3_drafter_model(w); + free_qwen3_model(w); } // Truncate inside the tensor-data section. The header, kv block, and tensor @@ -8494,13 +10376,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 = luce_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); diff --git a/server/test/test_skip_park_guard.cpp b/server/test/test_skip_park_guard.cpp index 695c6a663..103b30f11 100644 --- a/server/test/test_skip_park_guard.cpp +++ b/server/test/test_skip_park_guard.cpp @@ -1,34 +1,390 @@ -// Unit tests for skip_park_allowed — pure, GPU-free. +// Unit tests for the skip-park policy — GPU-free. +// +// Covers: the VMM crash guard (skip_park_allowed) and its VMM-pool scope, the +// auto|on|off mode parse, the startup estimator (skip_park_required_bytes, +// inspect_drafter_footprint on a synthetic GGUF header), the precedence +// contract of resolve_skip_park — explicit off/on beat the estimate, the +// guard beats an explicit on, auto follows the footprint vs. free VRAM and +// alone may keep the drafter loaded — and the runtime fail-safe +// (SkipParkFallback, run_skip_park_window): only out-of-memory windows retry +// parked, and the park latch backs off instead of lasting forever. #include "CppUnitTestFramework.hpp" +#include "common/gguf_inspect.h" #include "placement/skip_park_guard.h" +#include "gguf.h" +#include "kv_quant.h" + +#include +#include +#include +#include +#include + namespace { struct SkipParkGuardFixture {}; } +using namespace luce::common; + static constexpr size_t GiB = 1024ull * 1024 * 1024; TEST_CASE(SkipParkGuardFixture, T1_not_requested_stays_off) { - CHECK(!luce::common::skip_park_allowed(false, 24 * GiB, 32768)); + CHECK(!skip_park_allowed(false, 24 * GiB, 32768)); } TEST_CASE(SkipParkGuardFixture, T2_big_card_any_ctx) { - CHECK(luce::common::skip_park_allowed(true, 32 * GiB, 131072)); + CHECK(skip_park_allowed(true, 32 * GiB, 131072)); } TEST_CASE(SkipParkGuardFixture, T3_small_card_small_ctx_allowed) { - CHECK(luce::common::skip_park_allowed(true, 24 * GiB, 65536)); + CHECK(skip_park_allowed(true, 24 * GiB, 65536)); } TEST_CASE(SkipParkGuardFixture, T4_small_card_big_ctx_downgraded) { - CHECK(!luce::common::skip_park_allowed(true, 24 * GiB, 131072)); + CHECK(!skip_park_allowed(true, 24 * GiB, 131072)); } TEST_CASE(SkipParkGuardFixture, T5_boundary_ctx_one_over) { - CHECK(!luce::common::skip_park_allowed(true, 24 * GiB, 65537)); + CHECK(!skip_park_allowed(true, 24 * GiB, 65537)); } TEST_CASE(SkipParkGuardFixture, T6_boundary_vram_just_under_32g) { - CHECK(!luce::common::skip_park_allowed(true, 32 * GiB - 1, 131072)); + CHECK(!skip_park_allowed(true, 32 * GiB - 1, 131072)); +} + +// ── Mode parse ──────────────────────────────────────────────────────────── + +TEST_CASE(SkipParkGuardFixture, T7_mode_parse) { + SkipParkMode m = SkipParkMode::Auto; + CHECK(parse_skip_park_mode("auto", m) && m == SkipParkMode::Auto); + CHECK(parse_skip_park_mode("on", m) && m == SkipParkMode::On); + CHECK(parse_skip_park_mode("off", m) && m == SkipParkMode::Off); + CHECK(!parse_skip_park_mode("yes", m)); + CHECK(!parse_skip_park_mode("", m)); + CHECK(std::string(skip_park_mode_name(SkipParkMode::Auto)) == "auto"); + CHECK(std::string(skip_park_mode_name(SkipParkMode::On)) == "on"); + CHECK(std::string(skip_park_mode_name(SkipParkMode::Off)) == "off"); +} + +// ── Startup resolution ──────────────────────────────────────────────────── + +// Qwen3.5-0.8B-ish footprint: ~1.6 GiB weights, ~15 KB/token hybrid runtime. +static SkipParkDrafterInfo qwen35_like_info() { + SkipParkDrafterInfo i; + i.recognized = true; + i.weights_bytes = int64_t(1668 * 1024 * 1024); + i.runtime_bytes_per_token = 15000; + i.fixed_bytes = 384ll * 1024 * 1024; + i.context_length = 262144; + return i; +} + +TEST_CASE(SkipParkGuardFixture, T8_explicit_off_always_wins) { + const auto d = resolve_skip_park(SkipParkMode::Off, /*drafter=*/true, + qwen35_like_info(), + /*free=*/64 * GiB, /*total=*/64 * GiB, + /*max_ctx=*/131072); + CHECK(!d.enabled); +} + +TEST_CASE(SkipParkGuardFixture, T9_explicit_on_bypasses_estimate) { + // Free VRAM far too small for the estimate — explicit on wins anyway. + const auto d = resolve_skip_park(SkipParkMode::On, /*drafter=*/true, + qwen35_like_info(), + /*free=*/1, /*total=*/32 * GiB, + /*max_ctx=*/131072); + CHECK(d.enabled); +} + +TEST_CASE(SkipParkGuardFixture, T10_explicit_on_blocked_by_guard) { + // 24 GiB card with ctx > 64K: the crash guard beats even force-on. + const auto d = resolve_skip_park(SkipParkMode::On, /*drafter=*/true, + qwen35_like_info(), + /*free=*/16 * GiB, /*total=*/24 * GiB, + /*max_ctx=*/131072); + CHECK(!d.enabled); + // Same card at 64K ctx is under the guard — force-on honored. + const auto ok = resolve_skip_park(SkipParkMode::On, /*drafter=*/true, + qwen35_like_info(), + /*free=*/16 * GiB, /*total=*/24 * GiB, + /*max_ctx=*/65536); + CHECK(ok.enabled); +} + +TEST_CASE(SkipParkGuardFixture, T11_auto_fits_with_margin) { + // required = (1.63 GiB + 0.375 GiB + 15000·131072) × 1.25 ≈ 4.9 GiB. + const auto d = resolve_skip_park(SkipParkMode::Auto, /*drafter=*/true, + qwen35_like_info(), + /*free=*/8 * GiB, /*total=*/32 * GiB, + /*max_ctx=*/131072); + CHECK(d.enabled); + CHECK(d.window_tokens == 131072); + CHECK(d.required_bytes > 0); +} + +TEST_CASE(SkipParkGuardFixture, T12_auto_exceeds_free) { + const auto d = resolve_skip_park(SkipParkMode::Auto, /*drafter=*/true, + qwen35_like_info(), + /*free=*/1 * GiB, /*total=*/32 * GiB, + /*max_ctx=*/131072); + CHECK(!d.enabled); +} + +TEST_CASE(SkipParkGuardFixture, T13_auto_unknown_footprint_off) { + const auto d = resolve_skip_park(SkipParkMode::Auto, /*drafter=*/true, + SkipParkDrafterInfo{}, + /*free=*/32 * GiB, /*total=*/32 * GiB, + /*max_ctx=*/131072); + CHECK(!d.enabled); +} + +TEST_CASE(SkipParkGuardFixture, T14_no_drafter_off_in_all_modes) { + CHECK(!resolve_skip_park(SkipParkMode::Auto, false, qwen35_like_info(), + 32 * GiB, 32 * GiB, 131072).enabled); + CHECK(!resolve_skip_park(SkipParkMode::On, false, qwen35_like_info(), + 32 * GiB, 32 * GiB, 131072).enabled); +} + +TEST_CASE(SkipParkGuardFixture, T15_window_capped_by_drafter_ctx) { + // max_ctx above the drafter's native window must not inflate the estimate. + const auto d = resolve_skip_park(SkipParkMode::Auto, /*drafter=*/true, + qwen35_like_info(), + /*free=*/64 * GiB, /*total=*/80 * GiB, + /*max_ctx=*/1048576); + CHECK(d.window_tokens == 262144); +} + +TEST_CASE(SkipParkGuardFixture, T16_estimate_margin_and_guard_on_auto) { + const auto info = qwen35_like_info(); + const int64_t raw = info.weights_bytes + info.fixed_bytes + + info.runtime_bytes_per_token * 65536; + CHECK(skip_park_required_bytes(info, 65536) == + raw + raw * kSkipParkSafetyMarginPercent / 100); + // Guard still applies under auto: <32 GiB with ctx>64K resolves off even + // when the footprint would fit. + const auto d = resolve_skip_park(SkipParkMode::Auto, /*drafter=*/true, + info, /*free=*/20 * GiB, + /*total=*/24 * GiB, /*max_ctx=*/131072); + CHECK(!d.enabled); +} + +// ── VMM guard scope ─────────────────────────────────────────────────────── + +TEST_CASE(SkipParkGuardFixture, T17_guard_only_with_vmm_pool) { + // The R9700 reports 34.2 GB = 31.86 GiB: under the 32 GiB line. Without + // the VMM pool (HIP default) the guard does not apply. + const size_t r9700 = 34208743424ull; + CHECK(!skip_park_allowed(true, r9700, 131072, /*vmm_pool=*/true)); + CHECK(skip_park_allowed(true, r9700, 131072, /*vmm_pool=*/false)); + CHECK(!skip_park_allowed(false, r9700, 131072, /*vmm_pool=*/false)); + // A 24 GiB CUDA card with the VMM pool keeps the guard. + CHECK(!skip_park_allowed(true, 24 * GiB, 131072, /*vmm_pool=*/true)); + + const auto on_hip = resolve_skip_park(SkipParkMode::On, true, + qwen35_like_info(), 16 * GiB, + int64_t(r9700), 131072, + /*vmm_pool=*/false); + CHECK(on_hip.enabled); + const auto on_cuda = resolve_skip_park(SkipParkMode::On, true, + qwen35_like_info(), 16 * GiB, + int64_t(r9700), 131072, + /*vmm_pool=*/true); + CHECK(!on_cuda.enabled); +} + +// ── Keep-loaded decision ────────────────────────────────────────────────── + +TEST_CASE(SkipParkGuardFixture, T18_keep_loaded_only_from_auto_estimate) { + const auto info = qwen35_like_info(); + // Ample: window and keep-loaded footprints both fit. + const auto ample = resolve_skip_park(SkipParkMode::Auto, true, info, + /*free=*/16 * GiB, 32 * GiB, 32768); + CHECK(ample.enabled); + CHECK(ample.keep_drafter_loaded); + CHECK(ample.keep_loaded_bytes == + skip_park_with_margin(skip_park_raw_bytes(info, 32768) + + kTargetComputeReserveBytes)); + CHECK(ample.keep_loaded_bytes > ample.required_bytes); + + // Free VRAM between the two: skip-park on, drafter still released. + const int64_t between = (ample.required_bytes + ample.keep_loaded_bytes) / 2; + const auto tight = resolve_skip_park(SkipParkMode::Auto, true, info, + between, 32 * GiB, 32768); + CHECK(tight.enabled); + CHECK(!tight.keep_drafter_loaded); + + // Explicit on runs no estimate, so it never keeps the drafter loaded. + const auto forced = resolve_skip_park(SkipParkMode::On, true, info, + /*free=*/64 * GiB, 80 * GiB, 32768); + CHECK(forced.enabled); + CHECK(!forced.keep_drafter_loaded); + CHECK(!resolve_skip_park(SkipParkMode::Off, true, info, 64 * GiB, + 80 * GiB, 32768).keep_drafter_loaded); +} + +// ── Estimator on a synthetic drafter header ─────────────────────────────── + +namespace { + +std::string write_drafter_header(const char * arch, bool full) { + gguf_context * ctx = gguf_init_empty(); + gguf_set_val_str(ctx, "general.architecture", arch); + const std::string a = arch; + if (full) { + gguf_set_val_u32(ctx, (a + ".block_count").c_str(), 24); + gguf_set_val_u32(ctx, (a + ".embedding_length").c_str(), 1024); + gguf_set_val_u32(ctx, (a + ".attention.head_count").c_str(), 8); + gguf_set_val_u32(ctx, (a + ".attention.head_count_kv").c_str(), 2); + gguf_set_val_u32(ctx, (a + ".attention.key_length").c_str(), 256); + gguf_set_val_u32(ctx, (a + ".full_attention_interval").c_str(), 4); + gguf_set_val_u32(ctx, (a + ".context_length").c_str(), 262144); + } + const auto path = std::filesystem::temp_directory_path() / + ("luce-skippark-" + std::to_string(getpid()) + "-" + a + + (full ? "-full" : "-bare") + ".gguf"); + gguf_write_to_file(ctx, path.c_str(), false); + gguf_free(ctx); + return path.string(); +} + +} // namespace + +TEST_CASE(SkipParkGuardFixture, T19_estimator_qwen35_header) { + const std::string full = write_drafter_header("qwen35", true); + const std::string bare = write_drafter_header("qwen35", false); + const std::string other = write_drafter_header("qwen3", true); + + SkipParkDrafterInfo q8, q128, no_sessions, defaults, unknown; + CHECK(inspect_drafter_footprint(full, 8, 2, q8)); + CHECK(inspect_drafter_footprint(full, 128, 2, q128)); + CHECK(inspect_drafter_footprint(full, 128, 0, no_sessions)); + CHECK(inspect_drafter_footprint(bare, 128, 2, defaults)); + CHECK(!inspect_drafter_footprint(other, 128, 2, unknown)); + CHECK(!unknown.recognized); + + CHECK(q8.recognized && q8.context_length == 262144); + // The query window's logits cost (n_head + 1)·4 bytes per query token. + CHECK(q128.runtime_bytes_per_token - q8.runtime_bytes_per_token == + int64_t(128 - 8) * (8 + 1) * 4); + // Kept sessions (×2, sized ×1.5) dominate a scratch session. + CHECK(no_sessions.runtime_bytes_per_token < q128.runtime_bytes_per_token); + CHECK(no_sessions.fixed_bytes < q128.fixed_bytes); + // A header without dims falls back to the Qwen3.5-0.8B shape. + CHECK(defaults.runtime_bytes_per_token == q128.runtime_bytes_per_token); + + // Strict scorer per token: 2 sessions × 1.5 × (KV of blocks 0..14 + + // f32 keys) + activations + logits/mask + probe. + ggml_type k = GGML_TYPE_Q4_0, v = GGML_TYPE_Q4_0; + luce::resolve_kv_types(k, v); + const int64_t session = int64_t(luce::kv_reservation_bytes_per_token( + 15, 4, 2, k, 256, v, 256)) + 256 * 2 * 4; + const int64_t strict = 2 * session * 3 / 2 + 2 * 1024 * 4 + + 128 * 9 * 4 + 8; + CHECK(q128.runtime_bytes_per_token >= strict); + + for (const auto & p : {full, bare, other}) std::filesystem::remove(p); +} + +// ── Runtime fail-safe ───────────────────────────────────────────────────── + +namespace { + +// Fake window: each attempt pops the next outcome; records park flags. +struct FakeWindows { + std::vector script; + std::vector parks; + int drops = 0; + SkipParkWindowOutcome run(bool park) { + parks.push_back(park); + if (script.empty()) return SkipParkWindowOutcome::Ok; + const auto next = script.front(); + script.erase(script.begin()); + return next; + } +}; + +SkipParkWindowOutcome window(bool skip, SkipParkFallback & fb, + FakeWindows & fake) { + return run_skip_park_window( + skip, fb, [&](bool park) { return fake.run(park); }, + [](SkipParkWindowOutcome o) { return o; }, + [&]() { ++fake.drops; }, "[test]"); +} + +} // namespace + +TEST_CASE(SkipParkGuardFixture, T20_non_oom_failure_does_not_retry) { + // e.g. non-finite scoring-head scores: parking cannot fix it. + SkipParkFallback fb; + FakeWindows fake{{SkipParkWindowOutcome::Failed}}; + CHECK(window(true, fb, fake) == SkipParkWindowOutcome::Failed); + CHECK(fake.parks == std::vector{false}); + CHECK(fake.drops == 0); + CHECK(!fb.parking_forced()); + CHECK(!fb.memory_tight()); +} + +TEST_CASE(SkipParkGuardFixture, T21_oom_retries_parked_once) { + SkipParkFallback fb; + FakeWindows fake{{SkipParkWindowOutcome::OutOfMemory, + SkipParkWindowOutcome::Ok}}; + CHECK(window(true, fb, fake) == SkipParkWindowOutcome::Ok); + CHECK((fake.parks == std::vector{false, true})); + CHECK(fake.drops == 1); + CHECK(fb.memory_tight()); + CHECK(fb.forced_windows() == SkipParkFallback::kInitialBackoffWindows); + + // A parked retry that fails too returns its failure; no latch change + // beyond the earlier one, and no second retry. + SkipParkFallback fb2; + FakeWindows fake2{{SkipParkWindowOutcome::OutOfMemory, + SkipParkWindowOutcome::OutOfMemory}}; + CHECK(window(true, fb2, fake2) == SkipParkWindowOutcome::OutOfMemory); + CHECK(fake2.parks.size() == 2); + CHECK(!fb2.parking_forced()); +} + +TEST_CASE(SkipParkGuardFixture, T22_latch_backs_off_then_reprobes) { + SkipParkFallback fb; + FakeWindows fake{{SkipParkWindowOutcome::OutOfMemory, + SkipParkWindowOutcome::Ok}}; + window(true, fb, fake); + fake.parks.clear(); + // The next kInitialBackoffWindows windows park outright. + for (int i = 0; i < SkipParkFallback::kInitialBackoffWindows; ++i) { + CHECK(window(true, fb, fake) == SkipParkWindowOutcome::Ok); + } + CHECK(fake.parks == + std::vector(SkipParkFallback::kInitialBackoffWindows, true)); + CHECK(fb.memory_tight()); // no clean no-park window yet + // Then skip-park is probed again; OOM again doubles the backoff. + fake.parks.clear(); + fake.script = {SkipParkWindowOutcome::OutOfMemory, + SkipParkWindowOutcome::Ok}; + window(true, fb, fake); + CHECK((fake.parks == std::vector{false, true})); + CHECK(fb.forced_windows() == 2 * SkipParkFallback::kInitialBackoffWindows); + // Drain, then a clean no-park window resets everything. + for (int i = 0; i < 2 * SkipParkFallback::kInitialBackoffWindows; ++i) { + window(true, fb, fake); + } + fake.parks.clear(); + CHECK(window(true, fb, fake) == SkipParkWindowOutcome::Ok); + CHECK(fake.parks == std::vector{false}); + CHECK(!fb.memory_tight()); + CHECK(fb.backoff() == 0); +} + +TEST_CASE(SkipParkGuardFixture, T23_backoff_caps_and_park_requests_untouched) { + SkipParkFallback fb; + for (int i = 0; i < 10; ++i) fb.on_oom_recovered(); + CHECK(fb.backoff() == SkipParkFallback::kMaxBackoffWindows); + // A window that asks to park never consults or advances the latch. + FakeWindows fake; + const int before = fb.forced_windows(); + CHECK(window(false, fb, fake) == SkipParkWindowOutcome::Ok); + CHECK(fake.parks == std::vector{true}); + CHECK(fb.forced_windows() == before); }