From 5598b5b5f9433c8f8fe2edad9b5db7d2cff29d1e Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Mon, 21 Sep 2026 11:56:24 +0200 Subject: [PATCH] fix: size upper bound max num partial tiles manually Signed-off-by: simondanielsson --- vllm/v1/attention/backends/mla/rocm_aiter_mla.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/vllm/v1/attention/backends/mla/rocm_aiter_mla.py b/vllm/v1/attention/backends/mla/rocm_aiter_mla.py index 75173e59a3c4..b84a334eef3f 100644 --- a/vllm/v1/attention/backends/mla/rocm_aiter_mla.py +++ b/vllm/v1/attention/backends/mla/rocm_aiter_mla.py @@ -772,9 +772,18 @@ def _init_fp8_prefill_ps_buffers( device=device, ) + from vllm.platforms import current_platform from vllm.v1.worker.workspace import current_workspace_manager - max_num_partial_tiles = reduce_partial_map_size + # AITER metadata sizing assumes all requests can carry max_prefill_qlen tokens, + # which is a loose worst case in the number of QO tiles. Rather, the sum of + # qlens is bounded by max_num_batched_tokens. Manually compute the number of + # qo tiles to avoid OOM at startup. + # TODO: AITER should give us this budget constrained value + qo_tile_cnt = ( + cdiv(max_num_batched_tokens, _FP8_PREFILL_TILE_Q) + max_num_reqs - 1 + ) + max_num_partial_tiles = qo_tile_cnt + current_platform.num_compute_units() reservations: list[tuple[tuple[int, ...], torch.dtype]] = [ ( (max_num_partial_tiles * _FP8_PREFILL_TILE_Q, num_head_k, v_head_dim),