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, 02:15 UTC
[TPU] Optimize kv cache update kernel (#20415)
Signed-off-by: Yifei Teng <tengyifei88@gmail.com>
Yifei Teng committed
Jul 15, 2025 at 03:56 UTC
c586b55667fa3dc341c45fa55744791ae26fae5b
3 files changed
+63
-16
vllm/utils/__init__.py
+7
@@ -947,6 +947,13 @@ def next_power_of_2(n) -> int:
947
return 1 << (n - 1).bit_length()
948
949
950
+def prev_power_of_2(n: int) -> int:
951
+ """The previous power of 2 (inclusive)"""
952
+ if n <= 0:
953
+ return 0
954
+ return 1 << (n.bit_length() - 1)
955
+
956
+
957
def round_up(x: int, y: int) -> int:
958
return ((x + y - 1) // y) * y
959
vllm/v1/attention/backends/pallas.py
+6
@@ -324,3 +324,9 @@ def kv_cache_update_op_non_xla(kv: torch.Tensor, slot_mapping: torch.Tensor,
324
page_size: int,
325
num_slices_per_block: int) -> torch.Tensor:
326
return kv_cache
327
+
328
+
329
+def get_page_size_bytes(block_size: int, num_kv_heads: int, head_size: int,
330
+ kv_cache_dtype: torch.dtype) -> int:
331
+ """Returns the size in bytes of one page of the KV cache."""
332
+ return block_size * num_kv_heads * head_size * kv_cache_dtype.itemsize
vllm/v1/worker/tpu_model_runner.py
+50
-16
@@ -31,9 +31,10 @@ from vllm.multimodal.inputs import (BatchedTensorInputs, MultiModalKwargs,
31
from vllm.multimodal.utils import group_mm_inputs_by_modality
32
from vllm.sequence import IntermediateTensors
33
from vllm.utils import (STR_DTYPE_TO_TORCH_DTYPE, LayerBlockType, cdiv,
34
- is_pin_memory_available)
34
+ is_pin_memory_available, prev_power_of_2)
35
from vllm.v1.attention.backends.pallas import (PallasAttentionBackend,
36
- PallasMetadata)
36
+ PallasMetadata,
37
+ get_page_size_bytes)
38
from vllm.v1.core.encoder_cache_manager import compute_encoder_budget
39
from vllm.v1.kv_cache_interface import (AttentionSpec, FullAttentionSpec,
40
KVCacheConfig, KVCacheSpec,
@@ -56,8 +57,6 @@ logger = init_logger(__name__)
57
INVALID_TOKEN_ID = -1
58
# Smallest output size
59
MIN_NUM_SEQS = 8
59
-# Block size used for kv cache updating kernel
60
-NUM_SLICES_PER_KV_CACHE_UPDATE_BLOCK = 8
60
61
62
#########################################################
@@ -139,7 +138,11 @@ class TPUModelRunner(LoRAModelRunnerMixin):
138
self.pin_memory = is_pin_memory_available()
139
self.dtype = self.model_config.dtype
140
if cache_config.cache_dtype == "auto":
142
- self.kv_cache_dtype = self.dtype
141
+ model_dtype = self.dtype
142
+ if isinstance(model_dtype, str):
143
+ self.kv_cache_dtype = STR_DTYPE_TO_TORCH_DTYPE[model_dtype]
144
+ else:
145
+ self.kv_cache_dtype = model_dtype
146
else:
147
self.kv_cache_dtype = STR_DTYPE_TO_TORCH_DTYPE[
148
cache_config.cache_dtype]
@@ -192,6 +195,14 @@ class TPUModelRunner(LoRAModelRunnerMixin):
195
self.max_num_encoder_input_tokens = encoder_compute_budget
196
self.encoder_cache_size = encoder_cache_size
197
198
+ self._num_slices_per_kv_cache_update_block = \
199
+ _get_num_slices_per_kv_cache_update_block(get_page_size_bytes(
200
+ block_size=self.block_size,
201
+ num_kv_heads=self.num_kv_heads,
202
+ head_size=self.head_size,
203
+ kv_cache_dtype=self.kv_cache_dtype,
204
+ ))
205
+
206
# Lazy initialization
207
self.model: nn.Module # Set after load_model
208
self.kv_caches: list[torch.Tensor] = []
@@ -719,7 +730,7 @@ class TPUModelRunner(LoRAModelRunnerMixin):
730
num_kv_update_slices = slot_mapping_metadata.shape[0]
731
padded_num_slices = _get_padded_num_kv_cache_update_slices(
732
padded_total_num_scheduled_tokens, self.max_num_reqs,
722
- self.block_size)
733
+ self.block_size, self._num_slices_per_kv_cache_update_block)
734
slot_mapping_metadata = np.pad(
735
slot_mapping_metadata,
736
[[0, padded_num_slices - len(slot_mapping_metadata)], [0, 0]],
@@ -750,8 +761,8 @@ class TPUModelRunner(LoRAModelRunnerMixin):
761
num_kv_update_slices=torch.tensor([num_kv_update_slices],
762
dtype=torch.int32,
763
device=self.device),
753
- num_slices_per_kv_cache_update_block=
754
- NUM_SLICES_PER_KV_CACHE_UPDATE_BLOCK,
764
+ num_slices_per_kv_cache_update_block=self.
765
+ _num_slices_per_kv_cache_update_block,
766
)
767
# NOTE(woosuk): Due to chunked prefills, there can be at most 1 partial
768
# request in the batch. While we should not sample any token from this
@@ -1197,7 +1208,8 @@ class TPUModelRunner(LoRAModelRunnerMixin):
1208
position_ids = torch.zeros(num_tokens,
1209
dtype=torch.int32).to(self.device)
1210
padded_num_slices = _get_padded_num_kv_cache_update_slices(
1200
- num_tokens, self.max_num_reqs, self.block_size)
1211
+ num_tokens, self.max_num_reqs, self.block_size,
1212
+ self._num_slices_per_kv_cache_update_block)
1213
num_kv_update_slices = torch.tensor([padded_num_slices],
1214
dtype=torch.int32).to(self.device)
1215
slot_mapping = torch.zeros((3, padded_num_slices),
@@ -1220,8 +1232,8 @@ class TPUModelRunner(LoRAModelRunnerMixin):
1232
query_start_loc=query_start_loc,
1233
num_seqs=num_seqs,
1234
num_kv_update_slices=num_kv_update_slices,
1223
- num_slices_per_kv_cache_update_block=
1224
- NUM_SLICES_PER_KV_CACHE_UPDATE_BLOCK,
1235
+ num_slices_per_kv_cache_update_block=self.
1236
+ _num_slices_per_kv_cache_update_block,
1237
)
1238
1239
if self.is_multimodal_model:
@@ -1826,19 +1838,41 @@ def _get_padded_token_len(paddings: list[int], x: int) -> int:
1838
return paddings[index]
1839
1840
1829
-def _get_padded_num_kv_cache_update_slices(num_tokens: int, max_num_reqs: int,
1830
- page_size: int) -> int:
1841
+def _get_padded_num_kv_cache_update_slices(
1842
+ num_tokens: int, max_num_reqs: int, page_size: int,
1843
+ num_slices_per_kv_cache_update_block: int) -> int:
1844
"""Calculates the padded number of KV cache update slices to avoid
1845
recompilation."""
1846
padded_num_slices = 2 * max_num_reqs + num_tokens // page_size
1847
padded_num_slices = min(padded_num_slices, num_tokens)
1848
padded_num_slices = (
1836
- padded_num_slices + NUM_SLICES_PER_KV_CACHE_UPDATE_BLOCK - 1
1837
- ) // NUM_SLICES_PER_KV_CACHE_UPDATE_BLOCK * \
1838
- NUM_SLICES_PER_KV_CACHE_UPDATE_BLOCK
1849
+ padded_num_slices + num_slices_per_kv_cache_update_block - 1
1850
+ ) // num_slices_per_kv_cache_update_block * \
1851
+ num_slices_per_kv_cache_update_block
1852
return padded_num_slices
1853
1854
1855
+def _get_num_slices_per_kv_cache_update_block(page_size_bytes: int) -> int:
1856
+ """Find the optimum number of slices to copy per Pallas program instance.
1857
+
1858
+ Increasing the number of slices copied in one instance of the kernel program
1859
+ will increase HBM bandwidth utilization via more in-flight DMAs.
1860
+
1861
+ However, it will also use more VMEM, and experimentally, we observed
1862
+ performance regression at 128 slices on v6e, likely due to running
1863
+ out of scalar registers. Thus this function will limit the number of
1864
+ slices to 64.
1865
+ """
1866
+ # Conservative VMEM usage limit: 32 MiB
1867
+ vmem_limit = 32 * 1024 * 1024
1868
+ num_slices_per_block = vmem_limit // page_size_bytes
1869
+ assert num_slices_per_block > 0, "Number of slices should be positive"
1870
+ num_slices_per_block = prev_power_of_2(num_slices_per_block)
1871
+ if num_slices_per_block > 64:
1872
+ num_slices_per_block = 64
1873
+ return num_slices_per_block
1874
+
1875
+
1876
def replace_set_lora(model):
1877
1878
def _tpu_set_lora(