A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 11:00 UTC
Implements dual-chunk-flash-attn backend for dual chunk attention with sparse attention support (#11844)
Tao He committed
May 13, 2025 at 10:52 UTC
60f76243344d2d3deca5e5ecdade547acc7fed50
17 files changed
+2444
-32
CMakeLists.txt
+1
@@ -230,6 +230,7 @@ set(VLLM_EXT_SRC
230
"csrc/attention/paged_attention_v1.cu"
231
"csrc/attention/paged_attention_v2.cu"
232
"csrc/attention/merge_attn_states.cu"
233
+ "csrc/attention/vertical_slash_index.cu"
234
"csrc/pos_encoding_kernels.cu"
235
"csrc/activation_kernels.cu"
236
"csrc/layernorm_kernels.cu"
csrc/attention/vertical_slash_index.cu
new
+401
@@ -0,0 +1,401 @@
1
+// Copyright (c) Microsoft Corporation.
2
+// Licensed under the MIT license.
3
+
4
+#include <assert.h>
5
+
6
+#include <cuda.h>
7
+
8
+#include <torch/all.h>
9
+
10
+__device__ int64_t save_blocks(int* block_offset, int64_t range_start,
11
+ int64_t range_end, int64_t block_size,
12
+ int64_t input_block_count, int64_t kv_seqlen) {
13
+ if (range_start >= kv_seqlen) {
14
+ return input_block_count;
15
+ }
16
+ if (range_end > kv_seqlen) {
17
+ range_end = kv_seqlen;
18
+ }
19
+ int64_t current_block_count = input_block_count;
20
+ for (int idx = range_start; idx < range_end; idx += block_size) {
21
+ block_offset[current_block_count++] = idx;
22
+ }
23
+ return current_block_count;
24
+}
25
+
26
+__global__ void convert_vertical_slash_indexes_kernel(
27
+ const int* q_seqlens, // [BATCH, ]
28
+ const int* kv_seqlens, // [BATCH, ]
29
+ const int* vertical_indexes, // [BATCH, N_HEADS, NNZ_V]
30
+ const int* slash_indexes, // [BATCH, N_HEADS, NNZ_S]
31
+ int* block_count, // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M)]
32
+ int* block_offset, // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M), NNZ_S]
33
+ int* column_count, // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M)]
34
+ int* column_index, // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M), NNZ_V]
35
+ int64_t N_HEADS, int64_t N_ROWS, int64_t BLOCK_SIZE_M, int64_t BLOCK_SIZE_N,
36
+ int64_t NNZ_V, int64_t NNZ_S,
37
+ bool causal // True for intra, False for succ
38
+) {
39
+ const int batch_idx = blockIdx.y;
40
+ const int head_idx = blockIdx.x;
41
+ const int group_idx = blockIdx.z;
42
+
43
+ int64_t q_seqlen = q_seqlens[batch_idx];
44
+ int64_t kv_seqlen = kv_seqlens[batch_idx];
45
+ int64_t block_idx_m = group_idx * blockDim.x + threadIdx.x;
46
+ int64_t start_m = block_idx_m * BLOCK_SIZE_M;
47
+ if (start_m >= q_seqlen) {
48
+ return;
49
+ }
50
+ int64_t end_m = start_m + BLOCK_SIZE_M;
51
+ vertical_indexes += (batch_idx * N_HEADS + head_idx) * NNZ_V;
52
+ slash_indexes += (batch_idx * N_HEADS + head_idx) * NNZ_S;
53
+ int64_t row_offset = (batch_idx * N_HEADS + head_idx) * N_ROWS + block_idx_m;
54
+ block_count += row_offset;
55
+ block_offset += row_offset * NNZ_S;
56
+ column_count += row_offset;
57
+ column_index += row_offset * NNZ_V;
58
+
59
+ bool has_slash = true;
60
+ int64_t tmp_col_cnt = 0, tmp_blk_cnt = 0;
61
+ int64_t s = 0, v = 0;
62
+ int64_t v_idx = vertical_indexes[v++];
63
+ int64_t s_idx = slash_indexes[s++];
64
+ if (causal) {
65
+ while (s_idx >= end_m + (kv_seqlen - q_seqlen) && s < NNZ_S) {
66
+ s_idx = slash_indexes[s++];
67
+ }
68
+ if (s_idx > end_m + (kv_seqlen - q_seqlen)) has_slash = false;
69
+ s_idx = max((kv_seqlen - q_seqlen) + end_m - s_idx, BLOCK_SIZE_M);
70
+ } else {
71
+ while (s_idx >= end_m + kv_seqlen && s < NNZ_S) {
72
+ s_idx = slash_indexes[s++];
73
+ }
74
+ if (s_idx > end_m + kv_seqlen) has_slash = false;
75
+ s_idx = max(kv_seqlen + end_m - s_idx, BLOCK_SIZE_M);
76
+ }
77
+
78
+ int64_t range_start = s_idx - BLOCK_SIZE_M, range_end = s_idx;
79
+ if (!has_slash) {
80
+ if (causal) {
81
+ range_start = (kv_seqlen - q_seqlen) + end_m;
82
+ range_end = (kv_seqlen - q_seqlen) + end_m + BLOCK_SIZE_N;
83
+ } else {
84
+ range_start = kv_seqlen;
85
+ range_end = kv_seqlen + BLOCK_SIZE_N;
86
+ }
87
+ }
88
+
89
+ bool slash_finished = false;
90
+ while (1) {
91
+ if (v_idx < range_end) {
92
+ if (v_idx < range_start) {
93
+ column_index[tmp_col_cnt++] = v_idx;
94
+ }
95
+ if (v < NNZ_V) {
96
+ v_idx = vertical_indexes[v++];
97
+ } else {
98
+ if (causal)
99
+ v_idx = end_m + BLOCK_SIZE_N + (kv_seqlen - q_seqlen);
100
+ else
101
+ v_idx = end_m + BLOCK_SIZE_N + kv_seqlen;
102
+ }
103
+ } else {
104
+ if ((s < NNZ_S && causal) ||
105
+ (s < NNZ_S && !causal && slash_indexes[s] >= start_m)) {
106
+ if (causal)
107
+ s_idx = max((kv_seqlen - q_seqlen) + end_m - slash_indexes[s++],
108
+ BLOCK_SIZE_M);
109
+ else
110
+ s_idx = max(kv_seqlen + end_m - slash_indexes[s++], BLOCK_SIZE_M);
111
+ } else {
112
+ if (v == NNZ_V || (v_idx > range_start && causal)) {
113
+ // add the last vertical if no more slash
114
+ if (v == NNZ_V && !causal && v_idx < kv_seqlen) {
115
+ column_index[tmp_col_cnt++] = v_idx;
116
+ }
117
+ tmp_blk_cnt = save_blocks(block_offset, range_start, range_end,
118
+ BLOCK_SIZE_N, tmp_blk_cnt, kv_seqlen);
119
+ break;
120
+ } else {
121
+ if (causal) {
122
+ range_start = (kv_seqlen - q_seqlen) + end_m;
123
+ range_end = (kv_seqlen - q_seqlen) + end_m + BLOCK_SIZE_N;
124
+ } else {
125
+ // if slash_finished but there are vertical left, save current
126
+ // blocks
127
+ tmp_blk_cnt = save_blocks(block_offset, range_start, range_end,
128
+ BLOCK_SIZE_N, tmp_blk_cnt, kv_seqlen);
129
+ range_start = kv_seqlen;
130
+ range_end = kv_seqlen + BLOCK_SIZE_N;
131
+ }
132
+ slash_finished = true;
133
+ }
134
+ }
135
+ if (!slash_finished) {
136
+ if (s_idx > range_end + BLOCK_SIZE_M) {
137
+ tmp_blk_cnt = save_blocks(block_offset, range_start, range_end,
138
+ BLOCK_SIZE_N, tmp_blk_cnt, kv_seqlen);
139
+ range_start = s_idx - BLOCK_SIZE_M;
140
+ range_end = s_idx;
141
+ } else if (s_idx > range_end) {
142
+ range_end += BLOCK_SIZE_M;
143
+ }
144
+ }
145
+ }
146
+ }
147
+
148
+ block_count[0] = tmp_blk_cnt;
149
+ column_count[0] = tmp_col_cnt;
150
+}
151
+
152
+void convert_vertical_slash_indexes_64x64(
153
+ const int* q_seqlens, // [BATCH, ]
154
+ const int* kv_seqlens, // [BATCH, ]
155
+ const int* vertical_indexes, // [BATCH, N_HEADS, NNZ_V]
156
+ const int* slash_indexes, // [BATCH, N_HEADS, NNZ_S]
157
+ int* block_count, // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M)]
158
+ int* block_offset, // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M), NNZ_S]
159
+ int* column_count, // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M)]
160
+ int* column_index, // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M), NNZ_V]
161
+ int64_t BATCH_SIZE, int64_t N_HEADS, int64_t N_ROWS, int64_t BLOCK_SIZE_M,
162
+ int64_t BLOCK_SIZE_N, int64_t NNZ_V, int64_t NNZ_S, bool causal) {
163
+ const int N_THREADS = 64;
164
+ const dim3 dimBlock(N_THREADS);
165
+ const dim3 dimGrid(N_HEADS, BATCH_SIZE, (N_ROWS + N_THREADS - 1) / N_THREADS);
166
+ convert_vertical_slash_indexes_kernel<<<dimGrid, dimBlock>>>(
167
+ q_seqlens, kv_seqlens, vertical_indexes, slash_indexes, block_count,
168
+ block_offset, column_count, column_index, N_HEADS, N_ROWS, BLOCK_SIZE_M,
169
+ BLOCK_SIZE_N, NNZ_V, NNZ_S, causal);
170
+}
171
+
172
+/**
173
+ * Implements the Algorithm 4 in paper https://arxiv.org/abs/2407.02490.
174
+ *
175
+ * This function builds the index of each row of blocks from vertical indices
176
+ * and slash indices. The vertical indices are treated as points, while the
177
+ * slash indices are converted as ranges. The output consists of the merged
178
+ * ranges and separate column indices, where the ranges are represented by
179
+ * block indices.
180
+ *
181
+ * The implementation is referenced from the original MInference repo:
182
+ * https://github.com/microsoft/MInference/blob/main/csrc/vertical_slash_index.cu.
183
+ */
184
+void convert_vertical_slash_indexes(
185
+ torch::Tensor& block_count, // [BATCH, N_HEADS, NUM_ROWS]
186
+ torch::Tensor& block_offset, // [BATCH, N_HEADS, NUM_ROWS, NNZ_S]
187
+ torch::Tensor& column_count, // [BATCH, N_HEADS, NUM_ROWS]
188
+ torch::Tensor& column_index, // [BATCH, N_HEADS, NUM_ROWS, NNZ_V]
189
+ torch::Tensor q_seqlens, // [BATCH, ]
190
+ torch::Tensor kv_seqlens, // [BATCH, ]
191
+ torch::Tensor vertical_indexes, // [BATCH, N_HEADS, NNZ_V]
192
+ torch::Tensor slash_indexes, // [BATCH, N_HEADS, NNZ_S]
193
+ int64_t context_size, int64_t block_size_M, int64_t block_size_N,
194
+ bool causal) {
195
+ cudaSetDevice(q_seqlens.get_device());
196
+
197
+ int batch_size = slash_indexes.size(0);
198
+ int num_heads = slash_indexes.size(1);
199
+ int nnz_slash = slash_indexes.size(2);
200
+ int nnz_vertical = vertical_indexes.size(2);
201
+ int num_rows = (context_size + block_size_M - 1) / block_size_M;
202
+
203
+ convert_vertical_slash_indexes_64x64(
204
+ q_seqlens.data_ptr<int>(), kv_seqlens.data_ptr<int>(),
205
+ vertical_indexes.data_ptr<int>(), slash_indexes.data_ptr<int>(),
206
+ block_count.data_ptr<int>(), block_offset.data_ptr<int>(),
207
+ column_count.data_ptr<int>(), column_index.data_ptr<int>(), batch_size,
208
+ num_heads, num_rows, block_size_M, block_size_N, nnz_vertical, nnz_slash,
209
+ causal);
210
+}
211
+
212
+__global__ void convert_vertical_slash_indexes_kernel_mergehead(
213
+ const int* q_seqlens, // [BATCH, ]
214
+ const int* kv_seqlens, // [BATCH, ]
215
+ const int* vertical_indexes, // [BATCH, N_HEADS, NNZ_V]
216
+ const int* slash_indexes, // [BATCH, N_HEADS, NNZ_S]
217
+ const int* per_head_vertical_topkv, const int* per_head_slash_topkv,
218
+ int* block_count, // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M)]
219
+ int* block_offset, // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M), NNZ_S]
220
+ int* column_count, // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M)]
221
+ int* column_index, // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M), NNZ_V]
222
+ int64_t N_HEADS, int64_t N_ROWS, int64_t BLOCK_SIZE_M, int64_t BLOCK_SIZE_N,
223
+ int64_t NNZ_V, int64_t NNZ_S,
224
+ bool causal // True for intra, False for succ
225
+) {
226
+ const int batch_idx = blockIdx.y;
227
+ const int head_idx = blockIdx.x;
228
+ const int group_idx = blockIdx.z;
229
+
230
+ int64_t q_seqlen = q_seqlens[batch_idx];
231
+ int64_t kv_seqlen = kv_seqlens[batch_idx];
232
+ int64_t block_idx_m = group_idx * blockDim.x + threadIdx.x;
233
+ int64_t start_m = block_idx_m * BLOCK_SIZE_M;
234
+ if (start_m >= q_seqlen) {
235
+ return;
236
+ }
237
+ int64_t end_m = start_m + BLOCK_SIZE_M;
238
+ vertical_indexes += (batch_idx * N_HEADS + head_idx) * NNZ_V;
239
+ slash_indexes += (batch_idx * N_HEADS + head_idx) * NNZ_S;
240
+ int64_t row_offset = (batch_idx * N_HEADS + head_idx) * N_ROWS + block_idx_m;
241
+ block_count += row_offset;
242
+ block_offset += row_offset * NNZ_S;
243
+ column_count += row_offset;
244
+ column_index += row_offset * NNZ_V;
245
+
246
+ // MergeHead: each head has it's unique max topk NNZ_V，NNZ_S. (NNZ_V，NNZ_S
247
+ // above is buffer size, use to compute offset)
248
+ NNZ_S = per_head_slash_topkv[head_idx];
249
+ NNZ_V = per_head_vertical_topkv[head_idx];
250
+
251
+ bool has_slash = true;
252
+ int64_t tmp_col_cnt = 0, tmp_blk_cnt = 0;
253
+ int64_t s = 0, v = 0;
254
+ int64_t v_idx = vertical_indexes[v++];
255
+ int64_t s_idx = slash_indexes[s++];
256
+ if (causal) {
257
+ while (s_idx >= end_m + (kv_seqlen - q_seqlen) && s < NNZ_S) {
258
+ s_idx = slash_indexes[s++];
259
+ }
260
+ if (s_idx > end_m + (kv_seqlen - q_seqlen)) has_slash = false;
261
+ s_idx = max((kv_seqlen - q_seqlen) + end_m - s_idx, BLOCK_SIZE_M);
262
+ } else {
263
+ while (s_idx >= end_m + kv_seqlen && s < NNZ_S) {
264
+ s_idx = slash_indexes[s++];
265
+ }
266
+ if (s_idx > end_m + kv_seqlen) has_slash = false;
267
+ s_idx = max(kv_seqlen + end_m - s_idx, BLOCK_SIZE_M);
268
+ }
269
+
270
+ int64_t range_start = s_idx - BLOCK_SIZE_M, range_end = s_idx;
271
+ if (!has_slash) {
272
+ if (causal) {
273
+ range_start = (kv_seqlen - q_seqlen) + end_m;
274
+ range_end = (kv_seqlen - q_seqlen) + end_m + BLOCK_SIZE_N;
275
+ } else {
276
+ range_start = kv_seqlen;
277
+ range_end = kv_seqlen + BLOCK_SIZE_N;
278
+ }
279
+ }
280
+
281
+ bool slash_finished = false;
282
+ while (1) {
283
+ if (v_idx < range_end) {
284
+ if (v_idx < range_start) {
285
+ column_index[tmp_col_cnt++] = v_idx;
286
+ }
287
+ if (v < NNZ_V) {
288
+ v_idx = vertical_indexes[v++];
289
+ } else {
290
+ if (causal)
291
+ v_idx = end_m + BLOCK_SIZE_N + (kv_seqlen - q_seqlen);
292
+ else
293
+ v_idx = end_m + BLOCK_SIZE_N + kv_seqlen;
294
+ }
295
+ } else {
296
+ if ((s < NNZ_S && causal) ||
297
+ (s < NNZ_S && !causal && slash_indexes[s] >= start_m)) {
298
+ if (causal)
299
+ s_idx = max((kv_seqlen - q_seqlen) + end_m - slash_indexes[s++],
300
+ BLOCK_SIZE_M);
301
+ else
302
+ s_idx = max(kv_seqlen + end_m - slash_indexes[s++], BLOCK_SIZE_M);
303
+ } else {
304
+ if (v == NNZ_V || (v_idx > range_start && causal)) {
305
+ // add the last vertical if no more slash
306
+ if (v == NNZ_V && !causal && v_idx < kv_seqlen) {
307
+ column_index[tmp_col_cnt++] = v_idx;
308
+ }
309
+ tmp_blk_cnt = save_blocks(block_offset, range_start, range_end,
310
+ BLOCK_SIZE_N, tmp_blk_cnt, kv_seqlen);
311
+ break;
312
+ } else {
313
+ if (causal) {
314
+ range_start = (kv_seqlen - q_seqlen) + end_m;
315
+ range_end = (kv_seqlen - q_seqlen) + end_m + BLOCK_SIZE_N;
316
+ } else {
317
+ // if slash_finished but there are vertical left, save current
318
+ // blocks
319
+ tmp_blk_cnt = save_blocks(block_offset, range_start, range_end,
320
+ BLOCK_SIZE_N, tmp_blk_cnt, kv_seqlen);
321
+ range_start = kv_seqlen;
322
+ range_end = kv_seqlen + BLOCK_SIZE_N;
323
+ }
324
+ slash_finished = true;
325
+ }
326
+ }
327
+ if (!slash_finished) {
328
+ if (s_idx > range_end + BLOCK_SIZE_M) {
329
+ tmp_blk_cnt = save_blocks(block_offset, range_start, range_end,
330
+ BLOCK_SIZE_N, tmp_blk_cnt, kv_seqlen);
331
+ range_start = s_idx - BLOCK_SIZE_M;
332
+ range_end = s_idx;
333
+ } else if (s_idx > range_end) {
334
+ range_end += BLOCK_SIZE_M;
335
+ }
336
+ }
337
+ }
338
+ }
339
+
340
+ block_count[0] = tmp_blk_cnt;
341
+ column_count[0] = tmp_col_cnt;
342
+}
343
+
344
+void convert_vertical_slash_indexes_64x64_mergehead(
345
+ const int* q_seqlens, // [BATCH, ]
346
+ const int* kv_seqlens, // [BATCH, ]
347
+ const int* vertical_indexes, // [BATCH, N_HEADS, NNZ_V]
348
+ const int* slash_indexes, // [BATCH, N_HEADS, NNZ_S]
349
+ int* per_head_vertical_topkv, int* per_head_slash_topkv,
350
+ int* block_count, // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M)]
351
+ int* block_offset, // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M), NNZ_S]
352
+ int* column_count, // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M)]
353
+ int* column_index, // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M), NNZ_V]
354
+ int64_t BATCH_SIZE, int64_t N_HEADS, int64_t N_ROWS, int64_t BLOCK_SIZE_M,
355
+ int64_t BLOCK_SIZE_N, int64_t NNZ_V, int64_t NNZ_S, bool causal) {
356
+ const int N_THREADS = 64;
357
+ const dim3 dimBlock(N_THREADS);
358
+ const dim3 dimGrid(N_HEADS, BATCH_SIZE, (N_ROWS + N_THREADS - 1) / N_THREADS);
359
+ convert_vertical_slash_indexes_kernel_mergehead<<<dimGrid, dimBlock>>>(
360
+ q_seqlens, kv_seqlens, vertical_indexes, slash_indexes,
361
+ per_head_vertical_topkv, per_head_slash_topkv, block_count, block_offset,
362
+ column_count, column_index, N_HEADS, N_ROWS, BLOCK_SIZE_M, BLOCK_SIZE_N,
363
+ NNZ_V, NNZ_S, causal);
364
+}
365
+
366
+/**
367
+ * Implements the Algorithm 4 in paper https://arxiv.org/abs/2407.02490.
368
+ *
369
+ * Like the above convert_vertical_slash_indexes, but with
370
+ * pre-computed vertical and slash counts.
371
+ */
372
+void convert_vertical_slash_indexes_mergehead(
373
+ torch::Tensor& block_count, // [BATCH, N_HEADS, NUM_ROWS]
374
+ torch::Tensor& block_offset, // [BATCH, N_HEADS, NUM_ROWS, NNZ_S]
375
+ torch::Tensor& column_count, // [BATCH, N_HEADS, NUM_ROWS]
376
+ torch::Tensor& column_index, // [BATCH, N_HEADS, NUM_ROWS, NNZ_V]
377
+ torch::Tensor q_seqlens, // [BATCH, ]
378
+ torch::Tensor kv_seqlens, // [BATCH, ]
379
+ torch::Tensor vertical_indexes, // [BATCH, N_HEADS, NNZ_V]
380
+ torch::Tensor slash_indexes, // [BATCH, N_HEADS, NNZ_S]
381
+ torch::Tensor vertical_indices_count, // [N_HEADS, ]
382
+ torch::Tensor slash_indices_count, // [N_HEADS, ]
383
+ int64_t context_size, int64_t block_size_M, int64_t block_size_N,
384
+ bool causal) {
385
+ cudaSetDevice(q_seqlens.get_device());
386
+
387
+ int batch_size = slash_indexes.size(0);
388
+ int num_heads = slash_indexes.size(1);
389
+ int nnz_slash = slash_indexes.size(2);
390
+ int nnz_vertical = vertical_indexes.size(2);
391
+ int num_rows = (context_size + block_size_M - 1) / block_size_M;
392
+
393
+ convert_vertical_slash_indexes_64x64_mergehead(
394
+ q_seqlens.data_ptr<int>(), kv_seqlens.data_ptr<int>(),
395
+ vertical_indexes.data_ptr<int>(), slash_indexes.data_ptr<int>(),
396
+ vertical_indices_count.data_ptr<int>(),
397
+ slash_indices_count.data_ptr<int>(), block_count.data_ptr<int>(),
398
+ block_offset.data_ptr<int>(), column_count.data_ptr<int>(),
399
+ column_index.data_ptr<int>(), batch_size, num_heads, num_rows,
400
+ block_size_M, block_size_N, nnz_vertical, nnz_slash, causal);
401
+}
csrc/ops.h
+25
@@ -59,6 +59,31 @@ void merge_attn_states(torch::Tensor& output,
59
const torch::Tensor& prefix_lse,
60
const torch::Tensor& suffix_output,
61
const torch::Tensor& suffix_lse);
62
+
63
+void convert_vertical_slash_indexes(
64
+ torch::Tensor& block_count, // [BATCH, N_HEADS, NUM_ROWS]
65
+ torch::Tensor& block_offset, // [BATCH, N_HEADS, NUM_ROWS, NNZ_S]
66
+ torch::Tensor& column_count, // [BATCH, N_HEADS, NUM_ROWS]
67
+ torch::Tensor& column_index, // [BATCH, N_HEADS, NUM_ROWS, NNZ_V]
68
+ torch::Tensor q_seqlens, // [BATCH, ]
69
+ torch::Tensor kv_seqlens, // [BATCH, ]
70
+ torch::Tensor vertical_indexes, // [BATCH, N_HEADS, NNZ_V]
71
+ torch::Tensor slash_indexes, // [BATCH, N_HEADS, NNZ_S]
72
+ int64_t context_size, int64_t block_size_M, int64_t block_size_N,
73
+ bool causal);
74
+
75
+void convert_vertical_slash_indexes_mergehead(
76
+ torch::Tensor& block_count, // [BATCH, N_HEADS, NUM_ROWS]
77
+ torch::Tensor& block_offset, // [BATCH, N_HEADS, NUM_ROWS, NNZ_S]
78
+ torch::Tensor& column_count, // [BATCH, N_HEADS, NUM_ROWS]
79
+ torch::Tensor& column_index, // [BATCH, N_HEADS, NUM_ROWS, NNZ_V]
80
+ torch::Tensor q_seqlens, // [BATCH, ]
81
+ torch::Tensor kv_seqlens, // [BATCH, ]
82
+ torch::Tensor vertical_indexes, // [BATCH, N_HEADS, NNZ_V]
83
+ torch::Tensor slash_indexes, // [BATCH, N_HEADS, NNZ_S]
84
+ torch::Tensor vertical_indices_count, // [N_HEADS, ]
85
+ torch::Tensor slash_indices_count, int64_t context_size,
86
+ int64_t block_size_M, int64_t block_size_N, bool causal);
87
#endif
88
89
void rms_norm(torch::Tensor& out, torch::Tensor& input, torch::Tensor& weight,
csrc/torch_bindings.cpp
+23
@@ -77,6 +77,29 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
77
" Tensor suffix_output,"
78
" Tensor suffix_lse) -> ()");
79
ops.impl("merge_attn_states", torch::kCUDA, &merge_attn_states);
80
+
81
+ ops.def(
82
+ "convert_vertical_slash_indexes("
83
+ " Tensor! block_count, Tensor! block_offset, "
84
+ " Tensor! column_count, Tensor! column_index, "
85
+ " Tensor q_seqlens, Tensor q_seqlens, "
86
+ " Tensor vertical_indexes, Tensor slash_indexes, "
87
+ " int context_size, int block_size_M, int block_size_N, "
88
+ " bool causal) -> ()");
89
+ ops.impl("convert_vertical_slash_indexes", torch::kCUDA,
90
+ &convert_vertical_slash_indexes);
91
+
92
+ ops.def(
93
+ "convert_vertical_slash_indexes_mergehead("
94
+ " Tensor! block_count, Tensor! block_offset, "
95
+ " Tensor! column_count, Tensor! column_index, "
96
+ " Tensor q_seqlens, Tensor q_seqlens, "
97
+ " Tensor vertical_indexes, Tensor slash_indexes, "
98
+ " Tensor vertical_indices_count, Tensor slash_indices_count, "
99
+ " int context_size, int block_size_M, int block_size_N, "
100
+ " bool causal) -> ()");
101
+ ops.impl("convert_vertical_slash_indexes_mergehead", torch::kCUDA,
102
+ &convert_vertical_slash_indexes_mergehead);
103
#endif
104
105
// Activation ops
examples/offline_inference/qwen_1m.py
new
+66
@@ -0,0 +1,66 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+import os
3
+from urllib.request import urlopen
4
+
5
+from vllm import LLM, SamplingParams
6
+
7
+os.environ["VLLM_ATTENTION_BACKEND"] = "DUAL_CHUNK_FLASH_ATTN"
8
+os.environ["VLLM_ALLOW_LONG_MAX_MODEL_LEN"] = "1"
9
+
10
+
11
+def load_prompt() -> str:
12
+ # Test cases with various lengths can be found at:
13
+ #
14
+ # https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen2.5-1M/test-data/64k.txt
15
+ # https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen2.5-1M/test-data/200k.txt
16
+ # https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen2.5-1M/test-data/600k.txt
17
+ # https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen2.5-1M/test-data/1m.txt
18
+
19
+ with urlopen(
20
+ "https://qianwen-res.oss-cn-beijing.aliyuncs.com"
21
+ "/Qwen2.5-1M/test-data/600k.txt",
22
+ timeout=5) as response:
23
+ prompt = response.read().decode('utf-8')
24
+ return prompt
25
+
26
+
27
+# Processing the prompt.
28
+def process_requests(llm: LLM, prompts: list[str]) -> None:
29
+ # Create a sampling params object.
30
+ sampling_params = SamplingParams(
31
+ temperature=0.7,
32
+ top_p=0.8,
33
+ top_k=20,
34
+ repetition_penalty=1.05,
35
+ detokenize=True,
36
+ max_tokens=256,
37
+ )
38
+ # Generate texts from the prompts.
39
+ outputs = llm.generate(prompts, sampling_params)
40
+ # Print the outputs.
41
+ for output in outputs:
42
+ prompt_token_ids = output.prompt_token_ids
43
+ generated_text = output.outputs[0].text
44
+ print(f"Prompt length: {len(prompt_token_ids)}, "
45
+ f"Generated text: {generated_text!r}")
46
+
47
+
48
+# Create an LLM.
49
+def initialize_engine() -> LLM:
50
+ llm = LLM(model="Qwen/Qwen2.5-7B-Instruct-1M",
51
+ max_model_len=1048576,
52
+ tensor_parallel_size=4,
53
+ enforce_eager=True,
54
+ enable_chunked_prefill=True,
55
+ max_num_batched_tokens=131072)
56
+ return llm
57
+
58
+
59
+def main():
60
+ llm = initialize_engine()
61
+ prompt = load_prompt()
62
+ process_requests(llm, [prompt])
63
+
64
+
65
+if __name__ == '__main__':
66
+ main()
vllm/_custom_ops.py
+95
@@ -150,6 +150,101 @@ def merge_attn_states(output: torch.Tensor,
150
prefix_lse, suffix_output, suffix_lse)
151
152
153
+def convert_vertical_slash_indexes(
154
+ q_seqlens: torch.Tensor, # [BATCH, ]
155
+ kv_seqlens: torch.Tensor, # [BATCH, ]
156
+ vertical_indexes: torch.Tensor, # [BATCH, N_HEADS, NNZ_V]
157
+ slash_indexes: torch.Tensor, # [BATCH, N_HEADS, NNZ_S]
158
+ context_size: int,
159
+ block_size_M: int,
160
+ block_size_N: int,
161
+ causal: bool = True,
162
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
163
+ batch_size = slash_indexes.size(0)
164
+ num_heads = slash_indexes.size(1)
165
+ nnz_slash = slash_indexes.size(2)
166
+ nnz_vertical = vertical_indexes.size(2)
167
+ num_rows = (context_size + block_size_M - 1) // block_size_M
168
+
169
+ block_count = torch.zeros(batch_size,
170
+ num_heads,
171
+ num_rows,
172
+ dtype=q_seqlens.dtype,
173
+ device=q_seqlens.device)
174
+ block_offset = torch.zeros(batch_size,
175
+ num_heads,
176
+ num_rows,
177
+ nnz_slash,
178
+ dtype=q_seqlens.dtype,
179
+ device=q_seqlens.device)
180
+ column_count = torch.zeros(batch_size,
181
+ num_heads,
182
+ num_rows,
183
+ dtype=q_seqlens.dtype,
184
+ device=q_seqlens.device)
185
+ column_index = torch.zeros(batch_size,
186
+ num_heads,
187
+ num_rows,
188
+ nnz_vertical,
189
+ dtype=q_seqlens.dtype,
190
+ device=q_seqlens.device)
191
+
192
+ torch.ops._C.convert_vertical_slash_indexes(
193
+ block_count, block_offset, column_count, column_index, q_seqlens,
194
+ kv_seqlens, vertical_indexes, slash_indexes, context_size,
195
+ block_size_M, block_size_N, causal)
196
+ return block_count, block_offset, column_count, column_index
197
+
198
+
199
+def convert_vertical_slash_indexes_mergehead(
200
+ q_seqlens: torch.Tensor, # [BATCH, ]
201
+ kv_seqlens: torch.Tensor, # [BATCH, ]
202
+ vertical_indexes: torch.Tensor, # [BATCH, N_HEADS, NNZ_V]
203
+ slash_indexes: torch.Tensor, # [BATCH, N_HEADS, NNZ_S]
204
+ # [N_HEADS] : different head use different number of indices
205
+ vertical_indices_count: torch.Tensor,
206
+ slash_indices_count: torch.Tensor,
207
+ context_size: int,
208
+ block_size_M: int,
209
+ block_size_N: int,
210
+ causal: bool = True,
211
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
212
+ batch_size = slash_indexes.size(0)
213
+ num_heads = slash_indexes.size(1)
214
+ nnz_slash = slash_indexes.size(2)
215
+ nnz_vertical = vertical_indexes.size(2)
216
+ num_rows = (context_size + block_size_M - 1) // block_size_M
217
+
218
+ block_count = torch.empty(batch_size,
219
+ num_heads,
220
+ num_rows,
221
+ dtype=q_seqlens.dtype,
222
+ device=q_seqlens.device)
223
+ block_offset = torch.empty(batch_size,
224
+ num_heads,
225
+ num_rows,
226
+ nnz_slash,
227
+ dtype=q_seqlens.dtype,
228
+ device=q_seqlens.device)
229
+ column_count = torch.empty(batch_size,
230
+ num_heads,
231
+ num_rows,
232
+ dtype=q_seqlens.dtype,
233
+ device=q_seqlens.device)
234
+ column_index = torch.empty(batch_size,
235
+ num_heads,
236
+ num_rows,
237
+ nnz_vertical,
238
+ dtype=q_seqlens.dtype,
239
+ device=q_seqlens.device)
240
+
241
+ torch.ops._C.convert_vertical_slash_indexes_mergehead(
242
+ block_count, block_offset, column_count, column_index, q_seqlens,
243
+ kv_seqlens, vertical_indexes, slash_indexes, vertical_indices_count,
244
+ slash_indices_count, context_size, block_size_M, block_size_N, causal)
245
+ return block_count, block_offset, column_count, column_index
246
+
247
+
248
# pos encoding ops
249
def rotary_embedding(
250
positions: torch.Tensor,
vllm/attention/backends/dual_chunk_flash_attn.py
new
+1494
@@ -0,0 +1,1494 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+"""Attention layer with Dual chunk flash attention and sparse attention.
3
+"""
4
+import math
5
+from dataclasses import dataclass
6
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Type
7
+
8
+import torch
9
+import torch.distributed
10
+import torch.nn.functional as F
11
+
12
+from vllm import _custom_ops as ops
13
+from vllm.attention.backends.abstract import AttentionLayer, AttentionType
14
+from vllm.attention.backends.flash_attn import (FlashAttentionBackend,
15
+ FlashAttentionImpl,
16
+ FlashAttentionMetadata,
17
+ FlashAttentionMetadataBuilder)
18
+from vllm.distributed.parallel_state import get_tensor_model_parallel_rank
19
+from vllm.logger import init_logger
20
+from vllm.utils import async_tensor_h2d
21
+from vllm.vllm_flash_attn import (flash_attn_varlen_func,
22
+ flash_attn_with_kvcache, sparse_attn_func)
23
+
24
+if TYPE_CHECKING:
25
+ from vllm.worker.model_runner import ModelInputForGPUBuilder
26
+
27
+logger = init_logger(__name__)
28
+
29
+
30
+class DualChunkFlashAttentionBackend(FlashAttentionBackend):
31
+
32
+ accept_output_buffer: bool = False
33
+
34
+ @staticmethod
35
+ def get_name() -> str:
36
+ return "DUAL_CHUNK_FLASH_ATTN"
37
+
38
+ @staticmethod
39
+ def get_impl_cls() -> Type["DualChunkFlashAttentionImpl"]:
40
+ return DualChunkFlashAttentionImpl
41
+
42
+ @staticmethod
43
+ def get_metadata_cls() -> Type["DualChunkFlashAttentionMetadata"]:
44
+ return DualChunkFlashAttentionMetadata
45
+
46
+ @staticmethod
47
+ def get_builder_cls() -> Type["DualChunkFlashAttentionMetadataBuilder"]:
48
+ return DualChunkFlashAttentionMetadataBuilder
49
+
50
+
51
+@dataclass
52
+class DualChunkFlashAttentionMetadata(FlashAttentionMetadata):
53
+ # Block size of the paged kv cache.
54
+ block_size: int = 16
55
+
56
+ # Original max position embeddings.
57
+ original_max_position_embeddings: int = 0
58
+
59
+ # Chunk size
60
+ chunk_size: int = 8192
61
+
62
+ # Local size
63
+ local_size: int = 1024
64
+
65
+ # (batch_size,). The orig sequence length per sequence.
66
+ orig_seq_lens: Optional[List[int]] = None
67
+
68
+ # orig_seq_lens stored as a tensor.
69
+ orig_seq_lens_tensor: Optional[torch.Tensor] = None
70
+
71
+ # Length scaling factor
72
+ scaling_factor: Optional[torch.Tensor] = None
73
+
74
+ # (batch_size,). Sequence lengths for intra attention.
75
+ seq_lens_intra: Optional[torch.Tensor] = None
76
+
77
+ # Max sequence length for intra attention.
78
+ max_seq_len_intra: Optional[int] = None
79
+
80
+ # (batch_size, num_blocks). Block table for intra attention.
81
+ block_tables_intra: Optional[torch.Tensor] = None
82
+
83
+ # (batch_size,). Sequence lengths for succ attention.
84
+ seq_lens_succ: Optional[torch.Tensor] = None
85
+
86
+ # Max sequence length for succ attention.
87
+ max_seq_len_succ: Optional[int] = None
88
+
89
+ # (batch_size, num_blocks). Block table for succ attention.
90
+ block_tables_succ: Optional[torch.Tensor] = None
91
+
92
+ # (batch_size,). Sequence lengths for inter attention.
93
+ seq_lens_inter: Optional[torch.Tensor] = None
94
+
95
+ # Max sequence length for inter attention.
96
+ max_seq_len_inter: Optional[int] = None
97
+
98
+ _cached_prefill_metadata: Optional[
99
+ "DualChunkFlashAttentionMetadata"] = None
100
+ _cached_decode_metadata: Optional["DualChunkFlashAttentionMetadata"] = None
101
+
102
+ @property
103
+ def prefill_metadata(self) -> Optional["DualChunkFlashAttentionMetadata"]:
104
+ if self.num_prefills == 0:
105
+ return None
106
+
107
+ if self._cached_prefill_metadata is not None:
108
+ return self._cached_prefill_metadata
109
+
110
+ prefill_metadata = super().prefill_metadata
111
+ if prefill_metadata is None:
112
+ return None
113
+
114
+ prefill_metadata = DualChunkFlashAttentionMetadata(
115
+ **prefill_metadata.asdict_zerocopy())
116
+
117
+ prefill_metadata.orig_seq_lens = (
118
+ None if self.orig_seq_lens is None else
119
+ self.orig_seq_lens[:self.num_prefills])
120
+ prefill_metadata.orig_seq_lens_tensor = (
121
+ None if self.orig_seq_lens_tensor is None else
122
+ self.orig_seq_lens_tensor[:self.num_prefills])
123
+
124
+ if self.original_max_position_embeddings > 0:
125
+ assert prefill_metadata.orig_seq_lens_tensor is not None
126
+ prefill_metadata.scaling_factor = (
127
+ 0.1 * torch.log(prefill_metadata.orig_seq_lens_tensor /
128
+ self.original_max_position_embeddings) +
129
+ 1.0).clip(min=1)
130
+
131
+ self._cached_prefill_metadata = prefill_metadata
132
+ return prefill_metadata
133
+
134
+ @property
135
+ def decode_metadata(self) -> Optional["DualChunkFlashAttentionMetadata"]:
136
+ if self.num_decode_tokens == 0:
137
+ return None
138
+
139
+ if self._cached_decode_metadata is not None:
140
+ return self._cached_decode_metadata
141
+
142
+ decode_metadata = super().decode_metadata
143
+ if decode_metadata is None:
144
+ return None
145
+
146
+ decode_metadata = DualChunkFlashAttentionMetadata(
147
+ **decode_metadata.asdict_zerocopy())
148
+
149
+ decode_metadata.orig_seq_lens_tensor = (
150
+ None if self.orig_seq_lens_tensor is None else
151
+ self.orig_seq_lens_tensor[self.num_prefills:])
152
+
153
+ assert decode_metadata.orig_seq_lens_tensor is not None
154
+ assert decode_metadata.block_tables is not None
155
+
156
+ cache_seq_lens = decode_metadata.orig_seq_lens_tensor
157
+ chunk_len = self.chunk_size - self.local_size
158
+ chunk_num_curr = (cache_seq_lens - 1) // chunk_len
159
+ batch_size = decode_metadata.num_decode_tokens
160
+
161
+ if self.original_max_position_embeddings > 0:
162
+ decode_metadata.scaling_factor = (0.1 * torch.log(
163
+ cache_seq_lens / self.original_max_position_embeddings) +
164
+ 1.0).clip(min=1)
165
+
166
+ seq_lens_intra = cache_seq_lens - chunk_num_curr * chunk_len
167
+ max_seq_len_intra = seq_lens_intra.max().item()
168
+ decode_metadata.seq_lens_intra = seq_lens_intra
169
+ decode_metadata.max_seq_len_intra = max_seq_len_intra
170
+
171
+ block_tables_intra = torch.zeros(
172
+ batch_size,
173
+ (max_seq_len_intra - 1) // self.block_size + 1,
174
+ dtype=decode_metadata.block_tables.dtype,
175
+ device=decode_metadata.block_tables.device,
176
+ )
177
+ for i in range(batch_size):
178
+ st = chunk_num_curr[i] * chunk_len // self.block_size
179
+ ed = min(
180
+ st + (max_seq_len_intra - 1) // self.block_size + 1,
181
+ (cache_seq_lens[i] - 1) // self.block_size + 1,
182
+ )
183
+ block_tables_intra[i, :ed -
184
+ st] = decode_metadata.block_tables[i, st:ed]
185
+ decode_metadata.block_tables_intra = block_tables_intra
186
+
187
+ seq_lens_succ = (chunk_num_curr -
188
+ (chunk_num_curr - 1).clip(min=0)) * chunk_len
189
+ max_seq_len_succ = seq_lens_succ.max().item()
190
+ decode_metadata.seq_lens_succ = seq_lens_succ
191
+ decode_metadata.max_seq_len_succ = max_seq_len_succ
192
+ if max_seq_len_succ:
193
+ block_tables_succ = torch.zeros(
194
+ batch_size,
195
+ (max_seq_len_succ - 1) // self.block_size + 1,
196
+ dtype=decode_metadata.block_tables.dtype,
197
+ device=decode_metadata.block_tables.device,
198
+ )
199
+ for i in range(batch_size):
200
+ start = ((chunk_num_curr[i] - 1).clip(min=0) * chunk_len //
201
+ self.block_size)
202
+ end = min(
203
+ start + (max_seq_len_succ - 1) // self.block_size + 1,
204
+ (cache_seq_lens[i] - 1) // self.block_size + 1,
205
+ )
206
+ block_tables_succ[
207
+ i, :end - start] = decode_metadata.block_tables[i,
208
+ start:end]
209
+ decode_metadata.block_tables_succ = block_tables_succ
210
+
211
+ seq_lens_inter = (chunk_num_curr - 1).clip(min=0) * chunk_len
212
+ max_seq_len_inter = seq_lens_inter.max().item()
213
+ decode_metadata.seq_lens_inter = seq_lens_inter
214
+ decode_metadata.max_seq_len_inter = max_seq_len_inter
215
+
216
+ self._cached_decode_metadata = decode_metadata
217
+ return decode_metadata
218
+
219
+
220
+class DualChunkFlashAttentionMetadataBuilder(FlashAttentionMetadataBuilder):
221
+
222
+ def prepare(self):
223
+ super().prepare()
224
+ self.orig_seq_lens: List[int] = []
225
+
226
+ def _add_seq_group(
227
+ self, inter_data: "ModelInputForGPUBuilder.InterDataForSeqGroup",
228
+ chunked_prefill_enabled: bool, prefix_cache_hit: bool):
229
+ super()._add_seq_group(inter_data, chunked_prefill_enabled,
230
+ prefix_cache_hit)
231
+ for prompt_len, seq_len in zip(inter_data.prompt_lens,
232
+ inter_data.seq_lens):
233
+ self.orig_seq_lens.append(max(prompt_len, seq_len))
234
+
235
+ def build(self, seq_lens: List[int], query_lens: List[int],
236
+ cuda_graph_pad_size: int, batch_size: int):
237
+ attn_metadata = super().build(seq_lens, query_lens,
238
+ cuda_graph_pad_size, batch_size)
239
+ attn_metadata = DualChunkFlashAttentionMetadata(
240
+ **attn_metadata.asdict_zerocopy())
241
+
242
+ device = self.runner.device
243
+ attn_metadata.orig_seq_lens = self.orig_seq_lens
244
+ attn_metadata.orig_seq_lens_tensor = async_tensor_h2d(
245
+ self.orig_seq_lens, torch.int, device, self.runner.pin_memory)
246
+
247
+ attn_metadata.block_size = self.runner.block_size
248
+ dual_chunk_attn_config = getattr(self.runner.model_config.hf_config,
249
+ "dual_chunk_attention_config", {})
250
+ attn_metadata.original_max_position_embeddings = \
251
+ dual_chunk_attn_config.get("original_max_position_embeddings", 0)
252
+ attn_metadata.chunk_size = dual_chunk_attn_config.get(
253
+ "chunk_size", 8192)
254
+ attn_metadata.local_size = dual_chunk_attn_config.get(
255
+ "local_size", 1024)
256
+
257
+ return attn_metadata
258
+
259
+
260
+class DualChunkFlashAttentionImpl(FlashAttentionImpl):
261
+ """
262
+ If the input tensors contain prompt tokens, the layout is as follows:
263
+ |<--------------- num_prefill_tokens ----------------->|
264
+ |<--prefill_0-->|<--prefill_1-->|...|<--prefill_N-1--->|
265
+ Otherwise, the layout is as follows:
266
+ |<----------------- num_decode_tokens ------------------>|
267
+ |<--decode_0-->|..........|<--decode_M-1-->|<--padding-->|
268
+ Generation tokens can contain padding when cuda-graph is used.
269
+ Currently, prompt tokens don't contain any padding.
270
+ The prompts might have different lengths, while the generation tokens
271
+ always have length 1.
272
+ If chunked prefill is enabled, prefill tokens and decode tokens can be
273
+ batched together in a flattened 1D query.
274
+ |<----- num_prefill_tokens ---->|<------- num_decode_tokens --------->|
275
+ |<-prefill_0->|...|<-prefill_N-1->|<--decode_0-->|...|<--decode_M-1-->|
276
+ Currently, cuda graph is disabled for chunked prefill, meaning there's no
277
+ padding between prefill and decode tokens.
278
+ """
279
+
280
+ def __init__(
281
+ self,
282
+ num_heads: int,
283
+ head_size: int,
284
+ scale: float,
285
+ num_kv_heads: int,
286
+ alibi_slopes: Optional[List[float]],
287
+ sliding_window: Optional[int],
288
+ kv_cache_dtype: str,
289
+ blocksparse_params: Optional[Dict[str, Any]] = None,
290
+ logits_soft_cap: Optional[float] = None,
291
+ attn_type: str = AttentionType.DECODER,
292
+ layer_idx: int = -1,
293
+ dual_chunk_attention_config: Optional[Dict[str, Any]] = None,
294
+ ) -> None:
295
+ self.num_heads = num_heads
296
+ self.head_size = head_size
297
+ self.scale = float(scale)
298
+ self.num_kv_heads = num_kv_heads
299
+ if alibi_slopes is not None:
300
+ alibi_slopes = torch.tensor(alibi_slopes, dtype=torch.float32)
301
+ self.alibi_slopes = alibi_slopes
302
+ self.sliding_window = ((sliding_window, sliding_window)
303
+ if sliding_window is not None else (-1, -1))
304
+ self.kv_cache_dtype = kv_cache_dtype
305
+
306
+ assert self.num_heads % self.num_kv_heads == 0
307
+ self.num_queries_per_kv = self.num_heads // self.num_kv_heads
308
+ if sliding_window is not None:
309
+ # NOTE(woosuk): flash-attn's sliding window does not work with
310
+ # paged KV cache.
311
+ raise ValueError(
312
+ "Sliding window is not supported in FlashAttention.")
313
+
314
+ support_head_sizes = (
315
+ DualChunkFlashAttentionBackend.get_supported_head_sizes())
316
+
317
+ if head_size not in support_head_sizes:
318
+ raise ValueError(
319
+ f"Head size {head_size} is not supported by FlashAttention. "
320
+ f"Supported head sizes are: {support_head_sizes}.")
321
+
322
+ assert dual_chunk_attention_config is not None
323
+ self.chunk_size = dual_chunk_attention_config.get("chunk_size", 8192)
324
+ self.local_size = dual_chunk_attention_config.get("local_size", 1024)
325
+ self.original_max_position_embeddings = dual_chunk_attention_config.get(
326
+ "original_max_position_embeddings", 0)
327
+ self.sparse_attention_config = dual_chunk_attention_config.get(
328
+ "sparse_attention_config", None)
329
+ if not self.sparse_attention_config:
330
+ logger.warning_once("Sparse attention will not be enabled as "
331
+ "sparse attention config is not provided.")
332
+ self.sparse_attention_enabled = dual_chunk_attention_config.get(
333
+ "sparse_attention_enabled", self.sparse_attention_config
334
+ is not None)
335
+ self.sparse_attention_threshold = dual_chunk_attention_config.get(
336
+ "sparse_attention_threshold", 32768)
337
+ self.sparse_attention_last_q = dual_chunk_attention_config.get(
338
+ "sparse_attention_last_q", 64)
339
+ self.layer_idx = layer_idx
340
+ self.dual_chunk_attention_config = dual_chunk_attention_config
341
+
342
+ if self.sparse_attention_config:
343
+ self.sparse_attention_config = {
344
+ int(i): j
345
+ for i, j in self.sparse_attention_config[
346
+ self.layer_idx].items()
347
+ }
348
+ start_head = self.num_heads * get_tensor_model_parallel_rank()
349
+ end_head = start_head + self.num_heads
350
+ self.sparse_attention_config = [
351
+ self.sparse_attention_config[i]
352
+ for i in range(start_head, end_head)
353
+ ]
354
+
355
+ if self.sparse_attention_enabled:
356
+ self.arange = torch.arange(self.sparse_attention_last_q,
357
+ device="cuda")
358
+ self.last_q_mask = (self.arange[None, None, :, None]
359
+ >= self.arange[None, None, None, :])
360
+
361
+ def forward( # type: ignore
362
+ self,
363
+ layer: AttentionLayer,
364
+ query: torch.Tensor,
365
+ key: torch.Tensor,
366
+ value: torch.Tensor,
367
+ kv_cache: torch.Tensor,
368
+ attn_metadata: DualChunkFlashAttentionMetadata,
369
+ ) -> torch.Tensor:
370
+ """Forward pass with DualChunkFlashAttention.
371
+ Args:
372
+ query: shape = [num_tokens, num_heads * head_size]
373
+ query_succ: shape = [num_tokens, num_heads * head_size]
374
+ query_inter: shape = [num_tokens, num_heads * head_size]
375
+ key: shape = [num_tokens, num_kv_heads * head_size]
376
+ value: shape = [num_tokens, num_kv_heads * head_size]
377
+ kv_cache = [2, num_blocks, block_size, num_kv_heads * head_size]
378
+ attn_metadata: Metadata for attention.
379
+ Returns:
380
+ shape = [num_tokens, num_heads * head_size]
381
+ """
382
+ (
383
+ query,
384
+ query_succ,
385
+ query_inter,
386
+ query_succ_critical,
387
+ query_inter_critical,
388
+ ) = torch.split(query, query.shape[-1] // 5, dim=-1)
389
+
390
+ assert (
391
+ query_succ is not None and query_inter is not None
392
+ ), "query_succ and query_inter are required in Dual Chunk Attention."
393
+
394
+ num_tokens, hidden_size = query.shape
395
+
396
+ # Reshape the query, key, and value tensors.
397
+ query = query.view(-1, self.num_heads, self.head_size)
398
+ query_succ = query_succ.view(-1, self.num_heads, self.head_size)
399
+ query_inter = query_inter.view(-1, self.num_heads, self.head_size)
400
+ query_succ_critical = query_succ_critical.view(-1, self.num_heads,
401
+ self.head_size)
402
+ query_inter_critical = query_inter_critical.view(
403
+ -1, self.num_heads, self.head_size)
404
+ key = key.view(-1, self.num_kv_heads, self.head_size)
405
+ value = value.view(-1, self.num_kv_heads, self.head_size)
406
+
407
+ if self.original_max_position_embeddings > 0:
408
+ if prefill_meta := attn_metadata.prefill_metadata:
409
+ assert prefill_meta.scaling_factor is not None
410
+ assert prefill_meta.query_start_loc is not None
411
+ assert prefill_meta.orig_seq_lens is not None
412
+ current_start = 0
413
+ query_start_loc_cpu = prefill_meta.query_start_loc.cpu()
414
+ for i in range(len(prefill_meta.orig_seq_lens)):
415
+ current_end = (current_start +
416
+ (query_start_loc_cpu[i + 1] -
417
+ query_start_loc_cpu[i]).item())
418
+ key[current_start:current_end].mul_(
419
+ prefill_meta.scaling_factor[i])
420
+ current_start = current_end
421
+ assert current_end <= attn_metadata.num_prefill_tokens
422
+ if decode_meta := attn_metadata.decode_metadata:
423
+ assert decode_meta.scaling_factor is not None
424
+ scaling_factor = decode_meta.scaling_factor
425
+ key[attn_metadata.num_prefill_tokens:].mul_(
426
+ scaling_factor.unsqueeze(-1).unsqueeze(-1))
427
+
428
+ if kv_cache is not None and kv_cache.numel() > 0:
429
+ key_cache = kv_cache[0]
430
+ value_cache = kv_cache[1]
431
+
432
+ # Reshape the input keys and values and store them in the cache.
433
+ # If kv_cache is not provided, the new key and value tensors are
434
+ # not cached. This happens during the initial memory profiling run.
435
+ ops.reshape_and_cache_flash(
436
+ key,
437
+ value,
438
+ key_cache,
439
+ value_cache,
440
+ attn_metadata.slot_mapping.flatten(),
441
+ self.kv_cache_dtype,
442
+ layer._k_scale,
443
+ layer._v_scale,
444
+ )
445
+
446
+ num_prefill_tokens = attn_metadata.num_prefill_tokens
447
+ num_decode_tokens = attn_metadata.num_decode_tokens
448
+ assert key.shape[0] == num_prefill_tokens + num_decode_tokens
449
+ assert value.shape[0] == num_prefill_tokens + num_decode_tokens
450
+ output = torch.empty_like(query)
451
+
452
+ # Query for decode. KV is not needed because it is already cached.
453
+ decode_query = query[num_prefill_tokens:]
454
+ decode_query_succ = query_succ[num_prefill_tokens:]
455
+ decode_query_inter = query_inter[num_prefill_tokens:]
456
+
457
+ # QKV for prefill.
458
+ query = query[:num_prefill_tokens]
459
+ query_succ = query_succ[:num_prefill_tokens]
460
+ query_inter = query_inter[:num_prefill_tokens]
461
+ query_succ_critical = query_succ_critical[:num_prefill_tokens]
462
+ query_inter_critical = query_inter_critical[:num_prefill_tokens]
463
+ key = key[:num_prefill_tokens]
464
+ value = value[:num_prefill_tokens]
465
+ assert query.shape[0] == num_prefill_tokens
466
+ assert decode_query.shape[0] == num_decode_tokens
467
+
468
+ if prefill_meta := attn_metadata.prefill_metadata:
469
+ # Prompt run.
470
+ if (kv_cache is None or prefill_meta.block_tables is None
471
+ or prefill_meta.block_tables.numel() == 0):
472
+ # normal attention, called during the profiling run.
473
+ out = flash_attn_varlen_func(
474
+ q=query,
475
+ k=key,
476
+ v=value,
477
+ cu_seqlens_q=prefill_meta.seq_start_loc,
478
+ cu_seqlens_k=prefill_meta.seq_start_loc,
479
+ max_seqlen_q=prefill_meta.max_prefill_seq_len,
480
+ max_seqlen_k=prefill_meta.max_prefill_seq_len,
481
+ softmax_scale=self.scale,
482
+ causal=True,
483
+ window_size=self.sliding_window,
484
+ alibi_slopes=self.alibi_slopes,
485
+ )
486
+ assert output[:num_prefill_tokens].shape == out.shape
487
+ output[:num_prefill_tokens] = out
488
+ else:
489
+ # prefix-enabled attention
490
+ assert prefill_meta.seq_lens is not None
491
+ assert prefill_meta.orig_seq_lens is not None
492
+ output[:num_prefill_tokens] = (
493
+ self._dual_chunk_flash_attn_prefill(
494
+ q=query,
495
+ q_succ=query_succ,
496
+ q_inter=query_inter,
497
+ q_succ_critical=query_succ_critical,
498
+ q_inter_critical=query_inter_critical,
499
+ k=key_cache,
500
+ v=value_cache,
501
+ cu_seqlens_q=prefill_meta.query_start_loc,
502
+ cu_seqlens_k=prefill_meta.seq_start_loc,
503
+ orig_seq_lens=prefill_meta.orig_seq_lens,
504
+ scaling_factor=prefill_meta.scaling_factor,
505
+ softmax_scale=self.scale,
506
+ causal=True,
507
+ window_size=(-1, -1),
508
+ alibi_slopes=self.alibi_slopes,
509
+ block_table=prefill_meta.block_tables,
510
+ chunk_size=self.chunk_size,
511
+ local_size=self.local_size,
512
+ ))
513
+
514
+ if decode_meta := attn_metadata.decode_metadata:
515
+ # Decoding run.
516
+ output[num_prefill_tokens:] = (
517
+ self._dual_chunk_flash_attn_decoding(
518
+ decode_query.unsqueeze(1),
519
+ decode_query_succ.unsqueeze(1),
520
+ decode_query_inter.unsqueeze(1),
521
+ key_cache,
522
+ value_cache,
523
+ block_table=decode_meta.block_tables,
524
+ cache_seqlens=decode_meta.seq_lens_tensor,
525
+ softmax_scale=self.scale,
526
+ causal=True,
527
+ alibi_slopes=self.alibi_slopes,
528
+ chunk_size=self.chunk_size,
529
+ local_size=self.local_size,
530
+ original_max_position_embeddings=self.
531
+ original_max_position_embeddings,
532
+ decode_meta=decode_meta,
533
+ ).squeeze(1))
534
+ # Reshape the output tensor.
535
+ return output.view(num_tokens, hidden_size)
536
+
537
+ def _dual_chunk_flash_attn_prefill(
538
+ self,
539
+ q,
540
+ q_succ,
541
+ q_inter,
542
+ q_succ_critical,
543
+ q_inter_critical,
544
+ k,
545
+ v,
546
+ cu_seqlens_q,
547
+ cu_seqlens_k,
548
+ orig_seq_lens: List[int],
549
+ scaling_factor: torch.Tensor,
550
+ softmax_scale: float,
551
+ causal: Optional[bool] = True,
552
+ window_size: Tuple[int, int] = (-1, -1),
553
+ alibi_slopes: Optional[torch.Tensor] = None,
554
+ block_table: Optional[torch.Tensor] = None,
555
+ chunk_size: int = 8192,
556
+ local_size: int = 1024,
557
+ ):
558
+ if alibi_slopes is not None:
559
+ raise ValueError(
560
+ "Dual Chunk Attention does not support alibi_slopes")
561
+ if not causal:
562
+ raise ValueError(
563
+ "Dual Chunk Attention does not support causal=False")
564
+ if window_size != (-1, -1):
565
+ raise ValueError(
566
+ "Dual Chunk Attention does not support window_size")
567
+
568
+ cu_seqlens_q_cpu = cu_seqlens_q.cpu().tolist()
569
+ cu_seqlens_k_cpu = cu_seqlens_k.cpu().tolist()
570
+ all_outputs = []
571
+
572
+ for i in range(0, len(cu_seqlens_q_cpu) - 1):
573
+ qs = cu_seqlens_q_cpu[i]
574
+ qe = cu_seqlens_q_cpu[i:i + 2][-1]
575
+ ks = cu_seqlens_k_cpu[i]
576
+ ke = cu_seqlens_k_cpu[i:i + 2][-1]
577
+
578
+ current_q = q[qs:qe]
579
+ current_q_succ = q_succ[qs:qe]
580
+ current_q_inter = q_inter[qs:qe]
581
+ current_q_succ_critical = q_succ_critical[qs:qe]
582
+ current_q_inter_critical = q_inter_critical[qs:qe]
583
+
584
+ if block_table is None:
585
+ current_k = k[ks:ke]
586
+ current_v = v[ks:ke]
587
+ current_block_table = None
588
+ current_orig_seq_len = orig_seq_lens[i]
589
+ else:
590
+ current_block_table = block_table[i]
591
+ current_orig_seq_len = orig_seq_lens[i]
592
+ current_k = k
593
+ current_v = v
594
+ sparse_attn_enabled = (self.sparse_attention_enabled
595
+ and current_orig_seq_len
596
+ > self.sparse_attention_threshold)
597
+
598
+ if current_q.shape[0] == 0:
599
+ continue
600
+
601
+ if current_k.shape[0] == 0:
602
+ all_outputs.append(
603
+ torch.zeros(
604
+ (current_q.shape[0], current_q.shape[1], v.shape[2]),
605
+ device=q.device,
606
+ dtype=q.dtype,
607
+ ))
608
+ continue
609
+
610
+ current_output = torch.empty_like(current_q)
611
+ group_size = int(current_q.size(-2) / current_k.size(-2))
612
+
613
+ if sparse_attn_enabled:
614
+ num_device_q_heads = current_q.size(-2)
615
+ heads_vertical_size = torch.empty(size=(num_device_q_heads, ),
616
+ dtype=torch.int32)
617
+ heads_slash_size = torch.empty(size=(num_device_q_heads, ),
618
+ dtype=torch.int32)
619
+ for head_id in range(current_q.size(-2)):
620
+ (
621
+ ty,
622
+ vertical_size,
623
+ slash_size,
624
+ _,
625
+ ) = self.sparse_attention_config[head_id]
626
+ assert ty == "vertical_and_slash", "only support slash mode"
627
+
628
+ if vertical_size == 30:
629
+ vertical_size += 100
630
+ heads_vertical_size[head_id] = vertical_size
631
+ heads_slash_size[head_id] = slash_size
632
+
633
+ current_output = self._dual_chunk_flash_attn_prefill_func(
634
+ current_q, # allheads
635
+ current_q_succ,
636
+ current_q_inter,
637
+ current_q_succ_critical,
638
+ current_q_inter_critical,
639
+ current_k,
640
+ current_v,
641
+ current_block_table,
642
+ softmax_scale,
643
+ chunk_size,
644
+ local_size,
645
+ scaling_factor[i].item(),
646
+ ke - ks,
647
+ sparse_attn_enabled=sparse_attn_enabled,
648
+ heads_vertical_size=heads_vertical_size,
649
+ heads_slash_size=heads_slash_size,
650
+ group_size=group_size)
651
+ else:
652
+ for head_id in range(current_q.size(-2)):
653
+ # (seq_len, num_heads, head_size)
654
+ current_q_head = current_q[:, head_id, :].unsqueeze(1)
655
+ current_q_succ_head = \
656
+ current_q_succ[:, head_id, :].unsqueeze(1)
657
+ current_q_inter_head = \
658
+ current_q_inter[:, head_id, :].unsqueeze(1)
659
+ current_q_succ_head_critical = \
660
+ current_q_succ_critical[:, head_id, :].unsqueeze(1)
661
+ current_q_inter_head_critical = \
662
+ current_q_inter_critical[:, head_id, :].unsqueeze(1)
663
+ if block_table is not None:
664
+ current_k_head = current_k[..., head_id //
665
+ group_size, :].unsqueeze(2)
666
+ current_v_head = current_v[..., head_id //
667
+ group_size, :].unsqueeze(2)
668
+
669
+ else:
670
+ current_k_head = current_k[:, head_id, :].unsqueeze(1)
671
+ current_v_head = current_v[:, head_id, :].unsqueeze(1)
672
+
673
+ current_out = self._dual_chunk_flash_attn_prefill_func(
674
+ current_q_head,
675
+ current_q_succ_head,
676
+ current_q_inter_head,
677
+ current_q_succ_head_critical,
678
+ current_q_inter_head_critical,
679
+ current_k_head,
680
+ current_v_head,
681
+ current_block_table,
682
+ softmax_scale,
683
+ chunk_size,
684
+ local_size,
685
+ scaling_factor[i].item(),
686
+ ke - ks,
687
+ sparse_attn_enabled=sparse_attn_enabled,
688
+ )
689
+ current_output[:, head_id:head_id + 1, :] = current_out
690
+ all_outputs.append(current_output)
691
+ return torch.cat(all_outputs, dim=0)
692
+
693
+ def _dual_chunk_flash_attn_prefill_func(
694
+ self,
695
+ q,
696
+ q_succ,
697
+ q_inter,
698
+ q_succ_critical,
699
+ q_inter_critical,
700
+ k,
701
+ v,
702
+ block_table,
703
+ softmax_scale: float,
704
+ chunk_size: int,
705
+ local_size: int,
706
+ scaling_factor: float,
707
+ k_length: int,
708
+ sparse_attn_enabled: Optional[bool] = True,
709
+ heads_vertical_size=None,
710
+ heads_slash_size=None,
711
+ group_size=None,
712
+ ):
713
+ flash_results = []
714
+ chunk_len = chunk_size - local_size
715
+
716
+ if block_table is not None:
717
+ block_size = v.shape[1]
718
+ if chunk_len % block_size != 0:
719
+ raise ValueError("chunk_len must be divisible by block_size.")
720
+ else:
721
+ block_size = 1
722
+
723
+ if self.original_max_position_embeddings > 0:
724
+ softmax_scale = softmax_scale * scaling_factor
725
+
726
+ begin = k_length - q.shape[0]
727
+ while begin < k_length:
728
+ flash_per_chunk = []
729
+
730
+ prev_chunk_end_pos = (begin // chunk_len) * chunk_len
731
+ next_chunk_end_pos = prev_chunk_end_pos + chunk_len
732
+ end = min(next_chunk_end_pos, k_length)
733
+ qbegin = begin - (k_length - q.shape[0])
734
+ qend = end - (k_length - q.shape[0])
735
+
736
+ qk_chunks = []
737
+ q_states_intra = q[qbegin:qend]
738
+ # choose critical token
739
+ if block_table is not None:
740
+ block_tables_intra = _get_block(block_table, block_size,
741
+ prev_chunk_end_pos, end)
742
+ k_states_intra = k[block_tables_intra].view(
743
+ -1, *k.shape[-2:])[:(end - prev_chunk_end_pos)]
744
+ v_states_intra = v[block_tables_intra].view(
745
+ -1, *v.shape[-2:])[:(end - prev_chunk_end_pos)]
746
+ else:
747
+ block_tables_intra = None
748
+ k_states_intra = k[prev_chunk_end_pos:end]
749
+ v_states_intra = v[prev_chunk_end_pos:end]
750
+
751
+ if sparse_attn_enabled:
752
+ last_q_size = min(qend - qbegin, self.sparse_attention_last_q)
753
+ _, num_device_k_heads, head_dim = k_states_intra.shape
754
+ k_states_intra = (k_states_intra.unsqueeze(2).repeat(
755
+ 1, 1, group_size,
756
+ 1).reshape(-1, num_device_k_heads * group_size, head_dim))
757
+ v_states_intra = (v_states_intra.unsqueeze(2).repeat(
758
+ 1, 1, group_size,
759
+ 1).reshape(-1, num_device_k_heads * group_size, head_dim))
760
+ qk_chunks.append(
761
+ (q_states_intra.transpose(0, 1)[:, -last_q_size:] *
762
+ softmax_scale) @ k_states_intra.permute(1, 2, 0))
763
+
764
+ if prev_chunk_end_pos - chunk_len >= 0:
765
+ q_states_succ = q_succ[qbegin:qend]
766
+ q_states_succ_critical = q_succ_critical[qbegin:qend]
767
+ if block_table is not None:
768
+ block_tables_succ = _get_block(
769
+ block_table, block_size,
770
+ prev_chunk_end_pos - chunk_len, prev_chunk_end_pos)
771
+ k_states_succ = k[block_tables_succ].view(
772
+ -1, *k.shape[-2:])[:chunk_len]
773
+ v_states_succ = v[block_tables_succ].view(
774
+ -1, *v.shape[-2:])[:chunk_len]
775
+ else:
776
+ k_states_succ = k[prev_chunk_end_pos -
777
+ chunk_len:prev_chunk_end_pos]
778
+ v_states_succ = v[prev_chunk_end_pos -
779
+ chunk_len:prev_chunk_end_pos]
780
+
781
+ if sparse_attn_enabled:
782
+ k_states_succ = (k_states_succ.unsqueeze(2).repeat(
783
+ 1, 1, group_size,
784
+ 1).reshape(-1, num_device_k_heads * group_size,
785
+ head_dim))
786
+ v_states_succ = (v_states_succ.unsqueeze(2).repeat(
787
+ 1, 1, group_size,
788
+ 1).reshape(-1, num_device_k_heads * group_size,
789
+ head_dim))
790
+ qk_chunks.append((q_states_succ_critical.transpose(
791
+ 0, 1)[:, -last_q_size:] * softmax_scale)
792
+ @ k_states_succ.permute(1, 2, 0))
793
+
794
+ if prev_chunk_end_pos - chunk_len * 2 >= 0:
795
+ q_states_inter = q_inter[qbegin:qend]
796
+ q_states_inter_critical = q_inter_critical[qbegin:qend]
797
+ if block_table is not None:
798
+ block_tables_inter = _get_block(
799
+ block_table, block_size, 0,
800
+ prev_chunk_end_pos - chunk_len)
801
+ k_states_inter = k[block_tables_inter].view(
802
+ -1, *k.shape[-2:])[:(prev_chunk_end_pos - chunk_len)]
803
+ v_states_inter = v[block_tables_inter].view(
804
+ -1, *v.shape[-2:])[:(prev_chunk_end_pos - chunk_len)]
805
+ else:
806
+ k_states_inter = k[:prev_chunk_end_pos - chunk_len]
807
+ v_states_inter = v[:prev_chunk_end_pos - chunk_len]
808
+
809
+ if sparse_attn_enabled:
810
+ k_states_inter = (k_states_inter.unsqueeze(2).repeat(
811
+ 1, 1, group_size,
812
+ 1).reshape(-1, num_device_k_heads * group_size,
813
+ head_dim))
814
+ v_states_inter = (v_states_inter.unsqueeze(2).repeat(
815
+ 1, 1, group_size,
816
+ 1).reshape(-1, num_device_k_heads * group_size,
817
+ head_dim))
818
+ qk_chunks.append((q_states_inter_critical.transpose(
819
+ 0, 1)[:, -last_q_size:] * softmax_scale)
820
+ @ k_states_inter.permute(1, 2, 0))
821
+
822
+ if sparse_attn_enabled:
823
+ reversed_qk = qk_chunks[::-1]
824
+ qk = torch.cat(reversed_qk, dim=-1)
825
+
826
+ qk[:, :, -last_q_size:] = torch.where(
827
+ self.last_q_mask[..., -last_q_size:,
828
+ -last_q_size:].to(qk.device),
829
+ qk[:, :, -last_q_size:], -torch.inf)
830
+ qk = F.softmax(qk, dim=-1, dtype=torch.float32)
831
+
832
+ vertical = qk.sum(-2, keepdim=True)
833
+ vertical[..., :30] = torch.inf
834
+
835
+ # Avoid sorting by using the min/max ints to fill the indexer
836
+ # buffers.
837
+ int32_max = torch.iinfo(torch.int32).max
838
+ int32_min = torch.iinfo(torch.int32).min
839
+ n_heads = qk.size()[0]
840
+ max_slash_topk = torch.max(heads_slash_size).item()
841
+ max_vertical_topk = torch.max(heads_vertical_size).item()
842
+ # store each head's slash topk, vertical topk
843
+ vertical = vertical.reshape((n_heads, -1))
844
+ # prevent out of range when prompt size < max_vertical_topk
845
+ max_vertical_topk = min(vertical.shape[-1], max_vertical_topk)
846
+ vertical_topk_buffer = torch.topk(vertical, max_vertical_topk,
847
+ -1).indices
848
+ slash_topk_buffer = torch.empty(size=(n_heads, max_slash_topk),
849
+ dtype=torch.int64,
850
+ device=qk.device)
851
+ for head_i in range(n_heads):
852
+ # (nqheads=1, lastq, k_len)
853
+ head_score = qk[head_i:head_i + 1, :, :]
854
+ slash_scores = _sum_all_diagonal_matrix(head_score)
855
+ if head_score.size(1) != 1:
856
+ # drop right up corner
857
+ slash_scores = slash_scores[..., :-last_q_size + 1]
858
+ slash_scores[..., -100:] = torch.inf
859
+
860
+ head_slash_size = heads_slash_size[head_i]
861
+ head_slash_size = min(head_slash_size, vertical.size(-1))
862
+ slash_topk = torch.topk(slash_scores, head_slash_size,
863
+ -1).indices
864
+ #（nheads, max_topk）
865
+ slash_topk_buffer[head_i, :head_slash_size] = slash_topk
866
+
867
+ # reset heads topk
868
+ heads_slash_size[head_i] = head_slash_size
869
+ heads_vertical_size[head_i] = min(
870
+ heads_vertical_size[head_i], max_vertical_topk)
871
+
872
+ # store
873
+ vertical_buffer = torch.full((n_heads, max_vertical_topk),
874
+ int32_max,
875
+ dtype=torch.int64,
876
+ device=q.device)
877
+ slash_buffer = torch.full((n_heads, max_slash_topk),
878
+ int32_min,
879
+ dtype=torch.int64,
880
+ device=q.device)
881
+ succ_vertical_buffer = torch.full((n_heads, max_vertical_topk),
882
+ int32_max,
883
+ dtype=torch.int64,
884
+ device=q.device)
885
+ succ_slash_buffer = torch.full((n_heads, max_slash_topk),
886
+ int32_min,
887
+ dtype=torch.int64,
888
+ device=q.device)
889
+ inter_vertical_buffer = torch.full(
890
+ (n_heads, max_vertical_topk),
891
+ int32_max,
892
+ dtype=torch.int64,
893
+ device=q.device)
894
+ inter_slash_buffer = torch.full((n_heads, max_slash_topk),
895
+ int32_min,
896
+ dtype=torch.int64,
897
+ device=q.device)
898
+
899
+ vertical_size_buffer = torch.empty(size=(n_heads, ),
900
+ dtype=torch.int32,
901
+ device=q.device)
902
+ slash_sizes_buffer = torch.empty(size=(n_heads, ),
903
+ dtype=torch.int32,
904
+ device=q.device)
905
+ succ_vertical_size_buffer = torch.empty(size=(n_heads, ),
906
+ dtype=torch.int32,
907
+ device=q.device)
908
+ succ_slash_sizes_buffer = torch.empty(size=(n_heads, ),
909
+ dtype=torch.int32,
910
+ device=q.device)
911
+ inter_vertical_size_buffer = torch.empty(size=(n_heads, ),
912
+ dtype=torch.int32,
913
+ device=q.device)
914
+ inter_slash_sizes_buffer = torch.empty(size=(n_heads, ),
915
+ dtype=torch.int32,
916
+ device=q.device)
917
+
918
+ for head_i in range(n_heads):
919
+ vertical_topk = vertical_topk_buffer[
920
+ head_i, :heads_vertical_size[head_i]]
921
+ # intra
922
+ intra_vertical_indices = vertical_topk[
923
+ vertical_topk >=
924
+ prev_chunk_end_pos] - prev_chunk_end_pos
925
+ if intra_vertical_indices.nelement() == 0:
926
+ intra_vertical_indices = torch.cat([
927
+ intra_vertical_indices,
928
+ torch.arange(0,
929
+ k_states_intra.size(0),
930
+ max(1,
931
+ k_states_intra.size(0) / 5),
932
+ dtype=torch.int32,
933
+ device=intra_vertical_indices.device)
934
+ ])
935
+ slash_topk = slash_topk_buffer[
936
+ head_i, :heads_slash_size[head_i]]
937
+ intra_slash_indices = (
938
+ (qk.size(-1) - 1) -
939
+ slash_topk[slash_topk >= prev_chunk_end_pos])
940
+ # fill buffer
941
+ v_count = intra_vertical_indices.nelement()
942
+ s_count = intra_slash_indices.nelement()
943
+ vertical_size_buffer[head_i] = v_count
944
+ slash_sizes_buffer[head_i] = s_count
945
+ vertical_buffer[head_i, :v_count].copy_(
946
+ intra_vertical_indices)
947
+ slash_buffer[head_i, :s_count].copy_(intra_slash_indices)
948
+ # succ
949
+ if prev_chunk_end_pos - chunk_len >= 0:
950
+ succ_vertical_indices = vertical_topk[
951
+ (vertical_topk < prev_chunk_end_pos)
952
+ & (vertical_topk >= prev_chunk_end_pos -
953
+ chunk_len)] - (prev_chunk_end_pos - chunk_len)
954
+ # TODO: support no vertical
955
+ if succ_vertical_indices.nelement() == 0:
956
+ succ_vertical_indices = torch.cat([
957
+ succ_vertical_indices,
958
+ torch.arange(
959
+ 0,
960
+ k_states_succ.size(0),
961
+ max(1,
962
+ k_states_succ.size(0) / 5),
963
+ dtype=torch.int32,
964
+ device=intra_vertical_indices.device)
965
+ ])
966
+ succ_slash_indices = (
967
+ (prev_chunk_end_pos + (qend - qbegin) - 1) -
968
+ slash_topk[((slash_topk >=
969
+ (prev_chunk_end_pos - chunk_len)) &
970
+ (slash_topk < (prev_chunk_end_pos +
971
+ (qend - qbegin))))])
972
+ if succ_slash_indices.nelement() == 0:
973
+ succ_slash_indices = torch.cat([
974
+ succ_slash_indices,
975
+ torch.arange(
976
+ 0,
977
+ k_states_succ.size(0),
978
+ max(1,
979
+ k_states_succ.size(0) / 5),
980
+ dtype=torch.int32,
981
+ device=intra_vertical_indices.device)
982
+ ])
983
+ # fill buffer
984
+ v_count = succ_vertical_indices.nelement()
985
+ s_count = succ_slash_indices.nelement()
986
+ succ_vertical_size_buffer[head_i] = v_count
987
+ succ_slash_sizes_buffer[head_i] = s_count
988
+ succ_vertical_buffer[head_i, :v_count].copy_(
989
+ succ_vertical_indices)
990
+ succ_slash_buffer[head_i, :s_count].copy_(
991
+ succ_slash_indices)
992
+
993
+ if prev_chunk_end_pos - 2 * chunk_len >= 0:
994
+ inter_vertical_indices = vertical_topk[
995
+ vertical_topk < prev_chunk_end_pos - chunk_len]
996
+
997
+ if inter_vertical_indices.nelement() == 0:
998
+ inter_vertical_indices = torch.cat([
999
+ inter_vertical_indices,
1000
+ torch.arange(
1001
+ 0,
1002
+ k_states_inter.size(0),
1003
+ max(1,
1004
+ k_states_inter.size(0) / 5),
1005
+ dtype=torch.int32,
1006
+ device=intra_vertical_indices.device)
1007
+ ])
1008
+ inter_slash_indices = (
1009
+ (prev_chunk_end_pos - chunk_len +
1010
+ (qend - qbegin) - 1) -
1011
+ slash_topk[slash_topk < (prev_chunk_end_pos -
1012
+ chunk_len +
1013
+ (qend - qbegin))])
1014
+ if inter_slash_indices.nelement() == 0:
1015
+ inter_slash_indices = torch.cat([
1016
+ inter_slash_indices,
1017
+ torch.arange(
1018
+ 0,
1019
+ k_states_inter.size(0),
1020
+ max(1,
1021
+ k_states_inter.size(0) / 5),
1022
+ dtype=torch.int32,
1023
+ device=intra_vertical_indices.device)
1024
+ ])
1025
+ # fill buffer
1026
+ v_count = inter_vertical_indices.nelement()
1027
+ s_count = inter_slash_indices.nelement()
1028
+ inter_vertical_size_buffer[head_i] = v_count
1029
+ inter_slash_sizes_buffer[head_i] = s_count
1030
+ inter_vertical_buffer[head_i, :v_count].copy_(
1031
+ inter_vertical_indices)
1032
+ inter_slash_buffer[head_i, :s_count].copy_(
1033
+ inter_slash_indices)
1034
+ else:
1035
+ intra_vertical_indices, intra_slash_indices = None, None
1036
+ succ_vertical_indices, succ_slash_indices = None, None
1037
+ inter_vertical_indices, inter_slash_indices = None, None
1038
+
1039
+ if sparse_attn_enabled:
1040
+ flash_result = self._do_flash_attn(
1041
+ q_states_intra,
1042
+ k_states_intra,
1043
+ v_states_intra,
1044
+ softmax_scale=softmax_scale,
1045
+ causal=True,
1046
+ block_table=block_table,
1047
+ stage="intra",
1048
+ vertical_indices=vertical_buffer,
1049
+ slash_indices=slash_buffer,
1050
+ vertical_indices_count=vertical_size_buffer,
1051
+ slash_indices_count=slash_sizes_buffer,
1052
+ mergehead_softmax_scale=softmax_scale,
1053
+ sparse_attn_enabled=sparse_attn_enabled)
1054
+ else:
1055
+ flash_result = self._do_flash_attn(
1056
+ q_states_intra,
1057
+ k_states_intra,
1058
+ v_states_intra,
1059
+ softmax_scale=softmax_scale,
1060
+ causal=True,
1061
+ block_table=block_table,
1062
+ stage="intra",
1063
+ vertical_indices=intra_vertical_indices,
1064
+ slash_indices=intra_slash_indices,
1065
+ sparse_attn_enabled=sparse_attn_enabled)
1066
+ flash_per_chunk.append(flash_result)
1067
+
1068
+ if prev_chunk_end_pos - chunk_len >= 0:
1069
+ if sparse_attn_enabled:
1070
+ flash_result = self._do_flash_attn(
1071
+ q_states_succ,
1072
+ k_states_succ,
1073
+ v_states_succ,
1074
+ softmax_scale=softmax_scale,
1075
+ causal=False,
1076
+ block_table=block_table,
1077
+ stage="succ",
1078
+ vertical_indices=succ_vertical_buffer,
1079
+ slash_indices=succ_slash_buffer,
1080
+ vertical_indices_count=succ_vertical_size_buffer,
1081
+ slash_indices_count=succ_slash_sizes_buffer,
1082
+ mergehead_softmax_scale=softmax_scale,
1083
+ sparse_attn_enabled=sparse_attn_enabled)
1084
+ else:
1085
+ flash_result = self._do_flash_attn(
1086
+ q_states_succ,
1087
+ k_states_succ,
1088
+ v_states_succ,
1089
+ softmax_scale=softmax_scale,
1090
+ causal=False,
1091
+ block_table=block_table,
1092
+ stage="succ",
1093
+ vertical_indices=succ_vertical_indices,
1094
+ slash_indices=succ_slash_indices,
1095
+ sparse_attn_enabled=sparse_attn_enabled)
1096
+ flash_per_chunk.append(flash_result)
1097
+
1098
+ if prev_chunk_end_pos - chunk_len * 2 >= 0:
1099
+ if sparse_attn_enabled:
1100
+ flash_result = self._do_flash_attn(
1101
+ q_states_inter,
1102
+ k_states_inter,
1103
+ v_states_inter,
1104
+ softmax_scale=softmax_scale,
1105
+ causal=False,
1106
+ block_table=block_table,
1107
+ stage="inter",
1108
+ vertical_indices=inter_vertical_buffer,
1109
+ slash_indices=inter_slash_buffer,
1110
+ vertical_indices_count=inter_vertical_size_buffer,
1111
+ slash_indices_count=inter_slash_sizes_buffer,
1112
+ mergehead_softmax_scale=softmax_scale,
1113
+ sparse_attn_enabled=sparse_attn_enabled)
1114
+ else:
1115
+ flash_result = self._do_flash_attn(
1116
+ q_states_inter,
1117
+ k_states_inter,
1118
+ v_states_inter,
1119
+ softmax_scale=softmax_scale,
1120
+ causal=False,
1121
+ block_table=block_table,
1122
+ stage="inter",
1123
+ vertical_indices=inter_vertical_indices,
1124
+ slash_indices=inter_slash_indices,
1125
+ sparse_attn_enabled=sparse_attn_enabled)
1126
+ flash_per_chunk.append(flash_result)
1127
+
1128
+ flash_results.append(flash_per_chunk)
1129
+ begin = end
1130
+
1131
+ attn_output = self._merge_attn_outputs(flash_results)
1132
+ del flash_results
1133
+ return attn_output
1134
+
1135
+ def _do_flash_attn(
1136
+ self,
1137
+ query_states: torch.Tensor,
1138
+ key_states: torch.Tensor,
1139
+ value_states: torch.Tensor,
1140
+ softmax_scale: float,
1141
+ causal: bool = True,
1142
+ block_table: torch.Tensor = None,
1143
+ max_seqlen_k: Optional[int] = None,
1144
+ stage: str = "intra",
1145
+ vertical_indices: Optional[torch.Tensor] = None,
1146
+ slash_indices: Optional[torch.Tensor] = None,
1147
+ vertical_indices_count: Optional[torch.Tensor] = None,
1148
+ slash_indices_count: Optional[torch.Tensor] = None,
1149
+ mergehead_softmax_scale: Optional[float] = None,
1150
+ sparse_attn_enabled: Optional[bool] = False,
1151
+ ):
1152
+ if max_seqlen_k is None:
1153
+ max_seqlen_k = key_states.shape[0]
1154
+
1155
+ q_len = query_states.shape[0]
1156
+ q_heads = query_states.shape[1]
1157
+ h_dim = query_states.shape[-1]
1158
+
1159
+ if sparse_attn_enabled:
1160
+ assert slash_indices is not None
1161
+ if stage == "intra":
1162
+ assert causal
1163
+ else:
1164
+ assert not causal
1165
+
1166
+ query_states = query_states.unsqueeze(0).transpose(1, 2)
1167
+ key_states = key_states.unsqueeze(0).transpose(1, 2)
1168
+ value_states = value_states.unsqueeze(0).transpose(1, 2)
1169
+
1170
+ q = query_states
1171
+ k = key_states
1172
+ v = value_states
1173
+
1174
+ if (vertical_indices_count is not None and \
1175
+ slash_indices_count is not None):
1176
+ assert mergehead_softmax_scale is not None
1177
+
1178
+ res, s_lse = _vertical_slash_sparse_attention(
1179
+ q,
1180
+ k,
1181
+ v,
1182
+ vertical_indices,
1183
+ slash_indices,
1184
+ mergehead_softmax_scale,
1185
+ causal=causal,
1186
+ stage=stage,
1187
+ vertical_indices_count=vertical_indices_count,
1188
+ slash_indices_count=slash_indices_count)
1189
+ res = res.view(q_heads, q_len,
1190
+ h_dim).transpose(0, 1) # (qlen,nhead,h_dim)
1191
+ s_lse = s_lse.view(
1192
+ q_heads, q_len,
1193
+ 1).squeeze(-1).unsqueeze(0).float() # (1, nhead,qlen)
1194
+ else:
1195
+ res, s_lse = _vertical_slash_sparse_attention(q,
1196
+ k,
1197
+ v,
1198
+ vertical_indices,
1199
+ slash_indices,
1200
+ softmax_scale,
1201
+ causal=causal,
1202
+ stage=stage)
1203
+ res = res.view(q_len, q_heads, h_dim)
1204
+ s_lse = s_lse.view(q_len, q_heads, 1).transpose(0, 2).float()
1205
+ return res, s_lse
1206
+
1207
+ output, softmax_lse = flash_attn_varlen_func(
1208
+ q=query_states,
1209
+ k=key_states,
1210
+ v=value_states,
1211
+ softmax_scale=softmax_scale,
1212
+ cu_seqlens_q=torch.tensor([0, query_states.shape[0]],
1213
+ dtype=torch.int32,
1214
+ device=query_states.device),
1215
+ max_seqlen_q=query_states.shape[0],
1216
+ cu_seqlens_k=torch.tensor([0, max_seqlen_k],
1217
+ dtype=torch.int32,
1218
+ device=query_states.device),
1219
+ max_seqlen_k=max_seqlen_k,
1220
+ causal=causal,
1221
+ block_table=block_table.unsqueeze(0),
1222
+ return_softmax_lse=True,
1223
+ )
1224
+ softmax_lse = softmax_lse.view(q_len, q_heads, 1).transpose(0,
1225
+ 2).float()
1226
+ return output, softmax_lse
1227
+
1228
+ def _merge_attn_outputs(
1229
+ self,
1230
+ flash_results: List[List[Tuple[torch.Tensor, torch.Tensor]]],
1231
+ return_lse: Optional[bool] = False,
1232
+ ) -> torch.Tensor:
1233
+ attn_outputs_all = []
1234
+ logits_all = []
1235
+
1236
+ for flash_per_chunk in flash_results:
1237
+ if len(flash_per_chunk) == 1:
1238
+ attn_outputs_all.append(flash_per_chunk[0][0])
1239
+ if return_lse:
1240
+ logits_all.append(flash_per_chunk[0][1])
1241
+ continue
1242
+
1243
+ attn_outputs = torch.stack([
1244
+ flash_attn_output[0] for flash_attn_output in flash_per_chunk
1245
+ ])
1246
+ logits = torch.stack([
1247
+ flash_attn_output[1] for flash_attn_output in flash_per_chunk
1248
+ ])
1249
+ logits = logits.to(torch.float32)
1250
+
1251
+ if return_lse:
1252
+ max_val = torch.max(logits, dim=0).values
1253
+ diff = torch.abs(logits[0] - logits[1])
1254
+ log_sum_exp = max_val + torch.log1p(torch.exp(-diff))
1255
+ logits_all.append(log_sum_exp)
1256
+
1257
+ max_logits = torch.max(logits, dim=0).values
1258
+ stable_logits = logits - max_logits.unsqueeze(0)
1259
+ lse_s = torch.exp(stable_logits).detach()
1260
+ lse_sum = torch.sum(lse_s, dim=0)
1261
+ lse_s /= lse_sum
1262
+ attn_outputs *= lse_s.unsqueeze(-1).transpose(2, 3).squeeze(1)
1263
+ attn_outputs_all.append(attn_outputs.sum(dim=0))
1264
+
1265
+ if return_lse:
1266
+ return (torch.cat(attn_outputs_all,
1267
+ dim=0), torch.cat(logits_all, dim=-1))
1268
+ else:
1269
+ return torch.cat(attn_outputs_all, dim=0)
1270
+
1271
+ def _dual_chunk_flash_attn_decoding(
1272
+ self,
1273
+ query: torch.Tensor,
1274
+ query_succ: torch.Tensor,
1275
+ query_inter: torch.Tensor,
1276
+ key_cache: torch.Tensor,
1277
+ value_cache: torch.Tensor,
1278
+ block_table: torch.Tensor,
1279
+ cache_seqlens: torch.Tensor,
1280
+ softmax_scale: float,
1281
+ causal: bool,
1282
+ alibi_slopes: Optional[torch.Tensor],
1283
+ chunk_size: int,
1284
+ local_size: int,
1285
+ original_max_position_embeddings: int,
1286
+ decode_meta: DualChunkFlashAttentionMetadata,
1287
+ ):
1288
+ if not causal:
1289
+ raise ValueError(
1290
+ "Dual Chunk Attention does not support causal=False")
1291
+
1292
+ block_size = value_cache.shape[1]
1293
+ chunk_len = chunk_size - local_size
1294
+ if chunk_len % block_size != 0:
1295
+ raise ValueError("chunk_len must be divisible by block_size.")
1296
+ if original_max_position_embeddings > 0:
1297
+ assert decode_meta.scaling_factor is not None
1298
+ scaling_factor = decode_meta.scaling_factor
1299
+ query = (query * scaling_factor.view(-1, 1, 1, 1)).to(
1300
+ query.dtype
1301
+ ) # possible for numerical issue, need to fused in the kernel
1302
+ query_succ = (query_succ * scaling_factor.view(-1, 1, 1, 1)).to(
1303
+ query.dtype)
1304
+ query_inter = (query_inter * scaling_factor.view(-1, 1, 1, 1)).to(
1305
+ query.dtype)
1306
+ outputs_list = []
1307
+ softmax_lses_list = []
1308
+
1309
+ # intra-attention
1310
+ intra_output, intra_softmax_lse = (
1311
+ self._dual_chunk_flash_attn_decoding_with_exp_sums(
1312
+ query,
1313
+ key_cache,
1314
+ value_cache,
1315
+ decode_meta.block_tables_intra,
1316
+ decode_meta.seq_lens_intra,
1317
+ softmax_scale,
1318
+ alibi_slopes,
1319
+ causal=False,
1320
+ ))
1321
+ outputs_list.append(intra_output)
1322
+ softmax_lses_list.append(intra_softmax_lse)
1323
+
1324
+ # succ-attention
1325
+ if decode_meta.max_seq_len_succ:
1326
+ succ_output, succ_softmax_lse = (
1327
+ self._dual_chunk_flash_attn_decoding_with_exp_sums(
1328
+ query_succ,
1329
+ key_cache,
1330
+ value_cache,
1331
+ decode_meta.block_tables_succ,
1332
+ decode_meta.seq_lens_succ,
1333
+ softmax_scale,
1334
+ alibi_slopes,
1335
+ causal=False,
1336
+ ))
1337
+ outputs_list.append(succ_output)
1338
+ softmax_lses_list.append(succ_softmax_lse)
1339
+
1340
+ # inter-attention
1341
+ if decode_meta.max_seq_len_inter:
1342
+ inter_output, inter_softmax_lse = (
1343
+ self._dual_chunk_flash_attn_decoding_with_exp_sums(
1344
+ query_inter,
1345
+ key_cache,
1346
+ value_cache,
1347
+ block_table[:, :decode_meta.max_seq_len_inter],
1348
+ decode_meta.seq_lens_inter,
1349
+ softmax_scale,
1350
+ alibi_slopes,
1351
+ causal=False,
1352
+ ))
1353
+ outputs_list.append(inter_output)
1354
+ softmax_lses_list.append(inter_softmax_lse)
1355
+ outputs = torch.stack(outputs_list, dim=0)
1356
+ del outputs_list
1357
+ softmax_lses = torch.stack(softmax_lses_list, dim=0).to(torch.float32)
1358
+ del softmax_lses_list
1359
+ max_logits = torch.max(softmax_lses, dim=0).values
1360
+ stable_logits = softmax_lses - max_logits.unsqueeze(0)
1361
+ lse_s = torch.exp(stable_logits).detach()
1362
+ lse_sum = torch.sum(lse_s, dim=0)
1363
+ lse_s /= lse_sum
1364
+ outputs *= lse_s.unsqueeze(-1).transpose(2, 3)
1365
+ return outputs.sum(0)
1366
+
1367
+ def _dual_chunk_flash_attn_decoding_with_exp_sums(
1368
+ self,
1369
+ query: torch.Tensor,
1370
+ key_cache: torch.Tensor,
1371
+ value_cache: torch.Tensor,
1372
+ block_table: torch.Tensor,
1373
+ cache_seqlens: torch.Tensor,
1374
+ softmax_scale: float,
1375
+ alibi_slopes: Optional[torch.Tensor],
1376
+ causal: bool,
1377
+ ):
1378
+ out, softmax_lse = flash_attn_with_kvcache(
1379
+ q=query,
1380
+ k_cache=key_cache,
1381
+ v_cache=value_cache,
1382
+ block_table=block_table,
1383
+ cache_seqlens=cache_seqlens,
1384
+ softmax_scale=softmax_scale,
1385
+ alibi_slopes=alibi_slopes,
1386
+ causal=causal,
1387
+ return_softmax_lse=True,
1388
+ )
1389
+ mask = (cache_seqlens == 0)
1390
+ out[mask] = 0
1391
+ softmax_lse[mask] = -float("inf")
1392
+ return out, softmax_lse
1393
+
1394
+
1395
+def _vertical_slash_sparse_attention(
1396
+ query: torch.Tensor, # [BATCH, N_HEADS, N_CTX, D_HEAD]
1397
+ key: torch.Tensor, # [BATCH, N_HEADS, N_KV_CTX, D_HEAD]
1398
+ value: torch.Tensor, # [BATCH, N_HEADS, N_KV_CTX, D_HEAD]
1399
+ v_idx: torch.Tensor, # [BATCH, N_HEADS, NNZ_V]
1400
+ s_idx: torch.Tensor, # [BATCH, N_HEADS, NNZ_S]
1401
+ softmax_scale: float,
1402
+ causal: bool = True,
1403
+ stage: str = "intra",
1404
+ block_size_M: int = 64,
1405
+ block_size_N: int = 64,
1406
+ vertical_indices_count: torch.Tensor = None, # [N_HEADS,]
1407
+ slash_indices_count: torch.Tensor = None,
1408
+):
1409
+ if stage == "intra":
1410
+ assert causal
1411
+ else:
1412
+ assert not causal
1413
+
1414
+ batch_size, num_heads, context_size, head_dim = query.shape
1415
+ _, _, kv_seq_len, _ = key.shape
1416
+
1417
+ if head_dim not in [16, 32, 64, 128, 256, 512]:
1418
+ target_dim = 2**math.ceil(math.log2(head_dim)) - head_dim
1419
+ query = F.pad(query, [0, target_dim, 0, 0, 0, 0, 0, 0])
1420
+ key = F.pad(key, [0, target_dim, 0, 0, 0, 0, 0, 0])
1421
+ value = F.pad(value, [0, target_dim, 0, 0, 0, 0, 0, 0])
1422
+
1423
+ v_idx = v_idx.to(torch.int32).reshape(
1424
+ (batch_size, num_heads, -1)).sort(dim=-1, descending=False)[0]
1425
+ s_idx = s_idx.to(torch.int32).reshape(
1426
+ (batch_size, num_heads, -1)).sort(dim=-1, descending=True)[0]
1427
+ q_seqlens = torch.tensor([context_size],
1428
+ dtype=torch.int32,
1429
+ device=query.device)
1430
+ kv_seqlens = torch.tensor([kv_seq_len],
1431
+ dtype=torch.int32,
1432
+ device=query.device)
1433
+
1434
+ if vertical_indices_count is not None and slash_indices_count is not None:
1435
+ (
1436
+ block_count,
1437
+ block_offset,
1438
+ column_count,
1439
+ column_index,
1440
+ ) = ops.convert_vertical_slash_indexes_mergehead(
1441
+ q_seqlens, kv_seqlens, v_idx, s_idx, vertical_indices_count,
1442
+ slash_indices_count, context_size, block_size_M, block_size_N,
1443
+ causal)
1444
+ else:
1445
+ (
1446
+ block_count,
1447
+ block_offset,
1448
+ column_count,
1449
+ column_index,
1450
+ ) = ops.convert_vertical_slash_indexes(q_seqlens, kv_seqlens, v_idx,
1451
+ s_idx, context_size,
1452
+ block_size_M, block_size_N,
1453
+ causal)
1454
+
1455
+ q = query.transpose(1, 2).contiguous()
1456
+ k = key.transpose(1, 2).contiguous()
1457
+ v = value.transpose(1, 2).contiguous()
1458
+ out, lse = sparse_attn_func(
1459
+ q,
1460
+ k,
1461
+ v,
1462
+ block_count,
1463
+ block_offset,
1464
+ column_count,
1465
+ column_index,
1466
+ causal=causal,
1467
+ softmax_scale=softmax_scale,
1468
+ return_softmax_lse=True,
1469
+ )
1470
+ out = out.transpose(1, 2).contiguous()
1471
+ softmax_lse = lse.reshape(*lse.shape, 1)
1472
+ return (out[..., :context_size, :head_dim],
1473
+ softmax_lse[..., :context_size, :])
1474
+
1475
+
1476
+def _sum_all_diagonal_matrix(mat: torch.tensor):
1477
+ h, n, m = mat.shape
1478
+ # Zero matrix used for padding
1479
+ zero_mat = torch.zeros((h, n, n), device=mat.device)
1480
+ # pads the matrix on left and right
1481
+ mat_padded = torch.cat((zero_mat, mat, zero_mat), -1)
1482
+ # Change the strides
1483
+ mat_strided = mat_padded.as_strided((1, n, n + m),
1484
+ (n * (2 * n + m), 2 * n + m + 1, 1))
1485
+ # Sums the resulting matrix's columns
1486
+ sum_diags = torch.sum(mat_strided, 1)
1487
+ return sum_diags[:, 1:] # drop left bottom corner
1488
+
1489
+
1490
+def _get_block(block_table: torch.Tensor, block_size: int, begin: int,
1491
+ end: int):
1492
+ begin_block = begin // block_size
1493
+ end_block = (end - 1) // block_size + 1
1494
+ return block_table[begin_block:end_block]
vllm/config.py
+19
@@ -929,6 +929,23 @@ class ModelConfig:
929
"Number of experts in the model must be greater than 0 "
930
"when expert parallelism is enabled.")
931
932
+ def verify_dual_chunk_attention_config(
933
+ self,
934
+ load_config: "LoadConfig",
935
+ ) -> None:
936
+ if hasattr(self.hf_config, "dual_chunk_attention_config"):
937
+ # Try loading the sparse attention config
938
+ from vllm.model_executor.model_loader.weight_utils import (
939
+ get_sparse_attention_config)
940
+ sparse_attn_config = get_sparse_attention_config(self, load_config)
941
+ if sparse_attn_config:
942
+ self.hf_config.dual_chunk_attention_config[
943
+ "sparse_attention_config"] = sparse_attn_config
944
+ if "sparse_attention_enabled" not in \
945
+ self.hf_config.dual_chunk_attention_config:
946
+ self.hf_config.dual_chunk_attention_config[
947
+ "sparse_attention_enabled"] = True
948
+
949
def verify_async_output_proc(self, parallel_config, speculative_config,
950
device_config) -> None:
951
if not self.use_async_output_proc:
@@ -4187,6 +4204,8 @@ class VllmConfig:
4204
self.speculative_config,
4205
self.device_config)
4206
self.model_config.verify_with_parallel_config(self.parallel_config)
4207
+ self.model_config.verify_dual_chunk_attention_config(
4208
+ self.load_config)
4209
4210
if self.cache_config is not None:
4211
self.cache_config.verify_with_parallel_config(self.parallel_config)
vllm/engine/arg_utils.py
+13
-2
@@ -37,8 +37,8 @@ from vllm.reasoning import ReasoningParserManager
37
from vllm.test_utils import MODEL_WEIGHTS_S3_BUCKET, MODELS_ON_S3
38
from vllm.transformers_utils.utils import check_gguf_file
39
from vllm.usage.usage_lib import UsageContext
40
-from vllm.utils import (FlexibleArgumentParser, GiB_bytes, is_in_doc_build,
41
- is_in_ray_actor)
40
+from vllm.utils import (STR_DUAL_CHUNK_FLASH_ATTN_VAL, FlexibleArgumentParser,
41
+ GiB_bytes, is_in_doc_build, is_in_ray_actor)
42
43
# yapf: enable
44
@@ -983,6 +983,17 @@ class EngineArgs:
983
984
assert self.enable_chunked_prefill is not None
985
986
+ if envs.VLLM_ATTENTION_BACKEND in [STR_DUAL_CHUNK_FLASH_ATTN_VAL]:
987
+ assert self.enforce_eager, (
988
+ "Cuda graph is not supported with DualChunkFlashAttention. "
989
+ "To run the model in eager mode, set 'enforce_eager=True' "
990
+ "or use '--enforce-eager' in the CLI.")
991
+ assert current_platform.is_cuda(), (
992
+ "DualChunkFlashAttention is only supported on CUDA platform.")
993
+ assert not use_v1, (
994
+ "DualChunkFlashAttention is not supported on V1 engine. "
995
+ "To run the model in V0 engine, try set 'VLLM_USE_V1=0'")
996
+
997
cache_config = CacheConfig(
998
block_size=self.block_size,
999
gpu_memory_utilization=self.gpu_memory_utilization,
vllm/model_executor/layers/rotary_embedding.py
+202
-2
@@ -1486,6 +1486,184 @@ class MRotaryEmbedding(RotaryEmbedding):
1486
return updates
1487
1488
1489
+@CustomOp.register("dual_chunk_rotary_embedding")
1490
+class DualChunkRotaryEmbedding(CustomOp):
1491
+ """Rotary positional embedding for Dual Chunk Attention."""
1492
+
1493
+ def __init__(
1494
+ self,
1495
+ head_size: int,
1496
+ rotary_dim: int,
1497
+ max_position_embeddings: int,
1498
+ base: int,
1499
+ is_neox_style: bool,
1500
+ dtype: torch.dtype,
1501
+ chunk_size: int,
1502
+ local_size: int,
1503
+ ) -> None:
1504
+ super().__init__()
1505
+ self.head_size = head_size
1506
+ self.rotary_dim = rotary_dim
1507
+ self.max_position_embeddings = max_position_embeddings
1508
+ self.base = base
1509
+ self.is_neox_style = is_neox_style
1510
+ self.chunk_size = chunk_size
1511
+ self.local_size = local_size
1512
+ self.dtype = dtype
1513
+ self.device = torch.device(f"cuda:{torch.cuda.current_device()}")
1514
+ (q_cache, qc_cache, k_cache, qc_no_clamp_cache,
1515
+ q_inter_cache) = self._compute_cos_sin_cache()
1516
+
1517
+ self.register_buffer("cos_sin_q_cache", q_cache, persistent=False)
1518
+ self.register_buffer("cos_sin_qc_cache", qc_cache, persistent=False)
1519
+ self.register_buffer("cos_sin_k_cache", k_cache, persistent=False)
1520
+ self.register_buffer("cos_sin_qc_no_clamp_cache",
1521
+ qc_no_clamp_cache,
1522
+ persistent=False)
1523
+ self.register_buffer("cos_sin_q_inter_cache",
1524
+ q_inter_cache,
1525
+ persistent=False)
1526
+
1527
+ def _compute_inv_freq(self, base: Union[int, float]) -> torch.Tensor:
1528
+ """Compute the inverse frequency."""
1529
+ # NOTE(woosuk): The HF implementation uses `torch.arange(...).float()`.
1530
+ # However, we use `torch.arange(..., dtype=torch.float)` instead to
1531
+ # avoid numerical issues with large base values (e.g., 10000000).
1532
+ # This may cause a slight numerical difference between the HF
1533
+ # implementation and ours.
1534
+ # NOTE(woosuk): To exactly match the HF implementation, we need to
1535
+ # use CPU to compute the cache and then move it to GPU. However, we
1536
+ # create the cache on GPU for faster initialization. This may cause
1537
+ # a slight numerical difference between the HF implementation and ours.
1538
+ inv_freq = 1.0 / (base**(torch.arange(
1539
+ 0, self.rotary_dim, 2, dtype=torch.float) / self.rotary_dim))
1540
+ return inv_freq
1541
+
1542
+ def _compute_cos_sin_cache(self) -> torch.Tensor:
1543
+ """Compute the cos and sin cache."""
1544
+ inv_freq = self._compute_inv_freq(self.base)
1545
+ chunk_len = self.chunk_size - self.local_size
1546
+ q_t = torch.arange(chunk_len, dtype=torch.float)
1547
+ qc_t = (torch.arange(chunk_len, dtype=torch.float) +
1548
+ chunk_len).clamp(max=self.chunk_size)
1549
+ k_t = torch.arange(self.max_position_embeddings,
1550
+ dtype=torch.float) % chunk_len
1551
+
1552
+ # count from chunk_len, no clamp(self.chunk_size) restriction
1553
+ qc_no_clamp_t = torch.arange(chunk_len, dtype=torch.float) + chunk_len
1554
+ # count from self.chunk_size for q_inter's rope
1555
+ q_inter_t = torch.arange(chunk_len,
1556
+ dtype=torch.float) + self.chunk_size
1557
+
1558
+ q_freqs = torch.outer(q_t, inv_freq)
1559
+ qc_freqs = torch.outer(qc_t, inv_freq)
1560
+ k_freqs = torch.outer(k_t, inv_freq)
1561
+ qc_no_clamp_freqs = torch.outer(qc_no_clamp_t, inv_freq)
1562
+ q_inter_freqs = torch.outer(q_inter_t, inv_freq)
1563
+
1564
+ q_cos = q_freqs.cos()
1565
+ q_sin = q_freqs.sin()
1566
+ qc_cos = qc_freqs.cos()
1567
+ qc_sin = qc_freqs.sin()
1568
+ k_cos = k_freqs.cos()
1569
+ k_sin = k_freqs.sin()
1570
+
1571
+ qc_no_clamp_cos = qc_no_clamp_freqs.cos()
1572
+ qc_no_clamp_sin = qc_no_clamp_freqs.sin()
1573
+ q_inter_cos = q_inter_freqs.cos()
1574
+ q_inter_sin = q_inter_freqs.sin()
1575
+
1576
+ q_cache = torch.cat((q_cos, q_sin), dim=-1).to(dtype=self.dtype,
1577
+ device=self.device)
1578
+ qc_cache = torch.cat((qc_cos, qc_sin), dim=-1).to(dtype=self.dtype,
1579
+ device=self.device)
1580
+ k_cache = torch.cat((k_cos, k_sin), dim=-1).to(dtype=self.dtype,
1581
+ device=self.device)
1582
+ qc_no_clamp_cache = torch.cat((qc_no_clamp_cos, qc_no_clamp_sin),
1583
+ dim=-1).to(dtype=self.dtype,
1584
+ device=self.device)
1585
+ q_inter_cache = torch.cat((q_inter_cos, q_inter_sin),
1586
+ dim=-1).to(dtype=self.dtype,
1587
+ device=self.device)
1588
+ return q_cache, qc_cache, k_cache, qc_no_clamp_cache, q_inter_cache
1589
+
1590
+ def forward(
1591
+ self,
1592
+ positions: torch.Tensor,
1593
+ query: torch.Tensor,
1594
+ key: torch.Tensor,
1595
+ offsets: Optional[torch.Tensor] = None,
1596
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
1597
+ query = query.view(*query.shape[:-1], -1, self.head_size)
1598
+ key = key.view(*key.shape[:-1], -1, self.head_size)
1599
+ query_rot = query[..., :self.rotary_dim]
1600
+ key_rot = key[..., :self.rotary_dim]
1601
+ if self.rotary_dim < self.head_size:
1602
+ query_pass = query[..., self.rotary_dim:]
1603
+ key_pass = key[..., self.rotary_dim:]
1604
+ else:
1605
+ query_pass = None
1606
+ key_pass = None
1607
+
1608
+ positions_with_offsets = (torch.add(positions, offsets)
1609
+ if offsets is not None else positions)
1610
+ key = self._apply_rotary_embedding(
1611
+ self.cos_sin_k_cache[positions_with_offsets], key_rot, key_pass)
1612
+ chunk_len = self.chunk_size - self.local_size
1613
+ query = self._apply_rotary_embedding(
1614
+ self.cos_sin_q_cache[positions_with_offsets % chunk_len],
1615
+ query_rot, query_pass)
1616
+ query_succ = self._apply_rotary_embedding(
1617
+ self.cos_sin_qc_cache[positions_with_offsets % chunk_len],
1618
+ query_rot, query_pass)
1619
+ query_inter = self._apply_rotary_embedding(
1620
+ self.cos_sin_qc_cache[chunk_len - 1].repeat(positions.shape[0], 1),
1621
+ query_rot, query_pass)
1622
+ query_succ_critical = self._apply_rotary_embedding(
1623
+ self.cos_sin_qc_no_clamp_cache[positions_with_offsets % chunk_len],
1624
+ query_rot, query_pass)
1625
+ query_inter_critical = self._apply_rotary_embedding(
1626
+ self.cos_sin_q_inter_cache[positions_with_offsets % chunk_len],
1627
+ query_rot, query_pass)
1628
+
1629
+ # merge query into one tensor to simplify the interfaces
1630
+ query = torch.cat((
1631
+ query,
1632
+ query_succ,
1633
+ query_inter,
1634
+ query_succ_critical,
1635
+ query_inter_critical,
1636
+ ),
1637
+ dim=-1)
1638
+ return query, key
1639
+
1640
+ def _apply_rotary_embedding(self, cos_sin, hidden_rot, hidden_pass):
1641
+ cos, sin = cos_sin.chunk(2, dim=-1)
1642
+ if self.is_neox_style:
1643
+ # NOTE(woosuk): Here we assume that the positions tensor has the
1644
+ # shape [batch_size, seq_len].
1645
+ cos = cos.repeat(1, 1, 2).unsqueeze(-2)
1646
+ sin = sin.repeat(1, 1, 2).unsqueeze(-2)
1647
+ else:
1648
+ cos = cos.repeat_interleave(2, dim=-1).unsqueeze(-2)
1649
+ sin = sin.repeat_interleave(2, dim=-1).unsqueeze(-2)
1650
+ rotate_fn = _rotate_neox if self.is_neox_style else _rotate_gptj
1651
+ hidden_rot = hidden_rot * cos + rotate_fn(hidden_rot) * sin
1652
+
1653
+ if self.rotary_dim < self.head_size:
1654
+ hidden = torch.cat((hidden_rot, hidden_pass), dim=-1)
1655
+ else:
1656
+ hidden = hidden_rot
1657
+ return hidden.flatten(-2).squeeze(0)
1658
+
1659
+ def extra_repr(self) -> str:
1660
+ s = f"head_size={self.head_size}, rotary_dim={self.rotary_dim}"
1661
+ s += f", max_position_embeddings={self.max_position_embeddings}"
1662
+ s += f", base={self.base}, is_neox_style={self.is_neox_style}"
1663
+ s += f", chunk_size={self.chunk_size}, local_size={self.local_size}"
1664
+ return s
1665
+
1666
+
1667
_ROPE_DICT: Dict[Tuple, RotaryEmbedding] = {}
1668
1669
@@ -1498,6 +1676,7 @@ def get_rope(
1676
rope_scaling: Optional[Dict[str, Any]] = None,
1677
dtype: Optional[torch.dtype] = None,
1678
partial_rotary_factor: float = 1.0,
1679
+ dual_chunk_attention_config: Optional[Dict[str, Any]] = None,
1680
) -> RotaryEmbedding:
1681
if dtype is None:
1682
dtype = torch.get_default_dtype()
@@ -1510,14 +1689,35 @@ def get_rope(
1689
rope_scaling_args = tuple(rope_scaling_tuple.items())
1690
else:
1691
rope_scaling_args = None
1692
+
1693
+ if dual_chunk_attention_config is not None:
1694
+ dual_chunk_attention_tuple = {
1695
+ k: tuple(v) if isinstance(v, list) else v
1696
+ for k, v in dual_chunk_attention_config.items()
1697
+ if k != "sparse_attention_config"
1698
+ }
1699
+ dual_chunk_attention_args = tuple(dual_chunk_attention_tuple.items())
1700
+ else:
1701
+ dual_chunk_attention_args = None
1702
+
1703
if partial_rotary_factor < 1.0:
1704
rotary_dim = int(rotary_dim * partial_rotary_factor)
1705
key = (head_size, rotary_dim, max_position, base, is_neox_style,
1516
- rope_scaling_args, dtype)
1706
+ rope_scaling_args, dual_chunk_attention_args, dtype)
1707
if key in _ROPE_DICT:
1708
return _ROPE_DICT[key]
1709
1520
- if not rope_scaling:
1710
+ if dual_chunk_attention_config is not None:
1711
+ extra_kwargs = {
1712
+ k: v
1713
+ for k, v in dual_chunk_attention_config.items()
1714
+ if k in ("chunk_size", "local_size")
1715
+ }
1716
+ rotary_emb = DualChunkRotaryEmbedding(head_size, rotary_dim,
1717
+ max_position, base,
1718
+ is_neox_style, dtype,
1719
+ **extra_kwargs)
1720
+ elif not rope_scaling:
1721
rotary_emb = RotaryEmbedding(head_size, rotary_dim, max_position, base,
1722
is_neox_style, dtype)
1723
else:
vllm/model_executor/model_loader/weight_utils.py
+33
@@ -217,6 +217,39 @@ def get_quant_config(model_config: ModelConfig,
217
return quant_cls.from_config(config)
218
219
220
+def get_sparse_attention_config(
221
+ model_config: ModelConfig,
222
+ load_config: LoadConfig,
223
+ sparse_attention_config_filename: str = "sparse_attention_config.json",
224
+) -> Dict[str, Any]:
225
+ model_name_or_path = model_config.model
226
+ is_local = os.path.isdir(model_name_or_path)
227
+ if not is_local:
228
+ # Download the config files.
229
+ with get_lock(model_name_or_path, load_config.download_dir):
230
+ hf_folder = snapshot_download(
231
+ model_name_or_path,
232
+ revision=model_config.revision,
233
+ allow_patterns="*.json",
234
+ cache_dir=load_config.download_dir,
235
+ local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE,
236
+ tqdm_class=DisabledTqdm,
237
+ )
238
+ else:
239
+ hf_folder = model_name_or_path
240
+
241
+ config_file = os.path.join(hf_folder, sparse_attention_config_filename)
242
+ if not os.path.exists(config_file):
243
+ return {}
244
+
245
+ # Load the sparse attention config.
246
+ with open(config_file) as f:
247
+ config = json.load(f)
248
+ logger.info("Loaded sparse attention config from %s", config_file)
249
+
250
+ return config
251
+
252
+
253
def download_weights_from_hf(
254
model_name_or_path: str,
255
cache_dir: Optional[str],
vllm/model_executor/models/qwen2.py
+35
-21
@@ -23,7 +23,7 @@
23
# See the License for the specific language governing permissions and
24
# limitations under the License.
25
"""Inference-only Qwen2 model compatible with HuggingFace weights."""
26
-from typing import Iterable, Optional, Set, Tuple, Union
26
+from typing import Any, Iterable, Optional, Set, Tuple, Union
27
28
import torch
29
from torch import nn
@@ -53,7 +53,7 @@ from vllm.sequence import IntermediateTensors, PoolerOutput
53
54
from .interfaces import SupportsLoRA, SupportsPP
55
from .utils import (AutoWeightsLoader, PPMissingLayer, WeightsMapper,
56
- is_pp_missing_parameter,
56
+ extract_layer_index, is_pp_missing_parameter,
57
make_empty_intermediate_tensors_factory, make_layers,
58
maybe_prefix)
59
@@ -99,17 +99,20 @@ class Qwen2MLP(nn.Module):
99
100
class Qwen2Attention(nn.Module):
101
102
- def __init__(self,
103
- hidden_size: int,
104
- num_heads: int,
105
- num_kv_heads: int,
106
- max_position: int = 4096 * 32,
107
- rope_theta: float = 10000,
108
- cache_config: Optional[CacheConfig] = None,
109
- quant_config: Optional[QuantizationConfig] = None,
110
- rope_scaling: Optional[Tuple] = None,
111
- prefix: str = "",
112
- attn_type: str = AttentionType.DECODER) -> None:
102
+ def __init__(
103
+ self,
104
+ hidden_size: int,
105
+ num_heads: int,
106
+ num_kv_heads: int,
107
+ max_position: int = 4096 * 32,
108
+ rope_theta: float = 10000,
109
+ cache_config: Optional[CacheConfig] = None,
110
+ quant_config: Optional[QuantizationConfig] = None,
111
+ rope_scaling: Optional[Tuple] = None,
112
+ prefix: str = "",
113
+ attn_type: str = AttentionType.DECODER,
114
+ dual_chunk_attention_config: Optional[dict[str,
115
+ Any]] = None) -> None:
116
super().__init__()
117
self.hidden_size = hidden_size
118
tp_size = get_tensor_model_parallel_world_size()
@@ -131,6 +134,7 @@ class Qwen2Attention(nn.Module):
134
self.kv_size = self.num_kv_heads * self.head_dim
135
self.scaling = self.head_dim**-0.5
136
self.rope_theta = rope_theta
137
+ self.dual_chunk_attention_config = dual_chunk_attention_config
138
139
self.qkv_proj = QKVParallelLinear(
140
hidden_size,
@@ -155,15 +159,21 @@ class Qwen2Attention(nn.Module):
159
max_position=max_position,
160
base=self.rope_theta,
161
rope_scaling=rope_scaling,
162
+ dual_chunk_attention_config=dual_chunk_attention_config,
163
)
159
- self.attn = Attention(self.num_heads,
160
- self.head_dim,
161
- self.scaling,
162
- num_kv_heads=self.num_kv_heads,
163
- cache_config=cache_config,
164
- quant_config=quant_config,
165
- prefix=f"{prefix}.attn",
166
- attn_type=attn_type)
164
+ self.attn = Attention(
165
+ self.num_heads,
166
+ self.head_dim,
167
+ self.scaling,
168
+ num_kv_heads=self.num_kv_heads,
169
+ cache_config=cache_config,
170
+ quant_config=quant_config,
171
+ attn_type=attn_type,
172
+ prefix=f"{prefix}.attn",
173
+ **{
174
+ "layer_idx": extract_layer_index(prefix),
175
+ "dual_chunk_attention_config": dual_chunk_attention_config,
176
+ } if dual_chunk_attention_config else {})
177
178
def forward(
179
self,
@@ -192,6 +202,9 @@ class Qwen2DecoderLayer(nn.Module):
202
# Requires transformers > 4.32.0
203
rope_theta = getattr(config, "rope_theta", 1000000)
204
rope_scaling = getattr(config, "rope_scaling", None)
205
+ dual_chunk_attention_config = getattr(config,
206
+ "dual_chunk_attention_config",
207
+ None)
208
209
# By default, Qwen2 uses causal attention as it is a decoder-only model.
210
# You can override the HF config with `is_causal=False` to enable
@@ -213,6 +226,7 @@ class Qwen2DecoderLayer(nn.Module):
226
rope_scaling=rope_scaling,
227
prefix=f"{prefix}.self_attn",
228
attn_type=attn_type,
229
+ dual_chunk_attention_config=dual_chunk_attention_config,
230
)
231
self.mlp = Qwen2MLP(
232
hidden_size=self.hidden_size,
vllm/model_executor/models/qwen2_moe.py
+19
-7
@@ -175,6 +175,7 @@ class Qwen2MoeAttention(nn.Module):
175
cache_config: Optional[CacheConfig] = None,
176
quant_config: Optional[QuantizationConfig] = None,
177
prefix: str = "",
178
+ dual_chunk_attention_config: Optional[Dict[str, Any]] = None,
179
) -> None:
180
super().__init__()
181
self.hidden_size = hidden_size
@@ -198,6 +199,7 @@ class Qwen2MoeAttention(nn.Module):
199
self.scaling = self.head_dim**-0.5
200
self.rope_theta = rope_theta
201
self.max_position_embeddings = max_position_embeddings
202
+ self.dual_chunk_attention_config = dual_chunk_attention_config
203
204
self.qkv_proj = QKVParallelLinear(
205
hidden_size,
@@ -221,14 +223,20 @@ class Qwen2MoeAttention(nn.Module):
223
max_position=max_position_embeddings,
224
base=rope_theta,
225
rope_scaling=rope_scaling,
226
+ dual_chunk_attention_config=dual_chunk_attention_config,
227
)
225
- self.attn = Attention(self.num_heads,
226
- self.head_dim,
227
- self.scaling,
228
- num_kv_heads=self.num_kv_heads,
229
- cache_config=cache_config,
230
- quant_config=quant_config,
231
- prefix=f"{prefix}.attn")
228
+ self.attn = Attention(
229
+ self.num_heads,
230
+ self.head_dim,
231
+ self.scaling,
232
+ num_kv_heads=self.num_kv_heads,
233
+ cache_config=cache_config,
234
+ quant_config=quant_config,
235
+ prefix=f"{prefix}.attn",
236
+ **{
237
+ "layer_idx": extract_layer_index(prefix),
238
+ "dual_chunk_attention_config": dual_chunk_attention_config,
239
+ } if dual_chunk_attention_config else {})
240
241
def forward(
242
self,
@@ -256,6 +264,9 @@ class Qwen2MoeDecoderLayer(nn.Module):
264
self.hidden_size = config.hidden_size
265
rope_theta = getattr(config, "rope_theta", 10000)
266
rope_scaling = getattr(config, "rope_scaling", None)
267
+ dual_chunk_attention_config = getattr(config,
268
+ "dual_chunk_attention_config",
269
+ None)
270
max_position_embeddings = getattr(config, "max_position_embeddings",
271
8192)
272
self.self_attn = Qwen2MoeAttention(
@@ -268,6 +279,7 @@ class Qwen2MoeDecoderLayer(nn.Module):
279
cache_config=cache_config,
280
quant_config=quant_config,
281
prefix=f"{prefix}.self_attn",
282
+ dual_chunk_attention_config=dual_chunk_attention_config,
283
)
284
285
# Note: Qwen/Qwen2-57B-A14B-Instruct does not have
vllm/platforms/cuda.py
+4
@@ -222,6 +222,10 @@ class CudaPlatformBase(Platform):
222
elif selected_backend == _Backend.XFORMERS:
223
logger.info("Using XFormers backend.")
224
return "vllm.attention.backends.xformers.XFormersBackend"
225
+ elif selected_backend == _Backend.DUAL_CHUNK_FLASH_ATTN:
226
+ logger.info("Using DualChunkFlashAttention backend.")
227
+ return ("vllm.attention.backends.dual_chunk_flash_attn."
228
+ "DualChunkFlashAttentionBackend")
229
elif selected_backend == _Backend.FLASH_ATTN:
230
pass
231
elif selected_backend:
vllm/platforms/interface.py
+1
@@ -51,6 +51,7 @@ class _Backend(enum.Enum):
51
PALLAS_VLLM_V1 = enum.auto()
52
IPEX = enum.auto()
53
BLOCK_SPARSE_FLASH_ATTN = enum.auto()
54
+ DUAL_CHUNK_FLASH_ATTN = enum.auto()
55
NO_ATTENTION = enum.auto()
56
57
vllm/utils.py
+1
@@ -153,6 +153,7 @@ STR_TORCH_SDPA_ATTN_VAL: str = "TORCH_SDPA"
153
STR_ROCM_FLASH_ATTN_VAL: str = "ROCM_FLASH"
154
STR_XFORMERS_ATTN_VAL: str = "XFORMERS"
155
STR_FLASH_ATTN_VAL: str = "FLASH_ATTN"
156
+STR_DUAL_CHUNK_FLASH_ATTN_VAL: str = "DUAL_CHUNK_FLASH_ATTN"
157
STR_INVALID_VAL: str = "INVALID"
158
159
GB_bytes = 1_000_000_000
vllm/worker/model_runner.py
+12
@@ -204,6 +204,7 @@ class ModelInputForGPUBuilder(ModelRunnerInputBuilderBase[ModelInputForGPU]):
204
self.mrope_input_positions = None # type: ignore
205
self.seq_lens[0] = 0 # type: ignore
206
self.orig_seq_lens[0] = 0 # type: ignore
207
+ self.prompt_lens[0] = 0 # type: ignore
208
self.query_lens[0] = 0 # type: ignore
209
self.context_lens[0] = 0 # type: ignore
210
self.curr_sliding_window_blocks[0] = 0 # type: ignore
@@ -236,6 +237,8 @@ class ModelInputForGPUBuilder(ModelRunnerInputBuilderBase[ModelInputForGPU]):
237
# The original sequence length (before applying sliding window).
238
# This is used to compute slot mapping.
239
orig_seq_lens: Optional[List[int]] = None,
240
+ # This is used in the dual-chunk flash attention backend.
241
+ prompt_lens: Optional[List[int]] = None,
242
# The query length.
243
query_lens: Optional[List[int]] = None,
244
# The number of tokens that are already computed.
@@ -316,6 +319,12 @@ class ModelInputForGPUBuilder(ModelRunnerInputBuilderBase[ModelInputForGPU]):
319
for seq_id in range(len(self.seq_ids)):
320
self.orig_seq_lens[seq_id] = 0
321
322
+ if prompt_lens:
323
+ self.prompt_lens = prompt_lens
324
+ else:
325
+ for seq_id in range(len(self.seq_ids)):
326
+ self.prompt_lens[seq_id] = 0
327
+
328
if query_lens:
329
self.query_lens = query_lens
330
else:
@@ -370,6 +379,7 @@ class ModelInputForGPUBuilder(ModelRunnerInputBuilderBase[ModelInputForGPU]):
379
self.mrope_input_positions = mrope_input_positions or None
380
self.seq_lens = seq_lens or []
381
self.orig_seq_lens = orig_seq_lens or []
382
+ self.prompt_lens = prompt_lens or []
383
self.query_lens = query_lens or []
384
self.context_lens = context_lens or []
385
self.curr_sliding_window_blocks = \
@@ -403,6 +413,7 @@ class ModelInputForGPUBuilder(ModelRunnerInputBuilderBase[ModelInputForGPU]):
413
self.mrope_input_positions = None
414
self.seq_lens = [0] * self.n_seqs
415
self.orig_seq_lens = [0] * self.n_seqs
416
+ self.prompt_lens = [0] * self.n_seqs
417
self.query_lens = [0] * self.n_seqs
418
self.context_lens = [0] * self.n_seqs
419
self.curr_sliding_window_blocks = [0] * self.n_seqs
@@ -552,6 +563,7 @@ class ModelInputForGPUBuilder(ModelRunnerInputBuilderBase[ModelInputForGPU]):
563
564
inter_data.seq_lens[seq_idx] = seq_len
565
inter_data.orig_seq_lens[seq_idx] = seq_len
566
+ inter_data.prompt_lens[seq_idx] = seq_data.get_prompt_len()
567
inter_data.context_lens[seq_idx] = context_len
568
inter_data.input_tokens[seq_idx].extend(tokens)
569
inter_data.inputs_embeds = prompt_embeds