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, 19:45 UTC
[V1][TPU] Remove unnecessary padding for running on TPU. (#14467)
iefgnoix committed
Mar 8, 2025 at 18:56 UTC
10f755278953adec6c8e0916c0b2f636b4dd21e8
2 files changed
+6
-18
vllm/v1/attention/backends/pallas.py
+2
-2
@@ -12,8 +12,8 @@ from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,
12
from vllm.attention.backends.utils import CommonAttentionState
13
14
# These are the 2 tunable parameters of the paged attention Pallas kernel.
15
-NUM_QUERIES_PER_BLOCK = 16
16
-NUM_KV_PAGES_PER_BLOCK = 256
15
+NUM_QUERIES_PER_BLOCK = 32
16
+NUM_KV_PAGES_PER_BLOCK = 128
17
18
19
class PallasAttentionBackend(AttentionBackend):
vllm/v1/worker/tpu_model_runner.py
+4
-16
@@ -23,9 +23,7 @@ from vllm.multimodal.utils import group_mm_inputs_by_modality
23
from vllm.sampling_params import SamplingType
24
from vllm.sequence import IntermediateTensors
25
from vllm.utils import LayerBlockType, cdiv, is_pin_memory_available
26
-from vllm.v1.attention.backends.pallas import (NUM_KV_PAGES_PER_BLOCK,
27
- NUM_QUERIES_PER_BLOCK,
28
- PallasAttentionBackend,
26
+from vllm.v1.attention.backends.pallas import (PallasAttentionBackend,
27
PallasMetadata)
28
from vllm.v1.core.encoder_cache_manager import compute_encoder_budget
29
from vllm.v1.kv_cache_interface import (FullAttentionSpec, KVCacheConfig,
@@ -78,10 +76,8 @@ class TPUModelRunner:
76
self.block_size = cache_config.block_size
77
self.max_model_len = model_config.max_model_len
78
self.max_num_blocks_per_req = cdiv(self.max_model_len, self.block_size)
81
- self.max_num_tokens = _get_padded_number(
82
- scheduler_config.max_num_batched_tokens, NUM_QUERIES_PER_BLOCK)
83
- self.max_num_reqs = _get_padded_number(scheduler_config.max_num_seqs,
84
- NUM_QUERIES_PER_BLOCK)
79
+ self.max_num_tokens = scheduler_config.max_num_batched_tokens
80
+ self.max_num_reqs = scheduler_config.max_num_seqs
81
82
# Model-related.
83
self.num_attn_layers = model_config.get_num_layers_by_block_type(
@@ -142,16 +138,8 @@ class TPUModelRunner:
138
device="cpu")
139
self.slot_mapping_np = self.slot_mapping_cpu.numpy()
140
145
- # self.input_batch.block_table has a shape of [max_num_reqs,
146
- # max_num_blocks_per_req]. To reduce the number of recompilation,
147
- # we want the block_table.shape[0] to be num_tokens.
148
- # To make the block_table to be compatible with the paged attention
149
- # kernel, we want the block_table[1] to be multiple of
150
- # NUM_KV_PAGES_PER_BLOCK.
151
- padded_max_num_blocks_per_req = _get_padded_number(
152
- self.max_num_blocks_per_req, NUM_KV_PAGES_PER_BLOCK)
141
self.block_table_cpu = torch.zeros(
154
- (self.max_num_tokens, padded_max_num_blocks_per_req),
142
+ (self.max_num_tokens, self.max_num_blocks_per_req),
143
dtype=self.input_batch.block_table.get_cpu_tensor().dtype,
144
device="cpu")
145