acceptodds
Under review as a conference paper at ICLR 2027

JASP: A Joint Attention-Score Proxy for Long-Context KV-Cache Offloading

Abstract

Long-context LLM inference is increasingly important, but its growing KV cache consumes substantial GPU memory and slows autoregressive decoding. CPU offloading relieves this memory pressure, but sparse decoding must still identify important tokens from the complete history and retrieve their K/V states at every step. This process relies on a GPU-resident proxy whose accuracy determines how few tokens can be retrieved, while its footprint controls the history-dependent bytes read during scoring and is therefore a major contributor to scan latency. Existing projection-based proxies typically construct a separate index for each KV head, repeatedly storing and scanning related retrieval information. We show that KV heads share substantial query-conditioned attention-score structure across layers. Based on this observation, we introduce JASP, a joint attention-score proxy that learns one score-aligned code per historical token across KV heads. Unlike key reconstruction, its objective directly preserves the query–key scores that determine token rankings. Query-specific scoring preserves distinct token rankings, and each KV group retains its own selected token set. JASP therefore removes redundant historical proxy traffic while preserving head-specific retrieval behavior. We further combine GPU-driven selected-row retrieval with capacity-balanced CPU/GPU KV placement to translate the compact proxy into serving gains. Under saturated serving, JASP achieves 1.9–2.6× the aggregate decode throughput of FlashAttention-2 across 16K–128K contexts. JASP scores 75.9 on the complete 128K RULER suite, compared with dense attention’s 78.9, while matching dense quality on LongBench. Code is released at [Anonymous/JASP](https://anonymous.4open.science/r/JASP).

open until 14 Dec 2026

est. 32% chance this paper gets accepted at ICLR 2027.

Reject 68%Accept 32%

What do you think this paper will get?

All positions stay anonymous.

Related papers

Loading the map…

Discussion (0)

Sign in to comment.