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

## Abstract

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.
Meanwhile, it remains unclear wheter sparse attention can maintain the model’s quality at a scale of today’s large language models (LLMs), and how this can be achieved.
This paper presents Sparsely-Sharded(S2) Attention, a Triton library that provides kernel optimization for sparse attention customizable at both per-head and per-context-range levels.
S2-Attention enables the exploration of novel and high-performance sparse attention techniques, which we demonstrate through extensive ablations across a wide range of sparse attention deisngs at various model scales.
From these insights, we present several basic guidelines to design sparse attention that can achieve not only practical efficiency improvements, but also strong performance on downstream tasks.
To achieve high parallelization and optimized memory IO, sparse attention should shard the context heterogeneously across attention heads , where each head attends to a different subset of tokens while collectively covering the full context . Meanwhile, we find hybrid architectures combining sparse and dense attention particularly beneficial in practice.
These design choices lead to a novel sparse attention architecture,
which we evaluate with 1.3B, 7B models.
It achieves wall-clock speedup of 8.79X, 15.87X, 25.3X compared to the strong FlashAttention-2 baseline with strong downstream performance on-par with full attention and perfect retrieval performance at a 128k context length.
In inference, for 7B models, our model, with the help of our S2-Attention kernel, achieves 4.5x speed-up compared to dense counterparts.
S2-Attention will be released with easy-to-customize APIs for direct usage in Megatron and vLLM.
We hope they will help future research develop sparse attention algorithms to improve the efficiency of large language models.

## 1 Introduction

Transformer-based LLMs have opened up fresh opportunities to both research and applications (OpenAI, ; Touvron et al., ).
Their quadratic overhead imposes prohibitive overhead in training and serving these models.
For example, training Llama 2 (Touvron et al., ) 70B with a 4K context length on 2 trillion tokens takes 23 days on 2048 A100 GPUs Rucinski ().
When serving the model, the model’s KV cache consumes 343 GB GPU memory with a batch size 32 and 4K context length.
There is an urgent need to train LLMs efficiently and serve them cost-effectively.

Figure: Figure 1: Illustration of S2-Attention with four attention heads on a hypothetical GPU with 4 thread blocks. Each attention head is allocated with a shard of the context.
Refer to caption: x1.png

Figure: (a) Attention Latency Benchmark over FlashAttention-2.
Refer to caption: x2.png

Many works aim to improve efficiency of attention through various sparse attention techniques (Tay et al., ; Child et al., ; Beltagy et al., ; Zaheer et al., ),
where only a subset of the tokens in the context are attended to.
Despite their promising on-paper FLOP savings compared to full-context dense attention, these methods often fail to deliver real-world efficiency gains.
As pointed out by the seminal work FlashAttention (Dao et al., ; Dao, ),
GPU memory access, rather than computation, is the primary bottleneck for attention.
Dense attention has benefited from CUDA-level implementations specifically optimized for more efficient memory IO, an significant optimization that sparse attention methods have yet to receive.
The absence of a flexible, efficient, and easy-to-use library for optimized implementations for sparse attention has become a major roadblock for research in this area, slowing the progress in improving LLMs’ training and serving efficiency.

With Sparsely-Sharded(S2) Attention, we aim to bridge this gap. S2-Attention is 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.
The main challenge in implementing a fused kernel for general sparse attention arises from the sparsity itself. Tiling the Q, K, and V matrices in sparse attention results in idle threads within each tile, as parts of the context are not attended to, leading to reduced thread and SRAM usage while decreasing the FLOP reduction benefits.
S2-Attention addresses this by efficiently tracking the KV usage patterns and dynamically merging sparse query blocks that share common KVs into the same tile, ensuring each tile is fully utilized regardless of sparsity granularity. This approach substantially improves SRAM utilization and minimizes redundant KV loading.
S2-Attention is applicable in both training and inference, substantially lowering the barrier to exploring novel sparse attention architectures, which we explore in Sections 4 and 5.

Based on the insights from our systematic exploration of sparse attention strategies, we propose a novel architecture based on three key principles:
(1) Efficient sparse attention must be designed with hardware and software systems in mind.
Many sparse variants are incompatible with accelerator parallelization and efficient memory access patterns, making it challenging to translate their FLOP savings into real-world efficiency gains.
(2) Our approach uses a hardware-friendly sparse attention pattern that shards the context heterogeneously among attention heads.
Each head applies a unique strided attention pattern to cover a different subset of tokens, while, collectively, they cover the full context (Figure1).
(3) To achieve strong performance on challenging long-context tasks, the model must have direct access to all tokens, at least at certain layers. A hybrid design that combines sparse and dense attention balances efficiency and performance.

To evaluate S2-Attention and our proposed hybrid architecture,
we pre-train a suite of models at 1.3B, 7B scales with different sparse attention, and compare their quality to the dense attention baseline.
Our results show that the 7B hybrid architecture matches the performance of dense attention while achieving a 2.5X training speed-up and 4.5X inference speed-up.
Moreover, we extend the 1.3B models to a 32k context length, and 7B models to 128k.
We show that the hybrid structures can achieve perfect Needle in a Haystack retrieval performance.
Compared to FlashAttention-2 (Dao, ), S2-Attention can achieve as much as 8.79X, 15.87X attention speed-up for 1.3B, 7B scale respectively, and 2.83X,2.47 e2e training time reduction.

S2-Attention is compatible with commonly used LLM frameworks including PyTorch, Megatron, HuggingFace, and vLLM.
With its user-friendly APIs, importing and customizing S2-Attention take no more than several lines of code.
S2-Attention will be publicly released upon publication.

## 2 Related Works

Efficient Transformers Numerous of attempts have been made to make the training and serving of LLMs more efficient.
FlashAttention family (Dao et al., ; Dao, ) are the most widely-adopted attention acceleration framework.
Dao et al. () breaks down the attention computation into smaller blockwise computation to reduce the IO between SRAM and HBM.
Dao () further improve the performance by reducing the non-matmul operations, extra parallelization over sequence length, and better warp organization.
Our implementation is based on the FlashAttention kernel, and leverage its parallelization over the number of heads.

Sparse Self-Attention reduce the computational complexity of self-attention by only attending to a subset of tokens(Child et al., ; Katharopoulos et al., ; Kitaev et al., ; Zaheer et al., ; Beltagy et al., ).
Child et al. () factorized attention computation into local and stride patterns to reduce the computation.
Ho et al. () saves computation by applying multiple attentions along each axis of the input tensor.
Beltagy et al. () adds more sparsity on top of Child et al. () by taking advantage of the receptive field accumulation with dilated sliding window.
Zaheer et al. () replaces the one-to-all attention in transformers with sliding window attention and random attention.
Qiu et al. () chunks the sequence into blocks, which reduces FLOPs by performing attention on a larger granularity.
However, many of these methods can’t bring wall-clock speed-up due to ignorance of realistic memory access cost (Dao et al., ).
Meanwhile, we argue that some sparsity pattern are not inference-friendly for generative models due to inefficient KV caching.
Also, as they underutilized the parallelization over head, there’s actually great room to further reduce FLOPs while improving context coverage.
We will discuss this further in the following sections.

## 3 S2-Attention: Efficiency and Customization

In this section, we present the system implementation of S2-Attention.
We first briefly cover the basic GPU memory and execution hierarchy.
Based on the discussion, we show why sparse attention will have suboptimal performance when naively following the same tiling strategy of FlashAttention series.
We then propose MergeQ in Section 3.3, a technique that can significantly improve the kernel efficiency while allowing more fine-grained sparsity customization by adaptive merging query segments that shares KV.

Figure: Figure 3: Illustration of S2-Attention System Implementation. Left: Naive CSR-based strategy can only efficiently handle very coarse sparse pattern that keep or mask context chunks of block size. Right: MergeQ adaptively merge smaller query shards to tile size, significantly improving the system efficiency and flexibility.
Refer to caption: x4.png

### 3.1 Preliminaries

GPU threads can access several types of memory, including registers, on-chip shared-memory (SRAM), and the global high bandwidth memory (HBM).
Access times vary greatly, with registers being the fastest and the HBM being the slowest.
Inefficient memory access patterns, such as frequent I/O to HBM, can significantly hurt efficiency.
CUDA organizes threads into blocks, and blocks are further divided into warps, a group 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.
Lastly, a tile is a portion of the data assigned to a thread block to handle.
For clarity, we take tile size as block size.
FlashAttention imrpvoes efficiency by minimizing HBM I/O, tiling the Q, K, V matrices into chunks that fit into SRAM for efficient computation (Dao et al., ), a principle we follow below.

Why does sparse attention not fully benefit from fused kernels?
When tiling the Q, K, and V matrices,
some parts of the context are not attended to by the sparse attention.
As a result, some warps may be assigned with tiles where only a small subset of threads are active, reducing the utilization and overall throughput.
The issue is easier to handle if the sparsity is enforced at a granularity close to thread block size, e.g., keep or mask for every 64 tokens, as the non-sparse portion can occupy the whole block without warp divergence.
However, for finer granularities, this strategy will not work, which S2-Attention aims to address.

### 3.2 Overall Kernel Design

As shown in Figure3, for each query vector ${\bm{q}}$, we iterate through the ${\bm{K}}$ tiles in SRAM to compute $\operatorname{softmax}({\bm{q}}{\bm{K}}^{\top})$.
Unlike dense attention, which uses the entire ${\bm{K}}$ matrix, we only consider the keys specified by the sparse attention pattern, defined by an attention mask.
This mask is stored in a Compressed Sparse Row (CSR) format to optimize memory efficiency.

Given a sequence of $N$ tokens, without loss of generality, we first chunk context into a list of blocks of size $S\in[1,N]$, which gives us $B=ceil(N/S)$ blocks.
The attention mask, M, is then a B*B matrix.
Starting from here, to avoid ambiguity with thread blocks in kernel, we refer context blocks as shards.
We use $qshard_{i}$ and $kshard_{i}$ to denote the ith query shard and ith key shard.

For simplicity, we first consider a scenario where the shard size is the same as the tile size of Q, K, V.
As shown in the left side of Figure 3, given a attention mask M, we first transfer it into the CSR format.
Each row can be viewed as the attended KV positions for each qshard, where the corresponding index is specified in the column_idx arrary.
row_offset serves as pointers for each key to find its attended KV in column_idx.
Then when computing $qshard_{i}K^{T}$ in SRAM, we only need to loop the necessary ones according to column_idx.

### 3.3 Merge-Q

The implementation above works and can already generate expected speed-up.
However, the premise here is that the context shard size is the same as the tile size to avoid divergence.
Ideally, we want the context shard size to be smaller to customize more fine-grained sparsity, but kernel block size are usually required to be much larger to increase GPU efficiency.
Such conflicts would greatly limit the granularity of sparsity patterns S2-Attention can support.

Therefore, it is critical to decouple the context shard size and the kernel tile size in the implementation.
To achieve that, we design Merge-Q operation, which merges the context shards with the same K/V blocks but different q shards into a mega shard that forms a tile in the kernel.
Figure 3(b) illustrate a simpler case, where we only merge neighboring q shards.
Thanks the insight in section 4.1 on KV-efficient sparsity, we only need to know the number of
non-empty shards within the mega shard regardless where the empty shards are, as empty blocks are either at the end or out of the causal region.
For implementation, we twisted the CSR definition to use the data array to store how many times a $kshard[i]$ is used in the mega shard.
This help improve the reuse of those K/V shards, as now instead of reloading them into chip multiple times for different qshard, the kernel now only need to do it once.

Merge-Q allows us to support sparse block size as small as 16.
As shown in Figure 4, we see a significant improvement after merging multiple small blocks in the Q dimension.
If the sparse block size is larger than the best kernel block size, we simply break the larger sparse block to a few of smaller blocks that are optimal for GPU.

### 3.4 Split-D

FlashAttention load the whole head dimension(D) at once to SRAM, which is intuitive as the whole vector is needed to compute the attention scores.
However, in our experiment, we found that split over the D dimension in many cases is beneficial.
We suspect that this is because Split-D reduce SRAM usage and enable a larger degree of software pipelining.
Interestingly, we also find that with shard size of 64, Split-D only helps when head dimension is 128, the most commonly used case, while has no additional benefit when head dimension is 64 or 256.
We suspect that this is due to under-tuned hyper-parameters like kernel tile sizes.

Figure: (a) Forward w/ Merge-Q.
Refer to caption: x5.png

With the optimized kernel library, the community can customized fine-grained sparsity patterns while enjoying the speed-up.
However, what sparsity pattern can bring speed-up without quality degradation remains unknown.
We seek to answer this question and derisk architecture choices with out insights and experiments.

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

In this section, we first discuss which kind of sparse attention patterns are efficient for decoder-only LLMs.
Building on these insights, we discuss strategies to improve context coverage while achieving more savings in memory and FLOPs.
Lastly, we discuss general guidelines to design efficient and effective sparse attention.

### 4.1 KV-Efficient Sparsity

Figure: Figure 5: (a): The dilated attention as an example of sparse attention that is not KV-efficient. E.g., step 5 attends to KV at positions 1, 3, 5, while step 4 attends to 0, 2, 4. This results in requiring full KV cache. Although it suggests nearly 50% memory savings on paper, it actually requires storing the full KV cache in practice
Refer to caption: x9.png

KV caching is a primary memory bottlenecks for decoder-only LMs at inference time.
Many existing sparse attention determine which tokens to attend to based on relative distances (Child et al., ; Zaheer et al., ; Beltagy et al., ).
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 general principle of designing KV-efficient sparse attention. For $\forall j\geq i,\,l\geq 1$,

$$ $({\bm{k}}_{i},{\bm{v}}_{i})\mbox{ is attended by }{\bm{q}}_{j+l}\implies({\bm{ k}}_{i},{\bm{v}}_{i})\mbox{ must also be attended by }{\bm{q}}_{j}.$ $$

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, the attention pattern should include 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 windowed attention.

### 4.2 Heterogeneous Context Sharding

A simple pattern that meets the principle above is a combination of fixed positions, dubbed vertical stride, and a local window.
Each attention head accesses to a partial context at any given step.
We have different heads should attend to different subsets of equal-sized positions. They also, collectively, cover the full context.
We dub this method heterogeneous context sharding.
Compared to using the same pattern across all heads,
our heterogeneous sharding comes with the same overhead;
but it ensures that the model always has access to the full context.
The flexibility of our S2-Attention kernels facilitates us to implement this strategy efficiently.
As we will show in experiments, this ensures that the model can achieve strong long-context performance while maximizing the efficiency gains.
Figure 1 provides an illustrative example.

For a transformer with $H$ heads, we chunk a sequence of $N$ tokens into $B=ceil(N/S)$ shards.
Among the total $B$ blocks, We take the most recent $N_{l}$ blocks as local blocks and set the rest as remote blocks.
We denote the remote stride as $v_{r}$.
In practice, $v_{r}$ is set to be a factor number of attention heads to satisfy the “union as full context” constraint.
For each head $h\in\{1,...,H\}$, we set a unique position offset $o_{h}$. For simplicity, we set $o_{h}=h$ here, which means the offset of each head equals its head index.

Then for each head $h$, its block attention mask $M^{h}$ in $B\times B$ dimension is as follows:

$$ $M^{h}_{i,j}=\left\{\begin{array}[]{ll}1&q_{i}-k_{j}<N_{l},\\ 1&(k_{j}-o_{h})\in v_{r}\mathbb{Z}_{\geq 0}\,\&\,q_{i}-k_{j}\in[N_{l},B)\\ 0&\text{otherwise},\end{array}\right.$ (1) $$

where $x\in m\mathbb{Z}_{\geq 0}$ mean $x$ is 0 or a positive multiple of $m$.

The first constraint determines the shards within the local window.
The second constraint controls what shard each head attends to outside the local window. Note that if the distance is beyond $B$, the attention is dropped, thus extends similar to a sliding window.

This local-stride pattern can be easily extended to have multiple strides in different context ranges. With blocks $N_{l}<N_{r_{1}}<N_{r_{2}}<B$, and $v_{r_{1}}<v_{r_{2}}$, we define block attention mask as

$$ $M^{h}_{i,j}=\left\{\begin{array}[]{llll}1&q_{i}-k_{j}\in[0,N_{l}),&&\\ 1&q_{i}-k_{j}\in[N_{l},N_{r_{1}})&\&&(k_{j}-o^{1}_{h})\in v_{r_{1}}\mathbb{Z}_ {\geq 0},\\ 1&q_{i}-k_{j}\in[N_{r_{2}},B)&\&&(k_{j}-o^{2}_{h})\in v_{r_{2}}\mathbb{Z}_{ \geq 0},\\ 0&\text{otherwise}.\end{array}\right.$ (2) $$

To ensure that $M_{i,j}^{h}$ is KV-cache efficient, we requires that
$v_{r_{2}}\in v_{r_{1}}\mathbb{Z}_{\geq 0}$ and $(o^{2}_{h}-o^{1}_{h})\in v_{r_{1}}\mathbb{Z}_{\geq 0}$.

### 4.3 Union Completeness, Hybrid Architecture

Despite Heterogeneous Context Sharding, we introduce two other insights that we empirically find useful for better model quality.
With the formulation in Section 4.2, we can derive that the full context can be preserved if the vertical stride size is not larger than the number of attention heads.
We denote such property as Union Completeness, when the union of all the attention heads’ shards generate the full context.

As previous studies shows (Huang et al., ; ), 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.

## 5 Experiment

We evaluate the efficiency advantage and quality preserving S2-Attention can achieve.
We first study the pre-training quality is Section5.1 and Section5.2.
We then benchmark the kernel efficiency and end-to-end serving latency in Section5.3 and Section5.4.
Lastly, we ablate a few sparsity herustics for better understanding.

**Table 1: Training quality evaluation. SWA refers to sliding window attention, while S2 refers to S2-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.**
| Model | Passkey | WinoGrande | piqa | race | wikitext103(ppl) |
| --- | --- | --- | --- | --- | --- |
| Dense | 0.865 | 0.592 | 0.733 | 0.403 | 15.884 |
| Homogeneous |  |  |  |  |  |
| S2-L9(SWA) | 0.334 | 0.547 | 0.705 | 0.363 | 21.997 |
| S2-L9 + Dense | 0.62 | 0.575 | 0.714 | 0.373 | 20.450 |
| S2-L9+sink | 0.560 | 0.566 | 0.721 | 0.380 | 21.037 |
| S2-L9+sink+ Dense | 0.771 | 0.577 | 0.728 | 0.388 | 18.503 |
| S2-L1V15 | 0.542 | 0.541 | 0.716 | 0.352 | 21.035 |
| S2-L1V15 + Dense | 0.741 | 0.568 | 0.713 | 0.349 | 20.579 |
| Heterogeneous & Incomplete |  |  |  |  |  |
| S2-L2V18 | 0.630 | 0.565 | 0.728 | 0.357 | 20.502 |
| S2-L2V18 + Dense | 0.823 | 0.587 | 0.732 | 0.379 | 18.726 |
| S2-L4V25 | 0.612 | 0.542 | 0.720 | 0.352 | 20.875 |
| S2-L4V25 + Dense | 0.795 | 0.569 | 0.724 | 0.386 | 19.285 |
| Heterogeneous & Complete |  |  |  |  |  |
| S2-L1V15 | 0.782 | 0.571 | 0.724 | 0.361 | 19.551 |
| S2-L1V15 + Dense | 0.941 | 0.587 | 0.725 | 0.397 | 17.183 |

### 5.1 Benchmarking Model Pre-Training Quality

Experiment Settings. To study the S2-Attention, we first train a range of 1.3B Llama 2 structure models with 24 layers, 2048 hidden size with 16 heads, with max sequence length as 8192.
We use the open-source FineWeb-Edu-350B Penedo et al. () 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 Yang et al. () 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.

We use a model with all default 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 th sliding window attention(SWA), as the control set.
We add different heuristics to the SWA to see how they affect the pre-training result.
The treatment sets are grouped into Homogeneoeous, Heterogeneous&Incomplete, and Heterogeneous&Complete for comparison.

From Table 1, we can observe the hybrid architectures shows promising results.
As we can see from S2-L1V15+Dense 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, S2-Attention can achieve much better performance compared to the dense model.
This observation work as an initial demonstration of the context understanding ability of the S2-Attention design.
We’ll discuss more the continual pre-training section.

Across Homogeneous and Heterogeneous settings, adding two dense layers will translate to significant performance bonus.
Within the Homogeneous group, in addition to dense layers, we can observe adding attention sink can significantly boost training quality, compared to only using SWA.
In the Heterogeneous&Incomplete group, the vertical stride size is bigger than the number of attention heads, making the context incomplete after union.
For the Heterogeneous&Complete group, we tune the stride size and local window so that it’s just covering the full context while having the same FLOPs as the SWA control set.
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 Pre-training

Figure: (a) l31v32+dense15-17.
Refer to caption: x10.png

In this section, we further examine how to adapt sparse attention to longer contexts.
Following the settings in previous section, we extended the best 1.3B model, S2-L1V15 + Dense, to 32k context length.
We also continual pre-train Llama 2 to 128k context.
We respectively changed the RoPE base to 1,00,000 and 5,000,000 for the 1.3B model and 7B model.
Both models are continual trained with 10B tokens following the recipe in Fu et al. ().
We evaluate the extended models on the Needle In A Haystack (Kamradt, ).

As shown in Figure2(b), the 1.3B model can achieves perfect recall on the 32k context window without any further changes.
For the 7B models, however, we need to tune the parameters to get a perfect score.
Restricted by timing and computation resources, we mainly alternate the dense layer numbers and insertion depth, as we believe they are the quickest way to minimize the gap between sparse and dense architectures.
As shown in Figure6, for 128k context, with a vertical stride size of 32, the model can retrieve the full context with 8 dense layers but fails to do so with only 2 and 4 dense layers.
For a model with 24 vertical stride, the model can get perfect score with 12 dense layers.

The results validate the long context capability of S2-Attention design.
We believe the density could be further trimmed given more tuning, which we leave for future work.

### 5.3 Training Speed-up

#### 5.3.1 Attention Operation Benchmark

Benchmark Settings We measure the attention runtime of S2-Attention and FlashAttention-2 on an A100 80GB GPU for different context length, number of head, and head dimension settings.

Figure: (a) 1.3B Attention Speed-up.
Refer to caption: x13.png

In Figure7 and Figure2(a) We benchmark the speed-up brought by S2-Attention in 1.3B, 7B, 70B model sizes across different sequence lengths to showcase the scalability of our system.
For all the model sizes, S2-Attention can achieve multiple times of speed-up over FlashAttention-2.
For 70B models with 64 heads, S2-Attention can give 25.3X end-to-end speed-up.
For example, in 1.3B models with a vertical stride of 16, S2-Attention can achieve as much as 8.79X speed-up.
As the max sequence length grows longer, the speed-up gradually approximate the theoretical FLOPs reduction benefits.
The overall boost is hedged due to our less optimized backward kernel, which leaves room for further improvement.

#### 5.3.2 End-to-end Training Speed-up

Figure: (a) 1.3B end-to-end training speed-up.
Refer to caption: x15.png

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 models, S2-Attention can get 1.22X, 1.85X, 2.3X, and 2.83X token throughput on 8k to 128k context compared to FlashAttention-2.
For 7B models, S2-Attention can get 1.12X, 1.24, 1.47X, and 2.47X token throughput improvment.
Namely, we can half the time for training a 128k context length model.

### 5.4 Inference Speed-up

Figure: Figure 9: End-to-end latency speed-up for 7B size Llama architecture models.
Refer to caption: x17.png

In order to demonstrate the inference improvments of S2-Attention, we measure the end-to-end latency over different context length settings.
To make comparison realistic, our experiments are done on vLLM (Kwon et al., ).
We choose the FlashAttention-2 backend in vLLM as baseline for fair comparison, as the inference kernel of 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 9, S2-Attention can achieves 1.1X, 1.2X,1.5X, 1.9X, 2.9X, 4.5X speed-up on 8k, 16k, 128K, 256K context.

### 5.5 Ablation

In this section, we ablate a few herustics for a better understanding of effective sparsity design.
We do so by starting from the same S2-Attention with 8 local shards and a vertical stride of 16 shard.
We then compare the validation perplexity and passkey result after they are trained with 300B tokens.

#### 5.5.1 Dense Layer

**Table 2: Number of dense layers and insertion location Ablation**
| Model | L8V16 | +Dense1,2 | +Dense23,24 | +Dense12,13 | +Dense12,18,24 | +Dense8,16,24 | +Dense18-24 |
| --- | --- | --- | --- | --- | --- | --- | --- |
| PPL | 11.72 | 11.49 | 11.51 | 11.49 | 11.49 | 11.49 | 11.82 |
| Passkey | 0.620 | 0.840 | 0.916 | 0.945 | 0.964 | 0.923 | 0.824 |

We first study the effect of dense layers as shown in Table 2.
We can see adding dense layers can usually directly improve the valiation perplexity and passkey results, except when we put all the last 6 layers dense.
Different from previous studies  Huang et al. (), where attention layer are more dense at the start and last layers, we find it’s actually more beneficial to put dense layers around the middle during pre-training.

#### 5.5.2 Retrieval Head

**Table 3: 7B pre-training results. S2 refers to S2-Attention. L refers to number of local blocks. V refers to the vertical stride size.**
| Model | L8V16 | +Retrieval Head | +2*Retrieval Head |
| --- | --- | --- | --- |
| PPL | 11.72 | 11.61 | 11.61 |
| Passkey | 0.620 | 0.880 | 0.865 |

Wu et al. () propose the idea of retrieval head, which are attention head scattering attention around the context.
Retrieval head is found to be crucial to long context understanding tasks, and costs fewer compute and memory compared to dense layers.
In Table3, we find adding 1 retrieval head can significantly improve the passkey performance and beneficial to perplexity.
However, adding more of them doesn’t further improve the model quality.

## 6 Conclusion

In this paper, we introduce S2-Attention, a framework that sparsely shards the context for different attention heads to divide-and-conquer.
Experimental results show model trained with S2-Attention can achieve promising performance on long context tasks with reduced context for each head.
We back S2-Attention with a highly optimized kernel library, we can get equivalent model quality while achieving speed-up over FlashAttention-2 linearly increasing over the number of attention heads.
We open-sourced 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 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.

## 7 Acknowledgement

We would like to thank Liyuan Liu from Microsoft Research, and Yao Fu from University of Edinburgh for insightful discussions.
We also would like to thank Zhuohan Li, Simon Mo, and Kaichao You from UC Berkeley and Tsinghua University for sharing expertise on the vLLM codebase.

## Appendix A Appendix

You may include other additional sections here.

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 3
presentation: 2
questions: 1. How does the efficiency of the proposed sparse attention kernel compare to that of other sparse attention methods?
2. In what ways does the proposed sparse attention pattern adapt to varying input lengths, and what is its performance across different lengths?
3. Given that the kernel dedicatedly optimized for fine-grained sparse patterns, how does the fine-grained sparse pattern impact model accuracy? What are the performance-efficiency trade-offs at different granularities for the proposed kernel compared to others?
soundness: 2
strengths: 1. Useful libarary. The paper implements a practical sparse attention GPU kernel library that supports both training and inference. The flexibility to support fine-grained sparse patterns can benefit future research towards more effective and efficient sparse pattern design.
2. High efficiency. With the optimized sparse attention kernel, the paper shows speedups of up to 25.3 and 4.5 times for training and inference over the dense FlashAttention baseline.
summary: This paper presents a Triton-based GPU kernel library designed to enhance the efficiency of sparse attention training and inference. By merging query blocks and splitting along head dimensions, the library improves GPU warp utilization, particularly for fine-grained sparse attention patterns. With this library, the authors propose a KV-cache-efficient heterogeneous sparse attention method, showing the performance and efficiency benefits of the library.
weaknesses: 1. The main concern of the paper lies in the proposed sparse attention pattern design. The proposed KV-Cache design principle seems overly conclusive and conflicts with existing works.

    a. The principle itself is not novel; similar sparse pattern designs for KV-Cache optimization have been explored extensively in prior studies, such as [1, 2]. Furthermore, recent work on retrieval-based KV-Cache reduction [3] demonstrates high performance despite contradicting this principle. It would be beneficial for the authors to revise their claims to improve rigor and acknowledge alternative approaches.

    b. The deduction of the claim "the sparse patterns should be based on absolute positions rather than relative ones" is confusing. It is unclear why the "vertical line" sparse pattern leads to absolute positions, and why windowed attention is the only exception.

2. The advantages of the proposed heterogeneous context sharding pattern over existing designs are not clearly shown. Additionally, it is unclear how this pattern adapts to varying input lengths, as the ranges of context shards appear pre-defined and fixed.

3. The performance comparison with other sparse attention kernels is not shown. Beyond dense attention methods, it would be valuable to assess the proposed kernel’s performance against other sparse attention methods with GPU kernels, such as those optimize prefill [4] and decode [5].

4. Minor writing issues: 

    a. Gramma and typo: Spelling mistake in section 3.1: imrpvoes -> improves; in section 5.1: th -> the

    b. Clarity: In the abstract, it should be specified that the “8-25x speedup” refers to training time.

[1] Xiao, Guangxuan et al. “Efficient Streaming Language Models with Attention Sinks.”, ICLR'24
[2] Zhang, Zhenyu (Allen) et al. “H2O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models.”, NeurIPS'23
[3] Tang, Jiaming et al. “Quest: Query-Aware Sparsity for Efficient Long-Context LLM Inference.”, ICML'24
[4] Jiang, Huiqiang et al. “MInference 1.0: Accelerating Pre-filling for Long-Context LLMs via Dynamic Sparse Attention.”, arXiv'24
[5] Fu, Tianyu et al. “MoA: Mixture of Sparse Attention for Automatic Large Language Model Compression.”, arXiv'24

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: 1. I am also a researcher specializing in hardware optimizations for Transformers, and I recognize the importance of sparse attention. Due to the self-attention mechanism in Transformers, there is significant redundancy since token-to-token similarities vary in each head’s activation, and sparse attention can help reduce this. However, because operations across heads are independent in Transformers, applying different attention sparsity patterns to each head does not traditionally contribute much to overall model acceleration. Each head has a different computational overhead, leading to an imbalance. Ultimately, the runtime is dictated by the head with the highest computational complexity. In Figure 1, the 'S2-Attention' subfigure seems to illustrate this issue, but it doesn’t appear that the authors discuss this challenge in the library design. Could the authors clarify if this imbalance was considered and, if so, how it was addressed?
2. In Section 3.2, the authors mention that "retaining or masking every 64 tokens" allows the dense portion of sparse attention to fill each block. However, in Section 3.1, they state that "each warp contains 32 threads." The values of 32 and 64 seem contradictory. Is there a relationship between the 64-token blocks and the 32-thread warps, or are these independent design choices?
3. In Figure 2(b), the authors seem to indicate that with a 32K context length, the “depth perfect” remains at 100% regardless of the token limitation. How is “depth perfect” defined in this context? It doesn’t appear to be referenced or defined elsewhere in the paper. Could the authors clarify? The same question for Figure 6(c)(d).
4. In Figure 3, what is the control logic for data movement in MergeQ? Does it follow a fixed pattern? What kind of memory control logic is used in the SRAM for this process? Are the authors implementing in-memory computation within the SRAM, and is the softmax operation also executed within SRAM? Could you please provide a more detailed explanation of the MergeQ process, perhaps with a step-by-step breakdown or a flowchart illustrating the control logic and data movement?
5. Lastly, could the authors clarify how many hardware measurements were conducted to obtain the reported results? Given the inherent variability in hardware performance, multiple measurements are typically necessary to ensure stability and reliability. It would be helpful to understand the stability of the hardware results presented. Could you please provide details on the number of runs performed for each experiment, any measures taken to account for variability (e.g., averaging results, reporting standard deviations), and how you ensured the stability and reproducibility of your hardware measurements?
soundness: 3
strengths: 1. The paper presents a novel approach to improving the real-world efficiency of sparse attention mechanisms in LLMs through S2-Attention, a customizable, hardware-optimized library. Unlike prior sparse attention methods that often fail to deliver actual speedups, S2-Attention effectively addresses the GPU memory access bottleneck. Additionally, the hybrid architecture combining sparse and dense layers is an innovative solution to balance efficiency and model performance.
2. The paper demonstrates high-quality research with extensive benchmarking and thorough evaluations that underscore S2-Attention's performance benefits. The experiments cover a range of model sizes (e.g., 1.3B and 7B models) and contexts, showing both training and inference gains across different configurations. The authors also offer a clear implementation path in Triton, indicating that the work is robustly engineered for real-world applications.
3. By achieving up to 25.3x speedup over FlashAttention-2 and providing easy-to-use APIs for integration into frameworks like Megatron and vLLM, S2-Attention has the potential to become a standard tool for optimizing sparse attention in LLMs.
summary: This paper introduces S2-Attention, a Triton-based library designed to optimize sparse attention in large language models (LLMs). Traditional sparse attention methods reduce FLOPs but fail to deliver real-world speedups due to memory access inefficiencies. S2-Attention addresses this with customized kernel optimizations, enabling flexible and efficient sparse attention patterns at per-head and per-context levels. 

Key contributions include:
1. S2-Attention uses a novel method of context sharding across attention heads, maximizing parallelization and improving memory access efficiency.
2. By combining sparse and dense layers, S2-Attention achieves performance on par with dense attention models while accelerating training and inference.
3. The library achieves up to 25.3x speedup over FlashAttention-2 in 7B models and provides substantial efficiency gains across different model sizes and context lengths.
weaknesses: The paper is innovative in its approach and thorough experimentation. However, there are several critical questions that I raised in the "Question" section, which I believe are essential for the clarity and robustness of the findings. I hope the authors can provide insights on these points, and I look forward to further discussion.

---

confidence: 3
contribution: 2
initial_rating: 3
presentation: 2
questions: My questions are listed in the weakness section.
soundness: 2
strengths: + this work presents a flexible kernel implementation that supports finer-grained sparse attention. Previous work FlashAttention-2 requires the sparsity granularity to be same as the block size, while this work introduces Merge-Q technique to effectively decouple the granularity of sparsity pattern and attention computation while achieving the expected speedup.
+ this work provides a detailed accuracy comparison to demonstrate the effectiveness of heterogeneous context sharing and union completeness.
summary: This paper first introduces a Triton library that implements sparse attention customizable at both per-head and per-context-range levels, and then proposes Sparsely-Shardded (S2) Attention which shards the context heterogeneously across attention heads but collectively covers the full context. The evaluation show that the proposed S2-Attention can achieve 2.5X training speed-up and 4.5X inference speed-up for a 7B transformer architecture design.
weaknesses: - S2-Attention requires training models from scratch, raising concerns about its compatibility with pre-trained models. This limits its flexibility compared to other sparse attention methods  (e.g., QUEST, H2O) that support plug-and-play integration.
- the benefits of supporting finer-grained sparsity remain unclear; if existing block sparse attention methods suffice, the proposed library may be less practical.

---

confidence: 3
contribution: 3
initial_rating: 5
presentation: 2
questions: 1. Although this paper highlights its advantages for large-scale models, especially with high token lengths, it only provides task performance results on Llama2-1.3B with a maximum sequence length of 8192, which somewhat limits the validation of its overall effectiveness.

2. Table 1 lists various S2-Attention variants, each with different task performance. Could the authors clarify which S2-Attention variant was used to measure the speedup and throughput in the subsequent experiments?

3. The paper reports both backward latency and end-to-end training speedup, validating its effectiveness for a single step. However, I would like to know if the sparse models proposed here require the same number of training steps to converge as dense models, so that the benefits extend to the total training time as well.

4. Since memory reduction is another key advantage of sparsity, it would be helpful to include comparisons in terms of GPU memory reduction for the proposed S2-Attention.
soundness: 2
strengths: **Significant Performance Gains**: S2-Attention achieves impressive wall-clock speedups (8.79x, 15.87x, and 25.3x), outperforming FlashAttention-2 and maintaining strong downstream performance.

**Memory and Computation Efficiency**: The method offers higher parallelization and optimized memory I/O, crucial for large-scale models, especially when handling high token lengths.

**Real-World Applicability**: The library is designed to integrate with popular LLM frameworks (e.g., Megatron, vLLM) and provides user-friendly APIs, enhancing its usability for practitioners.
summary: This paper introduces Sparsely-Sharded (S2) Attention, a Triton-based library that optimizes sparse attention mechanisms for large language models (LLMs) by selectively attending to subsets of tokens. It proposes sharding context heterogeneously across attention heads to achieve higher parallelism and memory efficiency while maintaining strong downstream performance. Additionally, the paper demonstrates that hybrid architectures combining sparse and dense attention are particularly effective. Extensive experiments showcase the proposed attention mechanism’s significant wall-clock speedups and accuracy comparable to dense attention methods.
weaknesses: **Lack of Performance Comparison on Larger Models**: While the paper demonstrates hardware benefits on large-scale models, especially with high token lengths, it lacks the corresponding task performance comparisons.

**Paper Organization and Writing Issues**: There are several grammatical issues and unclear figure descriptions, which hinder the readability and understanding of the paper. For example:
1. The caption of Figure 5 only explains Figure 5(a) and lacks descriptions for Figure 5(b).
2. The legends in Figure 8 are too small to read. Additionally, there is an issue with the label on the y-axis of Figure 8(b).

Meta-review: **Summary:** The paper introduces S2-Attention, a Triton-based library for hardware-aware sparse attention optimization, featuring context sharding across attention heads to enhance memory efficiency and parallelization. Using this library, the work highlights the effectiveness of a hybrid architecture that combines sparse and dense attention. Evaluations of S2-Attention with model sizes ranging from 1.3B to 7B demonstrate significant speedups compared to FlashAttention-2, while maintaining competitive performance on downstream tasks.

**Strength:** 

1. The paper proposes a novel hardware-aware context sharding technique that effectively balances computation across attention heads, addressing the gap between theoretical efficiency gains of sparse attention methods and practical wall-clock speedups

2. It demonstrates notable wall-clock speedups compared to strong baselines, with practical applicability to long-context scenarios.

3. This work provides an open-source kernel library with user-friendly APIs, facilitating broader adoption in LLM training and inference pipelines.

**Weakness:**

1. Technical contributions and rationale: The proposed sparse attention pattern design lacks sufficient justification. As noted by Reviewer 8jB7, the principle behind the design is not novel. A more thorough analysis is needed to explain why the proposed design outperforms other designs within the existing design space.

2. Insufficient baseline comparisons: The paper does not provide adequate comparisons to other state-of-the-art sparse attention methods, such as Quest, MInference, and FlashInfer. These comparisons are essential to validate the claimed advantages of S2-Attention.

3. Training cost and applicability: S2-Attention requires training models from scratch, which significantly reduces its flexibility and limits its applicability compared to plug-and-play sparse attention methods. The paper does not sufficiently justify the additional training costs associated with this approach.

4. Clarity and writing issues: Grammatical errors, unclear figure descriptions, and inconsistent terminology detract from the paper's overall readability. Key aspects of the proposed method, such as the definition of "depth perfect" and the specifics of the MergeQ process, are not well explained, making the method less accessible to readers.

**Reasons for the decision:**

While S2-Attention provides a practical library for sparse attention, the lack of thorough analysis of the proposed attention pattern, missing baseline comparisons, and insufficient justification for the additional training costs diminish the paper's overall impact. Consequently, I recommend rejection.

Decision: Reject