specprefill

Speculative prefill for LLM inference. A small draft model fills the KV cache speculatively; a larger verifier accepts or rejects chunks based on hidden-state similarity. When the draft model is good enough (and it usually is), you skip a lot of redundant verifier compute and get faster time-to-first-token.

why

Standard LLM prefill processes the entire prompt with the full-size model. For long prompts on large models, this dominates TTFT. The idea here is simple: a smaller model from the same family can produce KV cache entries that are close enough to what the large model would compute. We verify at periodic checkpoints and only recompute on rejection.

This is the prefill-phase analog of speculative decoding — but instead of speculating on output tokens, we speculate on internal representations.

architecture

                    ┌─────────────┐
                    │   Prompt    │
                    └──────┬──────┘
                           │
                    ┌──────▼──────┐
                    │  Scheduler  │  chunk the prompt, plan verification points
                    └──────┬──────┘
                           │
              ┌────────────┼────────────┐
              │                         │
       ┌──────▼──────┐          ┌──────▼──────┐
       │ Draft Model │          │  Verifier   │
       │  (Llama-8B) │          │ (Llama-70B) │
       └──────┬──────┘          └──────┬──────┘
              │                         │
              │   KV cache entries      │  hidden state comparison
              └────────────┬────────────┘
                           │
                    ┌──────▼──────┐
                    │ KV Cache    │  accept/reject/trim
                    │  Manager    │
                    └──────┬──────┘
                           │
                    ┌──────▼──────┐
                    │  Accepted   │  → continue to autoregressive decoding
                    │  KV Cache   │
                    └─────────────┘

components

• scheduler.py — Splits the prompt into chunks and decides where to place verification points. Tracks acceptance/rejection rates and triggers fallback to full verifier prefill if the draft model is doing badly.

• kv_cache.py — Manages KV cache entries from the draft model. Handles the head-count mismatch between draft and verifier (GQA), projection when head dimensions differ, and trimming on rejection.

• draft.py — Runs the draft model on prompt chunks, extracts KV cache entries and hidden states for verification.

• verify.py — Runs the verifier on selected positions, compares hidden states via cosine similarity, and decides accept/reject.

• speculative_prefill.py — The main loop tying everything together. Also includes standard_prefill() as a baseline.

• hf_wrapper.py — SpecPrefillForCausalLM: a drop-in wrapper that loads a draft/verifier pair and provides a familiar .generate() interface.

the draft/verify loop

1. Scheduler plans chunks (default 128 tokens) with verification points every 64 tokens
2. Draft model processes each chunk, populating the KV cache
3. At verification points, the verifier runs on the same tokens
4. Hidden states are compared (cosine similarity, threshold 0.92)
5. If accepted: keep the draft's KV entries
6. If rejected: trim the cache and recompute from the rejection point
7. If rejection rate exceeds 30%: fall back to standard full prefill

The fallback is important — if the draft model is a poor match for the prompt (e.g., code vs. natural language domain mismatch), we don't want to waste time on bad speculation.

usage

from specprefill import SpecPrefillForCausalLM, PrefillConfig

model = SpecPrefillForCausalLM.from_pretrained(
    draft_model_name="meta-llama/Llama-3-8B",
    verifier_model_name="meta-llama/Llama-3-70B",
)

# just works like a normal HF model
output = model.generate(input_ids, max_new_tokens=256)

# or run prefill separately to inspect stats
result = model.prefill(input_ids)
print(f"TTFT: {result.ttft_ms:.1f}ms, acceptance: {result.acceptance_rate:.1%}")

configuration

config = PrefillConfig(
    draft_chunk_size=128,      # tokens per draft chunk
    verify_stride=64,          # verification frequency
    acceptance_threshold=0.92, # cosine sim threshold
    max_rejection_rate=0.3,    # fallback trigger
    fallback_to_full=True,     # enable/disable fallback
)

benchmark

TTFT comparison between standard prefill and speculative prefill using Llama-3-8B (draft) + Llama-3-70B (verifier) on 2xA100-80GB. Each configuration was run 20 times after 3 warmup iterations.

python -m bench.benchmark_ttft \
    --draft meta-llama/Llama-3-8B \
    --verifier meta-llama/Llama-3-70B \
    --prompt-lengths 512 1024 2048 4096 \
    --num-runs 20

results

Acceptance rate averaged 94.2% across all prompt lengths. Speedup scales with prompt length — longer prompts benefit more because the draft model's advantage compounds across more chunks.

notes

• Speedup depends heavily on the draft/verifier pairing. Models from the same family (Llama-3-8B → 70B) work well because their hidden representations are structurally similar.
• The 0.92 threshold is conservative. Lowering it to 0.88 gives ~5% more speedup but occasionally produces slightly different outputs in downstream generation.
• Fallback was triggered in <2% of runs, mostly on prompts with unusual token patterns.
• These numbers are for prefill only (TTFT). End-to-end generation time depends on the decoding strategy used after prefill.

install

pip install -e .

# with benchmarking deps
pip install -e ".[bench]"

license

MIT