SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Matteo Pagliardini*
matteo.pagliardini@epfl.ch
Daniele Paliotta*
daniele.paliotta@unige.ch
Martin Jaggi
martin.jaggi@epfl.ch
François Fleuret
francois.fleuret@unige.ch
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.
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 \times$  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 hash-based 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).
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.
Figure 1: Proposed sparsification of the attention matrix for a given attention head. In each depicted attention matrix, black areas indicate coefficients to 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 to 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 a 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 or query with the colors blue/green/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.
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 \times D$  query tensor  $\mathbf{Q}$  and a  $T \times D$  key tensor  $\mathbf{K}$ , we consider the matrix of dot-products  $\mathbf{Q} \mathbf{K}^{\top}$ , which can become impractical to compute for very long sequences. However, we are only interested in the row-wise  $\mathrm{softmax}(\mathbf{Q} \mathbf{K}^{\top})$ , 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 \times D$  query tensor  $Q$  by a  $T \times D$  key tensor  $K$ , to obtain a matching score matrix, which is then rescaled and row-normalized with softmax, to get a  $T \times T$  attention matrix  $A$ . This matrix is then multiplied by a  $T \times 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.
QK-sparse
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 in the case of irregular causal masking and band block sparsity due to hashing (right), FlashAttention does not provide means to compute a 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 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.
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\times H\times T\times 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  $\pmb{Q}$  of shape  $T_{Q}\times D$ , and a key  $\pmb{K}$  and value  $\pmb{V}$  tensors of shapes  $T_{KV}\times D$ . The algorithms described below will be run in parallel for all elements of the Cartesian product  $B\times H$ . We split tensors into blocks:  $\pmb{Q}\triangleq [Q_0,\dots ,Q_m],\pmb {K}\triangleq [\pmb {K}_0,\dots ,\pmb {K}_n]$ . We define a tile  $\mathcal{T}_{i,j}\triangleq \mathbf{Q}_i\mathbf{K}_j^\top$ , 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}, \ldots, 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} \in \mathbb{R}^{\widetilde{T}_{Q}}$  and  $k^{idx} \in \mathbb{R}^{T_{KV}}$  representing the indices of the queries and keys in the original uncompressed tensors. Those are similarly split into blocks:  $q^{idx} \triangleq [q_{0}^{idx}, \ldots, q_{m}^{idx}]$ . The condition for a tile  $T_{i,j}$  to be unnecessary to compute is now to have  $\max(q_{i}^{idx}) < \min(k_{j}^{idx})$ . When processing a block of queries  $Q_{i}$ , we iterate over the key indices  $k_{0}^{idx}, \ldots, k_{n}^{idx}$  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 \in [0, j_{stop}]$ . Within each tile  $T_{i,j}$ , we in addition apply a local causal mask by comparing indices in  $q_{i}^{idx}$  and  $k_{j}^{idx}$ . 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/query 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  $\pmb{q}^{hash}$  and  $\pmb{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  $\pmb{q}_i^{hash}$  are present in  $\pmb{k}_j^{hash}$ ,  $j_{stop}$  is the first index for which all indices in  $\pmb{k}_j^{hash}$  are strictly larger than indices in  $\pmb{q}_i^{hash}$ . In a second step we refine  $j_{stop}$  now based on the indices  $\pmb{k}^{idx}$  and  $\pmb{q}^{idx}$ , the updated  $\hat{j}_{stop}$  is the last index  $j\in [j_{start},j_{stop}[$  for which  $\max (\pmb {q}_i^{idx})\geq \min (\pmb {k}_j^{idx})$ . As shown in the last column of Fig. 2, we then only compute tiles  $T_{i,j}$  for  $j\in [j_{start},\hat{j}_{stop}]$ . As for the QK-sparse method, we use  $\pmb{q}^{idx}$  and  $\pmb{k}^{idx}$  to apply a causal mask locally for each tile. In addition to the causal mask, we use  $\pmb{q}^{hash}$  and  $\pmb{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.
Figure 3: Comparing several hash-based sparse attention implementations with FlashAttention. Similarly to QK-dropping-based sparsity in Fig. 7, due to the non-triangular causal mask resulting from re-ordering the tensors based on the hash buckets (see Fig. 1), a naive implementation would force the computation of the entire attention matrix before applying a 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 is required. While there is a 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.
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.
Datasets. We test our hash-based sparsity scheme on MNIST (LeCun et al., 1998) for autoregressive image generation, enwik8 (Hutter, 2012), and OpenText2 (Gao et al., 2020). We experiment with QK-dropping based sparsity on OpenText2.
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 \times 8 \times 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 \times 5 \times 3$ . The resulting models are of around 122M parameters. The goal not being to outperform the state-of-the-art perplexity, we train for  $15k$  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 data-parallelism 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
(a) Forward attn. runtimes
Figure 4: Comparing forward runtimes of attention modules alone. Fig.(a): Reformer attention ensures a linear computational complexity w.r.t. the sequence length, outperforming FlashAttention for longer sequences. Fig.(b): However, due to the fixed attention structure, the Reformer misses an increasing 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.
(b) Coverage of Reformer atm.
Figure 5: Comparing models using Reformer attention vs our Hash-sparse attention. On the simple sequential MNIST task (predicting pixels as a sequence), we obtain a comparable perplexity as the Reformer. On en-wik8 character language modeling, with  $T = 4096$ , we outperform the Reformer model by a margin.
Attention	MNIST (ppl ↓)	enwik8 (bits/c ↓)
Reformer	1.76	3.32
Hash-sparse	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 as the set of vectors that map to a certain hash.