diff --git a/examples/performance/speculative_prefill/TODO_technical_summary.md b/examples/performance/speculative_prefill/TODO_technical_summary.md index d642fb0220..1e2e39de64 100644 --- a/examples/performance/speculative_prefill/TODO_technical_summary.md +++ b/examples/performance/speculative_prefill/TODO_technical_summary.md @@ -4,15 +4,17 @@ Speculative prefill (SPECPREFILL) and your integration - High-level algorithm (look-ahead anchors supported; default look-ahead = 0): - Two-model pipeline: - Speculator (small model) runs full-prompt prefill to produce queries (Q) and per-layer keys (K). + - Q capture path (flag-gated): `QEffLlamaAttention.forward` emits last-token per-head queries (`q_last`), collected in `QEffLlamaDecoderLayer` and stacked in `QEffLlamaModel` as `prefill_queries`, surfaced via `QEffLlamaForCausalLM`. + - K capture path: retained-state `past_key.*_RetainedState` outputs enabled only on the final prefill chunk; harvested once from that chunk. - Importance computation: - - Per-head softmax over positions. - - Head-wise max, then layer-wise max. + - Host NumPy computes logits = Q·Kᵀ/√D per head (handles head_dim padding), softmax over positions. + - Head-wise max (or sum), then layer-wise max; look-ahead anchors averaged if enabled. - 1D smoothing (moving average). - - Block top‑K by mean score; force-keep last token. + - Block top‑K by mean score; force-keep last token; selection is percentage-based. - Restore absolute position_ids for the kept tokens. - Base model runs prefill only on selected tokens; decode continues from original context length S (full position indexing preserved). - Host-side scoring: - - Read Q and retained K to host; importance and selection computed on CPU to avoid modifying ONNX subgraphs. + - Read Q and retained K (final chunk) to host; importance and selection computed on CPU to avoid modifying ONNX subgraphs; retained-state K may be head-dim padded and is sliced to S before scoring. - Note: reading K to host can be DMA-heavy for very long contexts but simplifies deployment. - Integration into QEfficient (flag-gated, minimal disturbance to baseline): @@ -68,3 +70,24 @@ Base: ctx_len/prefill_seq_len must align with spec; CB supported when compiled w - Base model must align ctx_len and prefill_seq_len with speculator’s assumptions so position restoration remains consistent. - CB is supported for the base when compiled appropriately (full_batch_size, batch_index wiring). - The integration is opt-in; baseline generate code path untouched when flag off. + +Example: host-side scoring with tiny numbers (matches speculative_prefill_engine.py lines ~529–607) +- Dimensions: L=2 layers, H=2 attn heads, H_kv=1 KV head (GQA), D=2, S_total=3 tokens, no smoothing; head aggregation = max, layer aggregation = max. +- Inputs after slicing/padding fixes: + - Q_final [L,H,D]: + - Layer0: [[1.0,0.0],[0.0,1.0]] + - Layer1: [[0.5,0.5],[1.0,1.0]] + - K_global list length L; each [H_kv,S_total,D]: + - Layer0: [[[1.0,0.0],[0.0,1.0],[1.0,1.0]]] + - Layer1: [[[0.5,0.0],[0.0,0.5],[0.5,0.5]]] +- GQA mapping: group_size=H/H_kv=2; head_to_kv=[0,0]; K_sel per layer expands to [H,S,D]=[2,3,2] with both heads seeing the single KV head slice. +- Layer0 scoring (einsum hsd,hd->hs): + - Q0 head0 [1,0] vs K_sel0 → logits [1,0,1]; scale 1/sqrt(2) → [0.7071,0,0.7071]; softmax → [0.401,0.198,0.401]. + - Q0 head1 [0,1] → logits [0,1,1]; scale → [0,0.7071,0.7071]; softmax → [0.198,0.401,0.401]. + - Head max → importance_layer0 ≈ [0.401,0.401,0.401]. +- Layer1 scoring: + - Q1 head0 [0.5,0.5] vs K_sel1 → scaled logits [0.177,0.177,0.354]; softmax → [0.313,0.313,0.373]. + - Q1 head1 [1,1] → scaled logits [0.354,0.354,0.707]; softmax → [0.292,0.292,0.416]. + - Head max → importance_layer1 ≈ [0.313,0.313,0.416]. +- Layer aggregation (max across layers): importance ≈ [0.401,0.401,0.416] (shape [S_total]). + - Einsum detail: `np.einsum("hsd,hd->hs", K_sel, Q_l)` contracts over D, yielding [H,S]; order is K then Q but equivalent to q·kᵀ per head/token because D matches.