diff --git a/optimizations/kvflash/DESIGN.md b/optimizations/kvflash/DESIGN.md index 965dde7098..856c53a75f 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 ed3d6cb75a..9a86018ee3 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/unit/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 b53638c96a..fb35b2c802 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 ``` @@ -97,10 +97,12 @@ For an OpenAI-compatible server with transparent compression on long prompts, ru | Flag | Choices / type | Default | Effect | |---|---|:---:|---| | `--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-threshold` | int (tokens) | `32000` | Token threshold for `auto` mode. `16384` with a trained scoring head. | +| `--prefill-keep-ratio` | float `(0, 1]` | `0.05` | Fraction of source tokens to keep after compression. `0.02` for 128K, `0.10` for 32K. With a trained scoring head and no ratio flag, the release curve `16384:0.315 32768:0.270 65536:0.270 120000:0.225` applies. | +| `--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. | + +With `PFLASH_SCORING_HEAD_GGUF` set, every unset PFlash setting takes the measured release profile; see [PFlash configuration](../../server/docs/PFLASH_CONFIGURATION.md#release-profile). 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 +113,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 +157,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 +179,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 +207,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 +244,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 +264,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 89c2148924..942c6bba46 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 e289d25799..8f4dbb85a6 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 39db4f1cb9..c307e5588a 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 d7ab0fd957..bb995815f1 100644 --- a/server/CMakeLists.txt +++ b/server/CMakeLists.txt @@ -197,9 +197,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. @@ -570,6 +570,7 @@ set(LUCE_SRC_INCLUDE_DIRS ${CMAKE_CURRENT_SOURCE_DIR}/src/qwen4exp ${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 @@ -591,11 +592,21 @@ 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_fold.cpp + src/pflash/pflash_request_probe.cpp + src/pflash/pflash_request_stage2.cpp + src/pflash/pflash_query_candidates.cpp + src/pflash/pflash_scoring_resume.cpp + src/pflash/pflash_session_offload.cpp + src/pflash/pflash_ident_query.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 @@ -830,7 +841,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. @@ -1088,8 +1099,8 @@ endif() # PFlash phase-split harness -if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/smoke/pflash_daemon.cpp") - add_executable(pflash_daemon test/smoke/pflash_daemon.cpp) +if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/pflash/pflash_daemon.cpp") + add_executable(pflash_daemon test/pflash/pflash_daemon.cpp) target_include_directories(pflash_daemon PRIVATE ${LUCE_SRC_INCLUDE_DIRS}) target_link_libraries(pflash_daemon PRIVATE luce_common ggml ${LUCE_GGML_BACKEND_TARGET}) if(LUCE_GPU_BACKEND STREQUAL "cuda") @@ -1121,7 +1132,7 @@ if(LUCE_TESTS) set(_raw_unit_test_targets) enable_testing() - # Test sources live under test/{unit,smoke,bench}/ while the shared + # Test sources live under test/{unit,smoke,bench,pflash}/ while the shared # helpers (CppUnitTestFramework.hpp, host_check.h, ...) stay at test/. # Put the helper dir on every test target's include path. include_directories(${CMAKE_CURRENT_SOURCE_DIR}/test) @@ -1609,6 +1620,12 @@ if(LUCE_TESTS) SKIP_RETURN_CODE 77 ENVIRONMENT "GGML_CUDA_DISABLE_COPY_BATCH=1") endif() + # GPU F32 copy into a strided view (the PFlash scorer's chunked logits + # copy; a one-key-wide last chunk used to take the dense transpose + # kernel). CUDA and HIP; exits 77 without a device. + luce_add_ggml_gpu_executable(test_cpy_strided_dst test/unit/test_cpy_strided_dst.cpp) + target_link_libraries(test_cpy_strided_dst PRIVATE ggml-cpu) + list(APPEND _raw_unit_test_targets test_cpy_strided_dst) if(LUCE_GPU_BACKEND STREQUAL "hip") add_executable(test_ggml_host_mailbox_cuda test/unit/test_ggml_host_mailbox_cuda.cpp) set_source_files_properties(test/unit/test_ggml_host_mailbox_cuda.cpp PROPERTIES LANGUAGE HIP) @@ -2273,6 +2290,25 @@ if(LUCE_TESTS) endif() endif() + # PFlash fold replay: the runtime's fold selection over recorded traces + # (CPU only; the drafter GGUF supplies the vocabulary). + add_executable(pflash_fold_replay test/pflash/pflash_fold_replay.cpp) + target_include_directories(pflash_fold_replay PRIVATE ${LUCE_SRC_INCLUDE_DIRS}) + if(LUCE_GPU_BACKEND STREQUAL "hip") + target_compile_definitions(pflash_fold_replay PRIVATE LUCE_BACKEND_HIP=1 GGML_USE_HIP) + else() + target_compile_definitions(pflash_fold_replay PRIVATE + LUCE_BACKEND_CUDA=1 + LUCE_CUDA_MIN_SM=${_luce_cuda_min_sm}) + endif() + target_link_libraries(pflash_fold_replay PRIVATE luce_common ggml ${LUCE_GGML_BACKEND_TARGET}) + if(LUCE_GPU_BACKEND STREQUAL "cuda") + find_package(CUDAToolkit REQUIRED) + target_link_libraries(pflash_fold_replay PRIVATE CUDA::cudart) + else() + target_link_libraries(pflash_fold_replay PRIVATE hip::host) + endif() + # ─── Unit tests (no GPU, no model files) ──────────────────────────── # Shared image input: request transport and JPEG/PNG decoding. add_executable(test_image_input test/unit/test_image_input.cpp src/server/image_input.cpp) @@ -2443,8 +2479,16 @@ if(LUCE_TESTS) test/unit/test_chain_rollback_policy.cpp test/unit/test_ddtree_tau.cpp test/unit/test_anchor_transitive.cpp - test/unit/test_drafter_early_exit_score_range.cpp - test/unit/test_qwen3_buffer_plan.cpp + test/pflash/test_pflash_drafter_ipc.cpp + test/pflash/test_pflash_selection.cpp + test/pflash/test_pflash_fold.cpp + test/pflash/test_pflash_request_probe.cpp + test/pflash/test_pflash_request_stage2.cpp + test/pflash/test_pflash_query_candidates.cpp + test/pflash/test_pflash_scoring_resume.cpp + test/pflash/test_pflash_session_offload.cpp + test/pflash/test_pflash_ident_query.cpp + test/unit/test_qwen35_causal_mask.cpp test/unit/test_model_test_paths.cpp test/unit/test_gguf_mmap.cpp test/unit/test_kv_quant.cpp @@ -2469,7 +2513,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) target_include_directories(test_server_unit PRIVATE ${LUCE_SRC_INCLUDE_DIRS} ${CMAKE_CURRENT_SOURCE_DIR}/deps/llama.cpp/ggml/src diff --git a/server/README.md b/server/README.md index 4c1185fa37..539a06fdc3 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 \ @@ -387,17 +387,197 @@ the whole request's device footprint. `/status/json` reports ### PFlash +With a trained scoring head (`PFLASH_SCORING_HEAD_GGUF`), every PFlash setting +left unset takes the measured release profile: threshold 16384, the curve +`16384:0.315 32768:0.270 65536:0.270 120000:0.225`, budget selection over probe +segments and the request-probe query. Without a head the built-in defaults +below apply. The [PFlash configuration reference](docs/PFLASH_CONFIGURATION.md) +lists both, the validated overrides and the active experiments. + | Option | Default | Purpose | |---|---|---| | `--prefill-compression off\|auto\|always` | `off` | Select when PFlash compresses a prompt. | -| `--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-threshold ` | `32000` (`16384` with a head) | Token threshold used by auto mode. | +| `--prefill-keep-ratio ` | `0.05` (the release curve with a head) | Fraction of source tokens kept. | +| `--prefill-curve T:R [T:R ...]` | none (the release curve with a head) | Piecewise keep-ratio curve; overrides the flat ratio. | +| `--prefill-drafter ` | none | PFlash drafter GGUF (Qwen3.5-0.8B). | | `--prefill-skip-park` | off | Keep target and decode draft resident while PFlash runs. | | `--prefill-upstream-base ` | none | Enable compression-proxy mode. | | `--prefill-upstream-key ` | none | Bearer token for the upstream. | | `--prefill-upstream-model ` | none | Model name forwarded upstream. | +| `PFLASH_SELECT_MODE=top_k` + `PFLASH_SELECT_TOPK ` | unset: legacy selection (`budget_only` with a head) | Rank rule: keep the K highest-scoring optional segments in score order instead of filling the keep ratio, with the keep-ratio budget still a hard ceiling (min(K segments, the budget)). Use it where the evidence is compact and sits in the first few ranks -- needle retrieval, passage QA, code -- so a small K reaches it for a fraction of the budget's tokens. Do not use it where the answer needs a whole document identified, since the evidence there spans many segments and K cuts it off. | +| `PFLASH_SELECT_QUERY_HEAD=1` (with `PFLASH_SELECT_QUERY_PARSER=latest_user`) | off (on with a head) | The question may come before pasted material. The user's leading prose block (up to the first paste-like line: a code fence, heading, separator, record/file header, tag, JSON or table data, indented or code-shaped line) is scored as an extra query window and kept verbatim; the trailing prose block after the last paste becomes the query window instead of the fixed tail. Caps: `PFLASH_SELECT_QUERY_HEAD_TOKENS` / `PFLASH_SELECT_QUERY_TAIL_TOKENS` (256); a longer block contributes its first and last cap/2 tokens. Blocks under 16 tokens are ignored; a message with no paste structure keeps today's window, except that a short opening paragraph ending in `?` before a long plain text counts as the head. | +| `PFLASH_SELECT_STRATEGY=fold` (with `PFLASH_SELECT_MODE=budget_only`) | `segments` | Fold selection: each of the top 20 segments (`PFLASH_SELECT_FOLD_K`) is completed to its smallest enclosing block -- a function or class in code (Python indentation or brace-matched C/C++/JS/TS/Go/Rust bodies), a numbered record such as `Document N:`, a markdown section, else a blank-line paragraph (`PFLASH_SELECT_FOLD_PARAGRAPHS=0` turns that last fallback off). A block over 1500 tokens (`PFLASH_SELECT_FOLD_CAP`) keeps its first 150 (`PFLASH_SELECT_FOLD_HEAD`) plus the segment. Code-like prompts (most non-blank lines inside functions; JSON and markdown are not) rank segments by mass (density x length), others by density. The keep-ratio budget stays the ceiling. Needs the head scorer; the drafter loads its vocabulary to read the prompt's structure. | + +With compression on, `--max-ctx` is also the longest raw prompt PFlash takes, +and a longer one is refused with HTTP 400. A prompt that fits but leaves too +little room for its reply is compressed, and the reply is checked against the +compressed prompt. The drafter scores the whole raw prompt, and its scoring +memory grows with it, so set `--max-ctx` to what the card holds beside the +target: on a 32 GB R9700 with the 27B IQ4_XS target resident (q4_0 KV, +`--prefill-skip-park`), `--max-ctx 204800` fits, with a 204K-token prompt +peaking at 30.6 of 31.9 GiB. If the drafter does run out of memory, the +request is served from the raw prompt with full prefill (`[pflash-fallback] +... reason=compress_failed`). + +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 what a fresh +selection for the new question keeps that the view lacks: a question on the +view's topic misses little and is served by appending it, and a question +that needs more than a third of a fresh selection -- the conversation moved +to other material -- starts a new view from that selection instead. Kept pieces that were not adjacent in the prompt are joined by a +paragraph break when neither side has one (`PFLASH_SELECT_PARAGRAPH_JOIN=0` +turns it off; the breaks do not count against the token ceiling). +`PFLASH_VIEW_TRACE_PATH` +appends each compressed request's served prompt as JSONL, for evidence +checks in evaluations. + +Coding-agent sessions (with `PFLASH_ADMISSION=cache` and +`--prefill-compression auto`): in a conversation with tool activity (the +request defines `tools`, or its history has an assistant tool call or a tool +result) PFlash never compresses the agent's history or any tool output. By +default it compresses nothing in such a session: every request is served raw +with full prefill (`pflash_fallback.reason = keep_history`), and a raw prompt +over the context gets the context-overflow error full prefill gets. +`PFLASH_AGENT_PASTE=1` (default off) lets PFlash compress one thing: the +newest turn when it is a user message with at least `--prefill-threshold` +droppable tokens (a pasted document, on the first turn too). That content +alone is compressed, on its own keep budget, after the raw history (restored +from the prefix cache or prefilled) or after the conversation's view. A view +is then only continued: later turns and tool outputs are appended verbatim, a +new big paste is compressed as a region of its own, and recall adds the +passages of earlier pastes that the new question needs in the excerpt block +at the start of the newest user message, after the cached view. The view is +never rebuilt and never compressed again as a whole; when it cannot continue +(or no longer fits the context) the raw prompt is served. The switch is off +by default because follow-up questions after a compressed paste cost more +than full prefill, which keeps the whole raw prompt cached: a follow-up that +needs recall prefills the recalled passages again, and every user turn runs a +drafter selection (up to 48 s more per follow-up in a 10-turn coding session on +an R9700; the session took 281 s against 258 s with full prefill). Chats +without tools are not affected. + +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 at the last 8192-token grid point at or before 64 tokens from +its end. A prompt that shares that prefix runs only its new tokens through +the drafter, from the checkpoint, and the new query scores against every +stored key. Every call keeps a fresh call's partition (ubatches and key +chunks counted from token 0), so a resumed selection is the one a fresh pass +would make; the query step reads its rows from this call or from an exact F32 +store, and a request that cannot resume exactly is scored from token 0 with +the reason logged (`[qwen35-scorer] redo from token 0`). +With the default `--draft-residency auto` the drafter is released after each +compression; the sessions a call used are copied to host RAM first (plain +pageable memory; `PFLASH_SESSION_OFFLOAD_PINNED=1` uses the device's pinned +host memory instead, which on an R9700 raised peak VRAM by up to 1 GiB on +220K-token prompts, kept about 0.5 GiB of VRAM in use between requests and +slowed target prefill), only the positions a resume reads, and when the drafter loads again +only the session the new prompt can resume from comes back to the GPU. A +stored session is restored only under the drafter, head, probe and session +settings it was stored with; anything else drops it, and the call scores +from token 0. +`--draft-residency persistent` keeps the sessions on the GPU and copies +nothing. +A new session is sized with headroom (S + S/2 + 4096 tokens for an S-token +prompt) so later turns append without reallocating, but never past the +engine's `--max-ctx` or the drafter's trained context, and never below the +prompt itself. +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 @@ -542,9 +722,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' @@ -566,9 +746,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/deps/llama.cpp/ggml/src/ggml-cuda/cpy.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/cpy.cu index e3586763b6..c0dccc97f8 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/cpy.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/cpy.cu @@ -625,8 +625,13 @@ void ggml_cuda_cpy(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, gg char * src1_ddc = (char *) src1->data; const bool contiguous_srcs = ggml_is_contiguous(src0) && ggml_is_contiguous(src1); + // The tiled transpose kernel writes dst densely (dst[imat*n + row*ne00 + col]) and ignores + // src1's strides, so it is only valid for a contiguous destination. Without this guard a + // contiguous [1, n, m] source (ne00 == 1 also gives nb01 == element size) copied into a + // strided view lands in the wrong elements and leaves the right ones unwritten. const bool can_be_transposed = nb01 == (int64_t)ggml_element_size(src0) && - src0->ne[3] == 1 && nb02 == ne00 * ne01 * (int64_t)ggml_element_size(src0); + src0->ne[3] == 1 && nb02 == ne00 * ne01 * (int64_t)ggml_element_size(src0) && + ggml_is_contiguous(src1); if (src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_Q8_1) { // The q8_1 rows MMVQ reads, from the kernel it quantizes its own diff --git a/server/docs/DS4.md b/server/docs/DS4.md index 29b9a77299..46cde365a5 100644 --- a/server/docs/DS4.md +++ b/server/docs/DS4.md @@ -1005,7 +1005,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 4ea943c664..078215e328 100644 --- a/server/docs/ENVIRONMENT.md +++ b/server/docs/ENVIRONMENT.md @@ -12,6 +12,10 @@ Anything else in the inventory below is legacy surface: prefer the CLI flag where one exists, and treat undocumented variables as internal. The consolidation of this list into CLI flags is tracked as follow-up work. +PFlash's existing tuning, serving and evaluation overrides are documented in +[PFlash configuration](PFLASH_CONFIGURATION.md), including the measured release +profile and the requirements checked at startup. + ## Documented variables | Variable | Default | Purpose | @@ -435,7 +439,5 @@ consolidation of this list into CLI flags is tracked as follow-up work. - `LUCE_PC_DEEP_FIRST_MIN` - prefix_cache.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/PFLASH_CONFIGURATION.md b/server/docs/PFLASH_CONFIGURATION.md new file mode 100644 index 0000000000..2b8b941cc7 --- /dev/null +++ b/server/docs/PFLASH_CONFIGURATION.md @@ -0,0 +1,128 @@ +# PFlash configuration + +With a trained scoring head (`PFLASH_SCORING_HEAD_GGUF`) and prefill +compression on, every setting left unset takes the measured +[release profile](#release-profile); without a head the built-in defaults +apply. Fold, query-grid, stage-2 and oracle experiments remain available. +With prefill compression enabled, selection, query, admission and chat +settings are checked before model loading. Invalid settings stop startup and name the offending +variable. Model schemas and replay files receive their loader checks. + +The [development tools](../scripts/pflash/README.md) document request-probe +capture, validation and export; [PFlash tests](../test/pflash/README.md) describe +the fixtures and replay executables. + +Serving booleans accept `0` or `1`; unset or empty serving overrides retain their +defaults. Selection and query overrides reject empty values unless explicitly +allowed below. Valid integer overrides retain their existing upper clamps. + +## Serving defaults + +| Setting | Built-in default | Release profile | Values / behavior | +|---|---|---|---| +| `--prefill-compression` | `off` | | `off`, `auto`, `always` | +| `--prefill-threshold` | `32000` | `16384` | Auto admission threshold, in droppable tokens | +| `--prefill-keep-ratio` | `0.05` | curve | Flat ratio; `--prefill-curve tokens:ratio ...` overrides it. The release curve is `16384:0.315 32768:0.270 65536:0.270 120000:0.225`; either flag replaces it | +| `--prefill-skip-park` | off | | Keep the target resident while the drafter scores; see the release profile | +| `PFLASH_ADMISSION` | `length` | `cache` | `length` or `cache`; cache admission accounts for restorable target KV | +| `PFLASH_AGENT_PASTE` | `0` | | With cache admission in auto mode, `1` permits a large newest user paste in a tool session; earlier history stays raw | +| `PFLASH_CHAT_VIEW`, `PFLASH_CHAT_RECALL` | `1`, `1` | | Continue the served view; recall omitted material for a new question | +| `PFLASH_CHAT_COMPRESS_NEW_TOKENS` | `16384` | | Integer >= 1; smaller additions stay verbatim; capped at 2^30 | +| `PFLASH_CHAT_SKELETON_TOKENS` | `256` | | Integer >= 0; capped at 2^20 | +| `PFLASH_CHAT_HISTORY_QUERIES` | `3` | | Integer >= 0; capped at 8 earlier user turns | +| `PFLASH_CHAT_REBUILD_RATIO` | `2` | | Finite number >= 0; `0` rebuilds only at the context limit | +| `PFLASH_CHAT_REBUILD_MIN_TOKENS` | `0` | `65536` | Integer >= 0; smallest view eligible for drift rebuild; capped at 2^30 | +| `PFLASH_ARRIVE` | `0` | `1` | Score only a large newly arriving region | +| `PFLASH_SELECT_PARAGRAPH_JOIN` | `1` | `0` | Insert paragraph breaks between nonadjacent kept spans | +| `PFLASH_FALLBACK_SPREAD` | off | `0.11` | Positive finite threshold; unset, empty, `0` or `off` disables it. Applies to whole-prompt selections, not region plans | +| `PFLASH_FAIL_CLOSED` | `0` | | `0` serves raw after compression failure when it fits; `1` returns the error. Keep-history admission still serves raw | +| `PFLASH_DRAFTER_SESSIONS` | `2` | | Integer >= 0; capped at 16; `0` scores from scratch. Stage 2 uses its own rows and ignores this setting | +| `PFLASH_DRAFTER_TRIM_POOL` | `1` | | Release cached temporary allocations from a resident drafter's pool | +| `PFLASH_SESSION_OFFLOAD_PINNED` | `0` | | Plain host memory; `1` opts into pinned host allocations | + +## Selection and model assets + +| Variable | Built-in default | Release profile | Values / requirements | +|---|---|---|---| +| `PFLASH_SCORING_HEAD_GGUF` | unset | | Trained head file; otherwise built-in attention-mass scorer. Setting it selects the release profile | +| `PFLASH_SEGMENT_PROBE_GGUF` | unset | | Segment probe file; otherwise fixed chunks | +| `PFLASH_REQUEST_PROBE_GGUF` | unset | | Request probe file; required for query-probe selection outside the test stub | +| `PFLASH_SELECT_MODE` | unset: legacy | `budget_only` | `budget_only`, `top_p`, `top_k`; setting a mode enables strict selection | +| `PFLASH_SELECT_CHUNK_SIZE` | scheduled in strict mode | `256` | Positive integer; 128 below 500 input tokens, 512 below 3000, otherwise 1024 | +| `PFLASH_SELECT_QUERY_TOKENS` | `8` | `128` | Integer in [1, 512] | +| `PFLASH_SELECT_QUERY_PARSER` | `arbitrary_tail` | | `arbitrary_tail` or `latest_user`; default uses the latest user turn in chat | +| `PFLASH_SELECT_QUERY_HEAD` | `0` | `1` | `1` adds prose around pasted material; needs strict selection | +| `PFLASH_SELECT_QUERY_HEAD_TOKENS`, `PFLASH_SELECT_QUERY_TAIL_TOKENS` | `256` each | | Integers in [2, 512] | +| `PFLASH_SELECT_QUERY_PROBE` | `0` | `1` with the request probe | `1` replaces scoring windows with request-probe query tokens; needs strict head selection | +| `PFLASH_SELECT_SEGMENTS` | `auto` | `probe` with the segment probe | `auto`, `fixed`, `probe`; auto uses probe segments when loaded and the probe finds at least 4 boundaries; probe always does | +| `PFLASH_SELECT_SCORE` | `auto` | `density` | `auto`, `sum`, `density`; auto uses density for probe segments, sum for fixed chunks | +| `PFLASH_SELECT_SCORER`, `PFLASH_SELECT_SPLIT` | `head`, `0.5` | | `head`, `legacy`, `split`; split fraction finite and in (0, 1) | +| `PFLASH_SELECT_TOP_P`, `PFLASH_SELECT_TOPK` | `0.95`, unset | Top-p finite in (0, 1]; top-k positive integer, required in `top_k` mode | +| `PFLASH_SELECT_STRATEGY` | `segments` | `segments` or `fold`; fold needs `budget_only` and the head scorer | +| `PFLASH_SELECT_FOLD_K`, `PFLASH_SELECT_FOLD_CAP`, `PFLASH_SELECT_FOLD_HEAD` | `20`, `1500`, `150` | K and cap >= 1, head >= 0; with fold, head cannot exceed cap | +| `PFLASH_SELECT_FOLD_PARAGRAPHS` | `1` | Paragraph completion as fold's last fallback | + +Query-probe, identifier and spread switches modify selection; they do not +enable strict selection on their own. + +## Relative queries and retained experiments + +| Variable | Built-in default | Values / requirements | +|---|---|---| +| `PFLASH_QUERY_CANDIDATES`, `PFLASH_QUERY_PIN` | `0`, `0` | Candidates in [0, 8], pin tokens in [0, 4096]; positive enables the feature. The release profile sets `1` and `512` with both probes, unless request stage 2 is set or the query probe is off | +| `PFLASH_QUERY_UNIT_SCORE`, `PFLASH_QUERY_AGG` | `max`, `softmax` | `max` or `mean`; `softmax` or `wmax` | +| `PFLASH_QUERY_TEMP`, `PFLASH_QUERY_NULL_MARGIN` | `1`, `0` | Finite temperature > 0; finite null margin >= 0 | +| `PFLASH_QUERY_EXTENT`, `PFLASH_QUERY_PIN_MARGIN` | `gap`, `gap` | `gap` or finite number >= 0; pin margin also accepts `inf` | +| `PFLASH_QUERY_TW`, `PFLASH_QUERY_END` | `0.5`, `0.25` | Explicit shares in (0, 1), sum < 1; TW also accepts `score` | +| `PFLASH_REQUEST_PROBE_POLICY`, `PFLASH_SELECT_QUERY_PROBE_RECENCY_PIN` | GGUF policy | JSON policy file; recency pin in [0, 4096]; checked by the loader | +| `PFLASH_REQUEST_STAGE2_GGUF` | unset | Experimental look-back combiner; needs compatible request probe and trained head | +| `PFLASH_EVAL_REF_IDENT_QUERY` | `off` | `off` or `heur`; heuristic evaluation arm needs strict head selection and excludes query-probe selection | +| `PFLASH_QUERY_GRID`, `PFLASH_QUERY_GRID_OUT` | unset | Replay input and append-only result file; grid needs query-probe selection and an output path | +| `PFLASH_QUERY_ORACLE_DIR` | unset | Oracle evaluation output during grid replay | +| `PFLASH_QWEN35_LEGACY_SCORER` | `0` | Force the original Qwen3.5 scorer | +| `PFLASH_SELECT_QUERY_PROBE_STUB` | `0` | Test-only request-probe stub | +| `PFLASH_COMPRESS_ANCHOR_RADIUS`, `PFLASH_COMPRESS_MAX_ANCHOR_HITS` | tier default | Legacy anchor controls | +| `PFLASH_FREEZE_HOT_WINDOW` | `2` | FlowKV's recent uncompressed message window | + +Query tuning overrides require positive candidates or pinning. TW/END blend +overrides additionally require candidates. Candidates, pinning and grid replay +need the segment probe and cannot combine with request stage 2. + +`PFLASH_TRACE_PATH` and `PFLASH_VIEW_TRACE_PATH` append selection and served-view +JSONL. `PFLASH_TRACE_REQUEST_FEATURES`, `PFLASH_TRACE_REQUEST_LOGITS` and +`PFLASH_TRACE_STAGE2_FEATURES` retain the feature/logit diagnostics. + +## Release profile + +The PFlash-v2 measurements (Qwen3.8-27B target, Qwen3.5-0.8B drafter, R9700) +used the values in the release-profile columns above. They are the defaults +whenever prefill compression is on and `PFLASH_SCORING_HEAD_GGUF` is set, so +the measured configuration needs only the model files: + +```sh +export PFLASH_SCORING_HEAD_GGUF=/path/to/head.gguf +export PFLASH_SEGMENT_PROBE_GGUF=/path/to/segment-probe.gguf +export PFLASH_REQUEST_PROBE_GGUF=/path/to/request-probe.gguf + +luce_server /path/to/target.gguf \ + --prefill-drafter /path/to/Qwen3.5-0.8B.gguf \ + --prefill-compression auto --prefill-skip-park +``` + +The server installs the profile's environment variables at startup and logs +how many it applied and which ones it kept from the environment. A variable +that is already set keeps its value, and `--prefill-threshold`, +`--prefill-keep-ratio` or `--prefill-curve` on the command line replace the +profile's threshold or curve. `PFLASH_SELECT_QUERY_PROBE=1` needs the request +probe, and the query candidates and pin need both probes. They also yield to +the experiment arms they cannot combine with: another scorer, the identifier +query, request stage 2, or `PFLASH_SELECT_QUERY_PROBE=0`. Without a head the +built-in defaults apply, as before. + +`--prefill-skip-park` stays opt-in. It keeps the target resident while the +drafter scores, which the measurements used on a 32 GB card at +`--max-ctx 139264`; on smaller cards the two may not fit together. + +The obsolete `PFLASH_COLD_HISTORY` and `PFLASH_AGENT_HISTORY` names have been +removed: cache admission keeps tool history raw, with paste compression behind +`PFLASH_AGENT_PASTE`. The active experimental families above remain intact. diff --git a/server/docs/SPEC_PREFILL.md b/server/docs/SPEC_PREFILL.md index 171541778b..a4cc3a38f5 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 6883d10ff1..6032c6b302 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/smoke/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/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/unit/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/entrypoint.sh b/server/scripts/entrypoint.sh index 6bc7401d36..1c615064dc 100755 --- a/server/scripts/entrypoint.sh +++ b/server/scripts/entrypoint.sh @@ -356,9 +356,11 @@ if [ "$LUCE_PREFILL_MODE" != off ]; then [ -n "${LUCE_PREFILL_DRAFTER:-}" ] || die "LUCE_PREFILL_MODE=$LUCE_PREFILL_MODE requires LUCE_PREFILL_DRAFTER" [ -f "$LUCE_PREFILL_DRAFTER" ] || die "Prefill drafter not found at $LUCE_PREFILL_DRAFTER" CMD+=(--prefill-compression "$LUCE_PREFILL_MODE" - --prefill-keep-ratio "${LUCE_PREFILL_KEEP:-0.05}" - --prefill-threshold "${LUCE_PREFILL_THRESHOLD:-32000}" --prefill-drafter "$LUCE_PREFILL_DRAFTER") + # Unset, luce_server picks the ratio and threshold: the release profile + # with a trained scoring head, else 0.05 and 32000. + [ -n "${LUCE_PREFILL_KEEP:-}" ] && CMD+=(--prefill-keep-ratio "$LUCE_PREFILL_KEEP") + [ -n "${LUCE_PREFILL_THRESHOLD:-}" ] && CMD+=(--prefill-threshold "$LUCE_PREFILL_THRESHOLD") fi # Operator flags go last: luce_server keeps the last value of a repeated flag. diff --git a/server/scripts/laguna_pflash_niah.py b/server/scripts/laguna_pflash_niah.py index c936ab930a..a075e735d9 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/pflash/README.md b/server/scripts/pflash/README.md new file mode 100644 index 0000000000..44456f5f05 --- /dev/null +++ b/server/scripts/pflash/README.md @@ -0,0 +1,53 @@ +# PFlash development tools + +The three tools here cover request-probe capture, export and parity fixtures. +They preserve the existing probe schemas, capture configuration and seeded +fixtures. NumPy is required; GGUF commands use the vendored `gguf-py` package. + +## Capture and validate + +`request_probe_capture.py` combines the former capture and capture-validation +scripts. Its two modes retain the existing arguments, shard format, checks and +`CAPTURE.json` manifest: + +```sh +python server/scripts/pflash/request_probe_capture.py capture \ + --daemon server/build/pflash_daemon --drafter /path/to/drafter.gguf \ + --scoring-head /path/to/head.gguf --segment-probe /path/to/segment.gguf \ + --request-probe /path/to/request.gguf --input conversations.jsonl --out capture + +python server/scripts/pflash/request_probe_capture.py validate capture \ + --probe /path/to/request.gguf --recapture repeated-capture --trunk-convs 20 +``` + +Capture runs one fresh drafter session per conversation and stores user/tool +feature rows and runtime logits. Validation checks finite values, NumPy trunk +parity, labelled-message coverage and optional byte-identical recapture. Its +exit status is nonzero when the checks fail. Capture also retains partitioning, +shard sizing and remote transfer options; see `capture --help`. + +The old top-level `pflash_request_probe_capture.py` and +`pflash_request_probe_capture_validate.py` paths have been replaced by these +subcommands. + +## Export the probe + +```sh +python server/scripts/pflash/request_probe_gguf.py request.gguf \ + --drafter-sha256 HEX --tau THRESHOLD --weights probe.npz +``` + +`--random-seed` remains available for runtime plumbing tests. Those weights are +not a trained probe. The exporter retains the request-probe tensor shapes, +schema and extraction-version metadata. + +## Regenerate parity fixtures + +```sh +python server/scripts/pflash/gen_request_probe_v2_fixtures.py TRAINER_SRC_DIR \ + server/test/pflash/fixtures/request_probe_v2_parity.inc +``` + +The generator imports the trainer's offline rule, checks its candidate-rule +contract and emits the same 80 seed-0 cases. Keeping that offline implementation +as the reference makes the fixtures independent of the C++ rule under test. diff --git a/server/scripts/pflash/gen_request_probe_v2_fixtures.py b/server/scripts/pflash/gen_request_probe_v2_fixtures.py new file mode 100644 index 0000000000..bb4a199b11 --- /dev/null +++ b/server/scripts/pflash/gen_request_probe_v2_fixtures.py @@ -0,0 +1,76 @@ +#!/usr/bin/env python3 +"""Generate test/pflash/fixtures/request_probe_v2_parity.inc: replay-parity fixtures for the V2-A3 candidate rule against the +trainer's offline mirror rule.probe_query_v2 (luce_box lucebox_training/pflash, CANDIDATE_RULE +cand_v3_latest_human_prev8_plus_all_nonsystem_turns, V2-A5). Deterministic (seed 0). + python3 gen_request_probe_v2_fixtures.py TRAINER_SRC_DIR OUT.inc""" +import argparse +import random +import sys + +import numpy as np + + +def main(): + ap = argparse.ArgumentParser(description="Generate deterministic request-probe parity fixtures.") + ap.add_argument("trainer_src_dir", help="Directory containing the trainer's pflash_scorer package") + ap.add_argument("out", help="Output request_probe_v2_parity.inc") + args = ap.parse_args() + sys.path.insert(0, args.trainer_src_dir) + from pflash_scorer.request_probe import contract # noqa: E402 + from pflash_scorer.request_probe.rule import probe_query_v2 # noqa: E402 + + assert contract.CANDIDATE_RULE == "cand_v3_latest_human_prev8_plus_all_nonsystem_turns" + rng = random.Random(0) + cases = [] + for case in range(80): + roles = ["system"] if rng.random() < 0.7 else [] + n_msgs = rng.randint(1, 9) + for _ in range(n_msgs): + roles.append(rng.choice(["user", "user", "assistant", "tool", "tool"])) + big = case % 8 == 0 + spans, pos = [], 0 + for r in roles: + L = rng.randint(0, 900 if big else 60) + spans.append((pos, pos + L)) + pos += L + rng.randint(1, 4) + n = pos + 3 + hits = [] + lg = np.full(n, -1.0, np.float32) + for b, e in spans: + k = b + while k < e: + if rng.random() < (0.35 if big else 0.15): + L = rng.randint(1, 12 if not big else 400) + z = min(e, k + L) + hits.append((k, z)) + lg[k:z] = 1.0 + k = z + rng.randint(1, 5) + else: + k += rng.randint(1, 8) + q, _runs, info = probe_query_v2(lg, 0.0, list(zip(roles, spans))) + qs = [] + for t in q: + if qs and qs[-1][1] == t: + qs[-1][1] = t + 1 + else: + qs.append([t, t + 1]) + cases.append((roles, spans, n, hits, qs, info["candidates"])) + + with open(args.out, "w") as f: + import hashlib + import os + src = os.path.join(args.trainer_src_dir, "pflash_scorer", "request_probe") + shas = {n: hashlib.sha256(open(os.path.join(src, n), "rb").read()).hexdigest()[:16] for n in ("rule.py", "contract.py")} + f.write("// Generated by server/scripts/pflash/gen_request_probe_v2_fixtures.py from the trainer's rule.probe_query_v2\n") + f.write("// trainer sources: rule.py %s, contract.py %s\n" % (shas["rule.py"], shas["contract.py"])) + f.write("// (luce_box v2 branch, CANDIDATE_RULE cand_v3_latest_human_prev8_plus_all_nonsystem_turns). Do not edit.\n") + for i, (roles, spans, n, hits, qs, cands) in enumerate(cases): + f.write("{%d, {%s}, {%s}, {%s}, {%s}, {%s}},\n" % ( + n, ", ".join('"%s"' % r for r in roles), ", ".join("{%d, %d}" % s for s in spans), + ", ".join("{%d, %d}" % h for h in hits), ", ".join("{%d, %d}" % tuple(x) for x in qs), + ", ".join(str(c) for c in cands))) + print(args.out, len(cases), "cases") + + +if __name__ == "__main__": + main() diff --git a/server/scripts/pflash/request_probe_capture.py b/server/scripts/pflash/request_probe_capture.py new file mode 100644 index 0000000000..acef8f0f07 --- /dev/null +++ b/server/scripts/pflash/request_probe_capture.py @@ -0,0 +1,351 @@ +#!/usr/bin/env python3 +"""Request-probe FEATURE CAPTURE driver (SPEC A6): block-14 rows from the pinned runtime. + +Runs the drafter-only pflash_daemon in the SERVED config (Q4_0 KV default, segment probe + u314 head, strict +selection) with the request probe on, one conversation per FRESH session (PFLASH_DRAFTER_SESSIONS=0) over its full +causal prefix, and keeps the block-14 rows of user-role and tool tokens in F32 (exactly the rows the served trunk +reads, rather than bf16) plus the runtime's own request logits (f32, the A6 probe). The runtime writes each call's rows with PFLASH_TRACE_REQUEST_FEATURES +(record "RPF2": i32 resume, n_new, hidden, f32 rows, f32 logits). + +Input JSONL, one conversation per line (the A6 kit format works as is): + {"conversation_id" | "id": str, "ids": [drafter token ids, the full rendered conversation], + "messages": [{"index": int, "role": "user"|"tool"|"assistant"|"system"|..., "span": [b, e)}, ...]} + role "user" = human (the request probe's role rule); role "tool" = tool results / observations. + A message counts as labelled when it has "request_spans" or "labelled": true, else every user/tool message does. + +Output (OUT dir, possibly on another host via --remote host): shard-NNNNN.rows.f32 (float32 [rows, hidden]), +shard-NNNNN.logits.f32 ([rows]), shard-NNNNN.index.jsonl (per conversation: id, n_tokens, row_offset, rows, +messages [{index, role, span, row_offset}], tokens = merged captured token spans in row order), part-*.json summary. +Rows of a conversation are its user/tool tokens in ascending token order. + + request_probe_capture.py capture --daemon BIN --drafter GGUF --scoring-head GGUF --segment-probe GGUF --request-probe GGUF + --input convs.jsonl --out DIR [--part i/n] [--remote HOST] [--shard-mb 1024] +""" +import argparse, hashlib, json, os, shutil, struct, subprocess, sys, tempfile, time +from pathlib import Path +import numpy as np + +# Runtime artifacts, set from the command line in main(). +DAEMON = DRAFTER = HEAD = SEGPROBE = REQPROBE = None +CAPTURE_ROLES = ("user", "tool") + + +def sha256(path): + h = hashlib.sha256() + with open(path, "rb") as f: + for b in iter(lambda: f.read(1 << 22), b""): + h.update(b) + return h.hexdigest() + + +def merged(spans): + out = [] + for b, e in sorted(spans): + if out and b <= out[-1][1]: + out[-1][1] = max(out[-1][1], e) + elif e > b: + out.append([b, e]) + return out + + +class Daemon: + def __init__(self, workdir, env_extra): + self.dump = Path(workdir) / "features.rpf" + env = {k: v for k, v in os.environ.items() if not k.startswith(("PFLASH_", "LUCE_KV"))} + env.update({ + "PFLASH_SCORING_HEAD_GGUF": str(HEAD), "PFLASH_SEGMENT_PROBE_GGUF": str(SEGPROBE), + "PFLASH_SELECT_MODE": "budget_only", "PFLASH_SELECT_CHUNK_SIZE": "256", "PFLASH_SELECT_SCORE": "density", + "PFLASH_SELECT_SEGMENTS": "probe", "PFLASH_SELECT_QUERY_PARSER": "latest_user", + "PFLASH_SELECT_QUERY_TOKENS": "128", "PFLASH_SELECT_PARAGRAPH_JOIN": "0", + "PFLASH_SELECT_QUERY_PROBE": "1", "PFLASH_REQUEST_PROBE_GGUF": str(REQPROBE), + "PFLASH_DRAFTER_SESSIONS": "0", "PFLASH_DAEMON_KEEP_X1000": "1000", + "PFLASH_TRACE_REQUEST_FEATURES": str(self.dump), + }) + env.update(env_extra) + self.env = env + rfd, wfd = os.pipe() + self.log = open(Path(workdir) / "daemon.log", "a") + self.proc = subprocess.Popen([str(DAEMON), str(DRAFTER), f"--stream-fd={wfd}"], + stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=self.log, env=env, + pass_fds=(wfd,), text=True) + os.close(wfd) + self.stream = os.fdopen(rfd, "rb") + if "ready" not in self._line(["ready"]): + raise RuntimeError("pflash_daemon did not start (see daemon.log)") + + def _line(self, words): + while True: + line = self.proc.stdout.readline() + if not line: + return "" + if line.startswith("[pflash-daemon]") and any(w in line for w in words): + return line.strip() + + def capture(self, ids, human, path): + """Rows [S, hidden] f32 and request logits [S] of the whole prompt, from a fresh session.""" + self.dump.unlink(missing_ok=True) + with open(path, "wb") as f: + f.write(struct.pack("