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] Pad the block_table.shape[1] so the ragged paged attention can handle correctly (#14597)
iefgnoix committed
Mar 11, 2025 at 16:12 UTC
863d315c867e9cc455ac63a0372af71c8c63312b
1 file changed
+5
-2
vllm/v1/worker/tpu_model_runner.py
+5
-2
@@ -23,7 +23,8 @@ 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 (PallasAttentionBackend,
26
+from vllm.v1.attention.backends.pallas import (NUM_KV_PAGES_PER_BLOCK,
27
+ PallasAttentionBackend,
28
PallasMetadata)
29
from vllm.v1.core.encoder_cache_manager import compute_encoder_budget
30
from vllm.v1.kv_cache_interface import (FullAttentionSpec, KVCacheConfig,
@@ -138,8 +139,10 @@ class TPUModelRunner:
139
device="cpu")
140
self.slot_mapping_np = self.slot_mapping_cpu.numpy()
141
142
+ padded_max_num_blocks_per_req = _get_padded_number(
143
+ self.max_num_blocks_per_req, NUM_KV_PAGES_PER_BLOCK)
144
self.block_table_cpu = torch.zeros(
142
- (self.max_num_tokens, self.max_num_blocks_per_req),
145
+ (self.max_num_tokens, padded_max_num_blocks_per_req),
146
dtype=self.input_batch.block_table.get_cpu_tensor().dtype,
147
device="cpu")
148