From a9341ac8bcb81bab39864391812db65806334173 Mon Sep 17 00:00:00 2001 From: Advay Monga Date: Tue, 6 Oct 2026 23:33:26 -0400 Subject: [PATCH] docs: drop papers/; the reading notes moved to BlameGraph with the environment Co-Authored-By: Claude Fable 5.1 --- docs/papers/README.md | 14 ----- docs/papers/batchllm.md | 65 ------------------------ docs/papers/elis.md | 63 ----------------------- docs/papers/flashattention.md | 65 ------------------------ docs/papers/loki.md | 68 ------------------------- docs/papers/speculative-decoding.md | 79 ----------------------------- docs/papers/tokenweave.md | 64 ----------------------- docs/papers/vllm-paged-attention.md | 67 ------------------------ 8 files changed, 485 deletions(-) delete mode 100644 docs/papers/README.md delete mode 100644 docs/papers/batchllm.md delete mode 100644 docs/papers/elis.md delete mode 100644 docs/papers/flashattention.md delete mode 100644 docs/papers/loki.md delete mode 100644 docs/papers/speculative-decoding.md delete mode 100644 docs/papers/tokenweave.md delete mode 100644 docs/papers/vllm-paged-attention.md diff --git a/docs/papers/README.md b/docs/papers/README.md deleted file mode 100644 index 7de4d49..0000000 --- a/docs/papers/README.md +++ /dev/null @@ -1,14 +0,0 @@ -# papers/ — reading notes - -Notes on the papers this engine borrows from, each with the idea, what it buys, and what (if -anything) it changed here. - -| note | idea | status in this engine | -|---|---|---| -| `vllm-paged-attention.md` | block-paged KV, PagedAttention | implemented (`models/paged_kv_cache.py`, Triton kernels) | -| `flashattention.md` | IO-aware tiled attention | tiled prefill kernel built, measured, rejected end-to-end (head_dim 512) | -| `batchllm.md` | prefix-sharing-aware batching | radix prefix cache; wave grouping measured inert | -| `elis.md` | length-predictive scheduling | not pursued — TTFT tail is prefill compute, not queueing | -| `speculative-decoding.md` | draft-and-verify decode | not pursued yet | -| `loki.md` | low-rank keys for sparse attention | not pursued | -| `tokenweave.md` | compute/communication overlap | single-GPU engine; not applicable yet | diff --git a/docs/papers/batchllm.md b/docs/papers/batchllm.md deleted file mode 100644 index 593296b..0000000 --- a/docs/papers/batchllm.md +++ /dev/null @@ -1,65 +0,0 @@ -# BatchLLM: Optimizing Large Batched LLM Inference - -**Link:** https://arxiv.org/abs/2412.03594 -**Venue:** MLSys 2026 -**Authors:** Zhen Zheng, Xin Ji, Taosong Fang, Fanghao Zhou, Chuanjie Liu, Gang Peng -**Status:** not read -**Relevant phase:** Phase 5 (prefix sharing) + Phase 6 (continuous batching, scheduler) - -## Core idea - -Existing engines like vLLM and SGLang share KV prefixes opportunistically — typically via an LRU-managed prefix tree that catches reuse only when requests happen to arrive close in time. Under bursty multi-request workloads with many shared prefixes (system prompts, few-shot demonstrations, RAG templates), this opportunistic approach prematurely evicts prefixes that will be reused shortly and misses sharing opportunities entirely. BatchLLM treats prefix sharing as a **global scheduling problem**: it identifies common prefixes across the entire pending request set up front, groups requests by prefix, and reorders execution so requests that share a prefix run together. It also schedules requests with a higher decode-to-prefill token ratio first to keep the device fed with decode work while prefill chunks are interleaved. - -## Why it matters - -- Directly relevant to the project's radix-tree prefix sharing. The radix tree as designed catches prefix matches when they exist in the cache *at request arrival time*; BatchLLM's contribution is to *change the order of arrivals into the engine* so matches happen more often. -- The reported 1.3×–10.8× win over vLLM and SGLang on representative workloads is large enough to take seriously even after accounting for cherry-picked batches. The high end of that range is on workloads with heavy prefix overlap, which is exactly the multi-tenant / RAG / agent regime. -- The "decode-heavy first" scheduling rule is a concrete heuristic that fits cleanly into the planned `SchedulerInterface` — it is a policy, not an architectural change. - -## Key mechanism - -1. **Global prefix identification.** Across the current pending request set, common prefixes are extracted (the abstract describes this as global rather than per-arrival). Mechanically, this corresponds to building a prefix tree over the *queue* of pending requests, not just the cache, so the scheduler can see all sharing opportunities before dispatch. -2. **Prefix-grouped scheduling.** Requests sharing a prefix are scheduled together so the prefix's KV is computed once and reused by all members of the group while still resident. -3. **Decode-ratio prioritization.** Requests with a higher fraction of decode tokens (relative to remaining prefill chunks) are scheduled earlier. This keeps the in-flight batch dominated by decode work and lets newly-arrived prefill chunks interleave without bubbles. -4. **Memory-centric token batching.** Token-batch sizes are sized to fit the available KV/activation memory rather than a fixed shape, increasing GPU utilization on workloads with variable per-request memory demand. -5. **Anti-premature-eviction.** Because the scheduler knows which pending requests will reuse a prefix, it can pin or refcount that prefix's blocks so an LRU policy does not evict them out from under upcoming requests. - -## Implementation sketch - -``` -1. Pending-queue prefix index - - Build a radix tree over the queue (separate from the cache's radix tree, or unified) - - Maintain it incrementally as requests arrive and depart - verify: queries return the set of pending requests sharing a given prefix - -2. Prefix-grouped scheduler policy - - New SchedulerInterface implementation: GroupByPrefixScheduler - - At each dispatch tick, prefer requests whose prefix is already resident or whose - prefix has the largest pending-group size - verify: synthetic workload with N requests sharing one prefix executes the prefix - exactly once; baseline runs prefill N times - -3. Decode-ratio prioritization - - Per-request, track remaining_decode_tokens / remaining_prefill_chunks - - Sort the dispatch candidate set by this ratio descending before applying - fairness / priority-class constraints - verify: under a mixed-length workload, average decode tokens per step increases - vs FIFO - -4. Pinning to prevent premature eviction - - Refcount prefix blocks by the number of pending requests that will use them - - Eviction policy must respect refcount > 0 - verify: under cache-pressure load, prefixes with pending consumers are not evicted - -5. Benchmark - - RAG-style workload (heavy prefix overlap) and agent-style workload (shared - system prompt, varied user turns) - - Compare engine throughput vs the engine's own baseline (no global prefix scheduling) - verify: throughput improvement on prefix-heavy workloads, no regression on prefix-light ones -``` - -## Notes on integration with other engine components - -- **Radix tree:** the queue-side prefix index can either be a second radix tree over pending requests or an extension of the cache-side tree with a "pending consumers" annotation per node. The latter is cleaner because eviction already needs to read it. -- **FairScheduler:** prefix grouping and per-session fairness can conflict — grouping all requests sharing a prefix together can starve a session whose requests are not in any group. Resolution lives at the policy layer: apply fairness as the outer constraint, prefix grouping as the tie-breaker. -- **Continuous batching:** decode-ratio prioritization is per-step and slots into the iteration-level admit decision without architectural change. diff --git a/docs/papers/elis.md b/docs/papers/elis.md deleted file mode 100644 index 19f5ee4..0000000 --- a/docs/papers/elis.md +++ /dev/null @@ -1,63 +0,0 @@ -# ELIS: Efficient LLM Iterative Scheduling System with Response Length Predictor - -**Link:** https://arxiv.org/abs/2505.09142 -**Venue:** arXiv preprint, May 2025 -**Authors:** Seungbeom Choi, Jeonghoe Goo, Eunjoo Jeon, Mingyu Yang, Minsung Jang -**Status:** not read -**Relevant phase:** Phase 6 (scheduler / fairness) - -## Core idea - -Iteration-level batching (vLLM-style) admits new requests at every generation step, but if it admits them in FIFO order, a long-output request can sit at the head of the queue and force every shorter request behind it to wait or share a degraded slot — classic head-of-line blocking. ELIS sidesteps this by **predicting the output length of each request before it runs** using a trained encoder model (BGE-based), then running an **Iterative Shortest-Remaining-Time-First (ISRTF)** scheduler that prefers requests with fewer predicted-remaining tokens. Implemented as a Kubernetes-native serving system, ELIS reports ~19.6% reduction in average job completion time on real-world traces. - -## Why it matters - -- The project's planned `FairScheduler` needs at least one concrete admission/preemption policy. ISRTF is a credible second policy alongside per-session fairness — they target different workloads (ISRTF wins on JCT for mixed-length traffic, fairness wins when one session would starve others). -- A length predictor is a small, well-defined component that plugs in behind a clean interface (`predict_remaining_tokens(request) → int`). Different predictors (BGE, regex over prompt patterns, no predictor) become swappable without touching the scheduler. -- The 19.6% JCT reduction is on top of an already iteration-batched baseline, so the win is purely from scheduling, not from any kernel or memory work — clean to attribute and measure. - -## Key mechanism - -1. **Response length predictor.** A BGE encoder model is fine-tuned to map a prompt (and optionally request metadata) to an expected output token count. The predictor runs once per request at admission time, not per step. Training data comes from logged prompt → actual output length pairs. -2. **ISRTF scheduling.** At each iteration, the scheduler ranks pending and in-flight requests by predicted remaining tokens (`predicted_total - generated_so_far`) and prefers smaller values. New short requests can preempt or out-prioritize long in-flight requests when the engine has admit/preempt slots available. -3. **Iteration-batching aware.** Unlike classical SRTF on whole jobs, ISRTF operates at the iteration boundary. Each step the scheduler reconsiders the running set; long requests are not killed mid-flight, but newly admitted short requests are placed ahead of pending long ones. -4. **Prediction-error tolerance.** Predictor errors are unavoidable. The "remaining" count is updated each step using `max(0, predicted_total - generated_so_far)`. If a request runs past its predicted total, its remaining hits zero and it loses priority on subsequent ties — long-tail requests degrade gracefully instead of catastrophically. -5. **Kubernetes-native deployment.** ELIS ships as a K8s scheduler component, meaning request routing across replicas is part of the system, not just per-engine scheduling. For this project, only the per-engine ISRTF logic is directly applicable; the multi-replica routing is platform-layer (deferred per the Forward-Compatibility Constraint in CLAUDE.md). - -## Implementation sketch - -``` -1. LengthPredictor abstract interface - - predict(request) -> int (predicted total output tokens) - - Implementations: NaivePredictor (constant), BGEPredictor (trained encoder), - HeuristicPredictor (regex/feature rules over prompt) - verify: NaivePredictor returns the configured constant for any input - -2. ISRTFScheduler implementation behind SchedulerInterface - - At each tick: rank candidates by (predicted_total - generated_so_far) ascending - - Apply fairness / priority-class as outer constraint, ISRTF as inner ordering - verify: synthetic workload of mixed short/long requests shows shorter ones complete - sooner under ISRTF than under FIFO - -3. BGE-based predictor (optional, behind the interface) - - Load a BGE encoder via HF Transformers in the startup hook - - Add a calibration script: run prompts → log actual output lengths → fine-tune - - Predict once at admission, cache the result on the request object - verify: predicted vs actual length correlation > 0.5 on a held-out set - -4. Prediction-error handling - - When generated_so_far >= predicted_total, set remaining = 0 (do not go negative) - - Optionally: emit a metric for prediction error so the predictor can be retrained - verify: a request that runs past its prediction does not crash or hang the scheduler - -5. Benchmark - - Mixed short/long workload from a representative trace - - Compare avg JCT, p95 JCT, and throughput under ISRTF vs FIFO vs FairScheduler - verify: ISRTF reduces avg JCT vs FIFO without catastrophic p99 regression on long requests -``` - -## Notes on integration with other engine components - -- **FairScheduler vs ISRTF:** these are alternative `SchedulerInterface` implementations, not layers. Configuration selects one. A composite scheduler (fair-across-sessions, ISRTF-within-session) is a future option but unnecessary for a first cut. -- **Continuous batching:** ISRTF assumes iteration-level scheduling — it is an admission-ordering policy that runs at every batch step. No changes to batching mechanics required. -- **Modal deployment:** the predictor model loads at startup (consistent with the single-startup-hook constraint in CLAUDE.md). Predictor inference per-request is cheap (encoder forward pass on a short prompt) and lives in-process; no external service. diff --git a/docs/papers/flashattention.md b/docs/papers/flashattention.md deleted file mode 100644 index a0b9fc1..0000000 --- a/docs/papers/flashattention.md +++ /dev/null @@ -1,65 +0,0 @@ -# FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness - -**Link:** https://arxiv.org/abs/2205.14135 -**Venue:** NeurIPS 2022 -**Authors:** Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, Christopher Ré (Stanford) -**Status:** not read -**Relevant phase:** Phase 6 (CUDA / GPU optimization) - -## Core idea - -Attention's `O(N²)` memory footprint comes from materializing the full `N × N` score matrix in HBM. On modern GPUs, attention is bottlenecked by HBM reads and writes, not by FLOPs — the `softmax(QK^T) V` computation is fundamentally an IO problem. FlashAttention computes the same exact attention output without ever materializing the full score matrix. It tiles Q, K, and V into blocks that fit in on-chip SRAM, computes attention block-by-block using an online softmax that updates running max and sum statistics, and writes only the final output back to HBM. The result is the same numerical answer as standard attention, but with `O(N)` HBM traffic instead of `O(N²)` and substantially higher wall-clock throughput, especially at long context lengths. - -## Why it matters - -- Every modern serving stack uses some FlashAttention variant under the hood. vLLM's PagedAttention is itself a paged-memory variant of FlashAttention. -- The IO-awareness lens — count HBM accesses, not FLOPs — is the right way to reason about *every* attention kernel optimization that follows (FlashAttention-2, FlashAttention-3, FlashDecoding, paged variants). Reading this paper is the prerequisite for understanding the rest. -- Even on a non-CUDA dev box, the algorithm informs how a custom paged-attention kernel for the engine should be structured: tile-based, online softmax, single-pass. PyTorch's stock attention will become the limiting factor on long contexts well before the engine reaches vLLM-class throughput. - -## Key mechanism - -1. **Tile Q, K, V.** Split Q into blocks of `Br` rows; split K and V into blocks of `Bc` rows. Block sizes are chosen so that `Br × d` and `Bc × d` blocks fit in on-chip SRAM together with running statistics. -2. **Outer loop over K, V blocks; inner loop over Q blocks.** For each `(K_j, V_j)` block loaded into SRAM, iterate over Q blocks. For each `Q_i`, compute `S_ij = Q_i K_j^T` in SRAM, never writing it to HBM. -3. **Online softmax.** Maintain per-Q-row running statistics: max-so-far `m_i`, sum-of-exp `l_i`, and partial output `O_i`. When a new score block arrives, rescale the existing accumulators by `exp(m_i_old - m_i_new)` and add the new block's contribution. This produces the exact softmax output incrementally without ever holding the full `S_ij` row. -4. **Single HBM write.** Each tile's contribution to `O` is accumulated in registers/SRAM. Only the final `O` is written back to HBM. The softmax-normalization statistics never leave SRAM. -5. **Backward pass uses recomputation.** During training, the `S` matrix is recomputed from Q, K, V on the backward pass instead of being stored, trading FLOPs (cheap) for HBM traffic (expensive). Inference does not need this since there is no backward pass. -6. **Causal masking is fused.** For autoregressive attention, blocks above the diagonal are skipped entirely; the causal mask never materializes. - -## Implementation sketch - -``` -1. Verify FlashAttention availability via PyTorch SDPA - - torch.nn.functional.scaled_dot_product_attention auto-selects FA on CUDA - - Confirm on the target CUDA backend that FA path is used (not the math fallback) - verify: profiler shows FA kernel; HBM traffic scales with N, not N² - -2. Replace HF model's attention path with SDPA-based attention - - Most recent HF models support attn_implementation="sdpa" or "flash_attention_2" - - Configurable via config flag ATTN_IMPL - verify: end-to-end outputs match the eager-attention baseline within tolerance; - per-step latency drops at long context - -3. Long-context benchmark - - TTFT and per-token latency vs context length, eager vs SDPA-FA - - Sweep context from 512 to model max - verify: per-token latency at long contexts drops noticeably; eager path becomes - limiting factor before SDPA path does - -4. (Custom paged-attention kernel work) Adopt FA-style tiling and online softmax - - When implementing the custom paged attention for the block-paged KV cache, - structure the kernel as a FA tile loop over (K, V) blocks indexed through the - block table, with online softmax accumulators per Q row - verify: matches non-paged FA output within tolerance on equivalent inputs - -5. (FlashAttention-2 / 3) Track the variant available on the deployment hardware - - FA-2 reduces non-matmul FLOPs and improves parallelism; FA-3 uses Hopper-specific - features - - Use whatever the installed PyTorch + CUDA version exposes through SDPA - verify: deployment notes record which variant is active per hardware target -``` - -## Notes on integration with other engine components - -- **Paged KV cache:** the engine's eventual custom paged-attention kernel should be a FlashAttention kernel that gathers K, V through the block table inside the tile loop rather than from a contiguous tensor. The two ideas compose; this composition is exactly what vLLM's PagedAttention is. -- **Loki:** Loki's gather-then-attend final stage is a small attention call over the selected top-k tokens. That call should itself use FA. Loki's win is on top of FA, not instead of it. -- **MPS / non-CUDA:** FlashAttention is CUDA-only in practice. On the MPS dev box, SDPA falls back to a math implementation. Performance claims belong to the CUDA path. diff --git a/docs/papers/loki.md b/docs/papers/loki.md deleted file mode 100644 index 6c02dd2..0000000 --- a/docs/papers/loki.md +++ /dev/null @@ -1,68 +0,0 @@ -# Loki: Low-Rank Keys for Efficient Sparse Attention - -**Link:** https://proceedings.neurips.cc/paper_files/paper/2024/file/1e027da6bec9ceb2ec37951ceeccae93-Paper-Conference.pdf -**Venue:** NeurIPS 2024 -**Authors:** Prajwal Singhania et al. (UMD) -**Status:** read -**Relevant phase:** Phase 6 (CUDA / GPU optimization), composes with eviction in Phase 5 - -## Core idea - -Key vectors in a transformer's attention live in a much lower-dimensional subspace than their nominal `head_dim` would suggest. Loki exploits this by computing **approximate** attention scores in a low-rank PCA subspace, using those scores to pick a small top-k of relevant tokens, and only running full-precision attention over that top-k. The KV cache itself is not compressed and not evicted — Loki sparsifies the *compute*, not the storage. - -## Why it matters - -- Orthogonal to eviction policies (LRU, Attention-Sink-LRU, H2O): eviction decides what to drop from the cache, Loki decides what subset of the kept tokens to attend over per step. They compose. -- Attacks per-step attention cost rather than memory footprint, complementing the paged KV cache rather than competing with it. -- The reported ~45% throughput win on Llama2-13B with a Triton kernel makes it a credible Phase 6 lever once the engine is on CUDA. -- Quality degradation (~6.8% reported) is small enough that it's a viable always-on optimization rather than a long-context-only trick. - -## Key mechanism - -1. **Offline calibration.** A small calibration corpus (paper used ~16 WikiText sequences) is run through the model. K vectors are collected per layer per KV-head. PCA is run independently on each (layer, kv-head) bucket. All principal components are stored, but only the top `d` (≈25-50% of `head_dim`) will be used at inference. -2. **Per-step approximate scoring.** At each decode step, the current Q and the cached K are projected into the top-`d` PCA subspace. Approximate attention scores are computed in that low-dim space — cheap because both the matmul and the memory read are `d / head_dim` smaller. -3. **Top-k selection.** From the approximate scores, the top-`k` tokens (≈12.5-25% of the sequence) are selected. Special tokens (attention sinks, beginning-of-sequence) are always retained regardless of score. -4. **Full-precision attention on the subset.** Full-dimensional K and V are gathered for the selected `k` tokens, and the final softmax-weighted attention is computed over that subset. -5. **Triton kernel.** The paper's reported speedup comes from a fused Triton kernel that does the projection, top-k, gather, and final attention without materializing intermediate tensors. Pure-PyTorch implementations are correct but do not realize the speedup. - -## Implementation sketch - -``` -1. scripts/calibrate_loki.py - - Hook the model's attention modules; collect K per layer per kv-head over the calibration corpus - - Run PCA per (layer, kv-head) bucket via torch.linalg.svd - - Save components as a single .pt artifact: tensor of shape (num_layers, num_kv_heads, head_dim, head_dim) - verify: artifact loads, shapes match the model config - -2. src/inference_server/attention/loki.py - - LokiAttention forward replacing the model's attention path - - Loads PCA components at backend startup (runs in the single startup hook, see CLAUDE.md Modal constraint) - - Config flags: LOKI_ENABLED, LOKI_D (rank), LOKI_K (top-k count or fraction), LOKI_PROTECT_SINKS - verify: with d=head_dim and k=seq_len, output exactly matches baseline attention - -3. Backend integration - - Monkey-patch the loaded HF model's attention modules at startup when LOKI_ENABLED=true - - No changes to the scheduler, cache manager, or server - verify: end-to-end generation still produces sensible outputs - -4. Quality benchmark - - Perplexity on a small held-out eval set with and without Loki across (d, k) sweep - - Compare absolute PPL delta to the paper's ~6.8% claim - verify: quality degradation in line with paper for at least one (d, k) configuration - -5. (CUDA phase) Triton kernel - - Fused projection + top-k + gather + attention - - Throughput benchmark vs PyTorch baseline on the same hardware - verify: per-step attention latency reduction matches paper's order of magnitude - -6. Compose with eviction - - Run Loki on top of LRU and on top of Attention-Sink-LRU - - Confirm sink protection is consistent between Loki's "always keep sinks" and the eviction policy - verify: no quality regression beyond the sum of the two techniques' individual deltas -``` - -## Notes on integration with other engine components - -- **Paged KV cache:** Loki reads K through the same indirection path as full attention; the gather step naturally handles non-contiguous K. No extra work. -- **H2O eviction:** H2O's bookkeeping uses attention scores. If Loki only computes scores on its top-k subset, H2O's signal is incomplete. Two options: (a) track approximate scores from Loki's low-rank pass for H2O's per-token accumulator (cheap, slightly biased), or (b) periodically run a full-attention pass to refresh H2O's scores (expensive, accurate). Decision deferred until both are implemented. -- **Continuous batching:** Loki is per-request and per-layer; nothing in the scheduler needs to change. diff --git a/docs/papers/speculative-decoding.md b/docs/papers/speculative-decoding.md deleted file mode 100644 index 6856a6b..0000000 --- a/docs/papers/speculative-decoding.md +++ /dev/null @@ -1,79 +0,0 @@ -# Fast Inference from Transformers via Speculative Decoding - -**Link:** https://arxiv.org/abs/2211.17192 -**Venue:** ICML 2023 -**Authors:** Yaniv Leviathan, Matan Kalman, Yossi Matias (Google Research) -**Status:** not read -**Relevant phase:** Phase 6 / Optional Extensions (advanced optimization) - -## Core idea - -Autoregressive decoding generates one token per forward pass through the large target model, and each forward pass is bottlenecked by memory bandwidth, not by FLOPs — the GPU is mostly idle waiting on weight reads. Speculative decoding exploits the slack: a small, cheap **draft model** generates `K` candidate tokens in `K` cheap forward passes, and the large target model then verifies all `K` candidates in **one** forward pass (which is nearly the same cost as generating one token, because the bottleneck is per-call HBM traffic, not per-token compute). A rejection-sampling check guarantees the output distribution is **identical** to plain target-model sampling — no quality loss. When the draft model is right most of the time, every target forward pass yields multiple accepted tokens, multiplying decoding throughput by 2-3× without changing the output. - -## Why it matters - -- One of the cleanest "free lunches" in inference: speedup with provably no change in output distribution. No quality/throughput tradeoff to argue about. -- Plays well with everything else in the engine — it does not change the KV cache, the scheduler, or the attention kernel. It changes the decode loop. -- Practical wins depend on the draft-target acceptance rate, which depends on model pairing. With Gemma 4 E4B as target, a smaller distilled Gemma or a small generic draft model would be the natural pairing. -- The technique generalizes (Medusa, EAGLE, lookahead decoding, n-gram drafters). Implementing the basic version first creates the seam those variants slot into. - -## Key mechanism - -1. **Two models, one tokenizer.** A large **target** model `Mq` (the one whose distribution is desired) and a small fast **draft** model `Mp`. Both must share a tokenizer so token IDs are interchangeable. -2. **Draft phase.** Starting from the current prefix, the draft model autoregressively samples `K` candidate tokens `x_1 ... x_K` along with their probabilities `p(x_i | prefix, x_1 ... x_{i-1})`. This costs `K` cheap forward passes through `Mp`. -3. **Target verification phase.** The target model is called **once** on the prefix plus all `K` draft tokens, producing target probabilities `q(x_i | prefix, x_1 ... x_{i-1})` for every position in parallel. -4. **Rejection sampling per token.** For each `i` from 1 to `K`: accept `x_i` with probability `min(1, q(x_i) / p(x_i))`. If accepted, move on. If rejected at position `j`, throw away `x_j ... x_K`, sample a replacement token from the **adjusted distribution** `max(0, q - p)` normalized, and stop. The math guarantees the resulting token sequence has exactly the same distribution as sampling from `q` directly. -5. **Bonus token.** If all `K` drafts are accepted, the target model's parallel forward pass already produced logits for position `K+1`, so one extra "free" token can be sampled there. Best case: `K+1` tokens per target forward pass. -6. **Why it speeds up.** A target forward pass with `K+1` tokens of input costs only marginally more than one with 1 token (memory-bound regime), but produces `1` to `K+1` accepted output tokens. Average tokens-per-target-call ends up well above 1 for well-matched draft/target pairs. - -## Implementation sketch - -``` -1. SpeculativeDecoder behind a clean interface - - Wrap the existing single-model decode path; choose between standard and speculative - via config flag SPECULATIVE_ENABLED - - Config: DRAFT_MODEL_NAME, SPEC_K (draft length per round) - verify: with SPECULATIVE_ENABLED=false the path is bit-identical to today's decode - -2. Draft model loading - - Load draft model in the same startup hook as target (consistent with the single-startup - constraint in CLAUDE.md) - - Validate tokenizer compatibility (vocab IDs must match) at startup; fail fast if not - verify: draft and target tokenizers produce identical token IDs for a sanity prompt - -3. Draft + verify loop - - Per decode iteration: K cheap draft steps → one target forward pass with K tokens - of input → rejection sampling - - Accepted tokens stream out of the engine in order; on rejection, the bonus replacement - token streams next; loop continues from the new prefix - verify: output is statistically indistinguishable from non-speculative sampling on a - held-out prompt set (same temperature, same seed where possible) - -4. KV cache handling - - Both models maintain their own KV caches - - On rejection at position j, both caches must be rolled back to position j-1 - (drop entries appended for x_j ... x_K) - - This is straightforward with the paged cache: free the affected blocks - verify: cache state after rejection matches what plain decode from the accepted prefix - would produce - -5. Acceptance-rate metric - - Track per-request and engine-wide average accepted-tokens-per-target-call - - Expose on /metrics - verify: metric is positive and stable; compare against published numbers for similar - draft/target pairs - -6. Benchmark - - End-to-end TPOT and tokens/sec with and without speculation, same workload - verify: speedup proportional to acceptance rate; no change in output quality - -7. (Future) Tree-style speculation (Medusa / EAGLE) reuses the draft+verify scaffold; - only the draft proposal mechanism changes -``` - -## Notes on integration with other engine components - -- **Paged KV cache:** rejection-driven rollback is a block-table truncation, not a tensor copy. The block-paged design makes speculative decoding *easier* to implement, not harder. -- **Continuous batching:** speculative decoding changes the per-request decode unit from "1 token per step" to "1-(K+1) tokens per step" with variance. The scheduler can stay as-is, but per-step throughput accounting needs to handle variable token counts cleanly. -- **Streaming:** accepted tokens stream as they are produced in order. The user-facing SSE stream sees no difference from non-speculative output. -- **Quantization / Loki / eviction:** all of these apply to the target model independently. Speculative decoding does not interact with them. diff --git a/docs/papers/tokenweave.md b/docs/papers/tokenweave.md deleted file mode 100644 index fd626a0..0000000 --- a/docs/papers/tokenweave.md +++ /dev/null @@ -1,64 +0,0 @@ -# TokenWeave: Efficient Compute-Communication Overlap for Distributed LLM Inference - -**Link:** https://arxiv.org/abs/2505.11329 -**Venue:** arXiv preprint, May 2025 (v4 Oct 2025) -**Authors:** Raja Gond, Nipun Kwatra, Ramachandran Ramjee -**Status:** not read -**Relevant phase:** future / multi-GPU extension (out of scope for the current roadmap) - -## Core idea - -In tensor-parallel inference, every transformer layer ends with an `AllReduce` across GPUs to sum partial outputs. At low batch sizes — the regime that matters for low-latency serving — this collective sits on the critical path and cannot be hidden behind compute, because there is too little compute to overlap with. TokenWeave splits the tokens of each in-flight batch into two roughly equal subsets and pipelines them: while subset A is doing its `AllReduce`, subset B is doing its compute, and vice versa. The communication and the compute then overlap by construction, even at small batch sizes. A fused `AllReduce + RMSNorm` kernel using NVSHARP/Multimem (Hopper/Blackwell) does the comm step on 2-8 SMs, leaving the rest of the GPU free for the other subset's compute. Reported wins: 1.29× latency, 1.26× throughput. - -## Why it matters - -- Tensor parallelism is in the optional-extensions section of the roadmap, not a near-term priority. This paper is reference material for that future work, not something to implement now. -- It is the right paper to read *before* designing a TP path, because the standard "do compute, then AllReduce, then next layer" pipeline is what this work argues against. Building TP naively and then bolting on overlap later means rewriting the model code. -- The fused `AllReduce + RMSNorm` kernel is a hardware-specific optimization (Hopper/Blackwell with NVSHARP) — it is the speedup mechanism, but the *architectural* contribution (split-batch overlap) applies to any multi-GPU setup. -- Notable that vLLM, SGLang, and TensorRT-LLM all leave overlap off by default in TP serving — suggests the integration cost has historically been higher than the benefit, and that real wins require the kernel-level work the paper does. - -## Key mechanism - -1. **Token-subset splitting.** Each in-flight batch is partitioned into two subsets of roughly equal token count. The subsets are processed in an interleaved schedule rather than as one monolithic batch. -2. **Two-phase pipeline per layer.** Within one transformer layer, subset A's compute (Q/K/V projection, attention, FFN partials) runs while subset B's `AllReduce` is in flight, then subset B's compute runs while subset A's `AllReduce` is in flight. The naive pipeline (compute-all → comm-all) is replaced by an interleaved one. -3. **Fused AllReduce + RMSNorm kernel.** The next layer normalizes its inputs with RMSNorm immediately after `AllReduce`. Fusing these two operations avoids materializing the post-AllReduce tensor in HBM and reduces the comm step's footprint to 2-8 SMs via NVSHARP/Multimem (Hopper and Blackwell hardware features for in-network reduction). -4. **Low SM occupancy for comm.** Because the fused kernel uses only 2-8 SMs, the remaining ~100+ SMs on the GPU are free to run the other subset's compute. This is what makes overlap practical at small batch sizes — naive `AllReduce` saturates many more SMs and leaves no compute headroom. -5. **Code-level integration.** Each linear layer that ends in `AllReduce` is replaced with the fused variant, and the per-layer execution loop is rewritten to issue subset-A and subset-B operations on alternating streams. - -## Implementation sketch - -``` -0. Prerequisite: tensor-parallel inference path - - Not in the current roadmap; would require splitting model weights across GPUs, - adding NCCL/NVSHARP communicators, and a TP-aware backend - verify: a baseline TP path runs the model with naive compute → AllReduce → next-layer - -1. Token-subset splitter - - Given an in-flight batch, partition tokens into two subsets balanced by count - (and ideally by sequence boundary so attention masks stay contiguous) - verify: round-trip a batch through split/merge and recover identical tensors - -2. Two-stream pipeline executor - - Per layer: alternate compute(subset_A) || allreduce(subset_B) on two CUDA streams - - Synchronize at layer boundaries - verify: with overlap disabled, output bit-identical to baseline; with overlap enabled, - output matches within numerical tolerance - -3. Fused AllReduce + RMSNorm kernel - - Hopper/Blackwell only, NVSHARP/Multimem path - - Falls back to non-fused on older hardware (correctness preserved, no speedup) - verify: kernel output matches naive (allreduce → rmsnorm) within tolerance - -4. Benchmark - - TP=2, TP=4, TP=8 across batch sizes from low (latency regime) to high (throughput regime) - - Compare end-to-end latency and throughput vs the same TP path with overlap off - verify: latency improvement at low batch sizes consistent with paper's order of magnitude - -5. Defer until multi-GPU is on the project roadmap -``` - -## Notes on integration with other engine components - -- **Out of scope for the current engine.** This project is single-GPU through Phase 11. TokenWeave only matters once a TP path exists. -- **Architectural seam:** if and when TP is added, the per-layer execution loop in the backend must already be structured so that the comm and compute steps are separable operations — otherwise integrating overlap means rewriting the layer loop. Worth noting now even though no code is written. -- **Modal deployment:** Modal supports multi-GPU containers, so TokenWeave-style overlap is reachable from the planned deployment target without changing the deployment model. diff --git a/docs/papers/vllm-paged-attention.md b/docs/papers/vllm-paged-attention.md deleted file mode 100644 index e691859..0000000 --- a/docs/papers/vllm-paged-attention.md +++ /dev/null @@ -1,67 +0,0 @@ -# Efficient Memory Management for Large Language Model Serving with PagedAttention (vLLM) - -**Link:** https://arxiv.org/abs/2309.06180 -**Venue:** SOSP 2023 -**Authors:** Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph E. Gonzalez, Hao Zhang, Ion Stoica (UC Berkeley) -**Status:** read -**Relevant phase:** Phase 5 (KV cache integration) + Phase 6 (continuous batching, vLLM head-to-head) - -## Core idea - -Naive LLM serving wastes KV cache memory three ways: internal fragmentation (over-reserving for `max_tokens` that never generate), reservation (slots held for future tokens of an active sequence), and external fragmentation (gaps between sequences with different lifetimes). vLLM borrows the operating-system virtual-memory model: split the KV cache into fixed-size **blocks** (e.g., 16 tokens of K and V per block), give each sequence a **block table** mapping logical token positions to physical block IDs, and allocate blocks on demand. Blocks need not be contiguous, can be freed independently, and can be **shared** across sequences when prefixes match. A custom **PagedAttention** kernel reads K and V through the block-table indirection so attention runs directly on non-contiguous memory. - -## Why it matters - -- This is the architectural foundation Phase 5 is building toward. The planned `BlockManager`, `RadixTree`, and `CacheManager` are direct analogs of vLLM's block allocator, prefix-sharing index, and `BlockSpaceManager`. -- vLLM is the explicit benchmarking target for this project. Understanding the paper deeply is what justifies specific throughput and memory-utilization expectations rather than treating vLLM as an opaque baseline. -- Block-paging is what makes continuous batching practical at scale: without it, packing variable-length sequences into one batch wastes too much memory to be worth the compute gain. - -## Key mechanism - -1. **Paged KV cache.** Cache memory is a pool of fixed-size blocks. Each sequence holds a block table — an array of physical block IDs in logical-position order. New blocks are allocated only when the sequence's last block fills, eliminating internal fragmentation beyond at most one partially-filled block per sequence. -2. **PagedAttention kernel.** A custom CUDA kernel takes the block table plus the global block pool and computes attention by gathering through the indirection. Naive PyTorch indexing for this gather would dominate runtime; the kernel fuses the gather into the attention computation. This kernel is the part that does not exist as an off-the-shelf PyTorch primitive. -3. **Copy-on-write block sharing.** Multiple sequences can point to the same physical block — e.g., shared system prompts, beam-search siblings, parallel sampling from the same prompt. Each block carries a reference count. When a sequence needs to write into a shared block, the block is copied first and the writing sequence's block table is updated. -4. **Iteration-level scheduling (continuous batching).** At every generation step, the scheduler can: admit waiting sequences if memory allows, evict (swap-out or recompute) lower-priority sequences when memory is tight, and resume previously evicted sequences. Because admission and eviction operate on block IDs rather than tensor copies, these decisions are cheap. -5. **Swap vs recompute.** When the cache is full and a new sequence must run, vLLM either swaps blocks to CPU memory or drops them and rebuilds via prefill on resume. Swap wins for long prefixes when PCIe bandwidth is high; recompute wins for short prefixes and is simpler. -6. **Memory accounting.** Total concurrent capacity is `total_blocks * block_size / avg_seq_len`, not `max_concurrent * max_seq_len`. This is the single biggest reason vLLM serves more concurrent users on the same hardware than naive serving stacks. - -## Implementation sketch - -``` -0. Block manager, radix tree, cache manager (Phase 5 stages 1-3) - - Bookkeeping data structures only; not yet integrated with generation - verify: unit tests already in place - -1. Wire block-paged KV into the generation loop (Phase 5 stage 7) - - Replace HF's contiguous KV with the block pool - - Pure-PyTorch "gather K/V from block table → attend" path for correctness - - On MPS this will be slower than baseline; that is acceptable - verify: end-to-end outputs match baseline within numerical tolerance - -2. Reference counting + copy-on-write - - Hook into the radix tree on prefix-match: increment refcount on shared blocks - - On a write to a refcount > 1 block, copy and update the writing sequence's block table - verify: two sessions with a shared system prompt actually share blocks; mutation correctly forks - -3. Continuous batching scheduler (Phase 6) - - Iteration-level admit / preempt / resume - - Backpressure policy: queue when full, choose swap-vs-recompute on preemption - verify: load test with mixed-length requests does not OOM; batch utilization stays high - -4. (CUDA phase) Real PagedAttention kernel - - Triton port, or integrate FlashInfer / xformers paged variant - - Drops in behind the existing block-table abstraction; no scheduler changes required - verify: per-step attention latency competitive with vLLM on equivalent hardware - -5. Head-to-head vs vLLM (Phase 11) - - Same model, same hardware, same workload mix - - Metrics: throughput (req/s, tok/s), p50/p95/p99 TTFT and TPOT under concurrency, - KV cache utilization, memory footprint, saturation point - verify: numbers documented in BENCHMARKS.md alongside vLLM's -``` - -## Notes on integration with other engine components - -- **Eviction policies (LRU / H2O):** operate at the block level inside the block pool. The block manager is the right home for the eviction interface; no changes to scheduler or attention kernel required. -- **Loki sparse attention:** reads K and V through the same block-table indirection. A paged backend is strictly more flexible for Loki than a contiguous one because Loki's gather is already indirection-based. -- **Modal deployment:** the block pool is in-process and per-container, consistent with the constraint in CLAUDE.md. Multi-replica scaling would require sticky session routing on `session_id` rather than any change to the block layer.