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 1, 21:15 UTC
[ROCm][perf] Use workspace manager for sparse indexer allocations (#41002)
Signed-off-by: Stig-Arne Grönroos <stig-arne.gronroos@amd.com> Signed-off-by: Tuukka Sarvi <tuukka.sarvi@amd.com> Co-authored-by: Stig-Arne Grönroos <stig-arne.gronroos@amd.com> Co-authored-by: TJian <tunjian.tan@embeddedllm.com>
Tuukka Sarvi committed
Jun 5, 2026 at 09:46 UTC
b4a6f26c904c4dc1653c1bb623c4cb8186347175
1 file changed
+55
-24
vllm/v1/attention/ops/rocm_aiter_mla_sparse.py
+55
-24
@@ -8,6 +8,7 @@ from importlib.util import find_spec
8
import torch
9
import torch.nn.functional as F
10
11
+import vllm.envs as envs
12
from vllm.compilation.breakable_cudagraph import eager_break_during_capture
13
from vllm.forward_context import get_forward_context
14
from vllm.platforms import current_platform
@@ -15,6 +16,7 @@ from vllm.triton_utils import tl, triton
16
from vllm.utils.torch_utils import LayerNameType
17
from vllm.v1.attention.backends.mla.indexer import DeepseekV32IndexerMetadata
18
from vllm.v1.attention.ops.common import pack_seq_triton, unpack_seq_triton
19
+from vllm.v1.worker.workspace import current_workspace_manager
20
21
if current_platform.is_rocm():
22
from vllm.platforms.rocm import _ON_GFX942, _ON_GFX950
@@ -408,8 +410,8 @@ def rocm_fp8_paged_mqa_logits(
410
411
aiter_paged_mqa_logits_module = None
412
# if rocm_aiter_ops.is_enabled():
411
- batch_size, next_n, heads, head_dim = q_fp8.shape
412
- num_blocks, block_size, _, _ = kv_cache_fp8.shape
413
+ batch_size, next_n = q_fp8.shape[:2]
414
+ block_size = kv_cache_fp8.shape[1]
415
416
if rocm_aiter_ops.is_enabled():
417
aiter_paged_mqa_logits_module = paged_mqa_logits_module()
@@ -420,12 +422,10 @@ def rocm_fp8_paged_mqa_logits(
422
aiter_paged_mqa_logits_module.deepgemm_fp8_paged_mqa_logits
423
)
424
batch_size, next_n, heads, _ = q_fp8.shape
423
- out_logits = torch.full(
424
- [batch_size * next_n, max_model_len],
425
- float("-inf"),
426
- device="cuda",
427
- dtype=torch.float32,
425
+ (out_logits,) = current_workspace_manager().get_simultaneous(
426
+ ((batch_size * next_n, max_model_len), torch.float32),
427
)
428
+ out_logits.fill_(float("-inf"))
429
deepgemm_fp8_paged_mqa_logits(
430
q_fp8,
431
kv_cache_fp8,
@@ -444,12 +444,10 @@ def rocm_fp8_paged_mqa_logits(
444
aiter_paged_mqa_logits_module.deepgemm_fp8_paged_mqa_logits_stage1
445
)
446
batch_size, next_n, heads, _ = q_fp8.shape
447
- out_qk = torch.full(
448
- (heads, batch_size * next_n, max_model_len),
449
- float("-inf"),
450
- device="cuda",
451
- dtype=torch.float32,
447
+ (out_qk,) = current_workspace_manager().get_simultaneous(
448
+ ((heads, batch_size * next_n, max_model_len), torch.float32),
449
)
450
+ out_qk.fill_(float("-inf"))
451
deepgemm_fp8_paged_mqa_logits_stage1(
452
q_fp8,
453
kv_cache_fp8,
@@ -647,6 +645,43 @@ def rocm_aiter_sparse_attn_indexer(
645
k_cache_prefix = _resolve_layer_name(k_cache_prefix)
646
# assert isinstance(attn_metadata, dict)
647
if not isinstance(attn_metadata, dict):
648
+ # Profiling early-exit: reserve memory to account for runtime
649
+ # allocations. Must be in the real impl, not the fake impl —
650
+ # torch.compile calls the fake impl under FakeTensor mode where
651
+ # workspace manager operations on the locked real workspace
652
+ # would corrupt PyTorch's dispatch state.
653
+ workspace_manager = current_workspace_manager()
654
+
655
+ # Prefill k_fp8 and k_scale buffers, used by
656
+ # rocm_aiter_sparse_attn_indexer's prefill path
657
+ workspace_manager.get_simultaneous(
658
+ ((total_seq_lens, head_dim), fp8_dtype),
659
+ ((total_seq_lens, 4), torch.uint8),
660
+ )
661
+
662
+ # Decode logits buffer, used by rocm_fp8_paged_mqa_logits.
663
+ # batch_size * next_n <= hidden_states.shape[0] == max_num_batched_tokens
664
+ if _ON_GFX942 or _ON_GFX950:
665
+ workspace_manager.get_simultaneous(
666
+ ((hidden_states.shape[0], max_model_len), torch.float32),
667
+ )
668
+ else:
669
+ workspace_manager.get_simultaneous(
670
+ (
671
+ (q_fp8.shape[1], hidden_states.shape[0], max_model_len),
672
+ torch.float32,
673
+ ),
674
+ )
675
+ # Transient logits tensor peak memory, produced by
676
+ # rocm_fp8_mqa_logits (prefill) and rocm_fp8_paged_mqa_logits
677
+ # (decode). Prefill logits are bounded by
678
+ # VLLM_SPARSE_INDEXER_MAX_LOGITS_MB via chunking in
679
+ # split_indexer_prefill_chunks; decode logits are smaller.
680
+ max_logits_elems = envs.VLLM_SPARSE_INDEXER_MAX_LOGITS_MB * 1024 * 1024
681
+ _ = torch.empty(
682
+ max_logits_elems, dtype=torch.uint8, device=hidden_states.device
683
+ )
684
+
685
return rocm_aiter_sparse_attn_indexer_fake(
686
hidden_states,
687
k_cache_prefix,
@@ -671,7 +706,6 @@ def rocm_aiter_sparse_attn_indexer(
706
has_decode = layer_attn_metadata.num_decodes > 0
707
has_prefill = layer_attn_metadata.num_prefills > 0
708
num_decode_tokens = layer_attn_metadata.num_decode_tokens
674
- device = hidden_states.device if k is None else k.device
709
710
# during speculative decoding, k may be padded to the CUDA graph batch
711
# size while slot_mapping only covers actual tokens.
@@ -703,17 +737,15 @@ def rocm_aiter_sparse_attn_indexer(
737
if has_prefill:
738
prefill_metadata = layer_attn_metadata.prefill
739
assert prefill_metadata is not None
740
+
741
+ workspace_manager = current_workspace_manager()
742
+ k_fp8_full, k_scale_full = workspace_manager.get_simultaneous(
743
+ ((total_seq_lens, head_dim), fp8_dtype),
744
+ ((total_seq_lens, 4), torch.uint8),
745
+ )
746
for chunk in prefill_metadata.chunks:
707
- k_fp8 = torch.empty(
708
- [chunk.total_seq_lens, head_dim],
709
- device=device,
710
- dtype=fp8_dtype,
711
- )
712
- k_scale = torch.empty(
713
- [chunk.total_seq_lens, 4],
714
- device=device,
715
- dtype=torch.uint8,
716
- )
747
+ k_fp8 = k_fp8_full[: chunk.total_seq_lens]
748
+ k_scale = k_scale_full[: chunk.total_seq_lens]
749
if _ON_GFX942:
750
ops.cp_gather_indexer_k_quant_cache(
751
kv_cache,
@@ -731,7 +763,6 @@ def rocm_aiter_sparse_attn_indexer(
763
chunk.cu_seq_lens,
764
token_to_seq=chunk.token_to_seq,
765
)
734
-
766
logits = rocm_fp8_mqa_logits(
767
q_fp8[chunk.token_start : chunk.token_end],
768
(k_fp8, k_scale.view(torch.float32)),