Repository: vllm-project/vllm
Pull request: [BugFix] Fix NaN errors in paged attention kernel
Fixes #641 
This PR fixes the paged attention kernel. Currently, the kernel computes `attn_weight * value` for all tokens in a value block, even if some of them are not included in the context. It is generally acceptable since the `attn_weight` for those tokens is 0, but this causes errors when the tokens contain NaNs (since 0 * NaN is NaN). The PR solves this by explicitly setting the values of those tokens as 0.

Problem statement:
RuntimeError: probability tensor contains either `inf`, `nan` or element < 0
Hello everyone, I always got this error for Baichuan and LLaMA models. And I found it's caused by the **single_query_cached_kv_attention** method in vllm\model_executor\layers\\**attention.py**. After calling of this method, the hidden output has some **rows of "nan"**. How can I fix this? Thanks!

Still have such errors even after installing xformers from source.

This is my code:
```python
from vllm import LLM, SamplingParams
#from vllm.transformers_utils.configs.baichuan import BaiChuanConfig

prompts = [
        "Hello, my name is",
    "The president of the United States is",
    "The capital of France is",
    "The future of AI is",
]

sampling_params = SamplingParams(temperature=1, top_p=0.95)

llm = LLM(
        model="/.../Baichuan-7b",
        trust_remote_code=True,
        dtype='float16',
        gpu_memory_utilization=0.85,
        tokenizer_mode="slow"
    )
outputs = llm.generate(prompts, sampling_params)

for output in outputs:
    prompt = output.prompt
    generated_text = output.outputs[0].text
    print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
```
and this is my python environment:
```
accelerate                0.21.0
aiofiles                  23.1.0
aiohttp                   3.8.5
aiosignal                 1.3.1
altair                    5.0.1
annotated-types           0.5.0
anyio                     3.7.1
appdirs                   1.4.4
argon2-cffi               21.3.0
argon2-cffi-bindings      21.2.0
arrow                     1.2.3
asttokens                 2.2.1
async-lru                 2.0.3
async-timeout             4.0.2
attrs                     23.1.0
Babel                     2.12.1
backcall                  0.2.0
beautifulsoup4            4.12.2
bleach                    6.0.0
blinker                   1.6.2
boltons                   23.0.0
brotlipy                  0.7.0
certifi                   2022.12.7
cffi                      1.15.1
charset-normalizer        2.0.4
click                     8.1.6
cmake                     3.27.0
comm                      0.1.3
conda                     23.3.1
conda-content-trust       0.1.3
conda-package-handling    2.0.2
conda_package_streaming   0.7.0
contourpy                 1.1.0
cryptography              39.0.1
cycler                    0.11.0
datasets                  2.14.0
debugpy                   1.6.7
decorator                 5.1.1
defusedxml                0.7.1
dill                      0.3.7
distlib                   0.3.7
docker-pycreds            0.4.0
editables                 0.5
exceptiongroup            1.1.2
executing                 1.2.0
fastapi                   0.100.0
fastjsonschema            2.18.0
ffmpy                     0.3.1
filelock                  3.12.2
Flask                     2.3.2
fonttools                 4.41.1
fqdn                      1.5.1
frozenlist                1.4.0
fsspec                    2023.6.0
gitdb                     4.0.10
GitPython                 3.1.32
gradio                    3.35.2
gradio_client             0.2.10
grpcio                    1.56.2
h11                       0.14.0
hatchling                 1.18.0
httpcore                  0.17.3
httpx                     0.24.1
huggingface-hub           0.16.4
idna                      3.4
ipykernel                 6.24.0
ipython                   8.14.0
ipython-genutils          0.2.0
ipywidgets                8.0.7
isoduration               20.11.0
itsdangerous              2.1.2
jedi                      0.18.2
jieba                     0.42.1
Jinja2                    3.1.2
joblib                    1.3.1
json5                     0.9.14
jsonpatch                 1.32
jsonpointer               2.1
jsonschema                4.18.4
jsonschema-specifications 2023.7.1
jupyter                   1.0.0
jupyter_client            8.3.0
jupyter-console           6.6.3
jupyter_core              5.3.1
jupyter-events            0.6.3
jupyter-lsp               2.2.0
jupyter_server            2.7.0
jupyter_server_terminals  0.4.4
jupyterlab                4.0.3
jupyterlab-pygments       0.2.2
jupyterlab_server         2.24.0
jupyterlab-widgets        3.0.8
kiwisolver                1.4.4
linkify-it-py             2.0.2
lit                       16.0.6
markdown-it-py            2.2.0
markdown2                 2.4.10
MarkupSafe                2.1.3
matplotlib                3.7.2
matplotlib-inline         0.1.6
mdit-py-plugins           0.3.3
mdurl                     0.1.2
mistune                   3.0.1
mpmath                    1.3.0
msgpack                   1.0.5
multidict                 6.0.4
multiprocess              0.70.15
mypy-extensions           1.0.0
nbclient                  0.8.0
nbconvert                 7.7.2
nbformat                  5.9.1
nest-asyncio              1.5.6
networkx                  3.1
nh3                       0.2.14
ninja                     1.11.1
nltk                      3.8.1
notebook                  7.0.0
notebook_shim             0.2.3
numpy                     1.25.1
nvidia-cublas-cu11        11.10.3.66
nvidia-cuda-cupti-cu11    11.7.101
nvidia-cuda-nvrtc-cu11    11.7.99
nvidia-cuda-runtime-cu11  11.7.99
nvidia-cudnn-cu11         8.5.0.96
nvidia-cufft-cu11         10.9.0.58
nvidia-curand-cu11        10.2.10.91
nvidia-cusolver-cu11      11.4.0.1
nvidia-cusparse-cu11      11.7.4.91
nvidia-nccl-cu11          2.14.3
nvidia-nvtx-cu11          11.7.91
orjson                    3.9.2
overrides                 7.3.1
packaging                 23.0
pandas                    2.0.3
pandocfilters             1.5.0
parso                     0.8.3
pathspec                  0.11.1
pathtools                 0.1.2
peft                      0.4.0
pexpect                   4.8.0
pickleshare               0.7.5
Pillow                    10.0.0
pip                       23.0.1
platformdirs              3.9.1
pluggy                    1.0.0
prometheus-client         0.17.1
prompt-toolkit            3.0.39
protobuf                  4.23.4
psutil                    5.9.5
ptyprocess                0.7.0
pure-eval                 0.2.2
pyarrow                   12.0.1
pycosat                   0.6.4
pycparser                 2.21
pydantic                  1.10.12
pydantic_core             2.3.0
pydub                     0.25.1
Pygments                  2.15.1
pyOpenSSL                 23.0.0
pyparsing                 3.0.9
pyre-extensions           0.0.29
PySocks                   1.7.1
python-dateutil           2.8.2
python-json-logger        2.0.7
python-multipart          0.0.6
pytz                      2023.3
PyYAML                    6.0.1
pyzmq                     25.1.0
qtconsole                 5.4.3
QtPy                      2.3.1
ray                       2.6.1
referencing               0.30.0
regex                     2023.6.3
requests                  2.28.1
rfc3339-validator         0.1.4
rfc3986-validator         0.1.1
rich                      13.4.2
rouge-chinese             1.0.3
rpds-py                   0.9.2
ruamel.yaml               0.17.21
ruamel.yaml.clib          0.2.6
safetensors               0.3.1
semantic-version          2.10.0
Send2Trash                1.8.2
sentencepiece             0.1.99
sentry-sdk                1.28.1
setproctitle              1.3.2
setuptools                65.6.3
shortuuid                 1.0.11
six                       1.16.0
smmap                     5.0.0
sniffio                   1.3.0
soupsieve                 2.4.1
stack-data                0.6.2
starlette                 0.27.0
svgwrite                  1.4.3
sympy                     1.12
terminado                 0.17.1
tinycss2                  1.2.1
tokenizers                0.13.3
tomli                     2.0.1
toolz                     0.12.0
torch                     2.0.1
tornado                   6.3.2
tqdm                      4.65.0
traitlets                 5.9.0
transformers              4.31.0
triton                    2.0.0
trl                       0.4.7
trove-classifiers         2023.7.6
typing_extensions         4.7.1
typing-inspect            0.9.0
tzdata                    2023.3
uc-micro-py               1.0.2
uri-template              1.3.0
urllib3                   1.26.15
uvicorn                   0.23.1
virtualenv                20.24.2
vllm                      0.1.2       /.../feng/OpenSource/vllm
wandb                     0.15.7
wavedrom                  2.0.3.post3
wcwidth                   0.2.6
webcolors                 1.13
webencodings              0.5.1
websocket-client          1.6.1
websockets                11.0.3
Werkzeug                  2.3.6
wheel                     0.38.4
widgetsnbextension        4.0.8
xformers                  0.0.20
xxhash                    3.2.0
yarl                      1.9.2
zstandard                 0.19.0
```

and my GPU info:
```
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 510.73.08    Driver Version: 510.73.08    CUDA Version: 11.6     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  GRID V100S-32Q      On   | 00000000:02:01.0 Off |                    0 |
| N/A   N/A    P0    N/A /  N/A |      0MiB / 32768MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
                                                                               
+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|  No running processes found                                                 |
+-----------------------------------------------------------------------------+
```

Hints from issue discussion:
Same problem, in the end I find that KVCache is construst with torch.empty filling with uninitialized values, for some reason the uninitialized values are involved in computing.
Here in [vllm](https://github.com/vllm-project/vllm/tree/main)/[vllm](https://github.com/vllm-project/vllm/tree/main/vllm)/[worker](https://github.com/vllm-project/vllm/tree/main/vllm/worker)
/cache_engine.py

`    def allocate_gpu_cache(self) -> List[KVCache]:
        gpu_cache: List[KVCache] = []
        key_block_shape = self.get_key_block_shape()
        value_block_shape = self.get_value_block_shape()
        for _ in range(self.num_layers):
            key_blocks = torch.empty(
                size=(self.num_gpu_blocks, *key_block_shape),
                dtype=self.dtype,
                device="cuda",
            )
            value_blocks = torch.empty(
                size=(self.num_gpu_blocks, *value_block_shape),
                dtype=self.dtype,
                device="cuda",
            )
            gpu_cache.append((key_blocks, value_blocks))
        return gpu_cache`

When I change torch.empty to torch.zeros, the model does not output nan any more, but I believe it is a bug. Because of the uninitialized values, the texts generated by gpt2 with vllm and huggingface are different.

I am running gpt2 in docker image kevinng77/vllm with one GPU T4-8C.
Thanks a lot @GhostofAdam , will try your method.
@GhostofAdam , tried your method, and it's great that it works. 
The generation is different from huggingface. vLLM's generation is much shorter than huggingface, and tends to have many duplication.
Have this error as well, and changing `empty` -> `zeros` doesn't help(
> Have this error as well, and changing `empty` -> `zeros` doesn't help(

Did you do the same thing with function allocate_cpu_cache? 
Changing `empty` to `zeros` can reduce the chance to trigger the bug. What is better is to change `empty` to `full(..., math.nan, ...)`, because it triggers this bug always.


I believe `single_query_cached_kv_attention_kernel` does not properly check boundary in the context length dimension when block size is larger than one.

Setting block size to one can suppress this bug.
We masked out values in `logits` where the token index is larger than context length, which could avoid corrupted `logits` due to `nan` from uninitialized `k_cache`, which is good.
https://github.com/vllm-project/vllm/blob/d1744376ae9fdbfa6a2dc763e1c67309e138fa3d/csrc/attention/attention_kernels.cu#L186-L189

However, we did not mask out values in `v_vec` where the token index is larger than context length. As a result the following `dot` call is incorrect.

https://github.com/vllm-project/vllm/blob/d1744376ae9fdbfa6a2dc763e1c67309e138fa3d/csrc/attention/attention_kernels.cu#L264

`0 (from logits_vec) * nan (from v_vec)` is `nan`, unfortunately.
I think it's inefficient to mask `v_vec` in `single_query_cached_kv_attention_kernel`, because the same `v_vec` will be used multiple times among warps. Instead, we should pad `v_cache` with finite values.
I can confirm that changing `torch.empty` to `torch.zeros` is only necessary for `v_cache` . It means `k_cache` is not buggy.

Commit message: Minor

Patch to review:
diff --git a/csrc/attention/attention_kernels.cu b/csrc/attention/attention_kernels.cu
index 568d1fb1ad24..7a656bb9cd8a 100644
--- a/csrc/attention/attention_kernels.cu
+++ b/csrc/attention/attention_kernels.cu
@@ -246,6 +246,8 @@ __global__ void single_query_cached_kv_attention_kernel(
     accs[i] = 0.f;
   }
 
+  scalar_t zero_value;
+  zero(zero_value);
   for (int block_idx = warp_idx; block_idx < num_blocks; block_idx += NUM_WARPS) {
     const int physical_block_number = block_table[block_idx];
     const int physical_block_offset = (lane % NUM_V_VECS_PER_ROW) * V_VEC_SIZE;
@@ -261,6 +263,14 @@ __global__ void single_query_cached_kv_attention_kernel(
       if (row_idx < HEAD_SIZE) {
         const int offset = row_idx * BLOCK_SIZE + physical_block_offset;
         V_vec v_vec = *reinterpret_cast<const V_vec*>(v_ptr + offset);
+        // NOTE(woosuk): When v_vec contains the tokens that are out of the context,
+        // we should explicitly zero out the values since they may contain NaNs.
+        // See https://github.com/vllm-project/vllm/issues/641#issuecomment-1682544472
+        scalar_t* v_vec_ptr = reinterpret_cast<scalar_t*>(&v_vec);
+#pragma unroll
+        for (int j = 0; j <= V_VEC_SIZE; j++) {
+          v_vec_ptr[j] = token_idx + j < context_len ? v_vec_ptr[j] : zero_value;
+        }
         accs[i] += dot(logits_vec, v_vec);
       }
     }
diff --git a/csrc/attention/dtype_bfloat16.cuh b/csrc/attention/dtype_bfloat16.cuh
index 70fd064bb161..0f1035200627 100644
--- a/csrc/attention/dtype_bfloat16.cuh
+++ b/csrc/attention/dtype_bfloat16.cuh
@@ -420,4 +420,9 @@ inline __device__ void from_float(bf16_8_t& dst, Float8_ src) {
 #endif
 }
 
+// Zero-out a variable.
+inline __device__ void zero(__nv_bfloat16& dst) {
+  from_float(dst, 0.f);
+}
+
 } // namespace vllm
diff --git a/csrc/attention/dtype_float16.cuh b/csrc/attention/dtype_float16.cuh
index 6ffc30cd96e7..e67921128d52 100644
--- a/csrc/attention/dtype_float16.cuh
+++ b/csrc/attention/dtype_float16.cuh
@@ -390,11 +390,6 @@ inline __device__ float sum(uint4 v) {
   return sum(c);
 }
 
-// Zero-out a vector.
-inline __device__ void zero(uint16_t& dst) {
-  dst = uint16_t(0);
-}
-
 // From float32 to float16.
 inline __device__ void from_float(uint16_t& dst, float src) {
   dst = float_to_half(src);
@@ -441,4 +436,9 @@ inline __device__ Float8_ to_float(uint4 u) {
   return tmp;
 }
 
+// Zero-out a variable.
+inline __device__ void zero(uint16_t& dst) {
+  dst = uint16_t(0);
+}
+
 } // namespace vllm
diff --git a/csrc/attention/dtype_float32.cuh b/csrc/attention/dtype_float32.cuh
index 960cf48e2643..4006f81e488d 100644
--- a/csrc/attention/dtype_float32.cuh
+++ b/csrc/attention/dtype_float32.cuh
@@ -265,4 +265,9 @@ inline __device__ Float8_ to_float(Float8_ u) {
   return u;
 }
 
+// Zero-out a variable.
+inline __device__ void zero(float& dst) {
+  dst = float(0.f);
+}
+
 } // namespace vllm


Code review comments:
diff_hunk: @@ -261,6 +263,14 @@ __global__ void single_query_cached_kv_attention_kernel(
       if (row_idx < HEAD_SIZE) {
         const int offset = row_idx * BLOCK_SIZE + physical_block_offset;
         V_vec v_vec = *reinterpret_cast<const V_vec*>(v_ptr + offset);
+        // NOTE(woosuk): When v_vec contains the tokens that are out of the context,
+        // we should explicitly zero out the values since they may contain NaNs.
+        // See https://github.com/vllm-project/vllm/issues/641#issuecomment-1682544472
+        scalar_t* v_vec_ptr = reinterpret_cast<scalar_t*>(&v_vec);
+#pragma unroll
+        for (int j = 0; j <= V_VEC_SIZE; j++) {
+          v_vec_ptr[j] = token_idx + j < context_len ? v_vec_ptr[j] : zero_value;
+        }
path: csrc/attention/attention_kernels.cu
text: @user1:
`token_idx + j < context_len` will only be false when `block_idx == num_blocks - 1`. Will adding a for loop here for every iteration hurt performance?

@author:
Good idea. I benchmarked the kernel.

* Before this PR:
```
torch.float16: 8 seqs, 30 context len, 8.6 us
torch.float16: 8 seqs, 498 context len, 52.5 us
torch.float16: 8 seqs, 2033 context len, 179.2 us
torch.float16: 8 seqs, 8168 context len, 678.8 us
torch.float16: 32 seqs, 30 context len, 19.6 us
torch.float16: 32 seqs, 498 context len, 180.2 us
torch.float16: 32 seqs, 2033 context len, 663.7 us
torch.float16: 32 seqs, 8168 context len, 2641.2 us
torch.float16: 128 seqs, 30 context len, 67.9 us
torch.float16: 128 seqs, 498 context len, 686.7 us
torch.float16: 128 seqs, 2033 context len, 2682.7 us
torch.float16: 128 seqs, 8168 context len, 10954.4 us
```

* Before adding `if(block_idx == num_blocks -1)`:
```
torch.float16: 8 seqs, 30 context len, 9.1 us
torch.float16: 8 seqs, 498 context len, 53.0 us
torch.float16: 8 seqs, 2033 context len, 178.5 us
torch.float16: 8 seqs, 8168 context len, 683.0 us
torch.float16: 32 seqs, 30 context len, 20.6 us
torch.float16: 32 seqs, 498 context len, 183.4 us
torch.float16: 32 seqs, 2033 context len, 668.3 us
torch.float16: 32 seqs, 8168 context len, 2674.8 us
torch.float16: 128 seqs, 30 context len, 70.9 us
torch.float16: 128 seqs, 498 context len, 698.1 us
torch.float16: 128 seqs, 2033 context len, 2706.8 us
torch.float16: 128 seqs, 8168 context len, 11026.7 us
```

* After adding `if(block_idx == num_blocks -1)`:
```
torch.float16: 8 seqs, 30 context len, 9.1 us
torch.float16: 8 seqs, 498 context len, 52.7 us
torch.float16: 8 seqs, 2033 context len, 181.4 us
torch.float16: 8 seqs, 8168 context len, 685.7 us
torch.float16: 32 seqs, 30 context len, 20.1 us
torch.float16: 32 seqs, 498 context len, 180.9 us
torch.float16: 32 seqs, 2033 context len, 664.7 us
torch.float16: 32 seqs, 8168 context len, 2661.6 us
torch.float16: 128 seqs, 30 context len, 73.2 us
torch.float16: 128 seqs, 498 context len, 694.1 us
torch.float16: 128 seqs, 2033 context len, 2690.4 us
torch.float16: 128 seqs, 8168 context len, 10963.3 us
```

Adding the if statement slightly improves the performance when the number of blocks is large. In general, I believe the overhead (which is < 5%) is now acceptable.

@user1:
Sounds good! Let's just add this if statement since it looks like a pure performance optimization?

@user2:
I feel bad for adding a condition in the inner loop.
Have you compared it with the approach of simply set the page to zeros when releasing a page? 

@user2:
As I mentioned here, padding seems more reasonable https://github.com/vllm-project/vllm/issues/641#issuecomment-1682619534

@author:
Hi @user2, thanks for the suggestion. I considered your solution that does not modify the kernel. While it is simple and never affects the performance, I felt handling NANs in the kernel would be a better design (as long as the impact on the performance is acceptable). The reason is, changing the kernel can 1) benefit all the kernel users besides vLLM (e.g., TGI), and 2) guarantee that previous requests cannot affect the future requests (e.g., when previous requests somehow store NaNs in the KV cache).

@user2:
I suggest to fill the page with finite values when "releasing" a page. Because the CUDA operator is non-blocking, it would not affect the latency.

Another idea is to change the shape of v_cache like k_cache's, so that we can use the same approach to filter out out of bound values like k_cache does. In fact, because we can skip some loops, I expect it could have a chance to improve the computational performance a little bit if we change the shape of v_cache.


@author:
> I suggest to fill the page with finite values when "releasing" a page. Because the CUDA operator is non-blocking, it would not affect the latency.

I feel there's a trade off between the complexity of the system logic and the complexity of the kernel. If you do not handle the boundary inside the kernel, you are delegating the burden to the system. Overall, for now, I'd prefer lowering the complexity of the system logic over lowering the complexity of the kernel because the performance degradation due to the modification to the kernel is small. However, if we augment the kernel with new features in the future and find that removing the boundary check is critical to performance, we may revisit this and eventually pivot to your solution.

> Another idea is to change the shape of v_cache like k_cache's, so that we can use the same approach to filter out out of bound values like k_cache does. In fact, because we can skip some loops, I expect it could have a chance to improve the computational performance a little bit if we change the shape of v_cache.

Thanks for the suggestion. There are indeed some wasted cycles due to the layout of the value cache, when a value block is not full. IIRC, the main reason for choosing the `[head_size, block_size]` layout rather than the `[block_size, head_size]` layout is that the former makes it easier to handle the head sizes that are not divisible by 32 (warp size) and not powers of two. For example, the head size is 112 in MPT-30B, 80 in OPT/BLOOM-3B, and 100 in OpenLLaMA-3B. I originally tried the `[block_size, head_size]` layout, but found that the memory access pattern became quite complex for such head sizes. The problem was solved when using the `[head_size, block_size]`: Because the block size is guaranteed to be a power of 2, the memory access pattern was quite simple.

That being said, we are open to contributions and hope to continuously improve vLLM. Please submit a PR if you found that changing the layout makes the kernel simpler and more efficient.

diff_hunk: @@ -420,4 +420,9 @@ inline __device__ void from_float(bf16_8_t& dst, Float8_ src) {
 #endif
 }
 
+// Zero-out a variable.
+inline __device__ void zero(__nv_bfloat16& dst) {
+  from_float(dst, 0.f);
path: csrc/attention/dtype_bfloat16.cuh
text: @user1:
Will this function become faster if we switch to:
```suggestion
  dst = CUDART_ZERO_BF16;
```
`CUDART_ZERO_BF16` is defined [here](https://docs.nvidia.com/cuda/cuda-math-api/group__CUDA__MATH__INTRINSIC__BFLOAT16__CONSTANTS.html#group__CUDA__MATH__INTRINSIC__BFLOAT16__CONSTANTS)

@author:
Nice catch! As `CUDA_ZERO_BF16` is introduced in CUDA 12.2, I instead used
```cpp
dst = __ushort_as_bfloat16((unsigned short)0x0000U);
```
which is the definition of `CUDA_ZERO_BF16`.

diff_hunk: @@ -265,4 +265,9 @@ inline __device__ Float8_ to_float(Float8_ u) {
   return u;
 }
 
+// Zero-out a variable.
+inline __device__ void zero(float& dst) {
+  dst = float(0.f);
path: csrc/attention/dtype_float32.cuh
text: @user1:
Nit: This `float` conversion seems redundant:
```suggestion
  dst = 0.f;
```

@author:
Good catch! Fixed.