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, 17:15 UTC
[Kernel][TPU][ragged-paged-attn] vLLM code change for PR#8896 (#15659)
Signed-off-by: Yarong Mu <ymu@google.com>
yarongmu-google committed
Mar 28, 2025 at 21:13 UTC
7c1f7600248a0a0497a5c512ef0ee262577c5f7a
4 files changed
+37
-37
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.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"
20
+torch @ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch-2.8.0.dev20250328-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.dev20250328-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.dev20250328-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.dev20250328-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.dev20250328-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.dev20250328-cp311-cp311-linux_x86_64.whl ; python_version == "3.11"
vllm/v1/attention/backends/pallas.py
+22
-21
@@ -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 * 2, head_size)
45
46
@staticmethod
47
def swap_blocks(
@@ -132,7 +132,7 @@ class PallasAttentionBackendImpl(AttentionImpl):
132
query: torch.Tensor,
133
key: torch.Tensor,
134
value: torch.Tensor,
135
- kv_cache: tuple[torch.Tensor, torch.Tensor],
135
+ kv_cache: torch.Tensor,
136
attn_metadata: PallasMetadata,
137
output: Optional[torch.Tensor] = None,
138
) -> torch.Tensor:
@@ -142,14 +142,13 @@ 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 * 2, head_size]
146
attn_metadata: Metadata for attention.
147
Returns:
148
shape = [num_tokens, num_heads * head_size]
149
"""
150
# For determine_available_memory case.
152
- if kv_cache[0].numel() == 0:
151
+ if kv_cache.numel() == 0:
152
if output is None:
153
output = torch.ones_like(query)
154
return output
@@ -158,15 +157,13 @@ class PallasAttentionBackendImpl(AttentionImpl):
157
num_tokens, hidden_size = query.shape
158
query = query.view(num_tokens, self.num_heads, self.head_size)
159
161
- key_cache, value_cache = kv_cache
162
- if kv_cache[0].numel() > 0:
160
+ if kv_cache.numel() > 0:
161
slot_mapping = attn_metadata.slot_mapping
164
- write_to_kv_cache(key, value, key_cache, value_cache, slot_mapping)
162
+ write_to_kv_cache(key, value, kv_cache, slot_mapping)
163
164
output = torch.ops.xla.ragged_paged_attention(
165
query,
168
- key_cache,
169
- value_cache,
166
+ kv_cache,
167
attn_metadata.context_lens,
168
attn_metadata.block_tables,
169
attn_metadata.query_start_loc,
@@ -183,23 +180,27 @@ class PallasAttentionBackendImpl(AttentionImpl):
180
def write_to_kv_cache(
181
key: torch.Tensor,
182
value: torch.Tensor,
186
- key_cache: torch.Tensor,
187
- value_cache: torch.Tensor,
183
+ kv_cache: torch.Tensor,
184
slot_mapping: torch.Tensor,
185
) -> None:
186
""" Write the key and values to the KV cache.
187
188
Args:
189
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]
190
+ value: shape = [num_tokens, num_kv_heads * head_size]
191
+ kv_cache = [num_blocks, block_size, num_kv_heads * 2, head_size]
192
193
"""
199
- torch.ops.xla.dynamo_set_buffer_donor_(key_cache, True)
200
- torch.ops.xla.dynamo_set_buffer_donor_(value_cache, True)
194
+ _, _, num_combined_kv_heads, head_size = kv_cache.shape
195
+ num_kv_heads = num_combined_kv_heads // 2
196
202
- key_cache = key_cache.flatten(0, 1)
203
- value_cache = value_cache.flatten(0, 1)
204
- key_cache.index_copy_(0, slot_mapping, key)
205
- value_cache.index_copy_(0, slot_mapping, value)
197
+ key = key.view(-1, num_kv_heads, head_size)
198
+ value = value.view(-1, num_kv_heads, head_size)
199
+
200
+ kv = torch.cat([key, value], axis=-1).reshape(-1, num_combined_kv_heads,
201
+ head_size)
202
+
203
+ torch.ops.xla.dynamo_set_buffer_donor_(kv_cache, True)
204
+
205
+ kv_cache = kv_cache.flatten(0, 1)
206
+ kv_cache.index_copy_(0, slot_mapping, kv)
vllm/v1/worker/tpu_model_runner.py
+5
-6
@@ -861,12 +861,11 @@ class TPUModelRunner:
861
kv_cache_spec.num_kv_heads, kv_cache_spec.head_size)
862
dtype = kv_cache_spec.dtype
863
864
- tpu_k_cache = torch.zeros(kv_cache_shape,
865
- dtype=dtype,
866
- device=self.device)
867
- tpu_v_cache = torch.zeros_like(tpu_k_cache)
864
+ tpu_kv_cache = torch.zeros(kv_cache_shape,
865
+ dtype=dtype,
866
+ device=self.device)
867
869
- kv_caches[layer_name] = (tpu_k_cache, tpu_v_cache)
868
+ kv_caches[layer_name] = tpu_kv_cache
869
else:
870
raise NotImplementedError
871
@@ -893,7 +892,7 @@ class ModelWrapperV1(nn.Module):
892
self,
893
input_ids: torch.Tensor,
894
positions: torch.Tensor,
896
- kv_caches: list[tuple[torch.Tensor, torch.Tensor]],
895
+ kv_caches: list[torch.Tensor],
896
inputs_embeds: Optional[torch.Tensor] = None,
897
) -> torch.Tensor:
898
"""Executes the forward pass of the model.
vllm/v1/worker/tpu_worker.py
+4
-4
@@ -136,10 +136,10 @@ class TPUWorker:
136
137
# Use an empty tensor instead of `None`` to force Dynamo to pass
138
# it by reference, rather by specializing on the value ``None``.
139
- tpu_k_cache = torch.tensor([], dtype=dtype, device=self.device)
140
- tpu_v_cache = torch.tensor([], dtype=dtype, device=self.device)
141
-
142
- kv_caches[layer_name] = (tpu_k_cache, tpu_v_cache)
139
+ tpu_kv_cache = torch.tensor([],
140
+ dtype=dtype,
141
+ device=self.device)
142
+ kv_caches[layer_name] = tpu_kv_cache
143
else:
144
raise NotImplementedError
145