sliding-specprefill/specprefill.md

SpecPrefill — Implementation Guide

Paper: Speculative Prefill: Turbocharging TTFT with Lightweight and Training-Free Token Importance Estimation (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

---

---

Table of Contents

1. Overview
2. How It Works
3. Sliding Window
4. Benchmarks
5. Configuration Reference
6. Implementation Details
7. GPU Memory Analysis
8. Limitations
9. 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 top ~30% of tokens, reducing TTFT by 1.4-1.6×.

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
    ↓
Group N tokens into 32-token contiguous chunks
    ↓
Keep top ~30% of chunks by average importance
    ↓
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):

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)

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).

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:

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 (llama-server --specprefill)

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×
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

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.

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.

---

Configuration Reference

All configuration is via environment variables:

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_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)

Recommended Usage

# 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

# Low latency mode (more aggressive sparsification)
SP_KEEP_PCT=0.2 SP_LOOKAHEAD=1 ./bin/llama-specprefill ...

# High quality mode (less sparsification)
SP_KEEP_PCT=0.5 SP_LOOKAHEAD=4 ./bin/llama-specprefill ...

---

Implementation Details

File Structure

File	Purpose
examples/specprefill/specprefill.cpp	Main implementation (~470 lines)
examples/specprefill/CMakeLists.txt	Build config (3 lines)
docs/optimization/specprefill.md	This document

Build

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.
• select_chunks — Groups into chunks, ranks by average score, selects top-K.
• load_model — Creates llama context with cb_eval set and flash attention disabled (so kq_soft_max tensors are materialized).

Sliding Window Loop

for each chunk:
    llama_memory_clear(mem, false);  // clear KV cache (fast, ~0.05s)
    cap = CaptureCtx();
    cap.init(n_layers, chunk_size);
    
    // Prefill
    cap.active = true;
    llama_decode(ctx, batch_get_one(chunk_tokens));
    cap.active = false;
    cap.fini();
    
    // Lookahead decode steps
    for lookahead step:
        sample token
        cap.active = true;
        llama_decode(ctx, batch_get_one(&token));
        cap.active = false;
        cap.fini();
    
    // Compute importance for this chunk
    imp = compute_imp(cap.steps, chunk_size, n_layers);
    
    // 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;

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.

---

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. 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 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 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%.

5. 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.

---

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
2. Original vLLM implementation: github.com/Jingyu6/speculative_prefill
3. vLLM feature request: github.com/vllm-project/vllm/issues/39060
4. vLLM-mlx implementation: github.com/waybarrios/vllm-mlx/issues/179

---

Full Source Code

examples/specprefill/specprefill.cpp (473 lines):

// 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 <algorithm>
#include <cfloat>
#include <clocale>
#include <cmath>
#include <cstdint>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <numeric>
#include <string>
#include <vector>

// ============================================================================
// Backend: kq_soft_max capture via cb_eval
// ============================================================================

struct CapturedAttn {
    int il;
    int64_t n_k;
    int64_t n_h;
    std::vector<float> data;
};

struct CaptureCtx {
    bool active = false;
    int n_layers = 0;
    int n_tokens = 0;
    std::vector<CapturedAttn> layers;
    std::vector<std::vector<CapturedAttn>> 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<float> scores;
    int n_tokens;
};

static ImpResult compute_imp(
    const std::vector<std::vector<CapturedAttn>> & 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<std::vector<float>> smax(ns, std::vector<float>(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<int> select_chunks(const std::vector<float> & 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<int> a(nt); std::iota(a.begin(), a.end(), 0); return a;
    }

    int nc = (nt + csize - 1) / csize;
    std::vector<double> cs(nc, 0);
    std::vector<int> cc(nc, 0);
    for (int i = 0; i < nt; i++) { int c = i / csize; cs[c] += scores[i]; cc[c]++; }
    std::vector<float> 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<float> sv = ca; std::sort(sv.begin(), sv.end(), std::greater<float>());
    float th = kc <= (int)sv.size() ? sv[kc-1] : -FLT_MAX;

    std::vector<int> 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<int> 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<llama_token> 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<float> all_scores(n_input, 0.0f);
    std::vector<int>   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<llama_token> 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<int> 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<int> 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;
}