xinyun/vllm
1
0
Fork 0
mirror of https://git.datalinker.icu/vllm-project/vllm.git synced 2026-01-29 03:57:14 +08:00
Code Issues Packages Projects Releases Wiki Activity

[TPU] Temporary fix vmem oom for long model len by reducing page size (#20278)

Browse Source 
Signed-off-by: Chenyaaang <chenyangli@google.com>
This commit is contained in:
Chenyaaang 2025-07-07 22:16:16 -07:00 committed by GitHub
parent 7721ef1786
commit e34d130c16
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 6 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
vllm/v1/attention/backends/pallas.py
View File

@ -86,6 +86,12 @@ class PallasAttentionBackend(AttentionBackend):
# spill less likely. Meanwhile we make sure the page size is in [16, 256].
@staticmethod
def get_page_size(vllm_config: VllmConfig) -> int:
# TODO: This is a temporary fix for vmem OOM.
# For long model length, we use 16 page-size to avoid too much
# VMEM spill. A more robust solution should be implemented to
# handle VREG spills.
if vllm_config.model_config.max_model_len > 8192:
return 16
page_size = next_power_of_2(
vllm_config.model_config.max_model_len) // 16
if page_size <= 16: