# SpecPrefill — Implementation Guide > **Paper:** [Speculative Prefill: Turbocharging TTFT with Lightweight and > Training-Free Token Importance Estimation](https://arxiv.org/abs/2502.02789) > (Liu, Chen, Zhang — PMLR 2025 / ICML 2025) > > **llama.cpp implementation:** `examples/specprefill/specprefill.cpp` > **Built binary:** `bin/llama-specprefill` --- ## Version History | Version | Date | Changes | |:-------:|:----:|---------| | **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 | --- --- ## Table of Contents 1. [Overview](#overview) 2. [How It Works](#how-it-works) 3. [Sliding Window](#sliding-window) 4. [Benchmarks](#benchmarks) 5. [Configuration Reference](#configuration-reference) 6. [Implementation Details](#implementation-details) 7. [GPU Memory Analysis](#gpu-memory-analysis) 8. [Limitations](#limitations) 9. [References](#references) --- ## Overview 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 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):** | Metric | Value | |--------|-------| | Baseline prefill | ~270 t/s (4B model) | | Speculator prefill (GPU) | ~995 t/s (0.8B model) | | Sparse prefill (4B, 30% tokens) | ~340-400 t/s | | Speedup | 1.4-1.6× across all lengths | | Max prompt length | Unlimited (sliding window) | --- ## How It Works ### Core Algorithm ``` Prompt (N tokens) ↓ Small speculator (0.8B) ← prefill all N tokens + run L lookahead decode steps ↓ 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 ↓ [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) ↓ Normal decode ``` ### Attention Capture (cb_eval) The `llama_context_params.cb_eval` callback fires after each graph operation during evaluation. We match tensors named `kq_soft_max-{il}` (one per full-attention layer) and read the **last query position's** attention scores from each head. Qwen3.5 is a hybrid architecture (3:1 Gated DeltaNet + Attention). Only full-attention layers (every 4th) produce `kq_soft_max` tensors. For the 0.8B (24 layers), these are layers 3, 7, 11, 15, 19, 23. The row read from each head's attention matrix represents "how much does the last prompt token attend to each earlier token." After lookahead, it becomes "how much does the predicted next token attend to each prompt token." ### Token Importance For each lookahead step: 1. Read one row (last query position) from `kq_soft_max` per head 2. Max-pool over heads: take the maximum attention score across all heads 3. Max-pool over layers: take the maximum across the 6 attention layers Then mean-pool over all lookahead steps. Result: one importance score per prompt token. Tokens that the speculator attends to more are deemed more important. ### Chunk Selection Prompt tokens are grouped into contiguous 32-token chunks. Chunks are ranked by average importance score. The top ~30% of chunks are selected. First 4 and last 2 tokens are always kept for context. ### Sparse Prefill The target model processes only the selected tokens, with **original position IDs** preserved. This ensures correct RoPE-relative attention positions during decode. The target's KV cache is built from only the important tokens, making subsequent decode attention cheaper. --- ## Sliding Window ### The OOM Problem Standard attention materializes the full `Q @ K^T` matrix. On RADV/Vulkan, flash attention is not available, so this matrix is always materialized. Its size: ``` Attention matrix = KV_cache_cap × n_tokens × n_heads × 4 bytes ``` With the default speculator context (8192) and 2000 tokens: ``` 8192 × 2000 × 8 × 4 = 524 MB per attention layer × 6 attention layers = 3.1 GB (scheduler allocates all at once) ``` At ~2000 tokens, the combined GPU memory (speculator matrices + 4B weights + speculator weights + KV caches) hits the 16 GB limit. Beyond ~2000 tokens → OOM. ### 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). Each window contributes its first `stride` tokens for contiguous non-overlapping coverage: ``` 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, 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). ### Auto Window Size The tool auto-computes the optimal window size: ```cpp max_safe_window = 500 MB / (spec_ctx × 32) ``` For `SP_SPEC_CTX=8192` (default): 500M / 262144 = ~2000 tokens per chunk. The overlap is always 50% (stride = window / 2). --- ## Benchmarks ### System - **GPU:** Radeon 780M iGPU (RADV/Vulkan, 16 GB shared DDR5) - **Target model:** Qwen3.5-4B-UD-Q4_K_XL (2.8 GB on GPU) - **Speculator model:** Qwen3.5-0.8B-UD-Q4_K_XL (0.5 GB on GPU) - **llama.cpp build:** Build 286+ (`build_fresh`) - **Flags:** `-ngl 99 -t 4 --no-mmap` ### Prefill Throughput | Approach | 2K prompt | 6K prompt | 10K prompt | |----------|:---------:|:---------:|:----------:| | Baseline (4B only) | 270 t/s | 270 t/s | 270 t/s | | Spec GPU (no sliding) | **995 t/s** | ❌ OOM | ❌ OOM | | Spec GPU (sliding) | 740 t/s | 635 t/s | 630 t/s | | Sparse prefill (capped 2K) | 340 t/s | 385 t/s | 397 t/s | ### TTFT (Time-to-First-Token) | Prompt | Baseline | Spec GPU no sliding | Spec GPU sliding | Best speedup | |:------:|:-------:|:------------------:|:----------------:|:------------:| | **2K** | **7.4s** | **3.8s** (1.9×) | 5.1s (1.4×) | **1.9×** | | **6K** | **22.2s** | ❌ | **16.3s (1.4×)** | **1.4×** | | **10K** | **37.0s** | ❌ | **22.5s (1.6×)** | **1.6×** | ### Per-Phase Breakdown (Sliding Window) | Prompt | Windows | Spec time | Sparse prefill | Decode | Total | |:------:|:------:|:---------:|:--------------:|:-----:|:-----:| | 2K | 1 | 2.7s | 1.6s | 0.8s | 5.1s | | 6K | 7 | 9.4s | 5.2s | 1.7s | 16.3s | | 10K | 10 | 15.9s | 5.0s | 1.7s | 22.5s | Note: Sparse prefill is consistently 5s because `SP_MAX_KEEP=2000` caps the number of tokens sent to the 4B model. Decode time is ~1.7s for 5 tokens (warmup overhead). --- ### 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. > Server flags: `--specprefill --specprefill-min-tokens 0`. | Prompt | Tokens | Baseline | SpecPrefill | ⏱ Time Saved | 🚀 Speed Boost | |:------:|:-----:|:--------:|:-----------:|:------------:|:--------------:| | **4K** 🔬 | 5,022 | **18.5s** | **16.9s** | **1.6s faster** | **1.10×** | ### v4 — Real Measured Benchmarks with Flash Attention (2026-06-29) > **⚠️ 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. | 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. --- ## Configuration Reference All configuration is via environment variables: ### Legacy Variables (both modes) | Variable | Default | Description | |----------|:-------:|-------------| | `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) | ### v3 Variables | 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). | --- ## Implementation Details ### File Structure | File | Purpose | |------|---------| | `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 ```bash cd llama.cpp mkdir -p build && cd build cmake .. -DLLAMA_VULKAN=ON make -j$(nproc) llama-specprefill cp bin/llama-specprefill /path/to/deploy/ ``` ### Key Functions - **`capture_cb`** — `cb_eval` callback. Matches `kq_soft_max-{il}`, reads the last 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. - **`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 (v3.1 — fixed) ```cpp 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, n_chunk_tok); // Prefill cap.active = true; llama_decode(ctx, batch_get_one(chunk_tokens)); cap.active = false; cap.fini(); // Lookahead for each lookahead step: sample token greedily cap.active = true; llama_decode(ctx, batch_get_one(&tok)); cap.active = false; cap.fini(); // 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]++ 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 When the speculator's KV cache is kept alive across turns (not yet implemented as a server), each subsequent turn only processes delta tokens: ``` Turn 1: prefill 2000 tokens → cache has 2000 Turn 2: prefill 100 delta tokens → Q(100) @ K(2100)^T = tiny matrix ``` This makes each additional turn cost ~1ms instead of re-prefilling the full history. The sliding window is still needed for the first turn if the initial prompt is large. --- ## GPU Memory Analysis ### Per-Layer Attention Matrix Size | Spec ctx | Window | Per-layer matrix | 6 layers | + Weights (3.3 GB) | Fits 16 GB? | |:-------:|:------:|:----------------:|:--------:|:-----------------:|:-----------:| | 8192 | 2000 | 524 MB | 3.1 GB | 6.4 GB | ✅ | | 8192 | 2500 | 655 MB | 3.9 GB | 7.2 GB | ✅ but target OOMs | | 16384 | 2000 | 1.0 GB | 6.0 GB | 9.3 GB | ✅ | | 32768 | 2000 | 2.1 GB | 12.6 GB | 15.9 GB | ⚠️ borderline | | 32768 | 3000 | 3.1 GB | 18.6 GB | 21.9 GB | ❌ | ### Why No-Sliding OOMs at ~2000 Tokens With `SP_SPEC_CTX=32768` (or when the speculator uses the target's default ctx): ``` Matrix per layer = 32768 × n_tokens × 8 × 4 = 1,048,576 × n_tokens bytes = ~1 MB per token At 2000 tokens: 2.1 GB/layer × 6 = 12.6 GB → + 2.8 GB (4B) + 0.5 GB (0.8B) = 15.9 GB → just fits At 2200 tokens: 2.3 GB/layer × 6 = 13.8 GB → + 3.3 GB = 17.1 GB → OOM ``` The graph scheduler allocates all 6 attention layers' matrices simultaneously, 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 materializes the full QK^T matrix. A custom Vulkan flash attention shader would eliminate the OOM issue entirely and allow single-window processing at 2.0× speedup for any prompt length. 2. **Qwen3.5 hybrid architecture.** Only 6 out of 24 layers produce attention scores. A pure-attention model would give more granular importance signals. 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 fraction. 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 1. Liu, Chen, Zhang. *"Speculative Prefill: Turbocharging TTFT with Lightweight and Training-Free Token Importance Estimation."* PMLR 2025. [https://arxiv.org/abs/2502.02789](https://arxiv.org/abs/2502.02789) 2. Original vLLM implementation: [github.com/Jingyu6/speculative_prefill](https://github.com/Jingyu6/speculative_prefill) 3. vLLM feature request: [github.com/vllm-project/vllm/issues/39060](https://github.com/vllm-project/vllm/issues/39060) 4. vLLM-mlx implementation: [github.com/waybarrios/vllm-mlx/issues/179](https://github.com/waybarrios/vllm-mlx/issues/179) --- ## Full Source Code `examples/specprefill/specprefill.cpp` (473 lines): ```cpp // SPDX-FileCopyrightText: 2026 // SPDX-License-Identifier: MIT // // SpecPrefill with Sliding Window — Attention-Based Sparse Prefill for TTFT // // Processes long prompts by sliding a window over the prompt in overlapping // chunks. Each chunk fits in the speculator's context budget (e.g., 8K). // Attention scores from all chunks are merged, then top-K chunks selected. // // Usage: // SP_KEEP_PCT=0.3 SP_LOOKAHEAD=4 SP_WINDOW=8192 \ // ./bin/llama-specprefill -m model.gguf --model-draft draft.gguf \ // -p "prompt..." -ngl 99 -t 4 --no-mmap // #include "arg.h" #include "common.h" #include "sampling.h" #include "log.h" #include "llama.h" #include #include #include #include #include #include #include #include #include #include #include // ============================================================================ // Backend: kq_soft_max capture via cb_eval // ============================================================================ struct CapturedAttn { int il; int64_t n_k; int64_t n_h; std::vector data; }; struct CaptureCtx { bool active = false; int n_layers = 0; int n_tokens = 0; std::vector layers; std::vector> steps; void init(int nl, int nt) { n_layers = nl; n_tokens = nt; layers.resize(nl); steps.clear(); } void reset() { for (auto & l : layers) { l.data.clear(); l.n_k = 0; l.n_h = 0; } } void fini() { steps.push_back(layers); } }; static bool capture_cb(struct ggml_tensor * t, bool ask, void * ud) { if (ask || !ud) return true; auto * cc = (CaptureCtx *) ud; if (!cc->active) return true; const char * name = ggml_get_name(t); if (!name) return true; int il = -1; if (sscanf(name, "kq_soft_max-%d", &il) != 1) return true; if (il < 0 || il >= cc->n_layers) return true; const int64_t nk = t->ne[0], nq = t->ne[1], nh = t->ne[2]; if (nk <= 0 || nq <= 0 || nh <= 0) return true; const int64_t qi = (nq > 1) ? (nq - 1) : 0; const int64_t nr = std::min(nk, (int64_t) cc->n_tokens); if (nr <= 0) return true; auto & l = cc->layers[il]; l.data.resize((size_t) nh * nr, 0.0f); l.il = il; l.n_k = nr; l.n_h = nh; const size_t es = ggml_type_size(t->type); if (es == 0) return true; for (int h = 0; h < nh; h++) { const uint64_t off = ((uint64_t)qi * nk + (uint64_t)h * nq * nk) * es; const uint64_t sz = (uint64_t)nr * es; if (sz == 0) continue; ggml_backend_tensor_get(t, l.data.data() + (size_t)h * nr, (size_t)off, (size_t)sz); } return true; } // ============================================================================ // Importance computation // ============================================================================ struct ImpResult { std::vector scores; int n_tokens; }; static ImpResult compute_imp( const std::vector> & steps, int nt, int nl) { ImpResult r; r.n_tokens = nt; r.scores.resize(nt, 0.0f); if (steps.empty() || nt == 0) { std::fill(r.scores.begin(), r.scores.end(), 1.0f); return r; } int ns = (int)steps.size(); std::vector> smax(ns, std::vector(nt, -FLT_MAX)); for (int s = 0; s < ns; s++) { for (auto & cap : steps[s]) { if (cap.il >= nl || cap.data.empty() || cap.n_k <= 0) continue; int nu = std::min((int)cap.n_k, nt); for (int k = 0; k < nu; k++) { float mh = -FLT_MAX; for (int h = 0; h < cap.n_h; h++) mh = std::max(mh, cap.data[(size_t)h * cap.n_k + k]); if (mh > smax[s][k]) smax[s][k] = mh; } } } int vs = 0; for (int s = 0; s < ns; s++) { for (int k = 0; k < nt; k++) { if (smax[s][k] > -FLT_MAX/2) { vs++; break; } } } if (vs == 0) { std::fill(r.scores.begin(), r.scores.end(), 1.0f); return r; } for (int k = 0; k < nt; k++) { double sum = 0; int c = 0; for (int s = 0; s < ns; s++) { if (smax[s][k] > -FLT_MAX/2) { sum += smax[s][k]; c++; } } r.scores[k] = c > 0 ? (float)(sum / c) : 0.0f; } return r; } // ============================================================================ // Chunk selection // ============================================================================ static std::vector select_chunks(const std::vector & scores, int nt, int csize, float keep_pct, int min_tok) { if (csize <= 0) csize = 32; int eff_min = std::min(min_tok, std::max(4, nt / 10)); if (keep_pct >= 1.0f || nt <= eff_min) { std::vector a(nt); std::iota(a.begin(), a.end(), 0); return a; } int nc = (nt + csize - 1) / csize; std::vector cs(nc, 0); std::vector cc(nc, 0); for (int i = 0; i < nt; i++) { int c = i / csize; cs[c] += scores[i]; cc[c]++; } std::vector ca(nc); for (int c = 0; c < nc; c++) ca[c] = cc[c] > 0 ? (float)(cs[c] / cc[c]) : 0.0f; int kc = std::max(1, (int)std::ceil(nc * keep_pct)); std::vector sv = ca; std::sort(sv.begin(), sv.end(), std::greater()); float th = kc <= (int)sv.size() ? sv[kc-1] : -FLT_MAX; std::vector kept; for (int c = 0; c < nc; c++) { if (ca[c] >= th) { int s = c * csize, e = std::min(s + csize, nt); for (int i = s; i < e; i++) kept.push_back(i); } } std::vector always = {0,1,2,3}; if (nt > 4) { always.push_back(nt-2); always.push_back(nt-1); } for (int a : always) { if (a < nt && std::find(kept.begin(), kept.end(), a) == kept.end()) kept.push_back(a); } if ((int)kept.size() < eff_min) { for (int c = 0; c < nc && (int)kept.size() < eff_min; c++) { if (ca[c] < th) { int s = c*csize, e = std::min(s+csize, nt); for (int i = s; i < e && (int)kept.size() < eff_min; i++) if (std::find(kept.begin(), kept.end(), i) == kept.end()) kept.push_back(i); } } } std::sort(kept.begin(), kept.end()); return kept; } // ============================================================================ // Model loading with cb_eval // ============================================================================ struct ModelCtx { llama_model * model = nullptr; llama_context * ctx = nullptr; }; static ModelCtx load_model(const char * path, const common_params & params, int ngl, int ctx, CaptureCtx * cap) { ModelCtx mc; llama_model_params mp = llama_model_default_params(); mp.n_gpu_layers = ngl; mp.use_mmap = params.use_mmap; mc.model = llama_model_load_from_file(path, mp); if (!mc.model) return mc; llama_context_params cp = llama_context_default_params(); cp.n_ctx = ctx > 0 ? (uint32_t)ctx : std::min((uint32_t)65536, (uint32_t)params.n_ctx); cp.n_batch = (uint32_t)std::min((int32_t)cp.n_ctx, params.n_batch); cp.n_ubatch = (uint32_t)std::min((int32_t)(cp.n_ctx / 4), params.n_ubatch); cp.n_threads = params.cpuparams.n_threads; cp.n_threads_batch = params.cpuparams_batch.n_threads; cp.cb_eval = capture_cb; cp.cb_eval_user_data = cap; cp.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_DISABLED; mc.ctx = llama_init_from_model(mc.model, cp); return mc; } // ============================================================================ // Main // ============================================================================ int main(int argc, char ** argv) { std::setlocale(LC_NUMERIC, "C"); common_params params; common_init(); if (!common_params_parse(argc, argv, params, LLAMA_EXAMPLE_SPECULATIVE)) return 1; if (params.speculative.draft.mparams.path.empty()) { LOG_ERR("error: --model-draft is required\n"); return 1; } // Config from env int sp_lookahead = 4; const char * e; if ((e=getenv("SP_LOOKAHEAD"))) sp_lookahead=atoi(e); float sp_keep_pct = 0.3f; if ((e=getenv("SP_KEEP_PCT"))) sp_keep_pct=atof(e); int sp_chunk = 32; if ((e=getenv("SP_CHUNK"))) sp_chunk=atoi(e); int sp_window = 0; if ((e=getenv("SP_WINDOW"))) sp_window=atoi(e); int sp_spec_ctx = 8192; if ((e=getenv("SP_SPEC_CTX"))) sp_spec_ctx=atoi(e); int sp_max_keep = 2000; if ((e=getenv("SP_MAX_KEEP"))) sp_max_keep=atoi(e); LOG("SpecPrefill — Sliding Window\n"); LOG("Config: lookahead=%d keep=%.0f%% chunk=%d window=%s spec_ctx=%d max_keep=%d\n", sp_lookahead, sp_keep_pct*100, sp_chunk, sp_window > 0 ? std::to_string(sp_window).c_str() : "auto", sp_spec_ctx, sp_max_keep); // Init backend llama_backend_init(); llama_numa_init(params.numa); // Load target model llama_model_params tgt_mp = llama_model_default_params(); tgt_mp.n_gpu_layers = params.n_gpu_layers; llama_model * tgt_m = llama_model_load_from_file(params.model.path.c_str(), tgt_mp); if (!tgt_m) return 1; llama_context_params tgt_cp = llama_context_default_params(); tgt_cp.n_ctx = params.n_ctx; tgt_cp.n_batch = params.n_batch; tgt_cp.n_ubatch = params.n_ubatch; tgt_cp.n_threads = params.cpuparams.n_threads; tgt_cp.n_threads_batch = params.cpuparams_batch.n_threads; llama_context * tgt_ctx = llama_init_from_model(tgt_m, tgt_cp); if (!tgt_ctx) return 1; const llama_vocab * vocab = llama_model_get_vocab(tgt_m); // Load speculator with capture int dngl = params.speculative.draft.n_gpu_layers >= 0 ? params.speculative.draft.n_gpu_layers : params.n_gpu_layers; CaptureCtx cap; auto dft = load_model(params.speculative.draft.mparams.path.c_str(), params, dngl, sp_spec_ctx, &cap); if (!dft.model || !dft.ctx) return 1; // Determine window size for speculator. // Use sp_spec_ctx (the speculator's context size) to compute max safe window. // The attention matrix is KV_cache_cap × n_tokens × n_heads. // With ctx=SP_SPEC_CTX, KV cache capacity = SP_SPEC_CTX. // At 2000 tokens × 8 heads × 4 bytes = 256 MB per attention layer, // we want window × SP_SPEC_CTX × 8 × 4 < ~500 MB per layer for safety. // So window < 500M / (SP_SPEC_CTX × 32). // For SP_SPEC_CTX=4096: window < 500M / 131072 = ~3800. // Use 50% of that for safety margin, or the user's explicit setting. int spec_n_ctx = llama_n_ctx(dft.ctx); int max_safe_window = (int)(500 * 1024 * 1024 / (uint64_t)(spec_n_ctx * 32)); if (sp_window <= 0) sp_window = std::min(spec_n_ctx - 512, max_safe_window); sp_window = std::max(sp_window, 512); sp_window = std::min(sp_window, spec_n_ctx - 128); int stride = sp_window / 2; // 50% overlap // Tokenize std::vector inp = common_tokenize(tgt_ctx, params.prompt, true, true); int n_input = (int)inp.size(); if (n_input == 0) { LOG_ERR("empty prompt\n"); return 1; } if (n_input > (int)llama_n_ctx(tgt_ctx) - 4) { LOG_ERR("prompt too long\n"); return 1; } LOG("Prompt: %d tokens, window=%d stride=%d\n", n_input, sp_window, stride); auto t0 = ggml_time_us(); // =================================================================== // Sliding window: process chunks, accumulate importance // =================================================================== std::vector all_scores(n_input, 0.0f); std::vector score_cnt(n_input, 0); int n_chunks = 0; int pos = 0; while (pos < n_input) { int chunk_end = std::min(pos + sp_window, n_input); int n_chunk_tok = chunk_end - pos; // Build chunk token list std::vector chunk_tok(inp.begin() + pos, inp.begin() + chunk_end); LOG(" Window %d: tokens [%d..%d) (%d tok)\n", n_chunks, pos, chunk_end, n_chunk_tok); // --- Clear speculator KV cache for new chunk --- // Much faster than freeing/recreating the context (~0.05s vs ~2s) llama_memory_clear(llama_get_memory(dft.ctx), false); const int n_layers = llama_model_n_layer(dft.model); cap = CaptureCtx(); cap.init(n_layers, n_chunk_tok); // --- Prefill --- cap.active = true; cap.reset(); llama_decode(dft.ctx, llama_batch_get_one(chunk_tok.data(), n_chunk_tok)); cap.active = false; cap.fini(); // --- Lookahead --- if (sp_lookahead > 0) { common_params_sampling sps = {}; sps.temp = 0.0f; common_sampler * smpl = common_sampler_init(dft.model, sps); for (int s = 0; s < sp_lookahead; s++) { llama_token tok = common_sampler_sample(smpl, dft.ctx, -1); common_sampler_accept(smpl, tok, true); cap.active = true; cap.reset(); llama_decode(dft.ctx, llama_batch_get_one(&tok, 1)); cap.active = false; cap.fini(); } common_sampler_free(smpl); } // --- Compute importance for this chunk --- auto imp = compute_imp(cap.steps, n_chunk_tok, n_layers); // --- Merge scores: keep the MIDDLE stride tokens (exclude overlap edges) --- // For the first chunk, keep [0..stride). For middle chunks, keep [stride..window-stride). // For the last chunk, keep all remaining. int keep_start = (n_chunks == 0) ? 0 : stride; int keep_end = (chunk_end >= n_input) ? n_chunk_tok : (sp_window - stride); // But keep_start/keep_end are relative to the chunk. Map to global positions. int gstart = pos + keep_start; int gend = pos + keep_end; if (gend > n_input) gend = n_input; for (int i = gstart; i < gend; i++) { int local = i - pos; if (local >= 0 && local < n_chunk_tok) { all_scores[i] += imp.scores[local]; score_cnt[i]++; } } LOG(" Kept scores for [%d..%d) (local [%d..%d))\n", gstart, gend, keep_start, keep_end); n_chunks++; pos += stride; } // Average overlapping scores for (int i = 0; i < n_input; i++) { if (score_cnt[i] > 0) all_scores[i] /= score_cnt[i]; else all_scores[i] = 1.0f; } auto t1 = ggml_time_us(); LOG(" Sliding window done: %d chunks, %.0f ms\n", n_chunks, (t1-t0)/1000.0); LOG(" Score range: [%f, %f]\n", *std::min_element(all_scores.begin(), all_scores.end()), *std::max_element(all_scores.begin(), all_scores.end())); // =================================================================== // Select important chunks (full prompt) // =================================================================== std::vector kept = select_chunks(all_scores, n_input, sp_chunk, sp_keep_pct, 128); int n_sparse = (int)kept.size(); // Cap sparse tokens to avoid OOM on target model's attention matrix. // When capped, keep every Nth token to maintain even coverage. if (n_sparse > sp_max_keep) { std::vector capped; int step = n_sparse / sp_max_keep; for (int i = 0; i < n_sparse && (int)capped.size() < sp_max_keep; i += step) capped.push_back(kept[i]); // Ensure last token is included if (capped.back() != kept.back()) { capped.back() = kept.back(); } kept = capped; n_sparse = (int)kept.size(); LOG(" Capped to %d tokens (max_keep)\n", n_sparse); } LOG("Selected %d/%d tokens (%.0f%%)\n", n_sparse, n_input, 100.0f*n_sparse/n_input); // =================================================================== // Sparse prefill on target // =================================================================== auto t2 = ggml_time_us(); llama_batch batch = llama_batch_init(n_sparse, 0, 1); common_batch_clear(batch); for (int i = 0; i < n_sparse; i++) common_batch_add(batch, inp[kept[i]], kept[i], {0}, (i == n_sparse-1) ? 1 : 0); int ret = llama_decode(tgt_ctx, batch); if (ret < 0) LOG_ERR("sparse prefill error: %d\n", ret); llama_batch_free(batch); auto t3 = ggml_time_us(); LOG("Sparse prefill: %.0f ms (%.0f t/s)\n", (t3-t2)/1000.0, n_sparse/((t3-t2)/1000.0*1000.0)); // =================================================================== // Decode // =================================================================== common_params_sampling smp = params.sampling; common_sampler * smpl = common_sampler_init(tgt_m, smp); for (int i = 0; i < n_input; i++) LOG("%s", common_token_to_piece(tgt_ctx, inp[i]).c_str()); LOG("\n"); int np = 0; auto t4 = ggml_time_us(); while (true) { llama_token id = common_sampler_sample(smpl, tgt_ctx, -1); common_sampler_accept(smpl, id, true); LOG("%s", common_token_to_piece(tgt_ctx, id).c_str()); np++; if (llama_vocab_is_eog(vocab, id)) break; if (params.n_predict >= 0 && np >= params.n_predict) break; llama_decode(tgt_ctx, llama_batch_get_one(&id, 1)); } auto t5 = ggml_time_us(); LOG("\n\n"); // =================================================================== // Results // =================================================================== double spec_ms = (t1-t0)/1000.0; double prefill_ms = (t3-t2)/1000.0; double decode_ms = (t5-t4)/1000.0; double total_ms = (t5-t0)/1000.0; double est_full = 1000.0 * n_input / 270.0; LOG("═ RESULTS ═══════════════════════════\n"); LOG("Prompt: %d → sparse: %d (%.0f%%)\n", n_input, n_sparse, 100.0f*n_sparse/n_input); LOG("Sliding: %7.0f ms (%5d chunks)\n", spec_ms, n_chunks); LOG("Sparse: %7.0f ms (%.0f t/s)\n", prefill_ms, n_sparse/(prefill_ms/1000.0)); LOG("Decode: %7.0f ms (%.0f t/s)\n", decode_ms, np/(decode_ms/1000.0)); LOG("Total: %7.0f ms\n", total_ms); LOG("Full (est): %7.0f ms\n", est_full); if (est_full > 0) LOG("Speedup: %.1f×\n", est_full / (total_ms + 1.0)); LOG("\n"); common_sampler_free(smpl); llama_free(tgt_ctx); llama_model_free(tgt_m); llama_free(dft.ctx); llama_model_free(dft.model); llama_backend_free(); return 0; } ```