[Bugfix][ROCm][MLA] Pass q/kv dtypes to get_mla_metadata_v1 in FP8 decode (#46997)
Signed-off-by: pei.zhang <pei.zhang@amd.com> Co-authored-by: TJian <tunjian.tan@embeddedllm.com>
peizhang56 committed
Jun 30, 2026 at 05:31 UTC
aab7af0bcbf99e2a9462e112dfbcd8a15528ff09
3 files changed
+209
.buildkite/test-amd.yaml
+1
@@ -2828,6 +2828,7 @@ steps:
2828
- rocm-smi
2829
- python3 examples/basic/offline_inference/chat.py --attention-backend TRITON_ATTN
2830
- pytest -v -s tests/kernels/attention/test_attention_selector.py
2831
+ - pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py
2832
2833
- label: Kernels Attention Test %N # TBD
2834
timeout_in_minutes: 60
tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py
new
+202
@@ -0,0 +1,202 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Regression test for AITER MLA persistent decode metadata dtypes.
4
+
5
+For the gfx950 fp8/fp8 nhead=32 qlen=1 fold path, the split/reduce metadata
6
+layout depends on the q/kv element size. The builder must forward dtype_q/dtype_kv
7
+to ``get_mla_metadata_v1``; omitting them lays out the work for the wrong dtype
8
+and corrupts decode output. The test pins the builder's metadata to a golden
9
+recomputed at runtime with the explicit correct dtypes.
10
+"""
11
+
12
+import types
13
+from unittest.mock import patch
14
+
15
+import pytest
16
+import torch
17
+
18
+from vllm._aiter_ops import is_aiter_found
19
+from vllm.platforms import current_platform
20
+
21
+
22
+def _on_gfx950() -> bool:
23
+ if not (current_platform.is_rocm() and is_aiter_found()):
24
+ return False
25
+ from vllm.platforms.rocm import on_gfx950
26
+
27
+ return on_gfx950()
28
+
29
+
30
+pytestmark = pytest.mark.skipif(
31
+ not _on_gfx950(),
32
+ reason="AITER MLA fp8 persistent decode metadata is gfx950-only",
33
+)
34
+
35
+# The fold path that the bug corrupted: fp8 query + fp8 KV-cache, 32 query
36
+# heads, single-token decode, batch 128, context 8192, page_size 1.
37
+NUM_QUERY_HEADS = 32
38
+DECODE_QLEN = 1
39
+BATCH_SIZE = 128
40
+CONTEXT_LEN = 8192
41
+PAGE_SIZE = 1
42
+
43
+# Expected dtypes for this fold path: bf16 model dtype -> bf16 query; fp8
44
+# KV-cache -> fp8_e4m3 kv.
45
+EXPECTED_Q_DTYPE = torch.bfloat16
46
+EXPECTED_KV_DTYPE = torch.float8_e4m3fn
47
+
48
+# The split/reduce content tensors filled by get_mla_metadata_v1. work_meta_data
49
+# is excluded: it holds raw device pointers, never equal across allocations.
50
+_CONTENT_METADATA_FIELDS = (
51
+ "work_indptr",
52
+ "work_info_set",
53
+ "reduce_indptr",
54
+ "reduce_final_map",
55
+ "reduce_partial_map",
56
+)
57
+
58
+# The builder's get_mla_metadata_v1 call passes 6 input args then 6 output
59
+# buffers (see AiterMLAMetadataBuilder._build_decode). Output order -> field.
60
+_NUM_INPUT_ARGS = 6
61
+_OUTPUT_ARG_FIELDS = (
62
+ "work_meta_data",
63
+ "work_info_set",
64
+ "work_indptr",
65
+ "reduce_indptr",
66
+ "reduce_final_map",
67
+ "reduce_partial_map",
68
+)
69
+
70
+
71
+def _build_decode_metadata():
72
+ """Build AITER MLA decode metadata for the fp8/fp8 nhead=32 fold path.
73
+
74
+ Returns ``(metadata, captured)`` where ``captured`` records the positional
75
+ args/kwargs the builder passed to ``get_mla_metadata_v1``, so the golden can
76
+ be recomputed from the identical inputs.
77
+ """
78
+ from tests.v1.attention.utils import (
79
+ BatchSpec,
80
+ create_common_attn_metadata,
81
+ create_vllm_config,
82
+ )
83
+ from vllm.config.vllm import set_current_vllm_config
84
+ from vllm.v1.attention.backends.registry import AttentionBackendEnum
85
+ from vllm.v1.kv_cache_interface import MLAAttentionSpec
86
+ from vllm.v1.worker.workspace import init_workspace_manager
87
+
88
+ device = torch.device("cuda:0")
89
+
90
+ vllm_config = create_vllm_config(
91
+ model_name="deepseek-ai/DeepSeek-R1",
92
+ max_model_len=CONTEXT_LEN,
93
+ # One flat page per token (page_size=1); +buffer for the null block.
94
+ num_gpu_blocks=BATCH_SIZE * CONTEXT_LEN + 200,
95
+ block_size=PAGE_SIZE,
96
+ max_num_seqs=BATCH_SIZE,
97
+ max_num_batched_tokens=8192,
98
+ hf_config_override={"num_attention_heads": NUM_QUERY_HEADS},
99
+ )
100
+ vllm_config.cache_config.cache_dtype = "fp8"
101
+
102
+ spec = MLAAttentionSpec(
103
+ block_size=PAGE_SIZE,
104
+ num_kv_heads=1,
105
+ head_size=vllm_config.model_config.get_head_size(),
106
+ dtype=vllm_config.model_config.dtype,
107
+ cache_dtype_str="fp8",
108
+ )
109
+
110
+ builder_cls = AttentionBackendEnum.ROCM_AITER_MLA.get_class().get_builder_cls()
111
+
112
+ # The builder reads layer.prefill_backend from static_forward_context; a
113
+ # stub with the attribute is enough for metadata construction.
114
+ layer_name = "placeholder"
115
+ vllm_config.compilation_config.static_forward_context[layer_name] = (
116
+ types.SimpleNamespace(prefill_backend=None)
117
+ )
118
+
119
+ init_workspace_manager(device)
120
+
121
+ batch_spec = BatchSpec(
122
+ seq_lens=[CONTEXT_LEN] * BATCH_SIZE,
123
+ query_lens=[DECODE_QLEN] * BATCH_SIZE,
124
+ )
125
+
126
+ captured: dict = {}
127
+
128
+ with set_current_vllm_config(vllm_config):
129
+ builder = builder_cls(spec, [layer_name], vllm_config, device)
130
+ common_attn_metadata = create_common_attn_metadata(
131
+ batch_spec, PAGE_SIZE, device, arange_block_indices=True
132
+ )
133
+
134
+ import aiter
135
+
136
+ real_get_mla_metadata_v1 = aiter.get_mla_metadata_v1
137
+
138
+ def spy(*args, **kwargs):
139
+ captured["args"] = args
140
+ captured["kwargs"] = dict(kwargs)
141
+ return real_get_mla_metadata_v1(*args, **kwargs)
142
+
143
+ with patch("aiter.get_mla_metadata_v1", spy):
144
+ metadata = builder.build(
145
+ common_prefix_len=0,
146
+ common_attn_metadata=common_attn_metadata,
147
+ )
148
+
149
+ return metadata, captured
150
+
151
+
152
+def _compute_golden_metadata(captured: dict) -> dict[str, torch.Tensor]:
153
+ """Recompute the persistent metadata with explicit fp8/bf16 dtypes.
154
+
155
+ Replays ``get_mla_metadata_v1`` on the builder's exact input tensors with
156
+ fresh output buffers and the explicitly-correct dtypes. This reference must
157
+ match the builder's output when the fix is in place.
158
+ """
159
+ import aiter
160
+
161
+ args = captured["args"]
162
+ inputs = args[:_NUM_INPUT_ARGS]
163
+ # Fresh copies so the golden does not alias the builder's persistent buffers.
164
+ fresh_outputs = [arg.clone() for arg in args[_NUM_INPUT_ARGS:]]
165
+
166
+ golden_kwargs = dict(captured["kwargs"])
167
+ golden_kwargs["dtype_q"] = EXPECTED_Q_DTYPE
168
+ golden_kwargs["dtype_kv"] = EXPECTED_KV_DTYPE
169
+
170
+ aiter.get_mla_metadata_v1(*inputs, *fresh_outputs, **golden_kwargs)
171
+
172
+ return dict(zip(_OUTPUT_ARG_FIELDS, fresh_outputs))
173
+
174
+
175
+def test_persistent_decode_metadata_matches_fp8_golden():
176
+ """The builder's metadata must match the dtype-correct golden.
177
+
178
+ Regression guard: the fixed builder forwards fp8/bf16 dtypes so its
179
+ split/reduce metadata matches the golden recomputed with those explicit
180
+ dtypes. Dropping the dtypes (the original bug) produces a different layout
181
+ and fails this test.
182
+ """
183
+ metadata, captured = _build_decode_metadata()
184
+
185
+ # qlen=1 must take the persistent-metadata path for this to be meaningful.
186
+ assert metadata.decode is not None
187
+ assert metadata.decode.has_persistent_metadata
188
+ assert metadata.work_meta_data is not None
189
+
190
+ golden = _compute_golden_metadata(captured)
191
+
192
+ mismatched = [
193
+ name
194
+ for name in _CONTENT_METADATA_FIELDS
195
+ if getattr(metadata, name).shape != golden[name].shape
196
+ or not torch.equal(getattr(metadata, name), golden[name])
197
+ ]
198
+ assert not mismatched, (
199
+ "AITER MLA persistent decode metadata does not match the fp8/bf16 "
200
+ f"golden for fields {mismatched}; the builder must forward "
201
+ "dtype_q/dtype_kv to get_mla_metadata_v1."
202
+ )
vllm/v1/attention/backends/mla/rocm_aiter_mla.py
+6
@@ -199,6 +199,10 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
199
else:
200
kv_cache_dtype_str = "bf16"
201
kv_dtype = dtypes.d_dtypes.get(kv_cache_dtype_str, dtypes.bf16)
202
+ # Persist for get_mla_metadata_v1 (decode build): omitting these causes
203
+ # wrong split/reduce metadata for the gfx950 fp8 nhead=32 fold path.
204
+ self._mla_q_dtype = q_dtype
205
+ self._mla_kv_dtype = kv_dtype
206
(
207
(work_meta_data_size, work_meta_data_type),
208
(work_indptr_size, work_indptr_type),
@@ -534,6 +538,8 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
538
max_seqlen_qo=max_qo_len,
539
uni_seqlen_qo=max_qo_len,
540
fast_mode=True,
541
+ dtype_q=self._mla_q_dtype,
542
+ dtype_kv=self._mla_kv_dtype,
543
)
544
has_persistent_metadata = True
545