Title: FlashMask: Reducing the Complexity of Attention Computation through Sparse Mask Representation
Venue: NeurIPS 2024
Accepted: no
Keywords: FlashMask

Efficient Attention Computation

Sparse Mask Representation

Linear Memory Complexity

Low Computational Complexity
Abstract: Recent advancements in Larger-Scale Transformers have significantly benefited from sophisticated attention mechanisms, which are critical for modeling long-context sequences. However, the computational and memory demands of conventional attention mask computations, typically scaling with an $\mathcal{O}(N^2)$ complexity where $N$ is the sequence length, pose significant challenges. This paper introduces FlashMask, a simple yet effective \emph{Exact} attention algorithm designed to substantially reduce both the computational complexity and memory requirements of attention computations. By adopting a novel column-wise sparse representation of attention masks, FlashMask achieves a linear memory complexity of $\mathcal{O}(N)$ and computational complexity of $\mathcal{O}(N)\sim\mathcal{O}(N^2)$. We assess the performance of FlashMask in a variety of masking scenarios, including causal and customized attention masks, demonstrating its versatility and robustness across a wide range of attention patterns and models. Our empirical analysis encompasses a variety of downstream training modalities, including Supervised Fine-Tuning (SFT), Direct Preference Optimization (DPO), and Reward Model (RM). We compare FlashMask against state-of-the-art techniques, including notably FlashAttention. In kernel-level assessments, FlashMask achieves substantial computational speedups, up to 6.7x (SFT), 6.9x (DPO), and 8.3x (RM). Furthermore, in end-to-end training, FlashMask consistently enhances training speed significantly, with accelerations up to 2.4x (SFT), 4.2x (LoRA), 2.5x (DPO), and 2.6x (RM) across these varied scenarios without sacrificing model accuracy. Additionally, when implemented in the LoRA scenario, FlashMask enables the LLaMA2-7B to process sequence lengths of up to 544k, significantly enhancing its capability for long-context input.

Abstract

Recent advancements in Larger-Scale Transformers have significantly benefited from sophisticated attention mechanisms, which are critical for modeling longcontext sequences. However, the computational and memory demands of conventional attention mask computations, typically scaling with an O( 𝑁 2 ) complexity where 𝑁 is the sequence length, pose significant challenges. This paper introduces FlashMask, a simple yet effective Exact attention algorithm designed to substantially reduce both the computational complexity and memory requirements of attention computations. By adopting a novel column-wise sparse representation of attention masks, FlashMask achieves a linear memory complexity of O( 𝑁 ) and computational complexity of O( 𝑁 ) ∼ O( 𝑁 2 ) . We assess the performance of FlashMask in a variety of masking scenarios, including causal and customized attention masks, demonstrating its versatility and robustness across a wide range of attention patterns and models. Our empirical analysis encompasses a variety of downstream training modalities, including Supervised Fine-Tuning (SFT), Direct Preference Optimization (DPO), and Reward Model (RM). We compare FlashMask against state-of-the-art techniques, including notably FlashAttention [1]. In kernel-level assessments, FlashMask achieves substantial computational speedups, up to 6.7x (SFT), 6.9x (DPO), and 8.3x (RM). Furthermore, in end-to-end training, FlashMask consistently enhances training speed significantly, with accelerations up to 2.4x (SFT), 4.2x (LoRA), 2.5x (DPO), and 2.6x (RM) across these varied scenarios without sacrificing model accuracy. Additionally, when implemented in the LoRA scenario, FlashMask enables the LLaMA2-7B to process sequence lengths of up to 544k, significantly enhancing its capability for long-context input.

1 Introduction 24

Transformers [2] , equipped with self-attention mechanisms, have revolutionized natural language 25

processing (NLP) by efficiently modeling data dependencies without the limitations of sequential 26

processing. This makes them ideal for handling long sequences. Large Language Models (LLMs), 27

which utilize training paradigms such as Supervised Fine-Tuning (SFT) [3, 4] and Reinforcement 28

Learning from Human Feedback (RLHF) [5, 6], critically rely on selective attention management 29

30

through masks.

Effective mask management is essential to focus selectively on pertinent data

segments, optimizing both performance and computational efficiency. 31

However, the conventional attention mechanism in Transformers entails a quadratic increase in 32

computational and memory demands O( 𝑁 2 ) , where 𝑁 denotes the sequence length. This exponential 33

growth presents substantial challenges as models scale to sequence lengths ranging from 128K to 34

1M in advanced systems like GPT-4 [7], Claude [8], and Gemini [9], necessitating more efficient 35

computational approaches. As sequence lengths extend, the memory load for masked attention 36

computations also grows quadratically, adversely affecting computational speed and the ability to 37 manage various mask configurations across different tasks. Current methodologies often resort to 38 approximate sparse attention strategies [10, 11, 12], which unfortunately trade off precision for 39 computational efficiency, underscoring an essential gap in achieving high precision with reduced 40 computational costs. 41

This paper introduces FlashMask, a novel approach utilizing a sparse mask representation to accelerate 42 attention computations in transformers, effectively addressing both computational and memory scala43 bility issues. Unlike previous methods that compromise accuracy for efficiency, FlashMask provides 44 precise computations without sacrificing accuracy, ensuring high fidelity in attention mechanisms. 45 The contributions of this work include: 46

47

48

49

50

Exact Computation. FlashMask uniquely ensures precise attention computations across varying sequence lengths and tasks. It employs a unique column-wise sparse mask representation, denoted by FlashMaskStart (FMS) and FlashMaskEnd (FME), to precisely mask specific rows within columns, ensuring computational efficiency and accuracy.

51

52

53

54

55

56

57

58

59

60

61

62

Long Context Modeling. FlashMask significantly reduces computational and memory demands, enabling efficient processing of extended sequences critical for deploying LLMs in resource-limited settings.

Efficient Mask Computation. FlashMask leverages strategic sparse masking to increase computational throughput, thereby improving processing speeds and broadening the practical utility of LLMs in diverse real-world scenarios.

Extensive Empirical Validation. Empirical studies validate FlashMask's efficiency in computation and storage. Its practical application in real-world scenarios and integration with existing frameworks underscore its potential impact. Moreover, a comprehensive comparison with state-of-the-art methods like FlashAttention-DenseMask, FlashAttention-Varlen highlights FlashMask's efficiency and versatility.

2 Background

The attention mechanism has revolutionized data handling in NLP by mimicking human selective 63 focus, allowing neural networks to prioritize parts of the input data. This addresses limitations 64 of traditional sequence-to-sequence models, enhancing context awareness in long sequences. The 65 Transformer model by Vaswani et al. [2] implements this mechanism centrally, using multiple parallel 66 attention heads instead of recurrent layers, thus improving efficiency and performance. 67

68

69

70

71

2.1 Attention Computation

Central to the Transformer architecture is the attention mechanism, which computes relevance scores between elements in a sequence to focus more on important aspects and less on others. This mechanism can be expressed as:

where 𝑄 , 𝐾 , 𝑉 , and 𝑀 represent the query, key, value, and mask matrices respectively, derived 72 from the input data, and 𝑑 𝑘 is the dimension of keys. The term 𝑀 incorporates constraints to 73 selectively consider certain parts of the input sequence during attention computation, enabling 74 functionality like masking future tokens in sequence-to-sequence modeling. One inherent challenge 75 with attention is its computational and memory complexity, both of which scale quadratically with 76 the length of the input sequence. Processing long sequences presents significant challenges, which 77 are exacerbated in the downstream pipeline of training large language models (LLMs). Different 78 training stages, such as Supervised Fine-Tuning (SFT/LoRA [3, 4, 13, 14, 15]), Direct Preference 79 Optimization (DPO) [16, 17, 18, 19, 20], Reward Model (RM) [5, 21, 22, 23, 24], and Proximal 80 Policy Optimization (PPO) [25, 6], place diverse demands on the attention mask. 81

82

2.2 Masking Variable-Length Sequences

The advent of large transformer-based models has marked substantial progression in handling 83 increased sequence lengths in natural language processing. Previously, models like BERT [26] and 84

GPT-2 [27] were limited to sequences of approximately 512 tokens, whereas more recent adaptations 85 such as the LLaMA [28, 29, 30], GPT-4 [7] and Claude series [8] stretched these limits to encompass 86 2K to 200K tokens, respectively. Innovations from Google's Gemini [9] have further shifted this 87 boundary, managing up to 1M tokens. Enhanced sequence management within these models employs 88 various masking techniques in the attention matrix, adapting to the length and diversity of input 89 sequences. Techniques such as the use of padding operations are illustrated in Figure 1(a), which help 90 maintain efficiency by allowing uniform processing of diverse input lengths through padding masks. 91 However, conventional padding can lead to inefficiencies due to the diverse sequence lengths typically 92 found in training data, often following a long-tail distribution. This issue is adeptly addressed by 93 dynamic token allocation technologies like InToken [31, 3, 32, 33, 34], which optimize computational 94 resources by adjusting the token count based on actual data needs, significantly improving the training 95 efficiency for datasets with various sequence lengths in Figure 1(b)(c). 96

(a)

(b)

(c)

(d)

E2

E2

E2

0

0

0

E2

S2

S2

S2

S2

E

E1

E1

E

S1

5

5

5

5

5

5

5

5

5

S1

7101010

S1

4

4

4

7101010

S110101010777101010

0

1

Sample1

0

1

Sample 1

0

1

Question

0

1

2

2

2

1

2

3

3

4

Sample2

4

Sample2

4

Answer1

4

5

5

1

6

6

6

6

7

Sample3

7

Sample3

7

Answer 2

7

8

8

8

1

9

の

9

6

0

6

7

8

6

0

1

456

7

8

6

0

1

5

6

7

8

6

0

1

2

3

4

56

7

8

6

Despite having extensive text-handling capabilities, the meticulous design of masking configurations 97 remains crucial for specific training scenarios. The illustrated scenarios in Figure 1(d) and Figure 2 98 depict various specialized masking mechanisms employed to enhance model training efficiency and 99 applicability. Figure 1(d) illustrates a scenario involving DPO/RM with two or more answers, where 100 each answer's tokens have visibility to the tokens of the question, and tokens from different answers 101 are not visible to each other. Multi-shot and in-context learning scenarios facilitated by extended 102 attention spans in configurations like Figure 2(a) are becoming prevalent, which allows the final 103 question in a series to receive comprehensive attention, enhancing contextual understanding [35, 104 36]. Furthermore, hybrid masking forms combining features from different methodologies are 105 demonstrated in Figure 2(b). These incorporate sink tokens [37] and a sliding window mask from the 106 Big Bird [38], facilitating a localized yet extensive context capture. Figure 2(c) is also derived from 107 Big Bird, showing a bi-directional global attention mask, which allows for a comprehensive global 108 context capture. Such innovative approaches in masking not only bolster the efficiency of training 109 large transformer models but also pave the way for advanced explorations into the capabilities of 110 attention mechanisms, such as simulating token eviction during inference as depicted in Figure 2(d). 111 These advancements underscore the dynamic and adaptable nature of transformer technology in 112 accommodating varying training needs and enhancing the overall performance of LLMs. 113

114

2.3 Attention Optimization Techniques

As aforementioned in Equation 1, the computational and memory demands of this mechanism, 115 particularly the computation of 𝑄𝐾 𝑇 , become significant as the sequence length 𝑁 increases. This 116 is due to the size of the resultant attention scores matrix, which scales quadratically with the 117 sequence length, leading to a complexity of O( 𝑁 2 ) . Several related works has been proposed to 118 alleviate the issue. In the realm of model training optimizations, Memory Efficient Attention [39] 119 (MEA) and FlashAttention [1] have been pivotal. MEA focuses on reducing the model's memory 120 demands by altering the self-attention mechanisms. This allows either for the use of larger models 121 or for the extension of maximum sequence lengths within existing hardware constraints. On the 122

123

124

125

126

127

128

129

130

131

132

133

134

135

136

137

138

139

140

141

142

143

144

145

146

147

148

149

150

151

152

153

154

155

(a)

(b)

(c)

(d)

E2

E2

0

0

0

0

6

7

8

6

E2

S2

S2

S2

0

0

0

0

3

3

3

3

3

3

S2

E1

7

7

7

10

10

10101010

E1

E1

E

S1

4

4

4

10

10

10101010

0

Sink Token

S1

10

56789101010

0

4

0

S1104861071091010

0

Shot1

1

1

1

1

2

2

2

2

3

3

3

3

4

4

4

4

5

5

5

6

6

6

6

Question

7

7

7

7

8

8

8

8

9

9

9

0

2

3

4

5

6

7

8

9

D

1234567

8

9

D

1

2

3

4

6

7

8

9

0123456789

other hand, FlashAttention enhances the efficiency of attention mechanisms with IO-Awareness to better utilize contemporary GPU architectures, resulting in faster computations and reduced energy consumption. This method reduces memory overhead to O( 𝑁 ) utilizing tiling techniques during the computation process, making it particularly effective in scenarios without the need for a custom mask. However, for specific training contexts requiring custom masking, the memory overhead with FlashAttention remains O( 𝑁 2 ) . Note that, in typical training setups like unidirectional causal attention or bidirectional full-context attention, the default mode of operation with FlashAttention does not involve passing a custom mask.

During the inference stage, optimizations such as FlashDecoding [40] and FlashDecoding++ [41] play crucial roles. FlashDecoding enhances the decoder in transformers to expedite the generation of sequences by optimizing state management and employing techniques that minimize computational waste. FlashDecoding++ further advances these improvements, incorporating sophisticated dynamic batching and more refined state management to significantly boost throughput and reduce latency. Concerning long sequence training, RingAttention [42] is notable for its efficiency in distributed training contexts, managing communication overhead and memory utilization effectively across multiple nodes.

Another class of study targets on the sparsity/low-rank of attention computation. The Sparse Transformer [10] revolutionizes sequence processing with log-linear complexity. Similarly, Reformer [43] optimizes memory via locality-sensitive hashing, while Big Bird [38] introduces a hybrid attention method to manage longer sequences efficiently. Linformer [44] reduces complexity using low-rank approximations, significantly economizing computation and storage requirements. Both of the previously discussed solutions either compromise precision or yield only marginal enhancements in efficiency. Conversely, our proposed FlashMask is capable of delivering an exact computations.

3 FlashMask: Algorithm and Analysis

In this section, we present the critical design of the column-wise sparse mask representation, implementation of the mask computation kernel, and a complexity analysis of the proposed FlashMask.

3.1 Column-wise Sparse Mask Representation

We introduce FlashMask, a column-wise sparse masking technique, represented using FMS , FME ∈ R 𝑁 (the row index of F lash M ask S tart and F lash M ask E nd), where FMS 𝑐 , FME 𝑐 denote that elements in the 𝑐 -th column of the attention score matrix S = QK 𝑇 within the interval [ FMS 𝑐 , FME 𝑐 ) are masked (set to -∞ ). As shown in Fig. 2(a), FMS = [ 4 , 4 , 4 , 4 , 10 , 10 , 10 , 10 , 10 , 10 ] , FME = [ 7 , 7 , 7 , 7 , 10 , 10 , 10 , 10 , 10 , 10 ] indicates that, for the first column, the 4 -th to 6 -th rows are masked.

3.2 Integration with FlashAttention

Unidirectional (causal) attention, commonly utilized in large language models, incorporates Flash156 Mask within the FlashAttention-2 algorithm, as detailed in Algorithm 1. This paper elaborates the 157 implementation of FlashMask using the lower triangular section of the mask for illustration, where 158 the blue section represents the computation by the dense mask method (for comparison and not 159

160

161

162

163

164

165

166

167

168

169

170

171

172

173

174

175

Algorithm 1 Optimized Forward Pass with FlashMask

present in FlashMask) and the red section indicates the FlashMask computation. FlashAttention Forward involves two nested loops; the outer loop iterates over each block Q 𝑖 of Q , and the inner loop iterates over all blocks K 𝑗 of K and V 𝑗 of V . In the inner loop, S ( 𝑗 ) 𝑖 = QK 𝑇 is computed on SRAM. Once S ( 𝑗 ) 𝑖 is generated, the corresponding dense mask is added as a bias (shown in line 20 of Algorithm 1), whereas FlashMask applies the column-wise sparse mask by setting elements beyond FMS 𝑐 but not exceeding FME 𝑐 to -∞ (as shown in lines 21 to 23 of Algorithm 1).

FlashMask further exploits the block computation feature of FlashAttention-2 to reduce computation. If all elements within a block are masked, the block's computation, including matrix multiplication and softmax operations, can be skipped. A block defined by rows [ 𝑟 0 , 𝑟 1 ) and columns [ 𝑐 0 , 𝑐 1 ) is skipped if 𝑟 0 ≥ max ( FMS 𝑐 0 : 𝑐 1 ) and 𝑟 1 ≤ min ( FME 𝑐 0 : 𝑐 1 ) . Considering that mask regions often exhibit continuity, most blocks are either completely masked or not at all, with only boundary blocks requiring fine-grained masking. A block is completely unmasked if every coordinate ( 𝑟, 𝑐 ) satisfies 𝑟 < FMS 𝑐 or 𝑟 ≥ FME 𝑐 , thus skipping fine-grained masking and avoiding extra masking overhead.

176

177

178

179

180

181

182

183

184

185

186

187

188

189

190

191

192

193

194

195

196

197

198

199

200

201

202

203

204

205

206

207

208

209

210

211

212

213

214

215

216

217

218

219

parallel reduction operation within each thread block then computes the maximum and minimum values, yielding 𝑇 𝑐 values. The additional space complexity introduced here is O( 𝑇 𝑐 ) . Similar computations are made for max ( FMS 𝑐 0 : 𝑐 1 ) , min ( FMS 𝑐 0 : 𝑐 1 ) ,.

The backward computation in FlashAttention-2, which is typically column-parallel, benefits more from the column sparse mask approach. Blocks for which j max ( FMS 𝑐 0 : 𝑐 1 ) 𝐵 𝑟 k < 𝑖 < j min ( FME 𝑐 0 : 𝑐 1 ) 𝐵 𝑟 k are fully masked, allowing skipping of these intervals directly. Only blocks satisfying j min ( FMS 𝑐 0 : 𝑐 1 ) 𝐵 𝑟 k ≤ 𝑖 ≤ j max ( FME 𝑐 0 : 𝑐 1 ) 𝐵 𝑟 k require fine-grained masking.

It is important to note that unlike various approximate attention algorithms, our method ensures that each effective element of the attention score matrix is computed identically to FlashAttention-2, with masked elements explicitly set to -∞ , thus maintaining the accuracy of the algorithm's results. Futhermore, FlashMask is easily extendable to bidirectional attention computations.

3.3 Complexity Analysis

We define sparsity as 𝜌 = 𝑝 𝑁 2 , where 𝑝 is the number of masked elements in the attention score matrix, and 𝑁 is the maximum sequence length of Q and K, 𝑁 2 being the total number of elements in the attention score matrix. For a causal mask, 𝜌 = 2 × 𝑝 𝑁 2 since half of the elements in the attention score matrix are already masked by the causal mask. The block sparsity 𝛼 is defined as 𝛼 = 𝑎 l 𝑁 𝐵𝑟 m × l 𝑁 𝐵𝑐 m , where 𝐵 𝑟 , 𝐵 𝑐 are block sizes, and 𝑎 is the number of completely masked blocks. For a causal mask, 𝛼 = 2 × 𝑎 l 𝑁 𝐵𝑟 m × l 𝑁 𝐵𝑐 m .

Space complexity. The dense mask is represented as D ∈ R 𝑁 × 𝑁 , with a space complexity of O( 𝑁 2 ) . FlashMask denotes as FMS , FME ∈ R 𝑁 , occupying O( 𝑁 ) space, along with four precomputed arrays maxFMS , minFMS , maxFME , minFME ∈ R 𝑁 𝐵𝑐 , also occupying O( 𝑁 ) space. Thus, the total space complexity for FlashMask is O( 𝑁 ) , significantly reducing memory usage and supporting

×

l m training on longer sequences.

Memory access complexity. The dense mask accesses the entire D ∈ R 𝑁 𝑁 matrix in line 20 of Algorithm 1, totaling 𝑁 2 memory accesses on HBM. FlashMask reads the FMS , FME ∈ R 𝑁 vectors from HBM as shown in lines 17 and 18 of Algorithm 1, with each Q 𝑖 reading the entire FMS , FME , totaling 2 × 𝑇 𝑟 × 𝑁 memory accesses. This reduces the memory access to approximately 2 × 𝑇 𝑟 × 𝑁 𝑁 2 ≈ 2 𝐵 𝑟 , significantly boosting performance. Due to FlashMask's smaller space usage, it is possible to preload FMS , FME into SRAM using only 2 × 𝐵 𝑐 SRAM, enhancing memory access efficiency. For the backward process, which uses a column-parallel approach, SRAM-stored FMS , FME can be well reused, further reducing the total memory access on HBM to 2 × 𝑁 .

Computational complexity. The attention computation process normally iterates over the entire attention score matrix, with a computational complexity of O( 𝑁 2 ) . By skipping entirely masked blocks, FlashMask leverages block sparsity to reduce computational complexity to O(( 1 -𝛼 ) 𝑁 2 ) .

4 Experiments

4.1 Setup

Experiments were conducted using GPU A800-SXM 80G, Intel(R) Xeon(R) Platinum 8350C CPUs, CUDA 12.0, and driver version 525.125.06. We evaluated FlashMask against various methods including Vanilla Attention, FlashAttention with dense mask (FA-DenseMask), variable length (FAVarlen), and sliding window (FA-Window) across different scenarios and sequence lengths. Both kernel-level and end-to-end performance demonstrated the effectiveness of our method.

4.2 Data Construction

As mentioned in the Background section, commercial large models now support sequences up to 128K in length. FlashMask, with its lower memory overhead, can facilitate training with even longer

220

221

222

223

224

225

226

227

228

229

230

231

232

233

234

235

236

2

4

8

16

32

64

128

Sequence Length(K)

10 0

10 1

10 2

10 3

10 4

Latency(ms)

SFT

FA-Varlen

FA-DenseMask

VanillaAttention

FlashMask

2

4

8

16

32

64

128

Sequence Length(K)

10 0

10 1

10 2

10 3

10 4

Latency(ms)

DPO

FA-DenseMask

VanillaAttention

FlashMask

2

4

8

16

32

64

128

Sequence Length(K)

10 0

10 1

10 2

10 3

10 4

Latency(ms)

RM

FA-DenseMask

VanillaAttention

FlashMask

contexts. However, currently available public datasets do not contain training data for scenarios exceeding 128K. For comprehensive testing of FlashMask, we constructed synthetic data to simulate long-sequence training.

For a given sequence length 𝐿 , sequences were generated by mimicking InToken method with several sub-sequences. Randomly selecting 𝑠 ∈ [ 1 , 10 ] split points uniformly within the range ( 0 , 𝐿 ) , the sequence was divided into 𝑠 sub-sequences. The segment from the last split point to the end of the sequence was considered as Padding. For the RM scenario, shorter sequence lengths used a smaller upper limit on the number of splits: 𝑠 ∈ [ 1 , 3 ] for 𝐿 ∈ ( 0 , 4096 ] and 𝑠 ∈ [ 1 , 4 ] for 𝐿 ∈ ( 4096 , 8192 ] . By discarding samples not meeting size requirements, we ensure each sub-sequence length was at least 128 (SFT, LoRA, DPO) or 512 (RM) and padding not exceeding 128 (SFT, LoRA, DPO) or 512 (RM). Suppose one sub-sequence with length 𝐿 ′ was further divided into a query and 𝑘 answers based on the scenario. The length of each answer was randomly determined from the range [ 0 . 1 𝐿 ′ 1 + 0 . 1 × 𝑘 , 0 . 2 𝐿 ′ 1 + 0 . 2 × 𝑘 ] , making the answer lengths approximately [ 0 . 1 , 0 . 2 ] of the query length. Therefore, the query length was equal to 𝐿 ′ minus the total answer lengths. A total of 240 valid samples per given sequence length 𝐿 were collected and binned into 10 categories by sparsity 𝜌 , as shown in Appendix A.2.

256

512

1024

2048

4096

8192

WindowSize

0

2

4

6

8

10

12

Latency(ms)

Sequence Length 8K

FA-Window

FlashMask

256

512

1024

2048

4096

8192

WindowSize

0

5

10

15

20

25

30

Latency(ms)

Sequence Length 16K

FA-Window

FlashMask

256

512

1024

2048

4096

8192

WindowSize

0

10

20

30

40

50

60

70

Latency(ms)

Sequence Length 32K

FA-Window

FlashMask

0

20

40

60

80

100

Sparsity(%)

0

20

40

60

80

100

120

140

160

180

Latency(ms)

SFT

FA-Varlen

FlashMask

0

20

40

60

80

100

Sparsity(%)

0

20

40

60

80

100

120

140

160

180

Latency(ms)

DPO

FlashMask

0

20

40

60

80

100

Sparsity(%)

0

20

40

60

80

100

120

140

160

180

Latency(ms)

RM

FlashMask

4.3 Kernel Experiments

We conducted tests with batch sizes of 1, 2, and 4 using Vanilla Attention, FA-DenseMask, and 237 FlashMask. Each experiment began with 5 warm-up runs followed by 50 measurements, totaling 55 238 runs with kernel latency as the performance metric. Additional comparisons were made with FA239 Varlen in the SFT scenario. Results for batch size 1 are shown in Figure 3 (results for batch sizes 2 and 240

241

242

243

244

245

246

247

248

249

250

251

252

253

254

255

256

257

258

259

260

261

262

263

264

265

266

267

268

269

270

271

272

273

4

8

16

32

64

128

256

Sequence Length(K)

600

800

1000

1200

1400

1600

1800

Tokens/Sec/GPU

Speed Up 2.46x

LLaMA-7B

VanillaAttention

FA-DenseMask

FA-Varlen

FlashMask

4

8

16

32

64

128

196

Sequence Length(K)

300

400

500

600

700

800

900

1000

Tokens/Sec/GPU

Speed Up 2.35x

LLaMA-13B

VanillaAttention

FA-DenseMask

FA-Varlen

FlashMask

4

8

16

32

64

96

Sequence Length(K)

100

120

140

160

180

200

Tokens/Sec/GPU

Speed Up 1.96x

LLaMA-70B

VanillaAttention

FA-DenseMask

FA-Varlen

FlashMask

(a) SFT

4

8

16

32

64

128

256

512

544

Sequence Length(K)

2000

3000

4000

5000

6000

7000

8000

Tokens/Sec/GPU

Speed Up 4.16x

LLaMA-7B

VanillaAttention

FA-DenseMask

FA-Varlen

FlashMask

4

8

16

32

64

128

224

Sequence Length(K)

400

600

800

1000

1200

Tokens/Sec/GPU

Speed Up 2.59x

LLaMA-13B

VanillaAttention

FA-DenseMask

FA-Varlen

FlashMask

4

8

16

32

64

128

Sequence Length(K)

100

120

140

160

180

200

220

240

Tokens/Sec/GPU

Speed Up 2.10x

LLaMA-70B

VanillaAttention

FA-DenseMask

FA-Varlen

FlashMask

(b) LoRA

4

8

16

32

64

96

Sequence Length(K)

600

800

1000

1200

1400

Tokens/Sec/GPU

Speed Up 2.49x

LLaMA-7B

VanillaAttention

FA-DenseMask

FlashMask

4

8

16

32

64

128

180

Sequence Length(K)

300

400

500

600

700

800

Tokens/Sec/GPU

Speed Up 2.46x

LLaMA-13B

VanillaAttention

FA-DenseMask

FlashMask

4

8

16

32

64

80

Sequence Length(K)

80

100

120

140

160

Tokens/Sec/GPU

Speed Up 2.02x

LLaMA-70B

VanillaAttention

FA-DenseMask

FlashMask

(c) DPO

4

8

16

32

64

128

256

Sequence Length(K)

600

800

1000

1200

1400

1600

1800

2000

Tokens/Sec/GPU

Speed Up 2.60x

LLaMA-7B

VanillaAttention

FA-DenseMask

FlashMask

4

8

16

32

64

128

196

Sequence Length(K)

300

400

500

600

700

800

900

1000

Tokens/Sec/GPU

Speed Up 2.51x

LLaMA-13B

VanillaAttention

FA-DenseMask

FlashMask

4

8

16

32

64

96

Sequence Length(K)

100

120

140

160

180

200

220

Tokens/Sec/GPU

Speed Up 2.04x

LLaMA-70B

VanillaAttention

FA-DenseMask

FlashMask

(d) RM

4 can be found in Appendix A.3). FlashMask demonstrated significant latency advantages across all lengths, up to 8.3-fold time saving compared to FA-DenseMask. Vanilla Attention was significantly more time-consuming and exceeded memory limits at lengths greater than 32K. The closest competitor to FlashMask, FA-Varlen, exhibited higher latencies as sequence lengths increased. Similar trends were observed in the DPO and RM scenarios, with FlashMask significantly outperforming FADenseMask and Vanilla Attention, especially in the RM scenario where higher sparsity levels further enhanced FlashMask's effectiveness. Performance benefits from varying sparsity levels were also quantified, with FlashMask showing linear negative correlation with increasing sparsity, demonstrating efficient utilization of sample sparsity for acceleration. FlashMask's capability to perform sliding window attention was further tested against FA-Window with window sizes of 256, 512, 1024, 2048, 4096, and 8192, as shown in Figure 4 Top. FlashMask matched FA-Window in latency across sequence lengths of 8K, 16K, and 32K, showing comparable delay performances at increasing window sizes.

4.4 End-to-End Experiments

The end-to-end performance 1 of the model was tested using synthetic datasets across three scales of the LLaMA2 model and four downstream scenarios (SFT, LoRA, DPO, RM) at various sequence lengths, measuring throughput in average Tokens/Sec/GPU. Each sequence length of 240 valid samples was trained for one epoch, with results presented in Figure 5. In the SFT scenario, FlashMask showed a clear throughput advantage over FA-DenseMask and Vanilla Attention, performing comparably to FAVarlen. As sequence lengths increased, the throughput advantage of FlashMask over FA-DenseMask and Vanilla Attention also enhanced, even enabling the completion of longer sequence tasks within the same computational resources. In LoRA, DPO, and RM scenarios, FlashMask consistently showed significant advantages. Notably, in the LoRA scenario at the LLaMA2-7B, FlashMask achieved a 4.16x throughput improvement over FA-DenseMask, supporting sequence lengths up to 544K. It's important to note that FA-Varlen was unable to support the DPO and RM scenarios with the answers sharing one question, whereas FlashMask was capable of handling various scenarios including DPO and RM.

Additional experiments were conducted on the open-source dataset LongBench [45], comparing the end-to-end performance of FA-DenseMask, FA-Varlen, and FlashMask at sequence lengths of 16K, 32K, and 64K. The performance improvements were consistent with those observed in the synthetic dataset. The detailed results are presented in Appendix A.3. Memory usage during the experiments was also recorded, showing significant reductions for FlashMask compared to FA-DenseMask, with detailed results presented in Appendix A.3.

274

275

276

277

278

279

280

281

282

283

284

285

286

287

288

289

290

291

292

293

294

295

296

297

298

299

300

301

302

303

304

305

306

307

308

309

310

311

312

313

314

315

316

317

318

319

320

321

322

323

324

325

326

5 Discussion

Several key topics emerge that are crucial for comprehending the full scope and implications of FlashMask. These include the rationale behind the design choices, adaptations for supporting bidirectional and other custom masks, and the necessity as well as limits of the current approach.

Necessity and Scope of the Study. The substantial advancement rendered by FlashMask in improving attention mask computation is a significant evolution over the current FlashAttention framework. Notably, FlashMask addresses and significantly mitigates the limitations observed with FlashAttention in handling conventional and custom mask computations. This enhancement not only broadens the applicative reach of FlashAttention but also signifies a key shift in efficiency metrics critical for Transformer architectures. More importantly, the flexibility of FlashMask extends beyond the proprietary boundaries of FlashAttention, offering potential benefits to a wider range of Transformerbased models. By facilitating more efficient computation of the attention mechanism, FlashMask enables innovations in processing vast datasets and complex models, thereby improving performances across varied applications in the LLM field. This cross-model adaptability confirms the robustness and utility of FlashMask as a universally applicable enhancement tool within and potentially outside the Transformer architecture spectrum, promising substantial gains in computational efficiency and model scalability.

Bidirectional and Custom Masks. In the exploration of attention mechanisms, the introduction of FlashMask as discussed in this study offers a significant leap in computational efficiency, particularly for masking processes in unidirectional attention mechanisms. By extending this approach to bidirectional networks through the simple addition of vectors indicating the start and end indices of the mask, FlashMask transcends conventional computational bounds, casting itself not just as a sparse attention methodology, but as a versatile computational paradigm. Its adaptability across various custom masking tasks and ability to effectively manage diverse types of mask combinations underscores its potential to greatly enhance the efficiency of attention computations. Moreover, the inherent sparsity of the attention mask during inference provides a robust justification for employing FlashMask, indicating its utility and effectiveness in practical applications. This paradigm shift highlights the importance of developing scalable and efficient computational strategies in the evolving landscape of transformer architectures, suggesting that future research should continue to leverage these innovations to tackle increasing computational demands.

Limitations and Future Directions. While FlashMask demonstrates impressive performance in handling long-context sequences, it is observed that the computational cost of training Transformers increases more than linearly as the sequence length grows-not only due to the computation of masked attention but also because of the extensive use of other operators. This scenario highlights the inevitable need for leveraging or integrating distributed computing strategies or further algorithmic enhancements to elevate training efficiency. Such advancements could be practical in managing the computationally intensive tasks involved in processing extended contexts efficiently. As a part of future research directions, exploring synergistic solutions that combine the strengths of both algorithmic innovation (like FlashMask) and distributed system designs stands as a promising venture. This approach is anticipated to address scalability challenges and could set the stage for breakthroughs in handling unprecedentedly large data sets and complex model architectures.

6 Conclusion

In this paper, we introduced FlashMask, a groundbreaking attention computation paradigm designed to tackle the high computational and memory demands inherent in conventional attention mechanisms in large-scale transformers. By implementing a novel column-wise sparse representation of attention masks, FlashMask substantially reduces the memory and computational complexity from quadratic to linear with the sequence length, thereby enhancing processing speeds and efficiency. Our algorithm demonstrates versatility across various masking scenarios and retains robust performance in different training pipelines. Extensive empirical analysis confirms that FlashMask accelerates computational speed significantly, achieving up to 8.3x speedup in common modalities comparable to state-of-the-art methods like FlashAttention. This advancement marks a significant leap forward in the design of attention computation, offering the potential for broader applications and setting a new benchmark in the efficiency of processing long-context sequences.

327

328

329

330

331

332

333

334

335

336

337

338

339

340

341

342

343

344

345

346

347

348

349

350

351

352

353

354

355

356

357

358

359

360

361

362

363

364

365

366

367

368

369

370

371

372

373

374

References

Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory-efficient exact attention with io-awareness. Advances in Neural Information Processing Systems , 35:16344-16359, 2022.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems , 30, 2017.

Srinivasan Iyer, Xi Victoria Lin, Ramakanth Pasunuru, Todor Mihaylov, Daniel Simig, Ping Yu, Kurt Shuster, Tianlu Wang, Qing Liu, Punit Singh Koura, et al. Opt-iml: Scaling language model instruction meta learning through the lens of generalization. arXiv preprint arXiv:2212.12017 , 2022.

Hyung Won Chung, Le Hou, Shayne Longpre, Barret Zoph, Yi Tay, William Fedus, Yunxuan Li, Xuezhi Wang, Mostafa Dehghani, Siddhartha Brahma, et al. Scaling instruction-finetuned language models. Journal of Machine Learning Research , 25(70):1-53, 2024.

Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, et al. Training language models to follow instructions with human feedback. Advances in neural information processing systems , 35:27730-27744, 2022.

Zhihong Shao, Peiyi Wang, Qihao Zhu, Runxin Xu, Junxiao Song, Mingchuan Zhang, YK Li, Y Wu, and Daya Guo. Deepseekmath: Pushing the limits of mathematical reasoning in open language models. arXiv preprint arXiv:2402.03300 , 2024.

Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. Gpt-4 technical report. arXiv preprint arXiv:2303.08774 , 2023.

Anthropic. Introducing claude. https://www.anthropic.com/news/introducing-claude , 2024. Accessed: May 20, 2024.

Machel Reid, Nikolay Savinov, Denis Teplyashin, Dmitry Lepikhin, Timothy Lillicrap, Jean-baptiste Alayrac, Radu Soricut, Angeliki Lazaridou, Orhan Firat, Julian Schrittwieser, et al. Gemini 1.5: Unlocking multimodal understanding across millions of tokens of context. arXiv preprint arXiv:2403.05530 , 2024.

Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. arXiv preprint arXiv:1904.10509 , 2019.

Beidi Chen, Tri Dao, Eric Winsor, Zhao Song, Atri Rudra, and Christopher Ré. Scatterbrain: Unifying sparse and low-rank attention. Advances in Neural Information Processing Systems , 34:17413-17426, 2021.

Zhiqing Sun, Yiming Yang, and Shinjae Yoo. Sparse attention with learning to hash. In International Conference on Learning Representations , 2021.

Edward J Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. Lora: Low-rank adaptation of large language models. arXiv preprint arXiv:2106.09685 , 2021.

Shih-Yang Liu, Chien-Yi Wang, Hongxu Yin, Pavlo Molchanov, Yu-Chiang Frank Wang, KwangTing Cheng, and Min-Hung Chen. Dora: Weight-decomposed low-rank adaptation. arXiv preprint arXiv:2402.09353 , 2024.

Yukang Chen, Shengju Qian, Haotian Tang, Xin Lai, Zhijian Liu, Song Han, and Jiaya Jia. Longlora: Efficient fine-tuning of long-context large language models. arXiv preprint arXiv:2309.12307 , 2023.

Rafael Rafailov, Archit Sharma, Eric Mitchell, Christopher D Manning, Stefano Ermon, and Chelsea Finn. Direct preference optimization: Your language model is secretly a reward model. Advances in Neural Information Processing Systems , 36, 2024.

Hanze Dong, Wei Xiong, Deepanshu Goyal, Yihan Zhang, Winnie Chow, Rui Pan, Shizhe Diao, Jipeng Zhang, Kashun Shum, and Tong Zhang. Raft: Reward ranked finetuning for generative foundation model alignment. arXiv preprint arXiv:2304.06767 , 2023.

Tianqi Liu, Yao Zhao, Rishabh Joshi, Misha Khalman, Mohammad Saleh, Peter J Liu, and Jialu Liu. 375 Statistical rejection sampling improves preference optimization. arXiv preprint arXiv:2309.06657 , 2023. 376

Kawin Ethayarajh, Winnie Xu, Niklas Muennighoff, Dan Jurafsky, and Douwe Kiela. Kto: Model 377 alignment as prospect theoretic optimization. arXiv preprint arXiv:2402.01306 , 2024. 378

379

380

381

Tianqi Liu, Zhen Qin, Junru Wu, Jiaming Shen, Misha Khalman, Rishabh Joshi, Yao Zhao, Mohammad Saleh, Simon Baumgartner, Jialu Liu, et al. Lipo: Listwise preference optimization through learning-to-rank. arXiv preprint arXiv:2402.01878 , 2024.

382

383

384

385

386

387

388

389

390

391

392

393

394

395

396

397

398

399

400

401

402

403

404

405

406

407

408

409

410

411

412

413

414

415

416

417

418

419

420

421

422

423

Minae Kwon, Sang Michael Xie, Kalesha Bullard, and Dorsa Sadigh. Reward design with language models. arXiv preprint arXiv:2303.00001 , 2023.

Hunter Lightman, Vineet Kosaraju, Yura Burda, Harri Edwards, Bowen Baker, Teddy Lee, Jan Leike, John Schulman, Ilya Sutskever, and Karl Cobbe. Let's verify step by step. arXiv preprint arXiv:2305.20050 , 2023.

Binghai Wang, Rui Zheng, Lu Chen, Yan Liu, Shihan Dou, Caishuang Huang, Wei Shen, Senjie Jin, Enyu Zhou, Chenyu Shi, et al. Secrets of rlhf in large language models part ii: Reward modeling. arXiv preprint arXiv:2401.06080 , 2024.

Alexandre Rame, Guillaume Couairon, Corentin Dancette, Jean-Baptiste Gaya, Mustafa Shukor, Laure Soulier, and Matthieu Cord. Rewarded soups: towards pareto-optimal alignment by interpolating weights fine-tuned on diverse rewards. Advances in Neural Information Processing Systems , 36, 2024.

John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347 , 2017.

Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805 , 2018.

Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAI blog , 1(8):9, 2019.

Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.13971 , 2023.

Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09288 , 2023.

Wei Huang, Xudong Ma, Haotong Qin, Xingyu Zheng, Chengtao Lv, Hong Chen, Jie Luo, Xiaojuan Qi, Xianglong Liu, and Michele Magno. How good are low-bit quantized llama3 models? an empirical study. arXiv preprint arXiv:2404.14047 , 2024.

Mario Michael Krell, Matej Kosec, Sergio P Perez, and Andrew Fitzgibbon. Efficient sequence packing without cross-contamination: Accelerating large language models without impacting performance. arXiv preprint arXiv:2107.02027 , 2021.

Mostafa Dehghani, Basil Mustafa, Josip Djolonga, Jonathan Heek, Matthias Minderer, Mathilde Caron, Andreas Steiner, Joan Puigcerver, Robert Geirhos, Ibrahim M Alabdulmohsin, et al. Patch n'pack: Navit, a vision transformer for any aspect ratio and resolution. Advances in Neural Information Processing Systems , 36, 2024.

PaddleNLP Contributors. Paddlenlp: An easy-to-use and high performance nlp library. https://github. com/PaddlePaddle/PaddleNLP , 2021.

BYTEDANCE INC. Effective transformer. https://github.com/bytedance/effective_ transformer , 2021.

Qingxiu Dong, Lei Li, Damai Dai, Ce Zheng, Zhiyong Wu, Baobao Chang, Xu Sun, Jingjing Xu, and Zhifang Sui. A survey on in-context learning. arXiv preprint arXiv:2301.00234 , 2022.

Amanda Bertsch, Maor Ivgi, Uri Alon, Jonathan Berant, Matthew R Gormley, and Graham Neubig. In-context learning with long-context models: An in-depth exploration. arXiv preprint arXiv:2405.00200 , 2024.

Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficient streaming language 424 models with attention sinks. arXiv preprint arXiv:2309.17453 , 2023. 425

Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, 426 Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, et al. Big bird: Transformers for longer sequences. 427 Advances in neural information processing systems , 33:17283-17297, 2020. 428

Markus N Rabe and Charles Staats. Self-attention does not need o (n2) memory. arXiv preprint 429 arXiv:2112.05682 , 2021. 430

Tri Dao, Daniel Haziza, Francisco Massa, and Grigory Sizov. Flash-decoding for long-context inference, 431 2023. 432

Ke Hong, Guohao Dai, Jiaming Xu, Qiuli Mao, Xiuhong Li, Jun Liu, Kangdi Chen, Hanyu Dong, 433 and Yu Wang. Flashdecoding++: Faster large language model inference on gpus. arXiv preprint 434 arXiv:2311.01282 , 2023. 435

Hao Liu, Matei Zaharia, and Pieter Abbeel. Ring attention with blockwise transformers for near-infinite 436 context. arXiv preprint arXiv:2310.01889 , 2023. 437

Nikita Kitaev, Łukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. arXiv preprint 438 arXiv:2001.04451 , 2020. 439

Sinong Wang, Belinda Z Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear 440 complexity. arXiv preprint arXiv:2006.04768 , 2020. 441

Yushi Bai, Xin Lv, Jiajie Zhang, Hongchang Lyu, Jiankai Tang, Zhidian Huang, Zhengxiao Du, Xiao 442 Liu, Aohan Zeng, Lei Hou, Yuxiao Dong, Jie Tang, and Juanzi Li. Longbench: A bilingual, multitask 443 benchmark for long context understanding. arXiv preprint arXiv:2308.14508 , 2023. 444

445

446

447

448

449

450

451

A Appendix / supplemental material

A.1 Algorithm Details

The detail implementation of FlashMask Backward Pass is presented in Algorithm 2. We do precomputations of max and min values of FMS and FME similar to the Forward Pass. Then the FMS 𝑗 and FME 𝑗 can be loaded to SRAM outside the inner loop (line 14-15), reducing the HBM accesses to 2 × 𝑁 . Then, we do inner loop on 𝑄 𝑖 (line 16), computing the two valid parts and bypassing the masked part 𝑖 ∈ ( j maxFMS 𝑗 𝐵 𝑟 k , j minFME 𝑗 𝐵 𝑟 k ) .

Algorithm 2 Optimized Backward Pass with FlashMask

A.2 Supplementary Experimental Details 452

All end-to-end training and testing in this paper were conducted on 4 servers, each equipped with 32 453 NVIDIA A800-SXM 80G GPUs. We comprehensively evaluated the performance of the LLaMA2 454

model across three different parameter scales, four downstream task scenarios, and various sequence 455 lengths. Given the diversity of experimental combinations and the specific distributed parallel 456 strategies required by models, in varying parameter scales, the primary goal of the experiments is 457 not to achieve optimal end-to-end training performance but to demonstrate the effectiveness of the 458 FlashMask method. Therefore, to ensure consistency, we set the following hyperparameters in Table 1 459 with the same hardware configuration. 460

To verify the representativeness of our synthetic dataset, sparsity distribution histograms of synthetic 461 dataset are presented in Figure 6. Then we use InToken method with max sequence length of 16K, 462 32K, 64K, and 128K on the open-source dataset LongBench, and compute the distribution histograms, 463 presented in Figure 7. Note that many long sentences are truncated for max sequence length 16K, 464 and 32K. Results indicate that the sparsity distributions of LongBench dataset and synthetic dataset 465 are similar. 466

0

20

40

60

80

100

Sparsity(%)

0

20

40

60

80

100

Count

SFT 2K

0

20

40

60

80

100

Sparsity(%)

0

20

40

60

80

Count

SFT 4K

0

20

40

60

80

100

Sparsity(%)

0

10

20

30

40

50

60

70

Count

SFT 8K

0

20

40

60

80

100

Sparsity(%)

0

10

20

30

40

50

60

Count

SFT 16K

0

20

40

60

80

100

Sparsity(%)

0

10

20

30

40

50

60

70

Count

SFT 32K

0

20

40

60

80

100

Sparsity(%)

0

10

20

30

40

50

60

70

Count

SFT 64K

0

20

40

60

80

100

Sparsity(%)

0

10

20

30

40

50

60

70

Count

SFT 128K

0

20

40

60

80

100

Sparsity(%)

0

20

40

60

80

100

Count

DPO 2K

0

20

40

60

80

100

Sparsity(%)

0

20

40

60

80

Count

DPO 4K

0

20

40

60

80

100

Sparsity(%)

0

10

20

30

40

50

60

70

80

Count

DPO 8K

0

20

40

60

80

100

Sparsity(%)

0

10

20

30

40

50

60

Count

DPO 16K

0

20

40

60

80

100

Sparsity(%)

0

10

20

30

40

50

60

70

80

Count

DPO 32K

0

20

40

60

80

100

Sparsity(%)

0

10

20

30

40

50

60

70

Count

DPO 64K

0

20

40

60

80

100

Sparsity(%)

0

10

20

30

40

50

60

Count

DPO 128K

0

20

40

60

80

100

Sparsity(%)

0

10

20

30

40

50

60

Count

RM 2K

0

20

40

60

80

100

Sparsity(%)

0

10

20

30

40

50

60

Count

RM 4K

0

20

40

60

80

100

Sparsity(%)

0

10

20

30

40

50

60

Count

RM 8K

0

20

40

60

80

100

Sparsity(%)

0

20

40

60

80

100

Count

RM 16K

0

20

40

60

80

100

Sparsity(%)

0

20

40

60

80

100

Count

RM 32K

0

20

40

60

80

100

Sparsity(%)

0

20

40

60

80

Count

RM 64K

0

20

40

60

80

100

Sparsity(%)

0

20

40

60

80

Count

RM 128K

0

20

40

60

80

100

Sparsity(%)

0

200

400

600

800

1000

1200

1400

1600

Count

SFT 16K

0

20

40

60

80

100

Sparsity(%)

0

100

200

300

400

500

600

Count

SFT 32K

0

20

40

60

80

100

Sparsity(%)

0

50

100

150

200

250

300

350

400

Count

SFT 64K

0

20

40

60

80

100

Sparsity(%)

0

50

100

150

200

250

300

350

Count

SFT 128K

A.3 Full Experiment Results

467

Kernel experiments are also conducted on batch sizes 4, and 8. FA-Varlen is excluded by default. 468 Results are presented in Figure 8 and 9. The trends are identical to Figure 3 in Section 4.3, except 469 memory exhaustion occurred with less sequence length, especially for FA-DenseMask and Vanilla 470 Attention which require 𝑂 ( 𝑁 2 ) memory to launch. 471

2

4

8

16

32

64

128

Sequence Length(K)

10 1

10 2

10 3

10 4

Latency(ms)

SFT

FA-DenseMask

VanillaAttention

FlashMask

2

4

8

16

32

64

128

Sequence Length(K)

10 1

10 2

10 3

10 4

Latency(ms)

DPO

FA-DenseMask

VanillaAttention

FlashMask

2

4

8

16

32

64

128

Sequence Length(K)

10 1

10 2

10 3

10 4

Latency(ms)

RM

FA-DenseMask

VanillaAttention

FlashMask

2

4

8

16

32

64

128

Sequence Length(K)

10 1

10 2

10 3

Latency(ms)

SFT

FA-DenseMask

VanillaAttention

FlashMask

2

4

8

16

32

64

128

Sequence Length(K)

10 1

10 2

10 3

Latency(ms)

DPO

FA-DenseMask

VanillaAttention

FlashMask

2

4

8

16

32

64

128

Sequence Length(K)

10 1

10 2

10 3

Latency(ms)

RM

FA-DenseMask

VanillaAttention

FlashMask

Weevaluate the effectiveness of FlashMask on the open-source dataset LongBench. The throughput of 472

LoRA fine-tuning for LLaMA2-7B are shown in Figure 10. FlashMask performed close to FA-Varlen, 473

showcasing 4.12x faster than FA-DenseMask, proving that FlashMask can deliver significant training 474

accelerations in generalized real-world scenarios. 475

16

32

64

Sequence Length(K)

2000

3000

4000

5000

6000

7000

Tokens/Sec/GPU

Speed Up 4.12x

LLaMA-7B

FA-DenseMask

FA-Varlen

FlashMask

Figure 11 presents the GPU memory consumption in End-to-End training. FlashMask showed linear 476

memory consumption with increasing sequence length, far less than FA-DenseMask. Therefore, 477

FlashMask supports training with much longer sequences in memory limits of 80G. 478

4

8

16

32

64

128

256

Sequence Length(K)

0

10

20

30

40

50

60

70

GPU Memory(GB)

LLaMA-7B

VanillaAttention

FA-DenseMask

FA-Varlen

FlashMask

4

8

16

32

64

128

196

Sequence Length(K)

0

10

20

30

40

50

60

70

GPU Memory(GB)

LLaMA-13B

VanillaAttention

FA-DenseMask

FA-Varlen

FlashMask

4

8

16

32

64

96

Sequence Length(K)

0

10

20

30

40

50

60

70

GPU Memory(GB)

LLaMA-70B

VanillaAttention

FA-DenseMask

FA-Varlen

FlashMask

(a) SFT

4

8

16

32

64

128

256

512

544

Sequence Length(K)

0

10

20

30

40

50

60

70

GPU Memory(GB)

LLaMA-7B

VanillaAttention

FA-DenseMask

FA-Varlen

FlashMask

4

8

16

32

64

128

224

Sequence Length(K)

0

10

20

30

40

50

60

70

GPU Memory(GB)

LLaMA-13B

VanillaAttention

FA-DenseMask

FA-Varlen

FlashMask

4

8

16

32

64

128

Sequence Length(K)

0

10

20

30

40

50

60

GPU Memory(GB)

LLaMA-70B

VanillaAttention

FA-DenseMask

FA-Varlen

FlashMask

(b) LoRA

4

8

16

32

64

96

Sequence Length(K)

0

10

20

30

40

50

60

70

GPU Memory(GB)

LLaMA-7B

VanillaAttention

FA-DenseMask

FlashMask

4

8

16

32

64

128

180

Sequence Length(K)

0

10

20

30

40

50

60

70

GPU Memory(GB)

LLaMA-13B

VanillaAttention

FA-DenseMask

FlashMask

4

8

16

32

64

80

Sequence Length(K)

0

10

20

30

40

50

60

70

GPU Memory(GB)

LLaMA-70B

VanillaAttention

FA-DenseMask

FlashMask

(c) DPO

4

8

16

32

64

128

256

Sequence Length(K)

0

10

20

30

40

50

60

70

GPU Memory(GB)

LLaMA-7B

VanillaAttention

FA-DenseMask

FlashMask

4

8

16

32

64

128

196

Sequence Length(K)

0

10

20

30

40

50

60

70

GPU Memory(GB)

LLaMA-13B

VanillaAttention

FA-DenseMask

FlashMask

4

8

16

32

64

96

Sequence Length(K)

0

10

20

30

40

50

60

70

GPU Memory(GB)

LLaMA-70B

VanillaAttention

FA-DenseMask

FlashMask

(d) RM

479

480

481

482

483

484

485

486

487

488

489

490

491

492

493

494

495

496

497

498

499

500

501

502

503

504

505

506

507

508

509

510

511

512

513

514

515

516

517

518

519

520

521

522

523

524

525

526

527

528

529

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: This paper's contributions and scope are described in Abstract and Introduction.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: The paper includes a discussion section about limitations.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: In sec 3.3 Complexity Analysis

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: The source code will be public available and can reproduce the results according README.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

585

586

587

588

589

590

591

592

593

594

595

596

597

598

599

600

601

602

603

604

605

606

607

608

609

610

611

612

613

614

615

616

617

618

619

620

621

622

623

624

625

626

627

628

629

630

631

632

633

634

635

636

Answer: [Yes]

Justification: The source code will be public available.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: Refer to Experiments section.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: All our experimental results are run multiple times and then averaged.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: Referring to the Experiments section, we provide a running environment.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: We follow the NeurIPS Code of Ethics properly.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [NA]

Justification: There is no societal impact of the work performed.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to

generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: The paper poses no such risks.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: CC BY-NC-ND 4.0

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

741

742

743

744

745

746

747

748

749

750

751

752

753

754

755

756

757

758

759

760

761

762

763

764

765

766

767

768

769

770

771

772

773

774

775

776

777

778

779

780

781

782

783

784

785

Answer: [NA]

Justification: The paper does not release new assets.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: The paper does not involve crowdsourcing nor research with human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: The paper does not involve crowdsourcing nor research with human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes a novel method to address the high computational and memory complexity of current large-scale transformers. By adopting a simple yet effective column-wise sparse representation of attention masks, the algorithm achieves reduced memory and computational complexity while maintaining the accuracy of attention computation.\n\nSTRENGTHS:\n1. This paper investigates a topic of interest, given the current trend toward increasing context lengths in LLMs.\n2. The method proposed in this paper is straightforward and easy to implement.\n3. The paper is well-written and clearly presented.\n\nWEAKNESSES:\n1. It is crucial to highlight the advantages of this method over related work to help readers fully understand its significance. However, in the subsection \"Attention Optimization Techniques,\" the authors only mention the drawback of FlashAttention and discuss its relationship to their work. The introduction of other related works is confusing and makes it difficult to comprehend their relevance to this paper. The overall conclusion, \"*Both of the previously discussed solutions either compromise precision or yield only marginal enhancements in efficiency. Conversely, our proposed FlashMask is capable of delivering exact computations.*\" is general and non-specific. It is unclear which methods compromise precision and which ones only offer marginal improvements.\n\n2. In the experiments, the baseline algorithms are limited to Vanilla Attention and FlashAttention. Are there more efficient Transformer algorithms that could be used for comparison? If not, the authors should explain the rationale behind the selection of these specific baselines.\n\n3. As a non-expert in this field, I found the writing of this paper confusing. For instance, the initialism \"HBN\" is introduced without any explanation or context.\n\nQUESTIONS:\n1. Has there been any prior work on efficient attention mask computation?\n2. Why was the FA-Varlen method not included in the DPO and RM scenarios? Could the authors provide an explanation in their paper?\n3. In Figure 4, the latency of both FA-Window and FlashMask is almost identical. If the authors aim to demonstrate the efficiency of FlashMask, could they explain the experimental results more clearly?\n4. In Figures 3 and 5, it appears that the performance and efficiency of FA-Varlen are comparable to FlashMask in the SFT setting. Could the authors clarify this comparison?\n\nLIMITATIONS:\nnone", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-11"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes FlashMask, which accelerates the masked attention mechanism that can reduce the original attention from O(N^2) to O(N) and simultaneously reduces the memory cost. Experimental results show that the proposed FlashMask significantly reduces training time without accuracy degradation.\n\nSTRENGTHS:\n+ This paper provides a comprehensive study and analysis about the sparse attention, in terms of their efficiency. Also, this paper includes existing attention optimization like FlashAttention, explaining the motivation of the proposed FlashMask, which lies in the lack of optimization for sparse attention.\n\n+ This paper proposes an optimization for column-based sparse attention, which significantly improves memory efficiency and reduces computational costs.\n\n+ This paper provides a comprehensive complexity analysis, evaluation, and comparison with existing methods. It seems the authors make a lot of efforts on the proposed approach.\n\nWEAKNESSES:\n- Even though FlashMask achieves significant improvement in the memory efficiency of sparse attention, the key idea is similar to FlashAttention, but it is just for sparse attention mechanisms. Based on this fact, the novelty of this paper is not strong. I recommend the authors explain why the red part in the algorithm is designed and why it is unique for sparse attention.\n\n- The authors only present optimization for column-based sparse attention. The performance for other types of sparse attention is unknown. If the proposed approach can be applied to all sparse attention, the contribution of this paper is extremely great. However, the existing version is not comprehensive.\n\n- Based on the experiments, the practical latency is not significantly reduced as compared to other methods, even though the theoretical complexity is from N^2 to N. Besides, the authors do not provide results for accuracy.\n\nQUESTIONS:\nSee Weaknesses.\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-11"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper introduces FlashMask, an innovative algorithm designed to address the computational and memory challenges associated with conventional attention mechanisms in large-scale Transformers. FlashMask employs a column-wise sparse representation for attention masks, significantly reducing the computational complexity from quadratic to linear with respect to sequence length. The authors demonstrate FlashMask's effectiveness across various masking scenarios and training modalities, including Supervised Fine-Tuning (SFT), Direct Preference Optimization (DPO), and Reward Model (RM).\n\nSTRENGTHS:\nThis paper presents a novel solution to a well-known problem in the field of natural language processing, offering a practical method to reduce the computational burden of attention mechanisms in Transformers.\n\nThe paper provides extensive empirical evidence to support its claims, including comparisons with state-of-the-art techniques like FlashAttention, demonstrating FlashMask's superiority in terms of speed and efficiency.\n\nFlashMask's performance across different masking scenarios and training modalities shows its versatility and robustness, indicating its potential applicability to a wide range of models and tasks.\n\nPractical Impact: The paper not only presents theoretical advancements but also demonstrates practical benefits, such as enabling the\n\nWEAKNESSES:\nThe scaling ability of the proposed method deserves further verified on large scale datasets.\n\nWhile the paper demonstrates FlashMask's effectiveness in specific scenarios, it may lack broader evidence on how it performs across different types of NLP tasks or diverse datasets.\n\nThe paper could provide more detailed insights into how FlashMask handles different sparsity levels and the impact on various model sizes and complexities.\n\nQUESTIONS:\nSee the weakness\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-11"}

comments: SUMMARY:
This paper proposes FlashMask, a modification of FlashAttention with fixed masks. The paper shows speedup of FlashAttention when using sparse masks in the attention matrix.

STRENGTHS:
FlashAttention is an important algorithm, and sparsity in the attention matrix is an important feature. Further study of these aspects is helpful.

WEAKNESSES:
The paper seems to make claims that are unsubstantiated by experiments. In the abstract and introduction, the paper claims speedup without sacrificing model quality. However, there is no experiment evaluating model quality in the experiments. This is a critical flaw.

Further, the contribution of the paper is unclear. Block-sparsity is already supported in FlashAttention (see section 3.3 of FlashAttention). It is unclear how this paper is different. There are also more recent works such as "Fast Attention Over Long Sequences With Dynamic Sparse Flash Attention" (NeurIPS 2023), which seem to be strictly more expressive in features than this paper.

QUESTIONS:
See weaknesses

LIMITATIONS:
The paper discusses superlinear scaling in sequence length as a limitation, but is lacking in discussion of model quality.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 2
DATE: 2024-07-14