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, 13:30 UTC
[TPU][V1] Implicitly adjust page size when there's SMEM OOM (#16871)
Signed-off-by: Chengji Yao <chengjiyao@google.com>
Chengji Yao committed
Apr 21, 2025 at 14:43 UTC
471fe65630e02ec7a85ca6c7566127b801677425
3 files changed
+34
-2
tests/v1/tpu/test_basic.py
+5
-2
@@ -22,6 +22,7 @@ MODELS = [
22
]
23
24
TENSOR_PARALLEL_SIZES = [1]
25
+MAX_NUM_REQS = [16, 1024]
26
27
# TODO: Enable when CI/CD will have a multi-tpu instance
28
# TENSOR_PARALLEL_SIZES = [1, 4]
@@ -32,12 +33,14 @@ TENSOR_PARALLEL_SIZES = [1]
33
@pytest.mark.parametrize("model", MODELS)
34
@pytest.mark.parametrize("max_tokens", [5])
35
@pytest.mark.parametrize("tensor_parallel_size", TENSOR_PARALLEL_SIZES)
36
+@pytest.mark.parametrize("max_num_seqs", MAX_NUM_REQS)
37
def test_basic(
38
vllm_runner: type[VllmRunner],
39
monkeypatch: pytest.MonkeyPatch,
40
model: str,
41
max_tokens: int,
42
tensor_parallel_size: int,
43
+ max_num_seqs: int,
44
) -> None:
45
prompt = "The next numbers of the sequence " + ", ".join(
46
str(i) for i in range(1024)) + " are:"
@@ -51,9 +54,9 @@ def test_basic(
54
# Note: max_num_batched_tokens == 1024 is needed here to
55
# actually test chunked prompt
56
max_num_batched_tokens=1024,
54
- max_model_len=8196,
57
+ max_model_len=8192,
58
gpu_memory_utilization=0.7,
56
- max_num_seqs=16,
59
+ max_num_seqs=max_num_seqs,
60
tensor_parallel_size=tensor_parallel_size) as vllm_model:
61
vllm_outputs = vllm_model.generate_greedy(example_prompts,
62
max_tokens)
vllm/platforms/tpu.py
+14
@@ -97,6 +97,20 @@ class TpuPlatform(Platform):
97
"Using bfloat16 instead.", vllm_config.model_config.dtype)
98
vllm_config.model_config.dtype = torch.bfloat16
99
100
+ if envs.VLLM_USE_V1:
101
+ from vllm.v1.attention.backends.pallas import (
102
+ PallasAttentionBackend)
103
+ min_page_size = PallasAttentionBackend.get_min_page_size(
104
+ vllm_config)
105
+ if min_page_size > vllm_config.cache_config.block_size:
106
+ logger.warning(
107
+ "Increase the page size from %s to %s to make sure there's"
108
+ "no SMEM OOM",
109
+ vllm_config.cache_config.block_size,
110
+ min_page_size,
111
+ )
112
+ vllm_config.cache_config.block_size = min_page_size
113
+
114
parallel_config = vllm_config.parallel_config
115
scheduler_config = vllm_config.scheduler_config
116
if parallel_config.worker_cls == "auto":
vllm/v1/attention/backends/pallas.py
+15
@@ -10,7 +10,9 @@ import torch_xla.experimental.custom_kernel # noqa: F401
10
from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,
11
AttentionLayer, AttentionType)
12
from vllm.attention.backends.utils import CommonAttentionState
13
+from vllm.config import VllmConfig
14
from vllm.logger import init_logger
15
+from vllm.utils import cdiv
16
17
logger = init_logger(__name__)
18
@@ -50,6 +52,19 @@ class PallasAttentionBackend(AttentionBackend):
52
) -> None:
53
raise RuntimeError("swap_blocks is not used for the TPU backend.")
54
55
+ # In recent TPU generations, up to v6e, the SMEM size is 1MB. The
56
+ # block_tables within the PallasMetadata constitute almost the entire SMEM
57
+ # requirement. Its size is max_num_seqs * num_page_per_seq * 4 (Int). Here
58
+ # we simply make sure that the size is smaller than half of SMEM capacity.
59
+ @staticmethod
60
+ def get_min_page_size(vllm_config: VllmConfig) -> int:
61
+ max_num_page_per_req = (1024 * 1024 // 2 //
62
+ vllm_config.scheduler_config.max_num_seqs // 4)
63
+ min_page_size = cdiv(vllm_config.model_config.max_model_len,
64
+ max_num_page_per_req)
65
+ min_page_size = 1 << (min_page_size - 1).bit_length()
66
+ return min_page_size
67
+
68
69
@dataclass
70
class PallasMetadata: