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, 16:00 UTC
[TPU][V1] Remove ragged attention kernel parameter hard coding (#16041)
Signed-off-by: Chengji Yao <chengjiyao@google.com>
Chengji Yao committed
Apr 4, 2025 at 04:48 UTC
fadc59c0e6b3c348581e18c90d642de40c853ade
2 files changed
+8
-20
vllm/v1/attention/backends/pallas.py
+6
-14
@@ -11,10 +11,6 @@ from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,
11
AttentionLayer, AttentionType)
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 = 32
16
-NUM_KV_PAGES_PER_BLOCK = 128
17
-
14
15
class PallasAttentionBackend(AttentionBackend):
16
@@ -115,13 +111,6 @@ class PallasAttentionBackendImpl(AttentionImpl):
111
tpu_version = torch_xla.tpu.version()
112
if tpu_version < 4:
113
raise NotImplementedError("TPU version must be 4 or higher.")
118
- # NOTE(chengjiyao): the TPU v4's vmem capacity is 16MB
119
- # TODO(chengjiyao): autotune NUM_QUERIES_PER_BLOCK,
120
- # NUM_KV_PAGES_PER_BLOCK and vmem_limit_bytes
121
- if tpu_version == 4:
122
- self.vmem_limit_bytes = 16 * 1024 * 1024
123
- else:
124
- self.vmem_limit_bytes = 64 * 1024 * 1024
114
115
def forward(
116
self,
@@ -165,9 +154,12 @@ class PallasAttentionBackendImpl(AttentionImpl):
154
attn_metadata.block_tables,
155
attn_metadata.query_start_loc,
156
attn_metadata.num_seqs,
168
- num_kv_pages_per_block=NUM_KV_PAGES_PER_BLOCK,
169
- num_queries_per_block=NUM_QUERIES_PER_BLOCK,
170
- vmem_limit_bytes=self.vmem_limit_bytes,
157
+ # By default, the system utilizes optimized block size and
158
+ # vmem_limit_bytes parameters from the kernel repository. However,
159
+ # these can be manually adjusted for debugging if necessary.
160
+ num_kv_pages_per_block=None,
161
+ num_queries_per_block=None,
162
+ vmem_limit_bytes=None,
163
use_kernel=True,
164
sm_scale=self.scale,
165
sliding_window=self.sliding_window,
vllm/v1/worker/tpu_model_runner.py
+2
-6
@@ -24,8 +24,7 @@ from vllm.multimodal.utils import group_mm_inputs_by_modality
24
from vllm.sampling_params import SamplingType
25
from vllm.sequence import IntermediateTensors
26
from vllm.utils import LayerBlockType, cdiv, is_pin_memory_available
27
-from vllm.v1.attention.backends.pallas import (NUM_KV_PAGES_PER_BLOCK,
28
- PallasAttentionBackend,
27
+from vllm.v1.attention.backends.pallas import (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,
@@ -155,11 +154,8 @@ class TPUModelRunner:
154
dtype=torch.int64,
155
device="cpu")
156
self.slot_mapping_np = self.slot_mapping_cpu.numpy()
158
-
159
- padded_max_num_blocks_per_req = _get_padded_number(
160
- self.max_num_blocks_per_req, NUM_KV_PAGES_PER_BLOCK)
157
self.block_table_cpu = torch.zeros(
162
- (self.max_num_tokens, padded_max_num_blocks_per_req),
158
+ (self.max_num_tokens, self.max_num_blocks_per_req),
159
dtype=self.input_batch.block_table.get_cpu_tensor().dtype,
160
device="cpu")
161