[Bugfix][ROCm] Fix startup OOM in AITER MLA FP8 prefill workspace sizing (#57923)
Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com>
Simon Danielsson committed
Sep 23, 2026 at 18:44 UTC
b5ea0c71078bdbb0f1f9afd934d55156273fe5f8
1 file changed
+10
-1
vllm/v1/attention/backends/mla/rocm_aiter_mla.py
+10
-1
@@ -772,9 +772,18 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
772
device=device,
773
)
774
775
+ from vllm.platforms import current_platform
776
from vllm.v1.worker.workspace import current_workspace_manager
777
777
- max_num_partial_tiles = reduce_partial_map_size
778
+ # AITER metadata sizing assumes all requests can carry max_prefill_qlen tokens,
779
+ # which is a loose worst case in the number of QO tiles. Rather, the sum of
780
+ # qlens is bounded by max_num_batched_tokens. Manually compute the number of
781
+ # qo tiles to avoid OOM at startup.
782
+ # TODO: AITER should give us this budget constrained value
783
+ qo_tile_cnt = (
784
+ cdiv(max_num_batched_tokens, _FP8_PREFILL_TILE_Q) + max_num_reqs - 1
785
+ )
786
+ max_num_partial_tiles = qo_tile_cnt + current_platform.num_compute_units()
787
reservations: list[tuple[tuple[int, ...], torch.dtype]] = [
788
(
789
(max_num_partial_tiles * _FP8_PREFILL_TILE_Q, num_head_k, v_head_dim),