Speculative prefill

[rockofox]

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?

[rockofox]

@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

[rockofox]

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