🤖 machine learning

PersistentKV: Page-Aware Decode Scheduling for Long-Context LLM Serving on Commodity GPUs

PersistentKV introduces a native block-table decode attention engine and adaptive, page-aware scheduling policy that optimizes long-context LLM serving on commodity GPUs by dynamically selecting between FlashInfer and specialized workqueue strategies based on batch size and workload characteristics, achieving significant throughput improvements over existing single-kernel approaches.

Original authors: Muhammad Ahmed

Published 2026-06-26

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Muhammad Ahmed

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are running a massive library where a single librarian (the AI) is trying to answer questions for many different people (the users) at the same time. To do this, the librarian has to keep a giant, constantly growing notebook of facts (the "KV cache") for every conversation.

The problem is that in modern libraries, these notebooks are huge. The librarian spends more time flipping pages and walking to the shelves to find the right notes than actually writing the answers. This is the "memory traffic" problem that slows down AI.

PersistentKV is a new way of organizing the librarian's workflow to make them faster, specifically on standard, off-the-shelf computers (like a gaming laptop) rather than super-expensive data center machines.

Here is the breakdown using simple analogies:

1. The Problem: The "One-Size-Fits-All" Mistake

Currently, most AI systems use a very efficient method called FlashInfer. Think of FlashInfer as a highly trained librarian who is amazing at handling a crowd of people who all have short, simple questions. They can process a whole group at once very quickly.

However, this method struggles when:

The crowd is small, but the questions are huge: If only one person is asking a very long, complex question (a "long-context" query), the librarian is underutilized. They are waiting for the next person to arrive, wasting time.
The crowd is mixed: If you have a mix of people asking short questions and people asking massive, long questions, the system tries to force everyone into the same "batch." This is like forcing a person with a 1-page essay to wait in line with someone writing a 100-page novel, or worse, padding the short essay with blank pages so it looks like a novel. This creates wasted effort.

2. The Solution: The "Smart Splitting" Strategy (PersistentKV)

The authors built a new system called PersistentKV. Instead of forcing everyone into one big group, this system acts like a smart manager who looks at the specific needs of each person and breaks the work down differently.

The "Split" Analogy: Imagine a long novel that needs to be read. Instead of one person reading the whole thing at once, the manager cuts the book into 32 smaller chapters. They assign different parts of the book to different assistants to read simultaneously.
- Why this helps: If you only have one person asking a long question, this "splitting" keeps the librarian's team busy by having them work on different chapters of that one long story at the same time. This fills up the "empty seats" in the computer's brain.
The "Work Queue" Analogy: In the old system, if you had 8 people with different length stories, the system might try to start 16 different tiny tasks (one for every length), which is chaotic and slow.
- PersistentKV's fix: It uses a "compact work queue." It looks at the 8 people, sees exactly who needs what, and creates a single, efficient list of tasks. It only sends work to the assistants that actually needs to be done, skipping the empty pages.

3. The "Adaptive Policy": The Smart Manager

The most important part of this paper isn't just the new tool; it's the decision-making rule. The authors realized that the "Splitting" strategy isn't always better.

Scenario A (Small Group, Long Story): If you have just 1 person with a long story, the new "Splitting" method is a winner. It speeds things up by 1.4x.
Scenario B (Medium Group, Mixed Stories): If you have 8 people with mixed story lengths, the "Compact Work Queue" is a winner. It speeds things up by about 1.2x.
Scenario C (The "Goldilocks" Zone - 4 People): If you have 4 people, the new method actually gets slower because the overhead of splitting and merging the work takes too much time.
- The Fix: The system is smart enough to say, "Hey, for 4 people, let's just use the old, reliable FlashInfer method." It switches tools automatically based on the situation.

4. The Results: What Actually Happened?

The researchers tested this on a standard RTX 3060 graphics card (a common consumer GPU, not a supercomputer).

Accuracy: The answers were just as correct as the standard method (within a tiny margin of error).
Speed:
- For single, very long conversations, they were 40% faster.
- For groups of 8 people with mixed conversation lengths, they were 6% to 26% faster.
- For groups of 4, they didn't try the new method; they stuck to the old one to avoid slowing down.

The Bottom Line

This paper doesn't claim that their new method is the "best" for every single situation. Instead, it proves that how you schedule the work is just as important as the math itself.

By treating the AI like a flexible manager who knows when to split a big job into pieces and when to stick to the old routine, they can make standard computers run long, complex AI conversations significantly faster. It's about finding the right tool for the specific size of the crowd, rather than using the same hammer for every nail.

Technical Summary: PersistentKV

Problem Statement

Autoregressive Large Language Model (LLM) serving is increasingly bottlenecked by Key-Value (KV) cache movement rather than dense matrix multiplication. While modern paged-attention systems (e.g., vLLM, FlashInfer) have successfully reduced KV-cache fragmentation and optimized native-paged decode kernels, a single "best" kernel does not exist for all serving scenarios.

Specifically, a tension exists between:

Low-active long-context decode: Where a single optimized kernel may under-utilize commodity GPUs due to insufficient parallel work.
Mixed sequence lengths: Where scheduling many exact-length launches creates high overhead, while coarse padded batches introduce inefficiency.

The paper investigates whether a native block-table decode engine, combined with page-aware scheduling, can improve serving latency on commodity GPUs in regimes where production-quality paged attention kernels are already strong.

Methodology

The authors propose PersistentKV, a system comprising a native block-table decode engine and a page-aware scheduling policy.

1. Native Block-Table Decode Engine

PersistentKV is designed for Grouped-Query Attention (GQA), where $G = H_q/H_{kv}$ query heads share one KV head.

Layout: It executes directly over native paged KV-cache layouts (block tables) without repacking K,V pages into contiguous tensors.
Sequence Splitting: To expose sufficient parallelism for low-active workloads, the sequence dimension is split into $S$ disjoint ranges. Each split computes a partial online-softmax state. A separate merge kernel combines these states exactly.
Row-Local Bounds: Unlike bucketed approaches that use padded lengths, PersistentKV uses true per-row sequence lengths ( $N^*_i$ ) to determine tile counts and split ranges. Empty splits write neutral states and exit early, avoiding padded KV reads.
Compact Workqueue: To handle ragged batches (requests with varying lengths) without excessive kernel launches, the scheduler emits a 1D grid of tasks only for non-empty row–KV-head–split intervals. This avoids the "launch fan-out" of exact-length bucketing while preserving sequence parallelism.

2. Calibrated Adaptive Policy

The system employs a manually calibrated policy to select the execution route based on batch size ( $B$ ) and maximum context length ( $N_{max}$ ):

$B=1$ (Long Context): Uses PersistentKV with length buckets and sequence splitting to expose enough Compute Thread Arrays (CTAs).
$B \ge 8$ (Long Context): Uses PersistentKV with the compact workqueue to amortize the two-kernel (split + merge) overhead.
$B=4$ (Boundary): Routes to FlashInfer. At this batch size, the overhead of splitting and merging in PersistentKV outweighs the benefits, causing a regression in wall-clock time.
Otherwise: Uses FlashInfer.

3. Evaluation Protocol

Hardware: NVIDIA RTX 3060 (28 SMs, 12GB VRAM).
Baselines: FlashInfer 0.2.5 (primary), vLLM, TensorRT-LLM.
Traces: Synthetic traces (Bucketed, Bimodal, Uniform, Zipf) designed to isolate context length, active batch size, and length raggedness.
Metrics: Correctness (FP16 tolerance against FlashInfer), CUDA-event timing (queued GPU work), and synchronized wall timing (end-to-end latency including host planning).

Key Contributions

Native Block-Table GQA Engine: A decode engine that exposes sequence-parallel work for low-active long-context decode without repacking KV caches.
Page-Aware Scheduling: Implementation of row-local sequence-length scheduling and a compact workqueue that aligns split boundaries and prefetches with true request lengths, avoiding rectangular empty grids.
Adaptive Methodology: A rigorous evaluation framework using calibration/held-out trace seeds, correctness checks, and scheduler-derived workload counters to demonstrate that work assignment is a decisive variable alongside attention math.

Results

On the held-out seed (20260623) with fixed thresholds derived from a calibration seed (20260622):

Single Request ( $B=1$ ): PersistentKV with length buckets achieves a 1.399× improvement in synchronized wall throughput over FlashInfer on bucketed long-context traces.
High Batch ( $B=8$ ):
- Bimodal: 1.265× improvement.
- Uniform: 1.082× improvement.
- Zipf-like: 1.063× improvement.
Boundary Case ( $B=4$ ): The adaptive policy correctly selects FlashInfer, avoiding a performance regression where PersistentKV would otherwise underperform.
Scheduling Overhead: The compact workqueue reduces launch fan-out from 16.00 launches/step (exact buckets) to 2.00 launches/step on bimodal $B=8$ traces, while reducing merge-state traffic by ~37%.

Significance and Claims

The paper makes a modest, systems-focused claim: PersistentKV identifies a concrete systems niche for adaptive page-aware decode scheduling.

Not a Universal Kernel: The authors explicitly state that PersistentKV is not universally faster; FlashInfer remains superior for small active batches where split-merge overhead dominates.
Work Assignment Matters: The results demonstrate that "work assignment, not only attention math, is a decisive serving-system variable."
Workshop-Level Contribution: The paper positions itself as a workshop-level contribution that provides a reproducible methodology and ablations. It highlights that in low-active, long-context regimes on commodity GPUs, splitting the sequence and using a compact workqueue can significantly outperform a single, highly optimized kernel.
Future Directions: The authors note that a production-grade system would require a learned cost model for thresholds, integration into a full serving runtime (including MLP and sampling), and testing on Hopper-generation hardware with different performance characteristics.