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, 18:30 UTC
[V1][TPU] Change kv cache shape. (#15145)
Signed-off-by: Xiongfei Wei <isaacwxf23@gmail.com>
iefgnoix committed
Mar 19, 2025 at 12:16 UTC
b0e96aaebbfbe8e70478e4192a5a13864ffdefa6
2 files changed
+13
-16
requirements/tpu.txt
+6
-6
@@ -17,9 +17,9 @@ ray[data]
17
--find-links https://storage.googleapis.com/libtpu-releases/index.html
18
--find-links https://storage.googleapis.com/jax-releases/jax_nightly_releases.html
19
--find-links https://storage.googleapis.com/jax-releases/jaxlib_nightly_releases.html
20
-torch @ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch-2.8.0.dev20250314%2Bcxx11-cp39-cp39-linux_x86_64.whl ; python_version == "3.9"
21
-torch @ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch-2.8.0.dev20250314%2Bcxx11-cp310-cp310-linux_x86_64.whl ; python_version == "3.10"
22
-torch @ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch-2.8.0.dev20250314%2Bcxx11-cp311-cp311-linux_x86_64.whl ; python_version == "3.11"
23
-torch_xla[tpu, pallas] @ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.8.0.dev20250314%2Bcxx11-cp39-cp39-linux_x86_64.whl ; python_version == "3.9"
24
-torch_xla[tpu, pallas] @ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.8.0.dev20250314%2Bcxx11-cp310-cp310-linux_x86_64.whl ; python_version == "3.10"
25
-torch_xla[tpu, pallas] @ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.8.0.dev20250314%2Bcxx11-cp311-cp311-linux_x86_64.whl ; python_version == "3.11"
20
+torch @ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch-2.8.0.dev20250319-cp39-cp39-linux_x86_64.whl ; python_version == "3.9"
21
+torch @ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch-2.8.0.dev20250319-cp310-cp310-linux_x86_64.whl ; python_version == "3.10"
22
+torch @ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch-2.8.0.dev20250319-cp311-cp311-linux_x86_64.whl ; python_version == "3.11"
23
+torch_xla[tpu, pallas] @ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.8.0.dev20250319-cp39-cp39-linux_x86_64.whl ; python_version == "3.9"
24
+torch_xla[tpu, pallas] @ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.8.0.dev20250319-cp310-cp310-linux_x86_64.whl ; python_version == "3.10"
25
+torch_xla[tpu, pallas] @ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.8.0.dev20250319-cp311-cp311-linux_x86_64.whl ; python_version == "3.11"
vllm/v1/attention/backends/pallas.py
+7
-10
@@ -41,7 +41,7 @@ class PallasAttentionBackend(AttentionBackend):
41
num_kv_heads: int,
42
head_size: int,
43
) -> tuple[int, ...]:
44
- return (num_blocks, block_size, num_kv_heads, head_size)
44
+ return (num_blocks, block_size, num_kv_heads * head_size)
45
46
@staticmethod
47
def swap_blocks(
@@ -142,8 +142,8 @@ class PallasAttentionBackendImpl(AttentionImpl):
142
query: shape = [num_tokens, num_heads * head_size]
143
key: shape = [num_tokens, num_kv_heads * head_size]
144
value: shape = [num_tokens, num_kv_heads * head_size]
145
- kv_cache = ([num_blocks, block_size, num_kv_heads, head_size],
146
- [num_blocks, block_size, num_kv_heads, head_size])
145
+ kv_cache = ([num_blocks, block_size, num_kv_heads * head_size],
146
+ [num_blocks, block_size, num_kv_heads * head_size])
147
attn_metadata: Metadata for attention.
148
Returns:
149
shape = [num_tokens, num_heads * head_size]
@@ -157,8 +157,6 @@ class PallasAttentionBackendImpl(AttentionImpl):
157
assert layer._k_scale_float == 1.0 and layer._v_scale_float == 1.0
158
num_tokens, hidden_size = query.shape
159
query = query.view(num_tokens, self.num_heads, self.head_size)
160
- key = key.view(num_tokens, self.num_kv_heads, self.head_size)
161
- value = value.view(num_tokens, self.num_kv_heads, self.head_size)
160
161
key_cache, value_cache = kv_cache
162
if kv_cache[0].numel() > 0:
@@ -192,10 +190,10 @@ def write_to_kv_cache(
190
""" Write the key and values to the KV cache.
191
192
Args:
195
- key: shape = [num_tokens, num_kv_heads, head_size]
196
- value: shape = [num_tokens, num_kv_heads, head_size]
197
- k_cache = [num_blocks, block_size, num_kv_heads, head_size]
198
- v_cache = [num_blocks, block_size, num_kv_heads, head_size]
193
+ key: shape = [num_tokens, num_kv_heads * head_size]
194
+ value: shape = [num_tokens, num_kv_heads * head_size]
195
+ k_cache = [num_blocks, block_size, num_kv_heads * head_size]
196
+ v_cache = [num_blocks, block_size, num_kv_heads * head_size]
197
198
"""
199
torch.ops.xla.dynamo_set_buffer_donor_(key_cache, True)
@@ -203,6 +201,5 @@ def write_to_kv_cache(
201
202
key_cache = key_cache.flatten(0, 1)
203
value_cache = value_cache.flatten(0, 1)
206
- slot_mapping = slot_mapping.flatten()
204
key_cache.index_copy_(0, slot_mapping, key)
205
value_cache.index_copy_(0, slot_mapping, value)