APB: Accelerating Distributed Long-Context Inference by Passing Compressed Context Blocks across GPUs
Authors: Yuxiang Huang, Mingye Li, Xu Han, Chaojun Xiao, Weilin Zhao, Ao Sun, Jie Zhou, Hao Zhou, Zhiyuan Liu, Maosong Sun
Accepted at: acl25

APB: Accelerating Distributed Long-Context Inference by Passing Compressed Context Blocks across GPUs

Yuxiang Huang1*, Mingye Li2†‡, Xu Han1‡, Chaojun Xiao1‡, Weilin Zhao1, Ao Sun3, Hao Zhou4, Jie Zhou4, Zhiyuan Liu1, Maosong Sun1
1NLP Group, DCST, IAI, BNRIST, Tsinghua University, Beijing, China.
2Department of CS&T, Central South University, Changsha, China.
3BUPT, Beijing, China. 4Pattern Recognition Center, WeChat AI, Tencent Inc.
huang-yx21@mails.tsinghua.edu.cn, lmy2004@csu.edu.cn, han-xu@tsinghua.edu.cn, xiaocj20@mails.tsinghua.edu.cn

Abstract

While long-context inference is crucial for advancing large language model (LLM) applications, its prefill speed remains a significant bottleneck. Current approaches, including sequence parallelism strategies and compute reduction through approximate attention mechanisms, still fall short of delivering optimal inference efficiency. This hinders scaling the inputs to longer sequences and processing long-context queries in a timely manner. To address this, we introduce APB, an efficient long-context inference framework that leverages multi-host approximate attention to enhance prefill speed by reducing compute and enhancing parallelism simultaneously. APB introduces a communication mechanism for essential key-value pairs within a sequence parallelism framework, enabling a faster inference speed while maintaining task performance. We implement APB by incorporating a tailored FLASHATTN kernel alongside optimized distribution strategies, supporting diverse models and parallelism configurations. APB achieves speedups of up to 9.2×, 4.2×, and 1.6× compared with FLASHATTN, RINGATTN, and STARATTN, respectively, without any observable task performance degradation. We provide the implementation and experiment code of APB in https://github.com/thunlp/APB.

1 Introduction

Large language models (LLMs) (OpenAI, 2024; Anthropic, 2024; DeepSeek-AI, 2024) have demonstrated unprecedented proficiencies, pushing the boundaries of artificial intelligence research and practical applications. Recent advancements are not only transforming usage paradigms but also empowering intelligent systems such as LLM-based agents (Li, 2025; Qin et al., 2024; Zhao et al., 2024), robotics (Zeng et al., 2023; Kim et al., 2024), and prompting methodologies (Chu et al., 2023; Sahoo et al., 2024). These systems often rely on extended context inference. To address the growing demand for longer inputs, contemporary foundation models have been increasingly designed to support extended context lengths. For instance, Llama-3.1 (Dubey et al., 2024) supports up to 128K tokens, Claude-3.5 (Anthropic, 2024) extends input capacity to 200K tokens, and MiniMax-01 (Li et al., 2025) can even process input sequences up to 4M tokens.

As context lengths grow, the quadratic computational cost of attention makes single-GPU inference both infeasible and inefficient for LLMs. To address this, various optimizations aim to enhance parallelism or reduce compute. Sequence parallelism (Li et al., 2023), which aims to enhance parallelism, partitions the sequence across devices (termed as hosts) and significantly improves the prefill speed, especially for extremely long inputs (Figure 1). However, the overall computations remain unchanged to ensure the accuracy of the attention results. On the other hand, approximate attention mechanisms (Zhang et al., 2024e; Li et al., 2024b; Jiang et al., 2024), which compute only those elements selected from the attention matrix, accelerate inference by reducing compute but face scalabil-

![The prefill speed of methods with and without sequence parallelism when processing different input lengths. "SP" indicates sequence parallelism. "x" represents that the setting triggers out-of-memory error.](page_1012_370_377_246.png)

Figure 1: The prefill speed of methods with and without sequence parallelism when processing different input lengths. “SP” indicates sequence parallelism. “x” represents that the setting triggers out-of-memory error.
ity challenges and performance degradation when processing longer inputs. To this end, designing an approximate attention mechanism that fits in sequence parallelism frameworks offers a promising way to further enhance efficiency, particularly in accelerating long-context prefill. However, *designing such systems demands system and algorithm optimizations to address the key challenges*.

**Challenge 1:** *Localized Attention Pruning.* Existing widely-used approximate attention mechanisms, such as H2O (Zhang et al., 2024e) and SNAPKV (Li et al., 2024b), typically depend on full sequence information, such as the attention scores computed over the entire sequence, to prune the redundant compute of the attention scores. This requirement directly conflicts with the distributed architecture of sequence parallelism, where individual hosts only maintain partial context visibility without heavy host-to-host communication.

**Challenge 2:** *Multi-host Scalability.* Traditional sequence parallelism approaches face inherent scalability limitations due to model-architectural constraints and performance degradation risks. While sequence parallelism with attention head splitting for computation (Jacobs et al., 2023) offers substantial throughput improvements, its scalability remains fundamentally bounded by the fixed number of attention heads. Existing solutions that simply combine approximate attention with sequence parallelism, such as STARATTN (Acharya et al., 2024), suffer from progressive performance degradation when the number of hosts increases, as these solutions merely extend Xiao et al. (2024b) to multiple hosts with a large proportion of invisible context.

To address these challenges, we propose APB, a distributed inference framework designed to leverage approximate attention to reduce redundant compute and communication overhead. For *localized attention pruning*, APB introduces a local key-value (KV) cache compression technique that operates independently on each host, eliminating the need for a global sequence view to prune redundant attention compute. For *multi-host scalability*, APB ensures that each host processes all attention heads within its local context and selectively shares compressed critical context across hosts. This design enables APB to maintain stable model performance even as the number of hosts scales up.

We implement APB using a customized FLASH-ATTN (Dao, 2024) kernel and an optimized distribution framework, enabling efficient scaling across diverse sequence lengths and multiple hosts. Comprehensive evaluations demonstrate that APB achieves an excellent trade-off between inference speed and model performance across a variety of tasks. Additionally, APB is compatible with various model sizes and distribution settings, making it a robust solution for scalable distributed inference. APB achieves speedups of up to \(9.2 \times\), \(4.2 \times\), and \(1.6 \times\) compared with FLASH-ATTN, RINGATtn (Li et al., 2023), and STARATTN, respectively, without any observable performance degradation.

2 **Related Works**

Existing approaches to the efficient long-context inference of LLMs focus on system and algorithm optimizations. Please refer to the surveys (Zhao et al., 2023; Wan et al., 2024) for more about LLMs.

**System Optimizations.** The efficiency of long-context inference can be enhanced by leveraging hardware architectures and preserving accurate LLM computations. Hardware-aware algorithms, such as FLASH-ATTN (Dao et al., 2022; Dao, 2024; Shah et al., 2024), utilize matrix tiling to optimize GPU memory usage and boost inference speed. Additionally, efficient parallelism methods support longer sequences with a higher processing speed. RINGATtn (Li et al., 2023) splits long sequences across multiple hosts using ring-style communication to preserve accurate attention computation, while ULYSSES (Jacobs et al., 2023) distributes attention heads across hosts to reduce communication overhead. Other parallelism strategies (Narayanan et al., 2021; Huang et al., 2019; Rajbhandari et al., 2020; Sun et al., 2025; DeepSeek-AI, 2024) enable longer sequences on larger models, and the mixture of various distribution strategies (Ren et al., 2021; Sun et al., 2024a) can further enhance long-context efficiency. However, most parallelism approaches are tailored for training and lack optimization for inference. Offloading-based methods (Xiao et al., 2024a; Sun et al., 2024b; Lee et al., 2024; Chen et al., 2025) leverage hierarchical memory systems to reduce hardware requirements by storing redundant KV cache in CPU memory and recalling only a small part to GPU memory.

**Algorithm Optimizations.** The burden of long-context inference can also be mitigated through algorithm optimizations. KV cache-centric optimization reduces the size of the KV cache, enabling faster decoding speed and lower GPU memory usage. Cache eviction methods (Zhang et al., 2024e; Li et al., 2024b; Yao et al., 2024; Huang et al., 2024) discard irrelevant or unimportant KV cache, allevi-
ating the memory bottlenecks. Quantization techniques (Liu et al., 2024; Hooper et al., 2024; Zhang et al., 2024a; He et al., 2024) decrease the memory footprint by utilizing low-bit representations to store the KV cache. Merging methods (Cai et al., 2024; Wang et al., 2024; Zhang et al., 2024c,d) consolidate redundant KV cache units across sequences or layers to alleviate storage overheads. Sparse mechanisms (Zaheer et al., 2020; Beltagy et al., 2020; Lou et al., 2024) decrease inference latency by reducing the attention computational load. Specifically, approaches like MINFERENCE (Jiang et al., 2024) and FASTGEN (Ge et al., 2024) assign specific patterns to attention heads, accelerating inference by computing only elements selected from the attention matrix. Moreover, algorithm optimizations can complement system optimizations. STARATTN (Acharya et al., 2024) and APE (Yang et al., 2025) linearize attention complexity by dividing context into parallelized blocks, but they struggle with tasks that require inter-context dependencies. More details on these algorithm optimizations are elaborated in Li et al. (2024a) and Shi et al. (2024). Apart from KV cache optimizations on Transformer-based (Vaswani, 2017) LLMs, approaches altering backbone architectures can also enhance long-context inference efficiency. Typically, RNN-based models (Gu and Dao, 2024; Dao and Gu, 2024) and hybrid models (Lieber et al., 2025; Dong et al., 2024; Li et al., 2025) reduce the computation complexity of long-context inference.

3 Methodology

3.1 Preliminaries

We primarily introduce the Transformer architecture and the prefill-decoding pattern.

Transformers. Transformer-based LLMs take a sequence \( \{t_1, t_2, \cdots, t_n\} \) as input. We denote the hidden state of the token \( t_k \) at the \( i \)-th layer as \( \mathbf{H}^{(i)}[k] \), where \( \mathbf{H}^{(0)} \) represents the embeddings of the input sequence. Assuming the model has \( L \) layers, each layer consists of an attention block and a feedforward network (FFN) block. For each attention head in an attention block, the attention score at layer \( i \in \{1, 2, \cdots, L\} \) is calculated as follows, where \( \mathbf{M} \) represents the attention mask.

\[
\mathbf{A}^{(i)} = \operatorname{softmax}\left(\mathbf{M} \odot \frac{\mathbf{Q}^{(i)} \mathbf{K}^{(i)\top}}{\sqrt{d_m}}\right) \cdot \mathbf{V}^{(i)},
\]

where \( \mathbf{Q}^{(i)} = \mathbf{H}^{(i-1)} \mathbf{W}_Q^{(i)}, \mathbf{K}^{(i)} = \mathbf{H}^{(i-1)} \mathbf{W}_K^{(i)}, \) and \( \mathbf{V}^{(i)} = \mathbf{H}^{(i-1)} \mathbf{W}_V^{(i)} \). The hidden size of each head is denoted by \( d_m \). Then, the attention scores are passed through the FFN: \( \mathbf{H}^{(i)} = \mathrm{FFN}(\mathbf{A}^{(i)}) \).

Prefill and Decoding. The procedure of an LLM \( \mathcal{M} \) processing an input can be divided into two stages: prefill and decoding. During the prefill stage, the model computes the KV cache for the input sequence \( \{t_1, t_2, \cdots, t_n\} \) and predicts the first answer token \( t_{n+1} \). This process can be formalized as \( (t_{n+1}; \mathbf{K}[1:n], \mathbf{V}[1:n]) = \mathcal{M}(t_1, \cdots, t_n) \). After the prefill stage, the model decodes one token at a time using the KV cache of previous tokens, which can be written as \( (t_{k+1}; \mathbf{K}[k], \mathbf{V}[k]) = \mathcal{M}(t_k; \mathbf{K}[1:k-1], \mathbf{V}[1:k-1]) \). The complexities of the prefill stage and a single step of decoding are \( \mathcal{O}(n^2) \) and \( \mathcal{O}(n) \), respectively.

3.2 Framework

We adopt sequence parallelism with approximate attention to accelerate the prefill of long-context inputs. The framework of APB is illustrated in Figure 2. The inference of APB can be split into 4 stages: context splitting, block compression, communication, and computation.

Context Splitting. Given a long-context input \( t = \{t_1, t_2, \cdots, t_n\} \), we first split the input sequence into a document \( d = \{d_1, \cdots, d_{l_d}\} \) and a query \( q = \{q_1, \cdots, q_{l_q}\} \), i.e., \( t = \{d, q\} \) and \( n = l_d + l_q \). Since \( l_d \gg l_q \), we primarily focus on optimizing the prefill speed of the document. We apply accurate attention with online softmax for processing the query and decoding the answer. The document is evenly split across all hosts, with an anchor block prepended at the start of each local context block on every host to preserve its visibility to the document’s initial part.

Block Compression. To reduce the attention compute and communication overhead, we first perform KV cache compression using a compressor \( \mathcal{C} \) on each host to shorten the block length and only retain the essential KV cache units for inter-host communication in the next step. These compressed blocks are crucial for capturing long-distance semantic dependency and maintaining task performance by reducing attention errors.

Communication. We design a specialized communication pattern to gather the compressed context blocks on each host for the awareness of previous essential context. After the communication, we form the passing blocks to encapsulate the essential KV cache units passed from the previous hosts.

Computation. Finally, each host obtains a context layout containing the local context block, an
Figure 2: The framework of APB. The input document d is split into blocks \( \mathbf{B}_1, \mathbf{B}_2, \mathbf{B}_3 \) and distributed across 3 hosts. The anchor block is denoted as “A”, the passing block as “P”, and the compressor as “C”. Each block is first prepended with an anchor block. When calculating attention, the context block B is compressed into \( \mathbf{B}^C \) using the compressor C. Subsequently, the passing block is constructed after an AllGather communication. Finally, attention is performed using a modified attention mask. Passing blocks are discarded after the attention computation.

anchor block, and a passing block. We perform attention using a modified attention mask, executed with a specialized FLASHATTN (Dao, 2024) kernel.

Overall, APB is designed to communicate and compute only the most critical context by aggressively compressing the KV cache within a sequence parallelism framework. It introduces an algorithm-aware, distributed system-level optimization specifically tailored to its unique communication pattern and approximate attention mechanism.

3.3 Context Splitting

At the beginning of APB inference, we first split the document d across H hosts, with each host holding a partial sequence with length \( l_b = l_d / H \), denoted as blocks \( \mathbf{B}_1, \mathbf{B}_2, \cdots, \mathbf{B}_H \). Each block is prepended with an anchor block at the front. Inspired by STARATTN (Acharya et al., 2024), we adopt the anchor block A containing the first \( l_a \) tokens \( d_1, \cdots, d_{l_a} \) of the input sequence. We also embed the query q in the front of the anchor block, to provide more query-related information for the compressors to identify the essential cache units. Therefore, the anchor block can be formally written as \( \mathbf{A} = \{q_1, \cdots, q_{l_q}, d_1, \cdots, d_{l_a}\} \). The positions assigned to the tokens in anchor blocks are the starting positions, i.e. \( 0, 1, \cdots, l_q + l_a - 1 \). Notably, we apply much smaller anchor blocks compared to STARATTN, where \( l_a = \frac{1}{4} l_b \) or \( l_a = \frac{1}{8} l_b \) in APB but \( l_a = l_b \) in STARATTN. Formally, each host except the first one holds a sequence of \( \{\mathbf{A}, \mathbf{B}_h\} = \{q_1, \cdots, q_{l_q}, d_1, \cdots, d_{l_a};\ d_{(h-1)l_b+1}, \cdots, d_{hl_b}\} \). The first host only holds \( \mathbf{B}_1 \) without anchor block.

3.4 Block Compression

Before starting the attention computation on the h-th host, we compress the KV cache of \( \mathbf{B}_h \) into \( \mathbf{B}_h^C \) using the compressors C. As calculating full attention scores or full KV cache on a single host is infeasible, existing methods like H2O (Zhang et al., 2024e) and SNAPKV (Li et al., 2024b) are incompatible in this scenario. The need to identify important KV cache units without a global view aligns with the design of LOCRET (Huang et al., 2024), where importance scores are based on each token’s query and KV cache. We implement the compressors C as LOCRET’s retaining heads \( \mathcal{R} \), which are small MLPs trained on additional long-context SFT data. These MLPs take [Q, K, V] as input and output importance scores for the corresponding cache units, reflecting the influence towards the calculation of subsequent KV cache. Formally, the compressed context block on host h with passing length \( l_p \) is defined as \( \mathbf{B}_h^C = (\mathbf{K}_h^C, \mathbf{V}_h^C) = (\{k_h[i_1], \cdots, k_h[i_{l_p}]\}, \{v_h[i_1], \cdots, v_h[i_{l_p}]\}) \), where \( \{s_{i_1}, \cdots, s_{i_{l_p}}\} = \mathrm{Top}_{l_p}(s_1, \cdots, s_{l_b}) \) and \( s_1, \cdots, s_{l_b} = \mathcal{R}([\mathbf{Q}_h, \mathbf{K}_h, \mathbf{V}_h]) \).

3.5 Communication

To acquire all the compressed context blocks sent from each host after block compression, we apply an AllGather communication on the compressed KV cache across all the hosts by \( (\mathbf{K}_{[1:H]}^C, \mathbf{V}_{[1:H]}^C) = \mathrm{AllGather}(\mathbf{K}_h^C, \mathbf{V}_h^C) \). Subsequently, the passing block is constructed by concatenating all compressed context blocks from
Figure 3: The inference speed and model performance of APB and all the baselines. The top-right direction represents the optimal tradeoff between speed and performance. APB achieves the best tradeoff of the two metrics. FLASHATTN, RINGATTN, and ULYSSES share the same performance as they are all FULLATTN methods.

previous hosts, i.e., \( \mathbf{P}_h = (\mathbf{K}_p^C, \mathbf{V}_p^C) = (\mathbf{K}_{[1:h-1]}^C, \mathbf{V}_{[1:h-1]}^C) \). We simply ignore the compressed context blocks sent by subsequent hosts.

3.6 Computation

After the communication and construction of the passing blocks, each host acquires the context layout required for attention computation. We place the passing block between the anchor block and the local context block on each host to perform the attention as Equation 2. The passing blocks are discarded after the attention calculation and do not participate in the FFN calculation. \( M' \) is the modified attention mask. We implement the attention calculation in a modified FLASHATTN (Dao, 2024) kernel with only the attention mask changed.

We present the complete inference procedure of APB in Algorithm 1. Running approximate attention of the document \( d \) is detailed in Algorithm 2, while the accurate attention with online softmax is presented in Algorithm 3.

\[
\begin{align*}
\mathbf{Q}^{(i)} &= [\mathbf{Q}_a^{(i)}, \mathbf{Q}_h^{(i)}], \\
\mathbf{K}^{(i)} &= [\mathbf{K}_a^{(i)}, \mathbf{K}_p^C, \mathbf{K}_h^{(i)}], \\
\mathbf{V}^{(i)} &= [\mathbf{V}_a^{(i)}, \mathbf{V}_p^C, \mathbf{V}_h^{(i)}], \\
\left[ \mathbf{A}_a^{(i)}, \mathbf{A}_h^{(i)} \right] &= \operatorname{softmax}\left( M' \odot \frac{\mathbf{Q}^{(i)} \mathbf{K}^{(i)\top}}{\sqrt{d_m}} \right) \cdot \mathbf{V}^{(i)}, \\
[\mathbf{H}_a^{(i)}, \mathbf{H}_h^{(i)}] &= \operatorname{FFN}\left( [\mathbf{A}_a^{(i)}, \mathbf{A}_h^{(i)}] \right).
\end{align*}
\]

4 Experiments

In this section, we present our benchmark results to evaluate APB, aiming at addressing the following research questions: (Q1) Can APB obtain higher end-to-end prefill speed without significant performance degradation? (Q2) How effective and efficient is APB on various context length? (Q3) Is the design of each component in APB effective?

4.1 Experimental Setup

In this section, we briefly introduce experimental settings, and more details are in Appendix B.

**Benchmarks.** We evaluate APB against selected baselines on two long-context evaluation benchmarks: \( \infty \)Bench (Zhang et al., 2024b) and RULER (Hsieh et al., 2024). \( \infty \)Bench is a benchmark with an average context length exceeding 100K tokens, encompassing a mix of both synthetic and real-world tasks. Following MINFERENCE (Jiang et al., 2024), we utilize 10 tasks, excluding Math.Calc and Code.Run due to their difficulty, which leads to failures for all methods. RULER is a benchmark featuring a variety of synthetic tasks with a controllable context length. We evaluate all 13 tasks of the benchmark. We provide the mapping of task abbreviations to the original task name in Appendix D. Additional experiments and results are placed in Appendix E.

**Models.** For the end-to-end benchmark (Q1), we test APB with the baselines on the instruct versions of Llama-3.1-8B, Qwen-2.5-14B, and Yi-34B-200K. When benchmarking across various input lengths (Q2), we utilize Llama-3-8B-1M\textsuperscript{1}, which supports up to 1M input tokens.

**Metrics.** When evaluating performance on the benchmark, we directly adopt the original performance metrics for each task. For speed measurement, we define inference speed as the total number

\footnotetext{1https://huggingface.co/gradientai/Llama-3-8B-Instruct-Gradient-1048k}
<table>
  <tr>
    <th>Method</th>
    <th>R.PassKey</th>
    <th>R.Number</th>
    <th>R.KV</th>
    <th>E.Sum</th>
    <th>E.QA</th>
    <th>E.MC</th>
    <th>E.Dia</th>
    <th>Z.QA</th>
    <th>C.Debug</th>
    <th>M.Find</th>
    <th>Avg.</th>
  </tr>
  <tr>
    <td colspan="12">Llama-3.1-BB-instruct</td>
  </tr>
  <tr>
    <td>FULLATTN</td>
    <td>100.00</td>
    <td>99.49</td>
    <td>51.00</td>
    <td>30.59</td>
    <td>29.04</td>
    <td>63.76</td>
    <td>11.00</td>
    <td>36.18</td>
    <td>24.62</td>
    <td>28.82</td>
    <td>47.45</td>
  </tr>
  <tr>
    <td>MINERENCE</td>
    <td>98.47</td>
    <td>98.81</td>
    <td>17.40</td>
    <td>30.06</td>
    <td>26.42</td>
    <td>55.46</td>
    <td>12.50</td>
    <td>36.51</td>
    <td>28.17</td>
    <td><b>32.29</b></td>
    <td>43.61</td>
  </tr>
  <tr>
    <td>STARATTN</td>
    <td><b>100.00</b></td>
    <td><b>98.98</b></td>
    <td>40.60</td>
    <td>30.55</td>
    <td><b>30.66</b></td>
    <td>61.57</td>
    <td><b>15.00</b></td>
    <td>36.26</td>
    <td>26.65</td>
    <td>24.57</td>
    <td>46.48</td>
  </tr>
  <tr>
    <td>APB</td>
    <td><b>100.00</b></td>
    <td><b>98.81</b></td>
    <td><b>81.8</b></td>
    <td><b>30.63</b></td>
    <td>29.25</td>
    <td><b>63.32</b></td>
    <td>11.00</td>
    <td><b>36.99</b></td>
    <td><b>30.46</b></td>
    <td>26.86</td>
    <td><b>50.91</b></td>
  </tr>
  <tr>
    <td colspan="12">Qwen-2.5-14B-instruct</td>
  </tr>
  <tr>
    <td>FULLATTN</td>
    <td>100.00</td>
    <td>100.00</td>
    <td>17.80</td>
    <td>27.80</td>
    <td>10.40</td>
    <td>52.84</td>
    <td>28.00</td>
    <td>10.21</td>
    <td>38.07</td>
    <td>42.57</td>
    <td>42.68</td>
  </tr>
  <tr>
    <td>MINERENCE</td>
    <td><b>100.00</b></td>
    <td><b>100.00</b></td>
    <td>5.20</td>
    <td>25.63</td>
    <td>12.04</td>
    <td>61.57</td>
    <td>16.50</td>
    <td>11.06</td>
    <td>41.62</td>
    <td><b>48.57</b></td>
    <td>42.22</td>
  </tr>
  <tr>
    <td>STARATTN</td>
    <td><b>100.00</b></td>
    <td><b>100.00</b></td>
    <td>18.00</td>
    <td><b>27.48</b></td>
    <td>12.44</td>
    <td>62.88</td>
    <td>23.50</td>
    <td>11.03</td>
    <td><b>43.40</b></td>
    <td>37.71</td>
    <td>43.69</td>
  </tr>
  <tr>
    <td>APB</td>
    <td><b>100.00</b></td>
    <td><b>100.00</b></td>
    <td><b>62.20</b></td>
    <td>26.51</td>
    <td><b>12.74</b></td>
    <td><b>66.38</b></td>
    <td><b>25.00</b></td>
    <td><b>11.22</b></td>
    <td>37.31</td>
    <td>35.71</td>
    <td><b>47.34</b></td>
  </tr>
  <tr>
    <td colspan="12">Yi-34B-200K</td>
  </tr>
  <tr>
    <td>FULLATTN</td>
    <td>100.00</td>
    <td>100.00</td>
    <td>49.00</td>
    <td>5.83</td>
    <td>17.57</td>
    <td>47.60</td>
    <td>2.00</td>
    <td>18.77</td>
    <td>25.13</td>
    <td>28.00</td>
    <td>39.39</td>
  </tr>
  <tr>
    <td>MINERENCE</td>
    <td><b>100.00</b></td>
    <td><b>100.00</b></td>
    <td>52.40</td>
    <td><b>7.19</b></td>
    <td><b>20.07</b></td>
    <td><b>63.32</b></td>
    <td><b>2.50</b></td>
    <td>25.40</td>
    <td><b>28.17</b></td>
    <td><b>31.40</b></td>
    <td>43.05</td>
  </tr>
  <tr>
    <td>STARATTN</td>
    <td><b>100.00</b></td>
    <td><b>100.00</b></td>
    <td>24.60</td>
    <td>4.38</td>
    <td>19.71</td>
    <td>60.26</td>
    <td>0.00</td>
    <td>23.40</td>
    <td>28.17</td>
    <td>21.14</td>
    <td>38.17</td>
  </tr>
  <tr>
    <td>APB</td>
    <td><b>100.00</b></td>
    <td><b>100.00</b></td>
    <td><b>65.20</b></td>
    <td>5.96</td>
    <td>19.81</td>
    <td>62.45</td>
    <td>1.00</td>
    <td><b>25.42</b></td>
    <td>27.92</td>
    <td>30.00</td>
    <td><b>43.78</b></td>
  </tr>
</table>

Table 1: The results of APB compared with all the baselines on \( \infty \)Bench, where higher score represents better performance. “Avg.” represents the average score. The highest score in each column is marked in bold. FULLATTN represents FLASHATTN, RINGATTN, and ULYSSES, as their computational results remain unchanged.

<table>
  <tr>
    <th>Method</th>
    <th>SG1</th>
    <th>SG2</th>
    <th>SG3</th>
    <th>MK1</th>
    <th>MK2</th>
    <th>MK3</th>
    <th>MV</th>
    <th>MQ</th>
    <th>VT</th>
    <th>CWE</th>
    <th>FWE</th>
    <th>QA1</th>
    <th>QA2</th>
    <th>Avg.</th>
  </tr>
  <tr>
    <td colspan="15">Llama-3.1-BB-instruct</td>
  </tr>
  <tr>
    <td>FULLATTN</td>
    <td>99.40</td>
    <td>99.80</td>
    <td>99.60</td>
    <td>98.20</td>
    <td>87.60</td>
    <td>67.00</td>
    <td>94.65</td>
    <td>98.00</td>
    <td>60.98</td>
    <td>71.40</td>
    <td>72.20</td>
    <td>78.20</td>
    <td>41.6</td>
    <td>82.20</td>
  </tr>
  <tr>
    <td>MINERENCE</td>
    <td><b>100.00</b></td>
    <td>98.60</td>
    <td>99.00</td>
    <td>95.40</td>
    <td>58.20</td>
    <td>23.80</td>
    <td>84.35</td>
    <td>95.70</td>
    <td><b>66.40</b></td>
    <td>45.94</td>
    <td>74.67</td>
    <td>67.80</td>
    <td>38.80</td>
    <td>72.97</td>
  </tr>
  <tr>
    <td>STARATTN</td>
    <td><b>100.00</b></td>
    <td>99.60</td>
    <td>99.60</td>
    <td><b>95.80</b></td>
    <td>73.30</td>
    <td>53.00</td>
    <td>72.80</td>
    <td>94.45</td>
    <td>59.40</td>
    <td><b>65.72</b></td>
    <td>76.53</td>
    <td>67.40</td>
    <td>41.00</td>
    <td>76.84</td>
  </tr>
  <tr>
    <td>APB</td>
    <td><b>100.00</b></td>
    <td><b>100.00</b></td>
    <td><b>99.80</b></td>
    <td>85.60</td>
    <td><b>91.00</b></td>
    <td><b>89.00</b></td>
    <td><b>95.05</b></td>
    <td><b>96.40</b></td>
    <td>51.96</td>
    <td>63.82</td>
    <td><b>77.33</b></td>
    <td><b>70.00</b></td>
    <td><b>41.20</b></td>
    <td><b>81.63</b></td>
  </tr>
  <tr>
    <td colspan="15">Qwen-2.5-14B-instruct</td>
  </tr>
  <tr>
    <td>FULLATTN</td>
    <td><b>100.00</b></td>
    <td>99.20</td>
    <td>99.80</td>
    <td>94.20</td>
    <td>47.80</td>
    <td>27.20</td>
    <td>75.10</td>
    <td>94.60</td>
    <td>89.52</td>
    <td>93.88</td>
    <td>76.13</td>
    <td>63.20</td>
    <td>43.40</td>
    <td>77.23</td>
  </tr>
  <tr>
    <td>MINERENCE</td>
    <td><b>100.00</b></td>
    <td>99.60</td>
    <td>100.00</td>
    <td>89.60</td>
    <td>32.60</td>
    <td>6.94</td>
    <td>61.25</td>
    <td><b>93.29</b></td>
    <td>89.44</td>
    <td>84.36</td>
    <td>77.33</td>
    <td>56.60</td>
    <td>44.40</td>
    <td>71.95</td>
  </tr>
  <tr>
    <td>STARATTN</td>
    <td><b>100.00</b></td>
    <td>99.20</td>
    <td>97.40</td>
    <td>80.00</td>
    <td>38.20</td>
    <td>20.80</td>
    <td>48.75</td>
    <td>77.10</td>
    <td>82.44</td>
    <td>93.86</td>
    <td><b>77.53</b></td>
    <td>52.80</td>
    <td>38.00</td>
    <td>69.64</td>
  </tr>
  <tr>
    <td>APB</td>
    <td><b>100.00</b></td>
    <td><b>100.00</b></td>
    <td><b>99.80</b></td>
    <td><b>93.00</b></td>
    <td><b>87.20</b></td>
    <td><b>85.80</b></td>
    <td><b>66.55</b></td>
    <td><b>93.00</b></td>
    <td><b>96.40</b></td>
    <td><b>94.94</b></td>
    <td>76.33</td>
    <td><b>60.00</b></td>
    <td><b>55.60</b></td>
    <td><b>85.28</b></td>
  </tr>
  <tr>
    <td colspan="15">Yi-34B-200K</td>
  </tr>
  <tr>
    <td>FULLATTN</td>
    <td>100.00</td>
    <td>100.00</td>
    <td>99.60</td>
    <td>95.20</td>
    <td>76.00</td>
    <td>55.40</td>
    <td>92.10</td>
    <td>97.05</td>
    <td>85.56</td>
    <td>51.84</td>
    <td>84.27</td>
    <td>65.20</td>
    <td>50.00</td>
    <td>80.94</td>
  </tr>
  <tr>
    <td>MINERENCE</td>
    <td><b>100.00</b></td>
    <td><b>100.00</b></td>
    <td><b>100.00</b></td>
    <td><b>96.00</b></td>
    <td>87.00</td>
    <td>59.60</td>
    <td><b>90.15</b></td>
    <td>96.05</td>
    <td>73.68</td>
    <td>79.20</td>
    <td><b>84.87</b></td>
    <td><b>67.40</b></td>
    <td>49.40</td>
    <td>83.33</td>
  </tr>
  <tr>
    <td>STARATTN</td>
    <td><b>100.00</b></td>
    <td><b>100.00</b></td>
    <td>99.60</td>
    <td>93.20</td>
    <td>78.20</td>
    <td>48.00</td>
    <td>77.05</td>
    <td>88.35</td>
    <td><b>83.96</b></td>
    <td>80.58</td>
    <td>78.33</td>
    <td>61.20</td>
    <td>48.80</td>
    <td>79.79</td>
  </tr>
  <tr>
    <td>APB</td>
    <td><b>100.00</b></td>
    <td><b>100.00</b></td>
    <td><b>100.00</b></td>
    <td>92.80</td>
    <td><b>88.00</b></td>
    <td><b>92.60</b></td>
    <td>86.05</td>
    <td><b>96.45</b></td>
    <td>60.48</td>
    <td><b>88.46</b></td>
    <td>84.53</td>
    <td>63.80</td>
    <td><b>52.20</b></td>
    <td><b>85.03</b></td>
  </tr>
</table>

Table 2: The results of APB compared with all the baselines on RULER, where higher score represents better performance. “Avg.” represents the average score. The highest score in each column is marked in bold. FULLATTN represents FLASHATTN, RINGATTN, and ULYSSES, as their computational results remain unchanged.

of tokens in both the prefill and decoding stages, divided by the total inference time, which includes both the prefill and decoding stages.

Baselines. Existing methods for accelerating long-context inference can be categorized into four types based on their use of sequence parallelism and approximate attention mechanisms. We select FLASHATTN, ULYSSES, RINGATTN, MINERENCE, and STARATTN as our baselines. FLASHATTN does not use sequence parallelism and maintains accurate attention computation. ULYSSES and RINGATTN are two major approaches for performing accurate attention with sequence parallelism, each with a different communication design. MINERENCE is an approximate attention mechanism that accelerates long-context inference by applying sparse attention without sequence parallelism. STARATTN combines sequence parallelism and approximate attention by prepending an anchor block to each context block and removing all communication. Further details on the baselines can be found in Appendix C.

4.2 End-to-End Results

To address Q1, we conduct an end-to-end benchmark focusing on the overall performance and long-context processing speed by comparing APB with all the baselines on both \( \infty \)Bench and RULER.

APB demonstrates superior task performance in both real-world scenarios (\( \infty \)Bench in Table 1) and synthetic benchmarks (RULER in Table 2), compared to all the selected baselines. In certain tasks, APB even surpasses FULLATTN. In contrast, MINERENCE and STARATTN exhibit noticeable
performance degradation relative to FULLATTN. Notably, APB significantly enhances performance in complex context retrieval tasks, such as R.KV in ∞Bench and MK3 in RULER. This improvement can be attributed to APB’s ability to trim noisy context within each block by our compressor, thereby constructing cleaner passing blocks with reduced noise and enabling more accurate context retrieval.

In terms of inference speed, APB also demonstrates significant advantages over existing methods. As illustrated in Figure 3, methods without sequence parallelism, such as FLASHATTN and MINFERENCE, exhibit slower inference speeds. In contrast, sequence parallelism-based methods, including RINGATTN and ULYSSES, achieve speedups ranging from 3× to 10× compared to FLASHATTN. By incorporating approximate attention, STARATTN achieves a faster inference speed than RINGATTN, though its improvement over ULYSSES remains limited. APB, on the other hand, delivers at least a 25% speed improvement over STARATTN without any performance decay.

Overall, APB establishes the best speed-performance trade-off, outperforming all competing methods in both efficiency and effectiveness.

4.3 Benchmarking Across Different Lengths

We evaluate APB and baselines with input lengths ranging from 32K to 512K to address Q2. All methods are tested on the RULER benchmark, where we report task performance, inference speed, and the compute under each input length in Figure 4.

As shown in Figure 4(a), APB consistently outperforms all the methods across all input lengths in terms of performance. Additionally, APB achieves the fastest inference speed among all methods, as depicted in Figure 4(b). The speed advantage of APB remains humble for shorter input sequences. However, as the input length increases, APB’s advantage becomes more pronounced, demonstrating a substantial speedup over all competitors at a 512K input length. Notably, both STARATTN and APB show an increase in inference speed from 32K to 128K inputs, whereas other methods continue to slow down. Since STARATTN and APB both incorporate approximate attention mechanisms, the overall compute is reduced, delaying the point where compute becomes the dominant bottleneck instead of memory access. Figure 4(c) further explains APB’s superior speed, as it yields significantly less compute than FULLATTN and STARATTN, especially for longer inputs. Further discussion on the compute of each method can be referred to Table 9.

4.4 Ablation Studies

Regarding Q3, we conduct ablation studies and focus on the following key points.

Each Component’s Contribution. We focus on four major components in the APB framework: the anchor block, the passing block, selecting important KV pairs via retaining heads, and embedding the query at the beginning of the anchor block. We conduct an ablation study on these four components. For the anchor block and passing block, they can be removed in the corresponding settings. To

<table>
  <tr>
    <th>No.</th>
    <th>A</th>
    <th>P</th>
    <th>C</th>
    <th>Q</th>
    <th>E.MC</th>
  </tr>
  <tr>
    <td>0</td>
    <td>✓</td>
    <td>✓</td>
    <td>✓</td>
    <td>✓</td>
    <td>72.00</td>
  </tr>
  <tr>
    <td>1</td>
    <td>✓</td>
    <td>✓</td>
    <td>✓</td>
    <td>Rd.</td>
    <td>68.00</td>
  </tr>
  <tr>
    <td>2</td>
    <td>✓</td>
    <td>✓</td>
    <td>Rd.</td>
    <td>✓</td>
    <td>66.00</td>
  </tr>
  <tr>
    <td>3</td>
    <td>✓</td>
    <td>✓</td>
    <td>Rd.</td>
    <td>✓</td>
    <td>66.00</td>
  </tr>
  <tr>
    <td>4</td>
    <td>✓</td>
    <td>X</td>
    <td>Rd.</td>
    <td>✓</td>
    <td>64.00</td>
  </tr>
  <tr>
    <td>5</td>
    <td>✓</td>
    <td>X</td>
    <td>Rd.</td>
    <td>✓</td>
    <td>62.00</td>
  </tr>
  <tr>
    <td>6</td>
    <td>X</td>
    <td>✓</td>
    <td>✓</td>
    <td>Rd.</td>
    <td>28.00</td>
  </tr>
  <tr>
    <td>7</td>
    <td>X</td>
    <td>✓</td>
    <td>Rd.</td>
    <td>✓</td>
    <td>28.00</td>
  </tr>
  <tr>
    <td>8</td>
    <td>X</td>
    <td>X</td>
    <td>Rd.</td>
    <td>✓</td>
    <td>22.00</td>
  </tr>
</table>

Table 3: Ablation studies of APB on E.MC. “A” and “P” represent the anchor block and passing block. We implement the compressors C as the retaining heads Rd. or random selectors “Rd.”. “Q” indicates embedding the query in the anchor block.
ablate the implementation of the compressors \( C \), we compare retaining heads \( R \) with random selectors picking the same number of KV pairs. When the embedded query is removed, the anchor block contains only the beginning of the document.

The ablation results in Table 3 show that all four components are effective, as removing any of them leads to performance degradation (No. 1, 2, 4, 6). The anchor block is the most critical component, as its removal causes task failure (No. 6, 7). The passing block is also essential for good performance, as a performance drop of more than 8% occurs in No. 4 and 5. The retaining heads are crucial for selecting important KV pairs to pass to subsequent hosts, as randomly selecting KV pairs results in poorer performance (No. 2). Embedding the query in the anchor block must be used in conjunction with the retaining heads, as activating these components separately does not improve performance (No. 1, 2, 3). Above all, when used together, they yield significantly better performance (No. 0).

Wall-time Breakdown Analysis. We begin by analyzing the prefill and decoding time, as detailed in Appendix E.1. Our findings indicate that prefill dominates the overall processing time, while decoding time is negligible in long-context tasks.

Next, we decompose the prefill time into seven components: the projection of QKV, the calculation of retaining heads, the communication time, the attention computation, the projection of O, the FFN computation, and others. The wall-time of each component is shown in Figure 5. Methods without sequence parallelism (FLASHATTN and MININFERENCE) exhibit substantial inference time, whereas sequence parallelism proves effective in reducing the time of both attention and FFN computations.

Compared to existing sequence parallelism methods, APB achieves even lower attention time without incurring significant overhead for retaining heads and communication. When compared with RINGATTN and ULYSSES, although the introduction of the anchor blocks increases the compute cost for FFN, the additional time is smaller than the reduction in attention time, thus resulting in a faster inference speed. APB outperforming STARATTN is due to smaller anchor blocks, whose overhead in the FFN time is significantly lower.

Distributed Settings. We evaluate the effectiveness of APB across various distributed settings. As mentioned in STARATTN (Acharya et al., 2024), when the context is split across more hosts, significant performance degradation occurs due to the increasing invisibility of the middle context. However, APB addresses this issue by using passing blocks, allowing each host to share the most essential KV pairs with subsequent hosts.

We test APB and STARATTN across {2, 4, 6, 8} hosts, and the results are shown in Table 4. When the sequence is long (128K tokens), both APB and STARATTN exhibit improved performance as the sequence parallelism size increases. With a shorter input length (32K tokens), the block size on each host is smaller, and APB demonstrates consistently stronger and more stable performance across all settings, while STARATTN shows significant performance loss as the host count increases.

<table>
  <tr>
    <th rowspan="2">Length</th>
    <th rowspan="2">Method</th>
    <th colspan="4">H = 2</th>
    <th colspan="4">H = 4</th>
    <th colspan="4">H = 6</th>
    <th colspan="4">H = 8</th>
  </tr>
  <tr>
    <th>H = 2</th>
    <th>H = 4</th>
    <th>H = 6</th>
    <th>H = 8</th>
    <th>H = 2</th>
    <th>H = 4</th>
    <th>H = 6</th>
    <th>H = 8</th>
    <th>H = 2</th>
    <th>H = 4</th>
    <th>H = 6</th>
    <th>H = 8</th>
    <th>H = 2</th>
    <th>H = 4</th>
    <th>H = 6</th>
    <th>H = 8</th>
  </tr>
  <tr>
    <td>128K</td>
    <td>APB</td>
    <td>84.00</td>
    <td>92.00</td>
    <td>92.00</td>
    <td>92.00</td>
    <td>82.00</td>
    <td>84.00</td>
    <td>86.00</td>
    <td>88.00</td>
  </tr>
  <tr>
    <td></td>
    <td>STARATTN</td>
    <td>82.00</td>
    <td>84.00</td>
    <td>86.00</td>
    <td>88.00</td>
    <td>92.00</td>
    <td>92.00</td>
    <td>94.00</td>
    <td>94.00</td>
    <td>94.00</td>
    <td>94.00</td>
    <td>94.00</td>
    <td>94.00</td>
    <td>94.00</td>
    <td>94.00</td>
    <td>94.00</td>
  </tr>
  <tr>
    <td>32K</td>
    <td>APB</td>
    <td>92.00</td>
    <td>92.00</td>
    <td>94.00</td>
    <td>94.00</td>
    <td>94.00</td>
    <td>94.00</td>
    <td>94.00</td>
    <td>94.00</td>
    <td>94.00</td>
    <td>94.00</td>
    <td>94.00</td>
    <td>94.00</td>
    <td>94.00</td>
    <td>94.00</td>
    <td>94.00</td>
  </tr>
  <tr>
    <td></td>
    <td>STARATTN</td>
    <td>94.00</td>
    <td>92.00</td>
    <td>92.00</td>
    <td>84.00</td>
    <td>94.00</td>
    <td>94.00</td>
    <td>94.00</td>
    <td>94.00</td>
    <td>94.00</td>
    <td>94.00</td>
    <td>94.00</td>
    <td>94.00</td>
    <td>94.00</td>
    <td>94.00</td>
    <td>94.00</td>
  </tr>
</table>

Table 4: The E.MC performance of APB and STARATTN tested under various distributed settings and different context length. “H” refers to the number of hosts.

Short-Context Performance. While APB exhibits outstanding performance and inference speed on long-context benchmarks, its effectiveness and efficiency on short-context inputs (i.e., 4K tokens) are equally important. To evaluate this, we conduct experiments on RULER using 4K input lengths, reporting both task performance and inference speed.

![The wall-time breakdown of prefill for various methods on 128K context.](page_184_104_1207_246.png)

Figure 5: The wall-time breakdown of prefill for various methods on 128K context.
We compare APB against FLASHATTN using the Llama-3-8B-instruct-1M model, and present the results in Table 5.

The results demonstrate that APB consistently achieves superior speed and performance on short-context inputs. Even in such scenarios, APB outperforms FULLATTN without any degradation in performance or slowdown in inference. Please refer to Appendix E.3 for more details.

<table>
  <tr>
    <th>Method</th>
    <th>RULER-4K</th>
    <th>Speed(tok/s)</th>
  </tr>
  <tr>
    <td>FLASHATTN</td>
    <td>94.54</td>
    <td>6247.58</td>
  </tr>
  <tr>
    <td>APB</td>
    <td>94.89</td>
    <td>6597.47</td>
  </tr>
</table>

Table 5: The task performance and inference speed of APB and FLASHATTN on RULER, where context length set to 4K tokens. We report the speed in “tok/s”.

Orthogonality to Quantization. APB exhibits strong orthogonality and can be integrated with other compression methods, such as KV quantization. To demonstrate this, we tested APB in combination with KV cache quantization techniques. We integrate the KV cache quantization method officially supported by Hugging Face Transformers\(^{2}\) into the APB framework. This method is inspired by KIVI (Liu et al., 2024) and implemented using the HQQ backend\(^{3}\). We evaluate the integration on the RULER benchmark using 50 entries per task with LLaMA-3.1-8B-Instruct as the backbone. The KV cache is quantized to 4 bits. For comparison, the results of FULLATTN, MINFERENCE, STARATTN, and standalone APB are taken from Table 2, which used 500 entries per task. The results are summarized in Table 6.

<table>
  <tr>
    <th>Method</th>
    <th>Original (M)</th>
    <th>Quantized (M-4bits)</th>
  </tr>
  <tr>
    <td>FULLATTN</td>
    <td>82.20</td>
    <td>81.73</td>
  </tr>
  <tr>
    <td>MINFERENCE</td>
    <td>72.97</td>
    <td>-</td>
  </tr>
  <tr>
    <td>STARATTN</td>
    <td>76.84</td>
    <td>-</td>
  </tr>
  <tr>
    <td>APB</td>
    <td>81.63</td>
    <td>79.13</td>
  </tr>
</table>

Table 6: Performance of FULLATTN, MINFERENCE, STARATTN, APB, and their quantized variants on RULER-128K. “M” denotes the original method and “M-4bits” represents its 4-bit quantized version. Since MINFERENCE and STARATTN cause more degradation than quantization, we skip their quantized versions.

their training data. Specifically, we train the retaining heads on datasets including LongAlpaca (Chen et al., 2024), LongAlign, and AntiHaystack\(^{4}\), using Llama-3.1-8B-Instruct as the backbone model. We then evaluate APB with these trained retaining heads on RULER across 20 samples per task. Other settings are aligned to Table 2. The results are presented in Table 7.

The similar results across different training data indicate that the retaining heads in APB are robust across different training datasets, demonstrating their generalizability. Consequently, various long-context training data recipes can be employed in the APB framework.

<table>
  <tr>
    <th>Training Data</th>
    <th>LongAlign</th>
    <th>LongAlpaca</th>
    <th>AntiHaystack</th>
  </tr>
  <tr>
    <td>RULER-128K</td>
    <td>81.17</td>
    <td>81.04</td>
    <td>80.40</td>
  </tr>
</table>

Table 7: The task performance tested on RULER-128K of APB trained on various datasets.

5 Conclusion

We propose APB, a distributed long-context inference framework based on approximate attention, achieving speedups up to 9.2×, 4.2×, and 1.6× compared to FLASHATTN, RINGATTN, and STARATTN, respectively, without observable performance degradation. APB constructs the passing blocks with minimal communication cost by sending only core KV pairs among hosts. The experimental results on ∞Bench and RULER, using various models and sequence lengths, demonstrate that APB delivers faster inference speeds while maintaining or exceeding performance. Furthermore, APB is adaptable to diverse distribution configurations and models of varying sizes. Next, we aim to accelerate the decoding process in APB, where the KV cache is distributed across different hosts.

2https://github.com/huggingface/transformers
3https://github.com/mobiusml/hqq
4https://huggingface.co/datasets/wenbopan/anti-haystack
Limitations

As APB is specifically optimized to minimize the prefill time for extremely long inputs, it is less effective for processing shorter inputs, particularly those under 32K tokens. When applying APB to short inputs, the optimal distributed setting is to run the inference on a single host, as the computation can be effectively parallelized within a single GPU. In such cases, APB falls back to vanilla FLASHATTN.

Acknowledgments

This work is supported by the National Key R&D Program of China (No.2022ZD0116312) and a grant from the Guoqiang Institute, Tsinghua University. Yuxiang Huang is supported by Beijing National Science Foundation (No. QY24253). We deeply thank the anonymous reviewers for their constructive reviews.

References

Shantanu Acharya, Fei Jia, and Boris Ginsburg. 2024. Star Attention: Efficient Ilm inference over long sequences. ArXiv:2411.17116.

Anthropic. 2024. Claude 3.5 Sonnet.

Yushi Bai, Xin Lv, Jiajie Zhang, Yuze He, Ji Qi, Lei Hou, Jie Tang, Yuxiao Dong, and Juanzi Li. 2024. Longalign: A recipe for long context alignment of large language models. Proceedings of EMNLP.

Iz Beltagy, Matthew E Peters, and Arman Cohan. 2020. Longformer: The long-document transformer. ArXiv:2004.05150.

Ruisi Cai, Yuandong Tian, Zhangyang Wang, and Beidi Chen. 2024. LoCoCo: Dropping in convolutions for long context compression. Proceedings of ICML.

Yukang Chen, Shengju Qian, Haotian Tang, Xin Lai, Zhijian Liu, Song Han, and Jiaya Jia. 2024. Longlora: Efficient fine-tuning of long-context large language models. Proceedings of ICLR.

Zhuoming Chen, Ranajoy Sadhukhan, Zihao Ye, Yang Zhou, Jianyu Zhang, Niklas Nolte, Yuandong Tian, Matthijs Douze, Leon Bottou, Zhihao Jia, et al. 2025. MagicPig: LSH sampling for efficient IIm generation. Proceedings of ICLR.

Zheng Chu, Jingchang Chen, Qianglong Chen, Weijiang Yu, Tao He, Haotian Wang, Weihua Peng, Ming Liu, Bing Qin, and Ting Liu. 2023. A survey of chain of thought reasoning: Advances, frontiers and future. ArXiv:2309.15402.

Tri Dao. 2024. FlashAttention-2: Faster attention with better parallelism and work partitioning. Proceedings of ICLR.

Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. 2022. FlashAttention: Fast and memory-efficient exact attention with io-awareness. Proceedings of NeurIPS.

Tri Dao and Albert Gu. 2024. Transformers are SSMs: Generalized models and efficient algorithms through structured state space duality. Proceedings of ICML.

DeepSeek-AI. 2024. Deepseek-V3 technical report. ArXiv:2412.19437.

Xin Dong, Yonggan Fu, Shizhe Diao, Wonmin Byeon, Zijia Chen, Ameya Sunil Mahabaleshwarkar, Shih-Yang Liu, Matthijs Van Keirsbilck, Min-Hung Chen, Yoshi Suhara, et al. 2024. Hymba: A hybrid-head architecture for small language models. ArXiv:2411.13676.

Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Ajesha Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, et al. 2024. The llama 3 herd of models. ArXiv:2407.21783.

Suyu Ge, Yunan Zhang, Liyuan Liu, Minjia Zhang, Jiawei Han, and Jianfeng Gao. 2024. Model tells you what to discard: Adaptive kv cache compression for llms. Proceedings of ICLR.

Albert Gu and Tri Dao. 2024. Mamba: Linear-time sequence modeling with selective state spaces. Proceedings of COLM.

Yefei He, Luoming Zhang, Weijia Wu, Jing Liu, Hong Zhou, and Bohan Zhuang. 2024. ZipCache: Accurate and efficient kv cache quantization with salient token identification. Proceedings of NeurIPS.

Coleman Hooper, Sehoon Kim, Hiva Mohammadzadeh, Michael W Mahoney, Yukun Sophia Shao, Kurt Keutzer, and Amir Gholami. 2024. KVQuant: Towards 10 million context length IIm inference with kv cache quantization. Proceedings of NeurIPS.

Cheng-Ping Hsieh, Simeng Sun, Samuel Kriman, Shantanu Acharya, Dima Rekesh, Fei Jia, Yang Zhang, and Boris Ginsburg. 2024. RULER: What’s the real context size of your long-context language models? Proceedings of COLM.

Yanping Huang, Youlong Cheng, Ankur Bapna, Orhan Firat, Mia Xu Chen, Dehao Chen, HyoukJoong Lee, Jiquan Ngiam, Quoc V Le, Yonghui Wu, and Zhifeng Chen. 2019. GPipe: Efficient training of giant neural networks using pipeline parallelism. Proceedings of NeurIPS.

Yuxiang Huang, Binhang Yuan, Xu Han, Chaojun Xiao, and Zhiyuan Liu. 2024. Locret: Enhancing eviction in long-context IIm inference with trained retaining heads. ArXiv:2410.01805.

Sam Ade Jacobs, Masahiro Tanaka, Chengming Zhang, Minjia Zhang, Shuaiwen Leon Song, Samyam Rajbhandari, and Yuxiong He. 2023. DeepSpeed
Ulysses: System optimizations for enabling training of extreme long sequence transformer models. ArXiv:2309.14509.

Huiqiang Jiang, Yucheng Li, Chengruidong Zhang, Qianhui Wu, Xufang Luo, Surin Ahn, Zhenhua Han, Amir H Abdi, Dongsheng Li, Chin-Yew Lin, et al. 2024. Minference 1.0: Accelerating pre-filling for long-context llms via dynamic sparse attention. Proceedings of ICML.

Yeuseung Kim, Dohyun Kim, Jieun Choi, Jisang Park, Nayoung Oh, and Daehyung Park. 2024. A survey on integration of large language models with intelligent robots. Intelligent Service Robotics.

Wonbeom Lee, Jungi Lee, Junghwan Seo, and Jaewoong Sim. 2024. InfiniGen: Efficient generative inference of large language models with dynamic KV cache management. Proceedings of OSDI.

Aonian Li, Bangwei Gong, Bo Yang, Boji Shan, Chang Liu, Cheng Zhu, Chunhao Zhang, Congchao Guo, Da Chen, Dong Li, et al. 2025. Minimax-01: Scaling foundation models with lightning attention. ArXiv preprint arXiv:2501.08313.

Shenggui Li, Fuzhao Xue, Chaitanya Baranwal, Yongbin Li, and Yang You. 2023. Sequence Parallelism: Long sequence training from system perspective. Proceedings of ACL.

Xinzhe Li. 2025. A review of prominent paradigms for llm-based agents: Tool use (including RAG), planning, and feedback learning. Proceedings of COLING.

Yucheng Li, Huiqiang Jiang, Qianhui Wu, Xufang Luo, Surin Ahn, Chengruidong Zhang, Amir H Abdi, Dongsheng Li, Jianfeng Gao, Yuqing Yang, et al. 2024a. SCBench: A kv cache-centric analysis of long-context methods. ArXiv:2412.10319.

Yuhong Li, Yingbing Huang, Bowen Yang, Bharat Venkitesh, Acyr Locatelli, Hanchen Ye, Tianle Cai, Patrick Lewis, and Deming Chen. 2024b. SnapKV: Llm knows what you are looking for before generation. Proceedings of NeurIPS.

Opher Lieber, Barak Lenz, Hofit Bata, Gal Cohen, Jhonathan Osin, Itay Dalmedigos, Erez Safahi, Shaked Meirton, Yonatan Belinkov, Shai Shalev-Shwartz, et al. 2025. Jamba: A hybrid transformer-mamba language model. Proceedings of ICLR.

Zirui Liu, Jiayi Yuan, Hongye Jin, Shaochen Zhong, Zhaozhou Xu, Vladimir Braverman, Beidi Chen, and Xia Hu. 2024. KIVI: A tuning-free asymmetric 2bit quantization for KV cache. Proceedings of ICML.

Chao Lou, Zixia Jia, Zilong Zheng, and Kewei Tu. 2024. Sparser is Faster and Less is More: Efficient sparse attention for long-range transformers. ArXiv:2406.16747.

Deepak Narayanan, Mohammad Shoeybi, Jared Casper, Patrick LeGresley, Mostofa Patwary, Vijay Korsthikanti, Dmitri Vainbrand, Prethvi Kashinkunti, Julie Bernauer, Bryan Catanzaro, et al. 2021. Efficient large-scale language model training on gpu clusters using Megatron-LM. Proceedings of SC.

OpenAI. 2024. OpenAI GPT-4o.

Yujiia Qin, Shengding Hu, Yankai Lin, Weize Chen, Ning Ding, Ganqu Cui, Zheni Zeng, Xuanhe Zhou, Yufei Huang, Chaojun Xiao, et al. 2024. Tool learning with foundation models. ACM Computing Surveys.

Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He. 2020. ZeRO: Memory optimizations toward training trillion parameter models. Proceedings of SC.

Jie Ren, Samyam Rajbhandari, Reza Yazdani Am-inabadi, Olatunji Ruwase, Shuangyan Yang, Minjia Zhang, Dong Li, and Yuxiong He. 2021. ZeRO-Offload: Democratizing billion-scale model training. Proceedings of ATC.

Pranab Sahoo, Ayush Kumar Singh, Sriparna Saha, Vinija Jain, Samrat Mondal, and Aman Chadha. 2024. A systematic survey of prompt engineering in large language models: Techniques and applications. ArXiv:2402.07927.

Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, and Tri Dao. 2024. FlashAttention-3: Fast and accurate attention with asynchrony and low-precision. ArXiv:2407.08608.

Luohe Shi, Hongyi Zhang, Yao Yao, Zuchao Li, and Hai Zhao. 2024. Keep the Cost Down: A review on methods to optimize llm’s kv-cache consumption. Proceedings of COLM.

Ao Sun, Weilin Zhao, Xu Han, Cheng Yang, Zhiyuan Liu, Chuan Shi, and Maosong Sun. 2024a. BurstAttention: An efficient distributed attention framework for extremely long sequences. ArXiv:2403.09347.

Ao Sun, Weilin Zhao, Xu Han, Cheng Yang, Xinrong Zhang, Zhiyuan Liu, Chuan Shi, and Maosong Sun. 2025. Seq1F1B: Efficient sequence-level pipeline parallelism for large language model training. Proceedings of NAACL.

Hanshi Sun, Li-Wen Chang, Wenlei Bao, Size Zheng, Ningxin Zheng, Xin Liu, Harry Dong, Yuejie Chi, and Beidi Chen. 2024b. ShadowKV: KV cache in shadows for high-throughput long-context llm inference. ArXiv:2410.21465.

Ashish Vaswani. 2017. Attention is all you need. Proceedings of NeurIPS.

Zhongwei Wan, Xin Wang, Che Liu, Samiul Alam, Yu Zheng, et al. 2024. Efficient large language models: A survey. Transactions on Machine Learning Research.
Zheng Wang, Boxiao Jin, Zhongzhi Yu, and Minjia Zhang. 2024. Model tells you where to merge: Adaptive kv cache merging for llms on long-context tasks. arXiv:2407.08454.

Chaojun Xiao, Pengle Zhang, Xu Han, Guangxuan Xiao, Yankai Lin, Zhengyan Zhang, Zhiyuan Liu, Song Han, and Maosong Sun. 2024a. InflLM: Unveiling the intrinsic capacity of llms for understanding extremely long sequences with training-free memory. Proceedings of NeurIPS.

Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. 2024b. Efficient streaming language models with attention sinks. Proceedings of ICLR.

Xinyu Yang, Tianqi Chen, and Beidi Chen. 2025. APE: Faster and longer context-augmented generation via adaptive parallel encoding. Proceedings of ICLR.

Yao Yao, Zuchao Li, and Hai Zhao. 2024. SirLLM: Streaming infinite retentive LLM. Proceedings of ACL.

Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, and Amr Ahmed. 2020. Big Bird: Transformers for longer sequences. Proceedings of NeurIPS.

Fanlong Zeng, Wensheng Gan, Yongheng Wang, Ning Liu, and Philip S Yu. 2023. Large language models for robotics: A survey. ArXiv:2311.07226.

Tianyi Zhang, Jonah Yi, Zhaozhou Xu, and Anshumali Shrivastava. 2024a. KV cache is 1 bit per channel: Efficient large language model inference with coupled quantization. Proceedings of NeurIPS.

Xinrong Zhang, Yingfa Chen, Shengding Hu, Zihang Xu, Junhao Chen, Moo Khai Hao, Xu Han, Zhen Leng Thai, Shuo Wang, Zhiyuan Liu, et al. 2024b. \( \infty \) bench: Extending long context evaluation beyond 100k tokens. Proceedings of ACL.

Xuan Zhang, Cuxiao Du, Chao Du, Tianyu Pang, Wei Gao, and Min Lin. 2024c. SimLayerKV: A simple framework for layer-level kv cache reduction. ArXiv:2410.13846.

Yuxin Zhang, Yuxuan Du, Gen Luo, Yunshan Zhong, Zhenyu Zhang, Shiwei Liu, and Rongrong Ji. 2024d. CaM: Cache merging for memory-efficient llms inference. In Proceedings of ICML.

Zhenyu Zhang, Ying Sheng, Tianyi Zhou, Tianlong Chen, Lianmin Zheng, Ruisi Cai, Zhao Song, Yuandong Tian, Christopher Ré, Clark Barrett, et al. 2024e. H₂O: Heavy-hitter oracle for efficient generative inference of large language models. Proceedings of NeurIPS.

Jun Zhao, Can Zu, Hao Xu, Yi Lu, Wei He, Yiwen Ding, Tao Gui, Qi Zhang, and Xuanjing Huang. 2024. LongAgent: Scaling language models to 128k context through multi-agent collaboration. Proceedings of EMNLP.

Wayne Xin Zhao, Kun Zhou, Junyi Li, Tianyi Tang, Xiaolei Wang, Yupeng Hou, Yingqian Min, Beichen Zhang, Junjie Zhang, Zican Dong, et al. 2023. A survey of large language models. ArXiv:2303.18223.
A Pseudocode of APB

Algorithm 1: APB Inference
Input: Model M, Input sequence t, Anchor length \( l_a \), Passing length \( l_p \), Host index h, Total host number H, Stop Criteria SC
Output: Generated tokens \( t_{gen} \)
1 d, q ← SplitDocumentQuery(t)
2 \( \mathbf{B}_h \) ← SplitContext(d, h)
// Make anchor block and embed query
3 \( \mathbf{A} \leftarrow [q_1, \cdots, q_q, d_1, \cdots, d_{l_h}] \)
4 if h = 1 then
    | \( \mathbf{A} \leftarrow \text{None} \)
end if
7 \( \mathbf{H}_a^{(0)}, \mathbf{H}_h^{(0)} \leftarrow \text{M.embedding}(\mathbf{A}, \mathbf{B}_h) \)
// Prefill
8 KV_cache ← []
9 for i, l ∈ enumerate(M.layers) do
10   \( \mathbf{H}_a^{(i)}, \mathbf{H}_h^{(i)}, \mathbf{K}_a^{(i)}, \mathbf{K}_h^{(i)}, \mathbf{V}_a^{(i)}, \mathbf{V}_h^{(i)} \leftarrow \text{APB}(1, \mathbf{H}_a^{(i-1)}, \mathbf{H}_h^{(i-1)}, l_p, h, H) \)
11   KV_cache.append(\(\mathbf{K}_h^{(i)}, \mathbf{V}_h^{(i)}\))
end for
// Decoding
13 \( t_{gen} \leftarrow [] \)
14 x ← q
15 while not SC(M, t, t_{gen}) do
16   \( \mathbf{H}^{(0)} \leftarrow \text{M.embedding}(x) \)
17   for i, l ∈ enumerate(M.layers) do
18     \( \mathbf{H}^{(i)}, \mathbf{K}^{(i)}, \mathbf{V}^{(i)} \leftarrow \text{Accu}(1, \mathbf{H}^{(i-1)}, \text{KV_cache}[i], h, H) \)
19     if \( h = H \) then
20       \( \text{KV_cache}[i].extend(\mathbf{K}^{(i)}, \mathbf{V}^{(i)}) \)
21     end if
22   end for
23   x ← GenerateNextToken(H^{(L)}, M.LM_Head)
24   t_{gen} ← [t_{gen}, x]
25 end while
return t_{gen}

Here, we present the complete inference process of APB in Algorithm 1.

For initialization, the input sequence t is first split into the document d and query q (line 1). Each host then takes a partial context \( \mathbf{B}_h \) based on the host index h (line 2) and sets the anchor block (lines 3–6). To begin prefill, the embeddings of the anchor block \( \mathbf{A} \) and the context block \( \mathbf{B}_h \) are computed (line 7). During the prefill stage (lines 8–12), we employ Algorithm 2 to perform sequence parallelism-aware approximate attention. In the decoding stage (lines 13–25), we directly apply accurate attention with online softmax (Algorithm 3), as introduced in StarATtn (Acharya et al., 2024) as stage-2, to generate new tokens. Finally, the generated tokens are returned (line 26).

The prefill stage function of APB is described in Algorithm 2. Before performing the attention calculation, the KV cache of the current host is compressed using the retaining heads \( \mathcal{R} \) (line 2–4). Next, two Allgather communications are performed on the compressed K and V caches, allowing each host to obtain a full view of the compressed KV cache from all hosts (line 5–6). Subsequently, the passing block is formed by concatenating the compressed KV caches sent from the previous hosts (line 7), followed by the modified attention mechanism (line 8). After the attention calculation, the passing blocks are discarded, and only the attention scores corresponding to the anchor block and the context block held by the current host are passed through the FFN (line 9).

The decoding function of APB, also known as stage-2 of StarATtn, is described in Algorithm 3. In the decoding stage, all hosts hold the hidden states \( \mathbf{H} \) of the new token. The \( \mathbf{Q}, \mathbf{K}, \) and \( \mathbf{V} \) of the

Algorithm 2: APB Prefill Function
Input: Layer Layer, Anchor hidden states \( \mathbf{H}_a \), Block hidden states \( \mathbf{H}_h \), Passing length \( l_p \), Host index h, Total host number H
Output: Output hidden states \( \mathbf{H}_a^{out}, \mathbf{H}_h^{out} \), KV cache \( \mathbf{K}_a, \mathbf{K}_h, \mathbf{V}_a, \mathbf{V}_h \)
1 \( [\mathbf{Q}_a, \mathbf{Q}_h], [\mathbf{K}_a, \mathbf{K}_h], [\mathbf{V}_a, \mathbf{V}_h] \leftarrow \text{layer.qkv_proj}([\mathbf{H}_a, \mathbf{H}_h]) \)
// Compress KV cache by Retaining Head
2 \( s_1, \cdots, s_{lh} \leftarrow \text{layer.R}([\mathbf{Q}_h, \mathbf{K}_h, \mathbf{V}_h]) \)
3 indices ← ArgTop-l_p(s_1, \cdots, s_{lh})
4 \( \mathbf{K}_h^C, \mathbf{V}_h^C \leftarrow \mathbf{K}_h[\text{indices}], \mathbf{V}_h[\text{indices}] \)
// Communication
5 \( \mathbf{K}_h^C, \cdots, \mathbf{K}_h^{H} \leftarrow \text{AllGather}(\mathbf{K}_h^C) \)
6 \( \mathbf{V}_h^C, \cdots, \mathbf{V}_h^{H} \leftarrow \text{AllGather}(\mathbf{V}_h^C) \)
7 \( \mathbf{K}_p \leftarrow \mathbf{K}_1^C, \cdots, \mathbf{K}_h^{H-1} \mathbf{V}_p \leftarrow \mathbf{V}_1^C, \cdots, \mathbf{V}_h^{H-1} \)
// Attention with A block and P block
8 \( \mathbf{A}_a, \mathbf{A}_h \leftarrow \text{Attention}([\mathbf{Q}_a, \mathbf{Q}_h], [\mathbf{K}_a, \mathbf{K}_p, \mathbf{K}_h], [\mathbf{V}_a, \mathbf{V}_p, \mathbf{V}_h]) \)
9 \( \mathbf{H}_a^{out}, \mathbf{H}_h^{out} \leftarrow \text{FFN}(\mathbf{A}_a, \mathbf{A}_h) \)
10 return \( \mathbf{H}_a^{out}, \mathbf{H}_h^{out}, \mathbf{K}_a, \mathbf{K}_h, \mathbf{V}_a, \mathbf{V}_h \)

Algorithm 3: APB Decoding Function
Input: Layer Layer, Hidden states \( \mathbf{H} \), KV cache cache, Host index h, Total host number H
Output: Output hidden states \( \mathbf{H}^{out} \), KV cache \( \mathbf{K}, \mathbf{V} \)
1 \( \mathbf{Q}, \mathbf{K}, \mathbf{V} \leftarrow \text{layer.qkv_proj}(\mathbf{H}) \)
2 \( \mathbf{K}_{cache}, \mathbf{V}_{cache} \leftarrow \text{cache} \)
3 if \( h < H \) then
    | \( \mathbf{A}.lse \leftarrow \text{Attention}(\mathbf{Q}, \mathbf{K}_{cache}, \mathbf{V}_{cache}) \)
end if
else
    | \( \mathbf{A}.lse \leftarrow \text{Attention}(\mathbf{Q}, [\mathbf{K}_{cache}, \mathbf{K}], [\mathbf{V}_{cache}, \mathbf{V}]) \)
end if
9 \( \mathbf{A}_1, \cdots, \mathbf{A}_H; lse_1, \cdots, lse_H \leftarrow \text{Gather}(\mathbf{A}.lse) \)
10 \( \mathbf{A} \leftarrow \text{MergeScore}(\mathbf{A}_1, \cdots, \mathbf{A}_H; lse_1, \cdots, lse_H) \)
11 \( \mathbf{H}^{out} \leftarrow \text{FFN}(\mathbf{A}) \)
12 return \( \mathbf{H}^{out}, \mathbf{K}, \mathbf{V} \)
new token are computed simultaneously on each host (line 1). For all hosts except the last one, the attention score is calculated by the new token’s query with the KV cache held by the current host (line 4). For the last host, the local KV cache is first concatenated with the new KV, and then attention calculation is performed with the query (line 7). A Gather communication is performed to collect \( \mathbf{A}_h \) and 1se calculated by each host (line 9), which are then merged to obtain the global attention score \( \mathbf{A} \) (line 10). Notably, the \( \mathbf{A} \) on each host should be identical. Finally, \( \mathbf{A} \) is passed through the FFN simultaneously on each host, generating identical \( \mathbf{H}^{out} \) on all hosts (line 11).

B Reproducibility Settings

B.1 Training of the Compressor

We use the retaining heads proposed in Huang et al. (2024) as our compressor. Following Huang et al. (2024), we set the intermediate size of the retaining heads to 1024. We utilize the first 3000 samples from the LongAlign (Bai et al., 2024) dataset to generate the training labels, and we train the retaining heads with the a frozen model backbone for 3000 steps in 1 epoch. The batch size set is 1 and the maximum input length is set to 10240. We use AdamW as the optimizer and the learning rate is set to 5e-4, with \( \beta_1 \) set to 0.9 and \( \beta_2 \) set to 0.95. We apply a linear scheduler whose number of warmup steps is set to 300. Notably, the training loss of the retaining head consists of a regression loss and a smoothing loss, following the setting of (Huang et al., 2024), and the balance factor \( \alpha \) is set to 0.0025. The gradient clipping value is set to 0.5 to avoid gradient explosion.

<table>
  <tr>
    <th>n</th>
    <th>32K</th>
    <th>64K</th>
    <th>128K</th>
    <th>256K</th>
    <th>512K</th>
  </tr>
  <tr>
    <th>\( l_b \)</th>
    <td>4K</td>
    <td>8K</td>
    <td>16K</td>
    <td>32K</td>
    <td>64K</td>
  </tr>
  <tr>
    <th>\( l_a \)</th>
    <td>1K</td>
    <td>2K</td>
    <td>4K</td>
    <td>8K</td>
    <td>8K</td>
  </tr>
  <tr>
    <th>\( l_p \)</th>
    <td>0.5K</td>
    <td>1K</td>
    <td>2K</td>
    <td>4K</td>
    <td>8K</td>
  </tr>
</table>

Table 8: Hyperparameters of APB used in Section 4.3. n stands for input length, \( l_b \) is the block size, \( l_a \) represents the anchor length, and \( l_p \) represents the passing length. “K” is an abbreviation for 1024.

B.2 Inference Hyperparameters

Here, we elaborate on the details of the inference hyperparameters used in Section 4.2, 4.3, and 4.4.

B.2.1 End-to-End Benchmark (Section 4.2)

In the end-to-end benchmark, we first test the performance of 3 baselines with the performance of our method. Since methods without approximate attention do not alter the computational outcome of the attention mechanism, we take the result of ULYSSES as the performance benchmark for FULLATTN. Then, we evaluate the inference speed of all baselines along with APB .

Performance Evaluation. In the \( \infty \)Bench experiments, we evaluate the performance of each method by running all the data from each task. Details of the tasks can be found in Zhang et al. (2024b). When conducting the RULER experiments, we generate 500 test samples for each RULER task to evaluate the performance, with details introduced in Hsieh et al. (2024). We report the performance of each method under a 128K input length for both benchmarks. For methods with sequence parallelism, we set the sequence parallel size to 8. We use a single machine with 8 GPUs for Llama-3.1-8B and Qwen-2.5-14B, while for Yi-34B, due to its large model size, we employ two machines with layers evenly distributed between them. All the answers are generated with greedy decoding, i.e. with a temperature set to 0. For STARATTN, we set both the block length and the anchor length to 16K. For MIINERENCE, we use the official head configuration for Llama-3.1-8B and generate the configurations for Qwen-2.5-14B and Yi-34B based on the first data entry of SG1 from RULER. For APB, we set anchor length \( l_a \) to 4K and passing length \( l_p \) to 2K. We place the query tokens after the system prompt to embed the query within the anchor blocks without disrupting the model’s chat template.

Speed Evaluation. We run the first 20 samples for each task and average the inference speed. To reflect the speed of LLMs in processing a query, we define the inference speed as follows.

\[
\text{speed} = \frac{\# \text{input tokens} + \# \text{output tokens}}{\text{prefill time} + \text{decoding time}}
\]

We keep all the hyperparameters and distribution settings same as the performance evaluation.

B.2.2 Benchmarking Across Input Lengths (Section 4.3)

For Section 4.3, we evaluate the performance, speed, and report the compute of all the methods across different input lengths using the Llama-3.1-8B model and the RULER benchmark.
We run 50 samples per task for performance evaluation, while for speed testing, we run 20 samples per task. For STARATTN, we set both the block size and the anchor length to \( n/H \), where \( n \) is the input length and \( H \) is the sequence parallelism size. For APB, we use the hyperparameters listed in Table 8. We align other hyperparameters to the end-to-end benchmark. We visualize the compute of each method following Table 9.

B.2.3 Ablation Studies (Section 4.4)

Here, we elaborate on the settings in Section 4.4.

Each Component’s Contribution. We conduct the ablation study of each component using Llama-3.1-8B. We conduct experiments of different settings on E.MC from ∞Bench using the first 50 samples to assess the impact of different combinations on performance. We use the hyperparameters of \( n = 128K, l_b = 32K, l_a = 4K, l_p = 2K \).

Wall-time Breakdown Analysis. The prefill and decoding time is tested with Llama-3.1-8B on SGI from RULER. For the breakdown analysis of the prefill time, we run all the methods on synthetic random input. All the experiments are conducted under a 128K input length, and the hyperparameters are kept the same as in previous experiments.

Distributed Settings. We test the performance under different host numbers (sequence parallelism size) with Llama-3.1-8B. We set \( H \) across {2, 4, 6, 8} and evaluate the performance of STARATTN and APB on E.MC from ∞Bench. We set the sequence length to 32K and 128K, and test on the first 50 samples.

B.3 System Environment

All training and evaluation are conducted on workstations equipped with 8× NVIDIA A800-80GB GPUs and 104 Intel® Xeon® Platinum 8470 CPUs, running CentOS Linux 7 (Core). The GPUs within each machine are interconnected using third-generation NVLink technology, while cross-machine communication is facilitated via HDR InfiniBand. We perform the training of the retaining head and all experiments—except for FLASHATTN and MINERENCE tests—on 8 GPUs. The experiments for FLASHATTN and MINERENCE are conducted on a single GPU.

C Baselines

In this section, we briefly introduce the design principles of each baseline.

FLASHATTN. Dao (2024) introduces an attention computation method with hardware awareness that leverages the architecture of modern GPUs. FLASHATTN applies tiling techniques to the calculation of attention scores and fully utilizes the high-bandwidth memory to accelerate computation. FLASHATTN is an accurate attention method that preserves the computation results, primarily focusing on accelerating attention on a single GPU.

MINERENCE. As introduced in Jiang et al. (2024), three approximate attention patterns can be utilized to approximate FULLATN computation. MINERENCE first searches for the head configuration by assigning an approximate pattern to each head. During inference, only a limited number of attention score entries are calculated, accelerating the prefill by reducing the computation in attention. MINERENCE focuses on applying approximate attention without sequence parallelism.

RINGATTN. Introduced by Li et al. (2023), RINGATTN distributes the context across multiple GPUs (hosts). Online softmax is applied to the attention calculation, where each host computes the partial attention score passed by the previous host, using local context for \( H - 1 \) rounds (\( H \) is the number of hosts). The core idea is to overlap the communication of passing partial attention scores with the attention calculation. RINGATTN is primarily used for training LLMs with extremely long contexts but is not optimized for inference.

ULYSSES. Similar to RINGATTN, ULYSSES (Jacobs et al., 2023) was also introduced for long-context training. ULYSSES applies sequence parallelism, with each host holding a partial context. During attention calculation, three AlltoAll communications on Q, K, and V are performed to distribute the full context of specific heads to the corresponding host. After the calculation, another AlltoAll communication is conducted on the attention score to revert the distribution from splitting heads to splitting the sequence.

STARATTN. Introduced by Acharya et al. (2024), STARATTN is a pioneer in combining approximate attention with sequence parallelism. STARATTN introduces anchor blocks, which contain the initial tokens of the input sequence and are the same size as the blocks on each host. Each host calculates only the partial attention between the anchor block and the current context block, with no communication required. Although STARATTN reduces communication and synchronization, its large anchor blocks introduce overhead in the FFN,
limiting the speedup gains.

Here, we present the computation per forward call for each method in Table 9. Notably, since FlashAttn, RingAttn, and ULYSSES are FULLATTN methods, they share the same computation formula. We do not include MINERENCE, as its computation depends on the search results of the head configurations. We provide a visualization of Table 9 in Figure 4(c).

D Task Abbreviations

Here, we provide a mapping between the benchmark task abbreviations and their full names used in the main text in Table 10 and 11.

<table>
  <tr>
    <th>Abbr.</th>
    <th>Full Name</th>
  </tr>
  <tr>
    <td>R.PassKey</td>
    <td>Retrieve.PassKey</td>
  </tr>
  <tr>
    <td>R.Number</td>
    <td>Retrieve.Number</td>
  </tr>
  <tr>
    <td>R.KV</td>
    <td>Retrieve.KV</td>
  </tr>
  <tr>
    <td>E.Sum</td>
    <td>En.Sum</td>
  </tr>
  <tr>
    <td>E.QA</td>
    <td>En.QA</td>
  </tr>
  <tr>
    <td>E.Dia</td>
    <td>En.Dia</td>
  </tr>
  <tr>
    <td>Z.QA</td>
    <td>Zh.QA</td>
  </tr>
  <tr>
    <td>C.Debug</td>
    <td>Code.Debug</td>
  </tr>
  <tr>
    <td>M.Find</td>
    <td>Math.Find</td>
  </tr>
</table>

Table 10: Task name mapping of \( \infty \)Bench.

<table>
  <tr>
    <th>Abbr.</th>
    <th>Full Name</th>
  </tr>
  <tr>
    <td>SG1</td>
    <td>Single NIAH 1</td>
  </tr>
  <tr>
    <td>SG2</td>
    <td>Single NIAH 2</td>
  </tr>
  <tr>
    <td>SG3</td>
    <td>Single NIAH 3</td>
  </tr>
  <tr>
    <td>MK1</td>
    <td>Multi-keys NIAH 1</td>
  </tr>
  <tr>
    <td>MK2</td>
    <td>Multi-keys NIAH 2</td>
  </tr>
  <tr>
    <td>MK3</td>
    <td>Multi-keys NIAH 3</td>
  </tr>
  <tr>
    <td>MV</td>
    <td>Multi-values NIAH</td>
  </tr>
  <tr>
    <td>MQ</td>
    <td>Multi-queries NIAH</td>
  </tr>
  <tr>
    <td>VT</td>
    <td>Variable Tracking</td>
  </tr>
  <tr>
    <td>CWE</td>
    <td>Common Words Extraction</td>
  </tr>
  <tr>
    <td>FWE</td>
    <td>Frequent Words Extraction</td>
  </tr>
  <tr>
    <td>QA1</td>
    <td>Question Answering 1</td>
  </tr>
  <tr>
    <td>QA2</td>
    <td>Question Answering 2</td>
  </tr>
</table>

Table 11: Task name mapping of RULER.

E Further Ablation Studies and Detailed Experimental Results

Here, we represent the further ablation studies in Section E.1 and E.2. We list all experimental data that are not included in the main text in Section E.5.

E.1 Breakdown Analysis: Prefill and Decoding

We conduct a wall-time breakdown analysis to further investigate how APB achieves a higher inference speed. First, we measure the prefill and decoding times for each method, and the results shown in Figure 6 reveal that the prefill stage is the bottleneck, while the decoding stage takes up significantly less time. Since APB optimizes the prefill stage, it is able to address the most time-consuming part of the process.

![Bar chart showing wall-time of prefill and decoding for various methods on 128K context. Prefill is the major bottleneck in long-context query processing.](page_1012_370_377_246.png)

Figure 6: The wall-time of prefill and decoding for various methods on 128K context. Prefill is the major bottleneck in long-context query processing.

E.2 Hyperparameter Stability

There are two hyperparameters introduced in the APB framework: the length of the anchor block \( l_a \), and the passing length \( l_p \). In this study, we examine the sensitivity of these two parameters by measuring the performance on the E.QA task from \( \infty \)Bench for various lengths of \( l_a \) and \( l_p \).

In this experiment, we use Llama-3.1-8B as the model backbone. We test on the first 50 samples and set the sequence length to 128K. We present the results with \( l_a, l_p \in \{1\mathrm{K}, 2\mathrm{K}, 3\mathrm{K}, 4\mathrm{K}\} \).

![Line charts showing E.QA performance of APB with various anchor block length \( l_a \) and passing length \( l_p \).](page_1012_1042_377_246.png)

Figure 7: The E.QA performance of APB with various anchor block length \( l_a \) and passing length \( l_p \). We select the hyperparameters from \( \{1024, 2048, 3072, 4096\} \).

The experimental results shown in Figure 7 indicate that both \( l_a \) and \( l_p \) are stable. While there is a slight performance improvement with larger \( l_a \), the variation remains insignificant. The performance change with varying \( l_p \) is also minimal. This suggests that it is not necessary to tune \( l_a \) and \( l_p \) delicately, making APB easy to use.
<table>
  <tr>
    <th>Method</th>
    <th>FLOPs / forward</th>
  </tr>
  <tr>
    <td>FULLATN</td>
    <td>\( L \times \left( 4nd^2 + \frac{4}{3}nd^2 + n^2d + 6ndI \right) \)</td>
  </tr>
  <tr>
    <td>STARATTN</td>
    <td>\( \frac{L}{\pi} \times \left[ (8H-4)nd^2 + \frac{8H-4}{9}nd^2 + \frac{8H-4}{9}n^2d + (12H-6)ndI \right] \)</td>
  </tr>
  <tr>
    <td>APB</td>
    <td>\( L \times \left[ 4 \left( 1 + \frac{1}{3} + \frac{0.5n}{nd} + \frac{1.5I}{d} \right) \frac{\pi}{\pi} d^2 + 4(H-1) \left( 1 + \frac{1}{3} + \frac{0.5(\frac{n}{d} + l_a)}{d} + \frac{1.5I}{d} \right) \left( \frac{\pi}{\pi} + l_a \right) d^2 + l_p H(H-1) \left( \frac{\pi}{\pi} + l_a \right) d \right] \)</td>
  </tr>
</table>

Table 9: The FLOPs per forward call of FULLATN, STARATTN and APB. “\(L\)” stands for the number of layers, “\(n\)” stands for the input length, “\(d\)” is the hidden size of the model, and “\(I\)” is the intermediate size of FFN. “\(H\)” stands for the number of hosts, “\(l_a\)” is the anchor length, and “\(l_p\)” is the passing length. We calculate the compute of the model without the input embedding, the language-modeling head, positional embeddings, and all the normalizations.

<table>
  <tr>
    <th>Method</th>
    <th>R.PassKey</th>
    <th>R.Number</th>
    <th>R.KV</th>
    <th>E.Sum</th>
    <th>E.QA</th>
    <th>E.MC</th>
    <th>E.Dia</th>
    <th>Z.QA</th>
    <th>C.Debug</th>
    <th>M.Find</th>
    <th>Avg.</th>
  </tr>
  <tr>
    <th colspan="13">Llama-3-1-8B-instruct</th>
  </tr>
  <tr>
    <td>FLASHATN</td>
    <td>4056</td>
    <td>3998</td>
    <td>3530</td>
    <td>609</td>
    <td>3985</td>
    <td>3924</td>
    <td>4280</td>
    <td>4137</td>
    <td>3959</td>
    <td>5209</td>
    <td>3769</td>
  </tr>
  <tr>
    <td>ULYSSES</td>
    <td>30999</td>
    <td>30426</td>
    <td>25266</td>
    <td>3423</td>
    <td>27616</td>
    <td>25860</td>
    <td>29654</td>
    <td>25332</td>
    <td>31289</td>
    <td>39967</td>
    <td>26983</td>
  </tr>
  <tr>
    <td>RINGATN</td>
    <td>20440</td>
    <td>20218</td>
    <td>17507</td>
    <td>3082</td>
    <td>19476</td>
    <td>18577</td>
    <td>20826</td>
    <td>18779</td>
    <td>20092</td>
    <td>25956</td>
    <td>18495</td>
  </tr>
  <tr>
    <td>MINFERENCE</td>
    <td>7910</td>
    <td>7686</td>
    <td>5259</td>
    <td>832</td>
    <td>6186</td>
    <td>7238</td>
    <td>7016</td>
    <td>6399</td>
    <td>7719</td>
    <td>8027</td>
    <td>6427</td>
  </tr>
  <tr>
    <td>STARATTN</td>
    <td>33968</td>
    <td>32632</td>
    <td>25840</td>
    <td>2605</td>
    <td>27896</td>
    <td>31986</td>
    <td>32267</td>
    <td>27266</td>
    <td>31942</td>
    <td>37051</td>
    <td>28345</td>
  </tr>
  <tr>
    <td><b>APB</b></td>
    <td><b>44995</b></td>
    <td><b>43111</b></td>
    <td><b>32086</b></td>
    <td><b>2848</b></td>
    <td><b>32930</b></td>
    <td><b>41040</b></td>
    <td><b>36553</b></td>
    <td><b>32057</b></td>
    <td><b>44780</b></td>
    <td><b>45117</b></td>
    <td><b>35552</b></td>
  </tr>
  <tr>
    <th colspan="13">Qwen-2.5-14B-instruct</th>
  </tr>
  <tr>
    <td>FLASHATN</td>
    <td>2155</td>
    <td>2082</td>
    <td>1736</td>
    <td>706</td>
    <td>1929</td>
    <td>1758</td>
    <td>1863</td>
    <td>1882</td>
    <td>2131</td>
    <td>2347</td>
    <td>1859</td>
  </tr>
  <tr>
    <td>ULYSSES</td>
    <td>17214</td>
    <td>16461</td>
    <td>12623</td>
    <td>1574</td>
    <td>13374</td>
    <td>13398</td>
    <td>14092</td>
    <td>11131</td>
    <td>16932</td>
    <td>18460</td>
    <td>13526</td>
  </tr>
  <tr>
    <td>RINGATN</td>
    <td>11056</td>
    <td>10589</td>
    <td>8658</td>
    <td>1314</td>
    <td>10140</td>
    <td>9539</td>
    <td>9982</td>
    <td>8977</td>
    <td>10919</td>
    <td>11676</td>
    <td>9285</td>
  </tr>
  <tr>
    <td>MINFERENCE</td>
    <td>3527</td>
    <td>3305</td>
    <td>2244</td>
    <td>618</td>
    <td>2645</td>
    <td>2498</td>
    <td>2480</td>
    <td>2222</td>
    <td>3669</td>
    <td>3785</td>
    <td>2699</td>
  </tr>
  <tr>
    <td>STARATTN</td>
    <td>18793</td>
    <td>17892</td>
    <td>13695</td>
    <td>1641</td>
    <td>14604</td>
    <td>13874</td>
    <td>14462</td>
    <td>12453</td>
    <td>18640</td>
    <td>19640</td>
    <td>14569</td>
  </tr>
  <tr>
    <td><b>APB</b></td>
    <td><b>25218</b></td>
    <td><b>23643</b></td>
    <td><b>17192</b></td>
    <td><b>1592</b></td>
    <td><b>16875</b></td>
    <td><b>17937</b></td>
    <td><b>17267</b></td>
    <td><b>13303</b></td>
    <td><b>25134</b></td>
    <td><b>25929</b></td>
    <td><b>18409</b></td>
  </tr>
  <tr>
    <th colspan="13">Yi-34B-200K</th>
  </tr>
  <tr>
    <td>FLASHATN</td>
    <td>1041</td>
    <td>1089</td>
    <td>907</td>
    <td>372</td>
    <td>1110</td>
    <td>1137</td>
    <td>1099</td>
    <td>1125</td>
    <td>1111</td>
    <td>1208</td>
    <td>1030</td>
  </tr>
  <tr>
    <td>ULYSSES</td>
    <td>9192</td>
    <td>8820</td>
    <td>7383</td>
    <td>1369</td>
    <td>8523</td>
    <td>8388</td>
    <td>8841</td>
    <td>8374</td>
    <td>9016</td>
    <td>9901</td>
    <td>7981</td>
  </tr>
  <tr>
    <td>RINGATN</td>
    <td>6203</td>
    <td>6032</td>
    <td>5296</td>
    <td>1281</td>
    <td>6021</td>
    <td>5785</td>
    <td>5951</td>
    <td>6024</td>
    <td>5989</td>
    <td>6552</td>
    <td>5513</td>
  </tr>
  <tr>
    <td>MINFERENCE</td>
    <td>1732</td>
    <td>1637</td>
    <td>1199</td>
    <td>400</td>
    <td>1327</td>
    <td>1760</td>
    <td>1523</td>
    <td>1062</td>
    <td>1766</td>
    <td>1806</td>
    <td>1421</td>
  </tr>
  <tr>
    <td>STARATTN</td>
    <td>9039</td>
    <td>8675</td>
    <td>7284</td>
    <td>1893</td>
    <td>8419</td>
    <td>9007</td>
    <td>8866</td>
    <td>7877</td>
    <td>9046</td>
    <td>9373</td>
    <td>7948</td>
  </tr>
  <tr>
    <td><b>APB</b></td>
    <td><b>12441</b></td>
    <td><b>9404</b></td>
    <td><b>2059</b></td>
    <td><b>10798</b></td>
    <td><b>12726</b></td>
    <td><b>12233</b></td>
    <td><b>9327</b></td>
    <td><b>12744</b></td>
    <td><b>9328</b></td>
    <td><b>12744</b></td>
    <td><b>10624</b></td>
  </tr>
</table>

Table 12: The inference speed of APB compared with all the baselines on \( \infty \)Bench. The “Avg.” represents the average speed. The highest score in each column is marked in bold. We report the speed in “tok/s”.

E.3 Short-Context Performance of APB

We report the details and hyperparameters of short-context performance experiment here.

As reported in Section 4.4, we compare task performance and inference speed of RULER using FULLATN (implemented with FLASHATN) and APB, with the context length set to 4K tokens, utilizing Llama-3-8B-instruct-1M as the backbone model. In addition, the number of hosts \( H \) is set to 4, and both anchor length \( l_a \) and passing length \( l_p \) are set to 256. We evaluate each method on 50 samples per task. We align the settings for inference speed measurement with those used in Figure 3. The detailed results are reported in Table 19.

E.4 Accuracy Experiment on LoCoCo

To present a comparison of APB against optimization strategies that target KV cache size reduction to enable faster inference, we select LoCoCo as a baseline here for comparison with APB. We report its performance in this section. As Table 20 shows, LoCoCo achieves low performance on RULER, as its fixed KV cache size leads to significant information loss, hindering its ability to handle complex long-context tasks effectively. The performance of FULLATN and APB corresponds to the results reported in Table 2 of the main text, which were evaluated using the Llama-3.1-8B-instruct model. Meanwhile, for the evaluation of LoCoCo, we test only 20 samples for each task.

E.5 The Detailed Experimental Results

Here, we present the detailed experimental results to extensively illustrate each figure in the main text, Section E.1 and E.2. In Section 4.2, the detailed experimental results of Figure 3 are provided in Table 12 and Table 15. Table 12 provides the accurate speed for all methods on \( \infty \)Bench, while the speed is the average value across all tasks. By contrast, Table 15 reports the precise speed measurements for all methods on RULER, with the values similarly averaged over all tasks. The performance of each method in Figure 3 is provided in Table 1
and Table 2. In Section 4.3, Figure 17 and Figure 18 present the detailed results corresponding to Figure 4. Specifically, Table 17 aligns with Figure 4(a) while Table 18 corresponds to Figure 4(b). In Section 4.4, the detailed experimental results of Figure 5 are provided in Table 16. In Section E.1, we add the breakdown analysis of the prefill and decoding wall-time for various methods, while the results are illustrated in Figure 6. The detailed experimental results are provided in Table 13. In Section 1, the detailed experimental results of Figure 1 are provided in Table 14. The anchor length and passing length are set to 1024 and 512 for 32K input length, 2048 and 1024 for 64K input length, 4096 and 2048 for 128K-1024K input length.

<table>
  <tr>
    <th></th>
    <th>Prefill Time</th>
    <th>Decoding Time</th>
  </tr>
  <tr>
    <td>FLASHATTN</td>
    <td>30137.03</td>
    <td>422.31</td>
  </tr>
  <tr>
    <td>ULYSSES</td>
    <td>4028.66</td>
    <td>263.52</td>
  </tr>
  <tr>
    <td>RINGATTN</td>
    <td>6317.03</td>
    <td>185.33</td>
  </tr>
  <tr>
    <td>MINFERENCE</td>
    <td>18067.32</td>
    <td>2275.61</td>
  </tr>
  <tr>
    <td>STARATTN</td>
    <td>3556.60</td>
    <td>208.23</td>
  </tr>
  <tr>
    <td><b>APB</b></td>
    <td>2554.77</td>
    <td>284.21</td>
  </tr>
</table>

Table 13: The accurate time of Figure 6. All the time is reported in “ms”. We breakdown the wall-time of Transformer inference into prefill and decoding time.

<table>
  <tr>
    <th></th>
    <th>32K</th>
    <th>64K</th>
    <th>128K</th>
    <th>256K</th>
    <th>512K</th>
    <th>1024K</th>
  </tr>
  <tr>
    <td>FLASHATTN</td>
    <td>3.46</td>
    <td>9.51</td>
    <td>30.01</td>
    <td>OOM</td>
    <td>OOM</td>
    <td>OOM</td>
  </tr>
  <tr>
    <td>ULYSSES</td>
    <td>0.51</td>
    <td>1.30</td>
    <td>3.95</td>
    <td>13.49</td>
    <td>49.55</td>
    <td>OOM</td>
  </tr>
  <tr>
    <td>RINGATTN</td>
    <td>0.72</td>
    <td>2.00</td>
    <td>6.54</td>
    <td>21.80</td>
    <td>81.62</td>
    <td>OOM</td>
  </tr>
  <tr>
    <td>MINFERENCE</td>
    <td>4.95</td>
    <td>8.72</td>
    <td>15.16</td>
    <td>OOM</td>
    <td>OOM</td>
    <td>OOM</td>
  </tr>
  <tr>
    <td>STARATTN</td>
    <td>0.67</td>
    <td>1.43</td>
    <td>3.50</td>
    <td>9.63</td>
    <td>30.43</td>
    <td>OOM</td>
  </tr>
  <tr>
    <td><b>APB</b></td>
    <td><b>0.49</b></td>
    <td><b>1.09</b></td>
    <td><b>2.79</b></td>
    <td><b>5.53</b></td>
    <td><b>13.39</b></td>
    <td><b>37.60</b></td>
  </tr>
</table>

Table 14: The prefill time of Figure 1. We report the time in “s”. “OOM” represents out-of-memory error.
<table>
  <tr>
    <th>Method</th>
    <th>SG1</th>
    <th>SG2</th>
    <th>SG3</th>
    <th>MK1</th>
    <th>MK2</th>
    <th>MK3</th>
    <th>MV</th>
    <th>MQ</th>
    <th>VT</th>
    <th>CWE</th>
    <th>FWE</th>
    <th>QA1</th>
    <th>QA2</th>
    <th>Avg.</th>
  </tr>
  <tr>
    <td colspan="15">Llama-3.1-BB-instruct</td>
  </tr>
  <tr>
    <td>FLASHATTN</td><td>4277</td><td>4282</td><td>4062</td><td>4317</td><td>4293</td><td>4145</td><td>3908</td><td>3899</td><td>3847</td><td>4203</td><td>4056</td><td>4361</td><td>4382</td><td>4156</td>
  </tr>
  <tr>
    <td>ULYSSES</td><td>30138</td><td>29537</td><td>25214</td><td>29611</td><td>29542</td><td>24749</td><td>20721</td><td>20212</td><td>27411</td><td>19994</td><td>21721</td><td>29547</td><td>31609</td><td>26154</td>
  </tr>
  <tr>
    <td>RINGATTN</td><td>19642</td><td>19521</td><td>17318</td><td>19760</td><td>19622</td><td>17298</td><td>14582</td><td>15078</td><td>18716</td><td>14731</td><td>16244</td><td>20022</td><td>19862</td><td>17876</td>
  </tr>
  <tr>
    <td>MINERENCE</td><td>7666</td><td>7535</td><td>7362</td><td>5780</td><td>5780</td><td>5780</td><td>5780</td><td>5780</td><td>5780</td><td>5780</td><td>5780</td><td>5780</td><td>5780</td><td>6117</td>
  </tr>
  <tr>
    <td>STARATTN</td><td>31836</td><td>32532</td><td>29900</td><td>32214</td><td>25378</td><td>17494</td><td>20770</td><td>22919</td><td>25738</td><td>29301</td><td>19696</td><td>33661</td><td>36765</td><td>36765</td>
  </tr>
  <tr>
    <td>APB</td><td>43930</td><td>44504</td><td>33106</td><td>44502</td><td>43717</td><td>33882</td><td>31328</td><td>25349</td><td>39680</td><td>25580</td><td>27828</td><td>42931</td><td>45664</td><td>37077</td>
  </tr>
  <tr>
    <td colspan="15">Qwen-2.5-14B-instruct</td>
  </tr>
  <tr>
    <td>FLASHATTN</td><td>2175</td><td>1958</td><td>1819</td><td>2054</td><td>1925</td><td>1438</td><td>1426</td><td>1315</td><td>1882</td><td>1352</td><td>1982</td><td>2038</td><td>2045</td><td>1801</td>
  </tr>
  <tr>
    <td>ULYSSES</td><td>16221</td><td>13968</td><td>13761</td><td>15047</td><td>14713</td><td>11345</td><td>10552</td><td>9863</td><td>14312</td><td>11420</td><td>15174</td><td>15642</td><td>15464</td><td>13652</td>
  </tr>
  <tr>
    <td>RINGATTN</td><td>10571</td><td>9558</td><td>9322</td><td>10070</td><td>9892</td><td>7771</td><td>7791</td><td>7487</td><td>9816</td><td>8267</td><td>10505</td><td>10464</td><td>10358</td><td>9375</td>
  </tr>
  <tr>
    <td>MINERENCE</td><td>3391</td><td>2097</td><td>2637</td><td>2408</td><td>2389</td><td>1476</td><td>1845</td><td>1634</td><td>2721</td><td>1465</td><td>2756</td><td>3134</td><td>3010</td><td>2382</td>
  </tr>
  <tr>
    <td>STARATTN</td><td>18226</td><td>16314</td><td>15403</td><td>18285</td><td>15941</td><td>11293</td><td>13469</td><td>12576</td><td>11039</td><td>16059</td><td>16118</td><td>15965</td><td>14821</td><td>14821</td>
  </tr>
  <tr>
    <td>APB</td><td>24576</td><td>21588</td><td>19580</td><td>20732</td><td>23305</td><td>14292</td><td>15638</td><td>10786</td><td>18601</td><td>13469</td><td>19301</td><td>21758</td><td>20339</td><td>18767</td>
  </tr>
  <tr>
    <td colspan="15">Yi-34B-200K</td>
  </tr>
  <tr>
    <td>FLASHATTN</td><td>1194</td><td>1146</td><td>1122</td><td>1147</td><td>1182</td><td>1144</td><td>1078</td><td>1099</td><td>1174</td><td>1140</td><td>1224</td><td>1167</td><td>1186</td><td>1154</td>
  </tr>
  <tr>
    <td>ULYSSES</td><td>8466</td><td>5513</td><td>4427</td><td>5110</td><td>8659</td><td>7498</td><td>6144</td><td>6674</td><td>8152</td><td>7083</td><td>8099</td><td>8130</td><td>8130</td><td>7083</td>
  </tr>
  <tr>
    <td>RINGATTN</td><td>4478</td><td>3895</td><td>3915</td><td>5158</td><td>5814</td><td>5286</td><td>4657</td><td>4862</td><td>5615</td><td>5092</td><td>5685</td><td>5620</td><td>5625</td><td>5054</td>
  </tr>
  <tr>
    <td>MINERENCE</td><td>1644</td><td>1279</td><td>1199</td><td>1340</td><td>1585</td><td>1293</td><td>1012</td><td>1101</td><td>1450</td><td>1210</td><td>1361</td><td>1484</td><td>1461</td><td>1340</td>
  </tr>
  <tr>
    <td>STARATTN</td><td>8554</td><td>7569</td><td>6939</td><td>7690</td><td>7813</td><td>7442</td><td>6933</td><td>6704</td><td>8108</td><td>7086</td><td>7826</td><td>8160</td><td>8151</td><td>7683</td>
  </tr>
  <tr>
    <td>APB</td><td>11989</td><td>10176</td><td>8902</td><td>10030</td><td>12061</td><td>9768</td><td>7910</td><td>8396</td><td>10252</td><td>9135</td><td>10264</td><td>10990</td><td>10777</td><td>10050</td>
  </tr>
</table>

Table 15: The inference speed of APB compared with all the baselines on RULER. The “Avg.” represents the average speed. The highest score in each column is marked in bold. We report the speed in “tok/s”.

<table>
  <tr>
    <th></th>
    <th>QKV Projection</th>
    <th>Retaining Head</th>
    <th>Communication</th>
    <th>Attention</th>
    <th>O Projection</th>
    <th>FFN</th>
    <th>Others</th>
    <th>Transformer Block</th>
  </tr>
  <tr>
    <td>FLASHATTN</td><td>25.33</td><td>–</td><td>–</td><td>664.01</td><td>17.42</td><td>201.44</td><td>32.67</td><td>940.86</td>
  </tr>
  <tr>
    <td>ULYSSES</td><td>3.31</td><td>–</td><td>3.90</td><td>84.53</td><td>2.27</td><td>25.88</td><td>4.62</td><td>124.51</td>
  </tr>
  <tr>
    <td>RINGATTN</td><td>3.31</td><td>–</td><td>18.40</td><td>452.12</td><td>2.09</td><td>24.40</td><td>4.42</td><td>203.19</td>
  </tr>
  <tr>
    <td>MINERENCE</td><td>24.45</td><td>–</td><td>–</td><td>281.39</td><td>15.88</td><td>201.56</td><td>40.80</td><td>564.07</td>
  </tr>
  <tr>
    <td>STARATTN</td><td>6.67</td><td>–</td><td>–</td><td>41.84</td><td>4.29</td><td>50.01</td><td>8.56</td><td>111.37</td>
  </tr>
  <tr>
    <td>APB</td><td>4.01</td><td>1.72</td><td>0.62</td><td>34.07</td><td>2.67</td><td>30.76</td><td>6.33</td><td>80.18</td>
  </tr>
</table>

Table 16: The accurate time of Figure 5. All the time is reported in “ms”. We breakdown the wall-time of each Transformer block into 7 components. “–” indicates that the time of this component does not exist in the corresponding method.

<table>
  <tr>
    <th>Method</th>
    <th>SG1</th>
    <th>SG2</th>
    <th>SG3</th>
    <th>MK1</th>
    <th>MK2</th>
    <th>MK3</th>
    <th>MV</th>
    <th>MQ</th>
    <th>VT</th>
    <th>CWE</th>
    <th>FWE</th>
    <th>QA1</th>
    <th>QA2</th>
    <th>Avg.</th>
  </tr>
  <tr>
    <td colspan="15">n = 32K</td>
  </tr>
  <tr>
    <td>FULLATTN</td><td>100.00</td><td>100.00</td><td>98.00</td><td>100.00</td><td>96.00</td><td>82.00</td><td>97.00</td><td>98.50</td><td>92.00</td><td>40.20</td><td>88.00</td><td>82.00</td><td>64.00</td><td>87.52</td>
  </tr>
  <tr>
    <td>MINERENCE</td><td>100.00</td><td>100.00</td><td>100.00</td><td>100.00</td><td>96.00</td><td>76.00</td><td>95.50</td><td>99.00</td><td>90.40</td><td>59.40</td><td>88.00</td><td>80.00</td><td>62.00</td><td>88.18</td>
  </tr>
  <tr>
    <td>STARATTN</td><td>100.00</td><td>100.00</td><td>100.00</td><td>96.00</td><td>98.00</td><td>83.50</td><td>93.50</td><td>88.80</td><td>76.20</td><td>90.67</td><td>78.00</td><td>62.00</td><td>89.44</td><td>89.44</td>
  </tr>
  <tr>
    <td>APB</td><td>100.00</td><td>98.00</td><td>100.00</td><td>100.00</td><td>98.00</td><td>100.00</td><td>89.50</td><td>98.50</td><td>86.40</td><td>78.00</td><td>90.00</td><td>74.00</td><td>60.00</td><td>90.18</td>
  </tr>
  <tr>
    <td colspan="15">n = 64K</td>
  </tr>
  <tr>
    <td>FULLATTN</td><td>100.00</td><td>100.00</td><td>98.00</td><td>100.00</td><td>98.00</td><td>56.00</td><td>99.00</td><td>98.00</td><td>84.40</td><td>1.20</td><td>78.67</td><td>68.00</td><td>54.00</td><td>79.64</td>
  </tr>
  <tr>
    <td>MINERENCE</td><td>100.00</td><td>100.00</td><td>100.00</td><td>100.00</td><td>98.00</td><td>54.00</td><td>97.50</td><td>99.50</td><td>78.00</td><td>6.20</td><td>81.33</td><td>64.00</td><td>58.00</td><td>79.73</td>
  </tr>
  <tr>
    <td>STARATTN</td><td>100.00</td><td>94.00</td><td>100.00</td><td>96.00</td><td>96.00</td><td>86.00</td><td>81.50</td><td>94.50</td><td>82.40</td><td>16.00</td><td>82.00</td><td>70.00</td><td>52.00</td><td>80.80</td>
  </tr>
  <tr>
    <td>APB</td><td>100.00</td><td>96.00</td><td>100.00</td><td>100.00</td><td>98.00</td><td>100.00</td><td>95.00</td><td>97.50</td><td>80.40</td><td>11.40</td><td>85.33</td><td>72.00</td><td>52.00</td><td>83.97</td>
  </tr>
  <tr>
    <td colspan="15">n = 128K</td>
  </tr>
  <tr>
    <td>FULLATTN</td><td>100.00</td><td>100.00</td><td>100.00</td><td>98.00</td><td>100.00</td><td>36.00</td><td>98.50</td><td>95.50</td><td>77.20</td><td>0.00</td><td>72.00</td><td>68.00</td><td>46.00</td><td>76.25</td>
  </tr>
  <tr>
    <td>MINERENCE</td><td>100.00</td><td>100.00</td><td>100.00</td><td>100.00</td><td>100.00</td><td>46.00</td><td>96.50</td><td>99.00</td><td>74.80</td><td>0.20</td><td>80.00</td><td>76.00</td><td>54.00</td><td>78.96</td>
  </tr>
  <tr>
    <td>STARATTN</td><td>100.00</td><td>100.00</td><td>100.00</td><td>96.00</td><td>90.00</td><td>79.50</td><td>90.50</td><td>84.40</td><td>0.20</td><td>72.00</td><td>68.00</td><td>52.00</td><td>79.12</td><td>79.12</td>
  </tr>
  <tr>
    <td>APB</td><td>100.00</td><td>100.00</td><td>100.00</td><td>98.00</td><td>94.00</td><td>98.00</td><td>97.00</td><td>98.00</td><td>73.20</td><td>0.20</td><td>70.67</td><td>68.00</td><td>52.00</td><td>80.70</td>
  </tr>
  <tr>
    <td colspan="15">n = 256K</td>
  </tr>
  <tr>
    <td>FULLATTN</td><td>100.00</td><td>100.00</td><td>96.00</td><td>94.00</td><td>97.22</td><td>22.00</td><td>92.50</td><td>95.00</td><td>64.00</td><td>0.60</td><td>76.67</td><td>78.00</td><td>44.00</td><td>73.85</td>
  </tr>
  <tr>
    <td>MINERENCE</td><td>100.00</td><td>100.00</td><td>100.00</td><td>96.00</td><td>94.00</td><td>46.00</td><td>91.85</td><td>95.50</td><td>77.20</td><td>0.20</td><td>79.33</td><td>72.00</td><td>48.00</td><td>76.93</td>
  </tr>
  <tr>
    <td>STARATTN</td><td>100.00</td><td>96.00</td><td>100.00</td><td>94.00</td><td>86.00</td><td>66.00</td><td>78.50</td><td>83.00</td><td>70.00</td><td>0.20</td><td>80.67</td><td>68.00</td><td>42.00</td><td>74.18</td>
  </tr>
  <tr>
    <td>APB</td><td>100.00</td><td>96.00</td><td>100.00</td><td>94.00</td><td>94.00</td><td>93.00</td><td>98.00</td><td>47.60</td><td>0.40</td><td>78.00</td><td>70.00</td><td>46.00</td><td>77.77</td><td>77.77</td>
  </tr>
  <tr>
    <td colspan="15">n = 512K</td>
  </tr>
  <tr>
    <td>FULLATTN</td><td>98.00</td><td>98.00</td><td>100.00</td><td>94.00</td><td>76.00</td><td>10.00</td><td>90.50</td><td>96.00</td><td>46.80</td><td>0.60</td><td>86.67</td><td>70.00</td><td>46.00</td><td>70.20</td>
  </tr>
  <tr>
    <td>MINERENCE</td><td>100.00</td><td>100.00</td><td>100.00</td><td>98.00</td><td>76.00</td><td>12.00</td><td>85.50</td><td>91.00</td><td>77.20</td><td>0.40</td><td>83.33</td><td>62.00</td><td>42.00</td><td>71.34</td>
  </tr>
  <tr>
    <td>STARATTN</td><td>100.00</td><td>98.00</td><td>100.00</td><td>88.00</td><td>74.00</td><td>20.00</td><td>71.50</td><td>80.00</td><td>45.20</td><td>0.40</td><td>88.00</td><td>60.00</td><td>38.00</td><td>66.39</td>
  </tr>
  <tr>
    <td>APB</td><td>100.00</td><td>100.00</td><td>100.00</td><td>90.00</td><td>88.80</td><td>80.00</td><td>88.80</td><td>95.00</td><td>25.20</td><td>0.80</td><td>83.33</td><td>72.00</td><td>44.00</td><td>74.33</td>
  </tr>
</table>

Table 17: The task performance of APB compared with all the baselines on RULER across different input length n, where higher score represents better performance. “Avg.” represents the average score.
<table>
  <tr>
    <th>Method</th>
    <th>SG1</th>
    <th>SG2</th>
    <th>SG3</th>
    <th>MK1</th>
    <th>MK2</th>
    <th>MK3</th>
    <th>MV</th>
    <th>MQ</th>
    <th>VT</th>
    <th>CWE</th>
    <th>FWE</th>
    <th>QA1</th>
    <th>QA2</th>
    <th>Avg.</th>
  </tr>
  <tr>
    <td colspan="15">n = 32K</td>
  </tr>
  <tr>
    <td>FLASHATTN</td><td>8955</td><td>8764</td><td>8754</td><td>7369</td><td>8818</td><td>7534</td><td>6259</td><td>6751</td><td>8005</td><td>6614</td><td>8200</td><td>8333</td><td>8121</td><td>7883</td>
  </tr>
  <tr>
    <td>ULYSSES</td><td>42479</td><td>40140</td><td>22599</td><td>39735</td><td>45166</td><td>23239</td><td>14300</td><td>15644</td><td>27447</td><td>15465</td><td>26981</td><td>32685</td><td>25030</td><td>28532</td>
  </tr>
  <tr>
    <td>RINGATTN</td><td>31432</td><td>32855</td><td>16703</td><td>29762</td><td>32016</td><td>19901</td><td>14227</td><td>15243</td><td>21899</td><td>12687</td><td>20532</td><td>26173</td><td>22230</td><td>22743</td>
  </tr>
  <tr>
    <td>MINFERENCE</td><td>4940</td><td>4714</td><td>3116</td><td>4644</td><td>4725</td><td>3018</td><td>1771</td><td>2097</td><td>3765</td><td>2094</td><td>3608</td><td>3874</td><td>3657</td><td>3540</td>
  </tr>
  <tr>
    <td>STARATTN</td><td>33909</td><td>34256</td><td>17630</td><td>34951</td><td>36137</td><td>20604</td><td>14362</td><td>12704</td><td>23505</td><td>14008</td><td>23792</td><td>27637</td><td>24730</td><td>24479</td>
  </tr>
  <tr>
    <td><b>APB</b></td><td>50778</td><td>43160</td><td>18264</td><td>41896</td><td>45641</td><td>21149</td><td>24893</td><td>13909</td><td>30407</td><td>13457</td><td>25016</td><td>29156</td><td>30863</td><td>29892</td>
  </tr>
  <tr>
    <td colspan="15">n = 64K</td>
  </tr>
  <tr>
    <td>FLASHATTN</td><td>6647</td><td>6664</td><td>6026</td><td>6606</td><td>6627</td><td>6208</td><td>5006</td><td>5526</td><td>6334</td><td>5619</td><td>6411</td><td>6344</td><td>6468</td><td>6191</td>
  </tr>
  <tr>
    <td>ULYSSES</td><td>38672</td><td>37904</td><td>22392</td><td>38690</td><td>38257</td><td>24925</td><td>13190</td><td>18511</td><td>28011</td><td>17790</td><td>27585</td><td>28760</td><td>30551</td><td>27095</td>
  </tr>
  <tr>
    <td>RINGATTN</td><td>27484</td><td>27209</td><td>20279</td><td>20889</td><td>20761</td><td>21167</td><td>11823</td><td>16726</td><td>21220</td><td>15897</td><td>22340</td><td>23076</td><td>23959</td><td>22062</td>
  </tr>
  <tr>
    <td>MINFERENCE</td><td>5539</td><td>5384</td><td>3054</td><td>5521</td><td>5171</td><td>4012</td><td>2899</td><td>3531</td><td>4609</td><td>3123</td><td>4364</td><td>4574</td><td>4864</td><td>4287</td>
  </tr>
  <tr>
    <td>STARATTN</td><td>35871</td><td>35840</td><td>21658</td><td>33590</td><td>33620</td><td>22313</td><td>17342</td><td>16487</td><td>28760</td><td>19524</td><td>27363</td><td>26143</td><td>26049</td><td>26643</td>
  </tr>
  <tr>
    <td><b>APB</b></td><td>48822</td><td>46944</td><td>32005</td><td>50754</td><td>51244</td><td>29646</td><td>19554</td><td>19655</td><td>30925</td><td>21003</td><td>28469</td><td>29516</td><td>30056</td><td>33738</td>
  </tr>
  <tr>
    <td colspan="15">n = 128K</td>
  </tr>
  <tr>
    <td>FLASHATTN</td><td>4262</td><td>4233</td><td>3984</td><td>4285</td><td>4209</td><td>4062</td><td>3463</td><td>3766</td><td>4161</td><td>3910</td><td>4330</td><td>4184</td><td>4266</td><td>4086</td>
  </tr>
  <tr>
    <td>ULYSSES</td><td>30195</td><td>29300</td><td>25446</td><td>30020</td><td>30070</td><td>25720</td><td>16690</td><td>20717</td><td>28030</td><td>22484</td><td>28064</td><td>26920</td><td>26951</td><td>26200</td>
  </tr>
  <tr>
    <td>RINGATTN</td><td>20013</td><td>19221</td><td>17226</td><td>19559</td><td>19565</td><td>17264</td><td>13162</td><td>15344</td><td>18501</td><td>16042</td><td>18622</td><td>18715</td><td>18447</td><td>17822</td>
  </tr>
  <tr>
    <td>MINFERENCE</td><td>5389</td><td>5385</td><td>4459</td><td>5427</td><td>5097</td><td>4597</td><td>2275</td><td>3167</td><td>5149</td><td>3762</td><td>4851</td><td>4789</td><td>4734</td><td>4545</td>
  </tr>
  <tr>
    <td>STARATTN</td><td>34734</td><td>34416</td><td>29086</td><td>34711</td><td>33984</td><td>28468</td><td>19095</td><td>20370</td><td>31682</td><td>25456</td><td>31313</td><td>31339</td><td>30678</td><td>29600</td>
  </tr>
  <tr>
    <td><b>APB</b></td><td>46644</td><td>46110</td><td>36421</td><td>42653</td><td>46052</td><td>35107</td><td>18480</td><td>27517</td><td>40831</td><td>30095</td><td>39476</td><td>40254</td><td>37931</td><td>37575</td>
  </tr>
  <tr>
    <td colspan="15">n = 256K</td>
  </tr>
  <tr>
    <td>FLASHATTN</td><td>OOM</td><td>OOM</td><td>OOM</td><td>OOM</td><td>OOM</td><td>OOM</td><td>OOM</td><td>OOM</td><td>OOM</td><td>OOM</td><td>OOM</td><td>OOM</td><td>OOM</td><td>OOM</td>
  </tr>
  <tr>
    <td>ULYSSES</td><td>18365</td><td>18295</td><td>16987</td><td>18385</td><td>18265</td><td>17231</td><td>12866</td><td>15530</td><td>17689</td><td>16117</td><td>17894</td><td>17710</td><td>17735</td><td>17159</td>
  </tr>
  <tr>
    <td>RINGATTN</td><td>11847</td><td>11748</td><td>11335</td><td>11777</td><td>11849</td><td>11431</td><td>9325</td><td>10654</td><td>11642</td><td>10925</td><td>11799</td><td>11661</td><td>11663</td><td>11358</td>
  </tr>
  <tr>
    <td>MINFERENCE</td><td>5046</td><td>5223</td><td>4262</td><td>5027</td><td>4744</td><td>4411</td><td>2315</td><td>2711</td><td>4930</td><td>3705</td><td>4682</td><td>4519</td><td>4559</td><td>4318</td>
  </tr>
  <tr>
    <td>STARATTN</td><td>26455</td><td>26174</td><td>23985</td><td>26096</td><td>26304</td><td>24073</td><td>18374</td><td>18399</td><td>25383</td><td>22303</td><td>24674</td><td>25413</td><td>24953</td><td>24045</td>
  </tr>
  <tr>
    <td><b>APB</b></td><td>34126</td><td>33613</td><td>29983</td><td>34000</td><td>33914</td><td>30404</td><td>19710</td><td>25390</td><td>32728</td><td>26596</td><td>31437</td><td>31574</td><td>32469</td><td>30457</td>
  </tr>
  <tr>
    <td colspan="15">n = 512K</td>
  </tr>
  <tr>
    <td>FLASHATTN</td><td>OOM</td><td>OOM</td><td>OOM</td><td>OOM</td><td>OOM</td><td>OOM</td><td>OOM</td><td>OOM</td><td>OOM</td><td>OOM</td><td>OOM</td><td>OOM</td><td>OOM</td><td>OOM</td>
  </tr>
  <tr>
    <td>ULYSSES</td><td>10227</td><td>10249</td><td>10089</td><td>10299</td><td>10317</td><td>9993</td><td>8910</td><td>9334</td><td>10324</td><td>9797</td><td>10458</td><td>10192</td><td>10259</td><td>10042</td>
  </tr>
  <tr>
    <td>RINGATTN</td><td>6395</td><td>6337</td><td>6246</td><td>6331</td><td>6383</td><td>5253</td><td>5783</td><td>5937</td><td>6342</td><td>6138</td><td>6508</td><td>6332</td><td>6332</td><td>6255</td>
  </tr>
  <tr>
    <td>MINFERENCE</td><td>4796</td><td>4720</td><td>3999</td><td>4468</td><td>4465</td><td>4134</td><td>2261</td><td>2314</td><td>4585</td><td>3613</td><td>4562</td><td>4367</td><td>4332</td><td>4047</td>
  </tr>
  <tr>
    <td>STARATTN</td><td>16841</td><td>16661</td><td>16130</td><td>16556</td><td>16746</td><td>15988</td><td>13936</td><td>14365</td><td>16969</td><td>15737</td><td>16969</td><td>16798</td><td>16600</td><td>16177</td>
  </tr>
  <tr>
    <td><b>APB</b></td><td>28100</td><td>27614</td><td>25938</td><td>28315</td><td>28306</td><td>25773</td><td>19435</td><td>22188</td><td>27541</td><td>24470</td><td>27484</td><td>27087</td><td>27082</td><td>26102</td>
  </tr>
</table>

Table 18: The inference speed of APB compared with all the baselines on RULER across different input length n. “Avg.” represents the average speed. We report the speed in “tok/s”. “OOM” represents out-of-memory error.

<table>
  <tr>
    <th>Method</th>
    <th>SG1</th>
    <th>SG2</th>
    <th>SG3</th>
    <th>MK1</th>
    <th>MK2</th>
    <th>MK3</th>
    <th>MV</th>
    <th>MQ</th>
    <th>VT</th>
    <th>CWE</th>
    <th>FWE</th>
    <th>QA1</th>
    <th>QA2</th>
    <th>Avg.</th>
    <th>tok/s.</th>
  </tr>
  <tr>
    <td>FULLATTN</td><td>100.00</td><td>98.00</td><td>100.00</td><td>98.00</td><td>98.00</td><td>100.00</td><td>100.00</td><td>97.00</td><td>98.80</td><td>97.20</td><td>94.00</td><td>84.00</td><td>64.00</td><td>94.54</td><td>6247.58</td>
  </tr>
  <tr>
    <td>APB</td><td>100.00</td><td>92.00</td><td>100.00</td><td>98.00</td><td>100.00</td><td>94.00</td><td>100.00</td><td>99.50</td><td>98.80</td><td>98.00</td><td>95.33</td><td>90.00</td><td>68.00</td><td>94.89</td><td>6597.47</td>
  </tr>
</table>

Table 19: The task performance and inference speed of APB and FULLATTN (FLASHATTN) on RULER, where input length is set to 4K tokens. We report the speed in “tok/s”. “Avg.” represents the average score.

<table>
  <tr>
    <th>Method</th>
    <th>SG1</th>
    <th>SG2</th>
    <th>SG3</th>
    <th>MK1</th>
    <th>MK2</th>
    <th>MK3</th>
    <th>MV</th>
    <th>MQ</th>
    <th>VT</th>
    <th>CWE</th>
    <th>FWE</th>
    <th>QA1</th>
    <th>QA2</th>
    <th>Avg.</th>
  </tr>
  <tr>
    <td>FULLATTN</td><td>99.40</td><td>99.80</td><td>99.60</td><td>98.20</td><td>87.60</td><td>67.00</td><td>94.65</td><td>98.00</td><td>60.98</td><td>71.40</td><td>72.20</td><td>78.20</td><td>41.60</td><td>82.20</td>
  </tr>
  <tr>
    <td>LoCoCo</td><td>30.00</td><td>10.00</td><td>0.00</td><td>0.00</td><td>10.00</td><td>10.00</td><td>5.00</td><td>7.50</td><td>24.00</td><td>56.50</td><td>75.00</td><td>40.00</td><td>20.00</td><td>22.15</td>
  </tr>
  <tr>
    <td><b>APB</b></td><td>100.00</td><td>100.00</td><td>99.80</td><td>85.60</td><td>91.00</td><td>89.00</td><td>95.05</td><td>96.40</td><td>51.96</td><td>63.82</td><td>77.33</td><td>70.00</td><td>41.20</td><td>81.63</td>
  </tr>
</table>

Table 20: The task performance of LoCoCo compared with APB and FullAttn on RULER. “Avg.” represents the average score. FULLATTN represents FLASHATTN, RINGATTN, and ULYSSES, as their computational results remain unchanged.

Reviewer paper summaries: They propose a framework to efficiently process long contexts that use approximate attention mechanism over multiple hosts. They increase prefill speed. Their results show that their approach not only decreases the inference time but also increases the performance.
APB (tehir approach) uses a communication mechanism for KV pairs within a sequence parallelism framework. They use FLASHATTN kernel and optimized distribution strategies.
Summary of strengths: - Their results are strong and outperforms previous baselines significantly.
- Besides decreasing inference time they also improve accuracy which is interesting.
- Their approach is novel and builds on top of previous approaches.
- Their experiments and ablation studies are through enough to support their claims.
Summary of weaknesses: - It seems like, they do not compare the performance of the model with the original model (without any optimization). I think that should be added because it will show the accuracy tradeoff for accuracy
- although their approach is designed for long contexts, it could be good to see its effect on very short context to see its generality. In those cases, would it create overhead or still perform in a similar accuracy with similar inference time?
Comments, suggestions and typos: I think you should add accuracy results for the base models without any optimization or approximate attentions.
Overall assessment scores: 4