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] Support sliding window and logit soft capping in the paged attention kernel for TPU. (#15732)
Signed-off-by: Xiongfei Wei <isaacwxf23@gmail.com>
iefgnoix committed
Apr 3, 2025 at 14:23 UTC
b6be6f8d1e49d4aa884603e8675dc216be1cbd79
4 files changed
+128
-18
.buildkite/run-tpu-v1-test.sh
+4
-2
@@ -1,6 +1,6 @@
1
#!/bin/bash
2
3
-set -e
3
+set -xue
4
5
# Build the docker image.
6
docker build -f docker/Dockerfile.tpu -t vllm-tpu .
@@ -38,7 +38,9 @@ docker run --privileged --net host --shm-size=16G -it \
38
&& echo TEST_7 \
39
&& pytest -s -v /workspace/vllm/tests/v1/tpu/test_sampler.py \
40
&& echo TEST_8 \
41
- && pytest -s -v /workspace/vllm/tests/v1/tpu/test_topk_topp_sampler.py" \
41
+ && pytest -s -v /workspace/vllm/tests/v1/tpu/test_topk_topp_sampler.py \
42
+ && echo TEST_9 \
43
+ && pytest -s -v /workspace/vllm/tests/v1/tpu/test_pallas.py" \
44
45
46
# TODO: This test fails because it uses RANDOM_SEED sampling
tests/entrypoints/llm/test_accuracy.py
+20
-10
@@ -13,18 +13,24 @@ import pytest
13
14
from vllm.platforms import current_platform
15
16
-MODEL_NAME = "Qwen/Qwen2-1.5B-Instruct"
16
+MODEL_NAMES = [
17
+ "Qwen/Qwen2-1.5B-Instruct",
18
+ "google/gemma-3-1b-it",
19
+]
20
NUM_CONCURRENT = 500
21
TASK = "gsm8k"
22
FILTER = "exact_match,strict-match"
23
RTOL = 0.03
21
-EXPECTED_VALUE = 0.58
24
+EXPECTED_VALUES = {
25
+ "Qwen/Qwen2-1.5B-Instruct": 0.58,
26
+ "google/gemma-3-1b-it": 0.25,
27
+}
28
29
24
-def run_test(more_args=None):
30
+def run_test(model_name, more_args=None):
31
"""Run the end to end accuracy test."""
32
27
- model_args = f"pretrained={MODEL_NAME},max_model_len=4096"
33
+ model_args = f"pretrained={model_name},max_model_len=4096"
34
35
if more_args is not None:
36
model_args = "{},{}".format(model_args, more_args)
@@ -37,9 +43,12 @@ def run_test(more_args=None):
43
)
44
45
measured_value = results["results"][TASK][FILTER]
40
- assert (measured_value - RTOL < EXPECTED_VALUE
41
- and measured_value + RTOL > EXPECTED_VALUE
42
- ), f"Expected: {EXPECTED_VALUE} | Measured: {measured_value}"
46
+ assert model_name in EXPECTED_VALUES, (
47
+ f"Cannot find the expected value for the model {model_name=}")
48
+ expected_value = EXPECTED_VALUES[model_name]
49
+ assert (measured_value - RTOL < expected_value
50
+ and measured_value + RTOL > expected_value
51
+ ), f"Expected: {expected_value} | Measured: {measured_value}"
52
53
54
# TODO: [AlexM] Fix it with new CI/CD tests
@@ -49,7 +58,8 @@ TPU_TP_TEST_STR = "" #"tensor_parallel_size=4"
58
@pytest.mark.skipif(not current_platform.is_cuda()
59
and not current_platform.is_tpu(),
60
reason="V1 is currently only supported on CUDA and TPU")
52
-def test_lm_eval_accuracy_v1_engine(monkeypatch: pytest.MonkeyPatch):
61
+@pytest.mark.parametrize("model", MODEL_NAMES)
62
+def test_lm_eval_accuracy_v1_engine(model, monkeypatch: pytest.MonkeyPatch):
63
"""Run with the V1 Engine."""
64
65
with monkeypatch.context() as m:
@@ -64,7 +74,7 @@ def test_lm_eval_accuracy_v1_engine(monkeypatch: pytest.MonkeyPatch):
74
if TPU_TP_TEST_STR:
75
more_args += ",{}".format(TPU_TP_TEST_STR)
76
67
- run_test(more_args)
77
+ run_test(model, more_args)
78
79
80
def test_lm_eval_accuracy_v0_engine(monkeypatch: pytest.MonkeyPatch):
@@ -72,4 +82,4 @@ def test_lm_eval_accuracy_v0_engine(monkeypatch: pytest.MonkeyPatch):
82
83
with monkeypatch.context() as m:
84
m.setenv("VLLM_USE_V1", "0")
75
- run_test()
85
+ run_test("Qwen/Qwen2-1.5B-Instruct")
tests/v1/tpu/test_pallas.py
new
+98
@@ -0,0 +1,98 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+from unittest.mock import ANY, patch
3
+
4
+import torch
5
+
6
+from vllm.attention.backends.abstract import AttentionType
7
+from vllm.v1.attention.backends.pallas import (NUM_KV_PAGES_PER_BLOCK,
8
+ NUM_QUERIES_PER_BLOCK,
9
+ PallasAttentionBackendImpl,
10
+ PallasMetadata)
11
+
12
+
13
+def test_ragged_paged_attention():
14
+ # We verify that the kernel inputs such as sliding_window, etc. are passed
15
+ # in from the model correctly.
16
+ # The correctness of the paged attention kernel is tested in the kernel
17
+ # library.
18
+ num_heads = 4
19
+ head_size = 128
20
+ scale = 1.0
21
+ num_kv_heads = 4
22
+ sliding_window = 128
23
+ logits_soft_cap = 50.0
24
+ attn_impl = PallasAttentionBackendImpl(
25
+ num_heads=num_heads,
26
+ head_size=head_size,
27
+ scale=scale,
28
+ num_kv_heads=num_kv_heads,
29
+ alibi_slopes=None,
30
+ sliding_window=sliding_window,
31
+ kv_cache_dtype="auto",
32
+ logits_soft_cap=logits_soft_cap,
33
+ attn_type=AttentionType.DECODER,
34
+ )
35
+ mock_vmem_limit_bytes = 1024
36
+ attn_impl.vmem_limit_bytes = mock_vmem_limit_bytes
37
+
38
+ class FakeAttentionLayer:
39
+ _k_scale_float: float
40
+ _v_scale_float: float
41
+
42
+ layer = FakeAttentionLayer()
43
+ layer._k_scale_float = 1.0
44
+ layer._v_scale_float = 1.0
45
+
46
+ num_tokens = 16
47
+ num_blocks = 1024
48
+ block_size = 16
49
+ query = torch.zeros(num_tokens, num_heads * head_size)
50
+ key = torch.zeros(num_tokens, num_kv_heads * head_size)
51
+ value = torch.zeros(num_tokens, num_kv_heads * head_size)
52
+ kv_cache = torch.zeros(num_blocks, block_size, num_kv_heads * 2, head_size)
53
+ slot_mapping = torch.zeros(num_tokens, dtype=torch.int64)
54
+ max_num_reqs = 8
55
+ max_num_blocks_per_req = 8
56
+ block_tables = torch.zeros((max_num_reqs, max_num_blocks_per_req),
57
+ dtype=torch.int32)
58
+ context_lens = torch.ones((max_num_reqs, ), dtype=torch.int32)
59
+ query_lens = [1] * max_num_reqs
60
+ query_start_loc = torch.cumsum(torch.tensor([0] + query_lens,
61
+ dtype=torch.int32),
62
+ dim=0,
63
+ dtype=torch.int32)
64
+ num_seqs = torch.tensor([max_num_reqs], dtype=torch.int32)
65
+ attn_metadata = PallasMetadata(
66
+ slot_mapping=slot_mapping,
67
+ block_tables=block_tables,
68
+ context_lens=context_lens,
69
+ query_start_loc=query_start_loc,
70
+ num_seqs=num_seqs,
71
+ )
72
+
73
+ with patch("torch.ops.xla.ragged_paged_attention"
74
+ ) as mock_ragged_paged_attention:
75
+ attn_impl.forward(
76
+ layer=layer,
77
+ query=query,
78
+ key=key,
79
+ value=value,
80
+ kv_cache=kv_cache,
81
+ attn_metadata=attn_metadata,
82
+ )
83
+
84
+ mock_ragged_paged_attention.assert_called_once_with(
85
+ ANY, # query
86
+ ANY, # kv_cache
87
+ ANY, # context_lens
88
+ ANY, # block_tables
89
+ ANY, # query_start_loc
90
+ ANY, # num_seqs
91
+ num_kv_pages_per_block=NUM_KV_PAGES_PER_BLOCK,
92
+ num_queries_per_block=NUM_QUERIES_PER_BLOCK,
93
+ vmem_limit_bytes=mock_vmem_limit_bytes,
94
+ use_kernel=True,
95
+ sm_scale=scale,
96
+ sliding_window=sliding_window,
97
+ soft_cap=logits_soft_cap,
98
+ )
vllm/v1/attention/backends/pallas.py
+6
-6
@@ -92,6 +92,8 @@ class PallasAttentionBackendImpl(AttentionImpl):
92
self.head_size = head_size
93
self.scale = float(scale)
94
self.num_kv_heads = num_kv_heads
95
+ self.sliding_window = sliding_window
96
+ self.logits_soft_cap = logits_soft_cap
97
98
assert self.num_heads % self.num_kv_heads == 0
99
self.num_queries_per_kv = self.num_heads // self.num_kv_heads
@@ -99,15 +101,10 @@ class PallasAttentionBackendImpl(AttentionImpl):
101
raise NotImplementedError("Head size must be a multiple of 128.")
102
if alibi_slopes is not None:
103
raise NotImplementedError("Alibi slopes is not supported.")
102
- if sliding_window is not None:
103
- raise NotImplementedError("Sliding window is not supported.")
104
if kv_cache_dtype != "auto":
105
raise NotImplementedError("FP8 KV cache dtype is not supported.")
106
if blocksparse_params is not None:
107
raise NotImplementedError("Blocksparse is not supported.")
108
- if logits_soft_cap is not None:
109
- raise NotImplementedError(
110
- "Attention logits soft-capping is not supported.")
108
109
if attn_type != AttentionType.DECODER:
110
raise NotImplementedError("Encoder self-attention and "
@@ -172,7 +169,10 @@ class PallasAttentionBackendImpl(AttentionImpl):
169
num_queries_per_block=NUM_QUERIES_PER_BLOCK,
170
vmem_limit_bytes=self.vmem_limit_bytes,
171
use_kernel=True,
175
- sm_scale=self.scale)
172
+ sm_scale=self.scale,
173
+ sliding_window=self.sliding_window,
174
+ soft_cap=self.logits_soft_cap,
175
+ )
176
177
return output.reshape(num_tokens, hidden_size)
178