biondizzle/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 2 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
1 changed files with 6 additions and 0 deletions
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]. # spill less likely. Meanwhile we make sure the page size is in [16, 256].
@staticmethod @staticmethod
def get_page_size(vllm_config: VllmConfig) -> int: 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( page_size = next_power_of_2(
vllm_config.model_config.max_model_len) // 16 vllm_config.model_config.max_model_len) // 16
if page_size <= 16: if page_size <= 16: