@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
md 498 lines 20.2 KB
Rendered Raw
1 # Paged Attention
2
3 !!! warning
4 This is a historical document based on the [original paper for vLLM](https://arxiv.org/abs/2309.06180).
5 It no longer describes the code used in vLLM today.
6
7 Currently, vLLM utilizes its own implementation of a multi-head query
8 attention kernel (`csrc/attention/attention_kernels.cu`).
9 This kernel is designed to be compatible with
10 vLLM's paged KV caches, where the key and value cache are stored in
11 separate blocks (note that this block concept differs from the GPU
12 thread block. So in a later document, I will refer to vLLM paged
13 attention block as "block", while refer to GPU thread block as
14 "thread block").
15
16 To achieve high performance, this kernel relies on a specially
17 designed memory layout and access method, specifically when threads
18 read data from global memory to shared memory. The purpose of this
19 document is to provide a high-level explanation of the kernel
20 implementation step by step, aiding those who wish to learn about the
21 vLLM multi-head query attention kernel. After going through this
22 document, users will likely have a better understanding and feel easier
23 to follow the actual implementation.
24
25 Please note that this document may not cover all details, such as how
26 to calculate the correct index for the corresponding data or the dot
27 multiplication implementation. However, after reading this document
28 and becoming familiar with the high-level logic flow, it should be
29 easier for you to read the actual code and understand the details.
30
31 ## Inputs
32
33 The kernel function takes a list of arguments for the current thread
34 to perform its assigned work. The three most important arguments are
35 the input pointers `q`, `k_cache`, and `v_cache`, which point
36 to query, key, and value data on global memory that need to be read
37 and processed. The output pointer `out` points to global memory
38 where the result should be written. These four pointers actually
39 refer to multidimensional arrays, but each thread only accesses the
40 portion of data assigned to it. I have omitted all other runtime
41 parameters here for simplicity.
42
43 ```cpp
44 template<typename scalar_t, int HEAD_SIZE, int BLOCK_SIZE, int NUM_THREADS, int PARTITION_SIZE = 0>
45 __device__ void paged_attention_kernel(
46 ... // Other side args.
47 const scalar_t* __restrict__ out, // [num_seqs, num_heads, max_num_partitions, head_size]
48 const scalar_t* __restrict__ q, // [num_seqs, num_heads, head_size]
49 const scalar_t* __restrict__ k_cache, // [num_blocks, num_kv_heads, head_size/x, block_size, x]
50 const scalar_t* __restrict__ v_cache, // [num_blocks, num_kv_heads, head_size, block_size]
51 ... // Other side args.
52 )
53 ```
54
55 There is also a list of template arguments above the function
56 signature that are determined during compilation time. `scalar_t`
57 represents the data type of the query, key, and value data elements,
58 such as FP16. `HEAD_SIZE` indicates the number of elements in each
59 head. `BLOCK_SIZE` refers to the number of tokens in each block.
60 `NUM_THREADS` denotes the number of threads in each thread block.
61 `PARTITION_SIZE` represents the number of tensor parallel GPUs (For
62 simplicity, we assume this is 0 and tensor parallel is disabled).
63
64 With these arguments, we need to perform a sequence of preparations.
65 This includes calculating the current head index, block index, and
66 other necessary variables. However, for now, we can ignore these
67 preparations and proceed directly to the actual calculations. It will
68 be easier to understand them once we grasp the entire flow.
69
70 ## Concepts
71
72 Just before we dive into the calculation flow, I want to describe a
73 few concepts that are needed for later sections. However, you may
74 skip this section and return later if you encounter any confusing
75 terminologies.
76
77 - **Sequence**: A sequence represents a client request. For example,
78 the data pointed to by `q` has a shape of
79 `[num_seqs, num_heads, head_size]`. That represents there are total
80 `num_seqs` of query sequence data are pointed by `q`. Since this
81 kernel is a single query attention kernel, each sequence only has one
82 query token. Hence, the `num_seqs` equals the total number of tokens
83 that are processed in the batch.
84 - **Context**: The context consists of the generated tokens from the
85 sequence. For instance, `["What", "is", "your"]` are the context
86 tokens, and the input query token is `"name"`. The model might
87 generate the token `"?"`.
88 - **Vec**: The vec is a list of elements that are fetched and
89 calculated together. For query and key data, the vec size
90 (`VEC_SIZE`) is determined so that each thread group can fetch and
91 calculate 16 bytes of data at a time. For value data, the vec size
92 (`V_VEC_SIZE`) is determined so that each thread can fetch and
93 calculate 16 bytes of data at a time. For example, if the
94 `scalar_t` is FP16 (2 bytes) and `THREAD_GROUP_SIZE` is 2, the
95 `VEC_SIZE` will be 4, while the `V_VEC_SIZE` will be 8.
96 - **Thread group**: The thread group is a small group of
97 threads(`THREAD_GROUP_SIZE`) that fetches and calculates one
98 query token and one key token at a time. Each thread handles only a
99 portion of the token data. The total number of elements processed by
100 one thread group is referred as `x`. For example, if the thread
101 group contains 2 threads and the head size is 8, then thread 0
102 handles the query and key elements at index 0, 2, 4, 6, while thread
103 1 handles the elements at index 1, 3, 5, 7.
104 - **Block**: The key and value cache data in vLLM are split into
105 blocks. Each block stores data for a fixed number(`BLOCK_SIZE`)
106 of tokens at one head. Each block may contain only a portion of the
107 whole context tokens. For example, if the block size is 16 and the
108 head size is 128, then for one head, one block can store 16 * 128 =
109 2048 elements.
110 - **Warp**: A warp is a group of 32 threads(`WARP_SIZE`) that
111 execute simultaneously on a stream multiprocessor (SM). In this
112 kernel, each warp processes the calculation between one query token
113 and key tokens of one entire block at a time (it may process multiple
114 blocks in multiple iterations). For example, if there are 4 warps and
115 6 blocks for one context, the assignment would be like warp 0 handles
116 the 0th, 4th blocks, warp 1 handles the 1st, 5th blocks, warp 2
117 handles the 2nd block and warp 3 handles the 3rd block.
118 - **Thread block**: A thread block is a group of
119 threads(`NUM_THREADS`) that can access the same shared memory.
120 Each thread block contains multiple warps(`NUM_WARPS`), and in
121 this kernel, each thread block processes the calculation between one
122 query token and key tokens of a whole context.
123 - **Grid**: A grid is a collection of thread blocks and defines the
124 shape of the collection. In this kernel, the shape is
125 `(num_heads, num_seqs, max_num_partitions)`. Therefore, each thread
126 block only handles the calculation for one head, one sequence, and
127 one partition.
128
129 ## Query
130
131 This section will introduce how query data is stored in memory and
132 fetched by each thread. As mentioned above, each thread group fetches
133 one query token data, while each thread itself only handles a part of
134 one query token data. Within each warp, every thread group will fetch
135 the same query token data, but will multiply it with different key
136 token data.
137
138 ```cpp
139 const scalar_t* q_ptr = q + seq_idx * q_stride + head_idx * HEAD_SIZE;
140 ```
141
142 ![query](../assets/design/paged_attention/query.png)
143
144 Each thread defines its own `q_ptr` which points to the assigned
145 query token data on global memory. For example, if `VEC_SIZE` is 4
146 and `HEAD_SIZE` is 128, the `q_ptr` points to data that contains
147 total of 128 elements divided into 128 / 4 = 32 vecs.
148
149 ![q_vecs](../assets/design/paged_attention/q_vecs.png)
150
151 ```cpp
152 __shared__ Q_vec q_vecs[THREAD_GROUP_SIZE][NUM_VECS_PER_THREAD];
153 ```
154
155 Next, we need to read the global memory data pointed to by `q_ptr`
156 into shared memory as `q_vecs`. It is important to note that each
157 vecs is assigned to a different row. For example, if the
158 `THREAD_GROUP_SIZE` is 2, thread 0 will handle the 0th row vecs,
159 while thread 1 handles the 1st row vecs. By reading the query data in
160 this way, neighboring threads like thread 0 and thread 1 can read
161 neighbor memory, achieving the memory coalescing to improve
162 performance.
163
164 ## Key
165
166 Similar to the "Query" section, this section introduces memory layout
167 and assignment for keys. While each thread group only handle one
168 query token one kernel run, it may handle multiple key tokens across
169 multiple iterations. Meanwhile, each warp will process multiple blocks
170 of key tokens in multiple iterations, ensuring that all context
171 tokens are processed by the entire thread group after the kernel run.
172 In this context, "handle" refers to performing the dot multiplication
173 between query data and key data.
174
175 ```cpp
176 const scalar_t* k_ptr = k_cache + physical_block_number * kv_block_stride
177 + kv_head_idx * kv_head_stride
178 + physical_block_offset * x;
179 ```
180
181 Unlike `q_ptr`, `k_ptr` in each thread will point to different
182 key token at different iterations. As shown above, that `k_ptr`
183 points to key token data based on `k_cache` at assigned block,
184 assigned head and assigned token.
185
186 ![key](../assets/design/paged_attention/key.png)
187
188 The diagram above illustrates the memory layout for key data. It
189 assumes that the `BLOCK_SIZE` is 16, `HEAD_SIZE` is 128, `x` is
190 8, `THREAD_GROUP_SIZE` is 2, and there are a total of 4 warps. Each
191 rectangle represents all the elements for one key token at one head,
192 which will be processed by one thread group. The left half shows the
193 total 16 blocks of key token data for warp 0, while the right half
194 represents the remaining key token data for other warps or
195 iterations. Inside each rectangle, there are a total 32 vecs (128
196 elements for one token) that will be processed by 2 threads (one
197 thread group) separately.
198
199 ![k_vecs](../assets/design/paged_attention/k_vecs.png)
200
201 ```cpp
202 K_vec k_vecs[NUM_VECS_PER_THREAD]
203 ```
204
205 Next, we need to read the key token data from `k_ptr` and store
206 them on register memory as `k_vecs`. We use register memory for
207 `k_vecs` because it will only be accessed by one thread once,
208 whereas `q_vecs` will be accessed by multiple threads multiple
209 times. Each `k_vecs` will contain multiple vectors for later
210 calculation. Each vec will be set at each inner iteration. The
211 assignment of vecs allows neighboring threads in a warp to read
212 neighboring memory together, which again promotes the memory
213 coalescing. For instance, thread 0 will read vec 0, while thread 1
214 will read vec 1. In the next inner loop, thread 0 will read vec 2,
215 while thread 1 will read vec 3, and so on.
216
217 You may still be a little confused about the overall flow. Don't
218 worry, please keep reading the next "QK" section. It will illustrate
219 the query and key calculation flow in a clearer and higher-level
220 manner.
221
222 ## QK
223
224 As shown the pseudocode below, before the entire for loop block, we
225 fetch the query data for one token and store it in `q_vecs`. Then,
226 in the outer for loop, we iterate through different `k_ptrs` that
227 point to different tokens and prepare the `k_vecs` in the inner for
228 loop. Finally, we perform the dot multiplication between the
229 `q_vecs` and each `k_vecs`.
230
231 ```cpp
232 q_vecs = ...
233 for ... {
234 k_ptr = ...
235 for ... {
236 k_vecs[i] = ...
237 }
238 ...
239 float qk = scale * Qk_dot<scalar_t, THREAD_GROUP_SIZE>::dot(q_vecs[thread_group_offset], k_vecs);
240 }
241 ```
242
243 As mentioned before, for each thread, it only fetches part of the
244 query and key token data at a time. However, there will be a cross
245 thread group reduction happen in the `Qk_dot<>::dot` . So `qk`
246 returned here is not just between part of the query and key token dot
247 multiplication, but actually a full result between entire query and
248 key token data.
249
250 For example, if the value of `HEAD_SIZE` is 128 and
251 `THREAD_GROUP_SIZE` is 2, each thread's `k_vecs` will contain
252 total 64 elements. However, the returned `qk` is actually the
253 result of dot multiplication between 128 query elements and 128 key
254 elements. If you want to learn more about the details of the dot
255 multiplication and reduction, you may refer to the implementation of
256 `Qk_dot<>::dot`. However, for the sake of simplicity, I will not
257 cover it in this document.
258
259 ## Softmax
260
261 Next, we need to calculate the normalized softmax for all `qk`s,
262 as shown above, where each $x$ represents a `qk`. To do this,
263 we must obtain the reduced value of `qk_max`($m(x)$) and
264 the `exp_sum`($\ell(x)$) of all `qk`s. The reduction
265 should be performed across the entire thread block, encompassing
266 results between the query token and all context key tokens.
267
268 $$
269 \begin{gather*}
270 m(x):=\max _i \quad x_i \\ \quad f(x):=\left[\begin{array}{lll}e^{x_1-m(x)} & \ldots & e^{x_B-m(x)}\end{array}\right]\\ \quad \ell(x):=\sum_i f(x)_i \\
271 \quad \operatorname{softmax}(x):=\frac{f(x)}{\ell(x)}
272 \end{gather*}
273 $$
274
275 ### `qk_max` and `logits`
276
277 Just right after we get the `qk` result, we can set the temporary
278 `logits` result with `qk` (In the end, the `logits` should
279 store the normalized softmax result). Also we can compare and collect
280 the `qk_max` for all `qk`s that are calculated by current
281 thread group.
282
283 ```cpp
284 if (thread_group_offset == 0) {
285 const bool mask = token_idx >= context_len;
286 logits[token_idx - start_token_idx] = mask ? 0.f : qk;
287 qk_max = mask ? qk_max : fmaxf(qk_max, qk);
288 }
289 ```
290
291 Please note that the `logits` here is on shared memory, so each
292 thread group will set the fields for its own assigned context tokens.
293 Overall, the size of logits should be number of context tokens.
294
295 ```cpp
296 for (int mask = WARP_SIZE / 2; mask >= THREAD_GROUP_SIZE; mask /= 2) {
297 qk_max = fmaxf(qk_max, VLLM_SHFL_XOR_SYNC(qk_max, mask));
298 }
299
300 if (lane == 0) {
301 red_smem[warp_idx] = qk_max;
302 }
303 ```
304
305 Then we need to get the reduced `qk_max` across each warp. The main
306 idea is to make threads in warp to communicate with each other and
307 get the final max `qk` .
308
309 ```cpp
310 for (int mask = NUM_WARPS / 2; mask >= 1; mask /= 2) {
311 qk_max = fmaxf(qk_max, VLLM_SHFL_XOR_SYNC(qk_max, mask));
312 }
313 qk_max = VLLM_SHFL_SYNC(qk_max, 0);
314 ```
315
316 Finally, we can get the reduced `qk_max` from whole thread block by
317 compare the `qk_max` from all warps in this thread block. Then we
318 need to broadcast the final result to each thread.
319
320 ### `exp_sum`
321
322 Similar to `qk_max`, we need to get the reduced sum value from the
323 entire thread block too.
324
325 ```cpp
326 for (int i = thread_idx; i < num_tokens; i += NUM_THREADS) {
327 float val = __expf(logits[i] - qk_max);
328 logits[i] = val;
329 exp_sum += val;
330 }
331 ...
332 exp_sum = block_sum<NUM_WARPS>(&red_smem[NUM_WARPS], exp_sum);
333 ```
334
335 Firstly, sum all exp values from each thread group, and meanwhile,
336 convert each entry of `logits` from `qk` to `exp(qk - qk_max)`.
337 Please note, the `qk_max` here is already the max `qk` across the
338 whole thread block. And then we can do reduction for `exp_sum`
339 across whole thread block just like the `qk_max`.
340
341 ```cpp
342 const float inv_sum = __fdividef(1.f, exp_sum + 1e-6f);
343 for (int i = thread_idx; i < num_tokens; i += NUM_THREADS) {
344 logits[i] *= inv_sum;
345 }
346 ```
347
348 Finally, with the reduced `qk_max` and `exp_sum`, we can obtain
349 the final normalized softmax result as `logits`. This `logits`
350 variable will be used for dot multiplication with the value data in
351 later steps. Now, it should store the normalized softmax result of
352 `qk` for all assigned context tokens.
353
354 ## Value
355
356 ![value](../assets/design/paged_attention/value.png)
357
358 ![logits_vec](../assets/design/paged_attention/logits_vec.png)
359
360 ![v_vec](../assets/design/paged_attention/v_vec.png)
361
362 Now we need to retrieve the value data and perform dot multiplication
363 with `logits`. Unlike query and key, there is no thread group
364 concept for value data. As shown in diagram, different from key token
365 memory layout, elements from the same column correspond to the same
366 value token. For one block of value data, there are `HEAD_SIZE` of
367 rows and `BLOCK_SIZE` of columns that are split into multiple
368 `v_vecs`.
369
370 Each thread always fetches `V_VEC_SIZE` elements from the same
371 `V_VEC_SIZE` of tokens at a time. As a result, a single thread
372 retrieves multiple `v_vec`s from different rows and the same
373 columns through multiple inner iterations. For each `v_vec`, it
374 needs to be dot multiplied with the corresponding `logits_vec`,
375 which is also `V_VEC_SIZE` elements from `logits`. Overall, with
376 multiple inner iterations, each warp will process one block of value
377 tokens. And with multiple outer iterations, the whole context value
378 tokens are processed
379
380 ```cpp
381 float accs[NUM_ROWS_PER_THREAD];
382 for ... { // Iteration over different blocks.
383 logits_vec = ...
384 for ... { // Iteration over different rows.
385 v_vec = ...
386 ...
387 accs[i] += dot(logits_vec, v_vec);
388 }
389 }
390 ```
391
392 As shown in the above pseudocode, in the outer loop, similar to
393 `k_ptr`, `logits_vec` iterates over different blocks and reads
394 `V_VEC_SIZE` elements from `logits`. In the inner loop, each
395 thread reads `V_VEC_SIZE` elements from the same tokens as a
396 `v_vec` and performs dot multiplication. It is important to note
397 that in each inner iteration, the thread fetches different head
398 position elements for the same tokens. The dot result is then
399 accumulated in `accs`. Therefore, each entry of `accs` is mapped
400 to a head position assigned to the current thread.
401
402 For example, if `BLOCK_SIZE` is 16 and `V_VEC_SIZE` is 8, each
403 thread fetches 8 value elements for 8 tokens at a time. Each element
404 is from different tokens at the same head position. If `HEAD_SIZE`
405 is 128 and `WARP_SIZE` is 32, for each inner loop, a warp needs to
406 fetch `WARP_SIZE * V_VEC_SIZE = 256` elements. This means there are
407 a total of 128 * 16 / 256 = 8 inner iterations for a warp to handle
408 a whole block of value tokens. And each `accs` in each thread
409 contains 8 elements that accumulated at 8 different head positions.
410 For the thread 0, the `accs` variable will have 8 elements, which
411 are 0th, 32nd … 224th elements of a value head that are accumulated
412 from all assigned 8 tokens.
413
414 ## LV
415
416 Now, we need to perform reduction for `accs` within each warp. This
417 process allows each thread to accumulate the `accs` for the
418 assigned head positions of all tokens in one block.
419
420 ```cpp
421 for (int i = 0; i < NUM_ROWS_PER_THREAD; i++) {
422 float acc = accs[i];
423 for (int mask = NUM_V_VECS_PER_ROW / 2; mask >= 1; mask /= 2) {
424 acc += VLLM_SHFL_XOR_SYNC(acc, mask);
425 }
426 accs[i] = acc;
427 }
428 ```
429
430 Next, we perform reduction for `accs` across all warps, allowing
431 each thread to have the accumulation of `accs` for the assigned
432 head positions of all context tokens. Please note that each `accs`
433 in every thread only stores the accumulation for a portion of
434 elements of the entire head for all context tokens. However, overall,
435 all results for output have been calculated but are just stored in
436 different thread register memory.
437
438 ??? code
439
440 ```cpp
441 float* out_smem = reinterpret_cast<float*>(shared_mem);
442 for (int i = NUM_WARPS; i > 1; i /= 2) {
443 // Upper warps write to shared memory.
444 ...
445 float* dst = &out_smem[(warp_idx - mid) * HEAD_SIZE];
446 for (int i = 0; i < NUM_ROWS_PER_THREAD; i++) {
447 ...
448 dst[row_idx] = accs[i];
449 }
450
451 // Lower warps update the output.
452 const float* src = &out_smem[warp_idx * HEAD_SIZE];
453 for (int i = 0; i < NUM_ROWS_PER_THREAD; i++) {
454 ...
455 accs[i] += src[row_idx];
456 }
457
458 // Write out the accs.
459 }
460 ```
461
462 ## Output
463
464 Now we can write all of calculated result from local register memory
465 to final output global memory.
466
467 ```cpp
468 scalar_t* out_ptr = out + seq_idx * num_heads * max_num_partitions * HEAD_SIZE
469 + head_idx * max_num_partitions * HEAD_SIZE
470 + partition_idx * HEAD_SIZE;
471 ```
472
473 First, we need to define the `out_ptr` variable, which points to
474 the start address of the assigned sequence and assigned head.
475
476 ```cpp
477 for (int i = 0; i < NUM_ROWS_PER_THREAD; i++) {
478 const int row_idx = lane / NUM_V_VECS_PER_ROW + i * NUM_ROWS_PER_ITER;
479 if (row_idx < HEAD_SIZE && lane % NUM_V_VECS_PER_ROW == 0) {
480 from_float(*(out_ptr + row_idx), accs[i]);
481 }
482 }
483 ```
484
485 Finally, we need to iterate over different assigned head positions
486 and write out the corresponding accumulated result based on the
487 `out_ptr`.
488
489 ## Citation
490
491 ```bibtex
492 @inproceedings{kwon2023efficient,
493 title={Efficient Memory Management for Large Language Model Serving with PagedAttention},
494 author={Woosuk Kwon and Zhuohan Li and Siyuan Zhuang and Ying Sheng and Lianmin Zheng and Cody Hao Yu and Joseph E. Gonzalez and Hao Zhang and Ion Stoica},
495 booktitle={Proceedings of the ACM SIGOPS 29th Symposium on Operating Systems Principles},
496 year={2023}
497 }
498 ```