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, 04:45 UTC
[TPU] add kv cache update kernel (#19928)
Signed-off-by: Chengji Yao <chengjiyao@google.com>
Chengji Yao committed
Jun 26, 2025 at 10:01 UTC
04e1642e3251fc575d104c84782fafea348cfbaf
6 files changed
+342
-38
.buildkite/scripts/hardware_ci/run-tpu-v1-test.sh
+2
@@ -159,6 +159,8 @@ run_and_track_test 14 "test_tpu_qkv_linear.py" \
159
"python3 -m pytest -s -v /workspace/vllm/tests/v1/tpu/test_tpu_qkv_linear.py"
160
run_and_track_test 15 "test_spmd_model_weight_loading.py" \
161
"python3 -m pytest -s -v /workspace/vllm/tests/v1/tpu/test_spmd_model_weight_loading.py"
162
+run_and_track_test 16 "test_kv_cache_update_kernel.py" \
163
+ "python3 -m pytest -s -v /workspace/vllm/tests/v1/tpu/test_kv_cache_update_kernel.py"
164
165
# After all tests have been attempted, exit with the overall status.
166
if [ "$overall_script_exit_code" -ne 0 ]; then
tests/v1/tpu/test_kv_cache_update_kernel.py
new
+71
@@ -0,0 +1,71 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import numpy as np
5
+import pytest
6
+import torch
7
+import torch_xla
8
+
9
+import vllm.v1.attention.backends.pallas # noqa: F401
10
+from vllm.platforms import current_platform
11
+
12
+
13
+@pytest.mark.skipif(not current_platform.is_tpu(),
14
+ reason="This is a test for TPU only")
15
+@pytest.mark.parametrize("page_size", [32, 33])
16
+@pytest.mark.parametrize("combined_kv_head_num", [2, 16])
17
+@pytest.mark.parametrize("head_dim", [128, 256])
18
+@pytest.mark.parametrize("num_slices_per_block", [4, 8])
19
+def test_kv_cache_update_kernel(page_size: int, combined_kv_head_num: int,
20
+ head_dim: int, num_slices_per_block: int):
21
+ page_num = 1000
22
+ padded_num_tokens = 128
23
+ kv_cache_cpu = torch.zeros(
24
+ (page_num * page_size, combined_kv_head_num, head_dim),
25
+ dtype=torch.bfloat16,
26
+ device="cpu")
27
+ kv_cache_xla = kv_cache_cpu.to(torch_xla.device())
28
+ new_kv_cpu = torch.randn(
29
+ (padded_num_tokens, combined_kv_head_num, head_dim),
30
+ dtype=torch.bfloat16,
31
+ device="cpu")
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
+ kv_cache_start_indices = np.array([
36
+ page_size * 2 - 7, page_size * 2, page_size * 3, page_size * 4 + 6,
37
+ page_size * 5 + 7, page_size * 6 + 8, page_size * 15 + 3
38
+ ],
39
+ dtype=np.int32)
40
+ new_kv_cache_indices = np.concatenate(
41
+ [np.array([0], dtype=np.int32),
42
+ np.cumsum(slice_lens[:-1])])
43
+ slot_mapping = np.stack(
44
+ [kv_cache_start_indices, new_kv_cache_indices, slice_lens], axis=1)
45
+ padded_size = (slot_mapping.shape[0] + num_slices_per_block -
46
+ 1) // num_slices_per_block * num_slices_per_block
47
+ slot_mapping = np.pad(slot_mapping,
48
+ [[0, padded_size - slot_mapping.shape[0]], [0, 0]],
49
+ constant_values=0)
50
+ slot_mapping = np.transpose(slot_mapping)
51
+ slot_mapping_cpu = torch.tensor(slot_mapping,
52
+ device="cpu",
53
+ dtype=torch.int32)
54
+ slot_mapping_xla = slot_mapping_cpu.to(torch_xla.device())
55
+ torch_xla.sync()
56
+
57
+ torch.ops.xla.dynamo_set_buffer_donor_(kv_cache_xla, True)
58
+ 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)
61
+ kv_cache_xla.copy_(new_kv_cache_xla)
62
+ torch_xla.sync()
63
+
64
+ for ni, ci, sl in zip(new_kv_cache_indices, kv_cache_start_indices,
65
+ slice_lens):
66
+ kv_cache_cpu[ci:ci + sl, :, :] = new_kv_cpu[ni:ni + sl, :, :]
67
+
68
+ assert torch.allclose(kv_cache_xla.cpu(),
69
+ kv_cache_cpu,
70
+ atol=1e-4,
71
+ rtol=1e-4)
tests/v1/tpu/test_pallas.py
+2
-1
@@ -47,7 +47,7 @@ def test_ragged_paged_attention():
47
key = torch.zeros(num_tokens, num_kv_heads * head_size)
48
value = torch.zeros(num_tokens, num_kv_heads * head_size)
49
kv_cache = torch.zeros(num_blocks, block_size, num_kv_heads * 2, head_size)
50
- slot_mapping = torch.zeros(num_tokens, dtype=torch.int64)
50
+ slot_mapping = torch.zeros((3, num_tokens), dtype=torch.int64)
51
max_num_reqs = 8
52
max_num_blocks_per_req = 8
53
block_tables = torch.zeros((max_num_reqs, max_num_blocks_per_req),
@@ -65,6 +65,7 @@ def test_ragged_paged_attention():
65
context_lens=context_lens,
66
query_start_loc=query_start_loc,
67
num_seqs=num_seqs,
68
+ num_slices_per_kv_cache_update_block=8,
69
)
70
71
with patch("torch.ops.xla.ragged_paged_attention"
vllm/attention/ops/pallas_kv_cache_update.py
new
+117
@@ -0,0 +1,117 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import functools
5
+
6
+import jax
7
+from jax.experimental import pallas as pl
8
+from jax.experimental.pallas import tpu as pltpu
9
+
10
+
11
+def _kv_cache_update_kernel(
12
+ # Prefetch
13
+ slices_ref, # [3, num_slices], list of (kv_cache_start, new_kv_start,
14
+ # slice_len)
15
+ # Input
16
+ new_kv_hbm_ref, # [num_tokens, num_combined_kv_heads, head_dim]
17
+ kv_cache_hbm_ref, # [total_num_pages * page_size, num_combined_kv_heads,
18
+ # head_dim]
19
+ # Output
20
+ _, # [total_num_pages * page_size, num_combined_kv_heads, head_dim]
21
+ # Scratch
22
+ scratch, # [num_slices_per_block, page_size, num_combined_kv_heads,
23
+ # head_dim]
24
+ sem,
25
+):
26
+ async_copies = []
27
+ block_idx = pl.program_id(0)
28
+ num_slices_per_block = scratch.shape[0]
29
+
30
+ # Copy from new_kv_hbm_ref to scratch
31
+ for i in range(num_slices_per_block):
32
+ offset_i = i + block_idx * num_slices_per_block
33
+ new_kv_start = slices_ref[1, offset_i]
34
+ length = slices_ref[2, offset_i]
35
+ async_copy = pltpu.make_async_copy(
36
+ new_kv_hbm_ref.at[pl.ds(new_kv_start, length), ...],
37
+ scratch.at[i, pl.ds(0, length), ...],
38
+ sem,
39
+ )
40
+ async_copy.start()
41
+ async_copies.append(async_copy)
42
+
43
+ for async_copy in async_copies:
44
+ async_copy.wait()
45
+
46
+ # Copy from scratch to kv_cache_hbm_ref
47
+ async_copies.clear()
48
+ for i in range(num_slices_per_block):
49
+ offset_i = i + block_idx * num_slices_per_block
50
+ kv_cache_start = slices_ref[0, offset_i]
51
+ length = slices_ref[2, offset_i]
52
+ async_copy = pltpu.make_async_copy(
53
+ scratch.at[i, pl.ds(0, length), ...],
54
+ kv_cache_hbm_ref.at[pl.ds(kv_cache_start, length), ...],
55
+ sem,
56
+ )
57
+ async_copy.start()
58
+ async_copies.append(async_copy)
59
+ for async_copy in async_copies:
60
+ async_copy.wait()
61
+
62
+
63
+@functools.partial(
64
+ jax.jit,
65
+ static_argnames=["page_size", "num_slices_per_block"],
66
+)
67
+def kv_cache_update(
68
+ new_kv: jax.Array, # [total_num_token, num_combined_kv_heads, head_dim]
69
+ slices: jax.
70
+ Array, # [3, slices], list of (kv_cache_start, new_kv_start, slice_len)
71
+ kv_cache: jax.
72
+ Array, # [total_num_pages * page_size, num_combined_kv_heads, head_dim]
73
+ *,
74
+ page_size: int = 32,
75
+ num_slices_per_block: int = 8,
76
+):
77
+ assert slices.shape[1] % num_slices_per_block == 0
78
+ _, num_combined_kv_heads, head_dim = new_kv.shape
79
+ assert kv_cache.shape[1] == num_combined_kv_heads
80
+ assert kv_cache.shape[2] == head_dim
81
+ assert head_dim % 128 == 0
82
+ # TODO: Add dynamic check to make sure that the all the slice lengths are
83
+ # smaller or equal to page_size
84
+
85
+ in_specs = [
86
+ pl.BlockSpec(memory_space=pltpu.TPUMemorySpace.ANY),
87
+ pl.BlockSpec(memory_space=pltpu.TPUMemorySpace.ANY),
88
+ ]
89
+
90
+ out_specs = [pl.BlockSpec(memory_space=pltpu.TPUMemorySpace.ANY)]
91
+ out_shape = [jax.ShapeDtypeStruct(kv_cache.shape, dtype=kv_cache.dtype)]
92
+
93
+ scalar_prefetches = [slices]
94
+ scratch = pltpu.VMEM(
95
+ (num_slices_per_block, page_size, num_combined_kv_heads, head_dim),
96
+ new_kv.dtype,
97
+ )
98
+
99
+ scratch_shapes = [
100
+ scratch,
101
+ pltpu.SemaphoreType.DMA,
102
+ ]
103
+
104
+ kernel = pl.pallas_call(
105
+ _kv_cache_update_kernel,
106
+ grid_spec=pltpu.PrefetchScalarGridSpec(
107
+ num_scalar_prefetch=len(scalar_prefetches),
108
+ in_specs=in_specs,
109
+ out_specs=out_specs,
110
+ grid=(slices.shape[1] // num_slices_per_block, ),
111
+ scratch_shapes=scratch_shapes,
112
+ ),
113
+ out_shape=out_shape,
114
+ input_output_aliases={len(scalar_prefetches) + 1: 0},
115
+ )
116
+
117
+ return kernel(*scalar_prefetches, new_kv, kv_cache)[0]
vllm/v1/attention/backends/pallas.py
+50
-5
@@ -5,8 +5,12 @@ from dataclasses import dataclass
5
from typing import Any, Optional
6
7
import torch
8
-# Required to register custom ops.
8
+import torch_xla.core.xla_builder as xb
9
import torch_xla.experimental.custom_kernel # noqa: F401
10
+# Required to register custom ops.
11
+from torch.library import impl
12
+from torch_xla._internal.jax_workarounds import requires_jax
13
+from torch_xla.experimental.custom_kernel import XLA_LIB
14
15
from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,
16
AttentionLayer, AttentionType)
@@ -107,6 +111,7 @@ class PallasMetadata:
111
context_lens: torch.Tensor
112
query_start_loc: torch.Tensor
113
num_seqs: torch.Tensor
114
+ num_slices_per_kv_cache_update_block: int
115
116
117
class PallasAttentionBackendImpl(AttentionImpl):
@@ -212,7 +217,9 @@ class PallasAttentionBackendImpl(AttentionImpl):
217
# Write input keys and values to the KV cache.
218
# Skip this if sharing KV cache with an earlier attention layer.
219
slot_mapping = attn_metadata.slot_mapping
215
- write_to_kv_cache(key, value, kv_cache, slot_mapping)
220
+ write_to_kv_cache(
221
+ key, value, kv_cache, slot_mapping,
222
+ attn_metadata.num_slices_per_kv_cache_update_block)
223
224
output = torch.ops.xla.ragged_paged_attention(
225
query,
@@ -244,6 +251,7 @@ def write_to_kv_cache(
251
value: torch.Tensor,
252
kv_cache: torch.Tensor,
253
slot_mapping: torch.Tensor,
254
+ num_slices_per_kv_cache_update_block: int,
255
) -> None:
256
""" Write the key and values to the KV cache.
257
@@ -251,9 +259,9 @@ def write_to_kv_cache(
259
key: shape = [num_tokens, num_kv_heads * head_size]
260
value: shape = [num_tokens, num_kv_heads * head_size]
261
kv_cache = [num_blocks, block_size, num_kv_heads * 2, head_size]
254
-
262
+ num_slices_per_kv_cache_update_block: int
263
"""
256
- _, _, num_combined_kv_heads, head_size = kv_cache.shape
264
+ _, page_size, num_combined_kv_heads, head_size = kv_cache.shape
265
head_size = cdiv(head_size,
266
TPU_HEAD_SIZE_ALIGNMENT) * TPU_HEAD_SIZE_ALIGNMENT
267
kv = torch.cat([key, value], axis=-1).reshape(-1, num_combined_kv_heads,
@@ -262,4 +270,41 @@ def write_to_kv_cache(
270
torch.ops.xla.dynamo_set_buffer_donor_(kv_cache, True)
271
272
kv_cache = kv_cache.flatten(0, 1)
265
- kv_cache.index_copy_(0, slot_mapping, kv)
273
+ new_kv_cache = torch.ops.xla.kv_cache_update_op(
274
+ kv, slot_mapping, kv_cache, page_size,
275
+ num_slices_per_kv_cache_update_block)
276
+ # NOTE: the in-place copy will be optimized away by XLA compiler.
277
+ kv_cache.copy_(new_kv_cache)
278
+
279
+
280
+@requires_jax
281
+def kv_cache_update_op_impl(kv: torch.Tensor, slot_mapping: torch.Tensor,
282
+ kv_cache: torch.Tensor, page_size: int,
283
+ num_slices_per_block: int):
284
+ 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
+ return new_kv_cache
290
+
291
+
292
+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", )
295
+
296
+
297
+@impl(XLA_LIB, "kv_cache_update_op", "XLA")
298
+def kv_cache_update_op_xla(kv: torch.Tensor, slot_mapping: torch.Tensor,
299
+ kv_cache: torch.Tensor, page_size: int,
300
+ num_slices_per_block: int) -> torch.Tensor:
301
+ new_kv_cache = kv_cache_update_op_impl(kv, slot_mapping, kv_cache,
302
+ page_size, num_slices_per_block)
303
+ return new_kv_cache
304
+
305
+
306
+@impl(XLA_LIB, "kv_cache_update_op", "CompositeExplicitAutograd")
307
+def kv_cache_update_op_non_xla(kv: torch.Tensor, slot_mapping: torch.Tensor,
308
+ kv_cache: torch.Tensor, page_size: int,
309
+ num_slices_per_block: int) -> torch.Tensor:
310
+ return kv_cache
vllm/v1/worker/tpu_model_runner.py
+100
-32
@@ -53,12 +53,11 @@ if TYPE_CHECKING:
53
54
logger = init_logger(__name__)
55
56
-# Here we utilize the behavior that out-of-bound index is ignored.
57
-# FIXME(woosuk): Find a more reliable way to prevent possible bugs.
58
-_PAD_SLOT_ID = 1_000_000_000
56
INVALID_TOKEN_ID = -1
57
# Smallest output size
58
MIN_NUM_SEQS = 8
59
+# Block size used for kv cache updating kernel
60
+NUM_SLICES_PER_KV_CACHE_UPDATE_BLOCK = 8
61
62
63
#########################################################
@@ -526,6 +525,69 @@ class TPUModelRunner(LoRAModelRunnerMixin):
525
526
return kv_cache_spec
527
528
+ def _get_slot_mapping_metadata(self, num_reqs,
529
+ num_scheduled_tokens_per_req):
530
+ """
531
+ Computes metadata for mapping slots to blocks in the key-value (KV)
532
+ cache for a batch of requests.
533
+
534
+ This function determines, for each request in the batch, how the
535
+ scheduled tokens are distributed across memory blocks, and generates
536
+ metadata needed to map slices of tokens to their corresponding positions
537
+ in the KV cache.
538
+
539
+ Args:
540
+ num_reqs (int): Number of requests in the current batch.
541
+ num_scheduled_tokens_per_req (int or np.ndarray): Number of tokens
542
+ to be scheduled for each request.
543
+
544
+ Returns:
545
+ np.ndarray: A 2D array of shape (total_block_len, 3), where each row
546
+ contains:
547
+ - kv_cache_start_index (int): The starting index in the KV cache
548
+ for the corresponding slice.
549
+ - new_kv_start_index (int): The starting index in the new KV
550
+ cache for the corresponding slice.
551
+ - slice_len (int): The length of the slice.
552
+ """
553
+ slices_start = self.input_batch.num_computed_tokens_cpu[:num_reqs]
554
+ slices_end = self.input_batch.num_computed_tokens_cpu[:num_reqs] + \
555
+ num_scheduled_tokens_per_req
556
+ local_block_start_idx = slices_start // self.block_size
557
+ local_block_end_idx = (slices_end - 1) // self.block_size
558
+ no_repeat_req_indices = self.arange_np[:num_reqs]
559
+ global_block_start_idx = (
560
+ no_repeat_req_indices * self.max_num_blocks_per_req +
561
+ local_block_start_idx)
562
+ block_lens = local_block_end_idx - local_block_start_idx + 1
563
+ global_block_start_idx = np.repeat(global_block_start_idx, block_lens)
564
+ slice_arange = np.concatenate([self.arange_np[:n] for n in block_lens])
565
+ global_block_indices = global_block_start_idx + slice_arange
566
+ block_table_cpu = self.input_batch.block_table[0].get_cpu_tensor()
567
+ block_numbers = block_table_cpu.flatten()[global_block_indices].numpy()
568
+ total_block_len = np.sum(block_lens)
569
+ slot_mapping_slices = np.repeat(np.array([[0, self.block_size]],
570
+ dtype=np.int32),
571
+ total_block_len,
572
+ axis=0)
573
+ cu_block_lens = np.zeros(len(block_lens) + 1, dtype=np.int32)
574
+ np.cumsum(block_lens, out=cu_block_lens[1:])
575
+ for req_idx in range(num_reqs):
576
+ slot_mapping_slices[cu_block_lens[req_idx]][
577
+ 0] = slices_start[req_idx] % self.block_size
578
+ slot_mapping_slices[
579
+ cu_block_lens[req_idx + 1] -
580
+ 1][1] = (slices_end[req_idx] - 1) % self.block_size + 1
581
+ slice_lens = slot_mapping_slices[:, 1] - slot_mapping_slices[:, 0]
582
+ cu_slices_lens = np.zeros(len(slice_lens) + 1, dtype=np.int32)
583
+ np.cumsum(slice_lens, out=cu_slices_lens[1:])
584
+ kv_cache_start_indices = slot_mapping_slices[:, 0] + \
585
+ (block_numbers * self.block_size)
586
+ new_kv_start_indices = cu_slices_lens[:-1]
587
+ slot_mapping_metadata = np.stack(
588
+ [kv_cache_start_indices, new_kv_start_indices, slice_lens], axis=1)
589
+ return slot_mapping_metadata
590
+
591
def _prepare_inputs(self, scheduler_output: "SchedulerOutput",
592
start_index: int):
593
assert scheduler_output.total_num_scheduled_tokens > 0
@@ -603,26 +665,6 @@ class TPUModelRunner(LoRAModelRunnerMixin):
665
torch.from_numpy(token_indices),
666
out=self.input_ids_cpu[:total_num_scheduled_tokens])
667
606
- # Calculate the slot mapping.
607
- # E.g., [0, 1, 0, 1, 2, 3, 4, 0, 1, 2]
608
- # -> [0, 0, K, K, K + 1, K + 1, K + 2, 2 * K, 2 * K, 2 * K + 1]
609
- # where K is the max_num_blocks_per_req and the block size is 2.
610
- # NOTE(woosuk): We can't simply use `token_indices // block_size` here
611
- # because M (max_model_len) is not necessarily divisible by block_size.
612
- # req_indices: # E.g., [2, 5, 3] -> [0, 0, 1, 1, 1, 1, 1, 2, 2, 2]
613
- block_table_indices = (req_indices * self.max_num_blocks_per_req +
614
- positions_np // self.block_size)
615
- # NOTE(woosuk): We use torch.index_select instead of np.take here
616
- # because torch.index_select is much faster than np.take for large
617
- # tensors.
618
- block_table_cpu = self.input_batch.block_table[0].get_cpu_tensor()
619
- block_numbers = block_table_cpu.flatten()[block_table_indices].numpy()
620
- block_offsets = positions_np % self.block_size
621
- np.add(block_numbers * self.block_size,
622
- block_offsets,
623
- out=self.input_batch.block_table[0].
624
- slot_mapping_np[:total_num_scheduled_tokens])
625
-
668
# Prepare the attention metadata.
669
self.query_start_loc_np[0] = 0
670
np.cumsum(num_scheduled_tokens_per_req,
@@ -645,12 +687,6 @@ class TPUModelRunner(LoRAModelRunnerMixin):
687
self.position_ids = self.positions_cpu[:
688
padded_total_num_scheduled_tokens].to(
689
self.device)
648
- self.input_batch.block_table[0].slot_mapping_cpu[
649
- total_num_scheduled_tokens:] = _PAD_SLOT_ID
650
- slot_mapping = (
651
- self.input_batch.block_table[0].
652
- slot_mapping_cpu[:padded_total_num_scheduled_tokens].to(
653
- self.device))
690
if use_max_model_len:
691
block_tables = self.block_table_cpu[:self.num_reqs_max_model_len, :
692
self.max_num_blocks_per_req]
@@ -675,6 +711,19 @@ class TPUModelRunner(LoRAModelRunnerMixin):
711
self.device)
712
block_tables = block_tables.to(self.device)
713
714
+ slot_mapping_metadata = self._get_slot_mapping_metadata(
715
+ num_reqs, num_scheduled_tokens_per_req)
716
+ padded_num_slices = _get_padded_num_kv_cache_update_slices(
717
+ padded_total_num_scheduled_tokens, self.max_num_reqs,
718
+ self.block_size)
719
+ slot_mapping_metadata = np.pad(
720
+ slot_mapping_metadata,
721
+ [[0, padded_num_slices - len(slot_mapping_metadata)], [0, 0]],
722
+ constant_values=0)
723
+ slot_mapping_metadata = np.transpose(slot_mapping_metadata)
724
+ slot_mapping_metadata = torch.tensor(slot_mapping_metadata,
725
+ device=self.device)
726
+
727
if self.lora_config is not None:
728
# We need to respect padding when activating LoRA adapters
729
padded_num_scheduled_tokens_per_req = np.copy(
@@ -687,13 +736,15 @@ class TPUModelRunner(LoRAModelRunnerMixin):
736
padded_num_scheduled_tokens_per_req)
737
738
attn_metadata = PallasMetadata(
690
- slot_mapping=slot_mapping,
739
+ slot_mapping=slot_mapping_metadata,
740
block_tables=block_tables,
741
context_lens=seq_lens,
742
query_start_loc=query_start_loc,
743
num_seqs=torch.tensor([num_reqs],
744
dtype=torch.int32,
745
device=self.device),
746
+ num_slices_per_kv_cache_update_block=
747
+ NUM_SLICES_PER_KV_CACHE_UPDATE_BLOCK,
748
)
749
# NOTE(woosuk): Due to chunked prefills, there can be at most 1 partial
750
# request in the batch. While we should not sample any token from this
@@ -1119,8 +1170,10 @@ class TPUModelRunner(LoRAModelRunnerMixin):
1170
actual_num_reqs = min(num_tokens, num_reqs)
1171
position_ids = torch.zeros(num_tokens,
1172
dtype=torch.int32).to(self.device)
1122
- slot_mapping = torch.zeros(num_tokens,
1123
- dtype=torch.int64).to(self.device)
1173
+ padded_num_slices = _get_padded_num_kv_cache_update_slices(
1174
+ num_tokens, self.max_num_reqs, self.block_size)
1175
+ slot_mapping = torch.zeros((3, padded_num_slices),
1176
+ dtype=torch.int32).to(self.device)
1177
block_tables = torch.zeros((num_reqs, num_blocks),
1178
dtype=torch.int32).to(self.device)
1179
query_lens = [1] * num_reqs
@@ -1138,6 +1191,8 @@ class TPUModelRunner(LoRAModelRunnerMixin):
1191
context_lens=context_lens,
1192
query_start_loc=query_start_loc,
1193
num_seqs=num_seqs,
1194
+ num_slices_per_kv_cache_update_block=
1195
+ NUM_SLICES_PER_KV_CACHE_UPDATE_BLOCK,
1196
)
1197
1198
if self.is_multimodal_model:
@@ -1742,6 +1797,19 @@ def _get_padded_token_len(paddings: list[int], x: int) -> int:
1797
return paddings[index]
1798
1799
1800
+def _get_padded_num_kv_cache_update_slices(num_tokens: int, max_num_reqs: int,
1801
+ page_size: int) -> int:
1802
+ """Calculates the padded number of KV cache update slices to avoid
1803
+ recompilation."""
1804
+ padded_num_slices = 2 * max_num_reqs + num_tokens // page_size
1805
+ padded_num_slices = min(padded_num_slices, num_tokens)
1806
+ padded_num_slices = (
1807
+ padded_num_slices + NUM_SLICES_PER_KV_CACHE_UPDATE_BLOCK - 1
1808
+ ) // NUM_SLICES_PER_KV_CACHE_UPDATE_BLOCK * \
1809
+ NUM_SLICES_PER_KV_CACHE_UPDATE_BLOCK
1810
+ return padded_num_slices
1811
+
1812
+
1813
def replace_set_lora(model):
1814
1815
def _tpu_set_lora(