AB-Sparse: Sparse Attention with Adaptive Block Size for Accurate and Efficient Long-Context Inference
arXiv:2605.12110v1 [cs.DC] 12 May 2026
Di Liu1∗, Ruitian Wang1∗, Chen Chen1‡, Mingliang Gong2‡ Yongjie Yuan2 , Han Zhao1 , Yu Feng1 , Quan Chen1 , Minyi Guo1 1 Shanghai Jiao Tong University 2 Ant Group
Abstract As large language models scale to longer contexts, loading the growing KV cache during attention computation becomes a critical bottleneck. Previous work has shown that attention computation is dominated by a small subset of tokens. This motivates block sparse attention methods that partition the KV cache into fixed-size blocks and selectively compute attention over those blocks exhibiting high importance. However, these methods assign a uniform block size across all attention heads, implicitly assuming homogeneous behavior throughout the model. Our analysis reveals that this assumption is flawed: attention heads exhibit widely varying sensitivity to block granularity, and uniformity leads to suboptimal accuracy. We present AB-Sparse, a training-free algorithm-system co-designed framework that improves accuracy while preserving throughput. AB-Sparse introduces lightweight adaptive block size allocation across attention heads to improve accuracy. To compensate for the additional memory overhead, it further employs lossless block centroid quantization. In addition, custom GPU kernels are developed to support efficient execution with variable block sizes. Evaluation results demonstrate that AB-Sparse achieves an accuracy improvement of up to 5.43 % over existing block sparse attention baselines without throughput overhead.
1
Introduction
Large language models (LLMs) [1, 2, 3] are increasingly deployed in applications that demand long-context understanding, ranging from multi-document summarization [4] to repository-level code analysis [5] and long-form reasoning [6]. While larger context windows enable these capabilities, they introduce significant challenges for efficient model serving. At every decoding step, loading the entire KV cache from memory becomes a bottleneck that scales linearly with context length. For instance, the KV cache for a single request of 128K context length on Llama-3.1-8B [7] reaches 16GB, comparable to the model weights in size. The key to addressing this bottleneck lies in the inherent sparsity of attention: only a small subset of tokens dominates the attention output [8]. This property has inspired a growing body of sparse attention methods that can be categorized into three paradigms, as illustrated in Figure 1. While all three paradigms reduce KV cache loading to lower inference latency, they each face different trade-offs among efficiency, practicality, and accuracy. Token-based methods such as H2 O [9] and InfiniGen [10] estimate per-token importance at every decoding step and select the most relevant tokens, but incur high per-step selection overhead, compromising efficiency. Semantic-based methods such as Clusterkv [11] and RetroInfer [12] cluster tokens by key similarity and retrieve only the relevant clusters, but restructure the KV cache layout, compromising practicality with standard paged KV cache management [13, 14, 15]. Block-based methods such as Quest [16] and ArkVale [17] * Equal contribution. ‡ Corresponding authors: Chen Chen and Mingliang Gong.
Preprint.
Accuracy
Dot Product Estimation
Tokenbased
Semanticbased
Efficiency
Sparse Attention
Blockbased
Top-K Selection Importance Score Block Representation
Practicality
Query
Figure 1: Qualitative comparison of various sparse attention paradigms.
Centroid
Key
Figure 2: Illustration of block sparse attention workflow.
partition the KV cache into fixed-size blocks and load only the Top-K for attention computation, preserving both efficiency and practicality. This makes block-based methods a promising foundation, and the key to fully unlocking their potential lies in improving accuracy without sacrificing throughput. Our analysis reveals that the accuracy limitation of block-based methods stems from a fundamental yet overlooked assumption: a uniform block size is applied across all attention heads [16, 17]. However, attention heads are known to exhibit highly heterogeneous behaviors [18, 19]. As shown in Figure 3, our measurement study further reveals that this heterogeneity extends to their sensitivity to block granularity, with heads varying significantly in their block size preference. Forcing a uniform block size across all heads thus creates an inherent tension: for heads that require fine-grained resolution, an overly large block size coarsens selection granularity, causing critical tokens to be missed; for heads that are insensitive to granularity, an overly small block size unnecessarily increases the number of blocks, amplifying computation and memory cost. The attention head heterogeneity necessitates adaptive block size allocation: assigning finer granularity to sensitive heads to preserve accuracy, while allowing coarser blocks for insensitive heads to reduce overhead. However, realizing this in practical inference systems is non-trivial, with challenges arising from three aspects. First, adjusting block sizes at runtime requires recomputing centroids for all blocks, which is prohibitively expensive; a lightweight mechanism is needed to determine per-head block size assignments prior to deployment. Second, assigning smaller blocks to sensitive heads multiplies their centroid count, introducing significant memory overhead that calls for a lossless compression scheme to reduce centroid footprint. Third, heterogeneous block sizes lead to non-uniform centroid counts across heads and conflict with the uniform page size assumption in existing paged KV cache systems, necessitating custom GPU kernels for efficient execution. We present AB-Sparse, a training-free algorithm-system co-designed framework that addresses these challenges through three tightly integrated components. First, AB-Sparse introduces a lightweight calibration-driven profiling strategy; it exploits the stability of per-head block size sensitivity across diverse inputs to derive reliable assignments prior to deployment. Second, observing that block centroids are precision-insensitive as they are used only for ranking rather than for attention computation, AB-Sparse proposes lossless centroid quantization to reduce memory footprint. Third, AB-Sparse implements custom GPU kernels to support efficient execution with adaptive block sizes. An indexing mechanism enables variable-length batched execution across heads, while a page mapping mechanism maintains compatibility with standard paged KV cache management. We evaluate AB-Sparse on three widely used open source models across two long-context benchmarks. AB-Sparse consistently outperforms existing block sparse attention baselines, achieving up to 5.43% accuracy improvement without sacrificing throughput. Our contributions are summarized as follows: • We conduct a systematic measurement study on block size sensitivity, revealing that attention heads exhibit substantial heterogeneity in block granularity preference (§2.3). • We propose an adaptive block size allocation strategy based on lightweight calibration-driven profiling. We introduce lossless centroid quantization to reduce memory footprint. We design custom GPU kernels with two key mechanisms: indexing for variable-length batched execution, and page mapping for compatibility with standard paged KV cache management (§3). 2
• We evaluate AB-Sparse on three LLMs across two long-context benchmarks, demonstrating consistent accuracy improvements of up to 5.43% over uniform-block-size baselines without throughput overhead (§4).
2
Background and Motivation
2.1
LLMs and Attention Operation
The core component of LLMs is the attention operation [20]. At each decoding step t, the attention operation computes the dot product between the query vector qt ∈ R1×d (where d is the hidden dimension) and the key vectors of all preceding tokens ki ∈ R1×d (for i ≤ t). This product is scaled 1 by d− 2 and normalized through Softmax function to yield the attention score at,i . These scores then weight the value vectors vi , resulting in the final attention output ot . zt,i =
qt · kTi √ , d
at,i = P
ezt,i
j=1..t
ezt,j
,
ot =
X
at,i · vi
(1)
i=1..t
The attention module is typically composed of multiple components, each referred to as an attention head [20, 21]. Each head independently performs computation as in Equation (1) and captures diverse features from different subspaces. The results from all heads are then aggregated to yield the output. LLM inference consists of two stages: the prefill phase and the decoding phase. The prefill phase processes all prompt tokens simultaneously with O(n2 ) complexity. In the decoding phase, each newly generated token attends to all preceding tokens. A standard optimization is to cache these KV states (KV cache), reducing the complexity to O(n). However, loading the full KV cache becomes a bottleneck as context length grows, making the decoding phase memory-bound. 2.2
Block Sparse Attention
A promising approach to reducing KV cache loading is to exploit the inherent sparsity of attention, where only a small subset of tokens dominates the output [22, 23, 24]. Among various sparse attention methods, block-based approaches such as Quest [16] and ArkVale [17] have gained widespread adoption due to their efficiency and practicality with standard paged KV cache layouts [13, 14, 15]. Figure 2 illustrates the common workflow of block sparse attention. The KV cache is partitioned into blocks of equal size B, where each block is represented by a centroid ci 1 . During the estimation stage, the query vector qt computes the dot product with all centroids, yielding importance score rt,i = qt · c⊤ i . The Top-K blocks with the highest scores are then selected for approximate attention. The choice of block size B governs a fundamental trade-off between accuracy and efficiency. A larger B coarsens the centroid representation, degrading Top-K selection accuracy. Conversely, a smaller B increases the total centroid count, amplifying both memory and computation cost. 2.3
Adaptive Block Size Allocation
Existing block sparse attention methods fix Bh = B for all heads. However, attention heads exhibit substantial heterogeneity in how critical tokens are distributed across the KV cache: for some heads critical tokens are densely clustered, while for others they are sparsely scattered. This heterogeneity leads to varying sensitivity across heads, rendering a uniform block size inherently suboptimal. We conduct a systematic analysis of per-head block size sensitivity on Llama-3.1-8B [7] and Qwen38B [26] using the Wikipedia dataset [27] with context length of 32K tokens. For each attention head, we vary the block size over {16, 32, 64} while maintaining a fixed token budget of 4096 2 . We measure attention recall—the fraction of total attention score attributed to the tokens in the selected blocks—as a direct indicator of block selection quality. 1 Various methods have been proposed to compute block centroids, such as mean pooling [25] and max-min pooling [16]. The block representation strategy is orthogonal to this work. 2 Top-K blocks are chosen such that the total tokens across selected blocks equals 4096. This trend is consistent across different token budgets.
3
0.0 16
32 64 Block Size (a) Llama-3.1-8B
0.0 16
1 5 1015202530 Head (a) Llama-3.1-8B
32 64 Block Size (b) Qwen3-8B
Figure 3: Normalized recall curves across block sizes, where normalization is performed with respect to the recall at block size 16. Insensitive heads maintain near-perfect normalized recall across all block sizes, while sensitive heads degrade sharply as block size increases.
1 5 10 15 20 25 30 35
64 32 16
Block Size
0.5
1 5 10 15 20 25 30
Layer
0.5
1.0
Layer
1.0
Insensitive Heads Normalized Recall
Normalized Recall
Sensitive Heads
1 5 1015202530 Head (b) Qwen3-8B
Figure 4: Heatmap of the minimum block size required to retain 98% of peak recall for each attention head across layers. The wide variation across heads and layers indicates that no single uniform block size is simultaneously efficient and accurate.
Attention heads exhibit heterogeneous block size sensitivity. Figure 3 shows the normalized recall curves of representative attention heads for both models. Insensitive heads maintain near-perfect recall across all block sizes, while sensitive heads degrade sharply, dropping below 0.1 at block sizes as small as 32. This reveals that attention heads vary substantially in their sensitivity to block granularity, exhibiting distinct block-size preference. Adaptive allocation outperforms uniform block sizes. Figure 4 shows that the minimum block size to retain 98% of peak recall (i.e., recall at the smallest block size) varies widely across heads and layers, implying that no single uniform block size is simultaneously efficient and accurate. For instance, under a uniform block size of 32, the average recall is only 89.7% and 77.8% on Llama-3.18B and Qwen3-8B, whereas adaptive allocation achieves 98% with a larger average block size of 44.2 and 39.5. This demonstrates that adaptive per-head block size allocation has the potential to improve recall without reducing the average block size. These findings motivate the design of AB-Sparse, which adaptively assigns per-head block sizes to improve accuracy while maintaining system efficiency.
3
AB-Sparse Design
Our empirical findings in Section 2.3 uncover substantial heterogeneity in block size sensitivity across attention heads, which has been overlooked by existing block sparse attention methods. This highlights the potential for adaptive per-head block size allocation to improve accuracy without sacrificing system efficiency. Building on this insight, we first outline the key challenges and our system architecture in §3.1, then detail each component in the following subsections. 3.1
Overview
Adaptive block size allocation entails design challenges in three aspects of the practical inferBlock Size Centroid Physical KV ence system. First, adaptivity requires a block 2 Assignment Quantization Cache size assignment for each attention head; dynamically adjusting assignments at runtime is proSensitivity hibitively expensive, as it requires recomputing Analysis Page Varlen Varlen 3 centroids over all key vectors. Second, assigning Mapping Top-K Estimation 1 smaller blocks to sensitive heads significantly increases the number of centroids that must be Sparse Attention stored; as context length grows, this overhead Figure 5: Architecture of AB-Sparse. scales linearly with sequence length, threatening to bottleneck decoding throughput. Third, heterogeneous block sizes across heads break the execution uniformity assumed by standard batched kernels and are incompatible with existing inference systems that universally adopt fixed-size paged KV cache management. 4
S1 MK1 FWE QA1
(a) Llama-3.1-8B
40
0.95
0.90
10
#Centroids
Recall
Recall
0.95
0.90
Adaptive 1.00
0
#Centroids
Uniform 1.00
20 0 −20
−10
S1 MK1 FWE QA1
Channel (a) Llama-3.1-8B
(b) Qwen3-8B
Figure 6: Recall comparison between adaptive and uniform block size. The adaptive assignments are calibrated solely on wikipedia [27]. Despite this, they consistently outperform uniform block size across all RULER [28] tasks.
Channel (b) Qwen3-8B
−40
Figure 7: Centroid value distribution of Llama3.1-8B and Qwen3-8B. The column-wise patterns indicate that centroid values are tightly clustered per channel, supporting the use of perchannel quantization.
AB-Sparse addresses these challenges with three tightly integrated designs, as summarized in Figure 5. 1 : Observing that per-head block size sensitivity remains stable across diverse inputs, AB-Sparse profiles recall sensitivity on a small calibration set to derive reliable per-head block size assignments (§3.2). 2 : Recognizing that block centroids are precision-insensitive as they serve solely for ranking rather than attention computation, AB-Sparse applies lossless centroid quantization to reduce memory footprint without degrading block selection accuracy (§3.3). 3 : AB-Sparse implements dedicated GPU kernels with an indexing mechanism for variable-length batched execution and a page mapping mechanism for compatibility with standard paged KV cache management (§3.4). 3.2
Lightweight calibration-driven profiling
Determining per-head block size assignments is non-trivial. Adjusting block sizes dynamically requires recomputing centroids over the entire KV cache under each candidate block size, whose cost scales linearly with context length and is prohibitively expensive at inference time. The key insight is that per-head block size sensitivity is stable across diverse inputs. Previous work has shown that individual attention heads learn specialized roles, such as local pattern matching and long-range retrieval [18, 19, 22]. These roles are determined by learned parameters and thus remain consistent across inputs. Our finding that block size preference is similarly head-specific and input-invariant aligns with this understanding. Heads that are sensitive to block size remain sensitive regardless of the input, and vice versa. This suggests that a one-time offline calibration is sufficient to derive reliable assignments that generalize across requests. Concretely, AB-Sparse evaluates attention recall on 50 calibration samples from wikipedia [27]. For each head, the largest block size that satisfies a recall retention threshold τ is selected: Bh∗ = max{B | Recall(h, B) ≥ τ · Recall(h, Bmin )}
(2)
where Bmin is the smallest candidate block size, and τ serves as a knob to balance recall preservation and centroid overhead. To validate generalization, we evaluate the derived assignments on four tasks from RULER [28], covering diverse long-context scenarios. As shown in Figure 6, despite being calibrated solely on wikipedia, the assignments consistently outperform uniform block size across all tasks and models with a comparable average block size. This confirms that per-head block size sensitivity is stable across tasks, and that a one-time calibration is sufficient for reliable deployment. 3.3
Lossless centroid quantization
Adaptive block size allocation assigns smaller blocks to sensitive heads, which can significantly increase their centroid count and amplify memory overhead. To keep this overhead bounded, centroid compression is necessary. We observe that centroid vectors are used solely for ranking and selecting the Top-K blocks, rather than directly contributing to attention outputs. This precision-insensitive property makes quantization a natural fit for centroid compression. However, naively reducing bit width to very low precision risks degrading block selection accuracy. This necessitates a quantization scheme that maximizes compression while preserving accuracy. 5
Top-K Page Recall
INT2-SYM INT2-ASYM
INT4-SYM INT4-ASYM
INT8-SYM INT8-ASYM
Head A
Head B
Logical Blocks
TopKB=1, 3 StrideB=2
Mapping Strategy
1.0 0.8
TopKA=1 StrideA=4
0.6 0.4 0
4
8
12 16 20 Layer Index
24
28 31
Physical Pages
Figure 8: Top-K page recall across layers on Llama-3.1-8B under different quantization bit widths and strategies. INT4 asymmetric perchannel quantization consistently maintains recall above 0.9 across all layers.
Figure 9: Illustration of the page mapping process. Logical blocks of varying sizes are mapped to contiguous physical pages via a block-to-page stride, enabling variable block size to interface with standard paged KV cache management.
A closer examination reveals that for each position along the head dimension (i.e., each channel), centroid values across different blocks follow a concentrated distribution. As shown in Figure 7, centroid values exhibit clear column-wise patterns across both models, confirming that values within each channel are tightly clustered. This intra-channel similarity makes a single scaling factor per position sufficient to capture the value range without introducing large quantization error, enabling more aggressive compression while preserving ranking fidelity. To identify the optimal quantization scheme, we measure Top-K page recall across layers on Llama3.1-8B under different bit widths (INT2, INT4, INT8) and quantization strategies (symmetric and asymmetric3 ). As shown in Figure 8, lower bit widths (INT2) suffer significant recall degradation across layers. While INT4 symmetric quantization improves over INT2, it still fails to consistently maintain high recall. INT4 asymmetric per-channel quantization, on the other hand, achieves recall above 0.9 across all layers and both models. Although INT8 quantization yields slightly higher recall, INT4 asymmetric strikes a better balance between accuracy and memory efficiency. AB-Sparse therefore adopts INT4 asymmetric per-channel quantization. 3.4
Efficient custom GPU kernels
Modern GPU kernels achieve high throughput by batching all attention heads into a single kernel launch, which requires each head to have the same number of centroids for aligned execution. In addition, existing inference systems manage the KV cache in fixed-size physical pages, assuming a uniform block-to-page mapping across all heads. Heterogeneous block sizes break both assumptions. Different heads have varying centroid counts, forcing standard batched execution to resort to either wasteful padding or serial processing. Meanwhile, variable block sizes disrupt the uniform blockto-page mapping, forcing expensive KV gather operations before attention computation. AB-Sparse addresses these challenges with three dedicated GPU kernels. Kernel 1: Fused query-centroid estimation. Since heads with different block sizes have varying numbers of centroids for the same context length, AB-Sparse stores all centroids in a flattened 1D layout and uses a prefix-sum indexing array to delimit the centroid segment of each head. Specifically, if head h has Nh centroids, we define offseth+1 = offseth + Nh , so that the centroids of head h are stored in [offseth , offseth+1 ). This segmented layout enables all heads to be batched into a single kernel launch with fully vectorized execution and no padding overhead. We fuse dequantization into the kernel to avoid materializing dequantized centroids, reducing memory traffic. Kernel 2: Batched Top-K selection. Given the estimation scores from kernel 1, this kernel selects the Top-Kh blocks per head, where each head shares a fixed token budget T , and Kh = ⌈T /Bh ⌉ varies inversely with the assigned block size. This ensures that each head attends to the same number of tokens regardless of its block size, so that accuracy improvements stem from better block selection 3 Symmetric quantization maps values to a zero-centered range, while asymmetric quantization additionally uses a zero-point
offset to handle skewed distributions.
6
rather than increased token coverage. The prefix-sum indexing array from kernel 1 is reused to partition the scores by head, avoiding redundant computation. Kernel 3: Heterogeneous paged attention. The final kernel computes attention over the selected blocks per head. The key challenge is that different heads have different block sizes, making standard paged attention kernels inapplicable without a costly gather step. We avoid this by exploiting the hierarchical divisibility property between logical blocks and physical pages: any block naturally decomposes into an integer number of the finest-granularity pages. As illustrated in Figure 9, each head’s selected blocks are represented as a strided index view with no data movement, remaining fully compatible with existing paged attention kernels.
4
Evaluation
In this section, we perform quantitative experiments to demonstrate that AB-Sparse improves accuracy over existing block sparse attention baselines while preserving throughput. We present accuracy results in §4.2, efficiency results in §4.3, micro study in §4.4, and ablation studies in §4.5. 4.1
Experimental Setup
Hardware and models. We conduct throughput experiments on two hardware platforms: NVIDIA A100-80GB and NVIDIA H800-80GB GPUs. We evaluate AB-Sparse on three representative opensource LLMs: Llama-3.1-8B [7], Qwen3-8B [26], and Qwen3-32B [29], spanning two architecture families and natively supporting context lengths up to 128K tokens. Benchmarks. We employ two complementary benchmarks for accuracy evaluation: RULER [28] and LongBench [30]. RULER is a synthetic benchmark designed to systematically probe long-context capabilities. It encompasses four task categories: retrieval, multi-hop reasoning, aggregation, and question answering, covering 13 tasks in total. We evaluate at context lengths from 16K to 96K to assess performance scaling with sequence length. LongBench provides a more realistic evaluation suite comprising real-world long-document understanding tasks across six diverse categories: singledocument QA, multi-document QA, summarization, few-shot learning, synthetic tasks, and code completion. This benchmark complements RULER by evaluating AB-Sparse on natural text with practical downstream tasks. Baselines. We compare AB-Sparse against full attention [31] and two state-of-the-art block sparse attention methods: Quest [16] and ArkVale [17]. Quest estimates block importance using per-block min-max pooling centroids, while ArkVale employs bounding-volume centroids for tighter block representation. AB-Sparse is applied on top of Quest and ArkVale as a drop-in replacement for their uniform block size assignment, with centroid quantization enabled. For all sparse methods, we fix the KV budget at 4% and the average block size at 32, following the settings adopted in common practice [32, 33]. 4.2
Accuracy Evaluation
Table 1 reports accuracy results on RULER (left) and LongBench (right). AB-Sparse consistently outperforms both baselines across all models and benchmarks. AB-Sparse-Quest improves over Quest by 3.51%/2.19%/5.43% on RULER and 2.47%/2.61%/2.62% on LongBench for Llama-3.18B/Qwen3-8B/Qwen3-32B, respectively. AB-Sparse-ArkVale achieves similar gains, surpassing ArkVale by 3.80%/2.17%/3.16% on RULER and 1.92%/2.19%/1.52% on LongBench. These results confirm that adaptive block size allocation recovers a substantial fraction of the accuracy gap between sparse and full attention without increasing the average KV cache budget. Notably, AB-Sparse consistently improves over both Quest and ArkVale, despite their fundamentally different block representation strategies. This suggests that adaptive block size allocation is agnostic to the underlying block representation, offering a general and pluggable enhancement for block sparse attention methods. 7
4.3
Efficiency Evaluation
We evaluate the decoding efficiency of AB-Sparse against baselines across context lengths from 64K to 256K tokens. Since ArkVale differs from Quest only in block representation method, their latency characteristics are largely identical. We therefore exclude ArkVale from the efficiency comparison and use Quest as a representative block sparse attention baseline. Figure 10 presents the decoding attention latency of all methods on A100 and H800 GPUs across all three models. AB-Sparse matches Quest in latency at shorter contexts and becomes increasingly faster as context length grows. This is because INT4 centroid quantization reduces memory traffic during the estimation stage, an advantage that scales with context length.
Latency (ms)
Full Attention 20
Latency (ms)
40 20
10
64K
128K 256K Context Length (a) Llama-3.1-8B (A100)
20
0
64K
128K 256K Context Length (b) Qwen3-8B (A100)
128K 256K Context Length (d) Llama-3.1-8B (H800)
0
64K
128K 256K Context Length (c) Qwen3-32B (A100)
20
10
64K
0
40
20
10
0
AB-Sparse
20
10 0
Quest
64K
128K 256K Context Length (e) Qwen3-8B (H800)
0
64K
128K 256K Context Length (f) Qwen3-32B (H800)
Figure 10: Decoding attention latency (ms) across three models with varying context lengths on A100 and H800 GPUs. AB-Sparse achieves increasingly lower latency as context length grows.
4.4
Microscopic Study
Throughput vs. batch size. Figure 11 reports throughput of Llama-3.1-8B on A100 with 64K context length and batch sizes of {1, 2, 4}. At batch size 1, AB-Sparse achieves throughput comparable to Quest; at batch size 4, it reaches 1.59× the throughput of Quest. This improvement stems from two factors: INT4 centroid quantization reduces memory traffic during the estimation stage, and the prefix-sum indexing enables padding-free batched execution across heads with heterogeneous centroid counts, allowing AB-Sparse to scale more efficiently as batch size increases.
Methods
16K
32K
64K
96K
Avg.
SQA
MQA
SUM
FL
ST
CC
Avg.
Qwen3-8B Llama-3.1-8B
Full Attention Quest AB-Sparse-Quest ArkVale AB-Sparse-ArkVale
94.15 81.57 85.04 82.77 88.98
92.30 81.02 85.56 83.50 88.43
86.41 77.48 81.49 80.69 82.84
81.99 72.83 74.87 75.43 77.34
88.71 78.23 81.74 80.60 84.40
23.15 20.35 22.19 22.40 22.02
19.79 17.93 18.03 18.11 19.73
25.25 23.40 25.04 23.57 24.87
62.31 58.64 60.85 59.21 62.50
61.46 55.59 60.24 56.88 59.86
58.19 48.69 53.05 53.18 55.87
41.69 37.43 39.90 38.89 40.81
Full Attention Quest AB-Sparse-Quest ArkVale AB-Sparse-ArkVale
91.49 81.16 84.00 84.51 87.29
91.19 82.28 84.01 84.52 87.38
74.44 66.88 69.50 68.99 70.62
71.74 66.49 68.03 67.90 69.29
82.22 74.20 76.39 76.48 78.65
15.64 11.97 13.68 13.35 14.74
13.30 10.88 12.13 11.93 12.44
21.78 20.29 20.86 20.34 21.29
62.30 56.45 59.09 57.18 59.84
65.11 56.67 61.30 57.26 61.20
67.33 52.77 57.61 58.37 62.08
40.91 34.84 37.45 36.41 38.60
Qwen3-32B
Table 1: Accuracy (%) comparison on RULER (left) and LongBench (right) across three models. AB-Sparse consistently outperforms baselines across all tasks.
Full Attention Quest AB-Sparse-Quest ArkVale AB-Sparse-ArkVale
86.25 77.78 86.98 84.84 89.88
87.75 80.09 84.86 84.79 86.37
84.58 75.02 79.34 79.19 81.19
79.47 72.07 75.50 74.25 78.27
84.51 76.24 81.67 80.77 83.93
18.61 14.47 16.78 15.56 16.41
16.64 15.39 15.87 15.60 16.89
21.75 20.56 20.72 20.19 21.47
54.03 49.23 52.36 49.89 54.68
54.58 47.60 56.27 52.79 51.98
18.31 15.50 16.43 15.46 17.17
30.65 27.12 29.74 28.25 29.77
8
85
200 0
Quest
AB-Sparse
AIME24 AMC23 MATH500
20.0 47.5 74.0
23.3 60.0 76.0
Avg.
47.2
53.1
Accuracy (%)
Throughput (token/s)
400
Full Attention Quest AB-Sparse
bsz=1 bsz=2 bsz=4 Batch Size
Figure 11: Throughput (tokens/s) with 64K context length and varying batch sizes on Llama-3.1-8B.
Table 2: Long generation accuracy (%) (pass@4) of Qwen3-8B on three reasoning benchmarks with 32K max generation length.
80 75
Quest AB-Sparse 70 0.02 0.04 0.06 0.08 Budget
Figure 12: RULER accuracy (%) at 64K context length with varying token budget on Llama-3.1-8B.
Long generation accuracy. We additionally evaluate AB-Sparse on long-generation tasks using Qwen3-8B [26] on three reasoning benchmarks: AIME24 [34], AMC23 [35], and MATH-500 [36], which feature short inputs with long outputs. We adopt the sampling parameters recommended by Qwen3-8B [26] (top_k = 20, top_p = 0.95, and temperature=0.6) and set the maximum generation length to 32K following DeepSeek-R1 [37]. We sample each input four times and report pass@4 as the accuracy metric. As shown in Table 2, AB-Sparse-Quest outperforms Quest across three benchmarks, improving the average pass@4 from 47.2% to 53.1%. This demonstrates that adaptive block size allocation is effective not only for long-input tasks but also for long-generation tasks. Dynamic token budget. We evaluate Llama-3.1-8B on RULER at 64K context length, varying the token budget ratio from 2% to 8%. As shown in Figure 12, AB-Sparse consistently outperforms Quest across all budget levels by 2.97–3.69%. The persistent gap as the budget increases suggests that adaptive block size allocation provides benefits complementary to simply enlarging the token budget. Additional results on Qwen3-8B are provided in §??. 4.5
Ablation Study
Effect of centroid quantization. Figure 13 reports RULER accuracy under different centroid precisions across two models, with BF16 as the unquantized baseline. INT4 quantization achieves accuracy comparable to BF16, confirming that per-channel asymmetric quantization preserves block ranking with negligible accuracy loss. Additional results on Qwen3-32B are provided in §??. INT4 INT8 BF16 AB-Sparse-Quest AB-Sparse-ArkVale 100
90
90
80
80
70
70
60
16K 32K 64K 96K Context Length (a) Llama-3.1-8B
60
Naive Latency (ms)
Accuracy (%)
100
AB-Sparse 6
10 5
4
4
2
2
64K 128K256K 64K 128K256K 64K 128K256K Context Length Context Length Context Length
16K 32K 64K 96K Context Length (b) Qwen3-8B
(a) Estimation
Figure 13: RULER accuracy (%) under different centroid precisions across two models. INT4 quantization achieves accuracy comparable to the unquantized BF16 baseline.
(b) Top-K
(c) Attention
Figure 14: Kernel latency (ms) comparison between the naive implementation and ABSparse with varying context length. AB-Sparse achieves consistently lower latency.
Effect of custom kernels. Figure 14 compares the latency of the three core operations between the naive implementation and AB-Sparse’s custom kernels across context lengths from 64K to 256K. The naive estimation and Top-K kernels loop over heads sequentially due to varying centroid counts, while the naive attention kernel gathers selected KV blocks into contiguous memory before 9
computation. Our kernels consistently achieve lower latency, with speedups of up to 5.6×/9.4×/3.1× for estimation/Top-K/attention, respectively.
5
Conclusion
We present AB-Sparse, a training-free framework that improves the accuracy of block sparse attention by exploiting the heterogeneous block size sensitivity across attention heads. Through lightweight calibration-driven profiling, lossless centroid quantization, and efficient custom GPU kernels, ABSparse achieves up to 5.43% accuracy improvement on RULER and 2.62% on LongBench over existing baselines, without throughput overhead.
References [1] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877–1901, 2020. [2] Aixin Liu, Bei Feng, Bing Xue, Bingxuan Wang, Bochao Wu, Chengda Lu, Chenggang Zhao, Chengqi Deng, Chenyu Zhang, Chong Ruan, et al. Deepseek-v3 technical report. arXiv preprint arXiv:2412.19437, 2024. [3] An Yang, Anfeng Li, Baosong Yang, Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chang Gao, Chengen Huang, Chenxu Lv, et al. Qwen3 technical report. arXiv preprint arXiv:2505.09388, 2025. [4] Zukang Yang, Zixuan Zhu, and Jennifer Zhu. Curiousllm: Elevating multi-document question answering with llm-enhanced knowledge graph reasoning. In Proceedings of the 2025 Conference of the Nations of the Americas Chapter of the Association for Computational Linguistics: Human Language Technologies (Volume 3: Industry Track), pages 274–286, 2025. [5] Qinyu Luo, Yining Ye, Shihao Liang, Zhong Zhang, Yujia Qin, Yaxi Lu, Yesai Wu, Xin Cong, Yankai Lin, Yingli Zhang, et al. Repoagent: An llm-powered open-source framework for repository-level code documentation generation. arXiv preprint arXiv:2402.16667, 2024. [6] Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Fei Xia, Ed Chi, Quoc V Le, Denny Zhou, et al. Chain-of-thought prompting elicits reasoning in large language models. Advances in neural information processing systems, 35:24824–24837, 2022. Llama-3.1-8b-instruct. [7] Meta. 1-8B-Instruct, 2024.
https://huggingface.co/meta-llama/Llama-3.
[8] Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficient streaming language models with attention sinks. arXiv preprint arXiv:2309.17453, 2023. [9] Zhenyu Zhang, Ying Sheng, Tianyi Zhou, Tianlong Chen, Lianmin Zheng, Ruisi Cai, Zhao Song, Yuandong Tian, Christopher Ré, Clark Barrett, et al. H2o: Heavy-hitter oracle for efficient generative inference of large language models. Advances in Neural Information Processing Systems, 36:34661–34710, 2023. [10] Wonbeom Lee, Jungi Lee, Junghwan Seo, and Jaewoong Sim. {InfiniGen}: Efficient generative inference of large language models with dynamic {KV} cache management. In 18th USENIX Symposium on Operating Systems Design and Implementation (OSDI 24), pages 155–172, 2024. [11] Guangda Liu, Chengwei Li, Jieru Zhao, Chenqi Zhang, and Minyi Guo. Clusterkv: Manipulating llm kv cache in semantic space for recallable compression. arXiv preprint arXiv:2412.03213, 2024. [12] Yaoqi Chen, Jinkai Zhang, Baotong Lu, Qianxi Zhang, Chengruidong Zhang, Jingjia Luo, Di Liu, Huiqiang Jiang, Qi Chen, Jing Liu, et al. Retroinfer: A vector-storage approach for scalable long-context llm inference. arXiv preprint arXiv:2505.02922, 2025. 10
[13] Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph Gonzalez, Hao Zhang, and Ion Stoica. Efficient memory management for large language model serving with pagedattention. In Proceedings of the 29th symposium on operating systems principles, pages 611–626, 2023. [14] Lianmin Zheng, Liangsheng Yin, Zhiqiang Xie, Chuyue Sun, Jeff Huang, Cody H Yu, Shiyi Cao, Christos Kozyrakis, Ion Stoica, Joseph E Gonzalez, et al. Sglang: Efficient execution of structured language model programs. Advances in neural information processing systems, 37:62557–62583, 2024. [15] Zihao Ye, Lequn Chen, Ruihang Lai, Wuwei Lin, Yineng Zhang, Stephanie Wang, Tianqi Chen, Baris Kasikci, Vinod Grover, Arvind Krishnamurthy, et al. Flashinfer: Efficient and customizable attention engine for llm inference serving. Proceedings of Machine Learning and Systems, 7, 2025. [16] Jiaming Tang, Yilong Zhao, Kan Zhu, Guangxuan Xiao, Baris Kasikci, and Song Han. Quest: Query-aware sparsity for efficient long-context llm inference. arXiv preprint arXiv:2406.10774, 2024. [17] Renze Chen, Zhuofeng Wang, Beiquan Cao, Tong Wu, Size Zheng, Xiuhong Li, Xuechao Wei, Shengen Yan, Meng Li, and Yun Liang. Arkvale: Efficient generative llm inference with recallable key-value eviction. Advances in Neural Information Processing Systems, 37:113134– 113155, 2024. [18] Guangxuan Xiao, Jiaming Tang, Jingwei Zuo, Junxian Guo, Shang Yang, Haotian Tang, Yao Fu, and Song Han. Duoattention: Efficient long-context llm inference with retrieval and streaming heads. arXiv preprint arXiv:2410.10819, 2024. [19] Wenhao Wu, Yizhong Wang, Guangxuan Xiao, Hao Peng, and Yao Fu. Retrieval head mechanistically explains long-context factuality. arXiv preprint arXiv:2404.15574, 2024. [20] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Isabelle Guyon, Ulrike von Luxburg, Samy Bengio, Hanna M. Wallach, Rob Fergus, S. V. N. Vishwanathan, and Roman Garnett, editors, Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, pages 5998–6008, 2017. [21] Joshua Ainslie, James Lee-Thorp, Michiel de Jong, Yury Zemlyanskiy, Federico Lebrón, and Sumit Sanghai. GQA: training generalized multi-query transformer models from multi-head checkpoints. In Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, EMNLP 2023, Singapore, December 6-10, 2023, 2023. [22] Huiqiang Jiang, Yucheng Li, Chengruidong Zhang, Qianhui Wu, Xufang Luo, Surin Ahn, Zhenhua Han, Amir H Abdi, Dongsheng Li, Chin-Yew Lin, et al. Minference 1.0: Accelerating pre-filling for long-context llms via dynamic sparse attention. Advances in Neural Information Processing Systems, 37:52481–52515, 2024. [23] Di Liu, Meng Chen, Baotong Lu, Huiqiang Jiang, Zhenhua Han, Qianxi Zhang, Qi Chen, Chengruidong Zhang, Bailu Ding, Kai Zhang, et al. Retrievalattention: Accelerating longcontext llm inference via vector retrieval. arXiv preprint arXiv:2409.10516, 2024. [24] Zhuoming Chen, Ranajoy Sadhukhan, Zihao Ye, Yang Zhou, Jianyu Zhang, Niklas Nolte, Yuandong Tian, Matthijs Douze, Leon Bottou, Zhihao Jia, et al. Magicpig: Lsh sampling for efficient llm generation. arXiv preprint arXiv:2410.16179, 2024. [25] Enzhe Lu, Zhejun Jiang, Jingyuan Liu, Yulun Du, Tao Jiang, Chao Hong, Shaowei Liu, Weiran He, Enming Yuan, Yuzhi Wang, et al. Moba: Mixture of block attention for long-context llms. arXiv preprint arXiv:2502.13189, 2025. [26] Qwen. Qwen3-8b. https://huggingface.co/Qwen/Qwen3-8B, 2025. [27] wikipedia. 2025.
wikipedia.
https://huggingface.co/datasets/wikimedia/wikipedia,
11
[28] Cheng-Ping Hsieh, Simeng Sun, Samuel Kriman, Shantanu Acharya, Dima Rekesh, Fei Jia, Yang Zhang, and Boris Ginsburg. Ruler: What’s the real context size of your long-context language models? arXiv preprint arXiv:2404.06654, 2024. [29] Qwen. Qwen3-32b. https://huggingface.co/Qwen/Qwen3-32B, 2025. [30] Yushi Bai, Xin Lv, Jiajie Zhang, Hongchang Lyu, Jiankai Tang, Zhidian Huang, Zhengxiao Du, Xiao Liu, Aohan Zeng, Lei Hou, et al. Longbench: A bilingual, multitask benchmark for long context understanding. In Proceedings of the 62nd annual meeting of the association for computational linguistics (volume 1: Long papers), pages 3119–3137, 2024. [31] Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory-efficient exact attention with io-awareness. Advances in neural information processing systems, 35:16344–16359, 2022. [32] Guangda Liu, Chengwei Li, Zhenyu Ning, Jing Lin, Yiwu Yao, Danning Ke, Minyi Guo, and Jieru Zhao. Freekv: Boosting kv cache retrieval for efficient llm inference. arXiv preprint arXiv:2505.13109, 2025. [33] Fang Wu, Congming Gao, Weixi Zhu, and Jiwu Shu. PRKV:page restruct KV cache for high accuracy and efficiency LLM generation, 2026. [34] Art of Problem Solving. Aime problems and solutions. https://artofproblemsolving. com/wiki/index.php/AIME_Problems_and_Solutions, 2024. [35] Art of Problem Solving. Amc problems and solutions. https://artofproblemsolving. com/wiki/index.php?title=AMC_Problems_and_Solutions, 2023. [36] Hunter Lightman, Vineet Kosaraju, Yura Burda, Harri Edwards, Bowen Baker, Teddy Lee, Jan Leike, John Schulman, Ilya Sutskever, and Karl Cobbe. Let’s verify step by step, 2023. [37] Daya Guo, Dejian Yang, Haowei Zhang, Junxiao Song, Peiyi Wang, Qihao Zhu, Runxin Xu, Ruoyu Zhang, Shirong Ma, Xiao Bi, et al. Deepseek-r1: Incentivizing reasoning capability in llms via reinforcement learning. arXiv preprint arXiv:2501.12948, 2025.
12