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 11, 03:45 UTC
[Doc] Clarify FP8 KV cache computation workflow (#31071)
Signed-off-by: westers <steve.westerhouse@origami-analytics.com>
Steve Westerhouse committed
Dec 21, 2025 at 18:41 UTC
9d701e90d82f091395dc32634f2a26132c167bfd
2 files changed
+31
-21
docs/design/paged_attention.md
+21
-21
@@ -139,18 +139,18 @@ token data.
139
const scalar_t* q_ptr = q + seq_idx * q_stride + head_idx * HEAD_SIZE;
140
```
141
142
-<figure markdown="span">
143
- { align="center" alt="query" width="70%" }
144
-</figure>
142
+<p align="center">
143
+ <img src="../assets/design/paged_attention/query.png" alt="query" width="70%" />
144
+</p>
145
146
Each thread defines its own `q_ptr` which points to the assigned
147
query token data on global memory. For example, if `VEC_SIZE` is 4
148
and `HEAD_SIZE` is 128, the `q_ptr` points to data that contains
149
total of 128 elements divided into 128 / 4 = 32 vecs.
150
151
-<figure markdown="span">
152
- { align="center" alt="q_vecs" width="70%" }
153
-</figure>
151
+<p align="center">
152
+ <img src="../assets/design/paged_attention/q_vecs.png" alt="q_vecs" width="70%" />
153
+</p>
154
155
```cpp
156
__shared__ Q_vec q_vecs[THREAD_GROUP_SIZE][NUM_VECS_PER_THREAD];
@@ -187,9 +187,9 @@ key token at different iterations. As shown above, that `k_ptr`
187
points to key token data based on `k_cache` at assigned block,
188
assigned head and assigned token.
189
190
-<figure markdown="span">
191
- { align="center" alt="key" width="70%" }
192
-</figure>
190
+<p align="center">
191
+ <img src="../assets/design/paged_attention/key.png" alt="key" width="70%" />
192
+</p>
193
194
The diagram above illustrates the memory layout for key data. It
195
assumes that the `BLOCK_SIZE` is 16, `HEAD_SIZE` is 128, `x` is
@@ -202,9 +202,9 @@ iterations. Inside each rectangle, there are a total 32 vecs (128
202
elements for one token) that will be processed by 2 threads (one
203
thread group) separately.
204
205
-<figure markdown="span">
206
- { align="center" alt="k_vecs" width="70%" }
207
-</figure>
205
+<p align="center">
206
+ <img src="../assets/design/paged_attention/k_vecs.png" alt="k_vecs" width="70%" />
207
+</p>
208
209
```cpp
210
K_vec k_vecs[NUM_VECS_PER_THREAD]
@@ -361,17 +361,17 @@ later steps. Now, it should store the normalized softmax result of
361
362
## Value
363
364
-<figure markdown="span">
365
- { align="center" alt="value" width="70%" }
366
-</figure>
364
+<p align="center">
365
+ <img src="../assets/design/paged_attention/value.png" alt="value" width="70%" />
366
+</p>
367
368
-<figure markdown="span">
369
- { align="center" alt="logits_vec" width="50%" }
370
-</figure>
368
+<p align="center">
369
+ <img src="../assets/design/paged_attention/logits_vec.png" alt="logits_vec" width="50%" />
370
+</p>
371
372
-<figure markdown="span">
373
- { align="center" alt="v_vec" width="70%" }
374
-</figure>
372
+<p align="center">
373
+ <img src="../assets/design/paged_attention/v_vec.png" alt="v_vec" width="70%" />
374
+</p>
375
376
Now we need to retrieve the value data and perform dot multiplication
377
with `logits`. Unlike query and key, there is no thread group
docs/features/quantization/quantized_kvcache.md
+10
@@ -17,6 +17,16 @@ The E4M3 format offers higher precision compared to E5M2. However, due to its sm
17
18
For now, only per-tensor (scalar) scaling factors are supported. Development is ongoing to support scaling factors of a finer granularity (e.g. per-channel).
19
20
+### How FP8 KV Cache Works
21
+
22
+The FP8 KV cache implementation follows this workflow:
23
+
24
+1. **Storage**: Key and Value tensors are quantized to FP8 format using scaling factors before being stored in the KV cache
25
+2. **Retrieval**: When needed for attention computation, cached KV tensors are dequantized back to higher precision (FP16/BF16)
26
+3. **Attention**: The attention-value multiplication (softmax output × V) is performed using the dequantized higher-precision V tensor
27
+
28
+This means the final attention computation operates on dequantized values, not FP8 tensors. The quantization reduces memory usage during storage but maintains computation accuracy by using higher precision during the actual attention operations.
29
+
30
### Performance Impact
31
32
The current FP8 KV cache implementation primarily benefits throughput by allowing approximately double the amount of space for KV cache allocation. This enables either: