PaperlayerView exact revision on arXiv

Source-linked AI summary

GraniKV: Asymmetric Granularity KV-Cache Paging for Multi-Agent Systems with Long Shared Prefix

Jinhyun Jeon, Sungjoo Yoo

arXiv:2608.15584v1cs.LGcs.AI

TL;DR

Production paged-serving engines use one paging granularity for shared prefixes and per-request suffixes despite their opposing storage needs. GraniKV separates them into contiguous HOT and token-level COLD pools with per-step backend dispatch, reaching up to 2.16× baseline throughput at 16K shared tokens across three configurations. Under heterogeneous multi-prefix serving, it sustains 1.95× while batch-global cascade collapses to parity.

  • Problem

    Production paged-serving engines treat shared prefixes and unique suffixes identically even though prefixes require contiguity and suffixes require fine-grained allocation.

  • Method

    GraniKV uses a contiguous HOT pool for shared prefixes, a token-level COLD pool for suffixes, and a per-step dispatcher between dense and cascade attention backends.

  • Results

    At Lp=16K, GraniKV reaches 2.16×, 1.98×, and 1.57× baseline output-token throughput across three model and tensor-parallel configurations, while sustaining 1.95× under heterogeneous multi-prefix serving.

  • Takeaways & Limitations

    The asymmetric storage layer enables the batched-GEMM prefix backend and carries the gain under heterogeneous multi-prefix load when cascade attention provides no advantage.

  • Takeaways & Limitations

    GraniKV can regress by up to 30% in small-prefix, high-batch settings where prefix savings do not amortize its storage overhead.

Abstract

from arXiv · show

Production paged-serving engines apply uniform paging granularity to the KV cache, even though the two regions of a multi-agent workload have opposite storage requirements: a long shared prefix demands contiguity, while the per-request suffix demands fine-grained allocation. We present \textbf{GraniKV}, a KV-cache layer that allocates the shared prefix in a contiguous HOT pool and the suffix in a token-level COLD pool, combined with a per-step dispatcher which selects the appropriate backend among dual backends for each regime (compute-, memory-, or communication-bound). To the best of our knowledge, GraniKV is the first system to apply asymmetric paging granularity to the KV cache of a production paged-serving engine. At $L_p{=}16$\,K shared tokens GraniKV reaches $\mathbf{2.16\times}$, $\mathbf{1.98\times}$, and $\mathbf{1.57\times}$ output-token throughput over the production baseline on Llama-3.1-8B/TP=1, Qwen-2.5-14B/TP=2, and Qwen-2.5-32B/TP=4. The gain decomposes: cascade attention integration contributes the majority at saturation; the asymmetric storage layer adds $1.05$--$1.15\times$ end-to-end while being what makes the batched-GEMM prefix backend possible at all. Under heterogeneous multi-agent serving with \emph{distinct} prompts of different lengths, the attribution inverts: GraniKV sustains $\mathbf{1.95\times}$ while batch-global cascade collapses to parity --- the storage layer alone carries the win in the regime that motivates the paper.

1 Introduction

GraniKV addresses the mismatch between contiguous shared-prefix storage and fine-grained per-request suffix storage by splitting KV-cache paging into HOT and COLD pools. Integrated with production serving, it enables backend selection across regimes and improves throughput, especially under heterogeneous multi-agent workloads.

  • Motivation: Shared prefixes demand contiguous KV storage, whereas unique suffixes and decoded tokens require fine-grained allocation; production engines use one paging granularity for both.This mismatch affects paged-serving engines including SGLang, vLLM, and TensorRT-LLM.
  • Design: GraniKV splits KV storage into a contiguous HOT pool for shared prefixes and a token-level COLD pool for suffix and decoded tokens.The system is presented as the first asymmetric-granularity KV-cache layer for a production paged-serving engine.
  • Design: The contribution is production integration: one RadixAttention index, engine-integrated admission and eviction, paged fallback, and a hardware-fitted per-step backend dispatcher.The claim concerns coexistence and integration of these components rather than inventing the constituent kernels.
  • Results: At Lp=16K, GraniKV reaches 2.16×, 1.98×, and 1.57× the output-token throughput of the production baseline on Llama-3.1-8B/TP=1, Qwen-2.5-14B/TP=2, and Qwen-2.5-32B/TP=4.On ShareGPT-V3 under Poisson arrivals at saturation, gains grow monotonically with Lp and are at parity (≥0.97×) at short prefixes.
  • Results: The asymmetric storage layer adds 1.05–1.15× end-to-end, becomes the entire win at 1.95× under heterogeneous multi-prefix load, and enables the batched-GEMM prefix backend.Cascade attention provides most of the saturation gain on a single shared prefix, but collapses to parity under heterogeneous load.

2 Related Work

Related work establishes the KV cache as a dominant serving-memory cost and traces the evolution from fixed-size allocation to paged and prefix-aware reuse. Two complementary shared-prefix methods target compute utilization and memory bandwidth, respectively, motivating GraniKV’s asymmetric treatment of prefix and suffix regions.

  • KV-cache serving: KV-cache storage grows with batch size and context length, dominating HBM consumption during production autoregressive serving.Each generated token retains the keys and values of all previous tokens.
  • KV-cache serving: PagedAttention replaces over-reserved contiguous per-request regions with on-demand fixed-size pages and per-request block tables.This addresses memory waste from over-reservation and external fragmentation, improving serving concurrency.
  • Prefix-aware serving: Long shared prefixes recur across multi-agent workloads, while each concurrent request appends a unique suffix.Deployed traces report prefixes up to 23K tokens with 85−97% of prompt tokens shared, and prompt caching uses 1,024–4,096-token minimum prefixes.
  • Prefix-aware serving: RadixAttention indexes cached KV tensors by token-id sequences, enabling longest-prefix reuse and skipping prefill for matched tokens.Reference-counted radix-tree nodes retain cached prefixes while active requests depend on them.
  • Prefix–suffix attention: HydraGen batches shared-prefix queries into a matrix-matrix multiply to improve Tensor Core utilization, whereas cascade inference reads the prefix once per step to reduce redundant HBM traffic.Both separate shared-prefix and per-request suffix attention and merge their outputs with log-sum-exp; cascade reduces prefix HBM reads from O(B) to O(1).

3 GraniKV

GraniKV replaces uniform KV-cache paging with asymmetric HOT/COLD storage: contiguous allocation for shared prefixes and token-level allocation for per-request suffixes. A per-step dispatcher selects dense or cascade attention according to the realized workload regime, enabling up to 2.16× production-baseline throughput while preserving parity outside favorable conditions.

  • Asymmetric KV-cache paging: GraniKV uses contiguous HOT storage for shared prefixes and token-level COLD allocation for unique suffixes, matching their opposite reuse requirements.The HOT pool provides dense, indirection-free prefix reads, while page size = 1 eliminates intra-page waste for suffix tokens.
  • Performance regime: 2.16× is GraniKV’s peak throughput over the production baseline on Llama-3.1-8B/TP=1 when the shared prefix is long and the batch saturates the matmul tile.The improvement reaches this level in the production multi-agent regime of large shared contexts and high concurrency.
  • Dual attention backends: The dense backend targets compute-bound workloads, cascade targets memory-bound workloads, and both split communication-bound steps between saturated bursts and smaller follow-on steps.Dense uses fat batched prefix GEMMs, whereas cascade reuses the shared KV region through its two-level attention approach.
  • Per-step dispatch: GraniKV dispatches per decode step using the realized in-flight batch and an empirically measured A100 saturation threshold of Meff ≥ Msat = 1024.This avoids losing performance when batch shape changes across request arrivals, completions, and prefill bursts.
  • Per-step dispatch: < 0.2% of per-step wall time is added by the dispatch wrapper, making per-step backend selection operationally negligible.The wrapper adds only sub-microsecond CPU work per decode step.

4 Experiments

Under realistic Poisson multi-agent serving, GraniKV’s gains increase with shared-prefix length and peak at 2.16×, 1.98×, and 1.57× over SGLang default on the three evaluated configurations. With distinct prefixes, the storage layer carries the win, while kernel microbenchmarks attribute the mechanism to contiguous HOT storage enabling compute-bound dense prefix attention.

  • Headline results: 2.16×, 1.98×, and 1.57×: GraniKV’s throughput over SGLang default at Lp=16K on Llama-3.1-8B/TP=1, Qwen-2.5-14B/TP=2, and Qwen-2.5-32B/TP=4.Wins grow monotonically with shared-prefix length under Poisson arrivals.
  • Headline results: 0.97×: GraniKV’s worst headline throughput relative to SGLang default at Lp=512, with no regression as Lp increases.Per-step dispatch routes heterogeneous requests correctly, while small-Lp/high-B corners outside the effective regime are discussed as limitations.
  • Multi-prefix serving: 0.99×: SGLang-Cascade reaches parity with SGLang default for distinct system prompts, whereas GraniKV retains the win through per-prefix slabs and class-local dense batching.With distinct prefixes, the batch-global shared level is empty, so cascade attention provides no advantage.
  • Kernel mechanism: 15.6×: the storage layout alone speeds the fixed dense prefix kernel while moving it from HBM-bound to compute-bound.The HOT layout reads shared K/V once per step, unlike the unique paged-equivalent layout that rereads the prefix 256 times.
  • Kernel mechanism: Amdahl prediction with the kernel-isolated s=15.6 overshoots measured end-to-end speedup across all three configurations.The trace-measured residual includes host-side gaps and other terms documented in the decomposition.

5 Conclusion

GraniKV combines asymmetric KV-cache paging with per-step selection between two attention backends for long shared prefixes. At Lp=16 K, it achieves 2.16×, 1.98×, and 1.57× the baseline across three model/TP configurations, without short-prefix regression.

  • Conclusion: GraniKV uses a contiguous HOT pool for shared prefixes, a token-level COLD pool for suffixes, and a per-step dispatcher between two attention backends.
  • Conclusion: 2.16×, 1.98×, and 1.57× the baseline are achieved at Lp=16 K across three model/TP configurations.
  • Conclusion: No regression occurs at short prefixes, while gains increase as production shared prefixes lengthen.

Limitations

GraniKV’s benefits are limited outside its target small-k, long-Lp regime: it offers no specific advantage under high-tenancy prefix thrashing and can regress for short prefixes or small batches. Additional limitations include a static HOT/COLD pool split and evaluation restricted to five model/TP configurations on one A100 host.

  • High-tenancy multi-prefix deployments: At k=64 and Lp=8 K, the active prefix set exceeds the engine’s KV budget, causing all paged engines to thrash equally and GraniKV to offer no specific benefit.This limitation concerns multi-tenant deployments with many concurrent distinct prompts, beyond GraniKV’s target of one to a handful of active system prompts.
  • Static HOT/COLD pool fraction: The HOT-pool fraction fp is fixed at server boot, defaulting to 0.2, and does not adapt to workloads whose prefix-to-suffix ratios shift over time.Future work is a dynamic allocator that monitors HOT occupancy and COLD admission pressure to repartition the KV buffer online.
  • Limited model and hardware coverage: Measurements cover only 5 model/TP configurations on one A100-SXM4-40GB host, leaving broader model, hardware, and parallelism coverage unevaluated.The dispatch trend may generalize, but the M ≥1024 saturation threshold is hardware-specific and would require refitting for other GPUs.
  • Small-Lp / small-B regime: Up to 30% regression occurs at Lp=512, B=512 on Qwen-2.5-14B, where GraniKV reaches 0.70× SGLang default.At B≤256 short-prefix cells, performance is at parity in the worst case: 0.97× SGLang default on Llama-3.1-8B at Lp=512, B=256.

Ethics Statement · Appendix

GraniKV is an infrastructure-level KV-cache optimization that preserves the served model’s behavior while using public checkpoints and a public dataset for evaluation. Its efficiency benefits are dual-use, carrying societal-impact considerations shared with other inference optimizations.

  • Ethics Statement: GraniKV changes GPU-memory layout and attention-kernel selection without modifying model weights, training data, generation behavior, or externally visible outputs.It is an infrastructure-level optimization to the KV-cache layout used by LLM serving engines.
  • Ethics Statement: Bit-exactness against the unmodified baseline serving stack is preserved.This preserves equivalence with the original serving behavior despite the changed cache layout and attention computation.
  • Ethics Statement: Evaluation uses publicly released Llama-3.1-8B, Qwen-2.5-14B, and Qwen-2.5-32B checkpoints.No new model checkpoints are introduced for the evaluation.
  • Ethics Statement: Evaluation uses the publicly available ShareGPT V3 conversation dataset, with no new data collected, annotated, or released.The shared-prefix workload generator is synthetic.
  • Ethics Statement: More efficient LLM serving has dual-use societal impact: lower serving costs can broaden access while also reducing the cost of negative uses.The paper identifies both beneficial and harmful consequences of cheaper LLM-backed services.
  • Ethics Statement: GraniKV’s contribution is at the GPU-memory-layout level, comparable to PagedAttention and FlashInfer.Its societal-impact considerations therefore align with those of related inference-optimization work.

A Source Code and Reproducibility · B Full Evaluation Results · B.1 Full synthetic shared-prefix grid.

GraniKV provides an anonymized repository for reproducing its implementation and measurements, alongside a full synthetic shared-prefix sweep across lengths and batch sizes. The evaluation reports strong latency improvement at the headline cell but excludes Lp=16 K because the non-streaming injector produces invalid timing measurements.

  • A Source Code and Reproducibility: The anonymized repository includes the patched storage and attention implementations, per-step dispatcher, kernel microbenchmark, sweep drivers, workloads, and raw result files.It also contains scripts reproducing Table 5 and Appendix B hardware counters.
  • B.1 Full synthetic shared-prefix grid.: Tables 6–8 sweep Lp ∈ {512, 1K, 2K, 4K, 8K} against B ∈ {8, 32, 128, 256} for three winning configurations.Each configuration contains 20 cells, with throughput reported as 3-rep means in output tok/s and ratios relative to SGLang default.
  • B Full Evaluation Results: At B=256 and Lp=8 K, TPOT p50 falls from 109 ms with SGLang default to 42 ms with GraniKV on Llama-3.1-8B / TP=1.This corresponds to an approximately 2.6× per-token speedup.
  • B Full Evaluation Results: TTFT remains comparable across all three systems, staying within approximately 10% at dense-favourable settings.The passage reports this comparison for the synthetic-grid results.
  • B.1 Full synthetic shared-prefix grid.: Lp=16 K is omitted because the synthetic burst injector uses non-streaming requests whose short-prefix round trips otherwise capture total wall time correctly only through Lp=8 K.For prompts up to Lp=8 K, response times are short enough for this measurement approach.
  • B.1 Full synthetic shared-prefix grid.: At Lp=16 K, captured TTFT is 0 ms and per-cell durations are 4–10× shorter than the Lp=8 K row at the same B, making the measurements physically impossible.The reported cause is that the non-streaming injector returns the full completion in one HTTP response.

B.2 Full ShareGPT Lp×B matrix. · B.3 Cross-Configuration Validity. · B.4 Per-Step Dispatch Trace.

The ShareGPT sweep finds peak or near-peak GraniKV throughput at B=256, while cross-configuration traces reproduce cascade-attention effects end to end. Per-step dispatch follows realized effective batch size, with a clear dense-backend advantage emerging at Meff≥1024.

  • B.2 Full ShareGPT Lp×B matrix.: Tables 9–11 report the full ShareGPT Lp × B sweep at saturation rate across the three configurations.The sweep covers GraniKV output throughput over shared-prefix length and batch size.
  • B.2 Full ShareGPT Lp×B matrix.: B=256 is the throughput peak or near-peak across all three configurations, while higher concurrency adds no gain and can regress.At Lp=16K, Qwen-14B loses 3% and Qwen-32B/TP=4 regresses by 1–2% when B increases beyond 256.
  • B.2 Full ShareGPT Lp×B matrix.: The high-concurrency regression is attributed to simultaneously retaining both attention backends and pre-allocating CUDA-graph buffers, which consumes KV-cache budget.At B=1024, these allocations cost a few hundred MB and squeeze the available KV-cache budget.
  • B.3 Cross-Configuration Validity.: End-to-end torch.profiler traces at B=256 and Lp=8K validate kernel-isolated findings on both Llama-3.1-8B / TP=1 and Qwen-2.5-32B / TP=4.The traces compare all four systems, including structurally different model and tensor-parallel configurations.
  • B.3 Cross-Configuration Validity.: 3.66× lower prefix-attention kernel time is measured for SGLang-Cascade on Llama-3.1-8B / TP=1, versus 1.24× on Qwen-2.5-32B / TP=4.The corresponding times are 11058 → 3021 ms and 20298 → 16423 ms, respectively.
  • B.4 Per-Step Dispatch Trace.: Cells with Meff=B · nq/kv ≥1024 route almost entirely to matmul, whereas sub-threshold cells route to cascade, showing routing follows realized batch rather than a static choice.This trace is reported for Llama-3.1-8B / TP=1 across synthetic-grid (Lp, B) cells.
  • B.4 Per-Step Dispatch Trace.: At Meff=1024, dense-backend per-token TPOT falls to 0.67× cascade; below the threshold it remains within 3%, supporting the dispatcher cutoff.Below Meff=1024, the dispatcher routes ≤13% of steps to dense, while the crossover occurs between Meff=512 and Meff=1024.

B.5 Per-Step Dispatch Overhead.

Per-step kernel dispatch adds negligible overhead and does not lie on the critical path. In the dense-backend microbenchmark, the Python wrapper contributes less than 0.2% of per-step wall time.

  • B.5 Per-Step Dispatch Overhead.: The thin Python wrapper that selects the kernel each step is not on the critical path.This conclusion comes from an nsys capture using CUDA, NVTX, and OS runtime tracing.
  • B.5 Per-Step Dispatch Overhead.: 83 dense-backend prefix-kernel launches averaged 442 µs per launch.The measurements were collected with CUPTI direct capture over a 15-second window.
  • B.5 Per-Step Dispatch Overhead.: < 1 µs of CPU work from the Python wrapper represented under 0.2% of per-step wall time.This was measured at the headline cell of the dense-backend microbenchmark.

B.8 Multi-prefix sensitivity.

GraniKV remains advantageous with four distinct shared prefixes, but uniform multi-prefix workloads approach parity as the active prefix set exhausts the KV budget. With heterogeneous prefix lengths, GraniKV’s asymmetric storage layer preserves the throughput advantage when batch-global cascade sharing disappears.

  • Uniform-length multi-prefix: The sweep evaluates k ∈ {1, 4, 16, 64} distinct system prompts on Llama-3.1-8B/TP=1 at Lp=8 K and B=256.The k=1 case recovers the single-prefix headline.
  • Uniform-length multi-prefix: 1.36× over SGLang default and 1.38× over SGLang-Cascade at k=4, as the HOT pool holds all four prefixes despite reduced in-batch cascade sharing.The cascade margin is larger than at k=1 because the four-way prefix split reduces shared computation.
  • Uniform-length multi-prefix: ≈300 tok/s at k=64, an 87% drop from k=1, as every system thrashes after 64 × Lp = 512K tokens exceed the realizable KV working set.The collapse is a paged-engine-wide KV-budget phenomenon rather than a GraniKV-specific limitation.
  • Heterogeneous-length multi-prefix: 1.95× the production default at heterogeneous k=4, while SGLang-Cascade falls to 0.99× parity because four distinct prefixes eliminate batch-global shared levels.GraniKV independently tags each prefix, assigns contiguous HOT slabs, and batches each size class against its slab; it reaches 1.82× at k=8.
  • Heterogeneous-length multi-prefix: The asymmetric storage layer accounts for essentially the entire end-to-end win in heterogeneous serving, tapering toward the KV-budget wall rather than collapsing abruptly.Mixed slab sizes also exercise first-fit packing in the HOT pool.

B.10 Cross-pool admission and promotion policy.

GraniKV admits each shared prefix to the HOT pool first, requiring contiguous capacity for L_p. If that capacity is unavailable, it places the prefix in COLD without demoting existing HOT prefixes, preserving static HOT addressing during decode.

  • B.10 Cross-pool admission and promotion policy.: HOT-first admission attempts to place each prefix in the contiguous HOT pool.The allocator checks HOT before considering COLD.
  • B.10 Cross-pool admission and promotion policy.: If HOT lacks contiguous capacity for the requested L_p, the prefix is placed in COLD.Failure can result from a full pool or free fragments smaller than the request.
  • B.10 Cross-pool admission and promotion policy.: Existing HOT prefixes are never demoted when a new prefix cannot fit.This no-demotion rule keeps dense-backend stride-math addressing statically valid during decoding.
Loading 2608.15584v1…