[Bugfix][ROCm] Preserve AITER unified-attention metadata during graph replay (#53821)
Signed-off-by: andyluo7 <andy.luo@amd.com> Co-authored-by: OpenAI Codex <noreply@openai.com>
andyluo7 committed
Aug 31, 2026 at 06:53 UTC
76ff0cdff2defded8d912d550299c902ff33cdca
2 files changed
+59
-3
tests/v1/attention/test_rocm_attention_backends_selection.py
+40
@@ -34,6 +34,46 @@ def mock_get_cdna_version():
34
yield
35
36
37
+def test_aiter_unified_attention_uses_dedicated_metadata_builder():
38
+ from vllm.v1.attention.backends.rocm_aiter_unified_attn import (
39
+ RocmAiterUnifiedAttentionBackend,
40
+ RocmAiterUnifiedAttentionMetadataBuilder,
41
+ )
42
+ from vllm.v1.attention.backends.rocm_attn import (
43
+ RocmAttentionBackend,
44
+ RocmAttentionMetadataBuilder,
45
+ )
46
+
47
+ assert RocmAttentionBackend.get_builder_cls() is RocmAttentionMetadataBuilder
48
+ assert (
49
+ RocmAiterUnifiedAttentionBackend.get_builder_cls()
50
+ is RocmAiterUnifiedAttentionMetadataBuilder
51
+ )
52
+
53
+
54
+def test_aiter_unified_attention_capture_preserves_query_start_locations():
55
+ from vllm.v1.attention.backends.rocm_aiter_unified_attn import (
56
+ RocmAiterUnifiedAttentionMetadataBuilder,
57
+ )
58
+
59
+ builder = object.__new__(RocmAiterUnifiedAttentionMetadataBuilder)
60
+ metadata = MagicMock()
61
+ metadata.seq_lens = torch.tensor([1048576, 524288], dtype=torch.int32)
62
+ builder.build = MagicMock(return_value=metadata)
63
+ common = MagicMock()
64
+ expected_query_start_loc = torch.tensor([0, 2, 5], dtype=torch.int32)
65
+ common.query_start_loc = expected_query_start_loc.clone()
66
+ metadata.query_start_loc = common.query_start_loc
67
+
68
+ actual = builder.build_for_cudagraph_capture(common)
69
+
70
+ builder.build.assert_called_once_with(0, common)
71
+ assert actual is metadata
72
+ assert torch.equal(actual.seq_lens, torch.ones_like(actual.seq_lens))
73
+ assert actual.query_start_loc is common.query_start_loc
74
+ assert torch.equal(actual.query_start_loc, expected_query_start_loc)
75
+
76
+
77
@pytest.mark.parametrize(
78
"env_vars, selected_backend, expected_backend_path",
79
[
vllm/v1/attention/backends/rocm_aiter_unified_attn.py
+19
-3
@@ -15,7 +15,12 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
15
kFp8StaticTensorSym,
16
)
17
from vllm.utils.torch_utils import is_quantized_kv_cache
18
-from vllm.v1.attention.backend import AttentionLayer, AttentionType, MultipleOf
18
+from vllm.v1.attention.backend import (
19
+ AttentionLayer,
20
+ AttentionType,
21
+ CommonAttentionMetadata,
22
+ MultipleOf,
23
+)
24
from vllm.v1.attention.backends.rocm_attn import (
25
RocmAttentionBackend,
26
RocmAttentionImpl,
@@ -27,6 +32,17 @@ from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheLayout
32
logger = init_logger(__name__)
33
34
35
+class RocmAiterUnifiedAttentionMetadataBuilder(RocmAttentionMetadataBuilder):
36
+ def build_for_cudagraph_capture(
37
+ self, common_attn_metadata: CommonAttentionMetadata
38
+ ) -> RocmAttentionMetadata:
39
+ attn_metadata = self.build(0, common_attn_metadata)
40
+ # Avoid capturing the real maximum sequence length while preserving
41
+ # query_start_loc, which unified attention consumes during replay.
42
+ attn_metadata.seq_lens.fill_(1)
43
+ return attn_metadata
44
+
45
+
46
class RocmAiterUnifiedAttentionBackend(RocmAttentionBackend):
47
supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]
48
supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
@@ -101,8 +117,8 @@ class RocmAiterUnifiedAttentionBackend(RocmAttentionBackend):
117
return False
118
119
@staticmethod
104
- def get_builder_cls() -> type["RocmAttentionMetadataBuilder"]:
105
- return RocmAttentionMetadataBuilder
120
+ def get_builder_cls() -> type["RocmAiterUnifiedAttentionMetadataBuilder"]:
121
+ return RocmAiterUnifiedAttentionMetadataBuilder
122
123
@classmethod
124
def supports_attn_type(cls, attn_type: str) -> bool: