Repository: vllm-project/vllm

Issue:
Array Index Overflow for Large #Blocks
Suppose we have 64K blocks, and the key cache shape of a single layer is `[64K, 32, 16, 16, 8]` (LLaMA-7B). For slot mapping, the maximum value is `64K * 32 * 16 * 16 * 8 = 4194304K > INT_MAX`, which overflow.

## Env
`PyTorch 2.0.1 + vLLM 0.2.0 + cuda-11.7`

## Simple reproduce
We can reproduce this by changing the parameters of `tests/kernels/test_cache.py` to the following
```
DTYPES = [torch.half]
NUM_TOKENS = [7]  # Arbitrary values for testing
NUM_LAYERS = [1]  # Arbitrary values for testing
NUM_HEADS = [32]  # Arbitrary values for testing
HEAD_SIZES = [128]
BLOCK_SIZES = [16]
NUM_BLOCKS = [64000]  # Arbitrary values for testing
NUM_MAPPINGS = [32]  # Arbitrary values for testing
SEEDS = [0]
```
The `reshape_and_cache` kernel cannot pass the test.
If we change the dtype of `slot_mapping` to int64, we can pass the test.

Besides, change `NUM_BLOCKS` to 64K in `tests/kernels/test_attention.py`, I find it cannot pass the test, either.
It reports "illegal memory access".

@zhuohan123 @WoosukKwon Can you help look into this?




Discussion:

Resolving patch:
diff --git a/vllm/worker/worker.py b/vllm/worker/worker.py
--- a/vllm/worker/worker.py
+++ b/vllm/worker/worker.py
@@ -301,7 +301,7 @@ def _prepare_inputs(
                                         dtype=torch.long,
                                         device="cuda")
         slot_mapping_tensor = torch.tensor(padded_slot_mapping,
-                                           dtype=torch.int,
+                                           dtype=torch.long,
                                            device="cuda")
         context_lens_tensor = torch.tensor(context_lens,
                                            dtype=torch.int,


Test patch:
diff --git a/tests/kernels/test_attention.py b/tests/kernels/test_attention.py
--- a/tests/kernels/test_attention.py
+++ b/tests/kernels/test_attention.py
@@ -13,7 +13,7 @@
 # This will change depending on the compute capability.
 # - 512 as a buffer
 MAX_SEQ_LEN = get_max_shared_memory_bytes() // FLOAT32_BYTES - 512
-NUM_BLOCKS = 128  # Arbitrary values for testing
+NUM_BLOCKS = 40000  # Arbitrary values for testing
 PARTITION_SIZE = 512
 
 DTYPES = [torch.half, torch.bfloat16, torch.float]
diff --git a/tests/kernels/test_cache.py b/tests/kernels/test_cache.py
--- a/tests/kernels/test_cache.py
+++ b/tests/kernels/test_cache.py
@@ -6,13 +6,13 @@
 from vllm import cache_ops
 
 DTYPES = [torch.half, torch.bfloat16, torch.float]
-NUM_TOKENS = [7, 83, 2048]  # Arbitrary values for testing
-NUM_LAYERS = [5]  # Arbitrary values for testing
+NUM_TOKENS = [83]  # Arbitrary values for testing
+NUM_LAYERS = [1]  # Arbitrary values for testing
 NUM_HEADS = [8]  # Arbitrary values for testing
 HEAD_SIZES = [64, 80, 96, 112, 128, 256]
 BLOCK_SIZES = [8, 16, 32]
-NUM_BLOCKS = [1024]  # Arbitrary values for testing
-NUM_MAPPINGS = [32, 256]  # Arbitrary values for testing
+NUM_BLOCKS = [1024, 36000]  # Arbitrary values for testing
+NUM_MAPPINGS = [256]  # Arbitrary values for testing
 SEEDS = [0]
 
 
@@ -69,9 +69,9 @@ def test_copy_blocks(
     for src, dsts in block_mapping.items():
         for dst in dsts:
             for cloned_key_cache in cloned_key_caches:
-                cloned_key_cache[dst] = cloned_key_cache[src]
+                cloned_key_cache[dst].copy_(cloned_key_cache[src])
             for cloned_value_cache in cloned_value_caches:
-                cloned_value_cache[dst] = cloned_value_cache[src]
+                cloned_value_cache[dst].copy_(cloned_value_cache[src])
 
     # Compare the results.
     for key_cache, cloned_key_cache in zip(key_caches, cloned_key_caches):
@@ -106,7 +106,7 @@ def test_reshape_and_cache(
     # Create a random slot mapping.
     num_slots = block_size * num_blocks
     slot_mapping = random.sample(range(num_slots), num_tokens)
-    slot_mapping = torch.tensor(slot_mapping, dtype=torch.int, device="cuda")
+    slot_mapping = torch.tensor(slot_mapping, dtype=torch.long, device="cuda")
 
     qkv = torch.randn(num_tokens,
                       3,