[ROCm] Synchronize sparse MLA metadata before graph replay (#47404)
Signed-off-by: zihaomu <zmu@amd.com>
ZihaoMu committed
Jul 10, 2026 at 03:59 UTC
bbb0f945ff348d25d4107e935e54c63a9d3b3ddb
2 files changed
+128
tests/kernels/attention/test_rocm_aiter_mla_sparse_metadata_sync.py
new
+125
@@ -0,0 +1,125 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import sys
5
+from types import ModuleType, SimpleNamespace
6
+from unittest.mock import Mock
7
+
8
+import pytest
9
+import torch
10
+
11
+from vllm.platforms import current_platform
12
+
13
+if not current_platform.is_rocm():
14
+ pytest.skip(
15
+ "ROCm AITER sparse MLA metadata sync test requires ROCm.",
16
+ allow_module_level=True,
17
+ )
18
+
19
+from vllm._aiter_ops import is_aiter_found_and_supported
20
+
21
+if not is_aiter_found_and_supported():
22
+ pytest.skip(
23
+ "ROCm AITER sparse MLA metadata sync test requires a supported AITER "
24
+ "installation.",
25
+ allow_module_level=True,
26
+ )
27
+
28
+from vllm.v1.attention.backend import CommonAttentionMetadata
29
+from vllm.v1.attention.backends.mla import rocm_aiter_mla_sparse as sparse_mod
30
+
31
+
32
+class _FakeAiter(ModuleType):
33
+ get_mla_metadata_v1: Mock
34
+
35
+
36
+def _make_builder():
37
+ builder = object.__new__(sparse_mod.ROCMAiterMLASparseMetadataBuilder)
38
+ max_num_batched_tokens = 8
39
+ topk_tokens = 4
40
+
41
+ builder.device = torch.device("cpu")
42
+ builder.kv_cache_spec = SimpleNamespace(block_size=1)
43
+ builder.model_dtype = torch.bfloat16
44
+ builder.topk_tokens = topk_tokens
45
+ builder.req_id_per_token_buffer = torch.zeros(
46
+ max_num_batched_tokens, dtype=torch.int32, device="cpu"
47
+ )
48
+ builder.qo_indptr = torch.arange(
49
+ max_num_batched_tokens + 1, dtype=torch.int32, device="cpu"
50
+ )
51
+ builder.paged_kv_last_page_len = torch.ones(
52
+ max_num_batched_tokens, dtype=torch.int32, device="cpu"
53
+ )
54
+ builder.paged_kv_indices = torch.zeros(
55
+ max_num_batched_tokens * topk_tokens, dtype=torch.int32, device="cpu"
56
+ )
57
+ builder.paged_kv_indptr = torch.zeros(
58
+ max_num_batched_tokens + 1, dtype=torch.int32, device="cpu"
59
+ )
60
+ builder._num_attention_heads = 16
61
+ builder._mla_work_meta_data = torch.empty(1, dtype=torch.int32, device="cpu")
62
+ builder._mla_work_indptr = torch.empty(1, dtype=torch.int32, device="cpu")
63
+ builder._mla_work_info_set = torch.empty(1, dtype=torch.int32, device="cpu")
64
+ builder._mla_reduce_indptr = torch.empty(1, dtype=torch.int32, device="cpu")
65
+ builder._mla_reduce_final_map = torch.empty(1, dtype=torch.int32, device="cpu")
66
+ builder._mla_reduce_partial_map = torch.empty(1, dtype=torch.int32, device="cpu")
67
+ builder._prev_req_extent = 0
68
+ builder._prev_indices_extent = 0
69
+ builder._prev_metadata_key = None
70
+ return builder
71
+
72
+
73
+def _make_common_metadata():
74
+ query_start_loc = torch.tensor([0, 1, 2], dtype=torch.int32, device="cpu")
75
+ seq_lens = torch.tensor([16, 8], dtype=torch.int32, device="cpu")
76
+ return CommonAttentionMetadata(
77
+ query_start_loc=query_start_loc,
78
+ query_start_loc_cpu=query_start_loc,
79
+ seq_lens=seq_lens,
80
+ _seq_lens_cpu=seq_lens,
81
+ num_reqs=2,
82
+ num_actual_tokens=2,
83
+ max_query_len=1,
84
+ max_seq_len=16,
85
+ block_table_tensor=torch.arange(16, dtype=torch.int32, device="cpu").view(2, 8),
86
+ slot_mapping=torch.arange(2, dtype=torch.int64, device="cpu"),
87
+ )
88
+
89
+
90
+def test_sparse_persistent_metadata_syncs_only_after_recompute(monkeypatch):
91
+ builder = _make_builder()
92
+ common_metadata = _make_common_metadata()
93
+ events: list[str] = []
94
+
95
+ def fake_generate_sparse_seqlen_triton(*args, **kwargs):
96
+ return torch.tensor([1, 2], dtype=torch.int32, device="cpu")
97
+
98
+ fake_aiter = _FakeAiter("aiter")
99
+
100
+ def fake_get_mla_metadata_v1(*args, **kwargs):
101
+ events.append("metadata")
102
+
103
+ fake_get_mla_metadata_v1_mock = Mock(side_effect=fake_get_mla_metadata_v1)
104
+ fake_aiter.get_mla_metadata_v1 = fake_get_mla_metadata_v1_mock
105
+ monkeypatch.setitem(sys.modules, "aiter", fake_aiter)
106
+ monkeypatch.setattr(
107
+ sparse_mod, "generate_sparse_seqlen_triton", fake_generate_sparse_seqlen_triton
108
+ )
109
+ monkeypatch.setattr(
110
+ sparse_mod.torch.cuda,
111
+ "current_stream",
112
+ lambda device=None: SimpleNamespace(synchronize=lambda: events.append("sync")),
113
+ )
114
+
115
+ builder.build(common_prefix_len=0, common_attn_metadata=common_metadata)
116
+
117
+ assert events == ["metadata", "sync"]
118
+ assert fake_get_mla_metadata_v1_mock.call_count == 1
119
+
120
+ events.clear()
121
+
122
+ builder.build(common_prefix_len=0, common_attn_metadata=common_metadata)
123
+
124
+ assert events == []
125
+ assert fake_get_mla_metadata_v1_mock.call_count == 1
vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py
+3
@@ -566,6 +566,9 @@ class ROCMAiterMLASparseMetadataBuilder(
566
uni_seqlen_qo=1,
567
fast_mode=True,
568
)
569
+ # The persistent metadata buffers are read by graph replay. Order
570
+ # the async metadata write before the graph-captured decode kernel.
571
+ torch.cuda.current_stream(self.device).synchronize()
572
self._prev_metadata_key = metadata_key
573
574
metadata = ROCMAiterMLASparseMetadata(