Skip to content

Commit 1f25d60

Browse files
authored
[Fix] Cap max tokens to prevent potential OOM (#3720)
### What this PR does / why we need it? Caps the calculated maximum number of tokens at 512. This prevents allocating an excessively large buffer when a cudagraph capture size is not specified, mitigating the risk of out-of-memory errors. ### Does this PR introduce _any_ user-facing change? None. ### How was this patch tested? None. - vLLM version: v0.11.0rc3 - vLLM main: vllm-project/vllm@17c540a Signed-off-by: Yizhou Liu <[email protected]>
1 parent 63c363d commit 1f25d60

File tree

1 file changed

+3
-1
lines changed

1 file changed

+3
-1
lines changed

vllm_ascend/worker/model_runner_v1.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -586,7 +586,9 @@ def _init_mc2_tokens_capacity(self):
586586
if self.compilation_config.cudagraph_capture_sizes:
587587
max_num_tokens = self.compilation_config.cudagraph_capture_sizes[0]
588588
else:
589-
max_num_tokens = self.max_num_reqs * self.uniform_decode_query_len
589+
# NOTE: To save memory, we cap the max number of tokens to 512.
590+
max_num_tokens = min(
591+
self.max_num_reqs * self.uniform_decode_query_len, 512)
590592
tp_size = self.parallel_config.tensor_parallel_size
591593
# Use integer arithmetic for ceiling division.
592594
num_tokens_per_tp_rank = (max_num_tokens + tp_size - 1) // tp_size

0 commit comments

Comments
 (0)