ef1dd6870f · tduncan/Nvidia-Optimization-Stack-VLLM

@@ -3,78 +3,76 @@ title: vLLM Paged Attention

3 ---

4 [](){ #design-paged-attention }

5

6 -- Currently, vLLM utilizes its own implementation of a multi-head query

7 - attention kernel (`csrc/attention/attention_kernels.cu`).

8 - This kernel is designed to be compatible with

9 - vLLM's paged KV caches, where the key and value cache are stored in

10 - separate blocks (note that this block concept differs from the GPU

11 - thread block. So in a later document, I will refer to vLLM paged

12 - attention block as "block", while refer to GPU thread block as

13 - "thread block").

14 -- To achieve high performance, this kernel relies on a specially

15 - designed memory layout and access method, specifically when threads

16 - read data from global memory to shared memory. The purpose of this

17 - document is to provide a high-level explanation of the kernel

18 - implementation step by step, aiding those who wish to learn about the

19 - vLLM multi-head query attention kernel. After going through this

20 - document, users will likely have a better understanding and feel easier

21 - to follow the actual implementation.

22 -- Please note that this document may not cover all details, such as how

23 - to calculate the correct index for the corresponding data or the dot

24 - multiplication implementation. However, after reading this document

25 - and becoming familiar with the high-level logic flow, it should be

26 - easier for you to read the actual code and understand the details.

6 +Currently, vLLM utilizes its own implementation of a multi-head query

7 +attention kernel (`csrc/attention/attention_kernels.cu`).

8 +This kernel is designed to be compatible with

9 +vLLM's paged KV caches, where the key and value cache are stored in

10 +separate blocks (note that this block concept differs from the GPU

11 +thread block. So in a later document, I will refer to vLLM paged

12 +attention block as "block", while refer to GPU thread block as

13 +"thread block").

14 +

15 +To achieve high performance, this kernel relies on a specially

16 +designed memory layout and access method, specifically when threads

17 +read data from global memory to shared memory. The purpose of this

18 +document is to provide a high-level explanation of the kernel

19 +implementation step by step, aiding those who wish to learn about the

20 +vLLM multi-head query attention kernel. After going through this

21 +document, users will likely have a better understanding and feel easier

22 +to follow the actual implementation.

23 +

24 +Please note that this document may not cover all details, such as how

25 +to calculate the correct index for the corresponding data or the dot

26 +multiplication implementation. However, after reading this document

27 +and becoming familiar with the high-level logic flow, it should be

28 +easier for you to read the actual code and understand the details.

29

30 ## Inputs

31

30 -- The kernel function takes a list of arguments for the current thread

31 - to perform its assigned work. The three most important arguments are

32 - the input pointers `q`, `k_cache`, and `v_cache`, which point

33 - to query, key, and value data on global memory that need to be read

34 - and processed. The output pointer `out` points to global memory

35 - where the result should be written. These four pointers actually

36 - refer to multi-dimensional arrays, but each thread only accesses the

37 - portion of data assigned to it. I have omitted all other runtime

38 - parameters here for simplicity.

39 -

40 - ```cpp

41 - template<

42 - typename scalar_t,

43 - int HEAD_SIZE,

44 - int BLOCK_SIZE,

45 - int NUM_THREADS,

46 - int PARTITION_SIZE = 0>

47 - __device__ void paged_attention_kernel(

48 - ... // Other side args.

49 - const scalar_t* __restrict__ out, // [num_seqs, num_heads, max_num_partitions, head_size]

50 - const scalar_t* __restrict__ q, // [num_seqs, num_heads, head_size]

51 - const scalar_t* __restrict__ k_cache, // [num_blocks, num_kv_heads, head_size/x, block_size, x]

52 - const scalar_t* __restrict__ v_cache, // [num_blocks, num_kv_heads, head_size, block_size]

53 - ... // Other side args.

54 - )

55 - ```

56 -

57 -- There are also a list of template arguments above the function

58 - signature that are determined during compilation time. `scalar_t`

59 - represents the data type of the query, key, and value data elements,

60 - such as FP16. `HEAD_SIZE` indicates the number of elements in each

61 - head. `BLOCK_SIZE` refers to the number of tokens in each block.

62 - `NUM_THREADS` denotes the number of threads in each thread block.

63 - `PARTITION_SIZE` represents the number of tensor parallel GPUs (For

64 - simplicity, we assume this is 0 and tensor parallel is disabled).

65 -

66 -- With these arguments, we need to perform a sequence of preparations.

67 - This includes calculating the current head index, block index, and

68 - other necessary variables. However, for now, we can ignore these

69 - preparations and proceed directly to the actual calculations. It will

70 - be easier to understand them once we grasp the entire flow.

32 +The kernel function takes a list of arguments for the current thread

33 +to perform its assigned work. The three most important arguments are

34 +the input pointers `q`, `k_cache`, and `v_cache`, which point

35 +to query, key, and value data on global memory that need to be read

36 +and processed. The output pointer `out` points to global memory

37 +where the result should be written. These four pointers actually

38 +refer to multi-dimensional arrays, but each thread only accesses the

39 +portion of data assigned to it. I have omitted all other runtime

40 +parameters here for simplicity.

41 +

42 +```cpp

43 +template<typename scalar_t, int HEAD_SIZE, int BLOCK_SIZE, int NUM_THREADS, int PARTITION_SIZE = 0>

44 +__device__ void paged_attention_kernel(

45 + ... // Other side args.

46 + const scalar_t* __restrict__ out, // [num_seqs, num_heads, max_num_partitions, head_size]

47 + const scalar_t* __restrict__ q, // [num_seqs, num_heads, head_size]

48 + const scalar_t* __restrict__ k_cache, // [num_blocks, num_kv_heads, head_size/x, block_size, x]

49 + const scalar_t* __restrict__ v_cache, // [num_blocks, num_kv_heads, head_size, block_size]

50 + ... // Other side args.

51 +)

52 +```

53 +

54 +There are also a list of template arguments above the function

55 +signature that are determined during compilation time. `scalar_t`

56 +represents the data type of the query, key, and value data elements,

57 +such as FP16. `HEAD_SIZE` indicates the number of elements in each

58 +head. `BLOCK_SIZE` refers to the number of tokens in each block.

59 +`NUM_THREADS` denotes the number of threads in each thread block.

60 +`PARTITION_SIZE` represents the number of tensor parallel GPUs (For

61 +simplicity, we assume this is 0 and tensor parallel is disabled).

62 +

63 +With these arguments, we need to perform a sequence of preparations.

64 +This includes calculating the current head index, block index, and

65 +other necessary variables. However, for now, we can ignore these

66 +preparations and proceed directly to the actual calculations. It will

67 +be easier to understand them once we grasp the entire flow.

68

69 ## Concepts

70

74 -- Just before we dive into the calculation flow, I want to describe a

75 - few concepts that are needed for later sections. However, you may

76 - skip this section and return later if you encounter any confusing

77 - terminologies.

71 +Just before we dive into the calculation flow, I want to describe a

72 +few concepts that are needed for later sections. However, you may

73 +skip this section and return later if you encounter any confusing

74 +terminologies.

75 +

76 - **Sequence**: A sequence represents a client request. For example,

77 the data pointed to by `q` has a shape of

78 `[num_seqs, num_heads, head_size]`. That represents there are total

@@ -129,236 +127,236 @@ title: vLLM Paged Attention

127

128 ## Query

129

132 -- This section will introduce how query data is stored in memory and

133 - fetched by each thread. As mentioned above, each thread group fetches

134 - one query token data, while each thread itself only handles a part of

135 - one query token data. Within each warp, every thread group will fetch

136 - the same query token data, but will multiply it with different key

137 - token data.

130 +This section will introduce how query data is stored in memory and

131 +fetched by each thread. As mentioned above, each thread group fetches

132 +one query token data, while each thread itself only handles a part of

133 +one query token data. Within each warp, every thread group will fetch

134 +the same query token data, but will multiply it with different key

135 +token data.

136

139 - ```cpp

140 - const scalar_t* q_ptr = q + seq_idx * q_stride + head_idx * HEAD_SIZE;

141 - ```

137 +```cpp

138 +const scalar_t* q_ptr = q + seq_idx * q_stride + head_idx * HEAD_SIZE;

139 +```

140

141 <figure markdown="span">

142 ![](../../assets/kernel/query.png){ align="center" alt="query" width="70%" }

143 </figure>

144

147 -- Each thread defines its own `q_ptr` which points to the assigned

148 - query token data on global memory. For example, if `VEC_SIZE` is 4

149 - and `HEAD_SIZE` is 128, the `q_ptr` points to data that contains

150 - total of 128 elements divided into 128 / 4 = 32 vecs.

145 +Each thread defines its own `q_ptr` which points to the assigned

146 +query token data on global memory. For example, if `VEC_SIZE` is 4

147 +and `HEAD_SIZE` is 128, the `q_ptr` points to data that contains

148 +total of 128 elements divided into 128 / 4 = 32 vecs.

149

150 <figure markdown="span">

151 ![](../../assets/kernel/q_vecs.png){ align="center" alt="q_vecs" width="70%" }

152 </figure>

153

156 - ```cpp

157 - __shared__ Q_vec q_vecs[THREAD_GROUP_SIZE][NUM_VECS_PER_THREAD];

158 - ```

154 +```cpp

155 +__shared__ Q_vec q_vecs[THREAD_GROUP_SIZE][NUM_VECS_PER_THREAD];

156 +```

157

160 -- Next, we need to read the global memory data pointed to by `q_ptr`

161 - into shared memory as `q_vecs`. It is important to note that each

162 - vecs is assigned to a different row. For example, if the

163 - `THREAD_GROUP_SIZE` is 2, thread 0 will handle the 0th row vecs,

164 - while thread 1 handles the 1st row vecs. By reading the query data in

165 - this way, neighboring threads like thread 0 and thread 1 can read

166 - neighbor memory, achieving the memory coalescing to improve

167 - performance.

158 +Next, we need to read the global memory data pointed to by `q_ptr`

159 +into shared memory as `q_vecs`. It is important to note that each

160 +vecs is assigned to a different row. For example, if the

161 +`THREAD_GROUP_SIZE` is 2, thread 0 will handle the 0th row vecs,

162 +while thread 1 handles the 1st row vecs. By reading the query data in

163 +this way, neighboring threads like thread 0 and thread 1 can read

164 +neighbor memory, achieving the memory coalescing to improve

165 +performance.

166

167 ## Key

168

171 -- Similar to the "Query" section, this section introduces memory layout

172 - and assignment for keys. While each thread group only handle one

173 - query token one kernel run, it may handle multiple key tokens across

174 - multiple iterations. Meanwhile, each warp will process multiple blocks

175 - of key tokens in multiple iterations, ensuring that all context

176 - tokens are processed by the entire thread group after the kernel run.

177 - In this context, "handle" refers to performing the dot multiplication

178 - between query data and key data.

179 -

180 - ```cpp

181 - const scalar_t* k_ptr = k_cache + physical_block_number * kv_block_stride

182 - + kv_head_idx * kv_head_stride

183 - + physical_block_offset * x;

184 - ```

185 -

186 -- Unlike to `q_ptr`, `k_ptr` in each thread will point to different

187 - key token at different iterations. As shown above, that `k_ptr`

188 - points to key token data based on `k_cache` at assigned block,

189 - assigned head and assigned token.

169 +Similar to the "Query" section, this section introduces memory layout

170 +and assignment for keys. While each thread group only handle one

171 +query token one kernel run, it may handle multiple key tokens across

172 +multiple iterations. Meanwhile, each warp will process multiple blocks

173 +of key tokens in multiple iterations, ensuring that all context

174 +tokens are processed by the entire thread group after the kernel run.

175 +In this context, "handle" refers to performing the dot multiplication

176 +between query data and key data.

177 +

178 +```cpp

179 +const scalar_t* k_ptr = k_cache + physical_block_number * kv_block_stride

180 + + kv_head_idx * kv_head_stride

181 + + physical_block_offset * x;

182 +```

183 +

184 +Unlike to `q_ptr`, `k_ptr` in each thread will point to different

185 +key token at different iterations. As shown above, that `k_ptr`

186 +points to key token data based on `k_cache` at assigned block,

187 +assigned head and assigned token.

188

189 <figure markdown="span">

190 ![](../../assets/kernel/key.png){ align="center" alt="key" width="70%" }

191 </figure>

192

195 -- The diagram above illustrates the memory layout for key data. It

196 - assumes that the `BLOCK_SIZE` is 16, `HEAD_SIZE` is 128, `x` is

197 - 8, `THREAD_GROUP_SIZE` is 2, and there are a total of 4 warps. Each

198 - rectangle represents all the elements for one key token at one head,

199 - which will be processed by one thread group. The left half shows the

200 - total 16 blocks of key token data for warp 0, while the right half

201 - represents the remaining key token data for other warps or

202 - iterations. Inside each rectangle, there are a total 32 vecs (128

203 - elements for one token) that will be processed by 2 threads (one

204 - thread group) separately.

193 +The diagram above illustrates the memory layout for key data. It

194 +assumes that the `BLOCK_SIZE` is 16, `HEAD_SIZE` is 128, `x` is

195 +8, `THREAD_GROUP_SIZE` is 2, and there are a total of 4 warps. Each

196 +rectangle represents all the elements for one key token at one head,

197 +which will be processed by one thread group. The left half shows the

198 +total 16 blocks of key token data for warp 0, while the right half

199 +represents the remaining key token data for other warps or

200 +iterations. Inside each rectangle, there are a total 32 vecs (128

201 +elements for one token) that will be processed by 2 threads (one

202 +thread group) separately.

203

204 <figure markdown="span">

205 ![](../../assets/kernel/k_vecs.png){ align="center" alt="k_vecs" width="70%" }

206 </figure>

207

210 - ```cpp

211 - K_vec k_vecs[NUM_VECS_PER_THREAD]

212 - ```

213 -

214 -- Next, we need to read the key token data from `k_ptr` and store

215 - them on register memory as `k_vecs`. We use register memory for

216 - `k_vecs` because it will only be accessed by one thread once,

217 - whereas `q_vecs` will be accessed by multiple threads multiple

218 - times. Each `k_vecs` will contain multiple vectors for later

219 - calculation. Each vec will be set at each inner iteration. The

220 - assignment of vecs allows neighboring threads in a warp to read

221 - neighboring memory together, which again promotes the memory

222 - coalescing. For instance, thread 0 will read vec 0, while thread 1

223 - will read vec 1. In the next inner loop, thread 0 will read vec 2,

224 - while thread 1 will read vec 3, and so on.

225 -

226 -- You may still be a little confused about the overall flow. Don't

227 - worry, please keep reading the next "QK" section. It will illustrate

228 - the query and key calculation flow in a clearer and higher-level

229 - manner.

208 +```cpp

209 +K_vec k_vecs[NUM_VECS_PER_THREAD]

210 +```

211 +

212 +Next, we need to read the key token data from `k_ptr` and store

213 +them on register memory as `k_vecs`. We use register memory for

214 +`k_vecs` because it will only be accessed by one thread once,

215 +whereas `q_vecs` will be accessed by multiple threads multiple

216 +times. Each `k_vecs` will contain multiple vectors for later

217 +calculation. Each vec will be set at each inner iteration. The

218 +assignment of vecs allows neighboring threads in a warp to read

219 +neighboring memory together, which again promotes the memory

220 +coalescing. For instance, thread 0 will read vec 0, while thread 1

221 +will read vec 1. In the next inner loop, thread 0 will read vec 2,

222 +while thread 1 will read vec 3, and so on.

223 +

224 +You may still be a little confused about the overall flow. Don't

225 +worry, please keep reading the next "QK" section. It will illustrate

226 +the query and key calculation flow in a clearer and higher-level

227 +manner.

228

229 ## QK

230

233 -- As shown the pseudo code below, before the entire for loop block, we

234 - fetch the query data for one token and store it in `q_vecs`. Then,

235 - in the outer for loop, we iterate through different `k_ptrs` that

236 - point to different tokens and prepare the `k_vecs` in the inner for

237 - loop. Finally, we perform the dot multiplication between the

238 - `q_vecs` and each `k_vecs`.

239 -

240 - ```cpp

241 - q_vecs = ...

242 - for ... {

243 - k_ptr = ...

244 - for ... {

231 +As shown the pseudo code below, before the entire for loop block, we

232 +fetch the query data for one token and store it in `q_vecs`. Then,

233 +in the outer for loop, we iterate through different `k_ptrs` that

234 +point to different tokens and prepare the `k_vecs` in the inner for

235 +loop. Finally, we perform the dot multiplication between the

236 +`q_vecs` and each `k_vecs`.

237 +

238 +```cpp

239 +q_vecs = ...

240 +for ... {

241 + k_ptr = ...

242 + for ... {

243 k_vecs[i] = ...

246 - }

247 - ...

248 - float qk = scale * Qk_dot<scalar_t, THREAD_GROUP_SIZE>::dot(q_vecs[thread_group_offset], k_vecs);

249 - }

250 - ```

251 -

252 -- As mentioned before, for each thread, it only fetches part of the

253 - query and key token data at a time. However, there will be a cross

254 - thread group reduction happen in the `Qk_dot<>::dot` . So `qk`

255 - returned here is not just between part of the query and key token dot

256 - multiplication, but actually a full result between entire query and

257 - key token data.

258 -

259 -- For example, if the value of `HEAD_SIZE` is 128 and

260 - `THREAD_GROUP_SIZE` is 2, each thread's `k_vecs` will contain

261 - total 64 elements. However, the returned `qk` is actually the

262 - result of dot multiplication between 128 query elements and 128 key

263 - elements. If you want to learn more about the details of the dot

264 - multiplication and reduction, you may refer to the implementation of

265 - `Qk_dot<>::dot`. However, for the sake of simplicity, I will not

266 - cover it in this document.

244 + }

245 + ...

246 + float qk = scale * Qk_dot<scalar_t, THREAD_GROUP_SIZE>::dot(q_vecs[thread_group_offset], k_vecs);

247 +}

248 +```

249 +

250 +As mentioned before, for each thread, it only fetches part of the

251 +query and key token data at a time. However, there will be a cross

252 +thread group reduction happen in the `Qk_dot<>::dot` . So `qk`

253 +returned here is not just between part of the query and key token dot

254 +multiplication, but actually a full result between entire query and

255 +key token data.

256 +

257 +For example, if the value of `HEAD_SIZE` is 128 and

258 +`THREAD_GROUP_SIZE` is 2, each thread's `k_vecs` will contain

259 +total 64 elements. However, the returned `qk` is actually the

260 +result of dot multiplication between 128 query elements and 128 key

261 +elements. If you want to learn more about the details of the dot

262 +multiplication and reduction, you may refer to the implementation of

263 +`Qk_dot<>::dot`. However, for the sake of simplicity, I will not

264 +cover it in this document.

265

266 ## Softmax

267

270 -- Next, we need to calculate the normalized softmax for all `qk`s,

271 - as shown above, where each $x$ represents a `qk`. To do this,

272 - we must obtain the reduced value of `qk_max`($m(x)$) and

273 - the `exp_sum`($\ell(x)$) of all `qk`s. The reduction

274 - should be performed across the entire thread block, encompassing

275 - results between the query token and all context key tokens.

268 +Next, we need to calculate the normalized softmax for all `qk`s,

269 +as shown above, where each $x$ represents a `qk`. To do this,

270 +we must obtain the reduced value of `qk_max`($m(x)$) and

271 +the `exp_sum`($\ell(x)$) of all `qk`s. The reduction

272 +should be performed across the entire thread block, encompassing

273 +results between the query token and all context key tokens.

274

277 - $$

278 - \begin{gather*}

279 - 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 \\

280 - \quad \operatorname{softmax}(x):=\frac{f(x)}{\ell(x)}

281 - \end{gather*}

282 - $$

275 +$$

276 +\begin{gather*}

277 +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 \\

278 +\quad \operatorname{softmax}(x):=\frac{f(x)}{\ell(x)}

279 +\end{gather*}

280 +$$

281

282 ### `qk_max` and `logits`

283

286 -- Just right after we get the `qk` result, we can set the temporary

287 - `logits` result with `qk` (In the end, the `logits` should

288 - store the normalized softmax result). Also we can compare and collect

289 - the `qk_max` for all `qk`s that are calculated by current

290 - thread group.

291 -

292 - ```cpp

293 - if (thread_group_offset == 0) {

294 - const bool mask = token_idx >= context_len;

295 - logits[token_idx - start_token_idx] = mask ? 0.f : qk;

296 - qk_max = mask ? qk_max : fmaxf(qk_max, qk);

297 - }

298 - ```

299 -

300 -- Please note that the `logits` here is on shared memory, so each

301 - thread group will set the fields for its own assigned context tokens.

302 - Overall, the size of logits should be number of context tokens.

303 -

304 - ```cpp

305 - for (int mask = WARP_SIZE / 2; mask >= THREAD_GROUP_SIZE; mask /= 2) {

306 - qk_max = fmaxf(qk_max, VLLM_SHFL_XOR_SYNC(qk_max, mask));

307 - }

308 -

309 - if (lane == 0) {

310 - red_smem[warp_idx] = qk_max;

311 - }

312 - ```

313 -

314 -- Then we need to get the reduced `qk_max` across each warp. The main

315 - idea is to make threads in warp to communicate with each other and

316 - get the final max `qk` .

317 -

318 - ```cpp

319 - for (int mask = NUM_WARPS / 2; mask >= 1; mask /= 2) {

320 - qk_max = fmaxf(qk_max, VLLM_SHFL_XOR_SYNC(qk_max, mask));

321 - }

322 - qk_max = VLLM_SHFL_SYNC(qk_max, 0);

323 - ```

324 -

325 -- Finally, we can get the reduced `qk_max` from whole thread block by

326 - compare the `qk_max` from all warps in this thread block. Then we

327 - need to broadcast the final result to each thread.

284 +Just right after we get the `qk` result, we can set the temporary

285 +`logits` result with `qk` (In the end, the `logits` should

286 +store the normalized softmax result). Also we can compare and collect

287 +the `qk_max` for all `qk`s that are calculated by current

288 +thread group.

289 +

290 +```cpp

291 +if (thread_group_offset == 0) {

292 + const bool mask = token_idx >= context_len;

293 + logits[token_idx - start_token_idx] = mask ? 0.f : qk;

294 + qk_max = mask ? qk_max : fmaxf(qk_max, qk);

295 +}

296 +```

297 +

298 +Please note that the `logits` here is on shared memory, so each

299 +thread group will set the fields for its own assigned context tokens.

300 +Overall, the size of logits should be number of context tokens.

301 +

302 +```cpp

303 +for (int mask = WARP_SIZE / 2; mask >= THREAD_GROUP_SIZE; mask /= 2) {

304 + qk_max = fmaxf(qk_max, VLLM_SHFL_XOR_SYNC(qk_max, mask));

305 +}

306 +

307 +if (lane == 0) {

308 + red_smem[warp_idx] = qk_max;

309 +}

310 +```

311 +

312 +Then we need to get the reduced `qk_max` across each warp. The main

313 +idea is to make threads in warp to communicate with each other and

314 +get the final max `qk` .

315 +

316 +```cpp

317 +for (int mask = NUM_WARPS / 2; mask >= 1; mask /= 2) {

318 + qk_max = fmaxf(qk_max, VLLM_SHFL_XOR_SYNC(qk_max, mask));

319 +}

320 +qk_max = VLLM_SHFL_SYNC(qk_max, 0);

321 +```

322 +

323 +Finally, we can get the reduced `qk_max` from whole thread block by

324 +compare the `qk_max` from all warps in this thread block. Then we

325 +need to broadcast the final result to each thread.

326

327 ### `exp_sum`

328

331 -- Similar to `qk_max`, we need to get the reduced sum value from the

332 - entire thread block too.

333 -

334 - ```cpp

335 - for (int i = thread_idx; i < num_tokens; i += NUM_THREADS) {

336 - float val = __expf(logits[i] - qk_max);

337 - logits[i] = val;

338 - exp_sum += val;

339 - }

340 - ...

341 - exp_sum = block_sum<NUM_WARPS>(&red_smem[NUM_WARPS], exp_sum);

342 - ```

343 -

344 -- Firstly, sum all exp values from each thread group, and meanwhile,

345 - convert each entry of `logits` from `qk` to `exp(qk - qk_max)`.

346 - Please note, the `qk_max` here is already the max `qk` across the

347 - whole thread block. And then we can do reduction for `exp_sum`

348 - across whole thread block just like the `qk_max`.

349 -

350 - ```cpp

351 - const float inv_sum = __fdividef(1.f, exp_sum + 1e-6f);

352 - for (int i = thread_idx; i < num_tokens; i += NUM_THREADS) {

353 - logits[i] *= inv_sum;

354 - }

355 - ```

356 -

357 -- Finally, with the reduced `qk_max` and `exp_sum`, we can obtain

358 - the final normalized softmax result as `logits`. This `logits`

359 - variable will be used for dot multiplication with the value data in

360 - later steps. Now, it should store the normalized softmax result of

361 - `qk` for all assigned context tokens.

329 +Similar to `qk_max`, we need to get the reduced sum value from the

330 +entire thread block too.

331 +

332 +```cpp

333 +for (int i = thread_idx; i < num_tokens; i += NUM_THREADS) {

334 + float val = __expf(logits[i] - qk_max);

335 + logits[i] = val;

336 + exp_sum += val;

337 +}

338 +...

339 +exp_sum = block_sum<NUM_WARPS>(&red_smem[NUM_WARPS], exp_sum);

340 +```

341 +

342 +Firstly, sum all exp values from each thread group, and meanwhile,

343 +convert each entry of `logits` from `qk` to `exp(qk - qk_max)`.

344 +Please note, the `qk_max` here is already the max `qk` across the

345 +whole thread block. And then we can do reduction for `exp_sum`

346 +across whole thread block just like the `qk_max`.

347 +

348 +```cpp

349 +const float inv_sum = __fdividef(1.f, exp_sum + 1e-6f);

350 +for (int i = thread_idx; i < num_tokens; i += NUM_THREADS) {

351 + logits[i] *= inv_sum;

352 +}

353 +```

354 +

355 +Finally, with the reduced `qk_max` and `exp_sum`, we can obtain

356 +the final normalized softmax result as `logits`. This `logits`

357 +variable will be used for dot multiplication with the value data in

358 +later steps. Now, it should store the normalized softmax result of

359 +`qk` for all assigned context tokens.

360

361 ## Value

362

@@ -374,127 +372,127 @@ title: vLLM Paged Attention

372 ![](../../assets/kernel/v_vec.png){ align="center" alt="v_vec" width="70%" }

373 </figure>

374

377 -- Now we need to retrieve the value data and perform dot multiplication

378 - with `logits`. Unlike query and key, there is no thread group

379 - concept for value data. As shown in diagram, different from key token

380 - memory layout, elements from the same column correspond to the same

381 - value token. For one block of value data, there are `HEAD_SIZE` of

382 - rows and `BLOCK_SIZE` of columns that are split into multiple

383 - `v_vecs`.

384 -

385 -- Each thread always fetches `V_VEC_SIZE` elements from the same

386 - `V_VEC_SIZE` of tokens at a time. As a result, a single thread

387 - retrieves multiple `v_vec`s from different rows and the same

388 - columns through multiple inner iterations. For each `v_vec`, it

389 - needs to be dot multiplied with the corresponding `logits_vec`,

390 - which is also `V_VEC_SIZE` elements from `logits`. Overall, with

391 - multiple inner iterations, each warp will process one block of value

392 - tokens. And with multiple outer iterations, the whole context value

393 - tokens are processed

394 -

395 - ```cpp

396 - float accs[NUM_ROWS_PER_THREAD];

397 - for ... { // Iteration over different blocks.

398 - logits_vec = ...

399 - for ... { // Iteration over different rows.

400 - v_vec = ...

401 - ...

402 - accs[i] += dot(logits_vec, v_vec);

403 - }

404 - }

405 - ```

406 -

407 -- As shown in the above pseudo code, in the outer loop, similar to

408 - `k_ptr`, `logits_vec` iterates over different blocks and reads

409 - `V_VEC_SIZE` elements from `logits`. In the inner loop, each

410 - thread reads `V_VEC_SIZE` elements from the same tokens as a

411 - `v_vec` and performs dot multiplication. It is important to note

412 - that in each inner iteration, the thread fetches different head

413 - position elements for the same tokens. The dot result is then

414 - accumulated in `accs`. Therefore, each entry of `accs` is mapped

415 - to a head position assigned to the current thread.

416 -

417 -- For example, if `BLOCK_SIZE` is 16 and `V_VEC_SIZE` is 8, each

418 - thread fetches 8 value elements for 8 tokens at a time. Each element

419 - is from different tokens at the same head position. If `HEAD_SIZE`

420 - is 128 and `WARP_SIZE` is 32, for each inner loop, a warp needs to

421 - fetch `WARP_SIZE * V_VEC_SIZE = 256` elements. This means there are

422 - a total of 128 * 16 / 256 = 8 inner iterations for a warp to handle

423 - a whole block of value tokens. And each `accs` in each thread

424 - contains 8 elements that accumulated at 8 different head positions.

425 - For the thread 0, the `accs` variable will have 8 elements, which

426 - are 0th, 32th … 224th elements of a value head that are accumulated

427 - from all assigned 8 tokens.

375 +Now we need to retrieve the value data and perform dot multiplication

376 +with `logits`. Unlike query and key, there is no thread group

377 +concept for value data. As shown in diagram, different from key token

378 +memory layout, elements from the same column correspond to the same

379 +value token. For one block of value data, there are `HEAD_SIZE` of

380 +rows and `BLOCK_SIZE` of columns that are split into multiple

381 +`v_vecs`.

382 +

383 +Each thread always fetches `V_VEC_SIZE` elements from the same

384 +`V_VEC_SIZE` of tokens at a time. As a result, a single thread

385 +retrieves multiple `v_vec`s from different rows and the same

386 +columns through multiple inner iterations. For each `v_vec`, it

387 +needs to be dot multiplied with the corresponding `logits_vec`,

388 +which is also `V_VEC_SIZE` elements from `logits`. Overall, with

389 +multiple inner iterations, each warp will process one block of value

390 +tokens. And with multiple outer iterations, the whole context value

391 +tokens are processed

392 +

393 +```cpp

394 +float accs[NUM_ROWS_PER_THREAD];

395 +for ... { // Iteration over different blocks.

396 + logits_vec = ...

397 + for ... { // Iteration over different rows.

398 + v_vec = ...

399 + ...

400 + accs[i] += dot(logits_vec, v_vec);

401 + }

402 +}

403 +```

404 +

405 +As shown in the above pseudo code, in the outer loop, similar to

406 +`k_ptr`, `logits_vec` iterates over different blocks and reads

407 +`V_VEC_SIZE` elements from `logits`. In the inner loop, each

408 +thread reads `V_VEC_SIZE` elements from the same tokens as a

409 +`v_vec` and performs dot multiplication. It is important to note

410 +that in each inner iteration, the thread fetches different head

411 +position elements for the same tokens. The dot result is then

412 +accumulated in `accs`. Therefore, each entry of `accs` is mapped

413 +to a head position assigned to the current thread.

414 +

415 +For example, if `BLOCK_SIZE` is 16 and `V_VEC_SIZE` is 8, each

416 +thread fetches 8 value elements for 8 tokens at a time. Each element

417 +is from different tokens at the same head position. If `HEAD_SIZE`

418 +is 128 and `WARP_SIZE` is 32, for each inner loop, a warp needs to

419 +fetch `WARP_SIZE * V_VEC_SIZE = 256` elements. This means there are

420 +a total of 128 * 16 / 256 = 8 inner iterations for a warp to handle

421 +a whole block of value tokens. And each `accs` in each thread

422 +contains 8 elements that accumulated at 8 different head positions.

423 +For the thread 0, the `accs` variable will have 8 elements, which

424 +are 0th, 32th … 224th elements of a value head that are accumulated

425 +from all assigned 8 tokens.

426

427 ## LV

428

431 -- Now, we need to perform reduction for `accs` within each warp. This

432 - process allows each thread to accumulate the `accs` for the

433 - assigned head positions of all tokens in one block.

429 +Now, we need to perform reduction for `accs` within each warp. This

430 +process allows each thread to accumulate the `accs` for the

431 +assigned head positions of all tokens in one block.

432

435 - ```cpp

436 - for (int i = 0; i < NUM_ROWS_PER_THREAD; i++) {

437 - float acc = accs[i];

438 - for (int mask = NUM_V_VECS_PER_ROW / 2; mask >= 1; mask /= 2) {

433 +```cpp

434 +for (int i = 0; i < NUM_ROWS_PER_THREAD; i++) {

435 + float acc = accs[i];

436 + for (int mask = NUM_V_VECS_PER_ROW / 2; mask >= 1; mask /= 2) {

437 acc += VLLM_SHFL_XOR_SYNC(acc, mask);

440 - }

441 - accs[i] = acc;

442 - }

443 - ```

444 -

445 -- Next, we perform reduction for `accs` across all warps, allowing

446 - each thread to have the accumulation of `accs` for the assigned

447 - head positions of all context tokens. Please note that each `accs`

448 - in every thread only stores the accumulation for a portion of

449 - elements of the entire head for all context tokens. However, overall,

450 - all results for output have been calculated but are just stored in

451 - different thread register memory.

452 -

453 - ```cpp

454 - float* out_smem = reinterpret_cast<float*>(shared_mem);

455 - for (int i = NUM_WARPS; i > 1; i /= 2) {

456 - // Upper warps write to shared memory.

457 - ...

458 - float* dst = &out_smem[(warp_idx - mid) * HEAD_SIZE];

459 - for (int i = 0; i < NUM_ROWS_PER_THREAD; i++) {

460 - ...

461 - dst[row_idx] = accs[i];

462 - }

463 -

464 - // Lower warps update the output.

465 - const float* src = &out_smem[warp_idx * HEAD_SIZE];

466 - for (int i = 0; i < NUM_ROWS_PER_THREAD; i++) {

467 - ...

468 - accs[i] += src[row_idx];

469 - }

470 -

471 - // Write out the accs.

472 - }

473 - ```

438 + }

439 + accs[i] = acc;

440 +}

441 +```

442 +

443 +Next, we perform reduction for `accs` across all warps, allowing

444 +each thread to have the accumulation of `accs` for the assigned

445 +head positions of all context tokens. Please note that each `accs`

446 +in every thread only stores the accumulation for a portion of

447 +elements of the entire head for all context tokens. However, overall,

448 +all results for output have been calculated but are just stored in

449 +different thread register memory.

450 +

451 +```cpp

452 +float* out_smem = reinterpret_cast<float*>(shared_mem);

453 +for (int i = NUM_WARPS; i > 1; i /= 2) {

454 + // Upper warps write to shared memory.

455 + ...

456 + float* dst = &out_smem[(warp_idx - mid) * HEAD_SIZE];

457 + for (int i = 0; i < NUM_ROWS_PER_THREAD; i++) {

458 + ...

459 + dst[row_idx] = accs[i];

460 + }

461 +

462 + // Lower warps update the output.

463 + const float* src = &out_smem[warp_idx * HEAD_SIZE];

464 + for (int i = 0; i < NUM_ROWS_PER_THREAD; i++) {

465 + ...

466 + accs[i] += src[row_idx];

467 + }

468 +

469 + // Write out the accs.

470 +}

471 +```

472

473 ## Output

474

477 -- Now we can write all of calculated result from local register memory

478 - to final output global memory.

479 -

480 - ```cpp

481 - scalar_t* out_ptr = out + seq_idx * num_heads * max_num_partitions * HEAD_SIZE

482 - + head_idx * max_num_partitions * HEAD_SIZE

483 - + partition_idx * HEAD_SIZE;

484 - ```

485 -

486 -- First, we need to define the `out_ptr` variable, which points to

487 - the start address of the assigned sequence and assigned head.

488 -

489 - ```cpp

490 - for (int i = 0; i < NUM_ROWS_PER_THREAD; i++) {

491 - const int row_idx = lane / NUM_V_VECS_PER_ROW + i * NUM_ROWS_PER_ITER;

492 - if (row_idx < HEAD_SIZE && lane % NUM_V_VECS_PER_ROW == 0) {

493 - from_float(*(out_ptr + row_idx), accs[i]);

494 - }

495 - }

496 - ```

497 -

498 -- Finally, we need to iterate over different assigned head positions

499 - and write out the corresponding accumulated result based on the

500 - `out_ptr`.

475 +Now we can write all of calculated result from local register memory

476 +to final output global memory.

477 +

478 +```cpp

479 +scalar_t* out_ptr = out + seq_idx * num_heads * max_num_partitions * HEAD_SIZE

480 + + head_idx * max_num_partitions * HEAD_SIZE

481 + + partition_idx * HEAD_SIZE;

482 +```

483 +

484 +First, we need to define the `out_ptr` variable, which points to

485 +the start address of the assigned sequence and assigned head.

486 +

487 +```cpp

488 +for (int i = 0; i < NUM_ROWS_PER_THREAD; i++) {

489 + const int row_idx = lane / NUM_V_VECS_PER_ROW + i * NUM_ROWS_PER_ITER;

490 + if (row_idx < HEAD_SIZE && lane % NUM_V_VECS_PER_ROW == 0) {

491 + from_float(*(out_ptr + row_idx), accs[i]);

492 + }

493 +}

494 +```

495 +

496 +Finally, we need to iterate over different assigned head positions

497 +and write out the corresponding accumulated result based on the

498 +`out_ptr`.

[Doc] Fix indentation problems in V0 Paged Attention docs (#18659)