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, 03:30 UTC
[TPU] Temporary fix vmem oom for long model len by reducing page size (#20278)
Signed-off-by: Chenyaaang <chenyangli@google.com>
Chenyaaang committed
Jul 7, 2025 at 22:16 UTC
e34d130c1613dbabc708cd5f059506c887ac81b4
1 file changed
+6
vllm/v1/attention/backends/pallas.py
+6
@@ -86,6 +86,12 @@ class PallasAttentionBackend(AttentionBackend):
86
# spill less likely. Meanwhile we make sure the page size is in [16, 256].
87
@staticmethod
88
def get_page_size(vllm_config: VllmConfig) -> int:
89
+ # TODO: This is a temporary fix for vmem OOM.
90
+ # For long model length, we use 16 page-size to avoid too much
91
+ # VMEM spill. A more robust solution should be implemented to
92
+ # handle VREG spills.
93
+ if vllm_config.model_config.max_model_len > 8192:
94
+ return 16
95
page_size = next_power_of_2(
96
vllm_config.model_config.max_model_len) // 16
97
if page_size <= 16: