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, 18:30 UTC
[V1][TPU] Apply the ragged paged attention kernel fix and remove the padding. (#14846)
Signed-off-by: Xiongfei Wei <isaacwxf23@gmail.com>
iefgnoix committed
Mar 17, 2025 at 01:48 UTC
b4ad56c1bd2fd39028f64919a11a4c5af96bf0c5
2 files changed
+8
-11
requirements/tpu.txt
+6
-6
@@ -17,9 +17,9 @@ ray[data]
17
--find-links https://storage.googleapis.com/libtpu-releases/index.html
18
--find-links https://storage.googleapis.com/jax-releases/jax_nightly_releases.html
19
--find-links https://storage.googleapis.com/jax-releases/jaxlib_nightly_releases.html
20
-torch @ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch-2.7.0.dev20250306%2Bcxx11-cp39-cp39-linux_x86_64.whl ; python_version == "3.9"
21
-torch @ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch-2.7.0.dev20250306%2Bcxx11-cp310-cp310-linux_x86_64.whl ; python_version == "3.10"
22
-torch @ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch-2.7.0.dev20250306%2Bcxx11-cp311-cp311-linux_x86_64.whl ; python_version == "3.11"
23
-torch_xla[tpu, pallas] @ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.7.0.dev20250306%2Bcxx11-cp39-cp39-linux_x86_64.whl ; python_version == "3.9"
24
-torch_xla[tpu, pallas] @ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.7.0.dev20250306%2Bcxx11-cp310-cp310-linux_x86_64.whl ; python_version == "3.10"
25
-torch_xla[tpu, pallas] @ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.7.0.dev20250306%2Bcxx11-cp311-cp311-linux_x86_64.whl ; python_version == "3.11"
20
+torch @ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch-2.8.0.dev20250314%2Bcxx11-cp39-cp39-linux_x86_64.whl ; python_version == "3.9"
21
+torch @ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch-2.8.0.dev20250314%2Bcxx11-cp310-cp310-linux_x86_64.whl ; python_version == "3.10"
22
+torch @ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch-2.8.0.dev20250314%2Bcxx11-cp311-cp311-linux_x86_64.whl ; python_version == "3.11"
23
+torch_xla[tpu, pallas] @ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.8.0.dev20250314%2Bcxx11-cp39-cp39-linux_x86_64.whl ; python_version == "3.9"
24
+torch_xla[tpu, pallas] @ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.8.0.dev20250314%2Bcxx11-cp310-cp310-linux_x86_64.whl ; python_version == "3.10"
25
+torch_xla[tpu, pallas] @ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.8.0.dev20250314%2Bcxx11-cp311-cp311-linux_x86_64.whl ; python_version == "3.11"
vllm/v1/worker/tpu_model_runner.py
+2
-5
@@ -23,8 +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
- 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,
@@ -139,10 +138,8 @@ class TPUModelRunner:
138
device="cpu")
139
self.slot_mapping_np = self.slot_mapping_cpu.numpy()
140
142
- padded_max_num_blocks_per_req = _get_padded_number(
143
- self.max_num_blocks_per_req, NUM_KV_PAGES_PER_BLOCK)
141
self.block_table_cpu = torch.zeros(
145
- (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