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 14, 05:30 UTC
[TPU] kv cache update kernel doesn't need to be padded slices to multiple of num_slices_per_block (#22394)
Signed-off-by: Chengji Yao <chengjiyao@gmail.com> Co-authored-by: Chengji Yao <chengjiyao@gmail.com>
Chengji Yao committed
Aug 9, 2025 at 20:49 UTC
2a84fb422fc62ab29238dccbf7bdb214fc51c31e
3 files changed
+19
-21
tests/v1/tpu/test_kv_cache_update_kernel.py
-5
@@ -43,11 +43,6 @@ def test_kv_cache_update_kernel(page_size: int, combined_kv_head_num: int,
43
np.cumsum(slice_lens[:-1])])
44
slot_mapping = np.stack(
45
[kv_cache_start_indices, new_kv_cache_indices, slice_lens], axis=1)
46
- padded_size = (slot_mapping.shape[0] + num_slices_per_block -
47
- 1) // num_slices_per_block * num_slices_per_block
48
- slot_mapping = np.pad(slot_mapping,
49
- [[0, padded_size - slot_mapping.shape[0]], [0, 0]],
50
- constant_values=0)
46
slot_mapping = np.transpose(slot_mapping)
47
slot_mapping_cpu = torch.tensor(slot_mapping,
48
device="cpu",
vllm/attention/ops/pallas_kv_cache_update.py
+10
-6
@@ -14,6 +14,7 @@ def _kv_cache_update_kernel(
14
# Prefetch
15
slices_ref, # [3, padded_num_slices], list of (kv_cache_start,
16
# new_kv_start, slice_len)
17
+ num_slices_ref, # [1]
18
# Input
19
new_kv_hbm_ref, # [num_tokens, num_combined_kv_heads, head_dim]
20
kv_cache_hbm_ref, # [total_num_pages * page_size, num_combined_kv_heads,
@@ -32,8 +33,10 @@ def _kv_cache_update_kernel(
33
# Copy from new_kv_hbm_ref to scratch
34
for i in range(num_slices_per_block):
35
offset_i = i + block_idx * num_slices_per_block
35
- new_kv_start = slices_ref[1, offset_i]
36
- length = slices_ref[2, offset_i]
36
+ new_kv_start = jax.lax.select(offset_i < num_slices_ref[0],
37
+ slices_ref[1, offset_i], 0)
38
+ length = jax.lax.select(offset_i < num_slices_ref[0],
39
+ slices_ref[2, offset_i], 0)
40
async_copy = pltpu.make_async_copy(
41
new_kv_hbm_ref.at[pl.ds(new_kv_start, length), ...],
42
scratch.at[i, pl.ds(0, length), ...],
@@ -49,8 +52,10 @@ def _kv_cache_update_kernel(
52
async_copies.clear()
53
for i in range(num_slices_per_block):
54
offset_i = i + block_idx * num_slices_per_block
52
- kv_cache_start = slices_ref[0, offset_i]
53
- length = slices_ref[2, offset_i]
55
+ kv_cache_start = jax.lax.select(offset_i < num_slices_ref[0],
56
+ slices_ref[0, offset_i], 0)
57
+ length = jax.lax.select(offset_i < num_slices_ref[0],
58
+ slices_ref[2, offset_i], 0)
59
async_copy = pltpu.make_async_copy(
60
scratch.at[i, pl.ds(0, length), ...],
61
kv_cache_hbm_ref.at[pl.ds(kv_cache_start, length), ...],
@@ -77,7 +82,6 @@ def kv_cache_update(
82
page_size: int = 32,
83
num_slices_per_block: int = 8,
84
):
80
- assert slices.shape[1] % num_slices_per_block == 0
85
_, num_combined_kv_heads, head_dim = new_kv.shape
86
assert kv_cache.shape[1] == num_combined_kv_heads
87
assert kv_cache.shape[2] == head_dim
@@ -93,7 +97,7 @@ def kv_cache_update(
97
out_specs = [pl.BlockSpec(memory_space=pltpu.TPUMemorySpace.ANY)]
98
out_shape = [jax.ShapeDtypeStruct(kv_cache.shape, dtype=kv_cache.dtype)]
99
96
- scalar_prefetches = [slices]
100
+ scalar_prefetches = [slices, num_kv_update_slices]
101
scratch = pltpu.VMEM(
102
(num_slices_per_block, page_size, num_combined_kv_heads, head_dim),
103
new_kv.dtype,
vllm/v1/worker/tpu_model_runner.py
+9
-10
@@ -745,7 +745,7 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
745
num_kv_update_slices = slot_mapping_metadata.shape[0]
746
padded_num_slices = _get_padded_num_kv_cache_update_slices(
747
padded_total_num_scheduled_tokens, self.max_num_reqs,
748
- self.block_size, self._num_slices_per_kv_cache_update_block)
748
+ self.block_size)
749
slot_mapping_metadata = np.pad(
750
slot_mapping_metadata,
751
[[0, padded_num_slices - len(slot_mapping_metadata)], [0, 0]],
@@ -1244,8 +1244,7 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1244
position_ids = torch.zeros(num_tokens,
1245
dtype=torch.int32).to(self.device)
1246
padded_num_slices = _get_padded_num_kv_cache_update_slices(
1247
- num_tokens, self.max_num_reqs, self.block_size,
1248
- self._num_slices_per_kv_cache_update_block)
1247
+ num_tokens, self.max_num_reqs, self.block_size)
1248
num_kv_update_slices = torch.tensor([padded_num_slices],
1249
dtype=torch.int32).to(self.device)
1250
slot_mapping = torch.zeros((3, padded_num_slices),
@@ -1963,17 +1962,17 @@ def copy_kv_blocks(
1962
_copy_fn(src_tensor, dst_tensor, src_indices, dst_indices)
1963
1964
1966
-def _get_padded_num_kv_cache_update_slices(
1967
- num_tokens: int, max_num_reqs: int, page_size: int,
1968
- num_slices_per_kv_cache_update_block: int) -> int:
1965
+def _get_padded_num_kv_cache_update_slices(num_tokens: int, max_num_reqs: int,
1966
+ page_size: int) -> int:
1967
"""Calculates the padded number of KV cache update slices to avoid
1968
recompilation."""
1969
+ # NOTE(chengjiyao): let's say R_i is the token num for i-th request,
1970
+ # so it occupies most 2 + R_i // page_size pages. The total maximum
1971
+ # possible number of pages needed is sum(2 + R_i // page_size), which
1972
+ # is <= 2 * max_num_reqs + sum(R_i) // page_size
1973
+ # = 2 * max_num_reqs + num_tokens // page_size
1974
padded_num_slices = 2 * max_num_reqs + num_tokens // page_size
1975
padded_num_slices = min(padded_num_slices, num_tokens)
1973
- padded_num_slices = (
1974
- padded_num_slices + num_slices_per_kv_cache_update_block - 1
1975
- ) // num_slices_per_kv_cache_update_block * \
1976
- num_slices_per_kv_cache_update_block
1976
return padded_num_slices
1977
1978