FlashForge: Ultra-Efficient Prefix-Aware Attention for LLM Decoding

• URL: http://arxiv.org/abs/2505.17694v1
• Date: Fri, 23 May 2025 10:03:28 GMT
• Title: FlashForge: Ultra-Efficient Prefix-Aware Attention for LLM Decoding
• Authors: Zhibin Wang, Rui Ning, Chao Fang, Zhonghui Zhang, Xi Lin, Shaobo Ma, Mo Zhou, Xue Li, Zhongfeng Wang, Chengying Huan, Rong Gu, Kun Yang, Guihai Chen, Sheng Zhong, Chen Tian,
• Abstract summary: Prefix-sharing among multiple prompts presents opportunities to combine the operations of the shared prefix.<n>Decoding is a memory-intensive process requiring heavy memory access on the key-value ( KV) cache of the prefixes.<n>We propose a dedicated attention kernel to combine the memory access of shared KV cache in the decode stage, namely FlashForge.
• Score: 44.47821531299985
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: Prefix-sharing among multiple prompts presents opportunities to combine the operations of the shared prefix, while attention computation in the decode stage, which becomes a critical bottleneck with increasing context lengths, is a memory-intensive process requiring heavy memory access on the key-value (KV) cache of the prefixes. Therefore, in this paper, we explore the potential of prefix-sharing in the attention computation of the decode stage. However, the tree structure of the prefix-sharing mechanism presents significant challenges for attention computation in efficiently processing shared KV cache access patterns while managing complex dependencies and balancing irregular workloads. To address the above challenges, we propose a dedicated attention kernel to combine the memory access of shared prefixes in the decoding stage, namely FlashForge. FlashForge delivers two key innovations: a novel shared-prefix attention kernel that optimizes memory hierarchy and exploits both intra-block and inter-block parallelism, and a comprehensive workload balancing mechanism that efficiently estimates cost, divides tasks, and schedules execution. Experimental results show that FlashForge achieves an average 1.9x speedup and 120.9x memory access reduction compared to the state-of-the-art FlashDecoding kernel regarding attention computation in the decode stage and 3.8x end-to-end time per output token compared to the vLLM.

Related papers

• Sparse-dLLM: Accelerating Diffusion LLMs with Dynamic Cache Eviction [58.044803442346115]
  Diffusion Large Language Models (dLLMs) enable breakthroughs in reasoning and parallel decoding but suffer from prohibitive computational complexity and memory overhead during inference.<n>We propose Sparse-dLLM, the first training-free framework integrating dynamic cache eviction with sparse attention via delayed bidirectional sparse caching.
  arXiv  Detail & Related papers  (2025-08-04T16:14:03Z)
• Sparse Attention Remapping with Clustering for Efficient LLM Decoding on PIM [7.651654889371008]
  Transformer-based models are the foundation of modern machine learning, but their execution places significant pressure on memory systems.<n> processing-in-memory (PIM) architectures are a promising solution, offering high internal bandwidth and compute parallelism near memory.<n>Current PIM designs are primarily optimized for dense attention and struggle with the dynamic, irregular access patterns introduced by modern KV cache sparsity techniques.
  arXiv  Detail & Related papers  (2025-05-09T04:17:05Z)
• Injecting Adrenaline into LLM Serving: Boosting Resource Utilization and Throughput via Attention Disaggregation [23.130886760027586]
  In large language model (LLM) serving systems, executing each request consists of two phases: the compute-intensive prefill phase and the memory-intensive decoding phase.<n>This paper proposes Adrenaline, an attention disaggregation and offloading mechanism designed to enhance resource utilization and performance.<n> Experimental results show that Adrenaline achieves 2.28x higher memory capacity and 2.07x better memory bandwidth utilization in prefill instances.
  arXiv  Detail & Related papers  (2025-03-26T13:48:35Z)
• TopV: Compatible Token Pruning with Inference Time Optimization for Fast and Low-Memory Multimodal Vision Language Model [56.43860351559185]
  We introduce textbfTopV, a compatible textbfTOken textbfPruning with inference Time Optimization for fast and low-memory textbfVLM.<n>Our framework incorporates a visual-aware cost function to measure the importance of each source visual token, enabling effective pruning of low-importance tokens.
  arXiv  Detail & Related papers  (2025-03-24T01:47:26Z)
• A Universal Framework for Compressing Embeddings in CTR Prediction [68.27582084015044]
  We introduce a Model-agnostic Embedding Compression (MEC) framework that compresses embedding tables by quantizing pre-trained embeddings.<n>Our approach consists of two stages: first, we apply popularity-weighted regularization to balance code distribution between high- and low-frequency features.<n> Experiments on three datasets reveal that our method reduces memory usage by over 50x while maintaining or improving recommendation performance.
  arXiv  Detail & Related papers  (2025-02-21T10:12:34Z)
• APB: Accelerating Distributed Long-Context Inference by Passing Compressed Context Blocks across GPUs [81.5049387116454]
  We introduce APB, an efficient long-context inference framework.<n>APB uses multi-host approximate attention to enhance prefill speed.<n>APB achieves speeds of up to 9.2x, 4.2x, and 1.6x compared with FlashAttn, RingAttn, and StarAttn, respectively.
  arXiv  Detail & Related papers  (2025-02-17T17:59:56Z)
• S2-Attention: Hardware-Aware Context Sharding Among Attention Heads [49.1454481007861]
  Sparse attention selectively attends to a subset of tokens in the context.<n>It remains unclear whether sparse attention can maintain the model's quality at a scale of today's large language models.<n>This paper presents Sparsely-Sharded(S2) Attention, a Triton library that provides kernel optimization for sparse attention customizable at both per-head and per-context-range levels.
  arXiv  Detail & Related papers  (2024-07-25T00:27:07Z)
• DeFT: Decoding with Flash Tree-attention for Efficient Tree-structured LLM Inference [22.684773338989007]
  Large language models (LLMs) are increasingly employed for complex tasks that process multiple generation calls in a tree structure with shared prefixes of tokens.<n>Existing inference systems for tree-based applications are inefficient due to improper partitioning of queries and KV cache during attention calculation.<n>We propose DeFT, a hardware-efficient attention algorithm with prefix-aware and load-balanced KV cache partitions.
  arXiv  Detail & Related papers  (2024-03-30T04:34:54Z)
• Bifurcated Attention: Accelerating Massively Parallel Decoding with Shared Prefixes in LLMs [39.16152482491236]
  Bifurcated attention is a method designed to enhance language model inference in shared-context batch decoding scenarios. Our approach addresses the challenge of redundant memory IO costs, a critical factor contributing to latency in high batch sizes and extended context lengths.
  arXiv  Detail & Related papers  (2024-03-13T16:30:57Z)
• ChunkAttention: Efficient Self-Attention with Prefix-Aware KV Cache and Two-Phase Partition [3.659659889927316]
  ChunkAttention is a prefix-aware self-attention module for large language models. It can detect matching prompt prefixes across multiple requests and share their key/value tensors in memory at runtime. Experiments show that ChunkAttention can speed up the self-attention kernel by 3.2-4.8$times$ compared to the state-of-the-art implementation.
  arXiv  Detail & Related papers  (2024-02-23T09:29:19Z)
• SubGen: Token Generation in Sublinear Time and Memory [48.35076900702408]
  Large language models (LLMs) have extensive memory requirements for token generation. In this work, we focus on developing an efficient compression technique for the KV cache. We have devised a novel caching method with sublinear complexity, employing online clustering on key tokens and online $ell$ sampling on values. Not only does this algorithm ensure a sublinear memory footprint and sublinear time complexity, but we also establish a tight error bound for our approach.
  arXiv  Detail & Related papers  (2024-02-08T22:17:40Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.