Fast Attention Over Long Sequences With Dynamic Sparse Flash Attention
Authors: Matteo Pagliardini
Daniele Paliotta
Martin Jaggi
François Fleuret
Venue: NeurIPS 2023 poster (2023)
Keywords: self-attention
large language models
transformers
TL;DR: We propose to dynamic approaches to sparsify the attention matrix, they greatly reduce the attention runtime for large sequences.

Abstract: Transformer-based language models have found many diverse applications requiring them to process sequences of increasing length. For these applications, the causal self-attention---which is the only component scaling quadratically w.r.t. the sequence length---becomes a central concern. While many works have proposed schemes to sparsify the attention patterns and reduce the computational overhead of self-attention, those are often limited by implementation concerns and end up imposing a simple and static structure over the attention matrix. Conversely, implementing more dynamic sparse attention often results in runtimes significantly slower than computing the full attention using the Flash implementation from Dao et al. (2022). We extend FlashAttention to accommodate a large class of attention sparsity patterns that, in particular, encompass key/query dropping and hashing-based attention. This leads to implementations with no computational complexity overhead and a multi-fold runtime speedup on top of FlashAttention. Even with relatively low degrees of sparsity, our method improves visibly upon FlashAttention as the sequence length increases. Without sacrificing perplexity, we increase the training speed of a transformer language model by $2.0\times$ and $3.3\times$ for sequences of respectively $8k$ and $16k$ tokens.

Fast Attention Over Long Sequences With Dynamic Sparse Flash Attention
Matteo Pagliardini*
Daniele Paliotta* University of Geneva daniele, pal.ottaumit............
EPFL
matteo. matteo.pagli www.ipaddndnooppll..cccmm ch
Martin Jaggi EPFL
François Fleuret University of Geneva francois fleurettunige.ch ch
jaggin..aggieepll.coo
Abstract
Transformer-basel language models have found many diverse applications requir- ing them to process sequences of increasing length. For these applications, the causal self-attention-vve which the only component scaling quadratically w.r.t. the sequence length-becomes becomes central concern. While many works have proposed schemes to sparsify the attention patterns and reduce the computational overhead of self-attention, those are often limited by implementation concerns and end up imposing a simple and static structure over the attention matrix. Conversely, im- plementing more dynamic sparse attention often results in runtimes significantly slower than computing the full attention using the Flash implementation from Dao al. (2022). We extend FlashAttention to accommodate large class of attention sparsity patterns that, in particular, encompass key/query dropping and hashing-based attention. This leads implementations with computational complexity overhead and multi-fold runtime speedup on top of FlashAttention. Even with relatively low degrees sparsity, our method improves visibly upon Attentot as the sequence length increases. Without sacrificing perplexity, we increase the training speed of transformer language model by 2.0x and 3.3x for sequences of respectively 8k and 16k tokens.
Introduction
Many methods have been developed to mitigate the quadratic cost of self-attention in Transformers (Vaswani et al., 2017). Some methods attempt to linearize the attention (Beltagy et al., 2020; Wang et al., 2020) by for instance linearizing the softmax operator to take advantage of the associativity of matrix products (Katharopoulos al., 2020). Other methods rely predefined sparse masking of the attention matrix, e.g. to constrain the attention to local temporal neighborhood (Zaheer et al., 2020; Child et al., 2019). While the structure is fixed, it is assumed that information from arbitrary locations in the sequence can still flow through this structure over several layers. All those methods the attention matrix.
impose static implicit explicit constraints
Another promising line work consists in computing dynamic modulation sub-part of the attention matrix. They are based, for instance, dropping keys and queries (Kim et al., 2022) or using geometric hashing of the keys and queries to identify linear cost sub-blocks of the attention matrix that carry most of the weight (Kitaev et al., 2020).
Equal contribution.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

The promising theoretical computational complexity of these methods contrasts with the fact that today's 's most successfully deployed practical models instead rely on vanilla attention, in part thanks to the efficiency of the vanilla attention proposed by Vaswani et al. (2017) their seminal paper, but trades in additional compute for less memory While still avoiding a memory footprint quadratic with the sequence length, it delivers practical speedups of over 5x compared to naive implementation.
(Dao et al., 2022). This implementation is mathematically identical to
Using an attention layer in an autoregressive model-which has been key in the recent remarkable AI reeguurrouggss-reuurrs to make causal. This is achieved by applying a mask to the attention matrix, so that information cannot flow from the future to the past during training.
While FlashAttention can deal with vanilla causal masks, it does not provide enough flexibility to be used for situations where the causal attention mask is not perfectly regular, that is, lower triangular. This in particular prevents using it for models that dynamically drop keys and queries rely on geometric hashing, which results in irregular causal structures as illustrated in Fig. 1 and Fig. 2.
We propose extension of FlashAttention- Sparse Causal Flash Attention (SCFA)- that addresses this constraint. Our contribution is threefold:
We present the SCFA GPU kernel, which relaxes the constraint that the causal mask has to be triangular. This kernel can handle any sparsity pattern that can be expressed with a range of keys per query, and any causal masking in the resulting sub-blocks. See $ 3.
We show that SCFA permits to revisit the promising paradigm of dynamic hash-based attention. We devise an algorithm that builds upon the fundamental idea of Reformer (Kitaey et al., 2020) to restrict the computation of the attention matrix over 'hash collision blocks", but avoids both the high computational cost, and the approximate coverage of the hash collisions. See 3.2.
approach implemented with SCFA that reduces computation by dynami- cally selecting, for each head, keys and queries to be removed from the attention operation, superseding existing methods that limited pruning entire heads entire queries/keys, due to the lack of an efficient fine-grained kernel implementation. See & 3.1.
We propose
Experimental evaluations show that SCFA can efficiently be used for variety of sequence modeling tasks, and that our open-source implementation in the Triton language and compiler (Tillet et al., 2019) significantly outperforms FlashAttention as we increase the sparsity and for longer sequences. Moreover, unlike the hash-based attention introduced in Reformer (Kitaev et al., 2020), our hash- based SCFA not only implements the exact computation, but also has faster runtime (see $ 4.2). Finally, we show that a prototype of query and key dropping can be implemented thanks to SCFA, and that the computational reduction is proportional to the fraction of que pairs dropped (sce 4.3).
2 Related work
State-of-the-arr sequence models have very high computational requirements. As consequence, lot of effort has been invested into developing methods to reduce the memory footprint in Transformers. Many efficient Transformer variants have been developed, with the main goal of taming the quadratic complexity of the attention mechanism (Tay et al., 2020). Several methods rely on kernelized attention (Katharopoulos et al., 2020; Choromanski et al., 2020), while others endow the Transformer with some auxiliary memory to increase the context (Wu et al., 2022; Borgeaud et al., 2021).
In many cases, leveraging sparsity in the attention matrix has proven useful. The Sparse Transformer (Child et al., 2019) works with factorized sparse representation of the attention. They employ several sparse attention patterns, where each output position only computes weightings from a subset of input positions.
The Reformer (Kitaev et al., 2020) uses matrix and allow queries to restrict their context window to keys that collide with the same hash. However, to allow GPU-efficient processing, complex machinery has to be developed where the queries and keys are split into fixed-sized chunks, with the attention being applied only within the chunk and the immediate neighbor.
cunsitive-hashing (LSH) to sparsify the attention

---

QK-sparse
Hash-sparse
Figure Proposed sparsification of the attention matrix for a given attention head. In each depicted attention matrix, black areas indicate coefficients compute, patterned areas those forced to zero due to the causal masking, and white areas coefficients that are ignored. We consider two main dynamic strategies sparsify the left attention matrix. The QK-sparse attention consists of dropping some keys and queries (top, the discarded keys and queries are indicated in red), and the Hash-sparse attention computes hash code for each key and each query, and restricts the attention matrix to blocks of keys and queries of same hash code (bottom, the three hash values are indicated for each key query with the colors blue//reen/red). In both cases, the attention operation must be able to deal with sub-blocks of the attention matrix with a non-triangular causal mask.
Attention introduced by Dao et al. (2022) has recently gained a lot of popularity as an efficient, 10-aware exact attention implementation. FlashAttention uses tiling to avoid materializing the full attention matrix on slow GPU HBM, splitting the computation over blocks of query, key, and value vectors. FlashAttention has already reached wide adoption, as it's now available directly in Pytorch as of version 2.0. Additionally, FlashAttention supports very efficient block-sparse structures.
Bigbird (Zaheer et al., 2020) and Longformer (Beltagy al., 2020) two variants that work with sparsified version of the attention matrix. Both approaches rely on a fixed structure that is independent of the input values, using combination of local, global, and random attention.
Hash Attention. When computing the attention matrix for a T D query tensor Q and a x D key tensor K, we consider the matrix of dot-products QKT, which can become impractical to compute for very long sequences. However, we are only interested in the row-wise soft (K"') meaning that the contribution of the keys every query is dominated by the ones with the highest similarity. Thus, restricting the attention computation to queries and keys with high similarity a natural choice to reduce the computation.
Hash attention, introduced the Reformer (Kitaev et al., 2020), allows to quickly select the closest key vectors for each query using locality-sensitive sensitive-hashing (LSH). In general, the LSH mechanism assigns a hash code to vectors with the requirement that vectors that are close space are mapped to the same hash with high probability. For the hash attention, the Reformer assumes shared query-key space (Q- K). After computing the hashes, the queries are sorted according to their hash bucket. In the sorted attention matrix, pairs that fall into the same bucket cluster near the diagonal. In order to implement the LSH-attention scheme efficiently on GPU, the Reformer splits the queries into fixed-sized chunks. Queries belonging to the same chunk can attend each other and one chunk back. This results a suboptimal mechanism where there is no guarantee that the attention will capture exactly all of the elements that belong to the same bucket (See Fig. 4).
FlashAttention The standard self-atteni attention operation consists of multiplying T x query tensor Q by a T x D key tensor K. to obtain a matching score matrix, which is then rescaled and row- normalized with softmax, to get T attention matrix A. This matrix is then multiplied by T x D' value tensor V to obtain the final result. This is the core operation in standard Mulli-Head Attention layer, where additional operations take place to compute Q, K, and V from the layer's input, and multiple instances of this processing take place in parallel.

---

K
QK-spe
Hash-sparse
Figure 2: SCFA computation patterns. In each depicted attention matrix, black areas indicate coefficients to compute, patterned areas are those forced to zero due to the causal masking, and white areas coefficients that are ignored. The red squares in the bottom matrices show the tiles actually computed by our SCFA kernel. In the regular case (left), this coincides with the behavior of FlashAttention. However, in the case of irregular causal masking due to keys/queries dropping (center) or the case of irregular causal masking and band block sparsity due to hashing (right), FlashAttention does not provide means to compute fine-grain subset of the attention matrix.
The two key contributions of FlashAttention are (1) to compute the attention matrix block-wise, to minimize the transfer of keys and queries to the cache memory as much as possible, and (2) to compute the attention matrix on the fly during both the forward and the backward passes, which faster than retrieving it from memory, and avoids memory footprint quadratic with the sequence length T.
For the generalization that is of concern to this article, we focus on the block computation. In the implementation of FlashAttention, causal masking done by using the row and column indexes of the blocks, and the row and column indexes of the keys and queries in individual blocks: attention blocks are computed fully for any block with query index strictly larger than the key index. For the blocks for which the query index is equal to the key index, a regular lower triangular mask is applied. This is illustrated on Fig. 2, bottom left.
3 Method
develop an efficient CUDA kernel written Triton (Tillet et al., 2019) that maintains the careful memory management of FlashAttention but can handle causal structure defined through arbitrary indexing of the keys and queries. In the case where this indexing consists of a binary decision to drop not the head of a query/key, this corresponds to QK-sparse kernels as described in $ 3.1. In the case where the indexing corresponds to bucket indices e.g. obtained from hashing, this corresponds to our Hash-sparse kernel described in 3.2.
Notations. Input tensors for attention as in Vaswani et al. (2017) are of shape H 7 D, with B being the batch size, H the number of heads, T the sequence length, and D the dimension per head. In the following we take the view of single head and instead consider a query tensor Q of shape To D, and key K and value V tensors of shapes T TRv D. The algorithms described below will run in parallel for all elements of the Cartesian product x H. We split tensors into blocks:
Q A .Q...... Qm) K A Ko..... K,.)- We define tile Tr, A Q.KJ. which corresponds the dot products of subpart of the attention matrix (see Fig. 2).

---

3.1 QK-Sparse Attention
Shrinking the attention matrix. Our QK-sparse attention kernel is best summarized in the first row of Fig. 1. Independently for each head, we decide to keep or drop keys and queries. We then remove dropped keys and queries to create smaller Q", K and Ve tensors. Through this reduction we are left with smaller attention matrix A which still has a causal structure in that indices for the queries and keys are increasing monotonically.
Leveraging non-triangular causal attention structure. Despite the advantageous structure of the smaller attention matrix, existing implementations fail to take advantage of it. Especially, as shown in Fig. 2 bottom-left, FlashAttention can leverage the causal structure when the causal mask is triangular, but does not support any other shape. In the forward pass, FlashAttention is, for each block of queries Qs. processing blocks of keys Kj one after the other, moving along a row of tiles: Tr,,- ................................ Tr. Causality dictates that it is unnecessary to process a tile Ti,; when < j. We cannot follow this rule anymore when working with compact representations. To leverage the causal structure of A.. we build new kernel which gets as additional input vectors qide IR To and kd E RTKv representing the indices of the queries and keys in the original uncompressed tensors. Those are similarly split into blocks: q jdd A [gb :.... Im The condition for tile Tı, to be unnecessary to compute is now to have max (q. < min(kjde)) When processing a block of queries Qi we iterate over the key indices kidee ..... hide to find the index jstop of the first block satisfying that condition. We then know we need to process the tiles Ti.j for E [0. Jstopl. Within each tile Ti,j we in addition apply a local causal mask by comparing indices in q.d2 and kidd, By computing jstop in such a way we can leverage the causal structure and have runtimes matching those of FlashAttention. The backward pass can be adapted in a similar fashion, see App. B for more details.
Computing Q K and Vc requires sorting and allocating new tensors. Moreover, as we drop keys and queries for every attention head, and for every sequence in the minibatch, we are forced to consider the largest sequence of non dropped keys/queries and use padding. However, while reordering and reshaping tensors can be costly, this overhead grows linearly with the sequence length and is largely compensated for larger sequences as we show in 4.3.
Overhead.
Edge cases. Dropping keys and queries result in having stranded queries with no keys. This behaviour is undefined and results in NaNs when using the FlashAttention and naive Pytorch implementations. We solve this issue by modifying how softmax statistics are accumulated during the forward and backward passes and ensure stranded queries default to 0 vectors. see App. B for more details.
3.2 Hash-Sparse Attention
Restructuring attention based on hashes. Independently for each head, we associate a bucket identifier to each key and query. We then need to reorder Q, K, V by sorting them along the sequence length dimension. As shown in the bottom row of Fig. 1, this results in clusters of keys and queries with similar hash index close to the diagonal. If the sorting is stable, i... it preserves ordering of queries and keys when the hash index is the same, then those blocks have a local causal structure in which the original indices (original position in the sequence) of keys and queries is a monotonic function within the block. This brings us in a case very similar to the previous one in section $ 3.1, in that we now have the same structure but scattered by blocks within the full attention matrix.
Taking advantage of the new structure. We would like to take advantage of the block structure and only compute attention for queries and keys falling into the same block while at the same time respecting causality. We adapt the FlashAttention kernel in a very similar way as for our QK-sparse kernel. We now provide additional bucket indices gaash and hash to our kernel. Based on those hash indices, we now find not only the stopping index Jstop but also a starting index Jetart- Jstort is the first index for which some of the indices in gaash are present in khash Jstop is the first index for which all indices in khaahh are strictly larger than indices in hash In second step we refine jstop now based on the indices kids and qidz the updated jstop is the last index j E "istort:Jssto for which max(q.da) 2 min(kjde)) As shown in the last column of Fig. 2, we then only compute tiles Th, for j E [jstart; As for the QK-sparse method, we use gidz and kide to apply causal mask locally for each tile. In addition to the causal mask. we use q hash and k hash to mask interactions
5

---

between keys and queries of different buckets. See App. B for details and to see how to adapt the backward pass in a similar fashion.
Overhead. As for the previous method, sorting and re-ordering Q. and V is inducing some overhead increasing linearly with the sequence length. As shown in our experiments in $ 4.2, this overhead is by large compensated for as the sequence length increases.
100
1000
110
Hash-spose method)
600
Naive hasb-spasse
110
Flush artention
200
5000
10000
15000
20000
25000
30000
35000
Sequence lengish
Figure 3: Comparing several hash-based sparse attention implementations with FlashAttention. Similarly QK-dropping baseed sparsity in Fig. 7, due to the non contrriangular causal mask resulting from re-ordering the tensors based on the hash buckets (see Fig. 1). naive implementation would force the computation of the entire attention matrix before applying custom mask. This results in very large runtimes independent of the number of buckets. On the other hand, our implementation modifies the basic FlashAttention method to compute only what required. While there is cost to reordering the tensors based on the hash buckets, this cost is largely compensated for as the number of buckets nb increases, and as the sequence length increases.
4 Experiments & Results
In this section present experimental setup and results. We show that (i) unlike naive implemen- tations using existing libraries, our dynamic sparsity attention schemes can significantly improve over the Attentit runtime, (ii) this still holds in real-world sequence modeling tasks after factoring in all the non-attention operations, and (iii) it is possible to match sometimes outperform the baselines in terms of perplexity while significantly gaining in speed.
4.1 Experimental Setup
Datasets. We test our hash-based sparsity scheme on MNIST (LeCun et al., 1998) for autoregressive image generation, enwik8 (Hutter, 2012), and Open WebText2 (Gao al., 2020). We experiment with QK-dropping based sparsity on Open WexText2.
Models & Baselines. For our language modeling experiments OpenWebText2, we use a base autoregressive transformer architecture with 12 layers, hidden size of 768, 12 heads of 64 dimensions each. For experiments on sequence length T = 8192, we use batch size of 96 4 x 8 x 2 (batch size 4 with accumulation steps and data parallelism over node). When = 16384 we use a batch size of 30 = 2 x 5 x 3. The resulting models are of around 122M parameters. The goal not being to outperform the state-of-the-ar perplexity, we train for 15k iterations. The attention modules used are either using FlashA Atashtttntiinn for the baselines or of our sparse kernels for methods. To ensure a fair comparison, and similarly to Kitaev et al. (2020), we set the keys equal to normalized queries for all of our models. See App. B for more details.
Hardware. All of our timing experiments with random tensors are done on NVIDIA A100 GPUs, using bflo1116. For our language modeling tasks on OpenTeteext,, we trained using data- parallelism on two or three A 1100s for experiments with sequence lengths of respectively 8192 and 16384. When comparing runtimes in Fig 6 and Fig. 8, we normalize the times by multiplying by the

---

100
1.0
30
Reformace attestion
60
Coveagg
Hustrapara (ouns)
Figure 5: Comparing models us- ing Reformer attention vs our Hash-sparse attention. On the simple sequential MNIST task (predicting pixels sequence), we obtain a comparable perplex- ity the Reformer. On en- wik8 character language model- ing, with T=4096, we outperform the Reformer model by margin.
0.8
0.7
0
20000 40000
60000
20000
10000
60000
Sequence length
Sequomer length
(a) Forward attn. runtimes
(b) Coverage of Reformer attn.
Figure 4: Comparing forward runtimes of attention mod- ules alone. Fig.(a): Reformer attention ensures linear compu- tational complexity w.r.t. the sequence length, outperforming FlashAttention for longer sequences. Fig.(b): However, due to the fixed attention structure, the Reformer misses increas- ing fraction of hash collisions. Our approach outperforms both methods and maintains 100% exact coverage of collisions for all sequence lengths. See App. B.3 and App. C for more details.
Attention
MNIST
enwik8
(ppl L)
(bits/c 4
Reformer Hash-sparse
1.76
3.32
1.67
2.29
number of GPUs used. Comparisons with the Reformer are performed on a single A100 or a single NVIDIA RTX 4090 GPU.
4.2 Hash-based Attention
Hashing mechanism For our experiments, we adopt the same hashing procedure as Kitaev et al. (2020). Namely, we use a shared query-key space, and we disallow queries to attend to themselves. We also adopt the LSH scheme from Andoni et al. (2015). This allows us to pick the number of unique hash codes. We refer to bucket the set of vectors that map certain hash.
Runtime performances in a vacuum. We test our implementation with different numbers of buckets nb and random keys, queries, and values. In these tests, we assume hash bucket is provided for free for each head of each key and query (they are sampled uniformly at random (torch.. randint (0, nb)). In practice, runtime experiments on sequence modeling tasks show that obtaining the buckets can cheap and in no way prevents us from improving the attention runtime (see Fig. 6). We compare with causal FlashAttention over the entire sequence. Importantly, to ensure fair comparison, we take into account pre-processing and post-processing steps required to reshape the tensors for both methods. For our method this includes stable sorting by bucket index and transposing tensors, for the bascline only the transposition is required, see App. B.2 for detailed code. Fig. 3.b summarises our findings. We observe large improvements in runtime the number of buckets nb and the sequence length increases.
Language modeling on Open Open Text2. For sequences of length T = 8192 and T = 16384 we train transformer language models using FlashAttention (F-LM), and identical models replacing only the FlashAttention by our Hash-based sparse attention (H-LM) using nbb 16 hash buckets. In Fig. 6 we see that takes the same number of iterations However, H-LM iterations are 1.8x and 2.3x faster for respectively T 8192 and 7 16384. As result, H-LM models reach a given perplexity much faster than their F-LM counterpart. Interestingly, observe the H-LM models gain speed during training, see App. C for additional details.
H-LM and F-LMM to reach given perplexity.
Comparison with Reformer. We compare the speed and performance of our hash-sparse imple- mentation with the Reformer hashed attention. For all comparisons, we always equalize the average bucket size. Results are summarized in Fig. 4. Benchmarks with random inputs show that both our hash-sparse implementation and the Reformer, as expected, are linear with respect to the sequence length (Fig. 4.a). However, we still achieve a significant speedup thanks to our more efficient kernel. More importantly, Fig. 4.b shows that the fixed attention structure imposed by the Reformer does not allow to capture all the hash collisions, with the coverage decreasing steeply as the sequence length increases. On the contrary, our method is exact and covers every bucket collision in the attention
7

---

70
Valilation Puplexity
60
Validation Perplexity
60
Ite periiii oa GPU
8
50
x3.3
40-
4
30
300
20
T 8192
16381
5000
10000
15000
50
100
F-LM T Sk
H-LM 7 8k
Ferations
Time (houns)
(a) Iter to reach perplexity
(b) Time to reach perplexity
(c) Speed gains
Figure 6: Training Language Models (LM) on Open WebTeee2 (Gao et al., 2020) using our hash-based sparsity (H-LM) or FlashAttention over the entire sequence (F-LM). We train on sequences of length 8192 and 16384 and use 16 buckets for all of our H-LM models. We show that it is possible to use our proposed hash-based sparsity to significantly gain in speed while not compromising the perplexity. similarly as function of the iteration. In fact, H-LM even slightly outperform the baseline. (bg.(b): H-LM reach lower perplexity much faster than their F-LM counterpart. Fig.(b) and (c): H-LM models are significantly faster than F-LM models for a given sequence length. The gap widens as the sequence length increases.
(ag. we see for both sequence lengths the perplexity decreasing
matrix. This is reflected in Table 5: our sparee attention layer outperforms the Reformer attention even for shorter sequences.
4.3 Query/KeyyDrooping Based Attention
Q/K-dropping mechanism used. We show that naively dropping heads for each key and query at random can already yield competitive results while significantly improving the runtime. While better dropping schemes could be devised, they are outside of the scope of this work.
Runtime performances in vacuum. We test implementation with different sparsity ratios, corresponding to the probability of dropping some head associated to a given key or query. We assume that the tensors indicating the dropping of each head of each query and key are given for free, along with some random key, query, and value tensors. To ensure fair comparison, we take into account pre-processing and post-processing steps required to reshape the tensors for both methods, see App. B for more details. For our approach, we hope reducing the size of the key, query and value tensors and computing the attention on those would be faster than using FlashAttention over the entire sequence. For this, the time gained by computing the attention on smaller tensors should be larger than the overhead of re-ordering tensors build those smaller tensors. In Fig. 7.a, we show naive implementation using existing Py PyTocc functionalities only starts to provide speedup when dropping more than 70% of the keys and queries. Fig. 7.b shows that using our proposed implementation provides significant speedups even at relatively low sparsity levels. The linearly increasing cost of reshaping the tensors is rapidly compensated by large gains over the quadratic cost of self-attention.
Language modeling on OpenWebTexte For sequences of length 8192, we train transformer language models using FlashAttention (F-LM), as well as identical models replacing only the FlashAttention by our Q/K-dropping sparse attention (D-LM). We train with several sparsity ratios, dropping 30%. 50%, and 70% of the heads of keys and queries at random. In Fig. 8 we observe that while high sparsity can negatively affect the perplexity, lower sparsity D-LM models are matching LM models in perplexity per iterations, while training nearly twice fast. Importantly, the dropping pattern not static. An interesting approach similar curriculum learning in which we start the training with very large sparsity and reduce it linearly during training is studied in App. C.

---

60
QK-spprse method)
50
Naive QK spp-ee
Flash atsention
3076
1000
10
2075
40
800
20
3076
600
2000
1690
20
10
807
80%
25000
35000
Sequemoe length
Sequence length
(a) Naive implementation
(b) Our implementation
Figure 7: Runtimes the full Flash-attention of Dao al. (2022) and several implementations of Query/Key dropping based sparsity. For this figure we show total times for the forward and backward passes. For sparse methods, we drop random a percentage of keys and queries, this percentage is indicated on the right of each curve. (a). A naive implementation consisting in creating compact representations
the key, value, and query tensors by removing dropped keys and queries. As result, the attention matrix is no longer triangular (see Fig. 1). We call the attention method with a custom but still causal mask. The non- PyTorch scaled, sealedddott_.. ____ddo__ppooo _rrtt__rouuc...m www.uct_tttentttieeiieeeppppeeee triangular mask prevents FlashAttention to be used and only dropping more than 70% of the keys and queries seems to improve the runtime over attending the entire sequence using FlashAttention (b): Our modification of FlashAttention allows improve over the runtime. Similar to the naive implementation, reshaping the tensor induce an overhead which compensates the speed gain for shorter sequences. However, this offset is compensated by strong margin the sequence length increases. Our implementation allows significant gains over FlashAttention even for low levels of sparsity. The detailed runtimes for the forward and backward passes can be found in App. C.
70
Valddttion Parlexity
Validation Perploxity
Ite Iter/ /nin ea 1 GPU
15
x3.5
60
x2.6
50
x1.9
40
5
30-
20
5000
10000
15000
20
20
40
60
F-AM
D-LM
DALM
309t
D-LM
7095
Ferntions
Time (hours)
(a) Iter to reach perplexity
(b) Time to reach perplexity
(c) Speed gains
Figure 8: Training Language Models (LM) on OpenWeb Text2 (Gao et al., 2020) using random Query/Key dropping based sparsity (D-LM) or Ateshtnttntion over the entire sequence (F-LM). Dropping keys and queries randomly is naive and our point here is not to show that this approach good way to use the proposed Q/K-sparsity attention, rather we want to demonstrate that possible to significantly gain in speed while not losing too much in perplexity- even with naive approach, and a very dynamic way (two sequences are allowed to have completely different dropping patterns). For all methods we train over sequences of 8192 tokens. Fig.(a): While dropping large portions of keys and queries slows down the decrease of perplexity per iteration, dropping 30% seems to match the baseline F-LM. (bg. Our method is significantly faster to reach given perplexity. Interestingly, more sparsity does not necessarily mean decreasing the perplexity faster. (b) and (c): Using our O/K-sparse implementation train significantly faster than the bascline method.

---

Conclusion
We develop and validate efficient kernel that can make sparse attention based on dynamic patterns very fast. We hope that our contribution will inspire the community to research dynamic attention patterns way that is less constrained by tight computational budget.
The computational cost of large attention models remains both a practical issue in scaling up to very large contexts, and fundamental research question to close the gap between the energy usage of biological systems to that of GPU systems able to run very large models. Dynamically modulating the computation is obvious direction address this challenge.
6
Acknowledgments
The authors acknowledge support from the Swiss National Science Foundation under grant number CR- SII5-193716 "Robust Deep Density Models for High-Energy Particle Physics and Solar Flare Analysis (RODEM)". We also thank Igor Krawczuk for interesting discussions and suggesting using Triton.
References
Andoni, A., Indyk. P. Laarhoven, T., Razenshteyn, I., and Schmidt, L. Practical and optimal Ish for angular distance, 2015.
Behnke, M. and Heafield, K. Losing heads in the lottery: Pruning transformer attention in neural machine translation. In EMNLP (1), pp. 2664 2674. Association for Computational Linguisties, 2020.
Beltagy, I., Peters, M. E., and Cohan, A. Longformer: The long-document transformer. CoRR, abs/2004.05150, 2020. URL httts:://artiii org/ abs/////....////////////// 05150.
Bender, E. M., Gebru, T., McMillan-Major, A., and Shmitchell, S.
On the dangers of stochastic parrots: Can language models be too big? In Proceedings of the 2021 ACM Conference on Fairness, Accountability, and Transparency, FAccT '21, pp. 610-623, New York, NY, USA, 2021. Association for Computing Machinery, ISBN 978 1450383097. doi: 10. (11145/344218884445522 10. 00..1114/3442180 3444144..8488992..
URL https://dot org/
Borgeaud, S., Mensch, A., Hoffmaan,, J., Cai, T., Rutherford, E., Millican, K., van den Driessche, G., Lespiau, J., Damoc, B., Clark, A., de Las Casas, D., Guy, A., Menick, J., Ring, R., Hennigan, T., Huang, S., Maggiore, L., Jones, C., Cassirer, A., Brock, A., Paganini, M., Irving, G., Vinyals, O., Osindero, S., Simonyan, K., Rae, J. W., Elsen, E., and Sifre, L. Improving language models by retrieving from trillions of tokens. CoRR, abs/21 12.0442,, 2021. URL pestts://arree org/ abs/2112 2121111...44....... 04426.
Child, R., Gray, S., Radford, A., and Sutskever, I. Generating long sequences with sparse transformers. CoRR, 0tts//994..100 10500 2019. URL httt:://arriiii wwtp:/iwwwrrr///////// 1904 10509.
Choromanski, K., Likhosherstov, V., Dohan, D., Song, X., Gane, A., Sarlós, T., Hawkins, P., Davis, J.. Mohiuddin, A., Kaiser, I.., Belanger, D., Colwell, L. J., and Weller, A. Rethinking attention with performers. CORR, abs/2009. 14794, 2020. URL. htttp:://axii..
wwweeiivvrr...eeeeeeeeeeee abs//009. 14794.
Dao, T., Fu, D. Y., Ermon, S., Rudra, A., and Ré, C. Flashattention: Fast and memory- exact attention with io-awareness. In NeurIPS, 2022.
Gao, L., Biderman, S., Black, S., Golding, L., Hoppe, T., Foster, C., Phang, J., He, H., Thite, A., Nabeshima, N., Presser, S., and Leahy, C. OpenWetText dataset, as part of 'the Pile: An 800gb dataset of diverse text for language modeling'. arXiv preprini arXiv:2101.00027, 2020.
Goyal, S., Choudhury, A. R., Raje, S., Chakaravarthy, V. T., Sabharwal, Y., and Verma, A. Power-bert: Accelerating BERT inference via progressive word-vector elimination. In ICML, volume 119 of Proceedings of Machine Learning Research, pp. 3690-3699, PMLR, 2020.
Hutter, M. The human knowledge compression contest. URL http://prize. hutterl, net, 6, 2012.
10

---

Katharopoulos, A., Vyas, A., Pappas, N., and Fleuret, F. Transformers are RNNs: Fast autore- gressive transformers with linear attention. In Proceedings of the International Conference on Machine Learning (ICML), pp. 5294-5303, 2020. URI. https: ttttps://flee //fleuret.eeeeeee httpss///leuree..oo////aaaeeee wwtpttperr.ppp.a.eee. kkthhroppoloss-ete wwweeettallimmmme al-icm12020. pdf.
Kim, S., Shen, S., Thorsley, D., Gholami, A., Kwon, W., Hassoun, J., and Keutzer, K. Learned token pruning for transformers. In KDD, pp. 784-794. ACM, 2022.
Kitaev, N., Kaiser, L., and Levskaya, A. Reformer: The efficient transformer. In ICLR. OpenRe- view.net, 2020.
LeCun, Y., Bottou, L., Bengio, Y., and Haffner, P. Gradient-based learning applied to document recognition. Proc. IEEE, (6(11))2278-23) 1998.
Li, Cotterell, R., and Sachan, M. Differentiable subset pruning of transformer heads. Trans. Assoc. Comput. Linguistics, 9:1442-1453 2021.
Loshchilov, I. and Hutter, F. Decoupled weight decay regularization. In ICLR (Poster). OpenRe- view.net, 2019.
Michel, P., Levy, O., and Neubig, G. Are sixteen heads really better than one? In NeurlPS, pp. 14014-14024, 2019.
Peng, H., Schwartz, R., Li, D., and Smith, N. A. A mixture of h heads is better than h heads. In ACL, pp. 6566-6577. Association for Computational Linguistics, 2020.
Peng, H., Pappas, N., Yogatama, D., Schwartz, R., Smith, N. A., and Kong, L. Random feature attention. In ICLR. OpenReview.net, 2021.
Qin, Z., Sun, W., Deng, H., Li, D., Wei, Y., Lv, B., Yan, J., Kong, L., and Zhong, cosformer: Rethinking softmax in attention. In ICLR. OpenReview.net, 2022.
Radford, A., Wu, J., Child, R., Luan, D., Amodei, D., and Sutskever, Language models are unsupervised multitask learners, 2019.
Raganato, A., Scherrer, Y., and Tiedemann, J. Fixed encoder self-attention patterns in transformer- based machine translation. In EMNLP (Findings), volume EMNLP 2020 of Findings of ACL, pp. 556-568. Association for Computational Linguistics, 2020.
Tay, Y., Dehghani, M., Bahri, D., and Metzler, D. Efficient transformers: A survey. CoRR, 00s22009.06732, 2020. URL https: pestts://arrae //arxivv..i//////a org/ absa 06732..
Tay, Y., Bahri, D.. Metzler, D., Juan, D., Zhao, Z., and Zheng, Synthesizer: Rethinking self- attention for transformer models. In ICML, volume 139 of Proceedings of Machine Learning Research, pp. 10183-10192. PMLR, 2021a.
Tay, Y., Dehghani, M., Abnar, S., Shen, Y., Bahri, D., Pham, P., Rao, J., Yang, L., Ruder, S., and Metzler, D. Long range arena A benchmark for efficient transformers. ICLR. OpenReview.net 2021b.
Tillet, P., Kung, H. T., and Cox, D. Triton: An intermediate language and compiler for tiled neural network computations. In Proceedings of the 3rd ACM SIGPLAN International Workshop on Machine Learning and Programming Languages, MAPL 2019, pp. 10-19, New York, NY, USA, 2019. Association for Computing Machinery. ISBN 978 978145036718 196. doi: 10.1145/3315508. 3329973. URL pattts::/dol: http:///oooooo/////// orw.oogg./.oo 1145/3315500 3329973.
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L., and Polosukhin, I. Attention is all you need. CoRR, abs/1706.03762,, 2017. URL http://arxiv. org/abs/ /10/a//////////////////////eeee/ 03762.
1706.
Voita, E., Talbot, D., Moiseev, F., Sennrich, R., and Titov, I. Analyzing multi-head self-attention: Specialized heads do the heavy lifting, the rest can be pruned. In ACL (1), pp. 5797-5808. Association for Computational Linguistics, 2019.

---

Wang, H., Zhang, Z., and Han. S. Spatten: Efficient sparse attention architecture with cascade token and head pruning. In HPCA, pp. 97-110. IEEE, 2021.
Wang, S., Li, B. Z., Khabsa, M., Fang. H., and Ma, H. Linformer: Self-attention with linear complexity. CoRR, Abs/2006.04768, 2020,
Weidinger, L., Mellor, J., Rauh, M., Griffin, C., Uesato, J., Huang, P., Cheng, M., Glaese, M., Balle, B., Kasirzadeh, A., Kenton, Z., Brown, S., Hawkins, W. Stepleton, T., Biles, C., Birhane, A., Haas, J., Rimell, L., Hendricks, L. A., Isaac, W., Legassick, S., Irving, G., and Gabriel, I. Ethical and social risks of harm from language models. CoRR, abs/2112.04333 2021.
Wu, Y., Rabe, M. N., Hutchins, D., and Szegedy, C. Memorizing transformers, 2022.
Xia, M.. Zhong, Z.. and Chen, D. Structured pruning learns compact and accurate models. In ACL (1), pp. 1513-1528. Association for Computational Linguistics, 2022.
Zaheer, M., Guruganesh, G., Dubey, K. A., Ainslie, J., Alberti, C., Ontañón, S., Pham, P., Ravula, A., Wang, Q. Yang, L., and Ahmed, A. Big bird: Transformers for longer sequences. In Larochelle, H., Ranzato, M., Hadsell, R., Balcan, M., and Lin, H. (eds.). Advances in Neural Information Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual, 2020. URL https: //pooe/dgg...rreiinggg.. neurigss..e.ii.... cc/paper/2020 hhahh/aaacccee c85 165112d 22222222222224731335555r222 absstaaa.a.a......11aaaaaaaaaat html.
Processing Systems 33: Anmual Conference
Zheng, L.., Wang, C., and Kong, L. Linear complexity randomized self-attention mechanism. In ICMI, volume 162 of Proceedings of Machine Learning Research, pp. 27011-27041. PMLR, 2022.
Zhou, W., Ge, T., Wei. Zhou, M., and Xu, Scheduled drophead: A regularization method for transformer models. In EMNLP (Findings), volume EMNLP 2020 of Findings of ACL, pp. 1971-1980. Association for Computational Linguistics, 2020.
12

---

A Additional Background
A.1 Structured attention
With a focus on performances on downstream applications, Raganato et al. (2020) show that it is possible to replace all but one attention head with static attentive patterns only relying on the position, e.g. attending to the previous token. Similarly, Tay et al. (2021a) investigate removing entirely the dot product attention and instead let the input token alone predict its attention pattern over the entire sequence, or use a random attention pattern. While they show their approach can be competitive on certain downstream task, results on language modeling seems to indicate those attention patterns are not expressive enough to model natural language. Peng et al. (2021) propose Random Feature Attention (RFA) which relies on kernel methods to approximate the softmax atttention and achieve linear computational complexity w.r.t. the sequence length. A follow up work by Zheng et al. (2022) augments RFA with randomized attention build an unbiased estimator of the softmax attention. Instead of looking for linear approximation of the softmax, Qin et al. (2022) propose to replace the softmax by a linear projection kernel and cos-based weightiin mechanism which scales linearly with the sequence length. While kernel based method seem like good compromise in terms of speed vs. performance, they have been shown underperform on certain downstream tasks Tay et al. (2021b). Compared to those methods, we are not trying to replace or find an estimator of the softmax attention, we instead provide an efficient way to leverage different forms of sparsity of the softmax attention matrix. Our speed gains do not come from a conceptually different way to compute the attention but simply deciding not to compute certain regions of the softmax attention matrix.
A.2 Pruning
The large body of works on pruning mostly focus on faster inference on downstream tasks. As a result, many methods are modifying the training process to facilitate pruning, often resulting in slower training.
Pruning heads in transformers. Many works have investigated dropping entire attention heads in transformer architectures. Let the entire attention matrix A be of shape B H x T T with B the batch size, H the number of heads, and T the sequence length. Dropping entire heads imposes an implicit structure over A which is now of shape B x H' T x T, with H' the number of non-dropped heads. Michel et al. (2019) and Voita et al. (2019) both observe-in in the context of Neural Machine Translation (NMT)-that a large fraction of heads can be dropped without significantly impacting performance. While their goal is not primarily to speed up training, many of methods following on this insight try to incorporate sparsifying mechanism during training which facilitate dropping heads at test time (Behnke & Heafield, 2020; Peng et al., 2020; Xia et al., 2022; Li et al., 2021). Still limited to NMT and downstream classification tasks, Zhou et al. (2020) found a regularizing effect of dropping heads during training. In comparison, less works have investigated dropping head in the context of language modeling, with text generations applications mind. comparison to those works, we propose a way to take advantage of dynamic sparsity structures, which is much more general but still includes head dropping. We show the potential of our work to enable speedup during training. Our work is orthogonal to head-dropping mechanisms and could be used in addition to those.
Pruning tokens in transformers. While dropping heads can be seen as dropping the model parameters generating keys, queries and values for those heads, another more recent line of work looked into dropping tokens. Dropping entire tokens imposes an implicit structure over the attention matrix A which is now of shape B x И x Tq x TkV. with TQ and TKV the number remaining queries and keys. Goyal et al. (2020) obtain faster inference on downstream tasks by using an attention based scoring mechanism to eliminate redundant input vectors. Wang et al. (2021) develop a joint algorithm-architecuu architecture framework which speeds up inference for downstream classifications tasks and language generation tasks. Their method includes head and token pruning along with specific hardware optimizations. Our work can be used to implement those approaches and does not require custom hardware accelerators. We moreover allow dropping individual heads instead of entire tokens.
13

---

B Details On The Implementation
We here give implementation details for experiments of $ 4. First, in App. B.I, we describe in more details the implementations of our custom Triton kernels introduced in 3. Secondly in App. B.2 we provide the python code used in our runtime benchmarks of $ 4, including pre and post processing steps required to reshape and re-order tensors. Lastly in $ B.3 we give more details on the hyperparam- eters used in our sequence modeling experiments of $ 4. The code for all our experiments can be found via the following link: https: /ttp///gttttb/tt cww/ccomm//ccee /epfm/lm///// /dt///dyaammsaaaaaeee setarse-flade ffashhahhh aash-attention
B.1 Triton Kernels
QK-Sparse Triton Kernel. In Alg. we detail the core of the QK-sparse algorithm from $ 3.1. This algorithm is computing the softmax attention result only for one block of queries, corresponding to one head. In practice this algorithm would be run in parallel for all blocks of queries for all heads. We find the index of the last relevant tile to compute by itcrating over block of key indices and comparing them with the largest query index for the current block of queries. This works quda and kide have a monotonic structure thanks to the stable sort used when reshaping the tensors (see the mre-processing code in App. B.2 for more details). We apply causal masking locally by looking qida and kidz. query can only attend to past keys: mask q_idx[:, None] >= k_ idx [None, :] The backward pass relies exactly the same trick, first iterate over query indices to find the starting block of queries.
Algorithm 1 Forward Pass for the QK-sparsc kernel
Require: Matrix Q E Ng matrices K, V E RNNxv xd, index matrices Qidr E R R.Ngx4 Xd and Kids E RNNx ,softmax scaling constant T E IR, softmax statistics vectors M RR and L E IRNN No, output tensor O E IR Noxx, xd query block size Bm, key block size Bn, starting query index startm-
1: Initialize o + Bmxd E RB xd let (0) E B...., m t (-o0) E ERA Bm.
2: Load current block of queries q t Q[start start, + :]
3: Load current block of query indices qidz t- Qidı (start,, start, Bm)
4: end <-
5: for start,, in range(O, NKV. Bm) do
6: Load block of key indices: kidz +- -Strrsttrr,, start, + Bn] then
7: if miinkkan <
8: end +- end +
9: end if
10: end for
11: for in range(0, end) do
12: Start of current key block start, <- n * Bn
13: Load current block of values
V(start, start, + B,
14: Load current block of keys k < K [[start, : start, + Bn
15:
Load current block of key indices kido Kuto (start, [itn[start, start, + B,]
16: Compute inner product qk t-qq.k"
17:
Apply causal mask qk +- MASK (qk, qide: kide)
18:
Update softmax statistics L, m, t UPDATE STATS((, o, v)
19:
end for
20:
Write O Ojstartm start,, Bms t-
21:
Write L[start,, startm + Bm
22: Write M [start,, start, Bml f- m
Hash-Sparse Triton Kernel. In Alg. 2 we detail the core of the Hash-sparse algorithm from 3.2. This algorithm is computing the softmax attention result only for one block of queries, corresponding to one head. In practice this algorithm would be run in parallel for all blocks of queries for all heads. We find the index of the first and last relevant tiles to compute by iterating over block of key hashes and comparing them with the largest and smallest query hashes for the current block of queries. This works as we reshaped our Q, K, V tensors by sorting them by their hash values (see the procpsicg code in App. B.2 for more details). In addition to causal masking, also
14

---

enforce attention to happen within the same bucket: mask = {q_idx[: None] >= k_id [None, :]) & (q_hash[;, None] k k_hash [None, :)). In our experiments we often replace >= by > to prevent a query to attend to itself as in the Reformer. As side note, for most application it would also be fine to only enforce causal masking and allow attention across buckets within tile. While this could add some serendipity in the attention computation, some applications might require masking based on hash. The backward pass relies exactly on the same trick, we first iterate over query indices to find the starting and end blocks of queries.
Algorithm 2 Forward Pass for the Hash-sparse kernel
Require: Matrix Q E xd matrices K, V E &NNxvx xd index matrices Qidi E RRN xd and K E R.v xd matrices containing the hash values Qash E No and Khash R RNvxx ssoftmax scaling constant T E R. softmax statistics vectors M E IR and L E RNa, output tensor O IR RgNqxd, query block size Bm, key block size Bn, starting query index startm. , f- (0) E BB " mm 4- (-20) B.. E RBBm 2: Load current block of queries q +- [sttarm start + Bl
1: Initialize o f (0) B.., xd E IR Bm xd
3: Load current block of query indices lide &- Qide [start,, start, + Bm]
4: Load current block of query hashes qhash + hash [startm startm + Bm]]
5: start <- 0
6: endhash f-
7: for start, in range(O, NKV. Bm) do
8: Load block of key hashes: Khash <- K hash (start, : start, + Br)
9: if minnkaash) I max (hash) max(qhash) then
10:
endhash + endhash +
11: end if
12: if max (KnaSh) < minn(hash) then
13: start start
14: end if
15: end for
16: end + endbash
17: for in range(start, endhash) do
18: Load block of key indices: kidz KiduljBn (j + 1)
19: if "in(kkdd) < max then
20: end < j +1
21: end if
22: end for
for in range(start, end) do
24: Start of current key block start, + n Bn
Load current block of values v V /start, start, + BR
26: Load current block of keys & < K [start, start, + B,,il
27:
Load current block of key indices kids & Kide (start, start, + B,]
28:
Load current block of key hashes khash +- Khash (start, start, + B,
29:
Compute inner product qk +- &Tq.k Tq.k
Apply causal and bucket mask qk
MASK qk,, qide, kide, qhash knash) UPDATE (EATSS,,STT(,,, m, qk, o, v)
31:
Update softmax statistics L, m. o -
end for
33:
Write Olstartm start,, + Bm: 4-
34: Write L[start, start, + B, t
35: Write (start, start,, + Bm) +- m
Accumulating softmax statisties while avoiding NaNs. The following is brief summary of how the FlashAttention algorithm (Dao et al., 2022) proposes to accumulate softmax statistics when iterating over blocks of keys. Given an input vector x, our goal is to compute softmax(x) &
ce max(#:) E, ea -max(a)) Let f (ax m) A at-m, and E( m) 2 f(e, m). Hence:
softmax(x) f(x max(
E(x, max (x)
15

---

Given vector a - [21 22], let my = max(x) (global max), and mi max(x1), we notice:
f mg)
softmax(x)
E(x,, mg)
f - em- m 2 ( m1) + ((22, mg)
f(z22 mg)
+ em ( mi) ((22(112)))
Therefore, if we have my, ((11, m1), and r = ?(a)m)))
we can update the softmax statistics for
new block of entries 22 by following the following steps:
1. Compute new global max: my max(mi, max (22))
2. Compute f(22, m,y) e0a-ma
3. Compute f(#2, mg) = Ei (222, mg)
4. Compute new m,): E(1,, m2) = em ( ((222,
{(.)m)))
5. Correct running softmax result: r {(
6. Add contribution from 22 r: r f(æggmm)
f(m.
In case all the keys are masked for given query, we would have max(x²) = -0, given that the first my initialized to -00 as well (see Alg. and Alg. 2) the fourth and second steps above would be undefined and result in NaN values. We solve the problem by replacing -co values in my by Os when doing those two steps. Another potential issue in step three: when query has no matching key in 22 then {x2,, my) is now O, which generate +oo in step five. This is an issue as we process keys by blocks, and if there are no keys for a query in the current block, we might find matching keys in following blocks. Adding co values
would prevent us to accumulate statistics later on. To get our desired behaviour and have when queries have no matching key we replace oo values by 1s in ((x, mg) during steps five and six. We summarize those steps in Alg. 3.
Algorithm 3 UPDATE STATS method
Require: Vector & E R.,, vector m E RR-, matrix of masked inner products qk E N output buffer E Bm Bm block of values v E IR. Bm
1: Compute new global max (step 1) m.we <- max(rrwmaxx(n)) m))
2: Replace oo by 0: m &- WHERE( WHERE( (mnew == -0, 0, mnew)
3: Compute step 2: p f- e4k gak-mnnuul I. None
# row-wise subtraction
4: Compute step 3: &2 <- rowsum(p)
5: Compute step 4: knew <- m- 8
6: Compute
Enew
7: Replace Do by 1: WHERE(2 = 00, 1,
8: Update P p x None]
# row-wise multiplication
9: Correct running softmax output (step 5) o + o x (t,
# row-wise multiplication
Add contribution from current block (step 6) 0 p.v
11: return Enow, Mmew
Hyperparameters. We extend the implementation of Flash PlashAttention available in the Triton tutorial. In our benchmarks, we use batch size B 4, 48 heads of 64 dimensions each.
Effect of imbalanced dropping/hashing, Our QK-sparse pattern relies on building smaller "com- pact" tensors containing the keys and queries that have been dropped. In the worst case, one head can have many queries or keys dropped, while another head might not drop any. The batching would then force us to retain the entire sequence and the compact tensors would be of the same shape as the original tensors. The runtime would increase as we now have to iterate over all the blocks of keys for one of the heads. Due to the structure of Attentin which loops over blocks of keys, in parallel for each block of queries, the runtime is tied to the longest sequence of keys to process. A smart QK-dropping mechanism on top our QK-sparse pattern should be cautious of dropping imbalances across heads. Concerning the Hash-sparse pattern, having all the keys and queries falling in the same
16

---

bucket would equate to the normal FlashAttention. Therefore unbalanced buckets would increase the runtime. Interestingly, our experiments on language modeling, do not face imbalance problems and even notice a speedup early in the training, see Fig. 14.(a) in App. C.
B.2 Runtimes Vacuum
Baseline implementation. We use Pytorch's FlashAttention implementation provided by the torch. mn. wtw.nnn.uuuui... www.fffcionaa scaled_ dot_ cwtt_pprouut_aateeeeeeeteoeeeeee attention function. To ensure fairness, we as- sume that all benchmarked functions receive tensor Q of shape (BATCH, CTX_Q, and tensors K, of shapes (BATCH, CTX_KV, H, D D_HEAD), where BATCH is the batch size, CTX_Q is the number of gueries, CTX_KV is the number of keys and values, is the number of heads, and D_ HEAD number of dimensions per head. For this reason the only pre and post processing steps required are transposing the input and output tensors.
D_HEAD),
Listing 1: pytorch_fulle entire sequence (no sparsity).
function applying the FlashAttention algorithm the
def (tarch.ffll..aasattteeeeee
k, v):
BATCH, X_CTX, и, D_HEAD ahape
q transpoos (1, 2) (BATCH, N_CTX_Q, D_HEAD)
x k. k.transpose (1, 2) (BATCH, , H, X_CTX_KV, D_READ) (BATCH, X_CTX_XV, D_HEAD)
v.tranepoae
y torch. functional swwled_dol__aeeeeeeeeeeeeeeeeede ptr.dddddd...rrdcttttttttieeeee attn_mase Nitn_mask=Nome, cs_causal =True)
k, v. dropout.p-O.0
10
return y transpoee (1,2) contiguous ()
Our proposed interface. We propose the following interface orchestrate between the Hash- sparse and the QK-sparse implementations:
Listing dynamic, ler_s_ppparreeeeeeeee interface
def dyeamic__parre__ttten (q. k, v. q_idx, k_idx, en_scale-Kone
sparsity.node-/haaa
Xeyzord argunents:
query tensor of shape (BATCH, N_CTX_Q, 9, Н, D_HEAD)
k: key tensor of shape (BATCH, N_CTX_XV, H, D_HEAD)
value
of shape (BATCH, N_CTI_KV, B, D_HEAD)
q_idx: tensor of shape (BATCH, M_ H) for each sequence in the batch, for each head,
each query the sequence,
representa either the bucket index if sparsity.node=== hash! or the whether keep given The type should be torch. int32 1f Soaraitt_node==han modelttyoode=="gg
sparsity_mode---.
sorch.float float if and torch
10
x_idx: tensor of shape (BATCH, X_CTX_KV, H) for each each key in the sequence, for each head,
in the batch, for representa either the bucket index sparaity noda== hash to keep that given head if sparaity saaralty_node--.. mode=='gk type should poarsity_node=-gg
whather en_scale: normalieation constant, 1/agrt (D HEAD) unless specified
int32
spareity_node- 'hash'
13
14
sparoity_node: 'hash" select the bash aparee implenentation and 'gk/ for the gk-sparse inplenentation
.*. 15
16
17
sn_scale s0,acale 1.0 nath eqrt (q.zeeq
18
19
eparsity_node 'bash": return aaah__spar..
21
k, h_hashsq_idx, k_hashsk_idx,
an an.scale-am.scale
22
elif sparsity_mode =- 'qk':
return @t_spareo_atteationle k, ioep k_keep=k_idx, sn_scale=ea_ccaee)
24 clse:
KayError (f" Unknovn sparsity.node: (sparaityy.noee), should be P'hash', *q)
17

---

Pre & post processing steps for Hash-sparse. In addition to having to transpose the Q. Q. K. V tensors. The preprocessing steps consist re-ordering the Q. K and V tensors based on bucket indices in h_s and k hash. We keep track of the original position of the queries, keys and values by storing the indices given by the sorting operations. Importantly, we use stable sorts ensure the queries, keys and values are sorted within each bucket. The following code is showing how we implemented all those steps using Pytorch:
Listing 3: sparse algorithm.
function showing pre and post processing steps for the Hash-
v, q_hash, k_hash, sn_scale):
assert q_hash. dtype == torch. int32 and k_hash. dtype == torch. int32
BATCH, X_CTX_Q, D_HEAD g. shape
q g.tranepose (1,
(BATCH, , N_CTX_Q, D_HEAD)
& - k. k.transpore (1, 2) # (BATCH, H, X_CTX_XV, D_HEAD)
v.trranspoee 2) (BATCH, X_CTX_KV, D_BEAD) t..h.==.taaaaaaaaeeaa.aam thaash.trasspsscc (1, contiiguuua k_hash k_hash. traneooss 2). contiguous (BATCH, H, X_CTX_XV)
q_hash =
(BATCH, B, x_CTX_q)
10
12
* Re-order the , valuea according q_baab and k_hash
q_haeh sort
stable-Irue) g.aaee.sbape (BATCH, H, X_CTX_Q), stable sort to keep tine ordering within bucket
k_hash k.baah.sorttdiir.
stable-True) shape (BATCH, H, X_CTX_XV)
15
16
q_idx q_haeh, indices x_idx k_hash indices
18
19
q_haeh q.hash. values k_hach .chash.valuee
20
22
q. extended q_ unsqueeue (-1) (xpeeexaddd...dd....... expand_aaaa_eeceeeeeeeeee._.._..
23
k_ idiideextedde k_idx. unequeeze (-1) expanddaaadsss
25
torch, gather din--2, index-q.idk. idx_ exteededdd) contiguoue ()
torch. gather
torch.gather (k, dim--2, index=k_idx_ extendeed) contiguous
27
torch, gather dim--2,
contiguous
28
29
perriionneeeelleee (q, k, q_idx, k_idx, q_bash, k_bash, sm_scale) y zeros ((BATCH, H, X_CTX_0, D_HEAD), dtype..
30
device-q. device) scatter (din-22, index-q exteeddddd, src-y).trance transpose (1,2) contiguous
31
return y
Pre & post processing steps for QK-sparse. In addition to having to transpose the Q. K, V tensors. The preprocessing steps consist removing dropped keys, values and queries from K, V and Q. The sorting operations need to be stable to keep the original time ordering within the remaining keys and queries. Moreover, the index tensor has to padded so our kernel can rely on those indices to compute which tile it should and shouldn's 't compute.
Listing 4: tesssprss_aatentiin function showing pre and post processing steps for the QK-sparse algorithm.
def compact (keep_tensor, index=Mone): Build conрact representation of
Keyvord argunents:
head x: input tensor to conpact, x. shape (BATCH, ALCTX, H, D_HEAD) shape (BATCH, M_CTX, containing
keep_tensor:
is zept, else *na
BATCH, T, H, D_HEAD = x.shape
if index is None: with norch.oo_grad() 10
indices_per_head keep_tensor (dim=
buffer.eize indices
int firet computes
non con-killed elen per head, we take to max of that corting:
very important use causal masking 12
sorting stable, cannot
13
sorted keep_ tensor (dit-- descending-True. atable-Trus) sorted, indices shape (BATCH index sorted. indices
D) zorted sequence T indices sorted.indices bbffbbufforrzz .:] (BATCH buffer x B) expand dinensione
14
18

---

else:
16
indices_per _ead None
compact_x gathee (dim=-3, index=index unsqueeze (-1) (-1,, -1, (BATCH buffer_cize x D_HEAD) / expand indices to cover all the dimensions
for
18
return coupact_x, index, indices_per_beer
19
20
21 starch.no.gran go_
pad_index (index, indices. periccssper_.eee per_heed, idd_idx=-1);
23
Pad the index tensor to comply with the kernel, returns
24 Keyvord argunenta:
25
indez: originel index tenser
given by 'conpact', index shape (BATCH, buffer_sizo H. For each batch and tinestep, reprsents the head idx it's originating fron.
26
indices_per_hes
have head: of shape (BATCH, H), for each head, containa how many indices dropped.
27
28
BATCH, buffer_aize,
index.ahape
29
index_copy torch. clone (index). typ. (torch int32) corce. nask trane (buffer_ devicemindex dew.xx.ddiice....e...... viex()) (1,1 -1,,1).xppaaddd (BATCH, buffer_aize,
30
indices_pex pereeadppr_edd...... view 1,,
copy [nask] pad_idx
32
return index_copy
33
34
35 def e.t.pare__ateationn (g, k, v, g_keep, k_keop, sm_scale): q_keep. dtype torch. float k_keep, torch. float
37
BATCH, X_CTX_R, D_HEAD g.ehape
29
40
# Building conpact representations
41
s.cee (BATCH, conpact CT__T_____ k_c, k_idx, iph_k = compact(k_keee k) shaee (BATCH, compact_N_CTX_KE v.c, -1 conpact(x_.eep. v, iddex-k.idx) shape (BATCH, Conpact___cTx___ H)
g-c, g_idx, iph_q compact(q_
42
q_idx_padded iad_index(q_idk. iph_g, pad_idx--1) (B, compact N_CTX_Q, Q, k_idx_padded 45
(n ipb_k, pad_idx=le9) # (B, Sappatt___TTx_K0
46
# need transpose everything
q_s = q.c.transpose contiiigu () # (BATCH, H, compact_N_cTX, N_CTX_Q, HEAD) 48
k.c.tranapose contiiiguuuss)
compact.N. M.RRCtt_,_CTT_EEE N_CTX_VV D_HEAD) (BATCH,
v_c - tra.spaaspos 2) contigguuu # (BATCH, H, compact_N_ ______X__VVVV___________________ D_HEAD) 50
q_idx_padded q_idx_ppdded tranepose (1, contiiguuu
(BATCH, H,
compact_N_CTX,
52
k_idx_padded
trappp (1, 2) contiguous (BATCH, E. compact..N.Tx_.W
53
y_c sesparre__tteetton..keeeee .asscaalee.aaspssseeee..... (1,2)
k_c, v_c, q_idx_padded. k_ i_idk_psdãed.
55
y corch.zeren zeros_like(q) scatter (dim=1,
index=q_idx long (). viev (BATCH, -1,8,1) expand (BATCH, -1, H, D_HEAD), arcsy_c)
56
return y
B.3 Sequence Modeling Experiments
Language modeling OpenWebTert Our implementation is based NanoGPT We use the Adam optimizer (Loshchilov & Hutter, 2019). We
(github.com tatttbb.comkkappatyyaannPPPP used bfloat16 and NVIDIA A100-40GB GPUs for all our experiments. Here is list of hyperpa- rameters shared by all our language models (F-LM, H-LM, and D-LM):
Weight-decay: 0.1
Depth (number of transformer blocks): 12
Number of heads: 12
Dropout: 0.0
Learning rate: 0.001
Percentage of iterations for warmup: 2%. We use a cosine learning rate scheduler.
Adam betal: 0.9
Adam beta2: 0.95
19

---

Tokenizer: We use the GPT2
tokenizer provided
by the tiktoken library
Hidden dimensions: 768
Dimensions per head: 64
Sequential MNIST and enwik8. For the comparisons with Reformer, we use a standard GPT2 (Radford et al., 2019) implementation. For the language modeling on enwik8, the Transformer has 12 blocks with 768 hidden dimensions, 8 attention heads, and 64 dimensions per head. Dropout is set to 0.1 and the batch size is with gradient accumulation steps. The sequence length is 4096. For autoregressive image generation on MNIST, we use smaller model with 8 transformer blocks and a hidden dimension of 256. Dropout is 0.1 and the batch size is 10. We train for 25 epochs.
C Additional Details and Analysis
Quadratic computational cost of attention transformers. In Fig. we show the runtime (forward backward) of transformer language model as function of the sequence length. We separate the time taken by the attention operations from the time taken by the rest the model. We how the attention computation ends up dominating the runtime as the sequence length increases.
600
Attention
Feed-forward
500
400
300
200
5000
10000
15000
Sequencee length
Figure 9: Quadratic computational cost of self attention dominates for larger sequences. We measure the forward and backward runtimes for a 12 layers transformer (see App. B.3 for implementation details). Except for the attention operations, the rest of the transformer runtime grows linearly with the sequence length.
Additional runtimes performances in a vacuum. Fig. and Fig. 11 we show the runtime details for the forward and backward methods separately for respectively the Hash-sparse and QK- sparse methods. We also measure runtimes of the forward and backward passes when we assume pre and post-procsssing steps aree free, see Fig. 12.
Linear QK-dropping scheduler. In the main paper we show results dropping keys and queries at random with fixed pre-defined probability. In an additional experiment we start by dropping 80% of keys and queries at random and linearly decay this probability 20%. Our intuition is earlier iterations just aim to learn contextual cues which are very redundant (and therefore quite immune random dropping) before requiring more fine- mme-graieed representations. Fig. 13 we show and analyse the results of that experiment.
H-LM models speeding up during training. The speed of Hash-sparse attention is conditioned on the distribution of bucket indices over keys and queries-e.g if all the keys and queries were to fall in the same bucket then there would be speedup real data such as Open WebText2, we observe our Hash-sparse based models are speeding up during training. In Fig. 14 we plot the number of iterations reached after hours of training (normalizing the time by the number of GPUs used for training). In Fig. 14.(a) we see a speedup for H-LM models early in the training.
FlashAttention. Interestingly, when training
Different bucket sizes nb for H-LM. For sequences of 8192 tokens, we show the influence of increasing the number of buckets nó in Fig. 15. As nb increases the runtime decreases.
20

---

300
350
10
300
6000
200
9009
xine
150
300
[ach (ours)
Flash ottension
0-
0-
Sequence Jiugth
Sequesse length
(a) Forward
(b) Backward
Figure 10: Forward and Backward runtimes in vacuum for our Hash-sparse method. Fig.(a) Forward runtimes for our Hash-sparse method. Fig. Backward runtimes for our Hash-sparse method.
110
300
300
600
30
200
2015
300
1000
3055
150
403
300
QK-gure (mund)
100
305
200
005
80%
0-
0-
5000
10000
15000
20000
20003
30000
35000
5000
10000
15000
20000
25000
30000
35000
Suquinor lingia
Sequusaoi leugth
(a) Forward
(b) Backward
Figure 11: Forward and Backward runtimes vacuum for our QK-sparse method. Fig.(a) Forward runtimes for QK-sparse method. (b). Backward runtimes for QK-sparse method.
Training H-LMs for more iterations. In Fig, 8 and Fig. 6 we show results of language models trained on Open WebText2 for 15k iterations. To verify whether our finding are consistent when you train for more iterations we also try training for 50k iterations. For sequences of 8192 tokens, using the same hyperparameters described in App. B.3, we show in Fig. 16 that our findings for Hash-sparse based models do hold when training for more iterations match the perplexity per iterations of the baseline model using FlashAttention over the entire sequence while being significantly faster.
Visualizing the low coverage of Reformer. In Fig. 17 show the attention matrices for two different heads and show how the coverage-thh percentage of key-query interactions actually computed vs. what should be computed according to the hash indices computed for keys and queries-of of the Reformer LSH algorithm can be low.
D Limitations and Societal Impacts
The aim of our work is to develop a method for efficiently computing attention with several sparsity structures. We don't focus on developing the best method for sparsification, although, for example, we improve the hashing- based mechanism. Moreover, as already explained, on very small sequences we incur in some constant overhead which limits our gains.
Limitations.
21

---

1300
5000
500
0000
1000
400
200
200
0-
5000
10000
15000
20000
25000
300000
30000
5000
10000
10000
20000
20000
30000
13030
length
leugth
(a) QK-
(b) Hash-sparse
Figure 12: Forward+Backun Backward runtimes in vacuum when we assume the preprocessing is free. (a). Forward+backward runtimes for our QK-sparse method when the pre and postproceessing steps free. While the runtimes for large sequences stays relatively the same compared to Fig. 7, the offset for short sequences is now much smaller. Fig.(b): Forward+backwart runtimes for our Hash-sparse method when the pre and postproceessing steps sequences stays relatively the same compared to Fig. 3, the offset for short sequences is now smaller. Intrestingly, unlike for the QK-sparse method, the offset for smaller sequences stays significant. We believe this might be due to the influence of the block size used (128): when the sequence length large enough in comparison the block size, the block structure of the hash-sparse attention matrix cannot be efficiently leveraged and the number of tiles processed by the Hash-sparse method larger than that of the FlashAttention method. In contrast, for the QK-sparse method the good tiles all start from 0 (we know the first tile(s) will be efficiently packed).
free. While the runtimes for large
Iver/tii on IPPU
115
Validistion Puplexity
60
Validation Porplexity
60
x2.6
10
50
x1.9
x
40
30
20
60%
5000
100000
15000
D-LM 30%
D-LM 70%
Iserations
Time (hours)
(a) Iter reach perplexity
(b) Time reach perplexity
(c) Speed gains
Figure 13: Training Language Models (LM) on OpenWeh WehnTeeText' Text2 using random Query/Key dropping based sparsity (D-LM) or FlashAttention over the entire sequence (F-LM). Results are the same as Fig except for the addition of the D-LM model with linear decay of from 80% to 20%. In (b) and (cg. we observe the additional D-LM model in yellow is slower than other D-LM models we experimented with, but, as seen (ag, reaches slightly better perplexity.
22

---

15000
15000
Iverations
10000
Iterations
10000
5000
5000
0
20
40
60
20
40
60
Time (hours GPU)
Time (hours GPU)
F-LM
D-LM 70%
H-LM H-LM
F-LM
D-LM 30%
D-LM 50%
80% 20%
(a) Time comparison H-LM vs. F-LM
(b) Time comparison D-LM vs. F-LM
Figure 14: Training speed function of time for H-LM, D-LM and F-LM. All models trained on sequences of 8192 tokens. Fig.(h) we observe D-LMM models having a fixed speed throughout the experiments except for the one corresponding the linear decaying of the sparsity ratio s which slows down during training. In Fig.(a) H-LM models are speeding up early during training and then seem to keep constant speed.
10
X2.1
Iter GPU
x1.8
x3.3 x3.6
T=8192
T 16384
F-LM
H--.M nÀ
H-LM mb 16
H-LM uu
Figure 15: Influence of the number of buckets nó on H-LM training speed for T' 8192 and = 16384. In accordance with Fig. 3 increasing speeds up the training but the speed gain from doubling nó decreases as vab is increasing.
10
F-LM
25
Ite Iter/min on GPU
x2.1
Valiltion plesite
24
H-LM abb
X1.9
23
22
21
20-
2-
19-
18
F-LM
H-LM vb
10000
15000
2000
25000
Teeratones:
30000
35000
40000
45000
50000
LM nd
16
(a) Iter to reach perplexity
(b) Speed gains
Figure 16: Comparison between H-LM and F-LM when training for 50k iterations on sequences of 8192 tokens. Those curves are averaged over two seeds. In Fig. we zoom in on the perplexity per iteration for the baseline F-LM and two of our H-LM models using our Hash-sparse attention with and 16 buckets. We observe that our methods are matching the perplexity of the F-LM model which uses FlashAttention over the entire sequence. Fbg. We match the perplexity but gain in speed, our methods are 1.9x and 2.1x faster than the baseline for respectively 8 and 16 buckets.
23

---

Figure 17: Attention pattern of Reformer. The two images show the attention pattern for the Reformer LSH attention after reordering queries and keys according the bucket. The red squares are the positions for which attention is computed in the Reformer with bucket size 32. The back squares are positions for which the queries and keys map to the same hash bucket. Queries are sampled from normal distribution with μ 3 and vo 5. Each image refers to different attention head. Black regions which are not covered by red tiles are key-query interactions which should be computed but are missed by the Reformer LSH attention.
Societal impacts. The attention mechanism central to Large Language Models (LLMs). More- over, efficient attention mechanisms can make these models more powerful, by making them faster to train and by increasing their context length. The social impact and risks associated with our work, therefore, are included in the risks associated with the deployment of such systems (Bender et al., 2021; Weidinger et al., 2021).
24


Reviews:
[{"review_id": "HxI5ikisAC", "reviewer": "Reviewer_RPsB", "summary": "The paper proposes a method for fast causal attention by combining dynamic sparse attention (query & key dropping or from hashing) with FlashAttention, called SCFA (Sparse Causal FlashAttention). The paper shows that with the right preprocessing, attention with query & key dropping (QK-sparse) can also be done efficient in the style of FlashAttention, by only loading blocks of keys and values necessary for the computation. Similarly, with hashing-based sparsity (where only attention entries of queries and keys in the same hash buckets are computed), by careful manipulation of the indices, one can also load only the blocks of keys and values necessary, avoiding unnecessary memory reads/writes. These ideas are validated empirically, showing significant speedup compared to naive sparse attention, as well as speedup against FlashAttention for long sequence lengths. Training models with SCFA leads to similar quality but faster wall-clock time.\n", "questions": "1. Why is the focus on causal attention? Do these ideas apply to non-causal attention (e.g., in BERT, ViT)?\n\n2. For hashing-based sparsity, if there are multiple hashing rounds (e.g. as in Reformer), the same key & query could end up in the same buckets in multiple rounds. How does one avoid over-counting in this case? Can that be done efficiently?\n", "limitations": "Not necessary", "rating": 7, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Impressive wall-clock speedup, which could have significant impact on how sparse attention can be used in practice. One of the main issues with approximate attention (e.g., sparse attention) is that they don't necessarily bring wallclock speedup. With models being trained at much longer sequence lengths (e.g. 32k-100k), approximate attention might be a necessity, and this paper takes a step towards making this easier for practitioners.\n\n2. The hashing-based sparsity proposed in the paper seems to work better than the hashing method from Reformer, with higher coverage. This leads to higher model quality, as validated by the experiments.\n", "weaknesses": "1. The motivation for QK-sparse attention could have been explained better. I'm personally not as familiar with this method. Why would one want to drop query & key?\n\n2. The technical challenges could have been highlighted better. E.g. explaining why it is hard to efficient implement attention when the sparsity pattern is dynamic, would help readers understand the paper better.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the valuable time spent reviewing our paper!\n\n> The motivation for QK-sparse attention could have been explained better. I'm personally not as familiar with this method. Why would one want to drop query \\& key?}\n\nThere is a large body of work [1,2,3,4,5] that suggests it is possible to drop heads or tokens without degrading the performances on downstream tasks. However, those works only focus on speeding up the inference and usually add costly mechanisms during training. Our contribution is to propose a sparse pattern that is more general than those works, yet allows speed gains during training. We will add a paragraph to better justify the QK-sparse pattern in the next revision. \n\n> The technical challenges could have been highlighted better. E.g. explaining why it is hard to efficient implement attention when the sparsity pattern is dynamic, would help readers understand the paper better.\n\nThanks for raising this point. We will better highlight the difficulties of efficiently implementing dynamic sparse patterns in the next revision of our paper.\n\n> Why is the focus on causal attention? Do these ideas apply to non-causal attention (e.g., in BERT, ViT)?\n\nYou are right that the proposed methods can be extended outside of autoregressive models simply by removing the causal masking. We preferred to focus on causal attention as it's where a lot of the heavy, large model pre-training is currently done, and where efficiency gains seem to be most beneficial and relevant, especially on long sequences.\n\n> For hashing-based sparsity, if there are multiple hashing rounds (e.g. as in Reformer), the same key \\& query could end up in the same buckets in multiple rounds. How does one avoid over-counting in this case? Can that be done efficiently?\n\nFor this work, we focused on a single round of hashing. We already planned to extend this to multi-round hashing, and we believe that we can avoid the double-counting issue efficiently in our kernel, with the same mechanism that is used in the Reformer implementation [2].\n\n\nReferences:\n\n[1] Michel, P., Levy, O., and Neubig, G.: Are sixteen heads really better than one?\n\n[2] Voita, E., Talbot, D., Moiseev, F., Sennrich, R., and Titov, I.: Analyzing multi-head self-attention: Specialized heads do the heavy lifting, the rest can be pruned.\n\n[3] Behnke, M. and Heafield, K.: Losing heads in the lottery: Pruning transformer attention in neural machine translation.\n\n[4] Goyal, S., Choudhury, A. R., Raje, S., Chakaravarthy, V. T., Sabharwal, Y., and Verma, A.: Power-bert: Accelerating BERT inference via progressive word-vector elimination.\n\n[5] Wang, H., Zhang, Z., and Han S.: Spatten: Efficient sparse attention architecture with cascade token and head pruning.\n\n[6] Kitaev, N., Kaiser, L., and Levskaya, A.: Reformer: The efficient transformer\n"}, {"review_id": "Yu5BLjy06u", "reviewer": "Reviewer_jEHP", "summary": "The work proposes a new method to speed up and improve the causal self-attention of transformer-based language models for long sequences. The method uses a kernel called SCFA that can handle any sparsity pattern and causal mask in the attention matrix. The method also introduces two dynamic schemes to sparsify the attention: one based on hashing and one based on key/query dropping. The method achieves significant runtime gains and better performance than existing methods.", "questions": "When you run the comparsion, proposed solution running on 2 or 3 A100, while Reformer alwasy run on single GPU. Is it correct? if yes, why?", "limitations": "N/A", "rating": 7, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1: The work strives to speed up and improve the causal self-attention of transformer-based language models for long sequences, which is a relevant topic. The work conducts a thorough review of related work. And based on related work, it proposes two solutions based on hashing and dropping.\n\n2: The work provides a clear explanation of the two proposed solutions. Besides explaining the solutions, it also discusses the overhead and edge cases. In the appendix, it shares the algorithm and implementation in detail.\n\n3: The work explains the experiment in great detail, and also shows a significant speed improvement by applying the proposed solution.", "weaknesses": "1: One more experiment with large model will be great. \n\n2: One more experiment on other task, like text-classification task will be great. \n\n3: Agree, as the paper stated, a better dropping scheme is beyond the scope of this work. But it might improve the work further, such as achieving higher dropping rate and same speed performance.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank you for your review and your appreciation of our work!\n\n> When you run the comparsion, proposed solution running on 2 or 3 A100, while Reformer alwasy run on single GPU. Is it correct? if yes, why?\n\nGiven our limited resources, we decided to use our multi-GPUs infrastructure to experiment with language modeling tasks which typically rely on training large models. "}, {"review_id": "oVtWPbPRWu", "reviewer": "Reviewer_89Ny", "summary": "This paper proposes an improved version of FlashAttention to support irregular block sparsity due to queries/keys-dropping or hashing. The proposed method modifies the mechanism used in FlashAttention to arbitrary indexing of queries/keys, which can be viewed as combining both FlashAttention with either QK-dropping or hash-sparse methods.", "questions": "- Can you provide more details in the SCFA kernel design? Also, can you comment of the reproducibility of the kernel design?", "limitations": "No discussion.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The proposed method, being simple and effective, improves over prior methods (FlashAttention, Reformer, and QK-sparse).", "weaknesses": "- The paper is a bit low in novelty. The SCFA kernel design is based on a combination of prior methods.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your feedback.\n\n> Can you provide more details in the SCFA kernel design?\n\nOur kernel design relies on the FlashAttention algorithm as a starting point. On top of this, we (i) reshape input tensors to bring an interesting structure of the attention matrix, and (ii) exploit this structure by computing, for each block of queries, which tiles of keys should be considered.\nWe develop separate kernels for each type of sparsity (QK-sparse and Hash-sparse). Additionally, we provide a common Pytorch interface that abstracts away a lot of the complexity. The pseudo-code for the kernel implementations is provided in the Appendix. \nWe will add more details on the kernel design in the next revision.\n\n> Also, can you comment of the reproducibility of the kernel design?\n\nIn the supplementary material (App.B) we provide a link to an anonymized repository containing code and instructions to reproduce all of our results. In order to reproduce our results, one simply needs a GPU supporting the bfloat16 data type. We used one A100 40GB for all of our runtime experiments. "}, {"review_id": "ikpjVGHyZ9", "reviewer": "Reviewer_Z3Z8", "summary": "This paper extends Triton implementation of FlashAttention to support two forms of sparse attention: key/query dropping and hashing-based attention. Source code for the kernels are made available.", "questions": "Please see the weakness section.", "limitations": "N/A", "rating": 3, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The paper gives clear descriptions of the released kernels and provides comprehensive validation experiments of the kernels.", "weaknesses": "The main concern is that both types of sparse attention are very restricted. Although this is a good effort that benefits the researh community, I am not sure how significant the kernels are.\n\nThe K/Q dropping sparse attention seems not useful in training large language models. It seems very rare that we would mask certain tokens completely and do not let them contribute to attention at all.\n\nThe hashing-based sparse attention seems restrictive. We essentially need to separate tokens into groups and only allow attention within each group. The only useful scenario that I can think of is when a training sample is a concatenation of multiple unrelated pieces of text.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for taking the time to read our work. We hope those clarifications are addressing your concerns and may justify raising your score. \n\n> The K/Q dropping sparse attention seems not useful in training large language models. It seems very rare that we would mask certain tokens completely and do not let them contribute to attention at all.\n\nWe believe there might be a misunderstanding. We do not mask entire tokens, but single heads. Fig.1 and Fig.2 illustrate the sparsity patterns for a  single head. This means that if each token has $12$ heads for keys and $12$ for queries, then our QK-sparse method allows to efficiently drop e.g. $5$ heads for the query and $8$ for the key. All the heads need not to be dropped. Moreover, each layer can drop heads independently from other layers. Thus, we are not completely masking out tokens from the attention. There is a strong body of work suggesting it is possible to drop entire heads [1,2,3] or tokens [4,5] without losing too much accuracy on downstream tasks. Very recently, [6] use an elaborate query/key-dropping mechanism which can reduce the inference time without sacrificing perplexity. Our contribution is to show how these types of sparsity can be efficiently implemented to speed up the training, which to the best of our knowledge has never been done before.\n \n> The hashing-based sparse attention seems restrictive. We essentially need to separate tokens into groups and only allow attention within each group. The only useful scenario that I can think of is when a training sample is a concatenation of multiple unrelated pieces of text.\n\n\nSimilar to the point made above. We hash single heads, and not entire tokens. Therefore heads $1$ and $5$ can end up in bucket $8$ while head $3$ might end up in bucket $2$. Moreover, we perform this hashing for each layer independently, and therefore each token is not attending within a single group and has many chances of fetching various bits of information across the entire sequence. This pattern is conceptually similar to the LSH-attention implemented in the Reformer [7] but solves some of its serious limitations while being significantly faster. \n\n\nReferences:\n\n[1] Michel, P., Levy, O., and Neubig, G.: Are sixteen heads really better than one?\n\n[2] Voita, E., Talbot, D., Moiseev, F., Sennrich, R., and Titov, I.: Analyzing multi-head self-attention: Specialized heads do the heavy lifting, the rest can be pruned.\n\n[3] Behnke, M. and Heafield, K.: Losing heads in the lottery: Pruning transformer attention in neural machine translation.\n\n[4] Goyal, S., Choudhury, A. R., Raje, S., Chakaravarthy, V. T., Sabharwal, Y., and Verma, A.: Power-bert: Accelerating BERT inference via progressive word-vector elimination.\n\n[5] Wang, H., Zhang, Z., and Han S.: Spatten: Efficient sparse attention architecture with cascade token and head pruning.\n\n[6] Anagnostidis S., Pavllo D., Biggio L., Noci L., Lucchi A., Hofmann T.: Dynamic Context Pruning for Efficient and Interpretable Autoregressive Transformers\n\n[7] Kitaev, N., Kaiser, L., and Levskaya, A.: Reformer: The efficient transformer"}, {"review_id": "lSHhWftV7x", "reviewer": "Reviewer_kvnu", "summary": "This paper extends the FlashAttention to support the structured sparse attention, enabling attention to leverage sparsity to achieve further acceleration on the basis of FlashAttention. In this way, the prior QK-sparse attention and Hash-sparse attention methods can be further accelerated during training and evaluation.", "questions": "- Are the block sizes in Algorithm 2 and 3 related to the sparse patterns or sparsities?\n- QK-sparse attention and Hash-sparse attention adopt per-head sparse index and per-head hash, respectively, which leads to load unbalance during calculation. How much does this affect the GPU occupancy and utilization?\n- Can you add the theoretical speedup of the two sparse attention methods? I just want to see how much potential acceleration there was.\n- Can you show the time ratio of sorting and re-ordering costs with different nb and different sequence lengths in Figure 3? Can sorting and re-ordering be fused with attention?\n- typo: line 145, $j \\in [0, j_{stop}]$", "limitations": "The author discussed the limitation of the method and the societal impacts. Other limitations of my concern can refer to the weaknesses.", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The proposed method further accelerates QK-sparse attention and Hash-sparse attention, which helps to expand the future applications of sparse attention methods.\n- The proposed method can accelerate both model training and evaluation. Since current LLM training is extremely time-consuming, this may help the efficient training of LLMs.", "weaknesses": "- Experiments mainly compare with the FlashAttention, which is a dense baseline. But the sparse methods compared in the experiments are very naive implementations without any performance optimization, making it difficult to show the benefits that flashAttention brings to sparse computation. It is better to compare with the methods using efficient sparse acceleration libraries.\n- This work is too engineering, which implements the existing methods under the FlashAttention framework.\n- The application of this method will be limited, only suitable for two calculation modes, QK-sparse attention and Hash-sparse attention.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for taking the time to review our paper and for your feedback! We hope the following answers your questions and may bring you to raise your score. \n\n> The application of this method will be limited, only suitable for two calculation modes, QK-sparse attention and Hash-sparse attention.\n\n\nWe would like to emphasize the generality and relevance of our patterns. The Hash-sparse pattern is conceptually similar to the LSH-attention implemented in the Reformer [1]. We improve upon the LSH-attention by providing an implementation that is (i) guaranteeing $100\\%$ coverage of keys and queries falling in the same bucket, and (ii) is significantly faster. Moreover, the relevance of our QK-sparse pattern is supported by several works showing how a large number of tokens or heads can be dropped in transformer architectures [2,3,4,5,6] without detrimental effects on downstream tasks. While all of those works focus on more efficient inference, we show how our more general QK-sparse pattern can be used to speed up the training of language models.\n\n> Are the block sizes in Algorithm 2 and 3 related to the sparse patterns or sparsities?\n\nThe block sizes are hyperparameters used to parallelize efficiently the computation across the GPU cores. Those parameters can have an impact on the runtime but are not tied to sparse patterns or sparsities. In all of our experiments, we set the block sizes to $128$. \n\n> QK-sparse attention and Hash-sparse attention adopt per-head sparse index and per-head hash, respectively, which leads to load unbalance during calculation. How much does this affect the GPU occupancy and utilization?\n\nOur QK-sparse pattern relies on building smaller \"compact\" tensors containing the keys and queries which have not been dropped. In the worst case, one head can have many queries or keys dropped, while another head might not drop any. The batching would then force us to retain the entire sequence and the compact tensors would be of the same shape as the original $\\mathbf{Q},\\mathbf{K},\\mathbf{V}$ tensors. The runtime would increase as we now have to iterate over all the blocks of keys for one of the heads. As a small experiment, we consider random dropping $50\\%$ of the queries and keys, for some set of hyperparameters we get a runtime of $14$ms. Now we run the same operation but with one single head not being dropped at all, the runtime increases to $48$ms, which is close to the runtime with nothing being dropped. Due to the structure of FlashAttention which loops over blocks of keys, in parallel for each block of queries, the runtime is tied to the longest sequence of keys to process. Anyone implementing a smarter QK-dropping mechanism on top of our QK-sparse pattern should ensure that queries and keys are dropped relatively uniformly across heads.\n\nConcerning the Hash-sparse pattern, having all the keys and queries falling in the same bucket would equate to the normal FlashAttention. Therefore unbalanced buckets would increase the runtime as for the QK-sparse pattern. Interestingly, in our experiments on language modeling, we do not face imbalance problems and even notice a speedup after around $1$k iterations, see Fig.11.a in App.C. \n\nWe thank you for your comment and will add a section on imbalance in the next revision. \n\n\n> Can you add the theoretical speedup of the two sparse attention methods? I just want to see how much potential acceleration there was.\n\n\nFor the Hash-sparse pattern, assuming perfectly balanced $nb$ buckets and a sequence length of $T$, then we expect $\\mathcal{O} (\\frac{T^2}{nb})$ complexity as opposed to $\\mathcal{O} (T^2)$. For the QK-sparse pattern, given a balanced sparsity of $s \\in [0,1]$ we expect $\\mathcal{O} (T^2(1-s)^2)$ complexity for the attention mechanism.  \n\n\n> Can you show the time ratio of sorting and re-ordering costs with different nb and different sequence lengths in Figure 3? Can sorting and re-ordering be fused with attention?\n\nIn Fig.2 of the figures provided with this rebuttal, we show runtimes for the two kernels when the pre and post-processing steps are assumed to be free. We observe that, especially for the QK-sparse pattern, fusing those steps with attention could bring a significant speedup, especially for smaller sequences. We believe those figures could be further improved by a better tuning of GPU-related hyperparameters such as the block sizes, number of warps and stages used.\n\n\nReferences:\n\n[1] Kitaev, N., Kaiser, L., and Levskaya, A.: Reformer: The efficient transformer\n\n[2] Michel, P., Levy, O., and Neubig, G.: Are sixteen heads really better than one?\n\n[3] Voita, E., Talbot, D., Moiseev, F., Sennrich, R., and Titov, I.: Analyzing multi-head self-attention: Specialized heads do the heavy lifting, the rest can be pruned.\n\n[4] Behnke, M. and Heafield, K.: Losing heads in the lottery: Pruning transformer attention in neural machine translation.\n\n[5] Goyal, S., Choudhury, A. R., Raje, S., Chakaravarthy, V. T., Sabharwal, Y., and Verma, A.: Power-bert: Accelerating BERT inference via progressive word-vector elimination.\n\n[6] Wang, H., Zhang, Z., and Han S.: Spatten: Efficient sparse attention architecture with cascade token and head pruning."}]

Author rebuttal:
We thank all reviewers for taking the time to review our work. We improved and optimized our Hash-sparse kernel and new results are shown in the figures provided with this rebuttal. Compared to our previous version, still without sacrificing perplexity, we further increase the training speed of a transformer language model from $1.8\times$ to $2.0\times$ and from $2.3\times$ to $3.3\times$ for sequences of respectively $8k$ and $16k$ tokens.

Meta-review: The paper extends FlashAttention to accommodate a large class of attention sparsity patterns that, in particular, encompass key/query dropping and hashing-based attention. Half of the reviewers agree that the submission is borderline or tend to reject the paper, while the other two reviewers would like to accept the submission. According to the discussion between the authors and the reviewers, the work has its own merits, especially since attention has been almost everywhere. The improvement would be potentially useful for other researchers in this field.
Decision: Accept (poster)