LycheeDecode: Accelerating Long-Context LLM Inference via Hybrid-Head Sparse Decoding
Authors: Gang Lin
Dongfang Li
Zhuoen Chen
Yukun Shi
Xuhui Chen
Baotian Hu
Min Zhang
Venue: ICLR 2026 Poster (2026)
Keywords: Large Language Models
Long-Context Inference
Sparse Attention
Hybrid-Head Attention
Primary area: foundation or frontier models, including LLMs

Abstract: The proliferation of long-context large language models (LLMs) exposes a key bottleneck: the rapidly expanding key-value cache during decoding, which imposes heavy memory and latency costs. While recent approaches attempt to alleviate this by sharing a single set of crucial tokens across layers, such coarse-grained sharing undermines model performance by neglecting the functional diversity of attention heads. To address this, we propose LycheeDecode, an efficient decoding method centered on a fine-grained hybrid-head attention mechanism that employs a hardware-efficient top-$k$ selection strategy. Specifically, the novel HardKuma-based mechanism partitions attention heads into a small subset of retrieval heads that dynamically identify crucial tokens and a majority of sparse heads that reuse them for efficient computation. Through extensive experiments on leading models like Llama3 and Qwen3 across diverse benchmarks for long-context understanding (e.g., LongBench, RULER) and complex reasoning (e.g., AIME24, OlympiadBench), we demonstrate that LycheeDecode achieves generative quality comparable to, and at times surpassing even the full-attention baseline. Crucially, this is accomplished with up to a 2.7x speedup at a 128K context length. By preserving the functional diversity of attention heads, our fine-grained strategy overcomes the performance bottlenecks of existing methods, providing a powerful and validated pathway to both efficient and high-quality long-context LLM inference. The implementation code, kernels, and models will be publicly available.

Published as a conference paper at ICLR 2026
LYCHEEDECODE: ACCELERATING LONG-CONTEXT LLM INFERENCE VIA HYBRID-HEAL SPARSE DECOD- ING
Gang Lin, Dongfang Li, Zhuoen Chen, Yukun Shi, Xuhui Chen, Baotian Hu Min Zhang Harbin Institute of Technology, Shenzhen
glin3104 gmail. com (1 flidongfang, hubaotian, zhangmin20211 @hhit edu.
ABSTRACT
The proliferation of long-context large language models (LLMs) exposes key bottleneck: the rapidly expanding key-value cache during decoding, which im- poses heavy memory and latency costs. While recent approaches attempt to al- leviate this by sharing a single set of crucial tokens across layers, such coarse- grained sharing undermines model performance by neglecting the functional di- versity of attention heads. To address this, we propose LycheeDecode, an effi- cient decoding method centered on a fine-grained hybrid-head attention mech- top-k selection strategy. Specifically,
anism that employs a hardware-efficien mechanism partitions attention heads into a small
the novel HardKuma-base subset of retrieval heads that dynamically identify crucial tokens and major- ity of sparse heads that reuse them for efficient computation. Through extensive experiments on leading models like Llama3 and Qwen3 across diverse bench- marks for long-context understanding (c.g., LongBench, RULER) and complex reasoning (c.g., AIME24, OlympiadBench), achieves generative quality comparable to, and at times surpassing even the full- attention baseline. Crucially, this is accomplished with up to 2.7x speedup at 128K context length. By preserving the functional diversity of attention heads, our fine-grained strategy overcomes the performance bottlenecks of ex- isting methods, providing a powerful and validated pathway to both efficient and high-quality long-contexs LLM inference. Code is available [https: //github. com/HITsz TMG/TMGGNNNN Stt.ee////eeeeeeeeeeeeeeee http://eeee.ann/yceeeeeeeeee
demonstrate that Lyched
INTRODUCTION
Transformer-basee Large Language Models (LLMs) possess remarkable long-context capa- bilities. Leading models like GLM-4 (GLM et al., 2024), Qwen2.5-113 (Yang et al.] (025a) and Gemini-2.5 (Comanici al.] 2025) support up million tokens, enabling superior performance in various long-text tasks such as summarization (Huang et al.) 2021). question answering (Wei et al.) 2022). multi-turn dialogue (Li et al. 2025), and complex reasoning (Wang et 2024).
However, long-context processing is challenging. Due to the autoregressive nature of the Trans- former, for each new token generated, the model must perform attention calculations with the full key-value (KV) cache of previous tokens, leading to frequent memory access and increased I/O overhead. As the sequence grows, the KV cache expands linearly. leading to surge memory usage and a significant increase in computational latency. This severely constrains the deployment and scalability of long-context language models in practical applications. To address this chal- lenge, recent work has proposed sparse attention methods, which reduce computational overhead by computing attention on only a small subset critical tokens, exploiting the inherent sparsity of the attention mechanism. Typically, these methods are categorized into two types: eviction-based methods (Xiao et al. 2024 Li et al., 2024; Zhang et al., 2023). which compress the KV cache by permanently discarding tokens, and dllectionbbased methods (Gao al. 2025 Yang al. 2025b
Wu et al.) (025). which preserve the full KV cache while dynamically selecting a subset of tokens for computation each inference step. key observation is that recent work has identified a high

---

Published as a conference paper at ICLR 2026
Tip Ouedtp Roy for Atheri Stoe
135
13 L23
120 125
127
12 04 18-18-11 14
22 20. 27
Figure 1: Overlap rate of top-k (k- = 5) attention scores between corresponding heads in adjacent layers. The heatmap illustrates the functional diversity among attention heads. We input prompt Please directly output the final answer based on the given question. Question: In a world containing only squares, circles, and triangles, one shape defined by having no angles and being perfectly symmetrical from every point on its perimeter. What is the single name of the only shape that fits this description? Answer: and Llama-3 outputs circle. More cases can be found in Appendix E.3
degree of similarity in critical tokens across consecutive layers (Yang et al. 2025b Hao et al.. 2025) Consequently, they adopt a layer-level sharing strategy, where the same set of selected critical tokens is shared across all heads in subsequent layers. This hierarchical strategy forces all attention heads in the same layer to perform the same function. However, attention heads on the same layer do not exhibit highly similar patterns. As shown in Figure [ the top-k overlap rate of different heads in adjacent layers can vary significantly (c.g., the overlap rate of the 14th head of the last two layers is 0%, while the 24th head is 100%). This suggests that uniform, layer-wise sharing strategy may be overly simplistic, and a more fine-grained, head-based strategy is necessary.
Inspired by this, introduce LycheeDecode, simple and effective hybrid-head sparse decoding method that refines this sharing strategy to a more granular level. Specifically, we classify attention heads into few "retrieval heads" and a majority of "sparse heads". The retrieval heads are responsi- ble for performing full attention computation over the entire context to accurately identify the most important tokens. This selected tokens are then shared with the sparse heads in subsequent layers for efficient sparse attention computation. In this way, LycheeDecode can capture more diverse and relevant attention patterns with minimal precision loss. On the other hand, identifying the types of attention heads typically involves optimizing set of discrete binary variables. Previous work (Xiao) jet al. 2025) circumvents the challenge of discrete optimization by having each head learn a continu- ous variable. Although this variable is amenable to gradient-based methods during training, it must be rounded to a binary value for inference, which introduces a significant train-inference discrepancy that can degrade performance. To bridge this gap, we further introduce the Hard Kumaraswamy dis- tribution Kumaarawamy 1980, Bastings et al. 2019). The HardKuma distribution is specifically designed to produce values that naturally concentrated at O and 1, while remaining differen- tiable. By optimizing the distributional parameters of HardKuma during training, our model learns a near-binary selection mechanism directly, thus mitigating the ing to a more stable and robust head specialization. Evaluation with Llama3 and Qwen3 models on the long-context understanding (e.g., LongBench (Bai et al. 2024). RULER (Hsich et al. 2024)) and complex reasoning (e.g., AIME24, OlympiadBench) tasks demonstrate that LycheeDecode can achieve the best performance among other methods with the same sparsity. It can also achieve 2.7x the end-to-end decoding speedup compared to FlashAttention-2 implementation under 128k context length. Our contributions are summarized as follows:
discrepancy and lead-
We propose LycheeDecode, novel hybrid head sparse decoding method that delegates token selection to a small number of "retrieval heads", allowing for a more fine- grained and effective sparse attention mechanism.
2

---

Published as a conference paper at ICLR 2026
We introduce the Hard Kumaraswamy distribution to address the discrete optimization problem in end-to-end head type identification, reducing the train-inference gap and im- proving model robustness and performance.
We implement the hybrid head block-sparse decoding kernel using TileLang (Wang al. 2025 ) achieving up to 2.7x end-to-end decoding speedup.
RELATED WORK
Sparse attention methods These methods reduce computational and memory overhead during inference, falling into two main types: eviction-based and selection-based. Eviction-based sparse attention aims to lower KV cache memory usage by removing tokens considered less relevant (Xiao ct al.) 2024 Zhang et al. 2023 Li et al. 2024). In contrast, selection- based sparse attention preserves the full KV cache and selects the most important tokens for the attention mechanism to process (Gao 2025 Bastings et al. 2019 Liu et al. 2024) Recent works explored trainable mecha- nisms to further refine token selection. Methods such as Native Sparse Attention (Yuan et (025) and MiniCPM (Team al. (025)) demonstrate that extensive post-training with sparse constraints can yield efficient decoding while maintaining high performance. These methods effectively balance performance with efficiency, mitigating the risk of information loss.
al. 2024
Attention head functional specialization A key insight long-context inference is the functional specialization of attention heads, with a small subset of "retrieval heads" being crucial for recalling
information (Wu al.) 2025). Building on this, RazorAttention (Tang et al., 2025) introduced training-free compression technique that exclusively maintains a full KV cache for these crucial retrieval heads while discarding remote tokens in other heads. DuoAAtention (Xiao et al., (025) and PruLong (Bhaskar al. 2025) categorize heads as either "retrieval" "streaming" by learning continuous gating variable. However, these methods determine the role each head in isolation, lacking mechanism for direct collaboration. Unlike previous works, in our framework, retrieval heads not only perform full attention but also dynamically identify and propagate a curated subset of critical tokens for reuse by the majority of "sparse heads". This creates a fine-grained, cooperative mechanism. It differs from previous methods by enabling more direct and efficient sharing of contextual information between functionally distinct heads.
Cross-layer attention similarity Recent studies have identified high degree of similarity in important tokens and attention patterns across consecutive Transformer layers. This insight has inspired layer-level sharing strategies to improve inference efficiency. Approaches such as TidalDe- code (Yang et al., 20255) and OmniKV (Hao et al. (025) designate a few selector layers to identify critical tokens, which are then reused by subsequent layers for efficient sparse computation. Other methods, like LiSA (Mu et (0024) and PoD (Ma et al. 2024), leverage this redundancy by directly sharing attention weights or key states across layers to reduce redundant calculations. However, their layer-level nature can overlook the functional diversity of individual attention heads. In con- trast, our proposed LecoeeDecode framework introduces more fine- grained, head-level strategy, which preserves the functional diversity of attention heads, allowing for more precise and adaptive mechanism by enabling more efficient sharing of contextual information.
METHODOLOGY
This section introduces LycheeDecode, a head-level sparse decoding framework that leverages the functional specialization of Transformer attention heads, as illustrated in Figure 2 LecheeDecode assigns heterogeneous roles heads: Retrieval Heads that actively refresh critical tokens, and Sparse Heads that efficiently reuse them. By propagating token selections across layers, LycheeDe- code improves efficiency while maintaining model performance.
3.1 HEAD-LEVEL SPARSE DECODING
Retrieval Heads for Critical Token Identification. Certain attention heads are well-suited for capturing long-range dependencies such as co-reference resolution distant contextual links. We
3

---

Published as a conference paper at ICLR 2026
Layer Lilead
0 Hardkaam
5% .P
LayerLittadd Litead Sparse
Propagate suberqua heads
Sparsellood
Spors Load
x -
K
Topkinder
Lager L.Boad
2 HardKima (4)
Layer Rotrioval
Topk
Figure 2: Overall framework. Left: During the training phase, each head calculates full attention and sparse attention, weighted by HardKuma sampling values. Right: During inference, the retrieval head calculates the critical tokens for efficient calculation by the subsequent sparse heads.
designate these as Retrieval Heads (/ E H.R). A Retrieval Head performs standard dense attention over the full sequence:
(K
(9)
A softmax
In
dd
From the resulting attention map ( selects the indices of the top-k attended tokens:
s(-1)) argsTopK(A
where argsTopp returns the tokens with the highest attention weights. The updated set s((+1) is propagated to the head of the same index in the next layer, where it is then used by the subsequent attention heads for sparse attention computation. To initialize the critical token set s all heads in the first layer are designated as Retrieval Heads.
Sparse Heads for Efficient Computation. The other heads perform sparse attention computation on the critical token set, which we designate as Sparse Heads (h Hs). Sparse Head reuses the token set s() inherited from the previous layer and restricts attention computation accordingly:
Le (K)) (DD
0 softmax
(3)
ddee
where Ky (s and ( (1) denote the key and value matrices at head h restricted to the subset of tokens indexed by s( Since no new tokens are selected, the set is propagated unchanged, i.c., Sh s((+1) = S This mechanism reduces both the amount of computation and the K KV-cache loading cost, which constitutes the dominant efficiency gain during autoregressive decoding.
Retrieval-Sparse Synergy. The interaction between Retrieval and Sparse Heads forms a decod- ing pipeline that is both adaptive and efficient. Retrieval Heads periodically refresh the salient token set, ensuring responsiveness to new context, while Sparse Heads exploit these curated subsets for ficient computation across layers. This division of labor allows LeccheDecooee trade off adaptivity and efficiency a principled manner. The complete procedure is summarized in Appendix
3.2
HEAD SPECIALIZATION VIA HARDKUMA
The core challenge here lies effectively classifying each attention head as either Retrieval (H Sparse (H.s) head. This assignment is fundamentally discrete optimization problem over set
4

---

(1) Published as a conference paper at ICLR 2026
of binary variables. Prior work, such DuoAttention (Xiao al. 2025), addresses this by learning a continuous variable for each head. Although this continuous variable is easily optimized, it must be rounded to binary value for inference, which introduces the train- cein-infrrence discrepancy.
To bridge this gap, our approach leverages the Hard Kumaraswamy (HardKuma) distribution (Ku-) maraswamy 1980 Bastings al. 2019), differentiable proxy for binary variables. The HardKuma distribution is specifically designed to produce values that are naturally concentrated at O and 1, yet remains reparameterizable. By optimizing the distributional parameters of HardKuma during train- ing, our model learns a near-binary selection mechanism directly, thus mitigating the train-inference discrepancy and leading more stable and robust head specialization.
The HardKuma Distribution. The HardKuma distribution provides a reparameterizable way to model near-binary choices. sample 2 E [0, 1] is generated through a three-step process:
Sample:
First, a sample u drawn from uniform distribution, u ~ (1,0 1). Then using the inverse CDF of the Kumaraswamy distribution, u is transformed into sample
8 s= (1 (1-w//8)) 1///)1/a, where ~ Kuma(a, 38.
2. Stretch: This sample 8 E (0, 1) is then linearly stretched to a wider interval (p, q) where
p < O 10and 1: s' = s (q 10-2) p) +
3. Rectify: Finally, s' is passed through a hard-sigmoid function (i.c., clipping) to produce
the final sample: min(1. max(0, s')).
This process causes probability mass from the intervals (p. 0] and (1.g) collapse at exactly 0 and 1, respectively, making the output near-binary while the entire transformation from u remains differentiable almost everywhere.
To facilitate the learning of head roles, we introduce a differ- entiable training framework. Formally, for each head h in layer (for / > 0), we associate latent random variable =/ sampled from a HardKuma distribution, governed by learnable o% and g()
Identifying Attention Head Types.
2 HardKuma(a)) (1) (
(4)
During training, each head computes attention maps for both potential roles to create a fully differ- entiable learning path. It generates sparse attention map A using an inherited token set ((t) as
well as full attention map A The full attention map is also used to select the token set S/ ((++))
for the next layer (Fquation]) 2) These two attention maps are linearly combined to form a single hybrid attention map A using the stochastic sample z/2 as a weight:
1 = 2/
(1) R. (1- (0)) ) S,N-
(5)
It creates a fully differentiable path, allowing gradients from the final loss to flow back and update the distributional parameters 0%
(!)) and (8 thus enabling end-to-end learning of the head roles. During inference, this stochastic process is replaced by a deterministic assignment: a head is designated as Retrieval Head if its learned expectation B[2/)
(1) > 0.5, and as Sparse Head otherwise.
Loss Function and Sparsity Control. We optimize a distillation loss to align the logits of our hybrid-head student model with those of the full-attention teacher. Given sequence X partitioned into a prompt X prompp and target X Xungent the teacher encodes X Xprmmps to produce shared KV cache. Conditioned on this cache, both models compute logits over the target tokens. Let y (2) [j] and y ( []] denote the teacher and student logits, respectively, for the j-th target token in the i-th sequence of a batch of size N. The distillation loss is:
Cdistill
Σ j.
- bil
(6)
To enforce a strict sparsity budget on Retrieval Heads, we formulate training as a constrained opti- mization problem using Lagrangian relaxation. The objective is a min-max problem over the distri- butional parameters (a, 3) of the HardKuma selectors and a learnable Lagrange multiplier 2 0:
min max L(a, B, A) Lasanl + A (E[[|]]oo) a.3 X
Nunget)
5

---

Published as a conference paper at ICLR 2026
where the regularizer [|] to the expected number of active Retrieval Heads. The expectation can be expressed in closed form:
is the expected Lg norm of the selection variables, which corresponds
[| = Σ
(1-
( - -
(8)
120,6
where is the CDF function of the Kumaraswam distribution. The detailed derivation of Equation can be found in the Appendix A.3
During training process, (a, 3) are optimized via gradient descent to minimize the objective, while À is updated by gradient ascent according to the constraint violation: if the expected number of active heads exceeds Narrget À increases to strengthen the penalty: otherwise, decreases. This adaptive scheme automatically tunes the effective penalty strength, ensuring the desired sparsity without manual hyperparameter search.
4
EXPERIMENTS
4.1 EXPERIMENT SETTING
Benchmarks, Models, and Baselines We conduct experiments on both efficiency and perfor- mance of LycheeDecode. In Section[4.2 we analyze performance under two scenarios: long-context understanding and complex reasoning. For long-context understanding, we benchmark the Llama3- 8B and Qwen3-8B models on the LongBench dataset, comparing LycheeDecode against advanced sparse attention methods such as TidalDecode YYang et al. (025b). Quest (Tang et al. (024), DuoAt- tention (Xiao et al., 2025) and SeerAttention-R (Gao al. 2025). For complex reasoning, we as- sess the DeepSeek-RI-D benchmarks, including AIME24 and OlympiadBench. In Section we turn to efficiency analysis. Leveraging custom hybrid-head sparse attention kernels, we conduct a head-to head-to-head compar- ison with existing sparse attention methods, measuring both end-to-end speedup and kernel-level acceleration.
models on challenging mathematical reasoning
Training Setup for LycheeDecode To categorize the attention heads, we follow prior work (Xiao et al. 2025), inserting passkeys into the Booksum dataset and calculating distillation loss through passkey retrieval. In training phase, We trained for 3000 steps single NVIDIA A 100 80G GPU using a single batch size, which took only a few hours. The HardKuma distribution for each attention head is initialized to uniform distribution, i.e., parameters a and B are both initialized to 1. The critical token budget is set to 30% of the sequence length. For a fair comparison with TidalDecode, the retrieval head budget was set 32, matching the number of heads that perform full attention in TidalDecode (two full attention layer and two token selection layers, with 8 KV heads each).
4.2
PERFORMANCE EVALLATION
4.2.1
LONG CONTEXT UNDERSTANDING
We evaluate the model's ability to understand long contexts on the LongBench (Bai et al. 2024). a benchmark designed to evaluate LLMs on long-context tasks across diverse NLP domains. Fol- lowing previous work Yang al. 20256), we concentrate on eight tasks that span single/multi- document question answering, summarization, and retrieval: MultiFieldQA (MFQA), NarrativeQA (NrtАA), Qasper (Qas)). 2WikiMQA (2Wiki), HotpoiQA (HotQA), QMSum (QMSm). TriviaQA (TrQA), and Passage Retrieval (PRe).
The results, as detailed Table I. demonstrate that on the -lama-3-8B-Instruue Gratruct-Gradient- 1eteiientt10488 1048k model, LycheeDecode achieves an average score of 33.07 with 4096 token budget, not only outper- forms other sparse attention methods like TidalDecode and Quest but also surpasses the full-attention model in the average score. On the Qwen3-8B model, LycheeDecode outperforms TidalDecode with both 1024 and 4096 token budget, which demonstrates the clear advantage of Lychee head-level token sharing strategy over the layer-level sharing approach used by TidalDecode. Fur- thermore, compared to ScerAttention-R, which relies on a trainable gating network, LycheeDecode achieves comparable slightly superior performance. This demonstrates that our lightweight head

---

Published as a conference paper at ICLR 2026
Performance comparison LongBench benchmark. Decoe achieves the best average score in all settings, surpassing other sparse attention methods and full attention models. "*" indicates double the retrieval head budget. We bold the best-performing scores with the second- best underlined.
Table 1:
Method (Budget) Task
MFQA
NNQQ
Qasp
2Wiki
HoIQA
QMSm
TiQA
PRe
Avg.
Llama-3-aa --s-t--sttut-----eee Crassruct-rrdient-- 1048k
Full Attention
30.76
5.52
14.56
13.32
11.50
19.43
86.56
77.00
32.33
Quest
(1024)
26.21
4.08
12.19
12.61
10.75
19.56
83.47
63.84
29.09
DuoAttention
(1024)
19.02
7.36
8.60
9.68
8.77
17.75
41.92
13.25
15.79
DuoAttention"
(1024)
23.88
6.27
10.44
10.41
7.48
19.00
80.61
47.17
25.66
TidalDecode
(1024)
28.57
7.63
11.11
13.56
9.82
20.37
79.78
75.17
30.75
LycheeDecode
(1024)
28.28
6.12
14.89
14.42
12.81
19.05
82.69
69.92
31.02
Quest
(4096)
28.92
3.74
13.63
12.83
12.15
19.36
85,91
72.50
31.13
DuoAttention
(4096)
22.27
7.16
13.93
12.74
10.73
17.93
83.76
34.75
25.41
DuoAAtennion"
(4096)
23.74
6.63
13.80
13.67
10.40
17.93
86.03
61.00
29.15
TidalDecode
(4096)
30.94
6.19
13.85
14.40
13.71
19.48
86.30
78.00
32.86
Lecoee
(4096)
30.11
5.85
14.39
12.86
19.30
86.78
82.58
33.07
Qwen3-8B
Full Attention
25.84
3.43
10.96
11.97
11.74
20.90
90.21
89.08
33.02
SeerAttention-R (1024)
23,91
2.97
10.28
11.88
11.28
19.04
87.50
86.79
31.71
TidalDecode
(1024)
21.32
2.73
9.96
10.48
9.97
19.27
80.4
83.43
29.70
LycheeDecode (1024)
24.26
3.14
10.45
11.05
12.00
19.81
86.64
91.71
32.38
SeerAttention-R (4096)
24.85
3.30
11.15
12.42
11.35
20.61
90.19
93.17
33.38
TidalDecode (4096)
23.57
2.99
10.79
11.31
20.01
88.94
85.0
31.76
Lecoee
(4096)
24.90
3.32
10.88
12.74
11.68
20.71
90.34
93.25
33.48
identification strategy can effectively capture critical information without the complexity of training and deploying auxiliary gating network.
Table 2: Performance comparison on math reasoning tasks.
Method / Task
Gaokao2023HHn
Minerva
AIMH24
OlympiadHench Avg.
DeepSeek-RI-Disiill
Full Attention
68.8
39.1
23.3
10.2
35.4
TedalDecode
62.5
39.8
13.3
10.9
31.6
TidalDecode w/ Cache Correction
57.0
43.0
33.3
9.4
35.7
LycheeDecolk
68.8
40.6
26.7
10.9
36.8
Decde Cache Correction
68.8
40.0
10.9
40.3
DeepSeek R1 Distilll- Qwen-
Full Attention
74.2
47.7
40.0
10.2
43,0
TidalDecode
39.1
16.7
7.0
30.2
TidalDecode Cache Correction
63.3
26.7
35.0
SycheeDecode
74.2
48.4
43.3
10.9
44.2
Decoe w/ Cache Correction
72.7
47.7
46.7
12.5
44.9
4.2.2 COMPLEX REASONING TASK
To evaluate the reasoning capabilities of LycheeDecode, we conduct experiments on four challeng- ing math reasoning benchmarks: Gaokao2023En (Liao et al.] 2024), Minerva (Lewkowycz et al.) 2022], AIME24 (MAA 2024), and OlympiadBench (He et al. 2024). We compare our method against Full Attention and TidalDecode on two distilled models from the DeepSeek-R1. In our ex- perimental configuration, the number of tokens for sparse attention calculation is set to half of the sequence length, increasing linearly during decoding. Furthermore, mitigate the potential accu- mulation errors from sparse attention mechanisms, we incorporate a Cache Correction strat- egy (Yagg al. 2025b Sun al. 2025). Specifically, after every 32 decoded tokens, perform

---

Published as a conference paper at ICLR 2026
LycheeDecode
TidalDecode
Full Attention
batch size
batch site
botch Fi2
Figure 3: End-to-End Decoding Latency (TPOT) across various context lengths. Decoee and
TidalDecode fixed 4096 budget. Note that TedalDecode only support single batch.
a prefill step over these "polluted" tokens using dense attention to reconstruct and update their key- value (KV) representations within the eache.
As demonstrated in Table LycheeDecode outperforms both the TidalDecode and full attention baselines across both models. The introduction of the Cache Correction strategy further enhances the performance of LycheeDecode, solidifying its superiority. We hypothesize that this advantage over the full-attention model stems from our method's ability to capture more diverse attention patterns through head specialization, which allows LycheeDecode to more effectively focus on the key information crucial for the reasoning process while filtering out irrelevant context that may act noise, leading to robust and efficient inference.
4.3 EFFICIENCY EVALUATION
4.3.1 END-TO-END SPEEDUP
We evaluate the end-to-end decoding latency of LycheeDecode and compare it against TedalDecode and the full attention baseline across varying context lengths and batch sizes. We adopt TPOT (Time Per Output Token) as the primary evaluation metric. LycheeDecode and TidalDecode use fixed 4096 token budget. LycheeDecode leverages our efficient hybrid-head block-sparse decoding kernel, combined with auto-tuning to search for the optimal parameter settings in each layer, since different layers contain varying numbers of sparse heads.
As shown in Figure 3 as the context length grows, the latency of the full-attention model increases sharply. TidalDecode exhibits higher latency than full attention at shorter sequence lengths, but surpasses it in longer contexts (>64K). By comparison, LycheeDecode consistently maintains low latency as sequence length increases, achieving up to 2.7x speedup over full attention and 1.73x faster than TidalDecode under single batch size with 128K context. These results demonstrate that Decoee delivers robust end-to-end acceleration across different settings.
4.3.2 KERNEL-LEVEL SPEEDUP
Attention
sparse head
6/8 sp sparse head
sparse head
8/8 sparse head
batch
butch siae =2
batchaice
batcch slzw
2.30
1.5
0.8
1.2
1.00
0.22
0.4
0.12
8.50
5.0
223
6/4
1200
225
32x
545
32c
66k
128
Contrst Length
Figure 4: Latency comparison of our hybrid head kernel and the FlashAttention-2 kernel across different sparse head ratios, context lengths, and batch sizes.
This section evaluates our custom hybrid head block-sparse decoding kemel (detailed design is shown in Appendix C). We implement the kernel using Tilel.ang and select Flash Attention) (Dao.
8

---

Published as a conference paper at ICLR 2026
2024 our baseline. Experiments are conducted on single NVIDIA A800 GPU across different context lengths (16K to 128K) and batch sizes (1 to 8). We evaluate several configurations of our kernel, progressively increasing the ratio of sparse heads from 4/8 to 8/8 (out of 8 total key-value heads), with a fixed 90% sparsity ratio applied to the sparse heads and the block size set to 64.
The experimental results clearly validate the efficiency of our custom hybrid-head kernel. As shown in Figure while the configuration with 4/8 sparse heads exhibits latency comparable to or slightly underperforming the dense Attentoon baseline, all other settings with a higher degree of spar- sity consistently and significantly outperform it. This performance advantage becomes particularly pronounced as the input sequence length and batch size increase, which is expected, as the decoding kernel is primarily I/O-bound. When the KV cache size is sufficient to saturate memory bandwidth, the gains are substantial; for instance, at a 128K context length with a batch size of 8, our kernel achieves a peak speedup of up to 7x in the fully sparse (8/8) configuration. This evaluation confirms that our specialized kernel effectively translates the algorithmic gains of the hybrid-head strategy into significant kernel-level acceleration by minimizing redundant computation and memory access, serving as the fundamental enabler for the end-to-end speedups observed LycheeDecode.
4.4
ABLATION STUDY
4.4.1
DIFFERENT SPARSITY METHODS
0-70%
32
91%
32
90%
Table 3: Performance comparison of differ- ent head identification methods on different datasets. Scores are averaged across eight se- lected tasks in LongBench.
22
p- Ap2
Top-t
Method Dataset
Passkey Retrival HotpoiQA
Top-
Threshold
70-4
Direct Optimize
32.06
31.02
Hard Conerete
32.13
30.25
Sparsity (9)
HardKuma (Ours)
33.07
31.11
Figure 5: Results of LycheeDecode using dif- ferent sparse method on the LongBench.
To evaluate the effectiveness of different sparsity strategies, we conduct a comparative analysis of their performance-sparsitc trade-offs. We benchmark four distinet families of token selection methods, each with three different configurations: (1) Top-k, which retains a fixed-size set of tokens with the highest attention scores; (2) Top-p, which adaptively selects the smallest set of tokens whose cumulative attention probability exceeds a predefined threshold p; (3) Threshold, which preserves all tokens with attention scores surpassing specific value; and (4) Ratio, which selects a set of top tokens using a budget proportional to the sequence length, designed to increase gradually during the generation process.
For each configuration, we measure two key metrics: (1) Performance, quantified by the average Fl score on the LongBench benchmark, and (2) Sparsity, defined as the percentage of critical tokens identified by sparse heads to the total sequence length during inference.
The experimental results are shown in Figure 5 More details can be found in Appendix H.2 Increas- ing sparsity leads to decline in model performance. This is expected, as higher sparsity reduces the amount contextual information available. Top-p and Ratio perform robustly under low sparsity, sometimes even surpassing Top-k with comparable token budgets. However, their per- formance drops sharply under extreme sparsity. Notably, at equivalent sparsity levels, the Rat io method generally achieves the best performance. We hypothesize that training with fixed-sparsity objective endows the model with a general robustness to sparsity, which in turn allows it to effec- tively handle the dynamic adjustments made by the Ratio o method during inference.

---

Published as a conference paper at ICLR 2026
4.4.2 IDENTIFICATION METHODS & DATASET
To evaluate the advantages of our HardKuma distribution for identifying attention heads, we com- pare it with the direct optimization baseline from Xiao et al. (2025) and the HardConcrete distri- bution used by Bhaskar et al. (2025). The head identification process is performed on two distinct datasets: the previously mentioned Passkey Retrieval task and HotpotQA, which challenges the model to perform multi-hop reasoning over long contexts. For HotpotQA, the distillation loss is calculated based the logits of the answer tokens. Crucially, filter out questions that can be answered without relying on the provided context, thereby ensuring that the identification process specifically rewards heads capable of complex, long-range information integration. The specialized models are then evaluated on the LongBench benchmark with fixed 4096 token budget.
As shown in Table the HardKuma distribution achieves the best overall performance, outper- forming both the direct optimization baseline and HardConcrete distribution and demonstrating its superior ability to identify head type. Its score is slightly lower the HotpotQA dataset, which hypothesize this is because its answers are relatively short; calculating the loss over small number of tokens can lead to higher variance in the gradient estimate, making it difficult to accurately guide the specialization of attention heads. We leave the optimization of tasks where the supervision signal is sparse for future work. Refer to Appendix for more discussion of theoretical advantages.
5 CONCLUSION
We introduce LycheeDecode, a framework that speeds up long-context LLMs by specializing at- tention heads for different roles, enhancing efficiency while maintaining performance. This head specialization is enabled by the HardKuma distribution and custom Tilelang kernel, delivering significant end-to-end speedups. Our work highlights that treating attention heads as functionally specialized units, rather than monolithic block, is powerful and promising direction for LLMs.
ACKNOWLEDGMENTT
This work was supported by the National Natural Science Foundation of China (Grant No. 62406088, 62422603). Guangdong Basic and Applied Basic Research Foundation (Grant No. 2025A1515011376), Guangdong Basic and Applied Basic Research Foundation (Grant No. 2024801010500003, Baotian Hu is the corresponding author.
REFERENCES
Yushi Bai, Xin Lv, Jiajie Zhang, Hongchang Lyu, Jiankai Tang, Zhidian Huang, Zhengxiao Du, Xiao Liu, Aohan Zeng, Lei Hou, Yuxiao Dong, Jie Tang, and Juanzi Li. Longbench: A bilin- gual, multitask benchmark for long context understanding. In Lun-Wei Ku, Andre Martins, and Vivek Srikumar (eds.), Proceedings of the 62nd Annual Meeting of the Association for Com- putational Linguistics (Volume 1: Long Papers), ACL 2024, Bangkok, Thailand, August 1 11-16, 2024, pp. 3119-3137. Association for Computational Linguistics, 2024. doi: 10. 1/218664/V//224. .w:.22..lll.. long 172,
ACI-LONG. 172. URL wtt.ltttsss// //dtpp:://ooio http:/:://oooooooo.....oo// www.orgg110. 18653/665/111 22222/
Jasmijn Bastings, Wilker Aziz, and Ivan Titov. Interpretable neural predictions with diffcrentiable binary variables. In Anna Korhonen, David R. Traum, and Lluís Marquez (eds.), Proceedings of the 57th Conference of the Association for Computational Linguistics, ACL 2019, Florence, Italy, July 28- August 2, 2019, Volume 1: Long Papers, pp. 2963- Association for Computational Linguistics, 2019, doi: 10. L011853//1/P19-100 URL p19- 1284
Adithya Bhaskar, Alexander Wettig, Tianyu Gao, Yihe Dong, and Danqi Chen. Cache me if you can: How many kvs do you need for effective long-context Ims? arXiv preprint urXiv:2506. 17121, 2025.
Gheorghe Comanici, Eric Bieber, Mike Schackermann, Ice Pasupat, Noveen Sachdeva, Inder- jit S. Dhillon, Marcel Blistein, Ori Ram, Dan Zhang, Evan Rosen, Luke Marris, Sam Petulla,
10

---

Published as a conference paper at ICLR 2026
Colin Gaffney, Asaf Aharoni, Nathan Lintz, Tiago Cardal Pais, Henrik Jacobsson, Idan Szpek- tor, Nan-Jiang Jiang. Krishna Haridasan, Ahmed Omran, Nikunj Saunshi, Dara Bahri, Gaurav Mishra, Eric Chu, Toby Boyd, Brad Hekman, Aaron Parisi, Chaoyi Zhang, Kornraphop Kaw- intiranon, Tania Bedrax-Weiss Oliver Wang, Ya Xu, Ollie Purkiss, Uri Mendlovic, lai Deu- tel, Nam Nguyen, Adam Langley, Flip Korn, Lucia Rossazza, Alexandre Ramé, Sagar Wagh- mare, Helen Miller, Nathan Byrd, Ashrith Sheshan, Raia Hadsell Sangnie Bhardwaj, Pawel Janus, Tero Rissa, Dan Horgan, Sharon Silver, Ayzaan Wahid, Sergey Brin, Yves Raimond, Klemen Kloboves, Cindy Wang, Nitesh Bharadwaj Gundavarapu, Ilia Shumailov, Bo Wang, Mantas Pajarskas, Joe Heyward, Martin Nikoltchev, Maciej Kula, Hao Zhou, Zachary Garrett, Sushant Kafle, Sercan Arik, Ankita Goel, Mingyao Yang. Jiho Park. Koji Kojima, Parsa Mah- moudieh, Koray Kavukeuoglu, Grace Chen, Doug Fritz, Anton Bulyenov, Sudeshna Roy, Dim- itris Paparas, Hadar Shemtov, Bo-Juen Chen, Robin Strudel, David Reitter, Aurko Roy, An- drey Vlasov, Changwan Ryu, Chas Leichner, Haichuan Yang. Zelda Mariet, Denis Vnukov, Tim Sohn, Amy Stuart, Wei Liang. Minmin Chen, Praynaa Rawlani, Christy Koh, JD Co- Reyes, Guangda Lai, Praseem Banzal, Dimitrios Vytiniotis, Jieru Mei, and Mu Cai. Gemini 2.5: Pushing the frontier with advanced reasoning, multimodality, long context, and next genera- tion agentic capabilities. CoRR, abs/2507.06266 2025, doi: 10.48550/ARXIV.5507000 https::/doi. /tt////i..gg/...../ 0tt.:oo.rr/.....5//////..... 2507 06261
URL
Tri Dao. Flashattention-2: Faster attention with better parallelism and work partitioning. In The Twelfth International Conference on Learning Representations, 2024. URL https://
openreview net ftr/m/orrrmmm fooorummi-mm222yyn222
Yuan Feng, Junlin Lv, Yukun Cao, Xike Xie, and Kevin Zhou. Ada-kv: Optimizing kv cache evic- tion by adaptive budget allocation for efficient Ilm inference. arXiv preprint arXiv:2407.11550 2024.
Yizhao Gao, Zhichen Zeng, Dayou Du, Shijie Cao, Peiyuan Zhou, Jiaxing Qi. Junjie Lai, Hayden Kwok-Hay So, Ting Cao, Fan Yang, et al. Seerattention: Learning intrinsic sparse attention in your llms. arXiv preprint arXiv:2410. 13276, 2024.
Yizhao Gao, Shuming Guo, Shijie Cao. Yuqing Xia, Yu Cheng, Lei Wang. Lingxiao Ma. Yutao Sun, Tianzhu Ye, Li Dong, et al. Seerattention-r: Sparse attention adaptation for long reasoning. arXiv preprint arXiv:2506.00 088888,900.08889 2025.
Team GLM, Aohan Zeng, Bin Xu, Bowen Wang, Chenhui Zhang, Da Yin, Dan Zhang, Diego Rojas, Guanyu Feng, Hanlin Zhao, et al. Chatglm: A family of large language models from glm-130b to 4lm-4 all tools. arXiv preprint arXiv:2406. 12793, 2024.
Jitai Hao, Yuke Zhu, Tian Wang, Jun Yu, Xin Xin, Bo Zheng. Zhaochun Ren, and Sheng Guo. OmniKV: Dynamic context selection for efficient long-context LLMs. In The Thirteenth Interna- tional Conference on Learning Representations, 2025. URL
net/
Chaoqun He, Renjie Luo, Yuzhuo Bai, Shengding Hu, Zhen Leng Thai, Junhao Shen, Jinyi Hu, Xu Han, Yujie Huang, Yuxiang Zhang, Jie Liu, Lei Qi, Zhiyuan Liu, and Maosong Sun. Olympiadbench: A challenging benchmark for promoting AGI with olympiad-level bilingual multimodal scientific problems. In Lun- Wei Ku, Andre Martins, and Vivek Srikumar (eds.), Pro- ceedings of the 62nd Annual Meeting of the Association for Computational Linguisties (Volume 1: Long Papers). ACL 2024, Bangkok, Thailand, August 11-16, 2024, pp. 3828 3850. Asso- 1/265//1/202.AA11LL..00
ciation for Computational Linguistics, 2024. doi: ong 211
URI.
https: /ttps::/doii. woooorgg../ /00 O. 18653/ v1 /20222244 .cc2..ccl-
Cheng-Ping Hsieh, Simeng Sun, Samuel Kriman, Shantanu Acharya, Dima Rekesh, Fei Jia, and Boris Ginsburg. Ruler: What's the real context size of your long-context language models? In First Conference Language Modeling, 2024.
Luyang Huang, Shuyang Cao, Nikolaus Parulian, Heng Ji, and Lu Wang. Efficient attentions for long document summarization. In Kristina Toutanova, Anna Rumshisky, Luke Zettlemoyer, Dilek Hakkani-Tuu Iz Beltagy, Steven Bethard, Ryan Cotterell, Tanmoy Chakraborty, and Yichao Zhou (eds.). Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguisties: Human Language Technologies, pp. 1436, Online, June
11

---

Published as a conference paper at ICLR 2026
2021. Association for Computational Linguistics. doi: 10.18653//112221..mss mai.://:.....aammmmaaaaaa... 112. URL wttp/oogg.....o..oooooo www.:aaannalamaii.am..aa 112/
Ponnambalam Kumaraswamy, A generalized probability density function for double-bounded ran- 88, 1980.
dom processes. Journal of hydrology,
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 Jason Flinn, Margo I. Seltzer, Peter Druschel, Antoine Kaufmann, and Jonathan Mace (eds.), Proceedings of the 29th Symposium on Operating Systems Principles, SOSP 2023, Koblenz, Germany, October 23-26 2023, pp. 616-626. ACM, 2023. doi: 10.1145/ 3600006.3613165 httthtss://dee org 10. 1...1111 1145/114440 /600000 3613165
Aitor Lewkowycz, Anders Andreassen, David Dohan, Ethan Dyer, Henryk Michalewski, Vinay V. Ramasesh, Ambrose Slone, Cem Anil, Imanol Schlag, Theo Gutman-Sole Yuhuai Wu, Behnam Neyshabur, Guy Gur-Ari, and Vedant Misra. Solving quantitative reasoning problems with lan- guage models. In Sanmi Koyejo, S. Mohamed, A. Agarwal, Danielle Belgrave, K. Cho, and A. Oh (eds.), Advances in Neural Information Processing Systems 35: Annual Conference on Neural In- formation Processing Systems 2022, NeurlPS 2022, New Orleans, LA, USA, November 28 De- cember 9, 2022, 2022. URL http: /ttp::///pppeppeepe nips. cc cttccc///pppee wwwpc/aappprpaee www.fflll/ppoo wtt.e///e/esaap/22//221
hash/
html
Yucheng Li, Huiqiang Jiang, Qianhui Wu, Xufang Luo, Surin Ahn, Chengruidong Zhang, Amir H. Abdi, Dongsheng Li. Jianfeng Gao, Yuqing Yang, and Lili Qiu. SCBench: A KV cache-centric analysis of long-conin contect methods. In The Thirteenth International Conference on Learning Rep- resentations, 2025. URL https: wwtp///pwe.eeeeeeeeeeeee http://www...eerrrooooooooootooa
Yuhong Li, Yingbing Huang. Bowen Yang, Bharat Venkitesh, Acyr Locatelli, Hanchen Ye, Tianle Cai, Patrick Lewis, and Deming Chen. SnapKV: LLM knows what you are looking for before generation. In The Thirty-eigthh Annual Conference on Neural Information Processing Systems, 2024. URL BtLLtttss:// http::///ppeeeevveeee
net/ forum? i d= G0
Minpeng Liao, Chengxi Li, Wei Luo, Jing Wu, and Kai Fan. MARIO: math reasoning with code interpreter output A reproducible pipeline. In Lun-Wei Ku, Andre Martins, and Vivek Srikumar (eds.). Findings of the Association for Computational Linguistics, ACL 2024, Bangkok, Thai- land and virtual meeting, August 1 11-16, 2024, pp. 905-924, Association for Computational Lin- guistics. 2024. doi: 10. (0.18653/V 112///0AFIFIINNGGAAAIII 18653/v1/2028 http://222......../.... wwapanggsacc.... ac1.53
URL https: //dttps://oo.. org/1c..
Di Liu, Meng Chen, Baotong Lu, Huiqiang Jiang, Zhenhua Han, Qianxi Zhang, Qi Chen, Chen- gruidong Zhang, Bailu Ding, Kai Zhang, et al. Retrievalaatention: Accelerating long-context llm inference via vector retrieval. arXiv preprint arxiv:2409, 10516, 2024.
Da Ma, Lu Chen, Situo Zhang, Yuxun Miao, Su Zhu, Zhi Chen, Hongshen Xu, Hanqi Li, Shuai Fan, Lei Pan, et al. Compressing kv cache for cong-context Ilm inference with inter-layer attention similarity, arXiv preprint arXiv:2412.02252, 2024.
MAA. URL
American
invitational
mathematics
examination 2024, 2024.
http:s//errtoorrrooeessssllgggg com httmm//wiii///// /ieiiiiiiixx...ee/ Index /
Yongyu Mu, Yuzhang Wu, Yuchun Fan, Chenglong Wang. Hengyu Li. Qiaozhi He, Murun Yang. Tong Xiao, and Jingbo Zhu. Cross-layer attention sharing for large language models. arXiv preprint arXiv:2408.0 01 090, 2024.
Yutao Sun, Tianzhu Ye, Li Dong, Yuqing Xia, Jian Chen, Yizhao Gao, Shijie Cao, Jianyong Wang, and Furu Wei. Rectified sparse attention. arXiv preprint 7rXiv22506.0400 2025.
Hanlin Tang, Yang Lin, Jing Lin, Qingsen Han, Danning Ke, Shikuan Hong, Yiwu Yao, and Gongyi Wang. Razorateention: Efficient KV cache compression through retrieval heads. In The Thirteenth International Conference on Learning Representations, 2025. URL https: /ttss/////pp//i///////// net:/forum2idstkki22
12

---

Published as a conference paper at ICLR 2026
Jiaming Tang, Yilong Zhao, Kan Zhu, Guangxuan Xiao, Baris Kasikci, and Song Han. QUEST: Query-aware sparsity for efficient long-context LLM inference. In Forty-first International Conference on Machine Learning, 2024. URL
net httpr//vwwee...eteoooouu.aaee. forum?id
KZACYwOMTV
MiniCPM Team, Chaojun Xiao, Yuxuan Li, Xu Han, Yuzhuo Bai, Jie Cai, Haotian Chen, Wentong Chen, Xin Cong, Ganqu Cui, et al. Minicpm4: Ultra-efficient Ilms on end devices, arXiv preprint SXKv:22506009000 2025.
Lei Wang, Yu Cheng, Yining Shi. Zhengju Tang, Zhiwen Mo, Wenhao Xie, Lingxiao Ma. Yuqing Xia, Jilong Xue, Fan Yang, et al. Tilelang: A composable tiled programming model for systems. arXiv preprint arXiv:2504. /7577, 2025.
Minzheng Wang, Longze Chen, Fu Cheng, Shengyi Liao, Xinghua Zhang, Bingli Wu, Haiyang Yu, Nan Xu, Lei Zhang, Run Luo, Yunshui Li, Min Yang, Fei Huang, and Yongbin Li. Leave no document behind: Benchmarking long-contex. LLMs with extended multi-doc QA. In Yaser Al- Onaizan, Mohit Bansal, and Yun-Nung Chen (eds.), Proceedings of the 2024 Conference Em- pirical Methods in Natural Language Processing, pp. 5627-5646, Miami, Florida, USA, Novem- ber 2024. Association for Computational I Linguistics. doi: 10. URL htttpttss: wttp//acaaaaccccc wewellnnhholggg logy. org 2024.022222222 emn paannlp-maaa n. 322/
Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, brian ichter, Fei Xia, Ed Chi, Quoc V Le, and Denny Zhou. Chain-of-thought prompting elicits reasoning in large language models. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh (eds.), Advances in Neural Information Processing Systems, volume 35, pp. 24824-24837 Curran Associates, Inc., 2022. URL. httpp:://pptoeeeeetsetstt ngs. neuri cc wttp:///pppppp.eeeeeeeee File/ 0455056655554444555555500
Wenhao Wu, Yizhong Wang, Guangxuan Xiao, Hao Peng, and Yao Fu. Retrieval head mechanis- tically explains long-context factuality. In The Thirteenth International Conference on Learning Representations, 2025. URL
Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficient streaming language models with attention sinks. In The Twelfth International Conference on Learning Rep- resentations, 2024. URL WtR.httss:// http://://////eevv/eeeeee
net. forum? 6=NG7 SSNNG77551222
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. In The Thirteenth International Conference on Learning Representations, 2025. URL bttpps:/openreevew..o net/ /tt.:/oommmddddddeeeeee
An Yang, Bowen Yu, Chengyuan Li, Dayiheng Liu, Fei Huang, Haoyan Huang, Jiandong Jiang, Jianhong Tu, Jianwei Zhang, Jingren Zhou, et al. Qwen2. 5-1m technical report. arXiv preprint arXiv:2501.15383, 2025a.
Lijie Yang, Zhihao Zhang, Zhuofu Chen, Zikun Li, and Zhihao Jia. Tidaldecode: Fast and accurate LLM decoding with position persistent sparse attention. In The Thirteenth International Confer- ence on Learning Representations, 2025b. URL https::/openreviev..ot net.///wwweeeteeeeme fortt//orummmm ?id
Lkf.7Ckkk
Jingyang Yuan, Huazuo Gao, Damai Dai, Junyu Luo, Liang Zhao, Zhengyan Zhang, Zhenda Xie, Yuxing Wei, Lean Wang, Zhiping Xiao, Yuqing Wang, Chong Ruan, Ming Zhang, Wenfeng Liang, and Wangding Zeng. Native sparse attention: Hardware-aligned and natively trainable sparse attention. In Wanxiang Che, Joyce Nabende, Ekaterina Shutova, and Mohammad Taher Pilehvar (eds.), Proceedings of the 63rd Annual Meeting of the Association for Computational Linguistics (Volume Long Papers), pp. 23078-23097, Vienna, Austria, July 2025. Association for Computational Linguistics. ISBN 977-8-89976-222 URL
doi: 10. 121.33/v1122222lll1111 acl-lon 11267]
126.
Zhenyu Zhang, Ying Sheng, Tianyi Zhou, Tianlong Chen, Lianmin Zheng, Ruisi Cai, Zhao Song. Yuandong Tian, Christopher Re, Clark Barrett, Zhangyang Wang, and Beidi Chen. H2o: Heavy- hitter oracle for efficient generative inference of large language models. In shinth-eevnnth Confer- ence on Neural Information Processing Systems, 2023. URL forum? P=GKR
13

---

Published as a conference paper at ICLR 2026
A HARDKUMA DISTRIBUTION
A.1 KUMARASWAMY DISTRIBUTION
The Kumaraswamy (Kuma) distribution is continuous probability distribution defined on the in- terval (0.1). It is similar to the Beta distribution, its probability density function (PDF) and cumulative distribution function (CDF) simpler and have closed form expressions.
The PDF the Kumaraswamy distribution is given by:
f(x; a, B) aBx (1 (1-=2))--) 27)8-11
where x E (0, 1), a and 3 are positive shape parameters that control the distribution' 's shape.
The shape of the distribution can be unimodal, uniantimodal, increasing, decreasing, constant, depending on the values of a and B3.
The CDP of Kumaraswamy distribution can be defined as:
F( / (E;
(10)
| a05a-1(1 E (1
Let = -", then the differential is du -a5a-1d5 We also need to change the limits of integration: when & = 1, and when { = v, u x² Substituting these into the integral gives:
11-0"
F(x; -B
u B- du
/
1-x""
(11)
=
The PDF and CDF of the Kuma distribution with different parameters are shown in Figure 6
A.2 HARDKUMA DISTRIBUTION
The HardKuma distribution a modification of the Kumaraswamy distribution, engineered to create random variable the closed interval that exhibits both continuous and discrete behavior. It achieves this by having non-zero probability masses at the endpoints 0 and 1, while maintaining a continuous density the open interval (0, 1). This makes it particularly useful for applications like generating differentiable binary masks in machine learning.
The distribution is constructed follows. Let X be a random variable following the Kumaraswamy distribution, i.e., X ~ Kuma(a, 3). We define an intermediate stretched variable T by linearly transforming X to wider interval (p. 9). where p < 0 and q are fixed hyperparameters:
T = + (q - X
(12)
The Kuma random variable, which we denote Z, is then obtained by applying hard-sigmoid rectifier function to
Z = min(1, m0. 7))
(13)
variable Z constructed this way said to follow the HardKuma distribution, i.e., Z ~ HardKuma(o, B)
The key feature of this construction is that the discrete probabilities for Z = and Z = 1 can be computed in closed form, thanks the tractable CDF of the underlying Kumaraswamy distribution.
14

---

Published as a conference paper at ICLR 2026
0..0 P-
--10
P-22
1-10,
1-14, 1-..
Kumaraswamy Probability Density (PDF)
Kumaraswa my
(CDF)
50
1.0
0.8
Probability Denstyy
2.-
9.
Cumulative
04
0.4
60
3.6
0.0
6.0
02
3.6
3.8
Figure 6: PDF and CDF of Kuma distribution with different parameters.
The probability of sampling exactly 0 is the probability that the stretched variable T is less than or equal to O:
P(Z = 0) = P(T 0))
P(p + (8- -)x 5)
P X -
(14)
- p)
p p a
F
Similarly, the probability of sampling exactly 1 is the probability that T' is greater than or equal to 1:
1) P(T 2 1) = =1 - P(T < 1)
1 1 p
=1 P X
(15)
9 - p
1- P a, B
1 F
--p p
The remaining probability mass, 1 P(Z P(z = 1), is distributed continuously over the interval (0, 1). This mixed discrete-continuous nature allows the HardKuma distribution to model binary selections in a way that is amenable gradient- based optimization.
A.3 EXPECTED Lo NORM OF HARDKUMA
A primary application of the HardKuma distribution is to create sparse, differentiable masks. This involves generating vector of random variables = (Zi...., Zn). where each Z, is drawn inde- pendently from HardKuma distribution, Z, ~ HardKuma( (a; 8t). The sparsity of such a vector is measured by its Lg norm |Z]|,, which counts the number of non-zero elements.
A key result, which makes this distribution practical for optimization, is that the expected value of the Lo norm has a tractable, dooeed-frmm expression. We can derive it as follows.
First, we express the Ly norm using the indicator function I[:]]
IZ |b = Σ [[2, + 0
(16)
15

---

Published as a conference paper at ICLR 2026
By the linearity of expectation, the expectation of the sum is the sum of the expectations:
Z [b] = E
21 + 0]
Σ BIT 0
(17)
1=1
The expectation of an indicator function is simply the probability of the event it indicates:
E * 0]] P(Z: * 0)
(18)
Using the complement rule, the probability of being non-zero is one minus the probability of being zero:
Piz 0) = P
(19)
Combining these steps and Equation 14
arrive the final expression for the expected Lo norm:
E
|o|
(1 P-z- =0)
i=1
(20)
1 F (TP O1is Bi
1-
B ALGORITHM PSEUDOCODE
The complete procedure LycheeDeccode is shown Algorithm 1 each layer, the Key-Valuu (KV) cache is first updated with the key and value vectors of the current token. The algorithm then processes each attention head according to its designated type: Retrieval Heads perform full attention operation over the entire KV cache identify and select new of critical tokens. Conversely, Sparse Heads perform more efficient computation, calculating attention only on the sparse subset of tokens provided by the preceding layer. Following the attention step, the outputs from all heads are concatenated and passed through feed-forward network to produce the hidden state for the subsequent layer. This entire procedure repeated until the final logits are produced by the model's output layer.
Algorithm Decode
1: Input: Initial hidden state (0)), KV cache C, selected token set ) token budget
2: Output: Logits
3: for layer | = 0, 1,... L, do
4:
q, q,k. + z!1) "
5: (()) append(k,, v)
6:
K,V V. + C(I) Cc.,,kky, C(I) value
7:
for head h 0, o......, H do
8:
if 0 or 7 then
Retrieval Head
9:
A t softmax g.KT/ /Vd
Sn argTopK (An,, kn,
Select critical tokens
11:
on + An
12:
else
Sparse Head
On + softmax (1n (KN(S))))) /v)) Vn [Sa]
13:
14:
end
15:
end for
16:
0 + Concat Wn
17:
(.(+1) + FFN(()
18:
end for
19: logits + Im (ea.d (4-1)
20: return logits
16

---

Published as a conference paper at ICLR 2026
C KERNEL DESIGN
Algorithm 2 Hybrid-head Block-Sparse Decoding
1: Input: Query q. Key K, Value V. block indices
2: Output: Attention output O
3: for Grid indexed (b, s) by (batch_size, num_split) in parallel do
4: Calculate head id h and head-wise split id sp. base on the sparse head index and split id
5: Load corresponding query block 9b,h GQA group into shared memory
6: Opartial <- O, partial & -00, ppartial f- O
> Initialize accumulators
7:
for each block index E I within the current split do
8:
Load corresponding key block K; and value block V; into shared memory Compute score matrix via GEMM operation Update Opertial, mpartials lpartial with St, V, using online softmax algorithm
9:
S: = 9b. K?
10:
11:
end for
12:
Opartial (b, h, 81.) Opartial/ partial
Store partial output D Store partial log-sum-exp
13:
Ixpartial (b, h 8h. c log (parrtaaa))) + maprtiall
end for
15:
Combine Lpertial. O Opertials O)
Combine different splits
return
A critical challenge in designing efficient hybrid-head attention kernel is the inherent workload imbalance between the different head types. Retrieval heads, which must process the entire Key- Value cache, represent a substantial computational load. In contrast, sparse heads operate on only small, pre-selected subset of blocks, demanding significantly fewer resources. A naive scheduling approach that allocates an equal number of computational resources, such GPU thread blocks, to each head would result in severe performance bottleneck. Threads assigned to sparse heads would complete their tasks rapidly and remain idle, while the threads dedicated full-attention heads would dictate the critical path, leading to gross underutilization of the GPU's parallel architecture.
To overcome this, we implement a workload pooling strategy in our hybrid headid-head sparse decoding kernel that decouples resource allocation from individual heads. Instead of assigning work on per- head basis, we first aggregate the complete set of block computations required by heads (both full and sparse) into single, unified pool of work for each batch item. This aggregated workload is then partitioned into numerous smaller, uniform work units, which we term splits. These splits are subsequently distributed homogeneously among the available GPU thread blocks for execution. By aggregating the heterogeneous computations before partitioning. this approach ensures that every thread block receives workload of roughly equivalent size, maximizing hardware utilization and minimizing overall execution latency. See Algorithm 2 for detailed pseudo code.
D VISUALIZATION OF TRAINING PROCESS
To demonstrate the effectiveness of our proposed head identification strategy bridging the train- inference gap, we visualize the training dynamics of LycheeDecode alongside the baseline DuoAt- tention (Xiao et al. 2025). We conducted the comparison on the Barma-3-8B-Instreee training both for 1000 steps with identical learning rate of 0.01. Figure 7 presents the evolution of the probability that a specific attention head is identified as a "Retrieval Head" during training.
1048k model,
For DuoAttention, the heatmap values represent the continuous gating variables. As observed, DuoAttention exhibits noticeable "grey" areas (values hovering between 0.4 and 0.6) at step 1000. This indicates that simple continuous relaxation often fails to push parameters to the binary tremes. Consequently, rounding these ambiguous values to 0 or I during inference introduces a substantial train-inference discrepancy, potentially degrading performance.
For Lychee the heatmap values represent the expected value E|z(0) of the HardKuma distri- decisive polarization.
bution. In contrast to the DuoAttention, LycheeDecode demonstrates The values quickly converge to either 0 (Sparse Head) or (Retrieval Head), resulting in a clear "blue-and-red" pattern. This confirms that the HardKuma distribution effectively forces the model
17

---

Published as a conference paper at ICLR 2026
DusAttention
LycheeDecode
Kuma PDF (Layer 8, Head 0)
Kuma [Layer 10, Head 0)
Stap
1.42
1.29
1.24
425
6:0
Stnc 100
Step 1.C
2.50
Kuma FDF 1.08, p 0.921
Kuma PDF to 0.97 p 0.98)
1.79
96
1.00
1.50
432
025
9.2
14
90
02
1.20.
Kuma PDF 0.89 0.35)
1.5
1.5
1.50
0.55
0.70
Stea Sm0
Stop 106
Kuma (a. 1.52. m
Kuma PDF [yrm D. p 1.05)
2.0
2.00
1.55
0.54
0.5
1000
Stee 1303
Kuma PDF (c-m 1.53, C.0.11
Kamia PDF (arm 0.01, 1
1.75
1.50
1.00
9.79
6.7%
300
5.5
03
2.4 34
Figure 7: Visualization of head specialization dynamics Middle (Heatmaps): The probability of each head being identified as a Retrieval Head across training steps for DuoAttention (left) and LycheeDecode (middle). Right (PDFs): Evolution of the LycheeDecode Kuma distribution PDFs for specific heads at steps 0, 100, 200, 500, and 1000, showing how probability mass effectively concentrates at the boundaries.
Left &
18

---

Published as a conference paper at ICLR 2026
to make discrete decisions during the training phase itself, thereby minimizing the consistency gap between training and inference.
The two rightmost columns of Figure 7 provide microscopic view of this process by plotting the Probability Density Functions (PDFs) of the Kuma distribution for two representative heads (Layer 8 Head and Layer 10 Head 0) specific training steps. Initially uniform at Step O, the distributions undergo dramatic transformation. For the head specializing as a Retrieval Head, the probability mass shifts almost entirely the right, while for the Sparse Head, it collapses the left. This visualization corroborates that our optimization objective successfully shapes the underlying distribution be near-bine einarb.
E MORE EXPERIMENT RESULTS
E.1 RULER BENCHMARK
To assess the ability to comprehend longer contexts, we employ the RULER benchmark (Hsieh al. 2024), synthetic benchmark designed for thorough evaluation of long-context lan- guage models beyond simple retrieval tasks. RULER expands on the needle-in-a-haystack (NIAH) test by including more complex tasks like multihop tracing and aggregation, offering configurable sequence lengths and task difficulties. For our evaluation, we selected tasks including tiah-singlel, fiichmultikey/ 1, niah multivalue, niah maahiiultiu vt. fwe. qal, and qa2 to a wide range of long- context understanding capabilities. We configure LycheeDecode with a fixed budget of 4096 tokens and compare the full-attention Llama3-80 In-r-nnttutt--raaa Gradient- 1048k model.
The experimental results are shown in Table 4 As indicated, in shorter context scenarios, the per- formance of our method is highly competitive with the full attention model. For instance, at 8k con- text length, our approach achieves average score of 62.79, closely approaching the full-attention performance LycheeDecode de-
model's score of 63.30. As the context length increases, creases slightly. This performance degradation is acceptable trade-off, given that method operates on a fixed and significantly smaller 4096 token budget.
Table 4: Performance comparison of LycheeDecode and full attention model on RULER benchmark. Decoee uses a fixed budget of 4096.
Context
single
multikey
multivalue
multiquery
fwe
qui
qa2
Avg.
Full Attention
4k 100.0
89.6
87.8
79.2
17.4
0.1
79.8
56.4
63.7
100.0
70.0
16k
100.0
93.0
95.7
81.0
19.8
0.0
74.2
53.4
64.6
32k
99.2
97.4
96.5
81.9
19.8
0.0
70.6
51.6
65.9
64k
99.4
98.4
96.8
93.7
19.8
0.0
70.4
47.6
65.8
LychoeDecode
100.0
78.9
8k
100.0
94.4
90.6
65.9
19.4
0.4
75.4
56.2
62.8
168
100.0
81.8
96.3
68.7
19.6
0.0
71.0
53.4
61.4
32k
97.8
82.0
94.9
65.1
19.8
0.0
66.2
49.6
59.4
64k
99.6
73.2
90.3
81.7
19.9
0.0
63.4
44.4
59.0
E.2 DETAILED RESULTS OF DIFFERENT SPARSE METHODS
This section provides a detailed breakdown of the results from the ablation study on different sparsity methods, as discussed in Section 4.4.1 and visualized in Figure p the performance results of LycheeDecode on the LongBench benchmark when configured with different token selec- tion methods, including Top-k, Top-p, Threshold and Ratio, each with varying parameters. Complementing this, quantiiis the sparsity level (as a percentage of critical tokens selected) for each corresponding strategy and setting.
19

---

Published as a conference paper at ICLR 2026
Table 5: Performance comparison of LycheeDecode using different sparse strategies LongBench.
Method / Task
MFQA
NrQQ
Qasp
2Wiki
HotQA
QMSm
TrQA
PRc
Avg.
Top-kx-1021
28.28
6.12
14.89
14.42
12.81
19.05
82,69
69.92
31.02
Top-kx=20s
28. 13
5.78
14.72
12.76
11.82
19.14
84,98
78.42
31.97
Top-ki==1096
30.11
5.85
14.39
12.86
12.66
19.30
86.78
82.58
33.07
1op-Pp-0.90
30.01
11.05
12,65
13,48
12.91
21.05
74.27
42.67
27.26
10p-9p=0.995
33,42
9.47
13.84
15.55
13.70
20.12
80.93
65.25
31.54
ToP-Pp=0.999
31.02
6.30
13.77
14.00
12.25
19.99
85.77
74.67
32.22
Threshold, -10-4
25.22
8.17
11.63
13.78
10.86
19.75
60.08
19.93
21.18
Threshold,-100
28.64
6.41
14.91
15.02
13.48
19.31
78.16
$3.36
28.66
Threshold,
29.73
6.74
14.00
13.75
11.49
19.71
83.48
77.08
31.99
Raliog=70%
26.68
6.54
15.58
14.10
12.95
18.91
83.82
80.67
32.41
Ratio
26.65
6.88
13.59
15.70
11.97
19.19
81.94
77.17
31.63
Ratioy=00%
27.96
6.61
11.91
15.01
13.33
18.95
80.41
65.33
29.94
Table 6: Sparsity (%) of LycheeDecode under different settings on LongBench benchmark.
Method Task
MFQA
NrQQA
Qasp
2Wiki
HotQA
QMSm
TrQA
PRe
Avg.
Top-kx-1024
87.94
92.24
67.70
71.46
82.16
88.66
86.88
74.46
81.4
Top-kx-200
79.17
86.60
44.54
54.16
69.96
80.41
77.57
58.26
68.8
Top-kx==000
61.91
75.32
14.21
30.80
51.62
63.92
60.12
30.17
48.5
Top-Pp=0.90
81.81
84.59
76.34
79.85
82.12
85.78
79.71
79.17
81.1
Top-9p=0.906
76.22
78,50
70.42
75.06
76.67
81.50
74.99
73.80
75.9
Top-9p=0.999
59.67
61.02
53.66
58.40
58.32
60.76
57.29
56.37
58.2
Threshold,=10-
96.95
98.25
91.69
93.33
95.97
97.77
96.44
93.92
95.5
Threshold
86.20
90.38
73.22
78.60
85.05
88.44
85.12
78.87
83.2
Threshold 10
65.36
72.07
48.04
55.11
62.89
68.05
62.92
54.21
61.1
Ratiog-909
90.00
90.00
90.00
90.00
90.00
90.00
90.00
90.00
90.0
Ratiog-80%
80.00
80.00
80.00
80.00
80.00
80.00
80.00
80.00
80.0
Ratiog=70%
70.00
70.00
70.00
70.00
70.00
70.00
70.00
70.00
70.0
E.3
MORE CASES
In this section, we provide additional examples to illustrate the behavioral differences among various attention heads. We use prompts that require simple logical reasoning. For each attention head, we calculate the attention scores of the final answer token with respect to all previous tokens and identify the top-k crucial tokens with the highest scores. Subsequently, we compute the overlap rate of these crucial tokens for each attention head with those of the corresponding head in the adjacent layer. The results are presented Figures [0] Figure
Figure and Figure
E.4 ABLATION STUDY
To investigate the trade-offs between model performance and inference efficiency, we conducted an ablation study using the Ilama3-8B, -nnttruttttt -ee--structtGGraientttttt Grrdieent-10483 model. We evaluated the generative quality based on the average score across the LongBench benchmark, while efficiency was quantified by the end-to-end decoding speedup (measured via Time Per Output Token, TPOT) relative to the Full Attention baseline. In this experiment, we explored a range of sparsity configurations by varying two key hyperparameters: the critical token budget, which was to 1024, 2048, and 4096 tokens, and the ratio of retrieval heads, which was tested at 12.5%. 25.0%, and 50.0% the total attention heads.
As illustrated in Figure g the results demonstrate a clear trade-off between performance and effi- ciency. We observe that increasing the token budget from 1024 4096 consistently enhances the
20

---

Published as a conference paper at ICLR 2026
1024
12.5%
-
25 0%
50.0%
2046
120
Emelency
1024
1.15
2043
1.30
1024
1.00
4206
31.5
Performance
Figure 8: Performance and efficiency rrade-offs under different token budgets and retrieval head budgets.
LongBench average across settings by retaining more context, though this naturally leads to decrease decoding speedup. Regarding the proportion of retrieval heads, a smaller ratio yields the highest speedup by minimizing computational overhead. However, in terms of performance, simply maximizing the number of retrieval heads does not always lead to the best results. Notably, with larger token budgets (2048 and 4096). the 25.0% configuration outperforms the 50.0% setting. We hypothesize that this is due to the noise-filtering property of LycheeDecodee An excessive propor- tion of retrieval heads may introduce irrelevant context, whereas balanced configuration allows sparse heads to effectively focus on the most critical information.
E.5 ATTENTION VISUALIZATION
To analyze the behavior of Sparse Heads and investigate how they handle noisy context compared to full-attention, we conducted a case study using a logical reasoning prompt with irrelevant distractor text (Figure 9 We computed the attention weights of the final answer token with respect to the entire preceding context. We calculated the average attention scores across all Retrieval Heads (which execute full attention) and compared them against the average scores across all Sparse Heads (which execute sparse attention). This comparison allows us to directly observe the impact of the sparsity mechanism the attention distribution.
The comparative visualization is presented in Figure As shown in Figure a), the Retrieval Heads display a diffused attention pattern characteristic of full attention, where significant attention mass allocated to irrelevant distractor tokens (e.g., "West", "the"). In contrast, Figure (b) demonstrates that the Sparse Heads effectively eliminate this noise. Since Sparse Heads constitute the major- ity of the model's ' computation, this "denoising" effect explains the counter-intuttive finding that LecheeDecode can outperform the full-attention baseline, as filters out interference that would otherwise distract the model.
F IMPLEMENTATION DETAILS
We provide additional experimental details to ensure the reproducibility of our results. For the training phase, set the learning rate to 0.01. The stretching interval (p, for the HardKuma distribution was set to (-0.1, 1.1). For the Passkey Retrieval dataset, we follow the setup of Xiao et al. (2025) by inserting ten 32-word passkeys into the BookSum dataset, with the prompt length sampled from a range of 1k to 10k tokens. For the HotpotQA dataset, we filter out questions that could be answered without requiring the provided context. The prompt length for HotpotQA was
21

---

Published as a conference paper at ICLR 2026
(a) Full Attention
Question John father has four children: North, South, and East. What the name of the feurth child? <thinks
These the cardinal directions, We have North, South, and East. To complete the compass the remaining
direction West, logically follows the set of four points. So the final answer is West, Walt, am overlooking the
Introduction, The man is identified as "Johnis father".
Therefore Johnn is one of the children., The direction pattern is John
distraction, </thinks So the final answer
Sparse Attention
Question John's children: North, South,
child? <thinks
These are the cardinal directions. We have North, South, and East. To complete the compass rese, the remaining
direction is West. It logically follows the set of four points. So the final answer is West, Wait. am overlooking the
introduction. identified "Jonn's
Therefore, John
pattern
distraction) </thinks So the final answer
John
Figure 9: Visualization of attention scores for the final answer token in a noisy reasoning context. (a) Full Attention (Retrieval Heads) assigns significant attention weight to the irrelevant distraction text (e.g. "West", "the"), indicating susceptibility to noise. (b) Sparse Heads successfully filter out these distractions. By computing attention only on the propagated critical tokens, the Sparse Heads concentrate their focus solely on the relevant reasoning path, effectively denoising the context.
22

---

Published as a conference paper at ICLR 2026
vssss
u3
overaa atzentin
w 36
131
122 v a
121 79
135
18
00
Figure 10: Overlap rate of top-k attention scores between corresponding heads in adjacent layers. The heatmap illustrates the functional diversity among attention heads. We show the overlap rate (k 5) for the prompt: Please directly output the final answer based on the given question. Ques- tion: There are only two kinds of fruit in box: apples and bananas. All apples are sour, and all bananas are sweet. took a fruit from the box and tasted was sweet. What this fruit? An- swer: and Llama-3 outputs banana.
sampled from range of 1k 20k tokens. During inference, we preprocess the model by reordering the output channels of the Query, Key, and Value projection weights according to the attention head assignments, so to ensure that the retrieval head and the sparse head are grouped into two different continuous clusters. For Grouped Query Attention (GQA) models, we reduce the dimension of the Q heads to match that of the KV heads by applying average pooling, which allows for the calculation of the highest- scoring token set for each kv head. For all experiments, we employ the greedy decoding strategy.
G LIMITATION
Although LycheeDecode demonstrates a significant step towards efficient long-contess LLM infer- ence, acknowledge several limitations that present valuable avenues for future research. Cur- rently, we allocate a fixed budget for each sparse head. However, recent work (Feng et al., 2024) suggests that dynamically allocating the budget among attention heads lead better perfor- mance. Additionally, while we have achieved considerable speedup, our method is not yet integrated with highly optimized inference serving frameworks like vLLM (Kwon et 2023), which is left for future work.
23

---

Published as a conference paper at ICLR 2026
ww
Uwis
for ancentinn
1:4
U: w 18
1266
125 3
125
188
125 s
13 20
Haad Indox
Figure 11: Overlap rate of top-k attention scores between corresponding heads in adjacent layers. The heatmap illustrates the functional diversity among attention heads. We show the overlap rate (k = 5) for the prompt: Please directly output the final answer based on the given question. Ques- tion: you walk 10 meters north from starting point, then 10 meters east, and finally 10 meters west, what direction are you from the original position? Answer:, and Llama-3 outputs north.
u1
Tip Oweer for attenton
144
U3
168 16
11
125 124
125
(8
L50 20
U1
Head Eriles
Figure 12: Overlap rate of top-k attention scores between corresponding heads in adjacent layers. The heatmap illustrates the functional diversity among attention heads. We show the overlap rate (k = 5) for the prompt: Please directly output the final answer based on the given question. Ques- tion: You start facing east. You turn left 90 degrees, then turn right 180 degrees, and finally turn left 90 degrees. What direction are you facing now? Answer: and Llama-3 outputs east.
24

---

Published as a conference paper at ICLR 2026
urw
Actenten Score
u1
u1
120 w 19
i 20
122
135
136 26
130
131
Figure 13: Overlap rate of top-k attention scores between corresponding heads in adjacent layers. The heatmap illustrates the functional diversity among attention heads. We show the overlap rate (k = 5) for the prompt: Please directly output the final answer based on the given question. Ques- tion: If two days ago Monday, what day is tomorrow? Answer; and L Llama-3 outputs Thursday.
25


Reviews:
[{"review_id": "s0FJQlKtWJ", "reviewer": "Reviewer_vCSt", "summary": "This paper proposes LycheeDecode, a sparse attention method that categorizes attention heads into \"retrieval heads\" and \"sparse heads\" to accelerate long-context LLM inference. The retrieval heads perform full attention to identify critical tokens, which are then reused by sparse heads for efficient computation.", "questions": "Please refer to Weaknesses.", "limitations": "", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The paper presents a well-motivated approach. The experimental results demonstrate that LycheeDecode achieves performance comparable to full attention baselines on complex reasoning tasks.", "weaknesses": "1. Recent work has demonstrated that trainable sparse attention can also achieve efficient decoding [1-3]. The paper lacks discussion and empirical comparison with these methods.\n\n[1] Native sparse attention: Hardware-aligned and natively trainable sparse attention\n[2] Minicpm4: Ultra-efficient llms on end devices\n[3] SeerAttention-R: Sparse Attention Adaptation for Long Reasoning \n\n2. While the paper shows kernel-level speedup for different sparse head ratios, there is no corresponding analysis of how varying the proportion of retrieval heads affects model performance. An ablation study examining the performance-efficiency trade-off across different retrieval head budgets would strengthen the paper.\n\n3. The efficiency evaluation focuses on synthetic settings with fixed context lengths. For math reasoning tasks that involve long chain-of-thought generation, where the sequence length grows dynamically during decoding, end-to-end latency measurements would be more convincing.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "dlHRlUdpbh", "reviewer": "Reviewer_ZVhw", "summary": "This paper introduces LycheeDecode, which learns to partition attention heads into a few retrieval heads (full attention with top-k token selection) and many sparse heads (reusing selected tokens). The method uses HardKuma to mitigate train–inference mismatch and employs a custom TileLang block-sparse kernel. On long-context benchmarks (LongBench, RULER) and reasoning tasks (AIME24, OlympiadBench), it achieves competitive accuracy with up to 2.7× speedup at 128k context length.", "questions": "1. Do retrieval heads discovered by LycheeDecode overlap with previous methods' identified retrieval heads? (such as DuoAttention or QR head https://arxiv.org/abs/2506.09944)\n2. What is the training FLOPs overhead vs. DuoAttention or RazorAttention?\n3. Will head roles transfer to long-form summarization or narrative QA?\n4. How does training complexity grow with more layers and longer sequences?\n5. What is the Jaccard overlap between retrieval heads discovered by LycheeDecode and DuoAttention? Is the retrieval-head model-dependent or training method related?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. The paper provides compelling motivation that layer-level token sharing ignores significant functional diversity among heads, with empirical evidence showing highly variable top-k overlap across adjacent layers.\n2. Trainable retrieval and sparse head assignment, HardKuma offers a principled, differentiable relaxation that tends toward binary outcomes without rounding.\n3. The custom hybrid-head kernel yields substantial kernel-level improvements, this strengthens the practicality claim.", "weaknesses": "1. Lack of comparison with training-based long-context inference, DuoAttention is arguably the most directly comparable prior work, and the omission makes it difficult to isolate novelty beyond the use of top-k propagation.\n2. Head role consistency and interpretability are not evaluated, how stable head assignments are across different settings (random setting), whether specialization generalizes to unseen domains.\n3. Training complexity is under-characterized. The method computes both full and sparse attention per head during training, doubling attention paths. FLOPs, KV bandwidth, and convergence cost are not reported. Scaling to deeper models remains unclear.\n4. Propagation across layers may accumulate error.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "fiMplyZrta", "reviewer": "Reviewer_Zofo", "summary": "The paper proposes LycheeDecode, a sparse attention mechanism to accelerate long-context LLM inference by partitioning heads into a small set of \"Retrieval Heads\" that perform full attention and a majority of \"Sparse Heads\" that efficiently reuse the critical tokens identified by the retrieval heads. To learn this specialization, it introduces the HardKuma distribution, which mitigates the train-inference discrepancy common in discrete optimization. Experiments show LycheeDecode achieves up to 2.7x speedup at 128K context lengths while matching or even exceeding the performance of the full-attention baseline.", "questions": "1. Missing Baselines (Duo/Razor): Why are there no end to end comparisons against DuoAttention and RazorAttention? The paper's premise is that its cooperative architecture is a key advantage over DuoAttention's isolated head roles, but this central claim is not experimentally validated.\n\n2. Confounding Variable (Cache Correction): In Table 2, the \"Cache Correction\" (CC) strategy provides a large performance boost. Did you test a \"Full Attention + CC\" baseline? Your current results could be interpreted as the CC method being the main source of the performance gain over the baseline, not LycheeDecode.\n\n3. Evidence for \"Surpassing Full Attention\": Can you provide a specific, qualitative example (e.g., visualizing attention) that proves the full-attention model failed due to \"noise\" while LycheeDecode succeeded by filtering it?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper's primary innovation is the hybrid-head decoding mechanism. This mechanism establishes a head-indexed pipeline, where a Retrieval Head in one layer selects tokens specifically for its corresponding Sparse Head in the next layer to reuse. \n\n2. The use of the HardKuma distribution directly targets a known weakness (train-inference discrepancy) in prior training-based specialization methods, leading to a more stable and direct optimization of the discrete head roles.", "weaknesses": "1. Missing Key Baseline Comparisons: The paper's central claim is that its cooperative head-specialization architecture is superior. However, it fails to provide any end-to-end performance or speed comparisons against the most direct SOTA competitors in the head-specialization sub-field (e.g., DuoAttention, RazorAttention). It only compares against a layer-sharing method (TidalDecode). This makes the SOTA claim unsubstantiated, as we cannot see how it performs against other architectures with the same design philosophy.\n\n2. Weak Evidence: The claim of surpassing the full-attention baseline is remarkable. The paper's only explanation is a vague hypothesis about \"filtering out irrelevant context that may act as noise.\" This claim requires rigorous qualitative proof (e.g., visualizations, case studies), which is absent.\n\n3. Lack of Sensitivity Analysis: The method's performance hinges on two key hyperparameters: the retrieval head budget ($N_{target}$) and the token budget ($k$). The paper provides no ablation studies on how varying these budgets affects the performance/latency trade-off. The $N_{target}$ value was just set to \"match... TidalDecode,\" which is arbitrary and likely not optimal.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "RgHT8OgwV1", "reviewer": "Reviewer_GqdT", "summary": "This paper introduces LycheeDecode, a novel method for accelerating long-context inference by leveraging a hybrid sparse attention mechanism. The core idea is to specialize attention heads into retrieval heads and sparse heads\". To enable end-to-end training and robust head specialization, the authors propose the HardKuma distribution, a reparameterizable method for near-binary head selection that mitigates the train-inference discrepancy common in discrete optimization. Results on several datasets demonstrates effectiveness.", "questions": "Please see the weaknesses.", "limitations": "", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. The paper is easy to follow.\n2. The authors give thorough experiments to validate the proposed methods.", "weaknesses": "1. The proposed method does not appear novel to me. For example, the idea of parameterizing discrete variables has been explored in prior work. Techniques such as Gumbel softmax and STE have been widely used.\n2. Regarding the motivation in Figure 1, I question the significance of showing the overlap rate between corresponding heads in adjacent layers, as attention heads have no inherent ordering, and heads have no direct correspondence across layers.\n3. I am uncertain about the necessity of the HardKuma distribution to reduce the train-inference gap. As I understand it, its parameters are optimized during training, but head types are determined at inference using a fixed threshold—this still introduces a train-inference discrepancy.\n4. The method assigns fixed head types regardless of input, which may be suboptimal. A more robust approach would make head specialization adaptive to the specific input during inference.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "POEbyZ2WtD", "reviewer": "Reviewer_oLtn", "summary": "This paper proposes LycheeDecode, a head-wise sparse attention mechanism with head sampling using HardKuma probability. During training, the HardKuma takes the input of $\\alpha, \\beta$, which are trainable parameters, to generate the probability focused at 0 and 1 for head selection between full attention and sparse attention. The pattern of sparse attention is propagated across layers from the full attention in shallow layers to the sparse attention in deep layers. The training loss of LycheeDecode composes of two parts: the distillation loss (between the sparsified model and the original dense model) and the regularization loss (to limit the number of full attention heads). LycheeDecode is benchmarked on various datasets, including general-purpose long-context evaluation such as LongBench and RULER, and complex reasoning datasets such as AIME and Olympiad-Bench. LycheeDecode is able to outperform the baselines on all tested benchmarks.", "questions": "See above.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "1. The paper is well written. The statement of the problem and the description of the methods, especially the HardKuma part, are clear and easy to follow.\n\n2. The design of LycheeDecode is reasonable, efficient, and effective. In DuoAttention, head selection is conducted between the full attention head and the sliding window head with an attention sink. Therefore, the ability to capture long-context dependencies is reduced, since the sliding window has a relatively weaker ability to memorize context history. However, in LycheeDecode, this problem is well solved by using sparse attention instead of the sliding window. The selection process is performed by the head in the same position in previous layers, so no significant overhead for generating sparse patterns is introduced. This design is reasonable and makes a valuable contribution.\n\n3. Mainstream evaluation benchmark datasets are included in this paper, and LycheeDecode achieves promising results on these benchmarks.", "weaknesses": "1. Since DuoAttention [1] is an important preliminary work for this paper, comparisons between LycheeDecode and DuoAttention should be conducted on long-context benchmarks such as RULER and LongBench.\n\n2. \"By optimizing the distributional parameters of HardKuma during training, our model learns a near-binary selection mechanism directly, thus mitigating the train-inference discrepancy and leading to a more stable and robust head specialization.\" (line 211-213) During inference, a head is identified as a retrieval head if $\\mathbb{E}[z_h] \\geq 0.5$. In this case, since $z_h$ is generated from a probability distribution, there is still a chance that this head will be treated as a non-retrieval head during training. Could the authors provide more explanation on why such a design can eliminate the train–inference discrepancy? If possible, a formal proof would help clarify this point.\n\n---\n\n[1] DuoAttention: Efficient Long-Context LLM Inference with Retrieval and Streaming Heads", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)