Skip to content

Fix figures in design doc #18612

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
May 23, 2025
Merged

Fix figures in design doc #18612

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
38 changes: 12 additions & 26 deletions docs/design/kernel/paged_attention.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -140,22 +140,18 @@ title: vLLM Paged Attention
const scalar_t* q_ptr = q + seq_idx * q_stride + head_idx * HEAD_SIZE;
```

<figure markdown="span">
![](../../assets/kernel/query.png){ align="center" alt="query" width="70%" }
<figcaption>
</figcaption>
</figure>
<figure markdown="span">
![](../../assets/kernel/query.png){ align="center" alt="query" width="70%" }
</figure>

- Each thread defines its own `q_ptr` which points to the assigned
query token data on global memory. For example, if `VEC_SIZE` is 4
and `HEAD_SIZE` is 128, the `q_ptr` points to data that contains
total of 128 elements divided into 128 / 4 = 32 vecs.

<figure markdown="span">
![](../../assets/kernel/q_vecs.png){ align="center" alt="q_vecs" width="70%" }
<figcaption>
</figcaption>
</figure>
<figure markdown="span">
![](../../assets/kernel/q_vecs.png){ align="center" alt="q_vecs" width="70%" }
</figure>

```cpp
__shared__ Q_vec q_vecs[THREAD_GROUP_SIZE][NUM_VECS_PER_THREAD];
Expand Down Expand Up @@ -192,11 +188,9 @@ title: vLLM Paged Attention
points to key token data based on `k_cache` at assigned block,
assigned head and assigned token.

<figure markdown="span">
![](../../assets/kernel/key.png){ align="center" alt="key" width="70%" }
<figcaption>
</figcaption>
</figure>
<figure markdown="span">
![](../../assets/kernel/key.png){ align="center" alt="key" width="70%" }
</figure>

- The diagram above illustrates the memory layout for key data. It
assumes that the `BLOCK_SIZE` is 16, `HEAD_SIZE` is 128, `x` is
Expand All @@ -209,11 +203,9 @@ title: vLLM Paged Attention
elements for one token) that will be processed by 2 threads (one
thread group) separately.

<figure markdown="span">
![](../../assets/kernel/k_vecs.png){ align="center" alt="k_vecs" width="70%" }
<figcaption>
</figcaption>
</figure>
<figure markdown="span">
![](../../assets/kernel/k_vecs.png){ align="center" alt="k_vecs" width="70%" }
</figure>

```cpp
K_vec k_vecs[NUM_VECS_PER_THREAD]
Expand Down Expand Up @@ -372,20 +364,14 @@ title: vLLM Paged Attention

<figure markdown="span">
![](../../assets/kernel/value.png){ align="center" alt="value" width="70%" }
<figcaption>
</figcaption>
</figure>

<figure markdown="span">
![](../../assets/kernel/logits_vec.png){ align="center" alt="logits_vec" width="50%" }
<figcaption>
</figcaption>
</figure>

<figure markdown="span">
![](../../assets/kernel/v_vec.png){ align="center" alt="v_vec" width="70%" }
<figcaption>
</figcaption>
</figure>

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