S2-Attention: Hardware-Aware Context Sharding Among Attention Heads

Xihui Lin$^{1}$ $^{}$, Yunan Zhang$^{1}$ $^{}$, Suyu Ge$^{2}$ $^{*}$
Liliang Ren$^{1}$, Barun Patra$^{1}$, Vishrav Chaudhary$^{1}$, Hao Peng$^{2}$, Xia Song$^{1}$
$^{1}$Microsoft, $^{2}$UIUC
{xihlin,yunanzhang}@microsoft.com

Abstract

Sparse attention, which selectively attends to a subset of tokens in the context, has been an established approach to enhance the efficiency of Transformers. However, its theoretical reduction in FLOPs has rarely translated into wall-clock speed-up over its dense attention counterparts, mainly due to the lack of hardware-level optimizations like FlashAttention (Dao, 2023). Meanwhile, it remains unclear whether sparse attention can maintain the model's quality at the scale of today's large language models (LLMs), and how this can be achieved. This paper presents Sparsely-Sharded Attention (S2-ATTENTION), an optimized Triton kernel library providing a variety of customizable sparse attention implementations for both training and inference. S2-ATTENTION allows customizing the attention patterns at per head per context range level. The fresh insights from S2-ATTENTION inspire a novel sparse attention architecture that meets several desiderata that we find crucial for achieving both practical efficiency gains and strong accuracy on downstream tasks, called as Head-Heterogenous Strided Transformer (HHST). For higher sparsity, HHST shards the context heterogeneously across attention heads, where each head attends to a different subset of tokens while collectively covering the whole. We evaluate HHST by pretraining 1.3B and 7B sized models. For attention computation, HHST with S2-ATTENTION achieves 8.8$\times$ and 15.9$\times$ wall-clock attention speedup, as well as 2.8$\times$ and 2.5$\times$ training time reduction compared to a dense attention baseline implemented with FlashAttention-2. Moreover, HHST's downstream task performance is on-par with dense attention, and achieves a perfect retrieval accuracy at a 128K context length at 7B scale. At inference, our 7B HHST, achieves a 4.5$\times$ speed-up compared to the dense counterparts in vLLM. S2-ATTENTION is released with easy-to-customize APIs for direct usage in Megatron and vLLM.

$^{*}$Leading Authors. Xihui Lin, Yunan Zhang, and Suyu Ge contribute equally. Code is available at https://github.com/linxihui/dkernel

Executive Summary: Large language models face severe computational and memory bottlenecks during training and inference due to the quadratic cost of standard dense attention. While sparse attention methods reduce theoretical operations, they rarely achieve actual wall-clock speedups because they lack hardware-level memory optimizations and often degrade model accuracy on long-context tasks. Additionally, dynamic inference-time token eviction techniques create memory fragmentation and scheduling overhead in serving frameworks like PagedAttention. To resolve these challenges, the authors developed Sparsely-Sharded Attention (S2-Attention), a hardware-aware Triton kernel library that achieves practical wall-clock acceleration without sacrificing accuracy.

S2-Attention introduces a dynamic query-merging technique called Merge-Q, which groups query blocks that share key-value data into single computation tiles. This reduces redundant high-bandwidth memory access and maintains high shared memory utilization even at fine-grained shard sizes. Guided by hardware efficiency principles, the authors designed the Head-Heterogeneous Strided Transformer (HHST). HHST shards the input context heterogeneously across attention heads so that each head attends to a balanced, strided subset of tokens while all heads collectively cover the complete context at every layer. To retain strong retrieval capabilities, HHST adopts a hybrid model architecture that preserves full dense attention in a small fraction of layers.

The authors evaluated HHST by pretraining and continually training models at the 1.3-billion and 7-billion parameter scales across contexts up to 128,000 tokens. Compared to dense attention implemented with FlashAttention-2, HHST achieved an 8.8-fold attention speedup for 1.3-billion models and a 15.9-fold speedup for 7-billion models. End-to-end pretraining time dropped by up to 2.8-fold, while inference throughput in vLLM showed up to a 4.5-fold latency reduction on long sequences. On downstream benchmarks and challenging retrieval tests, HHST matched the quality of dense models and achieved perfect retrieval accuracy at a 128,000-token context length when configured with sufficient dense layers.

These findings demonstrate that hardware-conscious context sharding can deliver both large empirical speedups and state-of-the-art accuracy across model training and deployment. A noted qualification is that backward kernel optimization remains an ongoing area for improvement to fully capture theoretical gains during training. S2-Attention is integrated for direct use in existing distributed training and serving frameworks, providing an efficient foundation for scaling long-context models.

1. Introduction

Section Summary: Transformer-based large language models offer strong capabilities but face steep costs in training and serving due to their high memory demands and quadratic scaling with context length. Prior sparse attention approaches have struggled to deliver practical speedups because they lack the memory-efficient optimizations that benefit standard dense attention. This work introduces the S2-Attention library and the HHST method to enable flexible, hardware-friendly sparse attention that achieves major efficiency gains in both training and inference while preserving accuracy on long-context tasks.

Transformer-based LLMs have opened up fresh opportunities to both research and applications ([2, 3]). Their quadratic complexity imposes prohibitive cost in the training and serving these models. For example, training Llama 2 ([3]) 70B with a 4K context length on 2T tokens takes 23 days on 2048 A100 GPUs [4]. When serving, the model's KV cache consumes 343GB GPU memory with a 32 batch size and 4K context length. There is an urgent demand for training LLMs efficiently and serving them cost-effectively.

Many established works have managed to, at least on paper, improve the efficiency of these models through various sparse attention techniques ([5, 6, 7, 8]), where only a subset of the tokens in the context are attended to. However, their theoretical FLOP savings compared to full-context dense attention often fail to deliver real-world efficiency gains. As pointed out by the seminal work FlashAttention ([9, 1]), the major overhead in attention arises not from computation but from GPU memory access, especially the shared memory access (SRAM). Dense attention has benefited from CUDA-level implementations specifically optimized for efficient memory IO, a significant advantage that sparse attention methods have yet to receive. The absence of a flexible, efficient, and easy-to-use library for optimized sparse attention implementations has become a major roadblock, delaying progress in both research and applications in improving LLMs' training and serving efficiency.

We aim to bridge this gap with Sparsely-Sharded Attention ($\textsc{S2-Attention}$), a Triton library that provides kernel optimization for sparse attention. It is highly flexible, allowing practitioners to explore various sparse attention strategies and customize different attention patterns across attention heads and context ranges. Building a general-purpose fused kernel for sparse attention presents substantial challenges. In sparse attention, part of the context is not attended to. As a result, tiling the $Q$, $K$, $V$ tensors, a proven technique that divides large tensors into smaller ones for better parallelization and shared memory (SRAM) usage ([9, 1]), can often result in idling threads and inefficient SRAM usage when the tile size is small. $\textsc{S2-Attention}$ addresses this by efficiently tracking KV usage patterns and dynamically merging query blocks with shared KVs into the same tile. This ensures the IO efficiency, regardless of the sparsity granularity, significantly improving SRAM utilization and reducing redundant KV loading.

The insights from the development of $\textsc{S2-Attention}$ reveals that not all sparse attention mechanisms are efficient in practice. Many existing training-free sparse attention, including KV eviction methods such as LongGen ([10]), H2O ([11]), and MInference ([12]), are less compatible with foundational serving mechanisms like continuous batching ([13]), PagedAttention ([14]), 3D parallelism ([15]). For example, in PagedAttention ([14]), evicting tokens from KV blocks would only increase internal fragmentation, and bring extra overhead in scheduling, which in turn hurts serving throughput. Meanwhile, recent studies show that training free sparse attention would hurt model's long context capabilities ([16, 17, 18]). This has become a primary reason why they have limited adoption in industry serving and opens-source inference systems to date ([14, 19]).

These fresh insights lead to Head-Heterogenous Strided Transforme ($\textsc{HHST}$), a new sparse attention approach following key design principles (§ 4.1), which we find crucial for achieving efficiency gains in practice while maintaining strong accuracy on downstream tasks:

• (1) $\textsc{HHST}$ is designed with hardware and software systems in mind. It applies a novel hardware-friendly sharding strategy across attention heads, where each head attends to a distinct set of tokens following a strided pattern, while collectively covering the context in full (Figure 1; § 4.2).
• (2) In order to achieve strong performance on challenging long-context tasks, it is crucial to include direct access to all tokens, at least at certain layers. $\textsc{HHST}$ achieves this with a hybrid architecture that combines sparse and dense attention across layers, and balances efficiency and performance (§ 4.2).

$\textsc{S2-Attention}$ is applicable in both training and inference, substantially lowering the barrier to exploring novel sparse attention architectures, which we explore in § 4 and § 5. We pretrain a suite of models at 1.3B, 7B scales with different sparse attention, and compare them to the dense attention baseline. Our results show that our $\textsc{HHST}$-7B matches the performance of dense attention while achieving a 2.5 $\times$ training speed-up and 4.5 $\times$ inference speed-up. Moreover, we extend the 1.3B models to a 32K context length, and 7B models to 128K. We show that our $\textsc{HHST}$ can achieve perfect Needle in a Haystack retrieval ([20]). Compared to FlashAttention-2 ([1]), $\textsc{HHST}$ can achieve 8.8 $\times$ and 15.9 $\times$ attention speed-up for 1.3B, 7B scales, and 2.8 $\times$, 2.5 $\times$ training wall-clock time reduction.

$\textsc{S2-Attention}$ is compatible with commonly used LLM frameworks including PyTorch, Megatron, HuggingFace, and vLLM. With its user-friendly APIs, importing and customizing $\textsc{S2-Attention}$ take no more than several lines of code as shown in Appendix B.

2. Related Works

Section Summary: Existing attempts to speed up attention by focusing only on selected tokens have largely failed to deliver faster training times because they ignore real memory-access costs on hardware, unlike the widely used FlashAttention approach that processes data in efficient blocks. KV-eviction techniques meant to shrink memory use during inference create scattered free slots that existing paged-attention systems cannot reuse well, raising internal fragmentation and scheduling overhead across multiple devices. In addition, both training-based and training-free sparse methods tend to lose accuracy on long-context tasks and often require query-specific tuning that proves unreliable for varied real-world inputs.

We present our analysis and observations on existing sparse attention attempts in both training and inference.

2.1 Absence of Efficient Sparse Attention Kernel

There have been attempts to reduce the computational complexity of attention by only attending to a subset of tokens ([6, 21, 22, 8, 7]). However, these methods can't bring wall-clock speed-up in training due to the negligence of realistic memory access cost ([9]). [9] breaks down the attention computation into smaller block-wise computation to reduce the IO between SRAM and the high bandwidth memory (HBM). The hardware implementation of FlashAttention family ([9, 1]) make them the most widely-adopted attention acceleration framework. It remains unclear whether we can implement various sparse self-attention in such hardware-aware way, so that the training speed can be further boosted over FlashAttention.

2.2 Issues with Plug-in-and-Play KV Eviction Methods

Recently, plug-in-and-play KV eviction works thrives. More specifically, these methods dynamically drop KV vectors at inference to reduce the memory footprint based on certain criteria that designed to preserve model capabilities.

However, we observe such designs are hardly compatible with existing serving systems, which relies on PagedAttention and continuous batching for efficient memory management. As shown in Figure 3, during KV eviction, the corresponding tokens are release from the physical memory. However, as the token-wise eviction are not guaranteed to be contiguous, the released memory slots can not be effectively allocated for other requests, known as internal fragmentation. In this example, the internal fragmentation increases by 37.5%, which in turn hurts throughput.

Meanwhile, dynamic eviction also introduces overheads in scheduling. For example, if different heads have a different eviction policy/rate, the faster heads will have to wait for the slower ones, which is a classic load-unbalance scenario. The issue is more severe when serving larger models, where computations are distributed across devices and nodes with tensor parallel and pipeline parallel. Such drawbacks further prevent these algorithms from being integrated into real-world serving clusters with hundreds of nodes.

2.3 Performance Degradation

Existing studies points out both the training and training-free sparse attention methods have performance degradation compared to their dense counterparts, especially in long-context tasks. Furthermore, we also observe that some training-free methods([12, 23]) need benchmark-specific hyper parameters to maintain model quality. When applied to unseen requests, the same method can display unpredictable behavior. However, in real-world deployment, user queries often have long tail distribution. Thus, it's not feasible to pre-determine hyper-parameters for unseen user queries, which makes the deployment of such methods risky.

We discuss handling of these observations in sections below.

3. S2-Attention: Efficiency and Customization

Section Summary: S2-Attention improves sparse attention on GPUs by reducing costly data transfers between slow global memory and fast on-chip memory. It builds on FlashAttention-style tiling but adds a Merge-Q method that combines query groups attending the same key-value blocks, letting the system use smaller shards for finer control over which tokens are attended without extra memory loads. The result is customizable sparse patterns that deliver real speedups while preserving efficiency.

This section presents $\textsc{S2-Attention}$. We first briefly review the basics of GPU memory and execution hierarchy, and then introduce our Merge-Q technique, which significantly improves the kernel's efficiency while allowing more fine-grained customization of the sparse attention.

3.1 Preliminaries

GPU threads have access to a hierarchy of different types of memory. Global high-bandwidth memory (HBM) is the slowest but largest (roughly

Poorly-optimized implementations of attention can result in frequent I/O to HBM and significantly hurt the efficiency. CUDA organizes threads into thread blocks, which are further divided into warps, groups of 32 threads. Threads within a block share the data through SRAM. It is desirable that different threads in the same warp take the same execution path since otherwise efficiency will be hurt due to warp divergence. Besides, thread block size should be sufficiently large to achieve good utilization and load balancing. A tile is a portion of the $Q$, $K$, $V$ tensors assigned to a thread block to be processed. For clarity, we take tile size as block size. FlashAttention improves efficiency by minimizing HBM I/O, tiling the $Q$, $K$, $V$ tensors into chunks that fit into SRAM for efficient computation ([9]), a principle that this work follows.

3.2 S2-Attention

Warmup (Figure 4 left)

We first review a simple blocksparse implementation using the FlashAttention algorithm. A sequence of $N$ tokens is segmented into $B=\lceil N/S\rceil$ shards, each containing $S$ consecutive tokens. We use $Q_{[i]}$ to denote the query vectors for the $i$ th query shard, and similarly $K_{[i]}$ the key vectors for $i$ th key shard. Following [9], for each query vector $q$, we iterate through the $K$ tiles in SRAM to compute $\operatorname{softmax}(qK^\top)$. Unlike dense attention that uses the entire $K$ tensor, we only consider a subset of keys specified by a sparse attention mask $M$, which can be stored in a Compressed Sparse Row (CSR) format for memory efficiency.^1

To better understand the efficiency of such implementation, we can calculate the number of loading needed for each key/value shards. As shown in Figure 4 (left), the first key/value shards, $KV_1$, is attened by all the query shards, $q_1-q_8$. Thus, $KV_1$ is loaded 8 times from HBM to SRAM. If we double the shard size, the number of query shards attending $KV_1$ will be halved to 4. In this case, $KV_1$ only needs 4 loading which is more efficient. However, the IO efficiency comes at the cost of granularity of our sparse mask, as we now have to mask-or-keep $2S$ tokens instead of $S$. We then discuss how to achieve both IO efficiency and small mask granularity with Merge-Q.

Merge-Q

At a high level, the core idea is to merge the query shards attending the same $KV$ blocks into a single tile so that we don't need separately load the same KV blocks. In this way, even if the mask granularity becomes smaller, we can still maintain IO efficiency. Figure 4: right display a simpler case, where we merge the neighboring two query shards. Compared to the FlashAttention-2 baseline, this implementation only needs to load $KV_1$ 4 times instead of 8 times with the same mask granularity. Merge-Q helps $\textsc{S2-Attention}$ support shard sizes as small as 16, enabling a broader range of sparse attention patterns. Similar ideas can also be applied to merge KV blocks to further boost efficiency. We leave more detailed implementation discussion in the released code and Appendix D.

With $\textsc{S2-Attention}$, the community can customize fine-grained sparse attention patterns with wall-clock speed-up. However, it remains unclear what types of sparse attention can achieve speed-up without hurting the quality. We aim to answer this question in the following section.

4. S2-Attention: Insights, Formulation, and Sparsity Cookbook

Section Summary: The section explains that many sparse attention methods fail to deliver real memory savings during inference because they rely on relative token positions, forcing the full KV cache to remain stored even when unused. Instead, effective patterns must follow absolute positions so that KV entries can be safely evicted once no longer needed, following a simple implication rule that produces reusable vertical spans in the attention mask. Building on this, the authors introduce the Head-Heterogeneous Strided Transformer, which shards context across heads in a balanced, strided manner and mixes sparse layers with a few dense ones to maintain strong long-context performance while cutting both memory and compute.

In this section, we first discuss which kind of sparse attention patterns allow efficient kernel implementations in practice (§ 4.1). Building on these insights, we introduce Head-Heterogenous Strided Transformer ($\textsc{HHST}$), a novel sparse attention architecture (§ 4.2).

4.1 KV-Efficient Sparsity

KV cache is a primary memory bottlenecks for decoder-only LMs at inference time. Many existing sparse attentions determine which tokens to attend to based on relative distances ([6, 8, 7]). However, these approaches are not GPU memory-efficient during decoding, making it difficult to translate their FLOP savings into real-world efficiency gains. Figure 5(a) provides an illustrative example. The main issue is that, for such sparse attention, KV not used in earlier decoding steps might be required in later ones, making memory management more challenging. Despite the nearly 50% memory saving on paper, it actually requires storing the full KV cache in practice, resulting in zero memory savings.

In contrast, Figure 5(b) illustrates a sparse attention that can achieve memory saving in practice. The key is that the stored KV cache is reused across several decoding steps but is no longer needed in future steps, and thus can be evicted, freeing up the GPU memory.

The comparison between these two approaches leads to the following rule of thumb of designing KV-efficient sparse attention. For $\forall j \geq i, , l \geq 1$,

$ \begin{aligned} &({\bm{k}}_i, {\bm{v}}i) \text{ is attended by } {\bm{q}}{j+l} \ & \implies ({\bm{k}}_i, {\bm{v}}_i) \text{ must also be attended by } {\bm{q}}_j. \end{aligned}\tag{1} $

Otherwise, ${\bm{k}}_i$ and ${\bm{v}}_i$ need to be stored at step $j$ for future generations, even it is not used at step $j$. Intuitively, in the attention pattern matrix, we shall see continuous "vertical lines" as shown in Figure 5(b). This means the sparse patterns should be based on absolute positions rather than relative ones, except for consecutive local context (e.g., left figure in Figure 5(b)).

4.2 Head-Heterogenous Strided Transformer

This section introduces Head-Heterogenous Strided Transformer ($\textsc{HHST}$), a novel efficient sparse attention inspired by the insights we learned above. Core to its design are two design choices introduced below.

Heterogeneous Context Sharding across Attention Heads

To achieve balanced load across attention heads and enhance parallelization, each head should attend to an equal number of tokens. Additionally, $\textsc{HHST}$ ensures that different heads attend to different shards of the context while collectively covering the entire context. This design makes sure that $\textsc{HHST}$ always has direct access to the full context at each layer, without compromising parallelization. Figure 1 provides an illustrative diagram.

More formally, for context with $B$ shards, we take the most recent $B_{l}$ blocks as local blocks and set the rest as remote blocks. For attention head with index $h$, its $B \times B$ block attention mask $M^h$ is:

$ M^h_{i, j} = \left{ \begin{array}{lll} 1, & i-j < B_l, & \text{Local} \ 1, & j - o_h \in s \mathbb{Z}_{\geq 0}, \wedge, i - j \in [B_l, B) &\text{Stride} \ 0 & \text{otherwise} \end{array} \right.\tag{2} $

$s$ is the stride size, and $x \in m\mathbb{Z}_{\geq 0}$ mean $x$ is 0 or a positive multiple of $m$. Similarly to a sliding window, tokens beyond the $B$ shards are not attended to.

The flexibility of our $\textsc{S2-Attention}$ kernel enables an efficient implementation of this strategy. As shown in our experiments, this design allows the model to achieve strong long-context performance while maximizing efficiency gains.

Hybrid Architecture

As previous studies show ([24, 25]), some attention layers are significantly denser compared to the others, with attention weights distributed near uniformly across all positions. Therefore, it is particularly beneficial to retain dense attention in these layers. This motivates us to explore a hybrid architecture that combines our efficient sparse attention in most layers with dense attention in others. We empirically find that our sparse attention strategy is highly effective, requiring only 1/6 of the attention layers to be dense to achieve strong retrieval performance with 128K-long contexts. More exploration is presented in our experiments.

Discussion

It is important to point out that all eviction strategies targeting inference ([11, 26, 10]) are KV-cache efficient, since evicted KV will never be used by future queries. However, these strategies introduce sample-dependent sparsity patterns, making it computationally expensive to determine eviction timing during decoding. In contrast, our approach uses a fixed sparsity pattern across all samples, eliminating the overhead of deciding which tokens to evict. Besides, KV eviction approaches are post-hoc and often perform much poorly as compared to the original dense counterpart ([17]). Our $\textsc{HHST}$, as we will soon see in the experiments, adapts to the sparse attention during training (pre-training or post-training) performs comparably to dense baselines while reducing the training overhead.

5. Experiment

Section Summary: The experiments evaluate HHST by first pre-training 1.3B-parameter Llama-style models on a large text corpus and comparing various hybrid sparse attention designs against a standard dense baseline, finding that the full heterogeneous-and-complete version with a couple of dense layers matches dense performance on most tasks while using only 18% of the compute and even outperforms it on long-context retrieval. They then continually train a 7B model to handle 128K-length inputs, showing that increasing the number of dense layers enables reliable performance on needle-in-haystack tests. Separate kernel benchmarks on A100 GPUs demonstrate large speedups over FlashAttention-2, reaching over 25x for 70B-scale models as sequence length grows.

To evaluate $\textsc{HHST}$, we first study the pre-training quality in § 5.1 and § 5.2. We then benchmark the kernel efficiency and end-to-end serving latency in § 5.4 and §. Lastly, we conduct an ablation study of the design choices.

5.1 Benchmarking Model Training Quality

Settings

We first train a range of 1.3B model with the Llama 2 architecture, with 24 layers, 2048 hidden size with 16 heads, with max sequence length as 8192. We use the open-source FineWeb-Edu-350B [27] as the pre-training corpus. An OpenAI Tiktoken tokenizer with 100K vocabulary size is used to process the raw text. All model variations use batch size of 4M tokens for all sequence lengths and train for a total of 300 billion tokens. For hyperparameters, we use $\mu$ P [28] with a base shape of 256. A $\mu$ P learning rate of 0.02 is used with linear decay and 0.5% of total training tokens for warmup. All models are evaluated after training on the total 300B tokens for one epoch.

Downstream Tasks

We use a model with dense attention as our baseline, denoted as "Dense". To study our hybrid structure with heterogeneous sharding and union completeness, we control the FLOPs to be approximately equivalent. The total attended tokens is around 576 tokens, or 9 shards of 64 tokens. We use this to configure the sliding window attention (SWA), as the control set. We add different changes to SWA to see how they affect the training quality. The treatment sets are grouped into 1) Homogeneous (Different heads attend to the same shards); 2) Heterogeneous & Incomplete (Different heads attend to different shards but not covering the entire context), and 3) Heterogeneous & Complete (Different heads attend to the same shards and collectively cover the entire context).

::: {caption="Table 1: Pre-Training quality evaluation. 'SWA' refers to sliding window attention. '$L'$ refers to number of local blocks. '$V'$ refers to the vertical stride size. '+ Sink' refers attending to attention sink. '+ Dense' refers to making the first two attention layers dense."}

:::

From Table 1, we can observe the hybrid architectures shows promising results. As we can see from S2-L1V15 + Dense ($\textsc{HHST}$) in the last row, heterogeneous sharding with complete context and two dense layers give consistently best results across tasks, with minor gap from the default attention baseline while using only 18% FLOPs. Notably, in the Passkey Retrieval task, $\textsc{HHST}$ can achieve much better performance compared to the dense model. This observation works as an initial validation of the context understanding ability of the $\textsc{HHST}$ design. We'll further validate it in the long context continual training section.

We also found adding two dense layers generally leads to a significantly higher performance. Within the Homogeneous group, we can observe adding attention sink can significantly boost training quality, compared to only using the sliding window (SWA). In the Heterogeneous & Incomplete group, the vertical stride size is bigger than the number of attention heads, making the context incomplete after the union. For the Heterogeneous & Complete group, we tune the stride size and local window so that it just covers the full context while having the same FLOPs as others. When comparing the Incomplete group to the Complete group, we can see the benefits of making the union of context complete by limiting vertical stride size.

5.2 Long Context Continual Training

We further examine how to adapt sparse attention to longer contexts. We start from an existing densely pre-trained model and extend its context length by continually training it on a longer context length with $\textsc{HHST}$ sparse architecture. Specifically, we choose Llama-2-7B and continually train it on 128K context length. We change the RoPE base to 5M. Both models are continually trained with 10B tokens following the recipe in [29]. We evaluate the models on the Needle In A Haystack task [20].

To investigate how to achieve strong long-context performance, we modify the number of dense layers in $\textsc{HHST}$. We set the number of dense layers as 2, 4 and 8, respectively. We fix the number of local blocks to be 31 and vertical stride size to be 32. As shown in Figure 6, for 128K context, the model can retrieve the full context with 8 dense layers but fails to do so with only 2 and 4 dense layers. The results validate the long context capability of $\textsc{HHST}$ design.

5.3 Training Speed-up

5.3.1 Attention Operation Benchmark

Benchmark Settings We measure the attention runtime of $\textsc{HHST}$ with our $\textsc{S2-Attention}$ kernel, and FlashAttention-2 on an A100 80GB GPU for different context length, number of head, and head dimension settings.

In Figure 7 and Figure 2a We benchmark the speed-up brought by $\textsc{HHST}$ in 1.3B, 7B, 70B model sizes across different sequence lengths to showcase the scalability of our system. For all the model sizes, $\textsc{HHST}$ can achieve multiple times of speed-up over FlashAttention-2. For 70B models with 64 heads, $\textsc{HHST}$ can give 25.3 $\times$ end-to-end speed-up. For example, in 1.3B models with a vertical stride of 16, $\textsc{HHST}$ can achieve a 8.8 $\times$ speed-up. As the max sequence length grows longer, the speed-up gradually approximates the theoretical FLOPs reduction benefits. The overall boost is hedged a bit due to our less optimized backward kernel, which leaves room for further improvement.

5.4 Training and Inference Speed-up

We evaluate the end to end training speed-up of the 1.3B and 7B models by measuring the token throughput of both models. All models are trained on 256 A100, with a batch size of 8M tokens and activation checkpointing. For 1.3B, $\textsc{HHST}$ can get 1.2 $\times$, 1.8 $\times$, 2.3 $\times$, and 2.8 $\times$ token throughput on 8K to 128K context compared to FlashAttention-2. For 7B models, $\textsc{HHST}$ can get 1.1 $\times$, 1.2 $\times$, 1.5 $\times$, and 2.5 $\times$ token throughput improvement. In order to demonstrate the inference efficiency improvements of $\textsc{HHST}$, we measure the end-to-end latency over different context length settings. To make comparison realistic, our experiments are done on vLLM ([14]). We choose the FlashAttention-2 backend in vLLM as baseline for fair comparison, as the inference kernel of $\textsc{S2-Attention}$ is also based on vLLM. Both methods are deployed on a single node with 8 A100 80GPU, with tensor parallel size equals 4. We set output length as 128 and vary input length between 16K to 256K. As shown in Figure 8b, $\textsc{HHST}$ can achieves 1.1 $\times$, 1.2 $\times$, 2.9 $\times$, 4.5 $\times$ speed-up on 8K, 16K, 128K, 256K context.

6. Conclusion

Section Summary: The researchers introduced S2-Attention, a library of optimized tools that enable flexible sparse attention patterns for both training and running AI models. These tools helped create guidelines for effective sparse attention designs and led to a new hybrid model architecture called HHST that balances speed and accuracy. They plan to release the library openly so it can be easily used with popular frameworks and integrated into serving systems, encouraging further community exploration.

We presented $\textsc{S2-Attention}$, an optimized Triton kernel library that provides a variety of customizable sparse attention implementations for both training and inference. The insights from $\textsc{S2-Attention}$ led to several principles about designing choices of sparse attention methods to make them efficiency in practice. They inspired a novel hybrid sparse attention architecture that meets several desiderata that we find crucial for achieving both practical efficiency gains and strong accuracy on downstream tasks, called as Head-Heterogenous Strided Transformer ($\textsc{HHST}$). We will open-source our kernel library and make it a plug-in-and-play alternative for FlashAttention-2 module in popular training frameworks like Megatron and Pytorch. We also integrated $\textsc{S2-Attention}$ into vLLM backend for instant serving. Both the training and inference kernels allow users to freely customize their sparsity pattern, facilitating the whole community to study the topic in the future.

Appendix

You may include other additional sections here.

References

Section Summary: This section compiles a list of academic papers and technical reports focused on advancing large language models and transformer architectures. The references cover key developments in efficient attention mechanisms, memory optimization techniques, and systems for handling long sequences or faster model training and inference. They also include major model releases such as GPT-4 and Llama 2, along with surveys and specialized adaptations for improved performance.

[1] Tri Dao. Flashattention-2: Faster attention with better parallelism and work partitioning. CoRR, abs/2307.08691, 2023. doi:10.48550/ARXIV.2307.08691. URL https://doi.org/10.48550/arXiv.2307.08691.

[2] OpenAI. Gpt-4 technical report, 2023.

[3] Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, Dan Bikel, Lukas Blecher, Cristian Canton Ferrer, Moya Chen, Guillem Cucurull, David Esiobu, Jude Fernandes, Jeremy Fu, Wenyin Fu, Brian Fuller, Cynthia Gao, Vedanuj Goswami, Naman Goyal, Anthony Hartshorn, Saghar Hosseini, Rui Hou, Hakan Inan, Marcin Kardas, Viktor Kerkez, Madian Khabsa, Isabel Kloumann, Artem Korenev, Punit Singh Koura, Marie-Anne Lachaux, Thibaut Lavril, Jenya Lee, Diana Liskovich, Yinghai Lu, Yuning Mao, Xavier Martinet, Todor Mihaylov, Pushkar Mishra, Igor Molybog, Yixin Nie, Andrew Poulton, Jeremy Reizenstein, Rashi Rungta, Kalyan Saladi, Alan Schelten, Ruan Silva, Eric Michael Smith, Ranjan Subramanian, Xiaoqing Ellen Tan, Binh Tang, Ross Taylor, Adina Williams, Jian Xiang Kuan, Puxin Xu, Zheng Yan, Iliyan Zarov, Yuchen Zhang, Angela Fan, Melanie Kambadur, Sharan Narang, Aurelien Rodriguez, Robert Stojnic, Sergey Edunov, and Thomas Scialom. Llama 2: Open foundation and fine-tuned chat models, 2023.

[4] Szymon Rucinski. Efficient language adaptive pre-training: Extending state-of-the-art large language models for polish. CoRR, abs/2402.09759, 2024. doi:10.48550/ARXIV.2402.09759. URL https://doi.org/10.48550/arXiv.2402.09759.

[5] Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Metzler. Efficient transformers: A survey. ACM Comput. Surv., 55(6):109:1–109:28, 2023. doi:10.1145/3530811. URL https://doi.org/10.1145/3530811.

[6] Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. CoRR, abs/1904.10509, 2019. URL http://arxiv.org/abs/1904.10509.

[7] Iz Beltagy, Matthew E. Peters, and Arman Cohan. Longformer: The long-document transformer. CoRR, abs/2004.05150, 2020. URL https://arxiv.org/abs/2004.05150.

[8] Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontañón, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, and Amr Ahmed. Big bird: Transformers for longer sequences. In Hugo Larochelle, Marc'Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tien Lin (eds.), Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual, 2020. URL https://proceedings.neurips.cc/paper/2020/hash/c8512d142a2d849725f31a9a7a361ab9-Abstract.html.

[9] Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory-efficient exact attention with io-awareness. In Sanmi Koyejo, S. Mohamed, A. Agarwal, Danielle Belgrave, K. Cho, and A. Oh (eds.), Advances in Neural Information Processing Systems 35: Annual Conference on Neural Information Processing Systems 2022, NeurIPS 2022, New Orleans, LA, USA, November 28 - December 9, 2022, 2022. URL http://papers.nips.cc/paper_files/paper/2022/hash/67d57c32e20fd0a7a302cb81d36e40d5-Abstract-Conference.html.

[10] Suyu Ge, Yunan Zhang, Liyuan Liu, Minjia Zhang, Jiawei Han, and Jianfeng Gao. Model tells you what to discard: Adaptive KV cache compression for llms. In The Twelfth International Conference on Learning Representations, ICLR 2024, Vienna, Austria, May 7-11, 2024. OpenReview.net, 2024b. URL https://openreview.net/pdf?id=88nT0j5jAn.

[11] Zhenyu Zhang, Ying Sheng, Tianyi Zhou, Tianlong Chen, Lianmin Zheng, Ruisi Cai, Zhao Song, Yuandong Tian, Christopher Ré, Clark W. Barrett, Zhangyang Wang, and Beidi Chen. H${}_{2}$o: Heavy-hitter oracle for efficient generative inference of large language models. CoRR, abs/2306.14048, 2023. doi:10.48550/arXiv.2306.14048. URL https://doi.org/10.48550/arXiv.2306.14048.

[12] Huiqiang Jiang, Yucheng Li, Chengruidong Zhang, Qianhui Wu, Xufang Luo, Surin Ahn, Zhenhua Han, Amir H. Abdi, Dongsheng Li, Chin-Yew Lin, Yuqing Yang, and Lili Qiu. Minference 1.0: Accelerating pre-filling for long-context llms via dynamic sparse attention. CoRR, abs/2407.02490, 2024. doi:10.48550/ARXIV.2407.02490. URL https://doi.org/10.48550/arXiv.2407.02490.

[13] Gyeong-In Yu, Joo Seong Jeong, Geon-Woo Kim, Soojeong Kim, and Byung-Gon Chun. Orca: A distributed serving system for transformer-based generative models. In Marcos K. Aguilera and Hakim Weatherspoon (eds.), 16th USENIX Symposium on Operating Systems Design and Implementation, OSDI 2022, Carlsbad, CA, USA, July 11-13, 2022, pp. 521–538. USENIX Association, 2022. URL https://www.usenix.org/conference/osdi22/presentation/yu.

[14] Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph E. Gonzalez, Hao Zhang, and Ion Stoica. Efficient memory management for large language model serving with pagedattention. CoRR, abs/2309.06180, 2023. doi:10.48550/arXiv.2309.06180. URL https://doi.org/10.48550/arXiv.2309.06180.

[15] Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro. Megatron-lm: Training multi-billion parameter language models using model parallelism, 2020.

[16] Guangxuan Xiao, Jiaming Tang, Jingwei Zuo, Junxian Guo, Shang Yang, Haotian Tang, Yao Fu, and Song Han. Duoattention: Efficient long-context LLM inference with retrieval and streaming heads. CoRR, abs/2410.10819, 2024. doi:10.48550/ARXIV.2410.10819. URL https://doi.org/10.48550/arXiv.2410.10819.

[17] Suyu Ge, Xihui Lin, Yunan Zhang, Jiawei Han, and Hao Peng. A little goes a long way: Efficient long context training and inference with partial contexts. arXiv preprint arXiv:2410.01485, 2024a.

[18] Chi Han, Qifan Wang, Hao Peng, Wenhan Xiong, Yu Chen, Heng Ji, and Sinong Wang. LM-infinite: Zero-shot extreme length generalization for large language models. In Kevin Duh, Helena Gomez, and Steven Bethard (eds.), Proceedings of the 2024 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies (Volume 1: Long Papers), pp. 3991–4008, Mexico City, Mexico, June 2024. Association for Computational Linguistics. doi:10.18653/v1/2024.naacl-long.222. URL https://aclanthology.org/2024.naacl-long.222/.

[19] Lianmin Zheng, Liangsheng Yin, Zhiqiang Xie, Chuyue Sun, Jeff Huang, Cody Hao Yu, Shiyi Cao, Christos Kozyrakis, Ion Stoica, Joseph E Gonzalez, et al. Sglang: Efficient execution of structured language model programs, 2024. URL https://arxiv. org/abs/2312.07104, 2023.

[20] Greg Kamradt. Needle in a haystack-pressure testing llms. Github Repository, pp. 28, 2023.

[21] Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are rnns: Fast autoregressive transformers with linear attention. In Proceedings of the 37th International Conference on Machine Learning, ICML 2020, 13-18 July 2020, Virtual Event, volume 119 of Proceedings of Machine Learning Research, pp. 5156–5165. PMLR, 2020. URL http://proceedings.mlr.press/v119/katharopoulos20a.html.

[22] Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net, 2020. URL https://openreview.net/forum?id=rkgNKkHtvB.

[23] Jiaming Tang, Yilong Zhao, Kan Zhu, Guangxuan Xiao, Baris Kasikci, and Song Han. Quest: Query-aware sparsity for efficient long-context llm inference, 2024. URL https://arxiv.org/abs/2406.10774.

[24] Xin Huang, Ashish Khetan, Rene Bidart, and Zohar Karnin. Pyramid-bert: Reducing complexity via successive core-set based token selection. In Smaranda Muresan, Preslav Nakov, and Aline Villavicencio (eds.), Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), ACL 2022, Dublin, Ireland, May 22-27, 2022, pp. 8798–8817. Association for Computational Linguistics, 2022. doi:10.18653/v1/2022.acl-long.602. URL https://doi.org/10.18653/v1/2022.acl-long.602.

[25] Opher Lieber, Barak Lenz, Hofit Bata, Gal Cohen, Jhonathan Osin, Itay Dalmedigos, Erez Safahi, Shaked Meirom, Yonatan Belinkov, Shai Shalev-Shwartz, et al. Jamba: A hybrid transformer-mamba language model, 2024. URL https://arxiv. org/abs/2403.19887.

[26] Zichang Liu, Aditya Desai, Fangshuo Liao, Weitao Wang, Victor Xie, Zhaozhuo Xu, Anastasios Kyrillidis, and Anshumali Shrivastava. Scissorhands: Exploiting the persistence of importance hypothesis for LLM KV cache compression at test time. CoRR, abs/2305.17118, 2023. doi:10.48550/arXiv.2305.17118. URL https://doi.org/10.48550/arXiv.2305.17118.

[27] Guilherme Penedo, Hynek Kydlícek, Loubna Ben Allal, Anton Lozhkov, Margaret Mitchell, Colin Raffel, Leandro von Werra, and Thomas Wolf. The fineweb datasets: Decanting the web for the finest text data at scale. CoRR, abs/2406.17557, 2024. doi:10.48550/ARXIV.2406.17557. URL https://doi.org/10.48550/arXiv.2406.17557.

[28] Greg Yang, Edward J. Hu, Igor Babuschkin, Szymon Sidor, David Farhi, Jakub Pachocki, Xiaodong Liu, Weizhu Chen, and Jianfeng Gao. Tensor programs v: Tuning large neural networks via zero-shot hyperparameter transfer. In NeurIPS 2021, March 2022. URL https://www.microsoft.com/en-us/research/publication/tuning-large-neural-networks-via-zero-shot-hyperparameter-transfer/.

[29] Yao Fu, Rameswar Panda, Xinyao Niu, Xiang Yue, Hannaneh Hajishirzi, Yoon Kim, and Hao Peng. Data engineering for scaling language models to 128k context. In Forty-first International Conference on Machine Learning, ICML 2024, Vienna, Austria, July 21-27, 2024. OpenReview.net, 2024. URL https://openreview.net/forum?id=TaAqeo7lUh.