Skip to main content

Work / KV-cache LLM serving

KV-cache LLM serving

Graduate capstone on serving latency: I owned the KV-cache layer under a learned vLLM scheduler, where 4-bit KV fit 2.83 times the tokens and cut time to first token 6.8 times under load.

token capacity
2.83×
faster first token
6.8×
P99 time per token
−47%
Role
KV-cache layer owner, team of four
When
2026
  • Python
  • vLLM
  • Triton
  • PyTorch
  • Llama-3.1-8B
  • RTX 3090

A toy model of KV memory, not the measured system. Orange pages are a preemption.

Running
0
Waiting
0
Preempted
0

Measured on Llama-3.1-8B, one RTX 3090, 64 requests a second: mean time to first token 10.7 s → 1.6 s, preemptions 171 → 40.

The question

An LLM server keeps each request's attention keys and values, the KV cache, in GPU memory while it generates. Under load that memory runs out first. vLLM then preempts a request and has to swap or recompute its cache later, and the time to first token climbs for everyone queued behind it.

Our FDU capstone started from a learning-to-rank scheduler, which orders requests by predicted output length to get close to shortest-job-first. The question was whether a KV-cache layer underneath it could cut latency further without hurting accuracy, and in which load range each technique pays off. I owned that layer.

Where my work sat

Quantization makes each page of KV smaller, so the same memory holds more concurrent requests and fewer of them are preempted. Sparsity skips attention blocks that contribute little, so decode reads less KV. Both went into the same vLLM build, measured on three engine versions (0.4.1, 0.8.5 and 0.25) with accuracy checked on GSM8K and perplexity.

RequestsLMSYS-Chat-1M prompts, Poisson arrivalsLTR schedulera small ranker predicts output length; shortest firstKV cacheTurboQuant 4-bit pages: 2.83× the tokens in the same memorymy partDecode attentionBLASST block-sparse Triton kernels, patched into vLLMmy partGPUone RTX 3090, 24 GB, Llama-3.1-8B

What it measured

  1. 2.83×token capacity

    4-bit KV at equal memory

    TurboQuant rotates before it quantizes, so 4-bit pages keep accuracy while holding 2.83 times the tokens.

  2. 10.7 s → 1.6 smean time to first token

    6.8× faster under load

    At 64 requests a second, mean TTFT fell 6.8 times and preemptions fell from 171 to 40.

  3. 1.33–1.75×standalone kernel speed-up

    Block-sparse decode kernels in Triton

    I wrote BLASST attention decode kernels that skip low-weight blocks using the online softmax. They were lossless at 42% sparsity.

  4. −22% / −47%time per token / its P99

    Patched into the quantized path

    Runtime-patching vLLM's TurboQuant decode path with the sparse kernel cut time per output token by 22% and its P99 by 47%.

What did not work

Quantization is not free. Time per output token rose in the configuration we wrote up, and end-to-end latency got worse, so 4-bit KV is the right call when time to first token is what users feel. The same sparse kernel was a net loss on vLLM's shared GQA kernel: of the 42.1% of blocks that could be skipped, only 18.5% actually were. It paid off only on the per-head quantized kernel.

Two lessons I kept: read version metadata from the server, not the client, and run the reproduction commands in a clean environment before calling results final. The full story is in a post on the capstone.

Stack

Serving
  • vLLM 0.4.1, 0.8.5, 0.25
  • PagedAttention
  • TurboQuant
  • Llama-3.1-8B
Kernels
  • Triton
  • PyTorch
  • CUDA on WSL2
Evaluation
  • LMSYS-Chat-1M
  • Poisson load
  • GSM8K
  • perplexity

What's next

The original plan had three more pieces, left as future work: offloading KV by attention head, a lock-free Rust coordinator over vLLM's KV-connector API, and speculative decoding. Locally I run quantized 27B models on the same 3090; notes are in 27B on one 3090.