PQ-HSA: Reusing Product-Quantized Scores for Hybrid Sparse-Approximate Attention

Kunming Shao; Jierun Chen; Yanli Wang; Ruoyu Wang; Haoli Bai; Kwang-Ting Cheng; Chi Ying Tsui

cs.LGSep 27, 2026

PQ-HSA: Reusing Product-Quantized Scores for Hybrid Sparse-Approximate Attention

Authors: Kunming Shao, Jierun Chen, Yanli Wang, Ruoyu Wang, Haoli Bai, Kwang-Ting Cheng, Chi Ying Tsui

Organizations: The Hong Kong University of Science and Technology · Huawei Technologies Ltd. · Sun Yat-sen University · Nanyang Technological University

Abstract

At each decoding step a language model attends over the key-value (KV) cache of every earlier token, so at long context the attention call is bounded by memory bandwidth. Sparse attention reads only a subset of keys chosen by a cheap score estimate, and most methods give the unread tokens zero weight. The output then draws on only a small fraction of the KV cache, and accuracy drops at small budgets, most on tasks that aggregate information across the context. An inverted-file product-quantization (IVF-PQ) index over the cached keys computes an approximate score for every indexed token in order to rank them; after ranking, those scores approximate the attention logits of the tokens left out. PQ-HSA (hybrid sparse-approximate attention) attends the selected tokens with their original keys and values, and the unselected tokens, the background, enter the same softmax through those scores, summed per inverted list and multiplied by the list's mean value. At 128K and a 1-2% retrieval budget, PQ-HSA is more accurate than Quest and SnapKV on Llama-3.1-8B and Qwen3-30B-A3B and stays close to full attention in macro accuracy; with the same selector, the background term raises macro accuracy on the 8B model from 0.71 to 0.83. In the same 128K setting, inside vLLM on one NVIDIA H20, the decode attention call runs 1.6x faster than the FlashAttention-3 kernel; the speedup grows with context length, and a cost model fitted on 8B to 30B models gives the context length at which it begins. A vLLM plugin runs PQ-HSA on two engine versions without changes to the engine source; code is available at https://github.com/KunmingSHAO/pqhsa_release.

Figures & tables

(a) 1% budget
Method	RULER core-4	RULER ext-4	InfiniteBench
full attention †	0.985	0.784	0.655
PQ-HSA	0.950	0.777	0.646
Quest	0.915	0.641	0.575
SnapKV	0.885	0.621	0.512
sink + window	0.050	0.207	0.105

Table 1: Quality at 128K, higher is better. (a) Llama-3.1-8B-Instruct, macro accuracy at a 1% budget; bold marks the best budgeted method. (b, c) All methods on the same RULER prompts at a 2% budget, on Llama-3.1-8B (b) and Qwen3-30B-A3B (c);

\Delta

= PQ-HSA

-

method with its paired 95% interval, bold where the interval excludes zero.

\dagger

full attention; reprod. our reimplementation.

Variant	Selector	Background logits	Background values
A	top- $k$ , probed lists †	$-\infty$	—
T	top- $k$ , all lists	$-\infty$	—
B (PQ-HSA)	top- $k$ , all lists	$\widetilde{\ell}_{i}$	$\bar{v}_{j}$
C	top- $k$ , all lists	$\widetilde{\ell}_{i}$	$\bar{v}_{\mathrm{global}}$
C2	top- $k$ , all lists	$\operatorname{mean}_{i\in T}\widetilde{\ell}_{i}$	$\bar{v}_{\mathrm{global}}$
C2b	top- $k$ , all lists	$\min_{i\in T}\ell_{i}$	$\bar{v}_{\mathrm{global}}$

Table 2: Mechanism variants at 128K, in the notation of Section 3.2 ;

\bar{v}_{\mathrm{global}}

is the token-weighted global value mean, a logit of

-\infty

drops the background (truncation), and — marks no background term. † 128 probed lists (Appendix B ).

Appendix figures & tables7 assets

Supplementary material from the paper’s appendix.

Appendix

$\rho$	Variant	$n$ /task	MQ	MV	CWE	FWE	QA	Macro
1%	Dense	25	0.970	0.960	0.332	0.853	0.960	0.815
1%	B	25	0.930	0.920	0.492	0.827	0.960	0.826
1%	Bx	25	0.930	0.920	0.500	0.827	0.960	0.827
1%	T	25	0.780	0.720	0.360	0.760	0.920	0.708
1%	C	25	0.880	0.880	0.472	0.720	0.920	0.774
1%	C2b	25	0.910	0.860	0.036	0.733	1.000	0.708

Appendix

Table 3: Per-variant quality on Llama-3.1-8B at 128K (higher is better). A selects from IVF-probed candidates; T, B, Bx, C, C2 and C2b rank every indexed token by its PQ score, and T drops the background. MQ/MV: multi-query/multi-value needles; CWE/FWE: common-/frequent-word extraction; QA: question answering. All rows use the same 25 prompts per task; bold marks PQ-HSA (B).

Variant	$n$ /task	Single	MQ	MV	CWE	FWE	Macro
Dense	25	1.000	1.000	1.000	0.648	0.920	0.914
B	25	1.000	1.000	1.000	0.616	0.907	0.905
T	25	1.000	1.000	1.000	0.536	0.907	0.889

Appendix

Table 4: Per-variant quality on Qwen3-30B-A3B at 128K and 2% retrieval. All rows use the same 25 prompts per task. Single: single needle; other columns and bold as in Table 3 .

Quoted in	Model, input tokens	vLLM	KV	Steps	Ratio
Tab. 7	Llama-3.1-8B, 130,400	0.8.5	copy	$3\times 64$	$1.71\times$
Fig. 1 (b)	Llama-3.1-8B, 130,400	0.8.5	copy	32	$1.76\times$
Figs. 1 (b), 5 (a)	Llama / Gradient, 32K–250K	0.8.5	copy	32–64	0.76– $2.31\times$
§ 1 , § 4.5 , Tab. 7	Llama-3.1-8B, 130,400	0.8.5 plugin	paged	$3\times 64$	$\mathbf{1.58\times}$
Fig. 1 (b)	Gradient-1048k, 524,288	0.8.5	paged	32	$\mathbf{2.72\times}$
Fig. 5 (a)	Qwen2.5-14B, 195,936	0.8.5	paged	$3\times 64$	$1.52\times$

Appendix

Table 5: Protocol of every speedup quoted in the paper. All points use one H20, a single request, a 1% retrieval budget, fp16 weights and KV, and vLLM’s FlashAttention-3 kernel for dense attention. CUDA events time the attention call summed over layers, including all of PQ-HSA’s decode kernels. The last row reports time per output token for the whole decode step (median process). Rows with

3\times 64

steps report the median of three processes. “Copy” gathers from a separate copy of the KV cache held by PQ-HSA; “paged” reads the engine’s KV pages in place.

Context	Setting	Speedup
64K	1 request	$1.03\times$
64K	2 requests	$\mathbf{1.71\times}$
64K	4 requests	$1.39\times$
128K	1 request	$1.53\times$
128K	2 requests	$\mathbf{2.60\times}$
128K	FP8 KV for dense, 1 request	$1.37\times$

Appendix

Table 6: Attention-segment speedup (dense/PQ-HSA) per decode step on the concurrent-request path (Llama-3.1-8B, 1% budget, vLLM 0.8.5 with paged KV; median of three processes, lower median of four for 64K with 4 requests), and against dense attention over an FP8 KV cache at 128K (median of three processes per arm).

Model	$L$	$\rho$	$H_{\mathrm{kv}}$	$G$	KV	$R$ / $R_{\mathrm{model}}$
Llama-3.1-8B	128K	1%	8	4	copy	1.71 / 1.63
	128K	1%	8	4	paged	1.58 / 1.63
	128K	2%	8	4	paged	1.37 / 1.39
Qwen2.5-14B-1M	128K	1%	8	5	paged	1.30 / 1.26
	128K	2%	8	5	paged	1.05 / 1.08
	128K	2%	8	5	paged	1.10 / 1.08

Appendix

Table 7: Attention-segment speedup

R

(dense/PQ-HSA) on one H20 across model sizes, context lengths

L

and budgets

\rho

H_{\mathrm{kv}}

is the number of KV heads and

G

the GQA group size. KV: “copy” gathers from a separate copy of the KV cache and “paged” reads the engine’s pages in place (Table 5 ).

R_{\mathrm{model}}

is the cost model of Section 4.4 fitted to all rows; bold marks each model’s largest measured ratio. Rows at 192K and 250K and both Llama-3.1-8B rows at 1% report the median of three processes.

Method	Unit	Selection	Unselected tokens	Index
Quest	page	bound-based top- $k$	dropped	page bounds
ParisKV	token	low-bit pass, rerank	dropped	ANN index
PQCache	token	PQ top- $k$	dropped	PQ codes (host)
SparQ	token	$r$ of $d$ key components	one global mean value	second key copy a
RetroInfer	cluster	centroid top clusters	centroid logit, value sum; rest dropped	clusters (host KV)
Multipole	cluster	centroid top clusters	centroid logit, mean value	key clusters

Appendix

Table 8: Retrieval-family methods compared by design. Index is the state kept besides the KV cache; “host” marks state kept in host memory in the released stacks. a In SparQ’s fastest kernel.

Method	benchmark	$n$	method	PQ-HSA
full attention	LongBench-16	$16{\times}50$	48.8	49.1
PQCache	LongBench-16	$16{\times}50$	47.6	49.1
MagicPIG	LongBench-16	$16{\times}50$	48.6	49.1
RetrievalAttention	RULER core-4	$4{\times}50$	0.925	0.950
SparQ	RULER, 5 tasks	$5{\times}25$	0.777	0.826

Appendix

Table 9: Quality comparisons with retrieval-family baselines: the sixteen LongBench tasks at native lengths (up to 65.4K tokens; 0–100 scale, 50 prompts per task) in the BF16 pipeline, and RULER core-4 at 128K and a 1% budget in the primary FP16 pipeline, where RetrievalAttention and SparQ (at 1%) are our reimplementations under the protocol of Appendix A .

Explore similar work

Near-Oracle KV Selection via Pre-hoc Sparsity for Long-Context Inference
Feb 9, 2026Yifei Gao, Lei Wang, Rong-Cheng Tu +3Dynamic Sparse Attention LLM Inference Optimization
University College London, London, UK · Shenzhen Institutes of Advanced Technology, Chinese Academy of Sciences (CAS), Shenzhen 518055, China · The Chinese University of Hong Kong +1
Sparse Attention as a Range Searching Problem: Towards an Inference-Efficient Index for KV Cache
May 7, 2026Mohsen Dehghankar, Abolfazl AsudehDynamic Sparse Attention Key-Value Cache
Department of Computer Science University of Illinois Chicago
RaBitQCache: Rotated Binary Quantization for KVCache in Long Context LLM Inference
Jun 30, 2026Wenhao Li, Jinhao Dong, Hailin Zhang +3LLM Inference Optimization Dynamic Sparse Attention
School of Information, Renmin University of China, Beijing, China · Peking University, Beijing, China

Feb 9, 2026, cs.LGJ/K move · Enter open · S save

Near-Oracle KV Selection via Pre-hoc Sparsity for Long-Context Inference

Yifei Gao, Lei Wang, Rong-Cheng Tu, Qixin Zhang+2

University College London, London, UK · Shenzhen Institutes of Advanced Technology, Chinese Academy of Sciences (CAS), Shenzhen 518055, China · The Chinese University of Hong Kong+1

Dynamic Sparse Attention LLM Inference Optimization Kv-Cache Management Sparsity Bottlenecks Oracles

Open paper AskPDF arXiv

A core bottleneck in large language model (LLM) inference is the cost of attending over the ever-growing key-value (KV) cache. Although near-oracle top-k KV selection can preserve the quality of dense attention while sharply reducing computation and bandwidth, existing sparse methods generally rely on posterior heuristics, i.e., selectors conditioned on observed attention or proxy scores. Such conditioning introduces posterior bias: it tends to distort true token importance and miss salient tokens, thereby impairing long-range reasoning. To tackle this problem, we propose Pre-hoc Sparsity (PrHS), which selects KV entries before attention scoring and provides explicit accuracy control. Let the attention mass of discarded entries be delta (the dropped mass). Through a marginal-to-mutual-information analysis, we derive an upper bound on the mutual-information loss that depends only on the dropped mass. This relation explains failure modes of posterior heuristics and enables verifiable guarantees by controlling the dropped mass in advance. Within PrHS, we instantiate three orthogonal pre-hoc selectors along the axes of time, depth, and layer. Extensive experiments on LLaMA and Mistral families validate PrHS. Across GSM8K and CoQA, PrHS reduces retrieval overhead by over 90%, achieving 3x higher retrieval sparsity than HShare at matched or better accuracy. It incurs under 1% average degradation on LongBench, lowers attention FLOPs by about 15% versus prior sparse baselines, and yields a 9.9x speedup in attention-operator latency and 2.8x higher throughput on NVIDIA A100-80GB GPUs than the dense baseline.

Abstract

Explore similar work

Near-Oracle KV Selection via Pre-hoc Sparsity for Long-Context Inference

Sparse Attention as a Range Searching Problem: Towards an Inference-Efficient Index for KV Cache

RaBitQCache: Rotated Binary Quantization for KVCache in Long Context LLM Inference