hydran-lab/sliding-specprefill
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
71f9993b380e61daf4357e0633de2a6ea1a59eba
sliding-specprefill/specprefill.md
T

43 KiB
Raw Blame History

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

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

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)

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:

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

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
  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;
}
Reference in New Issue View Git Blame Copy Permalink