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, 21:00 UTC
[v1] Add comments to the new ragged paged attention Pallas kernel (#14155)
Signed-off-by: Xiongfei Wei <isaacwxf23@gmail.com> Co-authored-by: Michael Goin <mgoin64@gmail.com>
iefgnoix committed
Mar 3, 2025 at 15:00 UTC
79e4937c65d5f553f878293a0da50f83b3773141
1 file changed
+5
-1
vllm/v1/attention/backends/pallas.py
+5
-1
@@ -11,6 +11,7 @@ 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 = 16
16
NUM_KV_PAGES_PER_BLOCK = 128
17
@@ -154,6 +155,9 @@ class PallasAttentionBackendImpl(AttentionImpl):
155
write_to_kv_cache(key, value, key_cache, value_cache, slot_mapping)
156
157
query = query * self.scale
158
+ # use_kernel switches between using kernel or reference implementation
159
+ # (non kernel: https://github.com/pytorch/xla/blob/cee0820e78fc9675e2d0511db891fd44342e890d/torch_xla/experimental/custom_kernel.py#L890).
160
+ use_kernel = False
161
output = torch.ops.xla.ragged_paged_attention(
162
query,
163
key_cache,
@@ -164,7 +168,7 @@ class PallasAttentionBackendImpl(AttentionImpl):
168
attn_metadata.num_seqs,
169
num_kv_pages_per_block=NUM_KV_PAGES_PER_BLOCK,
170
num_queries_per_block=NUM_QUERIES_PER_BLOCK,
167
- use_kernel=False,
171
+ use_kernel=use_kernel,
172
)
173
174
return output.reshape(num_tokens, hidden_size)