[Test] Add ROCm AITER FP8 MLA prefill accuracy test (#51457)
Signed-off-by: Aarushi Jain <aarushi@amd.com> Co-authored-by: Aarushi Jain <aarushi@amd.com>
Aarushi Jain committed
Aug 8, 2026 at 19:42 UTC
7581c56c86ce8dd881a84739c0a1eb83f9f2fdd6
1 file changed
+161
tests/kernels/attention/test_rocm_aiter_mla_fp8_prefill.py
new
+161
@@ -0,0 +1,161 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Accuracy test for the ROCm AITER FP8 MLA prefill path.
4
+
5
+The FP8 MLA prefill path
6
+(``AiterMLAImpl._mla_fp8_prefill_attn`` -> ``mla_prefill_ps_asm_fwd`` +
7
+``mla_reduce_v1``) is auto-enabled on gfx950 (MI355) when AITER ships those
8
+kernels, and it has no other test coverage. This test drives the *real*
9
+metadata builder and impl methods (via ``object.__new__``, mirroring
10
+``test_rocm_aiter_mla_head_padding.py::test_h12_aiter_mla_decode_matches_reference``)
11
+so it exercises the actual persistent-scheduling metadata contract
12
+(``get_ps_metadata_v1``) plus both kernels -- rather than re-implementing them --
13
+and compares the output against a causal SDPA reference.
14
+
15
+gfx950-only: ``mla_prefill_ps_asm_fwd`` has no gfx942 build, so this is skipped
16
+on every other platform.
17
+"""
18
+
19
+import math
20
+from types import SimpleNamespace
21
+
22
+import pytest
23
+import torch
24
+import torch.nn.functional as F
25
+
26
+from vllm.platforms import current_platform
27
+
28
+
29
+def _fp8_prefill_available() -> bool:
30
+ if not (current_platform.is_rocm() and torch.cuda.is_available()):
31
+ return False
32
+ from vllm.v1.attention.backends.mla.rocm_aiter_mla import (
33
+ _fp8_mla_prefill_supported,
34
+ )
35
+
36
+ return _fp8_mla_prefill_supported()
37
+
38
+
39
+pytestmark = pytest.mark.skipif(
40
+ not _fp8_prefill_available(),
41
+ reason="FP8 MLA prefill needs gfx950 (MI355) + an AITER build exporting "
42
+ "mla_prefill_ps_asm_fwd / mla_reduce_v1",
43
+)
44
+
45
+# DeepSeek-style MLA head dims after kv_b_proj decompression (q/k carry the
46
+# nope+rope dims, v carries v_head_dim). FP8 prefill requires 16-aligned heads.
47
+NUM_HEADS = 16
48
+QK_NOPE_HEAD_DIM = 128
49
+QK_ROPE_HEAD_DIM = 64
50
+QK_HEAD_DIM = QK_NOPE_HEAD_DIM + QK_ROPE_HEAD_DIM # 192
51
+V_HEAD_DIM = 128
52
+SCALE = 1.0 / math.sqrt(QK_HEAD_DIM)
53
+
54
+# Dominated by fp8 (e4m3) rounding of q/k/v. Measured on MI355: worst-element
55
+# abs error ~0.054, so atol=0.1 is ~2x that noise floor; rtol=0.05 bounds larger
56
+# elements.
57
+ATOL, RTOL = 1e-1, 5e-2
58
+
59
+
60
+@pytest.fixture(autouse=True)
61
+def _workspace_manager():
62
+ """The FP8 prefill path draws per-call scratch from the global workspace
63
+ manager, which ``GPUModelRunner.__init__`` normally initializes. Set it up
64
+ for the test and tear it down so no global state leaks to other tests.
65
+ """
66
+ from vllm.v1.worker.workspace import (
67
+ init_workspace_manager,
68
+ reset_workspace_manager,
69
+ )
70
+
71
+ init_workspace_manager(torch.device("cuda"))
72
+ yield
73
+ reset_workspace_manager()
74
+
75
+
76
+def _make_impl():
77
+ """Minimal AiterMLAImpl exposing only what _mla_fp8_prefill_attn reads."""
78
+ from aiter import mla_prefill_ps_asm_fwd, mla_reduce_v1
79
+
80
+ from vllm.v1.attention.backends.mla.rocm_aiter_mla import AiterMLAImpl
81
+
82
+ impl = object.__new__(AiterMLAImpl)
83
+ impl.num_heads = NUM_HEADS
84
+ impl.v_head_dim = V_HEAD_DIM
85
+ impl.scale = SCALE
86
+ impl._mla_prefill_ps_asm_fwd = mla_prefill_ps_asm_fwd
87
+ impl._mla_reduce_v1 = mla_reduce_v1
88
+ return impl
89
+
90
+
91
+def _build_prefill_metadata(seq_lens: list[int], device: torch.device):
92
+ """Build fp8_prefill_* metadata via the real builder (no hand-rolling)."""
93
+ from vllm.v1.attention.backends.mla.rocm_aiter_mla import AiterMLAMetadataBuilder
94
+
95
+ qo_indptr_cpu = torch.zeros(len(seq_lens) + 1, dtype=torch.int32)
96
+ qo_indptr_cpu[1:] = torch.tensor(seq_lens, dtype=torch.int32).cumsum(0)
97
+ qo_indptr = qo_indptr_cpu.to(device)
98
+ total_q = int(qo_indptr_cpu[-1].item())
99
+ max_q = max(seq_lens)
100
+
101
+ builder = object.__new__(AiterMLAMetadataBuilder)
102
+ builder.num_heads = NUM_HEADS
103
+ builder.mla_dims = SimpleNamespace(v_head_dim=V_HEAD_DIM)
104
+ builder._init_fp8_prefill_ps_buffers(
105
+ max_num_reqs=len(seq_lens),
106
+ max_prefill_qlen=max_q,
107
+ max_num_batched_tokens=total_q,
108
+ device=device,
109
+ )
110
+
111
+ # SimpleNamespace stands in for the metadata dataclasses: the builder only
112
+ # reads prefill.{query_start_loc,max_query_len}, num_decodes, and
113
+ # common.query_start_loc_cpu, then populates metadata.fp8_prefill_*.
114
+ prefill = SimpleNamespace(query_start_loc=qo_indptr, max_query_len=max_q)
115
+ metadata = SimpleNamespace(prefill=prefill, num_decodes=0)
116
+ common = SimpleNamespace(query_start_loc_cpu=qo_indptr_cpu)
117
+ builder._build_fp8_prefill_ps_metadata(metadata, common)
118
+ return metadata, total_q
119
+
120
+
121
+def _reference(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
122
+ """Causal SDPA reference on fp8-quantized inputs (matches kernel dtype)."""
123
+ fp8 = current_platform.fp8_dtype()
124
+ qf = q.to(fp8).float().transpose(0, 1) # [H, S, Dqk]
125
+ kf = k.to(fp8).float().transpose(0, 1)
126
+ vf = v.to(fp8).float().transpose(0, 1) # [H, S, Dv]
127
+ out = F.scaled_dot_product_attention(qf, kf, vf, is_causal=True, scale=SCALE)
128
+ return out.transpose(0, 1) # [S, H, Dv]
129
+
130
+
131
+@pytest.mark.parametrize("seq_len", [128, 512])
132
+@torch.inference_mode()
133
+def test_fp8_prefill_matches_reference(seq_len: int) -> None:
134
+ device = torch.device("cuda")
135
+ torch.manual_seed(0)
136
+
137
+ metadata, total_q = _build_prefill_metadata([seq_len], device)
138
+
139
+ q = torch.randn(
140
+ total_q, NUM_HEADS, QK_HEAD_DIM, dtype=torch.bfloat16, device=device
141
+ )
142
+ k = torch.randn(
143
+ total_q, NUM_HEADS, QK_HEAD_DIM, dtype=torch.bfloat16, device=device
144
+ )
145
+ v = torch.randn(total_q, NUM_HEADS, V_HEAD_DIM, dtype=torch.bfloat16, device=device)
146
+ out = torch.zeros(
147
+ total_q, NUM_HEADS * V_HEAD_DIM, dtype=torch.bfloat16, device=device
148
+ )
149
+
150
+ impl = _make_impl()
151
+ impl._mla_fp8_prefill_attn(q, k, v, metadata, out)
152
+
153
+ out_ref = _reference(q, k, v)
154
+
155
+ assert torch.isfinite(out).all()
156
+ torch.testing.assert_close(
157
+ out.view(total_q, NUM_HEADS, V_HEAD_DIM).float(),
158
+ out_ref.float(),
159
+ atol=ATOL,
160
+ rtol=RTOL,
161
+ )