Title: Fast Attention Over Long Sequences With Dynamic Sparse Flash Attention
Venue: NeurIPS 2023
Accepted: yes
Keywords: self-attention

large language models

transformers
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

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 × and 3 . 3 × for sequences of respectively 8 k and 16 k tokens.

1 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 et al., 2020). Other methods rely on a predefined sparse masking of the attention matrix, e.g. to constrain the attention to a 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 impose static implicit or explicit constraints over the attention matrix.

Another promising line of work consists in computing a dynamic modulation of a sub-part of the attention matrix. They are based, for instance, on 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).

The promising theoretical computational complexity of these methods contrasts with the fact that today's most successfully deployed practical models instead rely on vanilla attention, in part thanks to the efficiency of FlashAttention (Dao et al., 2022). This implementation is mathematically identical to the vanilla attention proposed by Vaswani et al. (2017) in their seminal paper, but trades in additional compute for less memory I/O.While still avoiding a memory footprint quadratic with the sequence length, it delivers practical speedups of over 5 × compared to a naive implementation.

Using an attention layer in an autoregressive model-which has been key in the recent remarkable AI breakthroughs-requires to make it 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 or rely on geometric hashing, which results in irregular causal structures as illustrated in Fig. 1 and Fig. 2.

We propose an 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 (Kitaev 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.

We propose a new approach implemented with SCFA that reduces computation by dynamically selecting, for each head, keys and queries to be removed from the attention operation, superseding existing methods that limited pruning to entire heads or entire queries/keys, due to the lack of an efficient fine-grained kernel implementation. See § 3.1.

Experimental evaluations show that SCFA can efficiently be used for a 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 hashbased SCFA not only implements the exact computation, but also has a 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 query-key pairs dropped (see § 4.3).

2 Related work

State-of-the-art sequence models have very high computational requirements. As a consequence, a 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 a 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 locality-sensitive-hashing (LSH) to sparsify the attention 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.

K

Q

QK-sparse

Hash-sparse

FlashAttention introduced by Dao et al. (2022) has recently gained a lot of popularity as an efficient, IO-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 et al., 2020) are two more 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 a combination of local, global, and random attention.

Hash Attention. When computing the attention matrix for a T × D query tensor Q and a T × D key tensor K , we consider the matrix of dot-products QK ⊤ , which can become impractical to compute for very long sequences. However, we are only interested in the row-wise softmax( QK ⊤ ) , meaning that the contribution of the keys to every query is dominated by the ones with the highest similarity. Thus, restricting the attention computation to queries and keys with high similarity is a natural choice to reduce the computation.

Hash attention, introduced in the Reformer (Kitaev et al., 2020), allows to quickly select the closest key vectors for each query using locality-sensitive-hashing (LSH). In general, the LSH mechanism assigns a hash code to vectors with the requirement that vectors that are close in space are mapped to the same hash with high probability. For the hash attention, the Reformer assumes a 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 to each other and one chunk back. This results in 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-attention operation consists of multiplying a T × D query tensor Q by a T × D key tensor K , to obtain a matching score matrix, which is then rescaled and rownormalized with softmax, to get a T × T attention matrix A . This matrix is then multiplied by a T × D ′ value tensor V to obtain the final result. This is the core operation in a standard Multi-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.

Hash-sparse

QK-sparse

K

Q

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 is faster than retrieving it from memory, and avoids a 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 is 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 a 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

We develop an efficient CUDA kernel written in Triton (Tillet et al., 2019) that maintains the careful memory management of FlashAttention but can handle a causal structure defined through an arbitrary indexing of the keys and the queries. In the case where this indexing consists of a binary decision to drop or not the head of a query/key, this corresponds to our 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 B × H × T × 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 a single head and instead consider a query tensor Q of shape T Q × D , and a key K and value V tensors of shapes T KV × D . The algorithms described below will be run in parallel for all elements of the Cartesian product B × H . We split tensors into blocks: Q ≜ [ Q 0 , . . . , Q m ] , K ≜ [ K 0 , . . . , K n ] . We define a tile T i,j ≜ Q i K ⊤ j , which corresponds to the dot products of a 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 c , K c , and V c tensors. Through this reduction we are left with a smaller attention matrix A c 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 Q i , processing blocks of keys K j one after the other, moving along a row of tiles: T i, 0 , . . . , T i,n . Causality dictates that it is unnecessary to process a tile T i,j when i < j . We cannot follow this rule anymore when working with compact representations. To leverage the causal structure of A c , we build a new kernel which gets as additional input vectors q idx ∈ R T Q and k idx ∈ R T KV representing the indices of the queries and keys in the original uncompressed tensors. Those are similarly split into blocks: q idx ≜ [ q idx 0 , . . . , q idx m ] . The condition for a tile T i,j to be unnecessary to compute is now to have max( q idx i ) < min( k idx j ) . When processing a block of queries Q i , we iterate over the key indices k idx 0 , . . . , k idx n to find the index j stop of the first block satisfying that condition. We then know we need to process the tiles T i,j for j ∈ [0 , j stop [ . Within each tile T i,j , we in addition apply a local causal mask by comparing indices in q idx i and k idx j . By computing j stop 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.

Overhead. Computing Q c , K c , and V c 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.

Edge cases. Dropping keys and queries can 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 a similar hash index close to the diagonal. If the sorting is stable, i.e. 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 q hash and k hash to our kernel. Based on those hash indices, we now find not only the stopping index j stop but also a starting index j start . j start is the first index for which some of the indices in q hash i are present in k hash j , j stop is the first index for which all indices in k hash j are strictly larger than indices in q hash i . In a second step we refine j stop now based on the indices k idx and q idx , the updated ˆ j stop is the last index j ∈ [ j start , j stop [ for which max( q idx i ) ≥ min( k idx j ) . As shown in the last column of Fig. 2, we then only compute tiles T i,j for j ∈ [ j start , ˆ j stop ] . As for the QK-sparse method, we use q idx and k idx to apply a causal mask locally for each tile. In addition to the causal mask, we use q hash and k hash to mask interactions

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 , K 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.

ms

0

5000

10000

15000

20000

25000

30000

35000

Sequence length

0

200

400

600

800

1000

1200

nb

= 2

nb

= 4

nb

= 8

nb

= 16

nb

= 32

nb

= 64

Hash-sparse (our method)

Naive hash-sparse

Flash attention

0

2000

4000

0

10

20

30

4 Experiments & Results

In this section we present our experimental setup and results. We show that (i) unlike naive implementations using existing libraries, our dynamic sparsity attention schemes can significantly improve over the FlashAttention 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-and 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 OpenWebText2 (Gao et al., 2020). We experiment with QK-dropping based sparsity on OpenWebText2.

Models & Baselines. For our language modeling experiments on OpenWebText2, we use a base autoregressive transformer architecture with 12 layers, a hidden size of 768 , 12 heads of 64 dimensions each. For experiments on sequence length T = 8192 , we use a batch size of 96 = 4 × 8 × 2 (batch size 4 with 8 accumulation steps and data parallelism over 2 node). When T = 16384 we use a batch size of 30 = 2 × 5 × 3 . The resulting models are of around 122 Mparameters. The goal not being to outperform the state-of-the-art perplexity, we train for 15 k iterations. The attention modules used are either using FlashAttention for the baselines or one of our sparse kernels for our 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 bfloat16 . For our language modeling tasks on OpenWebText2, we trained using dataparallelism on two or three A100s 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

ms

0

20000

40000

60000

Sequence length

0

20

40

60

80

100

(a) Forward attn. runtimes

0

20000

40000

60000

Sequence length

0

.

6

0

.

7

0

.

8

0

.

9

1

.

0

Coverage

Reformer attention

Flash attention

Hash-sparse (ours)

(b) Coverage of Reformer attn.

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 as the set of vectors that map to a 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 a 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 be 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 a 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 baseline only the transposition is required, see App. B.2 for detailed code. Fig. 3.b summarises our findings. We observe large improvements in runtime as the number of buckets nb and the sequence length increases.

Language modeling on OpenWebText2. 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 nb = 16 hash buckets. In Fig. 6 we see that it takes the same number of iterations for H-LM and F-LM to reach a given perplexity. However, H-LM iterations are 1 . 8 × and 2 . 3 × faster for respectively T = 8192 and T = 16384 . As a result, H-LM models reach a given perplexity much faster than their F-LM counterpart. Interestingly, we observe the H-LM models gain speed during training, see App. C for additional details.

Comparison with Reformer. We compare the speed and performance of our hash-sparse implementation 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 of 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

70

0

5000

10000

15000

Iterations

20

30

40

50

60

Validation Perplexity

(a) Iter to reach perplexity

0

50

100

150

Time (hours)

20

30

40

50

60

Validation Perplexity

(b) Time to reach perplexity

T

= 8192

T

= 16384

0

2

4

6

8

10

Iter/min on 1 GPU

×

2

.

0

×

3

.

3

F-LM

T

= 8

k

F-LM

T

= 16

k

H-LM

T

= 8

k

H-LM

T

= 16

k

(c) Speed gains

70

matrix. This is reflected in Table 5: our hash-sparse attention layer outperforms the Reformer attention even for shorter sequences.

4.3 Query/Key-Dropping 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 a vacuum. We test our 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 a 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 to build those smaller tensors. In Fig. 7.a, we show a naive implementation using existing PyTorch functionalities only starts to provide a 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 OpenWebText2. For sequences of length T = 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 FLM models in perplexity per iterations, while training nearly twice as fast. Importantly, the dropping pattern is not static. An interesting approach similar to curriculum learning in which we start the training with very large sparsity and reduce it linearly during training is studied in App. C.

ms

0

1000

2000

3000

4000

Sequence length

0

10

20

30

40

50

50%

60%

70%

80%

QK-sparse (our method)

Naive QK-sparse

Flash attention

(a) Naive implementation

0

5000

10000

15000

20000

25000

30000

35000

Sequence length

0

200

400

600

800

1000

ms

20%

30%

40%

50%

60%

70%

80%

0

2000

4000

0

20

40

60

(b) Our implementation

70

0

5000

10000

15000

Iterations

20

30

40

50

60

70

Validation Perplexity

(a) Iter to reach perplexity

0

20

40

60

Time (hours)

20

30

40

50

60

Validation Perplexity

(b) Time to reach perplexity

F-LM

D-LM

s

= 30%

D-LM

s

= 50%

D-LM

s

= 70%

0

5

10

15

Iter/min on 1 GPU

×

1

.

9

×

2

.

6

×

3

.

5

(c) Speed gains

60

5 Conclusion

We develop and validate an 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 in a way that is less constrained by a tight computational budget.

The computational cost of large attention models remains both a practical issue in scaling up to very large contexts, and a 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 an obvious direction to 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 lsh 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 Linguistics, 2020.

Beltagy, I., Peters, M. E., and Cohan, A. Longformer: The long-document transformer. CoRR , abs/2004.05150, 2020. URL https://arxiv.org/abs/2004.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 9781450383097. doi: 10.1145/3442188.3445922. URL https://doi.org/10.1145/3442188.3445922 .

Borgeaud, S., Mensch, A., Hoffmann, 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/2112.04426, 2021. URL https://arxiv.org/ abs/2112.04426 .

Child, R., Gray, S., Radford, A., and Sutskever, I. Generating long sequences with sparse transformers. CoRR , abs/1904.10509, 2019. URL http://arxiv.org/abs/1904.10509 .

Choromanski, K., Likhosherstov, V., Dohan, D., Song, X., Gane, A., Sarlós, T., Hawkins, P., Davis, J., Mohiuddin, A., Kaiser, L., Belanger, D., Colwell, L. J., and Weller, A. Rethinking attention with performers. CoRR , abs/2009.14794, 2020. URL https://arxiv.org/abs/2009.14794 .

Dao, T., Fu, D. Y., Ermon, S., Rudra, A., and Ré, C. Flashattention: Fast and memory-efficient 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. OpenWebText2 dataset, as part of 'the Pile: An 800gb dataset of diverse text for language modeling'. arXiv preprint 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. hutter1. net , 6, 2012.

Katharopoulos, A., Vyas, A., Pappas, N., and Fleuret, F. Transformers are RNNs: Fast autoregressive transformers with linear attention. In Proceedings of the International Conference on Machine Learning (ICML) , pp. 5294-5303, 2020. URL https://fleuret.org/papers/ katharopoulos-et-al-icml2020.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 . OpenReview.net, 2020.

LeCun, Y., Bottou, L., Bengio, Y., and Haffner, P. Gradient-based learning applied to document recognition. Proc. IEEE , 86(11):2278-2324, 1998.

Li, J., Cotterell, R., and Sachan, M. Differentiable subset pruning of transformer heads. Trans. Assoc. Comput. Linguistics , 9:1442-1459, 2021.

Loshchilov, I. and Hutter, F. Decoupled weight decay regularization. In ICLR (Poster) . OpenReview.net, 2019.

Michel, P., Levy, O., and Neubig, G. Are sixteen heads really better than one? In NeurIPS , pp. 14014-14024, 2019.

Peng, H., Schwartz, R., Li, D., and Smith, N. A. A mixture of h - 1 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, Y. cosformer: Rethinking softmax in attention. In ICLR . OpenReview.net, 2022.

Radford, A., Wu, J., Child, R., Luan, D., Amodei, D., and Sutskever, I. Language models are unsupervised multitask learners, 2019.

Raganato, A., Scherrer, Y., and Tiedemann, J. Fixed encoder self-attention patterns in transformerbased 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 , abs/2009.06732, 2020. URL https://arxiv.org/abs/2009.06732 .

Tay, Y., Bahri, D., Metzler, D., Juan, D., Zhao, Z., and Zheng, C. Synthesizer: Rethinking selfattention 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. In 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 9781450367196. doi: 10.1145/3315508. 3329973. URL https://doi.org/10.1145/3315508.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/1706.03762 .

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.04359, 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 Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual , 2020. URL https://proceedings.neurips. cc/paper/2020/hash/c8512d142a2d849725f31a9a7a361ab9-Abstract.html .

Zheng, L., Wang, C., and Kong, L. Linear complexity randomized self-attention mechanism. In ICML , volume 162 of Proceedings of Machine Learning Research , pp. 27011-27041. PMLR, 2022.

Zhou, W., Ge, T., Wei, F., Zhou, M., and Xu, K. 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.

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 a linear computational complexity w.r.t. the sequence length. A follow up work by Zheng et al. (2022) augments RFA with randomized attention to build an unbiased estimator of the softmax attention. Instead of looking for a linear approximation of the softmax, Qin et al. (2022) propose to replace the softmax by a linear projection kernel and a cos-based re-weighting mechanism which scales linearly with the sequence length. While kernel based method seem like a good compromise in terms of speed vs. performance, they have been shown to 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 × 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 × H ′ × T × T , with H ′ the number of non-dropped heads. Michel et al. (2019) and Voita et al. (2019) both observe-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 a 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 in mind. In 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 × H × T Q × T KV , with T Q and T KV 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-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.

B Details On The Implementation

We here give implementation details for experiments of § 4. First, in App. B.1, 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 hyperparameters used in our sequence modeling experiments of § 4. The code for all our experiments can be found via the following link: https://github.com/epfml/dynamic-sparse-flash-attention .

B.1 Triton Kernels

QK-Sparse Triton Kernel. In Alg. 1 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 iterating over block of key indices and comparing them with the largest query index for the current block of queries. This works as q idx and k idx have a monotonic structure thanks to the stable sort used when reshaping the tensors (see the pre-processing code in App. B.2 for more details). We apply causal masking locally by looking at q idx and k idx , q query can only attend to past keys: mask = q_idx[:, None] >= k_idx[None, :] . The backward pass relies exactly on the same trick, we first iterate over query indices to find the starting block of queries.

Algorithm 1 Forward Pass for the QK-sparse kernel

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 pre-processing code in App. B.2 for more details). In addition to causal masking, we also

enforce attention to happen within the same bucket: mask = (q_idx[:, None] >= k_idx[None, :]) & (q_hash[:, None] == k_hash[None, :]) . In our experiments we often replace >= by > to prevent a query to attend to itself as in the Reformer. As a side note, for most application it would also be fine to only enforce causal masking and allow attention across buckets within a 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.

Accumulating softmax statistics while avoiding NaNs. The following is a 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 ) ≜ e x -max( x ) / ∑ i e x i -max( x ) . Let f ( x , m ) ≜ e x -m , and ℓ ( x , m ) ≜ ∑ i f ( x , m ) . Hence:

Given a vector x ≜ [ x 1 , x 2 ] , let m g = max( x ) (global max), and m 1 = max( x 1 ) , we notice:

Therefore, if we have m 1 , ℓ ( x 1 , m 1 ) , and r = f ( x 1 ,m 1 ) l ( x 1 ,m 1 ) , we can update the softmax statistics for a new block of entries x 2 by following the following steps:

Compute new global max: m g = max( m 1 , max( x 2 ))

Compute f ( x 2 , m g ) = e x 2 -m g

Compute ℓ ( x 2 , m g ) = ∑ i f ( x 2 , m g )

Compute new ℓ ( x , m g ) : ℓ ( x , m g ) = e m 1 -m g ℓ ( x 1 , m 1 ) + ℓ ( x 2 , m g )

Correct running softmax result: r = r ℓ ( x 1 ,m 1 ) ℓ ( x ,m g )

Add contribution from x 2 to r : r = r + f ( x 2 ,m g ) ℓ ( x ,m g )

In case all the keys are masked for a given query, we would have max( x 2 ) = -∞ , given that the first m 1 is initialized to -∞ as well (see Alg. 1 and Alg. 2) the fourth and second steps above would be undefined and result in NaN values. We solve the problem by replacing -∞ values in m g by 0 s when doing those two steps. Another potential issue is in step three: when a query has no matching key in x 2 then ℓ ( x 2 , m g ) is now 0 , which generate + ∞ 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 ∞ values to r would prevent us to accumulate statistics later on. To get our desired behaviour and have 0 s when queries have no matching key we replace ∞ values by 1 s in ℓ ( x , m g ) during steps five and six. We summarize those steps in Alg. 3.

Hyperparameters. We extend the implementation of FlashAttention available in the Triton tutorial. In our benchmarks, we use a batch size B = 4 , 48 heads of 64 dimensions each.

Effect of imbalanced dropping/hashing. Our QK-sparse pattern relies on building smaller "compact" tensors containing the keys and queries that 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 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 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. A smart QK-dropping mechanism on top of 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

bucket would equate to the normal FlashAttention. Therefore unbalanced buckets would increase the runtime. Interestingly, in our experiments on language modeling, we do not face imbalance problems and even notice a speedup early in the training, see Fig. 14.(a) in App. C.

B.2 Runtimes in a Vacuum

Baseline implementation. We use Pytorch's FlashAttention implementation provided by the torch.nn.functional.scaled_dot_product_attention function. To ensure fairness, we assume that all benchmarked functions receive a tensor Q of shape (BATCH, CTX_Q, H, D_HEAD) , and tensors K , V of shapes (BATCH, CTX_KV, H, D_HEAD) , where BATCH is the batch size, CTX_Q is the number of queries, CTX_KV is the number of keys and values, H is the number of heads, and D_HEAD is number of dimensions per head. For this reason the only pre and post processing steps required are transposing the input and output tensors.

Listing 1: pytorch_full_flashattention function applying the FlashAttention algorithm on the entire sequence (no sparsity).

Our proposed interface. We propose the following interface to orchestrate between the Hashsparse and the QK-sparse implementations:

Pre & post processing steps for Hash-sparse. In addition to having to transpose the Q , K , V tensors. The preprocessing steps consist in re-ordering the Q , K and V tensors based on bucket indices in q_hash 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 to ensure the queries, keys and values are sorted within each bucket. The following code is showing how we implemented all those steps using Pytorch:

Pre & post processing steps for QK-sparse. In addition to having to transpose the Q , K , V tensors. The preprocessing steps consist in 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 be padded so our kernel can rely on those indices to compute which tile it should and shouldn't compute.

B.3 Sequence Modeling Experiments

Language modeling on OpenWebText2. Our implementation is based on NanoGPT (github.com/karpathy/nanoGPT). We use the AdamW optimizer (Loshchilov & Hutter, 2019). We used bfloat16 and NVIDIA A100-40GB GPUs for all our experiments. Here is a list of hyperparameters 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 beta1: 0 . 9

Adam beta2: 0 . 95

Tokenizer: We use the GPT2 tokenizer provided by the tiktoken library (github.com/openai/tiktoken).

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 8 with 2 gradient accumulation steps. The sequence length is 4096 . For autoregressive image generation on MNIST, we use a smaller model with 8 transformer blocks and a hidden dimension of 256 . Dropout is set to 0 . 1 and the batch size is 10 . We train for 25 epochs.

C Additional Details and Analysis

Quadratic computational cost of attention in transformers. In Fig. 9 we show the runtime (forward + backward) of a transformer language model as a function of the sequence length. We separate the time taken by the attention operations from the time taken by the rest of the model. We see how the attention computation ends up dominating the runtime as the sequence length increases.

ms

0

5000

10000

15000

Sequence length

0

100

200

300

400

500

600

Attention

Feed-forward

Additional runtimes performances in a vacuum. In Fig. 10 and Fig. 11 we show the runtime details for the forward and backward methods separately for respectively the Hash-sparse and QKsparse methods. We also measure runtimes of the forward and backward passes when we assume the pre and post-processing steps aree free, see Fig. 12.

Linear QK-dropping scheduler. In the main paper we show results dropping keys and queries at random with a fixed pre-defined probability. In an additional experiment we start by dropping 80% of keys and queries at random and linearly decay this probability to 20% . Our intuition is earlier iterations just aim to learn contextual cues which are very redundant (and therefore quite immune to random dropping) before requiring more fine-grained representations. In 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 no speedup over FlashAttention. Interestingly, when training on real data such as OpenWebText2, we observe our Hash-sparse based models are speeding up during training. In Fig. 14 we plot the number of iterations reached after x 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.

Different bucket sizes nb for H-LM. For sequences of 8192 tokens, we show the influence of increasing the number of buckets nb in Fig. 15. As nb increases the runtime decreases.

400

0

5000

10000

15000

20000

25000

30000

35000

Sequence length

0

50

100

150

200

250

300

350

ms

nb

= 2

nb

= 4

nb

= 8

nb

= 16

nb

= 32

nb

= 64

Hash-sparse (our method)

Flash attention

0

2000

4000

0

2

4

6

8

10

(a) Forward

0

5000

10000

15000

20000

25000

30000

35000

Sequence length

0

100

200

300

400

500

600

700

800

ms

nb

= 2

nb

= 4

nb

= 8

nb

= 16

nb

= 32

nb

= 64

Hash-sparse (ours)

Flash attention

0

2000

4000

0

5

10

15

20

(b) Backward

400

0

5000

10000

15000

20000

25000

30000

35000

Sequence length

0

50

100

150

200

250

300

350

ms

20%

30%

40%

50%

60%

70%

80%

0

2000

4000

0

2

4

6

8

10

(a) Forward

0

5000

10000

15000

20000

25000

30000

35000

Sequence length

0

100

200

300

400

500

600

700

ms

20%

30%

40%

50%

60%

70%

80%

QK-sparse (ours)

Flash attention

0

2000

4000

0

5

10

15

20

(b) Backward

800

Training H-LMs for more iterations. In Fig. 8 and Fig. 6 we show results of language models trained on OpenWebText2 for 15 k iterations. To verify whether our finding are consistent when you train for more iterations we also try training for 50 k iterations. For sequences of 8192 tokens, using the same hyperparameters described in App. B.3, we show in Fig. 16 that our findings for our Hash-sparse based models do hold when training for more iterations-we 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 we show the attention matrices for two different heads and show how the coverage-the percentage of key-query interactions actually computed vs. what should be computed according to the hash indices computed for keys and queries-of the Reformer LSH algorithm can be low.

D Limitations and Societal Impacts

Limitations. 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.

ms

0

5000

10000

15000

20000

25000

30000

35000

Sequence length

0

200

400

600

800

1000

20%

30%

40%

50%

60%

70%

80%

0

2000

4000

0

5

10

15

20

25

(a) QK-sparse

0

5000

10000

15000

20000

25000

30000

35000

Sequence length

0

200

400

600

800

1000

1200

ms

nb

= 2

nb

= 4

nb

= 8

nb

= 16

nb

= 32

nb

= 64

0

2000

4000

0

10

20

(b) Hash-sparse

0

5000

10000

15000

Iterations

20

30

40

50

60

70

Validation Perplexity

(a) Iter to reach perplexity

0

20

40

60

Time (hours)

20

30

40

50

60

70

Validation Perplexity

(b) Time to reach perplexity

0

5

10

15

Iter/min on 1 GPU

×

1

.

9

×

1

.

5

×

2

.

6

×

3

.

5

F-LM

D-LM

s

= 30%

D-LM

s

= 80%

20%

D-LM

s

= 50%

D-LM

s

= 70%

(c) Speed gains

0

20

40

60

Time (hours on 1 GPU)

0

5000

10000

15000

Iterations

F-LM

H-LM

nb

= 8

H-LM

nb

= 16

H-LM

nb

= 32

(a) Time comparison H-LM vs. F-LM

0

20

40

60

Time (hours on 1 GPU)

0

5000

10000

15000

Iterations

F-LM

D-LM

s

= 30%

D-LM

s

= 80%

20%

D-LM

s

= 70%

D-LM

s

= 50%

(b) Time comparison D-LM vs. F-LM

10

T=8192

T

= 16384

0

2

4

6

8

Iter/min on 1 GPU

×

1

.

8

×

2

.

0

×

2

.

1

×

2

.

6

×

3

.

3

×

3

.

6

F-LM

H-LM

nb

= 16

H-LM

nb

= 8

H-LM

nb

= 32

10000

15000

20000

25000

30000

35000

40000

45000

50000

Iterations

18

19

20

21

22

23

24

25

26

Validation Perplexity

F-LM

H-LM

nb

= 16

H-LM

nb

= 8

(a) Iter to reach perplexity

0

2

4

6

8

10

Iter/min on 1 GPU

×

1

.

9

×

2

.

1

F-LM

H-LM

nb

= 16

H-LM

nb

= 8

(b) Speed gains

Societal impacts. The attention mechanism is central to Large Language Models (LLMs). Moreover, 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).

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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.\n\nQUESTIONS:\n- 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}]$\n\nLIMITATIONS:\nThe author discussed the limitation of the method and the societal impacts. Other limitations of my concern can refer to the weaknesses.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-05"}

comments: 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.

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.

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.

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.

QUESTIONS:
Please see the weakness section.

LIMITATIONS:
N/A
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-06

comments: 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.

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.

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.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-08

comments: 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.

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.

2: 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.

3: 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. 

2: One more experiment on other task, like text-classification task will be great. 

3: 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.

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
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-14

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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\nSTRENGTHS:\n1. 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\nWEAKNESSES:\n1. 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\nQUESTIONS:\n1. 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\nLIMITATIONS:\nNot necessary", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-24"}