A decode-latency planner for offloaded KV cache. It tells you when double-buffered prefetch hides the PCIe transfer of host-resident KV cache, and when you are simply bandwidth-bound.
Serving a long-context model on a GPU that cannot hold the whole KV cache is common: you keep the cache in host memory and stream each layer's K and V back over PCIe when that layer's attention runs. The naive policy fetches the KV for layer L, waits for it to land, then computes layer L. That wait is pure stall, and it repeats for every layer of every decoded token.
Everyone knows prefetching helps. What is rarely made explicit is the arithmetic of when it helps and by how much. Below some context length the per-layer transfer is smaller than the per-layer compute, so a prefetcher hides the transfer completely. Above it the transfer dominates and you are pinned to the PCIe link no matter what you do. Capacity planning needs that crossover point, not a vague "prefetching is good."
Given a model config (layers, hidden size, attention heads, grouped-query kv heads) and a hardware profile (HBM bandwidth, PCIe bandwidth), kvprefetch computes single-token decode latency with the KV cache offloaded to host memory, under two policies:
- on-demand: fetch each layer's KV just in time, then compute. No overlap.
- prefetch: double-buffered, layer-ahead transfer overlapped with compute.
It reports the decode step latency for each policy, the latency reduction, and the context-length crossover where the offload path stops being compute-bound and becomes PCIe-bound. It sweeps context length and writes a comparison chart.
The model is a two-term roofline. Single-token decode is a stack of GEMVs, so its compute cost is the time to stream a layer's weights from HBM once (weight_bytes / HBM_bandwidth), independent of context length. The transfer cost is the time to pull that layer's K and V back over PCIe (kv_bytes / PCIe_bandwidth), which grows with context length and batch and shrinks with grouped-query attention. On-demand pays their sum per layer. Prefetch pays a pipeline makespan of transfer + (layers - 1) * max(compute, transfer) + compute, which collapses to max(compute, transfer) per layer in steady state.
The result is not "prefetch is faster." It is that the benefit is unimodal in context length. It rises while the decode is compute-bound (more transfer to hide means more to gain), peaks exactly at the crossover where per-layer compute equals per-layer transfer, then falls as you become PCIe-bound and prefetch can only recover the overlapped compute. The planner locates that peak for your model and card, so you can decide whether offload plus prefetch is worth it before writing any transfer code. It also makes clear that look-ahead depth beyond one layer buys nothing in steady-state throughput; it only sets how much GPU memory the resident KV blocks cost, so depth is a memory knob, not a speed knob.
Diagram: https://www.figma.com/board/0WtgxEVAxmoraiPBxyU5js
The FigJam board shows how a model config and a hardware profile feed per-layer compute and transfer times, how the on-demand and prefetch schedulers turn those into a decode step latency, and how the crossover context and the sweep produce metrics.json and the comparison chart.
The metric is single-token decode step latency in milliseconds for Llama-3-8B on an A100-80GB PCIe card (effective HBM 2039 GB/s, effective Gen4 x16 PCIe 25 GB/s), batch 1, with the KV cache offloaded to host memory. The baseline is on-demand offload, which fetches each layer's KV then computes. The change is double-buffered prefetch, which overlaps layer L+1's transfer with layer L's compute.
At a 2048-token context the decode step drops from 17.58 ms to 10.95 ms, a 37.7 percent reduction. The reduction is largest near the crossover at 1306 tokens: at 1024 tokens it reaches 42.6 percent. Past the crossover the offload path is PCIe-bound and the win shrinks steadily, to 13.3 percent at 8k and 3.7 percent at 32k, because both policies are then limited by the same link and prefetch can only recover the overlapped compute. All numbers are produced by running the model in examples/demo.py; nothing is hand-entered.
pip install -r requirements.txt
python examples/demo.py # runs the sweep, writes docs/metrics.json
python examples/make_chart.py # renders docs/before_after.png from that json
python tests/test_scheduler.py # correctness checks on the latency model
Point it at a different model or card in a few lines:
from kvprefetch import compare, crossover_ctx, LLAMA2_13B, RTX4090
print(crossover_ctx(LLAMA2_13B, RTX4090)) # context where transfer == compute
c = compare(LLAMA2_13B, RTX4090, ctx_len=4096, batch=1)
print(c.ondemand.total_ms, c.prefetch.total_ms, c.reduction_pct)Built-in models: LLAMA3_8B, LLAMA2_13B, MISTRAL_7B. Built-in hardware: A100_PCIE, A100_PCIE_GEN3, RTX4090. Both are plain dataclasses, so a custom config is one constructor call.
Model: Llama-3-8B (6.979 B params)
Hardware: A100-80GB PCIe (Gen4 x16)
Batch: 1 Prefetch depth: 1
Crossover: ctx = 1306 tokens (compute-bound below, PCIe-bound above)
ctx on-demand prefetch reduction speedup bound
-------------------------------------------------------------
512 9.53ms 6.93ms 27.3% 1.38x compute
1024 12.21ms 7.01ms 42.6% 1.74x compute
2048 17.58ms 10.95ms 37.7% 1.61x pcie
4096 28.32ms 21.69ms 23.4% 1.31x pcie
8192 49.80ms 43.16ms 13.3% 1.15x pcie
16384 92.75ms 86.11ms 7.2% 1.08x pcie
32768 178.64ms 172.01ms 3.7% 1.04x pcie
graph TD
A["Model + hardware"] --> B["compute = weight_bytes / HBM"]
A --> C["transfer = kv_bytes(ctx, batch) / PCIe"]
B --> D{"transfer > compute ?"}
C --> D
D -->|"no, compute-bound"| E["prefetch hides transfer fully: latency = transfer + layers * compute"]
D -->|"yes, PCIe-bound"| F["prefetch pinned to link: latency = layers * transfer + compute"]
E --> G["reduction vs on-demand, peaks at crossover"]
F --> G
- Single-token decode is treated as HBM-bound on weight reads, which is the standard first-order model for autoregressive decode. Prefill and large-batch compute-bound regimes are out of scope; this is about the offload path during decode.
- PCIe bandwidth is the effective sustained one-direction figure after pinning and protocol overhead, not the nominal slot rate.
- The transfer assumes the full per-layer KV must cross the link each step. Keeping a resident recent-token window on the GPU reduces the transferred bytes; that is a straightforward extension of
kv_bytes_per_layer.