Speculative prefill

[kilofox]

Overview

This PR implements Speculative Prefill, based on the ICML 2025 paper: "Speculative Prefill: Turbocharging TTFT with Lightweight and Training-Free Token Importance Estimation" by Jingyu Liu, Beidi Chen, Ce Zhang (arXiv, GitHub)

Note

In contrast to speculative decoding, this is not lossless and can cause model performance degradation.

Approach

In long-context inference, prefill TTFT scales quadratically with prompt length. Speculative prefill reduces this compute by using a smaller draft model to estimate which context tokens are most relevant to the query before evaluating the main target model:

1. Draft lookahead: The draft model prefills the prompt and generates $K$ lookahead tokens (default: 8).
2. Attention extraction: Softmax attention matrices (kq_soft_max) are captured during draft generation via an evaluation callback.
3. Importance pooling & chunking: Attention scores are smoothed with a 1D average-pooling filter ($W = 13$), max-reduced across heads/layers, and aggregated into contiguous chunks ($C = 32$).
4. Sparse target prefill: Only the top-$p$ fraction of chunks (plus BOS and prompt boundary tokens) are fed into the target model's KV cache.
5. Generation: The target model generates tokens autoregressively as normal.

Benchmarks & Performance Evaluation

All benchmarks were run natively on Vulkan (AMD Radeon 8060S, RADV STRIX_HALO, 128 GiB unified LPDDR5X) using Qwen3.8-27B (Q6_K_XL) as target and Qwen3.5-2B (Q4_K_XL) as the primary draft model (unless noted otherwise).

1. Draft Model Sizing & Context Scaling

Draft Model Comparison ($N = 1609$, $p = 0.15$):

Configuration	Draft Eval Latency	Target Prefill	Total TTFT	Target Compute	End-to-End Speedup
Baseline (27B alone)	—	5,300.8 ms	5,300.8 ms	1.00x	1.00x
27B Target + 27B Draft	5,836.5 ms	789.2 ms	7,597.2 ms	6.72x faster	0.70x (draft overhead)
27B Target + 9B Draft	1,486.7 ms	781.2 ms	2,505.7 ms	6.79x faster	2.12x
27B Target + 2B Draft	333.9 ms	788.8 ms	1,196.4 ms	6.72x faster	4.43x

Context Length Scaling (27B Target + 2B Draft):

Keep Ratio ($p$)	$N = 803$ tokens	$N = 1607$ tokens	$N = 3212$ tokens
100% (Baseline)	2,882.0 ms (1.00x)	5,324.6 ms (1.00x)	10,787.7 ms (1.00x)
50%	1,795.3 ms (1.61x)	3,273.8 ms (1.63x)	6,174.8 ms (1.75x)
30%	1,026.3 ms (2.81x)	2,045.3 ms (2.60x)	4,004.3 ms (2.69x)
15%	685.1 ms (4.21x)	1,615.7 ms (3.30x)	2,418.1 ms (4.46x)

2. Synthetic Retrieval & QA Benchmarks

• Needle-in-a-Haystack (NIAH): 100.0% retrieval accuracy across depths (20%, 50%, 80%) down to $p = 0.15$ (298 / 2,058 tokens kept, 1,798.0 ms vs 6,863.2 ms baseline, 3.82x speedup).
• RULER Suite: 100.0% accuracy across Single Needle (25%, 75%), Multi-Key Retrieval, and Variable Tracking at $p = 0.15$ (4.11x speedup, 1,930.7 ms vs 7,925.4 ms). One miss at $p = 0.50$ (depth-75% needle; that row 75.0%).
• LongBench QA Latency (HotpotQA, Qasper, 2WikiMQA): Latency reduced from 4,821.5 ms baseline to 1,313.7 ms at $p = 0.15$ (3.67x speedup).

Failure Mode Analysis & Operational Limits

Because speculative prefill drops tokens permanently from target prefill, accuracy drops when the token budget ($p \cdot N$) cannot contain all non-redundant evidence across separate chunks.

1. Draft Model Comparison on Stress Tests (2B vs. 9B)

High-Entropy Multi-Key Retrieval (5 keys across 5 sections, ~1,100 tokens):

Keep Ratio ($p$)	Tokens Kept	Qwen 2B Draft	Qwen 9B Draft	Observed Behavior
100% (Baseline)	971 / 971	PASS (5/5)	PASS (5/5)	All keys present and retrieved correctly
50%	491 / 971	FAIL (0/5)	FAIL (2/5)	Partial chunk pruning causes model hesitation/evasion
30%	299 / 971	PASS (5/5)	FAIL (2/5)	2B focused on all 5 nodes; 9B diffused attention
15%	140 / 971	FAIL (0/5)	FAIL (3/5)	Budget allows ~4 chunks (5 needed) $\to$ information loss
8%	76–98 / 971	FAIL (2/5)	FAIL (1/5)	Severe pruning $\to$ both models hallucinate fake IDs

Distributed Summation (Box A: 15 + Box B: 25 + Box C: 60 = 100):

Keep Ratio ($p$)	Tokens Kept	Qwen 2B Draft	Qwen 9B Draft	Status / Observed Behavior
100% (Baseline)	796 / 796	100 units	100 units	PASS — All 3 box chunks retained
50%	412 / 796	100 units	100 units	PASS — All 3 box chunks retained
30%	252 / 796	100 units	100 units	PASS — All 3 box chunks retained
15%	125 / 796	100 units	100 units	PASS — Boundary + 3 box chunks fit in 125 tokens
8%	61 / 796	40 units	77 units	FAIL — Box C (60) pruned $\to$ silent partial calculation

Non-Parametric Multi-Hop QA (Target Year: 2348, 12 non-redundant documents):

Keep Ratio ($p$)	Tokens Kept	Qwen 2B Draft	Qwen 9B Draft	Status / Observed Behavior
100% (Baseline)	473 / 473	2348	2348	PASS — Both clue documents retained
50%	249 / 473	2348	Refusal / Evasion	PASS (2B) / FAIL (9B) — 9B lookahead diverted to distractor
30%	153 / 473	2348	2348	PASS — Both clue documents retained
15%	90 / 473	2348	2348	PASS — Both clue documents retained
8%	58 / 473	2280	2185	FAIL — Clue 2 pruned $\to$ hallucinated from distractor

2. Multi-Ratio Stress Matrix (Qwen 27B + 2B)

• Task 1 (Latent Multi-Hop): Clue 1 matches question; Clue 2 shares no tokens with question.
• Task 2 (Distributed Aggregation): Summing 6 distributed facility incident counts (total: 27).
• Task 3 (Temporal Revision): Conflicting notices with emergency override ZETA-9999.
• Task 4 (Distributed Variable Tracking): Sequential operations spread across distant chunks.

Keep %	Tokens Kept	Task 1	Task 2	Task 3	Task 4	Accuracy
100% (Baseline)	1170 / 1170	PASS	FAIL	PASS	FAIL	50.0%
50%	594 / 1170	PASS	PASS	PASS	FAIL	75.0%
30%	370 / 1170	PASS	PASS	PASS	FAIL	75.0%
15%	178 / 1170	PASS	FAIL	PASS	FAIL	50.0%
8%	83 / 1170	PASS	FAIL	PASS	FAIL	50.0%
5%	51 / 1170	PASS	FAIL	PASS	FAIL	50.0%

Practical Operating Guidelines

• Draft Model Selection: Qwen 2B is the recommended draft model for Qwen 27B. The 2B model evaluates ~4× faster than the 9B model (488 ms vs 2,004 ms) while maintaining comparable chunk selection accuracy. Larger draft models do not resolve token budget limits at $p \le 0.08$.
• $p = 0.25 - 0.35$: Recommended general-purpose default (preserves multi-chunk evidence, delivers ~2.0–2.3× speedup).
• $p = 0.15$: Effective for focused single-needle retrieval and localized QA (~3.5× speedup).
• $p \le 0.10$: High risk of information loss in non-redundant contexts with distributed evidence.
• Whole-document summarization: Not recommended for low keep ratios ($p < 0.50$) due to uniform information distribution across the document.

Example Usage

# Build
cmake -B build -DGGML_VULKAN=ON -DCMAKE_BUILD_TYPE=Release
cmake --build build --target llama-speculative-prefill -j

# Run
./build/bin/llama-speculative-prefill \
    -m models/Qwen3.8-27B-Q6_K.gguf \
    -mpd models/Qwen3.5-2B-Q4_K.gguf \
    -ngl 99 -nglpd 99 \
    -p "Long prompt context goes here..." \
    --spec-prefill-percentage 0.20 \
    --spec-prefill-chunk-size 32 \
    --spec-prefill-lookahead 8 \
    -n 128

Additional information

Has been mentioned as a feature request here: #19082

Requirements

• I have read and agree with the contributing guidelines: Yes
• AI usage disclosure: YES, mostly for code structure, debugging and writing some of the benchmark/testing code

[ggml-gh-bot]

Hi @rockofox, thanks for your contribution!

Per our contribution guidelines, the automated PR checker found the following issue(s) that need your attention:

• AI-generated content: While code is allowed to be generated by AI, please write the PR description and commit messages on your own without the help of AI.

• Large PR: Large changes require prior discussion (e.g. an issue or RFC) and maintainers may not be able to review this PR as-is. Consider splitting it into smaller, focused PRs.

---

Please note that maintainers reserve the right to make final decisions on PRs. If you believe there is a mistake, please comment below.

[tctien342]

Can we set which device to put prefill draft model on ? I have a 3060ti + 2x5060ti setup, it much more convenience to put that draft model on 3060ti and let dual 5060ti handle full 27b model instead.

Btw, thanks you so much 👏

[gopinath87607]

very interesting one i think this would help us to use the rpc gpu. right now with help of deflash2 i was able to decode fast almost near to the sm tensor level but when i am using the rpc gpu the prefill went to really down which hurt the using the rpc gpus this would help everyone if its woks well i will test it when i have the time.
i have rtx30602 and rtx5062 and via rpc gpu rtx 3080 16gb.

[gopinath87607]

Bug: DFlash2 draft decode crashes when combined with speculative prefill (non-contiguous KV positions)

Summary

Running --spec-type draft-dflash (PR #27342) together with --spec-prefill (PR #27692) in the same llama-server instance causes the DFlash2 draft model's decode to fail on the very first generation request, with:

E init: sequence 0 positions are not continuous
E decode: failed to initialize batch
E llama_decode: failed to decode, ret = -1
E process: llama_decode(ctx_dft) failed rc=-1 (n_tokens=80, offset=0)
E srv        decode: failed to process speculative batch

Each feature works correctly on its own (confirmed separately below). The crash only occurs when both are active at the same time.

Setup

• Base: master @ f280b2698 (includes model: add dots3-note #27060 dots3-note)
• Merged on top: PR spec : add DFlash2 support (local convolution + candidate selector) #27342 (DFlash2, commit range 64f765f5a..f7aadef09) and PR Speculative prefill #27692 (Speculative Prefill, a1f3ad0..ee7f301)
• Model: Qwen3.8-27B target (Q8_K_XL), DFlash2 draft (BF16), separate ~2B speculative-prefill draft model (Q4_K_M)
• 5x GPU tensor-split + 1 RPC backend

Relevant server flags

--model-draft <dflash2 draft gguf> -devd CUDA3 --spec-type draft-dflash --spec-draft-n-max 7 --spec-draft-p-min 0.01
--spec-prefill -mpd <prefill draft gguf> -nglpd 99 -devpd CUDA0 --spec-prefill-percentage 0.20 --spec-prefill-chunk-size 32 --spec-prefill-lookahead 8

Sequence of events from the log

1. Server loads target model, DFlash2 draft, and speculative-prefill draft — all three load successfully, no errors.
2. On the first request (372-token prompt), speculative prefill runs first and works as designed:

   I slot apply_spec_p: id  0 | task 0 | speculative prefill kept 84 / 372 tokens (22.6%)
   

   i.e. it drops ~77% of prompt tokens, keeping only the "important" ones per its attention-based selection — but keeps their original absolute sequence positions, so the retained tokens have gaps between them (contiguous runs separated by skipped ranges).
3. Immediately after, the DFlash2 draft context (ctx_dft) tries to decode against that same sparse sequence and fails:

   E init: sequence 0 positions are not continuous
   E decode: failed to initialize batch
   

4. This aborts the whole request with a 500 error.

Root cause hypothesis

Speculative prefill's sparse retention is designed to feed a target model with attention over the full theoretical context (i.e. it's fine with position gaps at the target-model level, since the target model's KV cache is built directly from the retained sparse token set). DFlash2's draft decode path (ctx_dft), however, appears to assume the sequence of positions it's asked to decode is contiguous — there's an explicit continuity check (init: sequence 0 positions are not continuous) that has no tolerance for the gaps that speculative prefill introduces upstream of it.

In other words: two independently-developed PRs each made a valid assumption about the KV cache/position handling that the other silently breaks. Neither PR was written with the other in mind — they were merged upstream only days apart, and I don't see evidence either was tested against the other.

Reproduction

1. Merge spec : add DFlash2 support (local convolution + candidate selector) #27342 and Speculative prefill #27692 on top of current master.
2. Start llama-server with both --spec-type draft-dflash (+ --model-draft) and --spec-prefill (+ -mpd) set simultaneously.
3. Send any chat completion request.
4. Server returns 500 with decode() failed: failed to process speculative batch on the first turn.

Workaround

Each feature works fine in isolation — running with only DFlash2 or only speculative prefill enabled, everything loads and serves correctly. The bug only manifests when both are active together.

Suggested fix direction (not verified, just a guess from the log)

Either:

• DFlash2's draft decode path needs to tolerate non-contiguous positions (mirror however the target model already handles the sparse sequence from speculative prefill), or
• Speculative prefill needs an option to remap retained tokens to contiguous positions for consumption by any secondary draft context, rather than passing through the original absolute positions.

[gopinath87607]

@rockofox can you check above error?

[kilofox]

@gopinath87607 I'm not sure if compatibility with DFlash2 is feasible right now, especially since that's not merged yet. Let's see

[gopinath87607]

@gopinath87607 I'm not sure if compatibility with DFlash2 is feasible right now, especially since that's not merged yet. Let's see

should i test it with dflash2 and rpc ?
btw the main idea is using dflash +preflash aka prefill is to make use of my ideal rpc gpu.

if possiple maybe you can change change name from prefill to preflash little bit attractive

[kilofox]

Like I mentioned earlier, DFlash2 currently cannot work with speculative prefill, since it needs a full, continous prefill context. It should be possible to make it work, but this is out of scope for this PR for now, especially since DFlash2 hasn't been merged to master yet. I can revisit compatibility with DFlash2 once both PRs are merged.

In the meantime, my personal advice is that you try using MTP instead of DFlash2. It's also quite fast (especially when combined with ngram-mod) and works with spec prefill.
And I haven't tested rpc, but it should work.

---

if possiple maybe you can change change name from prefill to preflash little bit attractive

I don't think that makes sense, since the name is based on the original paper for clarity. Also, the word "Flash" would imply that this has something to do with e.g DFlash, which it hasn't really

[johnkarlhill]

Nice work!

Found a bug while testing speculative prefill. The draft context's KV cache never gets freed between requests, so after 2-3 sequential requests the draft fills up, decode fails, and speculative prefill silently stops working until you restart the server.

How to reproduce (llama-server, --ctx-size 120000, default n_parallel=4):
Send request A with a ~50K prompt. Works, spec-prefill engages.
Send request B with a ~50K prompt from a different session. Works.
Send request C with a ~50K prompt from a different session. Fails with llama_decode: failed to decode, ret = -1 on the draft. Spec-prefill stays disabled for every request after that.

What's happening: the spec-prefill draft context (ctx_spf) is a single shared llama_context with n_ctx = params_base.n_ctx, and its KV is keyed by slot id. The only place its KV gets cleared is apply_spec_prefill() (server-context.cpp:1056), and that only removes the current slot's seq before re-decoding. Slot release, idle-slot purge (try_clear_idle_slots), and slot takeover never touch ctx_spf. So every completed slot leaves its draft KV behind. With 4 slots and ~50K prompts, the draft context accumulates 150K > 120K and decode fails. The target context has try_clear_idle_slots-style recovery when it fills up, but the draft has nothing like that, so it stays broken until restart. The draft KV is never reused across requests anyway (every request re-decodes the full prompt), so freeing it on slot release costs nothing.

Proposed fix (tested): clear the draft seq when a slot's KV is released. Add a guarded llama_memory_seq_rm(llama_get_memory(ctx_spf), seq_id, -1, -1) helper and call it from prompt_clear() (server-context.cpp:354), release() (:583), and copy_state_to() (:763). That's 6 lines in server-context.cpp. Verified with 4 sequential ~54K-token requests: all 4 slots engage spec-prefill at 30% kept, no decode failures.

Environment: Intel Arc B70/B50 (SYCL), Windows, Qwen3.8-27B target + Qwen3.5-0.8B draft, --spec-prefill --spec-prefill-chunk 64 --spec-prefill-lookahead 16.
Ran a 4-test synthetic known-answer suite (multi-key retrieval, aggregation, multi-hop, temporal-override) on 76K-token prompts at temp 0. All 4 stress tests produced identical answers with spec-prefill p=0.20 vs the true baseline, at 1.92x speedup (112.8s → 58.8s avg).

set SYCL_CACHE_PERSISTENT=1
set GGML_SYCL_FA_ONEDNN=1
REM  -ctkd q8_0
REM -dev SYCL0 pin main model to 0
REM -devpd SYCL1 pin spec-prefill model to 1
REM -devd SYCL0 pin MTP draft to 0
C:\llama.cpp-build-sycl\llama.cpp\build-x64-windows-sycl-release\bin\llama-server.exe ^
  -m C:\llama.cpp-prod\models\B70\Qwen3.8-27B-Q6_K.gguf ^
  -mpd C:\llama.cpp-prod-b50\models\B50\Qwen3.5-0.8B-UD-Q5_K_XL.gguf ^
  -ngl 999 -nglpd 99 ^
  -dev SYCL0 -devpd SYCL1 -devd SYCL0 ^
  --chat-template-file C:\llama.cpp-prod\models\B70\chat_template.jinja ^
  --temperature 0.6 --top-p 0.95 --top-k 20 --min-p 0.0 --presence-penalty 0.1 --repeat-penalty 1.06 ^
  --reasoning true --reasoning-budget 8192 --reasoning-preserve ^
  --ctx-size 163840 --flash-attn on^
  --cache-type-k q8_0 --cache-type-v q8_0 ^
  --kv-unified --cache-reuse 512 --context-shift --cache-idle-slots --jinja ^
  --spec-type draft-mtp --spec-draft-p-min 0.3 --spec-draft-p-split 0.15 --spec-draft-n-max 3 ^
  --spec-prefill --spec-prefill-p 0.20 --spec-prefill-chunk 64 --spec-prefill-lookahead 16 ^
  --batch-size 1024 --ubatch-size 1024 --parallel 1^
  --load-mode none ^
  --host 0.0.0.0 --port 8081 ^
  --log-timestamps

| model           |     test |            t/s |     peak t/s |        ttfr (ms) |     est_ppt (ms) |    e2e_ttft (ms) |
|:----------------|---------:|---------------:|-------------:|-----------------:|-----------------:|-----------------:|
| Qwen38-27B-Q6_K | pp100000 | 1316.87 ± 0.95 |              | 75939.35 ± 54.57 | 75938.00 ± 54.57 | 75939.35 ± 54.57 |
| Qwen38-27B-Q6_K |    tg512 |   34.24 ± 0.42 | 34.67 ± 0.47 |                  |                  |                  |

| model           |     test |            t/s |     peak t/s |        ttfr (ms) |     est_ppt (ms) |    e2e_ttft (ms) |
|:----------------|---------:|---------------:|-------------:|-----------------:|-----------------:|-----------------:|
| Qwen38-27B-Q6_K | pp135000 | 1097.97 ± 0.00 |              | 122956.59 ± 0.00 | 122955.26 ± 0.00 | 122956.59 ± 0.00 |
| Qwen38-27B-Q6_K |    tg512 |   32.12 ± 0.00 | 33.00 ± 0.00 |                  |                  |                  |

| model           |     test |           t/s |     peak t/s |        ttfr (ms) |     est_ppt (ms) |    e2e_ttft (ms) |
|:----------------|---------:|--------------:|-------------:|-----------------:|-----------------:|-----------------:|
| Qwen38-27B-Q6_K | pp160000 | 985.81 ± 0.00 |              | 162305.44 ± 0.00 | 162304.47 ± 0.00 | 162305.44 ± 0.00 |
| Qwen38-27B-Q6_K |    tg512 |  31.04 ± 0.00 | 32.00 ± 0.00 |                  |                  |                  |

[kilofox]

@johnkarlhill Thanks! I implemented your fix and it seems to work fine

[kilofox]

@gopinath87607 I managed to fix the issue regarding DFlash2. It was a different, simpler problem

[RealFascinated]

could this have a new max ctx option for the prefill. this is mostly impossible to use on small vram system where it's most important

[nr23730]

Hi
I tested with together with #27952 and it works like a charm!
One thing that I did notice is the following: Previously some applications/coding agent were running into timeouts and would then try again. Then the request would land in the same slot and the prefill would basically start were it stopped in the previous request. Now it doesn't do this anymore and would start all over again.

I noticed that when using Qwen3.8-27b with DeepSeek Harness when cold starting a large context window of around 150k tokens.
Is that intentional?

[HumerousGorgon]

Would it ever be possible to use the MTP layers in a model to do speculative prefill?

[gopinath87607]

CUDA_VISIBLE_DEVICES=2,3,0,1
~/llama.cpp/build/bin/llama-server
--model /mnt/nvme/Qwen3.8-27B-UD-Q8_K_XL.gguf
--model-draft /mnt/nvme/Qwen3.8-27B-DFlash2-BF16.0.gguf
-mpd /mnt/nvme/qwen3-5-2B-Q4_K_M.gguf
--chat-template-kwargs '{"reasoning_effort":"xhigh"}'
--spec-prefill-percentage 0.35
--spec-prefill-chunk-size 16
--spec-prefill-lookahead 32
--spec-type draft-dflash
--spec-draft-n-max 7
--reasoning-preserve
--spec-draft-p-min 0.01
--rpc 10.0.0.2:50053
-devpd CUDA2
-devd CUDA3
-dev RPC0,CUDA0,CUDA1,CUDA2,CUDA3
-ot '^output_norm.weight=CUDA3,^output.weight=CUDA3'
--tensor-split 2.7,1.8,2,0.3,0
--flash-attn 1
--fit off
-ngl 99
-ngld 99
-nglpd 99
--ctx-size 150000
--jinja
--parallel 1
--threads 28
--threads-batch 28
--host 0.0.0.0
--port 8082

i tested like this looks well but still there is a break some where like after 80k context like gpu is working but words is not shown in the ui and also sometime i feel the prompt might gets truncate the important tokens. maybe a specially designed fine tuned model might fix this issues as far i know its still need the work like if model wants to see some context its should ask small model to show it and also the speculative prefill need to have cpu moe option so that they can use big moe models offloading big layers to cpu and keeping low layer

like this way we can compain the qwen 3.8 30b and 3b and 27b dense and create cool one but this all is a complex design lol just sharing my thoughts