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, 03:30 UTC
[TPU] kv cache update kernel supports dynamic grid (#20235)
Signed-off-by: Chengji Yao <chengjiyao@google.com>
Chengji Yao committed
Jul 1, 2025 at 23:33 UTC
7da296be04933cfc29031f5bd1ba7cd28f376faa
4 files changed
+42
-17
tests/v1/tpu/test_kv_cache_update_kernel.py
+6
-2
@@ -32,6 +32,7 @@ def test_kv_cache_update_kernel(page_size: int, combined_kv_head_num: int,
32
new_kv_xla = new_kv_cpu.to(torch_xla.device())
33
slice_lens = np.array([7, page_size, page_size, 1, 1, 1, 9],
34
dtype=np.int32)
35
+ num_kv_update_slices = len(slice_lens)
36
kv_cache_start_indices = np.array([
37
page_size * 2 - 7, page_size * 2, page_size * 3, page_size * 4 + 6,
38
page_size * 5 + 7, page_size * 6 + 8, page_size * 15 + 3
@@ -52,12 +53,15 @@ def test_kv_cache_update_kernel(page_size: int, combined_kv_head_num: int,
53
device="cpu",
54
dtype=torch.int32)
55
slot_mapping_xla = slot_mapping_cpu.to(torch_xla.device())
56
+ num_kv_update_slices_xla = torch.tensor([num_kv_update_slices],
57
+ device=torch_xla.device(),
58
+ dtype=torch.int32)
59
torch_xla.sync()
60
61
torch.ops.xla.dynamo_set_buffer_donor_(kv_cache_xla, True)
62
new_kv_cache_xla = torch.ops.xla.kv_cache_update_op(
59
- new_kv_xla, slot_mapping_xla, kv_cache_xla, page_size,
60
- num_slices_per_block)
63
+ new_kv_xla, slot_mapping_xla, kv_cache_xla, num_kv_update_slices_xla,
64
+ page_size, num_slices_per_block)
65
kv_cache_xla.copy_(new_kv_cache_xla)
66
torch_xla.sync()
67
vllm/attention/ops/pallas_kv_cache_update.py
+6
-3
@@ -7,11 +7,13 @@ import jax
7
from jax.experimental import pallas as pl
8
from jax.experimental.pallas import tpu as pltpu
9
10
+from vllm.utils import cdiv
11
+
12
13
def _kv_cache_update_kernel(
14
# Prefetch
13
- slices_ref, # [3, num_slices], list of (kv_cache_start, new_kv_start,
14
- # slice_len)
15
+ slices_ref, # [3, padded_num_slices], list of (kv_cache_start,
16
+ # new_kv_start, slice_len)
17
# Input
18
new_kv_hbm_ref, # [num_tokens, num_combined_kv_heads, head_dim]
19
kv_cache_hbm_ref, # [total_num_pages * page_size, num_combined_kv_heads,
@@ -70,6 +72,7 @@ def kv_cache_update(
72
Array, # [3, slices], list of (kv_cache_start, new_kv_start, slice_len)
73
kv_cache: jax.
74
Array, # [total_num_pages * page_size, num_combined_kv_heads, head_dim]
75
+ num_kv_update_slices: jax.Array, # [1]
76
*,
77
page_size: int = 32,
78
num_slices_per_block: int = 8,
@@ -107,7 +110,7 @@ def kv_cache_update(
110
num_scalar_prefetch=len(scalar_prefetches),
111
in_specs=in_specs,
112
out_specs=out_specs,
110
- grid=(slices.shape[1] // num_slices_per_block, ),
113
+ grid=(cdiv(num_kv_update_slices[0], num_slices_per_block), ),
114
scratch_shapes=scratch_shapes,
115
),
116
out_shape=out_shape,
vllm/v1/attention/backends/pallas.py
+22
-12
@@ -111,6 +111,7 @@ class PallasMetadata:
111
context_lens: torch.Tensor
112
query_start_loc: torch.Tensor
113
num_seqs: torch.Tensor
114
+ num_kv_update_slices: torch.Tensor
115
num_slices_per_kv_cache_update_block: int
116
117
@@ -219,7 +220,8 @@ class PallasAttentionBackendImpl(AttentionImpl):
220
slot_mapping = attn_metadata.slot_mapping
221
write_to_kv_cache(
222
key, value, kv_cache, slot_mapping,
222
- attn_metadata.num_slices_per_kv_cache_update_block)
223
+ attn_metadata.num_slices_per_kv_cache_update_block,
224
+ attn_metadata.num_kv_update_slices)
225
226
output = torch.ops.xla.ragged_paged_attention(
227
query,
@@ -252,6 +254,7 @@ def write_to_kv_cache(
254
kv_cache: torch.Tensor,
255
slot_mapping: torch.Tensor,
256
num_slices_per_kv_cache_update_block: int,
257
+ num_kv_update_slices: torch.Tensor,
258
) -> None:
259
""" Write the key and values to the KV cache.
260
@@ -271,7 +274,7 @@ def write_to_kv_cache(
274
275
kv_cache = kv_cache.flatten(0, 1)
276
new_kv_cache = torch.ops.xla.kv_cache_update_op(
274
- kv, slot_mapping, kv_cache, page_size,
277
+ kv, slot_mapping, kv_cache, num_kv_update_slices, page_size,
278
num_slices_per_kv_cache_update_block)
279
# NOTE: the in-place copy will be optimized away by XLA compiler.
280
kv_cache.copy_(new_kv_cache)
@@ -279,32 +282,39 @@ def write_to_kv_cache(
282
283
@requires_jax
284
def kv_cache_update_op_impl(kv: torch.Tensor, slot_mapping: torch.Tensor,
282
- kv_cache: torch.Tensor, page_size: int,
285
+ kv_cache: torch.Tensor,
286
+ num_kv_update_slices: torch.Tensor, page_size: int,
287
num_slices_per_block: int):
288
from vllm.attention.ops.pallas_kv_cache_update import kv_cache_update
285
- new_kv_cache = xb.call_jax(kv_cache_update, (kv, slot_mapping, kv_cache), {
286
- "page_size": page_size,
287
- "num_slices_per_block": num_slices_per_block
288
- })
289
+ new_kv_cache = xb.call_jax(
290
+ kv_cache_update, (kv, slot_mapping, kv_cache, num_kv_update_slices), {
291
+ "page_size": page_size,
292
+ "num_slices_per_block": num_slices_per_block
293
+ })
294
return new_kv_cache
295
296
297
XLA_LIB.define(
293
- "kv_cache_update_op(Tensor kv, Tensor slot_mapping, Tensor kv_cache, "
294
- "int page_size, int num_slices_per_block) -> Tensor", )
298
+ "kv_cache_update_op(Tensor kv, Tensor slot_mapping, Tensor kv_cache," \
299
+ "Tensor num_kv_update_slices, int page_size, int num_slices_per_block)" \
300
+ "-> Tensor", )
301
302
303
@impl(XLA_LIB, "kv_cache_update_op", "XLA")
304
def kv_cache_update_op_xla(kv: torch.Tensor, slot_mapping: torch.Tensor,
299
- kv_cache: torch.Tensor, page_size: int,
305
+ kv_cache: torch.Tensor,
306
+ num_kv_update_slices: torch.Tensor, page_size: int,
307
num_slices_per_block: int) -> torch.Tensor:
308
new_kv_cache = kv_cache_update_op_impl(kv, slot_mapping, kv_cache,
302
- page_size, num_slices_per_block)
309
+ num_kv_update_slices, page_size,
310
+ num_slices_per_block)
311
return new_kv_cache
312
313
314
@impl(XLA_LIB, "kv_cache_update_op", "CompositeExplicitAutograd")
315
def kv_cache_update_op_non_xla(kv: torch.Tensor, slot_mapping: torch.Tensor,
308
- kv_cache: torch.Tensor, page_size: int,
316
+ kv_cache: torch.Tensor,
317
+ num_kv_update_slices: torch.Tensor,
318
+ page_size: int,
319
num_slices_per_block: int) -> torch.Tensor:
320
return kv_cache
vllm/v1/worker/tpu_model_runner.py
+8
@@ -713,8 +713,10 @@ class TPUModelRunner(LoRAModelRunnerMixin):
713
self.device)
714
block_tables = block_tables.to(self.device)
715
716
+ # Calculate the slot mapping
717
slot_mapping_metadata = self._get_slot_mapping_metadata(
718
num_reqs, num_scheduled_tokens_per_req)
719
+ num_kv_update_slices = slot_mapping_metadata.shape[0]
720
padded_num_slices = _get_padded_num_kv_cache_update_slices(
721
padded_total_num_scheduled_tokens, self.max_num_reqs,
722
self.block_size)
@@ -745,6 +747,9 @@ class TPUModelRunner(LoRAModelRunnerMixin):
747
num_seqs=torch.tensor([num_reqs],
748
dtype=torch.int32,
749
device=self.device),
750
+ num_kv_update_slices=torch.tensor([num_kv_update_slices],
751
+ dtype=torch.int32,
752
+ device=self.device),
753
num_slices_per_kv_cache_update_block=
754
NUM_SLICES_PER_KV_CACHE_UPDATE_BLOCK,
755
)
@@ -1174,6 +1179,8 @@ class TPUModelRunner(LoRAModelRunnerMixin):
1179
dtype=torch.int32).to(self.device)
1180
padded_num_slices = _get_padded_num_kv_cache_update_slices(
1181
num_tokens, self.max_num_reqs, self.block_size)
1182
+ num_kv_update_slices = torch.tensor([padded_num_slices],
1183
+ dtype=torch.int32).to(self.device)
1184
slot_mapping = torch.zeros((3, padded_num_slices),
1185
dtype=torch.int32).to(self.device)
1186
block_tables = torch.zeros((num_reqs, num_blocks),
@@ -1193,6 +1200,7 @@ class TPUModelRunner(LoRAModelRunnerMixin):
1200
context_lens=context_lens,
1201
query_start_loc=query_start_loc,
1202
num_seqs=num_seqs,
1203
+ num_kv_update_slices=num_kv_update_slices,
1204
num_slices_per_kv_cache_update_block=
1205
NUM_SLICES_PER_KV_CACHE_UPDATE_BLOCK,
1206
)