| 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). |
| 47 | # |
| 48 | # The kernel requires 16-aligned heads, so _mla_fp8_prefill_attn replicate-pads |
| 49 | # q/k/v up to get_fp8_prefill_num_heads(num_heads) and slices the output back. |
| 50 | # Cover all three regimes: |
| 51 | # 12 -> padded to 16 (non-divisor, tile+slice). Live case: a K3 rank at TP8. |
| 52 | # 16 -> no padding, output aliases the caller's buffer. |
| 53 | # 24 -> padded to 32. No current model and TP reaches this band (96 heads |
| 54 | # would need TP4, which does not fit in 288 GiB; 128-head models give |
| 55 | # 128/64/32/16/8), so this case exists to keep the general path honest |
| 56 | # for future architectures: it is the only one of the three where |
| 57 | # padding above 16 actually happens, so it is what proves the |
| 58 | # replicate-pad/slice-back argument holds for a target other than 16. |
| 59 | HEAD_COUNTS = [12, 16, 24] |
| 60 | QK_NOPE_HEAD_DIM = 128 |
| 61 | QK_ROPE_HEAD_DIM = 64 |
| 62 | QK_HEAD_DIM = QK_NOPE_HEAD_DIM + QK_ROPE_HEAD_DIM # 192 |
| 63 | V_HEAD_DIM = 128 |
| 64 | SCALE = 1.0 / math.sqrt(QK_HEAD_DIM) |
| 65 | |
| 66 | # Dominated by fp8 (e4m3) rounding of q/k/v. Measured on MI355: worst-element |
| 67 | # abs error ~0.054, so atol=0.1 is ~2x that noise floor; rtol=0.05 bounds larger |
| 68 | # elements. |
| 69 | ATOL, RTOL = 1e-1, 5e-2 |
| 70 | |
| 71 | # The assembly kernel writes bf16 through a raw output pointer. Keep its tensors |
| 72 | # and workspace reservation on one dtype. |
| 73 | ATTN_OUT_DTYPE = torch.bfloat16 |
| 74 | |
| 75 | |
| 76 | @pytest.fixture(autouse=True) |
| 77 | def _workspace_manager(): |
| 78 | """The FP8 prefill path draws per-call scratch from the global workspace |
| 79 | manager, which ``GPUModelRunner.__init__`` normally initializes. Set it up |
| 80 | for the test and tear it down so no global state leaks to other tests. |
| 81 | """ |
| 82 | from vllm.v1.worker.workspace import ( |
| 83 | init_workspace_manager, |
| 84 | reset_workspace_manager, |
| 85 | ) |
| 86 | |
| 87 | init_workspace_manager(torch.device("cuda")) |
| 88 | yield |
| 89 | reset_workspace_manager() |
| 90 | |
| 91 | |
| 92 | def _make_impl(num_heads: int): |
| 93 | """Minimal AiterMLAImpl exposing only what _mla_fp8_prefill_attn reads.""" |
| 94 | from aiter import mla_prefill_ps_asm_fwd, mla_reduce_v1 |
| 95 | |
| 96 | from vllm.v1.attention.backends.mla.rocm_aiter_mla import AiterMLAImpl |
| 97 | |
| 98 | impl = object.__new__(AiterMLAImpl) |
| 99 | impl.num_heads = num_heads |
| 100 | impl.v_head_dim = V_HEAD_DIM |
| 101 | impl.scale = SCALE |
| 102 | impl._mla_prefill_ps_asm_fwd = mla_prefill_ps_asm_fwd |
| 103 | impl._mla_reduce_v1 = mla_reduce_v1 |
| 104 | return impl |
| 105 | |
| 106 | |
| 107 | def _build_prefill_metadata( |
| 108 | seq_lens: list[int], |
| 109 | device: torch.device, |
| 110 | num_heads: int, |
| 111 | attn_out_dtype: torch.dtype = ATTN_OUT_DTYPE, |
| 112 | ): |
| 113 | """Build metadata with the production builder and output dtype.""" |
| 114 | from vllm.v1.attention.backends.mla.rocm_aiter_mla import AiterMLAMetadataBuilder |
| 115 | |
| 116 | qo_indptr_cpu = torch.zeros(len(seq_lens) + 1, dtype=torch.int32) |
| 117 | qo_indptr_cpu[1:] = torch.tensor(seq_lens, dtype=torch.int32).cumsum(0) |
| 118 | qo_indptr = qo_indptr_cpu.to(device) |
| 119 | total_q = int(qo_indptr_cpu[-1].item()) |
| 120 | max_q = max(seq_lens) |
| 121 | |
| 122 | builder = object.__new__(AiterMLAMetadataBuilder) |
| 123 | builder.num_heads = num_heads |
| 124 | builder.mla_dims = SimpleNamespace(v_head_dim=V_HEAD_DIM) |
| 125 | builder._init_fp8_prefill_ps_buffers( |
| 126 | max_num_reqs=len(seq_lens), |
| 127 | max_prefill_qlen=max_q, |
| 128 | max_num_batched_tokens=total_q, |
| 129 | attn_out_dtype=attn_out_dtype, |
| 130 | device=device, |
| 131 | ) |
| 132 | |
| 133 | # SimpleNamespace stands in for the metadata dataclasses: the builder only |
| 134 | # reads prefill.{query_start_loc,max_query_len}, num_decodes, and |
| 135 | # common.query_start_loc_cpu, then populates metadata.fp8_prefill_*. |
| 136 | prefill = SimpleNamespace(query_start_loc=qo_indptr, max_query_len=max_q) |
| 137 | metadata = SimpleNamespace(prefill=prefill, num_decodes=0) |
| 138 | common = SimpleNamespace(query_start_loc_cpu=qo_indptr_cpu) |
| 139 | builder._build_fp8_prefill_ps_metadata(metadata, common) |
| 140 | return metadata, total_q, builder |
| 141 | |
| 142 | |
| 143 | def _reference(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor: |
| 144 | """Causal SDPA reference on fp8-quantized inputs (matches kernel dtype).""" |
| 145 | fp8 = current_platform.fp8_dtype() |
| 146 | qf = q.to(fp8).float().transpose(0, 1) # [H, S, Dqk] |
| 147 | kf = k.to(fp8).float().transpose(0, 1) |
| 148 | vf = v.to(fp8).float().transpose(0, 1) # [H, S, Dv] |
| 149 | out = F.scaled_dot_product_attention(qf, kf, vf, is_causal=True, scale=SCALE) |
| 150 | return out.transpose(0, 1) # [S, H, Dv] |
| 151 | |
| 152 | |
| 153 | @pytest.mark.parametrize("num_heads", HEAD_COUNTS) |
| 154 | @pytest.mark.parametrize("seq_len", [128, 512]) |
| 155 | @torch.inference_mode() |
| 156 | def test_fp8_prefill_matches_reference(seq_len: int, num_heads: int) -> None: |
| 157 | device = torch.device("cuda") |
| 158 | torch.manual_seed(0) |
| 159 | |
| 160 | metadata, total_q, _ = _build_prefill_metadata([seq_len], device, num_heads) |
| 161 | |
| 162 | # Lock the workspace, as GPUModelRunner does after warmup, so the forward |
| 163 | # cannot quietly grow it. The builder reserves scratch at num_head_k and |
| 164 | # the forward requests it at its own padded width; if a future change makes |
| 165 | # the forward's width exceed the builder's, that is an under-reservation |
| 166 | # which is invisible on an unlocked workspace (it just grows) and fatal in |
| 167 | # a real serve. Locking turns it into a failure here. |
| 168 | # |
| 169 | # Note this does not catch the reverse -- a builder wider than the forward |
| 170 | # merely over-reserves -- so it would not have flagged the pre-PR |
| 171 | # max(16, num_heads) sizing, which over-reserves 2.8x at 24 heads. |
| 172 | from vllm.v1.worker.workspace import lock_workspace |
| 173 | |
| 174 | lock_workspace() |
| 175 | |
| 176 | qkv_kwargs = dict(dtype=ATTN_OUT_DTYPE, device=device) |
| 177 | q = torch.randn(total_q, num_heads, QK_HEAD_DIM, **qkv_kwargs) |
| 178 | k = torch.randn(total_q, num_heads, QK_HEAD_DIM, **qkv_kwargs) |
| 179 | v = torch.randn(total_q, num_heads, V_HEAD_DIM, **qkv_kwargs) |
| 180 | out = torch.zeros( |
| 181 | total_q, num_heads * V_HEAD_DIM, dtype=ATTN_OUT_DTYPE, device=device |
| 182 | ) |
| 183 | |
| 184 | impl = _make_impl(num_heads) |
| 185 | impl._mla_fp8_prefill_attn(q, k, v, metadata, out) |
| 186 | |
| 187 | out_ref = _reference(q, k, v) |
| 188 | |
| 189 | assert torch.isfinite(out).all() |
| 190 | torch.testing.assert_close( |
| 191 | out.view(total_q, num_heads, V_HEAD_DIM).float(), |
| 192 | out_ref.float(), |
| 193 | atol=ATOL, |
| 194 | rtol=RTOL, |
| 195 | ) |
| 196 | |
| 197 | |
| 198 | @pytest.mark.parametrize("num_heads", HEAD_COUNTS) |
| 199 | def test_fp8_prefill_metadata_width_matches_forward(num_heads: int) -> None: |
| 200 | """The PS metadata must be built for the head count the kernel is handed. |
| 201 | |
| 202 | ``_init_fp8_prefill_ps_buffers``/``build()`` size the work and reduce maps |
| 203 | from ``num_head_k``, while ``_mla_fp8_prefill_attn`` pads q/k/v to its own |
| 204 | target. Divergence is not a numerical bug in either direction -- maps built |
| 205 | narrower than the kernel width still cover the real heads, since padding |
| 206 | only ever appends duplicate heads above them. It is a sizing bug: narrower |
| 207 | maps mean a lower head alignment, hence more partial tiles and a larger |
| 208 | reservation, while wider ones under-reserve the forward's scratch. Pin the |
| 209 | two to one helper rather than to independent expressions. |
| 210 | """ |
| 211 | from vllm.v1.attention.backends.mla.rocm_aiter_mla import AiterMLAHelper |
| 212 | |
| 213 | width = AiterMLAHelper.get_fp8_prefill_num_heads(num_heads) |
| 214 | assert width % AiterMLAHelper._AITER_MIN_MLA_HEADS == 0 |
| 215 | assert width >= num_heads |
| 216 | assert width - num_heads < AiterMLAHelper._AITER_MIN_MLA_HEADS |
| 217 | # Every count the enable gate accepts must have a usable padded width. |
| 218 | assert AiterMLAHelper.is_valid_num_heads(num_heads) |
| 219 | |
| 220 | # The builder must actually size its maps at that width. Assert against the |
| 221 | # buffer it allocated rather than recomputing the width here, or the test |
| 222 | # passes no matter what _init_fp8_prefill_ps_buffers does. |
| 223 | # |
| 224 | # fp8_ps_reduce_partial_map is sized by the partial-tile count, which is |
| 225 | # gcd-driven and so is a sharp probe of the width used: a lower head |
| 226 | # alignment yields *more* tiles. At 24 heads, sizing at the unpadded count |
| 227 | # gives 4x the tiles (64 vs 16) and ~2.8x the scratch reservation (193.6 |
| 228 | # vs 68.6 MiB at batch=1/qlen=512) -- not wrong numerically (the 24-wide |
| 229 | # maps still cover the real heads) but the reason to pin the widths |
| 230 | # together. |
| 231 | from aiter import get_ps_metadata_info_v1 |
| 232 | |
| 233 | from vllm.v1.attention.backends.mla.rocm_aiter_mla import _FP8_PREFILL_TILE_Q |
| 234 | |
| 235 | def _tiles(nhk: int) -> int: |
| 236 | return get_ps_metadata_info_v1( |
| 237 | batch_size=1, |
| 238 | num_head_k=nhk, |
| 239 | max_qlen=512, |
| 240 | qlen_granularity=_FP8_PREFILL_TILE_Q, |
| 241 | )[5][0] |
| 242 | |
| 243 | _, _, builder = _build_prefill_metadata([512], torch.device("cuda"), num_heads) |
| 244 | assert builder.fp8_ps_reduce_partial_map.numel() == _tiles(width) |