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, 08:30 UTC
Fix figures in design doc (#18612)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Harry Mellor committed
May 23, 2025 at 17:09 UTC
3d28ad343f72c950be36aed7fb8c18ab39f14dd2
1 file changed
+12
-26
docs/design/kernel/paged_attention.md
+12
-26
@@ -140,22 +140,18 @@ title: vLLM Paged Attention
140
const scalar_t* q_ptr = q + seq_idx * q_stride + head_idx * HEAD_SIZE;
141
```
142
143
- <figure markdown="span">
144
- { align="center" alt="query" width="70%" }
145
- <figcaption>
146
-</figcaption>
147
- </figure>
143
+<figure markdown="span">
144
+ { align="center" alt="query" width="70%" }
145
+</figure>
146
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.
151
154
- <figure markdown="span">
155
- { align="center" alt="q_vecs" width="70%" }
156
- <figcaption>
157
-</figcaption>
158
- </figure>
152
+<figure markdown="span">
153
+ { align="center" alt="q_vecs" width="70%" }
154
+</figure>
155
156
```cpp
157
__shared__ Q_vec q_vecs[THREAD_GROUP_SIZE][NUM_VECS_PER_THREAD];
@@ -192,11 +188,9 @@ title: vLLM Paged Attention
188
points to key token data based on `k_cache` at assigned block,
189
assigned head and assigned token.
190
195
- <figure markdown="span">
196
- { align="center" alt="key" width="70%" }
197
- <figcaption>
198
-</figcaption>
199
- </figure>
191
+<figure markdown="span">
192
+ { align="center" alt="key" width="70%" }
193
+</figure>
194
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
@@ -209,11 +203,9 @@ title: vLLM Paged Attention
203
elements for one token) that will be processed by 2 threads (one
204
thread group) separately.
205
212
- <figure markdown="span">
213
- { align="center" alt="k_vecs" width="70%" }
214
- <figcaption>
215
-</figcaption>
216
- </figure>
206
+<figure markdown="span">
207
+ { align="center" alt="k_vecs" width="70%" }
208
+</figure>
209
210
```cpp
211
K_vec k_vecs[NUM_VECS_PER_THREAD]
@@ -372,20 +364,14 @@ title: vLLM Paged Attention
364
365
<figure markdown="span">
366
{ align="center" alt="value" width="70%" }
375
- <figcaption>
376
-</figcaption>
367
</figure>
368
369
<figure markdown="span">
370
{ align="center" alt="logits_vec" width="50%" }
381
- <figcaption>
382
-</figcaption>
371
</figure>
372
373
<figure markdown="span">
374
{ align="center" alt="v_vec" width="70%" }
387
- <figcaption>
388
-</figcaption>
375
</figure>
376
377
- Now we need to retrieve the value data and perform dot multiplication