PaperlayerView exact revision on arXiv

Source-linked AI summary

LycheeDecode: Accelerating Long-Context LLM Inference via Hybrid-Head Sparse Decoding

Gang Lin, Dongfang Li, Zhuoen Chen, Yukun Shi, Xuhui Chen, Baotian Hu, Min Zhang

arXiv:2602.04541v1cs.CLcs.AI

TL;DR

Long-context decoding is limited by the growing KV cache and by coarse token sharing that neglects attention-head diversity. LycheeDecode uses HardKuma-trained hybrid head roles, with retrieval heads selecting tokens for sparse heads to reuse. It reports generative quality comparable to, and sometimes surpassing, full attention, with up to 2.7× speedup at 128K context.

  • Problem

    Long-context decoding incurs increasing KV-cache memory, I/O, and latency costs, while coarse sharing can overlook the functional diversity of attention heads.

  • Method

    LycheeDecode partitions heads into retrieval heads that identify critical tokens and sparse heads that reuse them, using HardKuma for differentiable near-binary role selection.

  • Results

    Up to 2.7× end-to-end decoding speedup is achieved with the custom hybrid-head block-sparse kernel, while reported generative quality is comparable to or sometimes better than full attention.

  • Takeaways & Limitations

    Fine-grained attention-head specialization provides a supported pathway toward efficient long-context inference while maintaining model performance.

  • Takeaways & Limitations

    The method currently uses fixed per-head sparse budgets, is evaluated only on text-based language models, and lacks integration with optimized serving frameworks such as vLLM.

Abstract

from arXiv · show

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.

1. Introduction

Long-context decoding is bottlenecked by the growing KV cache, while coarse layer-wise token sharing overlooks attention-head diversity. LycheeDecode addresses this with hybrid head roles, HardKuma-based specialization, and a custom sparse kernel.

  • Motivation: The expanding KV cache increases memory use, I/O overhead, and decoding latency as context length grows.Autoregressive decoding repeatedly attends to previous tokens, making long-context deployment and scaling difficult.
  • Approach: LycheeDecode assigns a few retrieval heads to identify important tokens and lets most sparse heads reuse those selections.This fine-grained sharing strategy is designed to preserve diverse attention patterns while reducing computation.
  • Approach: Hard Kumaraswamy selectors address discrete head-role optimization while reducing the train-inference discrepancy.The mechanism supports end-to-end identification of retrieval and sparse heads without relying only on post-training rounding.
  • Results: 2.7× end-to-end decoding speedup is achieved with a custom hybrid-head block-sparse kernel implemented using TileLang.The reported speedup is an end-to-end result for long-context decoding.

2. Related Work

Prior sparse-attention work reduces inference cost through token eviction or selection, layer-wise sharing, and attention-head specialization. LycheeDecode builds on these ideas by enabling retrieval heads to cooperatively select tokens for sparse heads.

  • Sparse attention: Sparse attention methods reduce inference overhead either by evicting less relevant tokens or selecting important tokens while retaining the KV cache.These represent the two main sparse-attention design types described in the related work.
  • Head specialization: Attention-head specialization identifies retrieval heads as important for recalling information and distinguishes them from other head roles.Earlier approaches learn or maintain head roles largely without direct collaboration between specialized heads.
  • LycheeDecode: LycheeDecode differs by having retrieval heads dynamically propagate curated critical tokens for reuse by most sparse heads.This creates a fine-grained cooperative mechanism for sharing contextual information across functionally distinct heads.
  • Cross-layer sharing: Layer-level methods reuse critical-token selections across subsequent layers to exploit cross-layer similarity in attention patterns.Examples include TidalDecode and OmniKV, while LiSA and PoD also leverage related redundancy.

3. Methodology

LycheeDecode assigns retrieval and sparse roles to attention heads, propagates selected critical tokens across layers, and trains the assignment with differentiable HardKuma sampling. Retrieval heads refresh token sets, while sparse heads reuse them to reduce computation and KV-cache loading.

  • Head-level sparse decoding: Retrieval heads perform dense attention over the full sequence to identify the top-k attended tokens for the next layer.The selected token indices are propagated to the corresponding head in the following layer.
  • Head-level sparse decoding: Sparse heads restrict attention to the inherited critical-token set instead of the full sequence.This reduces computation and KV-cache loading, the dominant efficiency source during autoregressive decoding.
  • Head-level sparse decoding: Retrieval heads periodically refresh salient tokens, whereas sparse heads reuse curated subsets across layers.Their division of labor trades off adaptivity and efficiency within one decoding pipeline.
  • Head specialization via HardKuma: Head-role assignment is a discrete optimization problem whose continuous-variable alternatives create a train-inference discrepancy after rounding.LycheeDecode targets this mismatch directly rather than relying on a separately rounded continuous gate.
  • Head specialization via HardKuma: HardKuma produces near-binary, reparameterizable samples, enabling differentiable learning of head roles with more stable inference assignments.Its sampling process transforms a uniform draw through Kumaraswamy sampling, stretching, and hard-sigmoid rectification.
  • Head specialization via HardKuma: During training, each head combines full- and sparse-attention maps using a stochastic HardKuma sample, while inference uses a deterministic role assignment.Gradients from the final loss update the selector parameters; expectations above 0.5 designate retrieval heads, otherwise heads are sparse.

4. Experiments

LycheeDecode is evaluated for long-context understanding, complex reasoning, latency, kernel efficiency, and sparsity strategies. Across these experiments, it combines distillation-based head identification, constrained HardKuma optimization, hybrid-head decoding, and sparse token selection.

  • Training Setup: LycheeDecode’s head-identification training uses teacher–student logit distillation over target tokens, with a shared prompt KV cache.The teacher and student compute logits conditioned on the shared cache, and the distillation loss aligns their outputs.
  • Training Setup: A Lagrangian objective enforces a strict Retrieval Head sparsity budget by penalizing the expected number of active heads.HardKuma selector parameters are optimized against the distillation objective, while the multiplier adapts to constraint violations.
  • Long-Context Understanding: On LongBench, LycheeDecode achieves a 33.07 average score with a 4096-token budget on Llama-3-8B-Instruct-Gradient-1048k, surpassing sparse and full-attention comparisons.On Qwen3-8B, it outperforms TidalDecode with both 1024- and 4096-token budgets and is comparable or slightly superior to SeerAttention-R.
  • Complex Reasoning: On four math-reasoning benchmarks across two distilled DeepSeek-R1 models, LycheeDecode outperforms both TidalDecode and full attention.The experiments cover Gaokao2023En, Minerva, AIME24, and OlympiadBench; Cache Correction further enhances performance.
  • End-to-End Efficiency: At 128K context with a single batch, LycheeDecode achieves up to 2.7× speedup over full attention and is 1.73× faster than TidalDecode.Full-attention latency rises sharply with context length, while LycheeDecode maintains low latency using a fixed 4096-token budget.
  • Kernel Efficiency: At 128K context and batch size 8, the fully sparse 8/8 hybrid-head kernel reaches up to 7x speedup over FlashAttention-2.Higher-sparsity configurations outperform the dense baseline more consistently, especially at larger sequence lengths and batch sizes.

5. Conclusion

The experiments compare attention-head identification methods across Passkey Retrieval and HotpotQA, with HardKuma achieving the best overall performance while showing a limitation on sparse-supervision tasks.

  • The evaluation uses Passkey Retrieval and HotpotQA, with HotpotQA testing multi-hop reasoning over long contexts.
  • HardKuma achieves the best overall performance, outperforming direct optimization and HardConcrete for attention-head identification.The comparison covers Passkey Retrieval and HotpotQA.
  • HardKuma scores slightly lower on HotpotQA, which the authors hypothesize results from short answers producing higher-variance gradient estimates.The authors leave optimization for tasks with sparse supervision to future work.

5 Conclusion

This section presents the Kumaraswamy and HardKuma distributions, LycheeDecode’s hybrid-head decoding procedure, and the workload-pooling strategy used to address heterogeneous head costs.

  • Kumaraswamy distribution: The Kumaraswamy distribution is defined on (0, 1), with shape parameters controlling forms including unimodal, increasing, decreasing, and constant densities.Its PDF and CDF have closed-form expressions.
  • HardKuma distribution: HardKuma adds endpoint probability masses at 0 and 1 while retaining continuous density on (0, 1), enabling differentiable binary masks.Its mixed discrete-continuous structure supports sparse selection.
  • Hybrid-head decoding: LycheeDecode assigns a few retrieval heads to full-context attention and lets sparse heads reuse their selected critical tokens.Retrieval heads select tokens from the KV cache; sparse heads compute attention over the resulting subset.
  • Kernel design: Hybrid-head kernels face workload imbalance because retrieval heads process the full KV cache while sparse heads process only small subsets.Equal resource allocation can leave sparse-head threads idle while full-attention heads determine the critical path.
  • Kernel design: Workload pooling aggregates block computations across full and sparse heads, partitions them into uniform splits, and decouples resource allocation from individual heads.

D Visualization of training process

Training dynamics show that LycheeDecode polarizes head assignments toward retrieval or sparse roles, unlike DuoAttention’s persistent intermediate values.

  • DuoAttention retains many gating values between 0.4 and 0.6 after 1000 steps, creating a train-inference discrepancy when rounded.
  • LycheeDecode’s HardKuma assignments quickly converge toward 0 or 1, corresponding to sparse and retrieval heads.This directly reduces the consistency gap between training and inference.
  • Kuma PDFs evolve from initially uniform distributions toward boundary-concentrated forms for representative retrieval and sparse heads.Retrieval-head mass shifts right, whereas sparse-head mass collapses left.

E.1 RULER benchmark

On RULER, LycheeDecode remains close to full attention at shorter contexts, while performance declines slightly as context length increases under a fixed 4096-token budget.

  • RULER evaluates long-context comprehension with tasks extending beyond simple retrieval, including multihop tracing and aggregation.
  • At 8k context length, LycheeDecode achieves an average score of 62.79 versus 63.30 for full attention.The comparison uses a fixed LycheeDecode budget of 4096 tokens.
  • LycheeDecode’s performance decreases slightly as context length increases.The authors describe this as an acceptable trade-off for using a fixed, smaller token budget.

E.2 Detailed results of different sparse methods

This section details LongBench evaluations of LycheeDecode under multiple sparse token-selection strategies and reports their associated sparsity levels.

  • LongBench performance is compared across Top-k, Top-p, Threshold, and Ratio token-selection strategies with varying parameters.Table 5 presents the performance comparison, while Table 6 quantifies the percentage of critical tokens selected.
  • The experiments quantify how many critical tokens LycheeDecode selects under different LongBench settings.These sparsity measurements complement the performance results for the evaluated selection methods.
  • The section provides additional examples of attention-head behavior by comparing top-k crucial-token overlap across corresponding heads in adjacent layers.The analysis uses prompts requiring simple logical reasoning and computes attention scores for the final answer token.

E.4 Ablation Study

The ablation study examines performance–efficiency trade-offs across token budgets and retrieval-head ratios, alongside attention behavior under noisy reasoning contexts.

  • Performance and efficiency trade-offs: Increasing the token budget from 1024 to 4096 consistently improves LongBench average scores but decreases decoding speedup.The larger budget retains more context while increasing computation.
  • Performance and efficiency trade-offs: Smaller retrieval-head ratios provide the highest speedup by minimizing computational overhead.The speedup is evaluated using end-to-end decoding speedup measured by Time Per Output Token relative to Full Attention.
  • Performance and efficiency trade-offs: With token budgets of 2048 and 4096, the 25.0% retrieval-head configuration outperforms the 50.0% configuration.The authors hypothesize that excessive retrieval heads may introduce irrelevant context, whereas balanced configurations help Sparse Heads focus on critical information.
  • Performance and efficiency trade-offs: Figure 9 summarizes the trade-off between model performance and inference efficiency under different token and retrieval-head budgets.
  • Attention visualization: In noisy reasoning contexts, Retrieval Heads assign attention to irrelevant distractors, whereas Sparse Heads eliminate this noise.Sparse Heads compute attention only over propagated critical tokens, concentrating on the relevant reasoning path.

F Implementation details

The implementation details specify training settings for HardKuma and preprocessing steps used during inference.

  • The training learning rate is set to 0.01, and the HardKuma stretching interval is set to (−0.1, 1.1).
  • Passkey Retrieval inserts ten 32-word passkeys into BookSum with prompt lengths sampled from 1k to 10k tokens.
  • HotpotQA filters questions answerable without the provided context and samples prompt lengths from 1k to 20k tokens.
  • During inference, Query, Key, and Value projection weights are reordered according to attention-head assignments.

G Limitation & Future work

The paper identifies fixed per-head token budgets, text-only evaluation, and missing integration with optimized serving frameworks as limitations and future-work directions.

  • LycheeDecode currently uses a fixed budget for each Sparse Head, while dynamic allocation may yield better performance.
  • The experiments are confined to text-based language models, leaving multimodal validation for future work.The authors propose extending the mechanism to Multimodal Large Language Models such as Uni-MoE.
  • LycheeDecode is not yet integrated with highly optimized inference-serving frameworks such as vLLM.Integration with such frameworks is left for future work.
Loading 2602.04541v1…