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 10, 00:15 UTC
[Refactor] Remove unused file `pallas_kv_cache_update.py` (#32433)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Jan 18, 2026 at 15:46 UTC
16de822c719dedec51e5f2a3ee6ed5b0890e56cb
1 file changed
-130
vllm/v1/attention/ops/pallas_kv_cache_update.py
deleted
-130
@@ -1,130 +0,0 @@
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
-from vllm.utils.math_utils import cdiv
11
-
12
-
13
-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,
21
- # head_dim]
22
- # Output
23
- _, # [total_num_pages * page_size, num_combined_kv_heads, head_dim]
24
- # Scratch
25
- scratch, # [num_slices_per_block, page_size, num_combined_kv_heads,
26
- # head_dim]
27
- sem,
28
-):
29
- async_copies = []
30
- block_idx = pl.program_id(0)
31
- num_slices_per_block = scratch.shape[0]
32
-
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
36
- new_kv_start = jax.lax.select(
37
- offset_i < num_slices_ref[0], slices_ref[1, offset_i], 0
38
- )
39
- length = jax.lax.select(
40
- offset_i < num_slices_ref[0], slices_ref[2, offset_i], 0
41
- )
42
- async_copy = pltpu.make_async_copy(
43
- new_kv_hbm_ref.at[pl.ds(new_kv_start, length), ...],
44
- scratch.at[i, pl.ds(0, length), ...],
45
- sem,
46
- )
47
- async_copy.start()
48
- async_copies.append(async_copy)
49
-
50
- for async_copy in async_copies:
51
- async_copy.wait()
52
-
53
- # Copy from scratch to kv_cache_hbm_ref
54
- async_copies.clear()
55
- for i in range(num_slices_per_block):
56
- offset_i = i + block_idx * num_slices_per_block
57
- kv_cache_start = jax.lax.select(
58
- offset_i < num_slices_ref[0], slices_ref[0, offset_i], 0
59
- )
60
- length = jax.lax.select(
61
- offset_i < num_slices_ref[0], slices_ref[2, offset_i], 0
62
- )
63
- async_copy = pltpu.make_async_copy(
64
- scratch.at[i, pl.ds(0, length), ...],
65
- kv_cache_hbm_ref.at[pl.ds(kv_cache_start, length), ...],
66
- sem,
67
- )
68
- async_copy.start()
69
- async_copies.append(async_copy)
70
- for async_copy in async_copies:
71
- async_copy.wait()
72
-
73
-
74
-@functools.partial(
75
- jax.jit,
76
- static_argnames=["page_size", "num_slices_per_block"],
77
-)
78
-def kv_cache_update(
79
- # [total_num_token, num_combined_kv_heads, head_dim]
80
- new_kv: jax.Array,
81
- # [3, slices], list of (kv_cache_start, new_kv_start, slice_len)
82
- slices: jax.Array,
83
- # [total_num_pages * page_size, num_combined_kv_heads, head_dim]
84
- kv_cache: jax.Array,
85
- # [1]
86
- num_kv_update_slices: jax.Array,
87
- *,
88
- page_size: int = 32,
89
- num_slices_per_block: int = 8,
90
-):
91
- _, num_combined_kv_heads, head_dim = new_kv.shape
92
- assert kv_cache.shape[1] == num_combined_kv_heads
93
- assert kv_cache.shape[2] == head_dim
94
- assert head_dim % 128 == 0
95
- # TODO: Add dynamic check to make sure that the all the slice lengths are
96
- # smaller or equal to page_size
97
-
98
- in_specs = [
99
- pl.BlockSpec(memory_space=pltpu.TPUMemorySpace.ANY),
100
- pl.BlockSpec(memory_space=pltpu.TPUMemorySpace.ANY),
101
- ]
102
-
103
- out_specs = [pl.BlockSpec(memory_space=pltpu.TPUMemorySpace.ANY)]
104
- out_shape = [jax.ShapeDtypeStruct(kv_cache.shape, dtype=kv_cache.dtype)]
105
-
106
- scalar_prefetches = [slices, num_kv_update_slices]
107
- scratch = pltpu.VMEM(
108
- (num_slices_per_block, page_size, num_combined_kv_heads, head_dim),
109
- new_kv.dtype,
110
- )
111
-
112
- scratch_shapes = [
113
- scratch,
114
- pltpu.SemaphoreType.DMA,
115
- ]
116
-
117
- kernel = pl.pallas_call(
118
- _kv_cache_update_kernel,
119
- grid_spec=pltpu.PrefetchScalarGridSpec(
120
- num_scalar_prefetch=len(scalar_prefetches),
121
- in_specs=in_specs,
122
- out_specs=out_specs,
123
- grid=(cdiv(num_kv_update_slices[0], num_slices_per_block),),
124
- scratch_shapes=scratch_shapes,
125
- ),
126
- out_shape=out_shape,
127
- input_output_aliases={len(scalar_prefetches) + 1: 0},
128
- )
129
-
130
- return kernel(*scalar_prefetches, new_kv, kv_cache)[0]