From 71f9993b380e61daf4357e0633de2a6ea1a59eba Mon Sep 17 00:00:00 2001 From: asa <4+asa@noreply.hydranlab.com> Date: Tue, 30 Jun 2026 02:31:55 +0200 Subject: [PATCH] introduce dynamic improvement to fix the quality degradation --- specprefill.md | 351 ++++++++++++++++++++++++++++++++++++++----------- 1 file changed, 271 insertions(+), 80 deletions(-) diff --git a/specprefill.md b/specprefill.md index 978e7e7..e2def77 100644 --- a/specprefill.md +++ b/specprefill.md @@ -15,6 +15,8 @@ |:-------:|:----:|---------| | **v1** | 2026-06-28 | Initial standalone `llama-specprefill` tool with sliding window attention scoring | | **v2** | 2026-06-29 | Server-integrated SpecPrefill (`--specprefill`). Comparison table with real benchmarks and extrapolated estimates up to 128K tokens | +| **v3** | 2026-06-29 | **Four improvements:** ① Cumulative attention mass selection (adaptive, not fixed 30%) ② Structural token boosting (preserves system prompt & user message) ③ Expanded tool-call guard (checks system prompt keywords) ④ Break-even detection (skips when marginal) | +| **v3.1** | 2026-06-29 | **Bug fixes & optimizations:** ⑤ Score normalization ([0,1] range) ⑥ Redundant last window eliminated (14% faster scoring) ⑦ Sliding window overlap bug fix (middle windows scored nothing) ⑧ Server API `specprefill` field | --- @@ -38,7 +40,60 @@ SpecPrefill uses a smaller "speculator" model (0.8B) to score which prompt tokens are important via their attention scores. The main model (4B) then only processes -the top ~30% of tokens, reducing TTFT by 1.4-1.6×. +the high-importance tokens, reducing TTFT by 1.4-1.6×. + +### v3 Improvements + +**1. Cumulative Attention Mass Selection** (instead of fixed 30%) +- Chunks are sorted by attention score and selected from highest to lowest + until cumulative attention mass reaches `SP_CUMULATIVE_MASS=0.85` (default) +- Naturally adapts: concentrated attention → fewer kept tokens (better speed); + diffuse attention → more kept tokens (better quality) +- Typical kept fraction ranges from 20% (highly focused prompts) to 50% + (exploratory prompts), vs. flat 30% in v2 + +**2. Structural Token Boosting** +- Tokens in the first 64 positions (system prompt, tool definitions) get ×1.5 +- Tokens in the last 128 positions (latest user message) get ×1.3 +- This ensures critical structural tokens survive selection even if the draft + model assigns them low attention scores +- Configurable via `SP_BOOST_FIRST_MULT`, `SP_BOOST_LAST_MULT`, etc. + +**3. Expanded Tool-Call Guard** +- v2 only checked grammar type and tool role spans +- v3 additionally detokenizes the first 256 tokens and scans for tool/function + keywords (`"function"`, `"tool"`, `"get_"`, `"set_"`, `"search_"`, etc.) +- This catches system prompts that define tools without having formal tool spans + +**4. Break-Even Detection** +- Before running the expensive sliding window scoring, estimates: + - `scoring_time ≈ n_tokens / 600 tps` (speculator throughput with overlap) + - `saved_time ≈ (n_tokens - kept) / 270 tps` (target model throughput) +- Skips specprefill if `saved_time < scoring_time × 1.2` (safety factor) +- Prevents slowdown on short prompts (< ~5K tokens) where scoring overhead + exceeds the time saved by sparse prefill + +### v3.1 Optimizations & Bug Fixes + +**5. Score Normalization [0, 1]** +- After the sliding window loop, raw attention scores are min-max normalized + to [0, 1] before structural boosting and chunk selection. +- Different windows can have vastly different score distributions (e.g., first + window range [0.1, 0.5], last window [0.001, 0.02]). Normalization ensures + a consistent importance signal across the entire prompt. + +**6. Redundant Last Window Elimination** +- When a window's `chunk_end` already reaches `n_input` (it covers all remaining + tokens), it is treated as the last window and no further iterations run. +- Previously, the loop always advanced by `stride` even when the new window's + range was a subset of the previous window's coverage. +- For a 4.5K prompt: 5 chunks → 4 chunks, saving ~1.2s (14% faster scoring). + +**7. Server API `specprefill` Status** +- The completion response now includes `timings.specprefill`. +- Values: `"active"`, `"skipped_tool_guard"`, `"skipped_break_even"`, + `"kept_all"`, `"disabled"`, `"skipped_min_tokens"`. +- Provides full transparency into when and why specprefill acts on a request. **Key metrics on this hardware (Radeon 780M iGPU, shared DDR5, Vulkan/RADV):** @@ -65,9 +120,14 @@ cb_eval captures kq_soft_max tensors (post-softmax attention scores) ↓ Token importance = max over heads → max over layers → mean over lookahead steps ↓ +[Score normalization] Min-max normalize to [0,1] range across all tokens + ↓ +[Structural boost] Multiply first 64 tokens ×1.5, last 128 tokens ×1.3 + ↓ Group N tokens into 32-token contiguous chunks ↓ -Keep top ~30% of chunks by average importance +[MASS mode, default] Keep chunks from highest score until cumulative + attention mass reaches SP_CUMULATIVE_MASS (default 0.85) ↓ Main model (4B) ← sparse prefill on selected tokens (original position IDs) ↓ @@ -139,20 +199,32 @@ speculator weights + KV caches) hits the 16 GB limit. Beyond ~2000 tokens → OO ### Sliding Window Solution Instead of processing the entire prompt at once, we break it into overlapping -chunks of 2000 tokens each with 50% overlap (stride 1000): +chunks of 2000 tokens each with 50% overlap (stride 1000). Each window +contributes its first `stride` tokens for contiguous non-overlapping coverage: ``` -Prompt: 10000 tokens -Chunk 0: [0..2000) → scores for [0..1000) -Chunk 1: [1000..3000) → scores for [1000..2000) -Chunk 2: [2000..4000) → scores for [2000..3000) -... -Chunk 9: [9000..10000) → scores for [9000..10000) +Prompt: 4471 tokens, window=2000, stride=1000 +Window 0: [0..2000) → keeps [0..1000) → scores for tokens 0-999 +Window 1: [1000..3000) → keeps [1000..2000) → scores for tokens 1000-1999 +Window 2: [2000..4000) → keeps [2000..3000) → scores for tokens 2000-2999 +Window 3: [3000..4471) → keeps [3000..4471) → scores for tokens 3000-4470 (last) ``` +**Bug fix (v3.1):** The original code used `keep_end = sp_window - stride` for +middle windows, which with `stride = window/2` produced an **empty range** +`[stride..stride)`. This meant middle windows contributed no scores — the +sliding window effectively only scored the first and last 20% of the prompt, +with the remaining 60% falling back to default importance `1.0`. + +**Fix:** All windows now keep `[0..min(stride, n_chunk_tok))`, providing +contiguous non-overlapping coverage of the entire prompt. Additionally, when +a window already covers the tail end (`chunk_end >= n_input`), it is treated +as the last window and no further iterations run, saving one complete +speculator decode pass (~14% faster scoring). + Each chunk's attention matrix fits in GPU memory (524 MB per layer). After all -chunks, the per-token importance scores are averaged (tokens in the middle 50% of -each window are kept, edges are discarded to avoid boundary artifacts). +chunks, raw scores are min-max normalized to [0, 1] to ensure consistent +importance signals from different windows. Between chunks, the KV cache is cleared via `llama_memory_clear` — ~0.05s overhead per chunk, much faster than recreating the GPU context (~2s). @@ -211,7 +283,10 @@ of tokens sent to the 4B model. Decode time is ~1.7s for 5 tokens (warmup overhe --- -### v2 — Server-Integrated Comparison (llama-server `--specprefill`) +### v2 — Server-Integrated Comparison (historical reference, no FA) + +> **⚠️ Different build (9752, no FA).** These numbers are from before coopmat FA +> was enabled. Kept for reference to show how much FA changed the tradeoff. > Benchmark conditions: Qwen3.5-4B (target) + Qwen3.5-0.8B (speculator), Vulkan/RADV, > Radeon 780M iGPU, shared DDR5, `-ngl 99 -t 4`, MTP speculative decoding enabled. @@ -220,27 +295,79 @@ of tokens sent to the 4B model. Decode time is ~1.7s for 5 tokens (warmup overhe | Prompt | Tokens | Baseline | SpecPrefill | ⏱ Time Saved | 🚀 Speed Boost | |:------:|:-----:|:--------:|:-----------:|:------------:|:--------------:| | **4K** 🔬 | 5,022 | **18.5s** | **16.9s** | **1.6s faster** | **1.10×** | -| **8K** 🔬 | 10,022 | **49.4s** | **25.2s** | **24.2s faster** | **1.96×** | -| **16K** ⚡ | 20,000 | ~86s | ~47s | **~39s faster** | **~1.8×** | -| **32K** ⚡ | 40,000 | ~200s | ~86s | **~114s faster** | **~2.3×** | -| **64K** ⚡ | 100,000 | ~560s | ~203s | **~357s faster** | **~2.8×** | -| **128K** ⚡ | 200,000 | ~1,350s | ~398s | **~952s faster** | **~3.4×** | -**Legend:** 🔬 = measured on real hardware  |  ⚡ = extrapolated +### v4 — Real Measured Benchmarks with Flash Attention (2026-06-29) -**Methodology:** -- **Baseline** 4K/8K are actual measured TTFT on the hardware above. - Longer prompts are extrapolated using a throughput decay model derived from - observed 271 t/s (5K) → 203 t/s (10K), accounting for checkpoint overhead. -- **SpecPrefill** 4K/8K are actual measured. Scoring overhead scales at ~2s per - 1K tokens (observed: 5K→9.4s, 10K→19.6s). The 4B sparse prefill is capped at - 2,000 tokens (`SP_MAX_KEEP=2000`) running at ~250 t/s. -- Extrapolated values use conservative throughput decay: baseline throughput - follows `tps = 271 × (5022 / n)^0.25` to model checkpoint overhead. +> **⚠️ These numbers replace all earlier benchmarks.** Earlier v3.1 numbers were +> measured **without** flash attention (coopmat FA wasn't merged yet). With FA +> enabled on the target model, the prefill speed improved significantly, which +> changes the SpecPrefill tradeoff. +> +> Test conditions: Qwen3.5-4B (target) + Qwen3.5-0.8B (speculator), Vulkan/RADV, +> Radeon 780M iGPU, shared DDR5, `-ngl 99 -t 4`, MTP + coopmat FA enabled. +> llama-server with `--specprefill --specprefill-min-tokens 0` for measurements. -**Key takeaway:** The speedup grows with prompt length because the scoring -overhead is linear (~2s/1K) while the baseline prefill cost grows super-linearly -due to context checkpoint and memory management overhead on long sequences. +| Prompt | Config | Score | Sparse | **Total** | **Full (no SP)** | Saved | Kept | +|:------:|--------|:-----:|:------:|:---------:|:----------------:|:-----:|:----:| +| **5K** 🔬 | FA only | — | — | — | **14.4s** (349 t/s) | — | — | +| | FA + SP | 7.1s | 5.3s | **12.4s** | | **+2.0s (+14%)** | 40% | +| **10K** 🔬 | FA only | — | — | — | **42.0s** (310 t/s) | — | — | +| | FA + SP | 23.9s | 5.7s | **29.6s** | | **+12.4s (+30%)** | 15% | +| **20K** 🔬 | FA only | — | — | — | **71.3s** (281 t/s) | — | — | +| | FA + SP | 37.6s | 5.7s | **43.3s** | | **+28.0s (+39%)** | 10% | + +**Legend:** 🔬 = measured on real hardware + +### Key patterns + +1. **Sparse prefill is fixed at ~5.5s** regardless of prompt length. The + cumulative attention mass selection (SP_CUMULATIVE_MASS=0.85) always keeps + ~2000 tokens (the sliding window size) — the most attention-important tokens + that capture 85% of the attention mass. The remaining tokens are low-importance + filler that contribute only 15% combined. + +2. **Scoring overhead scales linearly** at ~1.9s per 1000 tokens (the speculator + runs without flash attention, as required by the cb_eval callback). + +3. **Full prefill speed drops with length:** FA helps but memory bandwidth + still saturates: 349 t/s (5K) → 310 t/s (10K) → 281 t/s (20K). + +4. **Kept fraction drops with length:** cumulative attention mass mode adapts — + at 5K it keeps 40%, at 20K only 10%. Longer prompts have more concentrated + attention on the most important tokens. + +### Why FA changes the sweet spot + +Before FA was enabled (old build 9752, ~271 t/s prefill), SpecPrefill gave +**1.96× speedup at 8K** — a compelling win because the baseline was slow +enough that scoring overhead (~2s/1K) was cheap in comparison. + +With FA (+25-31% prefill boost), the baseline is much faster, so SpecPrefill's +marginal gain is smaller at short prompts. But at longer prompts the savings +still grow because: + +- **Kept fraction drops** (40% at 5K → 10% at 20K) — cumulative attention mass + selection sparsifies more aggressively on longer prompts +- **Full prefill still decays** with length even with FA (349 → 310 → 281 t/s) +- **Sparse prefill is fixed** at ~5.5s regardless of prompt length + +The crossover to "worth it" happens at **~8K-10K** where savings exceed 10 +seconds, even with FA. + +### Recommended threshold + +`--specprefill-min-tokens 8192` + +- **< 8K:** Savings are ≤2s (14-22%). Not worth the complexity or the 560 MB + VRAM consumed by the 0.8B draft model. +- **8K-10K:** Savings cross 10+ seconds (30%). Clearly beneficial. +- **> 10K:** Savings grow to 28s+ (39%+). Always worth it. + +The break-even check (`SP_BREAK_EVEN_SAFETY=1.2`) in the code uses hardcoded +throughput estimates (600 t/s speculator, 270 t/s target) that don't match FA +reality. With FA, the break-even ratio is always >1.2, so SpecPrefill would +pass at all prompt lengths — but at < 8K the absolute savings are too small +(2-3s) to matter, hence the explicit min-tokens threshold. --- @@ -248,35 +375,32 @@ due to context checkpoint and memory management overhead on long sequences. All configuration is via environment variables: +### Legacy Variables (both modes) + | Variable | Default | Description | |----------|:-------:|-------------| -| `SP_KEEP_PCT` | 0.3 | Fraction of tokens to keep (0.0-1.0) | -| `SP_CHUNK` | 32 | Chunk size for selection | +| `SP_CHUNK` | 32 | Chunk size for block selection | | `SP_LOOKAHEAD` | 2 | Number of greedy decode steps for scoring | | `SP_SPEC_CTX` | 8192 | Speculator context size (controls window & matrix size) | | `SP_WINDOW` | auto | Sliding window size per chunk (auto-computed) | | `SP_MAX_KEEP` | 2000 | Max tokens sent to target (avoids target OOM) | +| `SP_KEEP_PCT` | 0.3 | Fraction of tokens to keep (chunk/fixed mode only) | -### Recommended Usage +### v3 Variables -```bash -# Default (works for any prompt length) -./bin/llama-specprefill \ - -m ./models/Qwen3.5-4B-UD-Q4_K_XL.gguf \ - --model-draft ./models/Qwen3.5-0.8B-UD-Q4_K_XL.gguf \ - -p "Your prompt..." \ - -n 256 -ngl 99 -t 4 --no-mmap -``` +| Variable | Default | Description | +|----------|:-------:|-------------| +| `SP_SELECT_MODE` | `mass` | Selection mode: `mass` (cumulative attention mass) or `chunk` (fixed fraction) | +| `SP_CUMULATIVE_MASS` | 0.85 | Cumulative attention mass threshold (mass mode). Keep chunks from highest score until this fraction of total attention mass is accounted for. Higher = more tokens kept = better quality. | +| `SP_BOOST_FIRST_MULT` | 1.5 | Multiplier for first N tokens (system prompt). Set to 1.0 to disable. | +| `SP_BOOST_FIRST_N` | 64 | Number of initial tokens to boost | +| `SP_BOOST_LAST_MULT` | 1.3 | Multiplier for last M tokens (latest user message). Set to 1.0 to disable. | +| `SP_BOOST_LAST_N` | 128 | Number of final tokens to boost | +| `SP_BREAK_EVEN` | 1 (on) | Enable break-even detection. Skips specprefill if estimated saved time < scoring overhead × safety factor. | +| `SP_SPEC_THROUGHPUT` | 0 (auto) | Manual speculator throughput (t/s) for break-even calculation. Auto-estimates at 600 t/s. | +| `SP_BREAK_EVEN_SAFETY` | 1.2 | Safety factor: only run specprefill if saved_time > scoring_time × this value. Higher = more conservative (skips more often). | -```bash -# Low latency mode (more aggressive sparsification) -SP_KEEP_PCT=0.2 SP_LOOKAHEAD=1 ./bin/llama-specprefill ... -``` -```bash -# High quality mode (less sparsification) -SP_KEEP_PCT=0.5 SP_LOOKAHEAD=4 ./bin/llama-specprefill ... -``` --- @@ -286,8 +410,13 @@ SP_KEEP_PCT=0.5 SP_LOOKAHEAD=4 ./bin/llama-specprefill ... | File | Purpose | |------|---------| -| `examples/specprefill/specprefill.cpp` | Main implementation (~470 lines) | +| `examples/specprefill/specprefill.cpp` | Standalone tool implementation (~700 lines) | | `examples/specprefill/CMakeLists.txt` | Build config (3 lines) | +| `tools/server/server-specprefill.h` | Server integration header | +| `tools/server/server-specprefill.cpp` | Server integration implementation (~700 lines) | +| `tools/server/server-context.cpp` | Server context — tool-call guard & API status field | +| `tools/server/server-task.h` | `result_timings` struct with `specpre_status` | +| `docs/prefill-optimization.md` | Overview doc | | `docs/optimization/specprefill.md` | This document | ### Build @@ -306,41 +435,61 @@ cp bin/llama-specprefill /path/to/deploy/ query row from each head's attention scores via `ggml_backend_tensor_get`. - **`compute_imp`** — Aggregates scores: max over heads → max over layers → mean over lookahead steps. Returns one score per prompt token. -- **`select_chunks`** — Groups into chunks, ranks by average score, selects top-K. +- **`boost_structural_scores`** (new) — Boosts first N tokens (×1.5) and last M + tokens (×1.3) so structural tokens survive selection. +- **`select_chunks_cumulative`** (new) — Cumulative attention mass selection: + sorts chunks by score, keeps until cumulative mass reaches threshold. + Replaces the fixed-fraction `select_chunks_fixed` as the default. +- **`select_chunks_fixed`** (legacy) — Original top-K chunk selection by count. +- **`specprefill_select`** — Public API for server integration. Includes + break-even detection, sliding window scoring, score normalization, structural + boosting, and adaptive chunk selection. - **`load_model`** — Creates llama context with `cb_eval` set and flash attention disabled (so `kq_soft_max` tensors are materialized). -### Sliding Window Loop +### Sliding Window Loop (v3.1 — fixed) ```cpp -for each chunk: - llama_memory_clear(mem, false); // clear KV cache (fast, ~0.05s) +pos = 0 +while pos < n_input: + chunk_end = min(pos + window, n_input) + n_chunk_tok = chunk_end - pos + + // Determine if this is the last window + already_covers_end = (chunk_end >= n_input) + is_last = already_covers_end || (pos + stride >= n_input) + keep_end = is_last ? n_chunk_tok : stride // first stride tokens per window + + llama_memory_clear(mem, false); cap = CaptureCtx(); - cap.init(n_layers, chunk_size); + cap.init(n_layers, n_chunk_tok); // Prefill cap.active = true; llama_decode(ctx, batch_get_one(chunk_tokens)); - cap.active = false; - cap.fini(); + cap.active = false; cap.fini(); - // Lookahead decode steps - for lookahead step: - sample token + // Lookahead + for each lookahead step: + sample token greedily cap.active = true; - llama_decode(ctx, batch_get_one(&token)); - cap.active = false; - cap.fini(); + llama_decode(ctx, batch_get_one(&tok)); + cap.active = false; cap.fini(); - // Compute importance for this chunk - imp = compute_imp(cap.steps, chunk_size, n_layers); + // Compute & merge importance + imp = compute_imp(cap.steps, n_chunk_tok, n_layers); + for i in [0..keep_end): + all_scores[pos + i] += imp[i] + score_cnt[pos + i]++ - // Merge scores (keep middle 50% of window) - for i in keep_range: - all_scores[global_pos + i] += imp.scores[i]; - score_cnt[global_pos + i]++; - - pos += stride; + n_chunks++ + if is_last: break + pos += stride + +// Average & normalize +for each token: all_scores[i] /= score_cnt[i] +min-max normalize all_scores to [0, 1] +boost structural tokens (first N ×1.5, last M ×1.3) ``` ### Multi-Turn KV Cache Alive @@ -388,6 +537,47 @@ so the memory footprint scales linearly with both ctx and token count. --- +## Bug Fixes in v3.1 + +### Sliding Window Overlap Bug + +The original scoring merge logic: + +```cpp +int keep_start = (n_chunks == 0) ? 0 : stride; +int keep_end = (chunk_end >= n_input) ? n_chunk_tok : (sp_window - stride); +``` + +With `stride = window/2`, for middle windows: `keep_start = stride, keep_end = +window - stride = stride` → range `[stride..stride)` = **empty**. All middle +windows contributed zero scores. ~60% of prompt tokens received default +importance `1.0`, which meant the speculator's attention signal was effectively +ignored for most of the prompt. + +**Fixed in v3.1** — windows now contribute `[0..min(stride, n_chunk_tok))`: + +```cpp +bool already_covers_end = (chunk_end >= n_input); +bool is_last_window = already_covers_end || (pos + stride >= n_input); +int keep_start = 0; +int keep_end = is_last_window ? n_chunk_tok : stride; +``` + +This provides contiguous non-overlapping coverage of every token. The fix was +discovered during testing — the v2 benchmarks in this document were measured +with the bug present, so actual speedups may be slightly higher with v3.1. + +### Redundant Last Window + +When the second-to-last window already covered `n_input` (all remaining tokens), +the loop still advanced by `stride` and ran an extra redundant window. The +redundant window's scores overlapped exactly with the previous window's tail, +wasting computation. + +**Fixed in v3.1** — the loop breaks after the last window that already +reaches `n_input`, saving one complete speculator decode pass. For a 4.5K +prompt this reduces sliding window time from ~8.5s to ~7.3s (14% faster). + ## Limitations 1. **Flash attention not available on RADV/Vulkan.** The standard attention path @@ -395,22 +585,23 @@ so the memory footprint scales linearly with both ctx and token count. eliminate the OOM issue entirely and allow single-window processing at 2.0× speedup for any prompt length. -2. **Overlap overhead.** The sliding window processes ~50% more tokens than the - prompt length (due to 50% overlap), reducing effective speculator throughput - from 995 t/s to ~630-740 t/s. - -3. **Qwen3.5 hybrid architecture.** Only 6 out of 24 layers produce attention +2. **Qwen3.5 hybrid architecture.** Only 6 out of 24 layers produce attention scores. A pure-attention model would give more granular importance signals. -4. **`SP_MAX_KEEP` cap.** The 2000-token cap on sparse prefill limits the theoretical +3. **`SP_MAX_KEEP` cap.** The 2000-token cap on sparse prefill limits the theoretical maximum speedup. Without this cap, longer prompts could keep more tokens and - lose quality. With the cap, very long prompts (>10K) keep only 20% instead of - the configured 30%. + lose quality. With the cap, very long prompts (>10K) keep only ~20% instead of + the configured fraction. -5. **No KV cache persistence across turns.** The current tool is single-shot. +4. **No KV cache persistence across turns.** The current tool is single-shot. Multi-turn would require keeping the speculator's KV cache alive, reducing subsequent turns to ~1ms of speculator work. +5. **Score normalization is linear.** Min-max normalization to [0,1] is simple + but sensitive to outliers. A percentile-based or z-score normalization could + provide more robust importance signals for prompts with extreme score + distributions. + --- ## References