A high-throughput and memory-efficient inference and serving engine for LLMs
[ROCm][MLA][Bugfix] Reserve FP8 prefill workspace before lock for Kimi-K2.5 (#42978)
Signed-off-by: Xavier Aguilar <xavier.aguilarfruto@amd.com> Co-authored-by: Douglas Lehr <91553416+dllehr-amd@users.noreply.github.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>
xaguilar-amd committed
Jun 9, 2026 at 07:25 UTC
baacbfcebf5eb9e04f66baf8efb72e3d24b9ef26
1 file changed
+34
-5
vllm/v1/attention/backends/mla/rocm_aiter_mla.py
+34
-5
@@ -242,7 +242,12 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
242
vllm_config.model_config.max_model_len,
243
vllm_config.scheduler_config.max_num_batched_tokens,
244
)
245
- self._init_fp8_prefill_ps_buffers(max_num_reqs, max_prefill_qlen, device)
245
+ self._init_fp8_prefill_ps_buffers(
246
+ max_num_reqs,
247
+ max_prefill_qlen,
248
+ vllm_config.scheduler_config.max_num_batched_tokens,
249
+ device,
250
+ )
251
252
if self.compilation_config.cudagraph_mode.has_full_cudagraphs():
253
self.paged_kv_indptr = torch.zeros(
@@ -257,21 +262,29 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
262
self,
263
max_num_reqs: int,
264
max_prefill_qlen: int,
265
+ max_num_batched_tokens: int,
266
device: torch.device,
267
) -> None:
268
"""Pre-allocate persistent buffers for FP8 MLA prefill PS metadata.
269
270
Uses ``get_ps_metadata_info_v1`` with max values so the buffers are
271
large enough for any batch. ``get_ps_metadata_v1`` fills them
266
- per-batch in ``build()``.
272
+ per-batch in ``build()``. The FP8 prefill forward path also uses the
273
+ global workspace manager for per-call scratch, so reserve its maximum
274
+ shape here before the workspace manager is locked after warmup.
275
276
Args:
277
max_num_reqs: Maximum number of concurrent requests.
278
max_prefill_qlen: Maximum Q-length for a single request in one
279
prefill batch. Should be ``min(max_model_len,
272
- max_num_batched_tokens)`` — the chunked-prefill scheduler
273
- never emits more than ``max_num_batched_tokens`` new tokens
274
- per batch.
280
+ max_num_batched_tokens)`` — a single request never exceeds
281
+ ``max_model_len`` tokens, nor the per-batch token budget.
282
+ max_num_batched_tokens: Maximum number of tokens scheduled in one
283
+ batch. The ``final_lse`` scratch is sized by ``total_q`` (the
284
+ summed Q-length over all prefill requests in the batch), which
285
+ is bounded by this budget rather than by a single request's
286
+ ``max_prefill_qlen`` — concurrent requests can sum to more than
287
+ ``max_model_len`` when ``max_model_len < max_num_batched_tokens``.
288
device: Target device for the buffers.
289
"""
290
from aiter import get_ps_metadata_info_v1
@@ -279,6 +292,7 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
292
# After kv_b_proj decompression, K has num_heads heads (same as Q).
293
# So gqa_ratio=1 and num_head_k=num_heads for the PS kernel.
294
num_head_k = self.num_heads
295
+ v_head_dim = self.mla_dims.v_head_dim
296
# gqa_ratio = 1
297
# qlen_granularity = _FP8_PREFILL_TILE_Q // max(gqa_ratio, 1)
298
qlen_granularity = _FP8_PREFILL_TILE_Q
@@ -318,6 +332,21 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
332
device=device,
333
)
334
335
+ from vllm.v1.worker.workspace import current_workspace_manager
336
+
337
+ max_num_partial_tiles = reduce_partial_map_size
338
+ current_workspace_manager().get_simultaneous(
339
+ (
340
+ (max_num_partial_tiles * _FP8_PREFILL_TILE_Q, num_head_k, v_head_dim),
341
+ torch.float32,
342
+ ),
343
+ (
344
+ (max_num_partial_tiles * _FP8_PREFILL_TILE_Q, num_head_k),
345
+ torch.float32,
346
+ ),
347
+ ((max_num_batched_tokens, num_head_k), torch.float32),
348
+ )
349
+
350
logger.info(
351
"FP8 MLA prefill PS buffers allocated "
352
"(max_batch=%d, max_qlen=%d, num_head_k=%d)",