@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 252 lines 7.79 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Tests for ROCm AITER MLA FP8 support detection."""
4
5 import sys
6 import types
7 from typing import Any
8 from unittest.mock import patch
9
10 import pytest
11
12 from vllm.platforms import current_platform
13
14 _SKIP_UNSUPPORTED_AITER_HARDWARE = True
15 if current_platform.is_rocm():
16 from vllm.platforms.rocm import get_cdna_version
17
18 _SKIP_UNSUPPORTED_AITER_HARDWARE = get_cdna_version() <= 2
19
20 pytestmark = pytest.mark.skipif(
21 not current_platform.is_rocm(), reason="ROCm-specific tests"
22 )
23
24
25 @pytest.fixture(autouse=True)
26 def reset_aiter_mla_support_cache(monkeypatch: pytest.MonkeyPatch) -> None:
27 import vllm._aiter_ops as aiter_ops
28
29 monkeypatch.setattr(aiter_ops, "_AITER_MLA_SUPPORTS_FP8", None)
30
31
32 def _install_fake_aiter_modules(
33 monkeypatch: pytest.MonkeyPatch,
34 *,
35 supports_fp8: bool,
36 supports_causal: bool = False,
37 ) -> None:
38 aiter_mod: Any = types.ModuleType("aiter")
39 mla_mod: Any = types.ModuleType("aiter.mla")
40
41 if supports_fp8:
42
43 def mla_decode_fwd_with_fp8(
44 q,
45 kv_buffer,
46 kv_indptr,
47 kv_indices,
48 o,
49 sm_scale,
50 q_scale=None,
51 kv_scale=None,
52 ):
53 return None
54
55 mla_decode_fwd: Any = mla_decode_fwd_with_fp8
56
57 else:
58
59 def mla_decode_fwd_without_fp8(
60 q,
61 kv_buffer,
62 kv_indptr,
63 kv_indices,
64 o,
65 sm_scale,
66 ):
67 return None
68
69 mla_decode_fwd = mla_decode_fwd_without_fp8
70
71 if supports_causal:
72 inner = mla_decode_fwd
73
74 def mla_decode_fwd_with_causal(*args, causal=True, **kwargs):
75 return inner(*args, **kwargs)
76
77 mla_decode_fwd = mla_decode_fwd_with_causal
78
79 mla_mod.mla_decode_fwd = mla_decode_fwd
80 aiter_mod.mla = mla_mod
81
82 monkeypatch.setitem(sys.modules, "aiter", aiter_mod)
83 monkeypatch.setitem(sys.modules, "aiter.mla", mla_mod)
84
85
86 def test_aiter_mla_fp8_support_detects_fp8_signature(monkeypatch):
87 """The support check should detect q_scale and kv_scale parameters."""
88 from vllm._aiter_ops import _check_aiter_mla_fp8_support
89
90 _install_fake_aiter_modules(monkeypatch, supports_fp8=True)
91
92 assert _check_aiter_mla_fp8_support() is True
93
94
95 def test_aiter_mla_fp8_support_rejects_missing_fp8_signature(monkeypatch):
96 """The support check should return False when FP8 params are absent."""
97 from vllm._aiter_ops import _check_aiter_mla_fp8_support
98
99 _install_fake_aiter_modules(monkeypatch, supports_fp8=False)
100
101 assert _check_aiter_mla_fp8_support() is False
102
103
104 @pytest.mark.skipif(
105 _SKIP_UNSUPPORTED_AITER_HARDWARE,
106 reason="Installed AITER MLA FP8 check requires CDNA 3 or newer",
107 )
108 def test_installed_aiter_mla_supports_fp8():
109 """Supported ROCm CI must provide AITER with MLA FP8 scaling."""
110 from vllm._aiter_ops import (
111 _check_aiter_mla_fp8_support,
112 is_aiter_found_and_supported,
113 )
114
115 assert is_aiter_found_and_supported(), (
116 "AITER must be installed on supported ROCm hardware"
117 )
118 assert _check_aiter_mla_fp8_support() is True
119
120
121 @pytest.mark.parametrize(
122 "error_type",
123 [ImportError, ModuleNotFoundError, AttributeError, ValueError, TypeError],
124 )
125 def test_aiter_mla_fp8_support_handles_signature_errors(monkeypatch, error_type):
126 """The support check should fail closed on import or signature problems."""
127 import vllm._aiter_ops as aiter_ops
128 from vllm._aiter_ops import _check_aiter_mla_fp8_support
129
130 _install_fake_aiter_modules(monkeypatch, supports_fp8=True)
131
132 with patch("inspect.signature", side_effect=error_type("boom")):
133 assert _check_aiter_mla_fp8_support() is False
134 assert aiter_ops._AITER_MLA_SUPPORTS_FP8 is False
135
136
137 def test_aiter_mla_fp8_support_result_is_cached(monkeypatch):
138 """The support check should reuse the cached result on later calls."""
139 import inspect
140
141 from vllm._aiter_ops import _check_aiter_mla_fp8_support
142
143 _install_fake_aiter_modules(monkeypatch, supports_fp8=True)
144
145 with patch("inspect.signature", wraps=inspect.signature) as signature_mock:
146 assert _check_aiter_mla_fp8_support() is True
147 assert _check_aiter_mla_fp8_support() is True
148 assert signature_mock.call_count == 1
149
150
151 def test_missing_causal_arg_fails_closed(monkeypatch):
152 """Old aiter wheels have no causal=; do not report that as a valid draft path."""
153 from vllm._aiter_ops import rocm_aiter_ops
154 from vllm.v1.attention.backends.mla.rocm_aiter_mla import AiterMLABackend
155
156 _install_fake_aiter_modules(monkeypatch, supports_fp8=True, supports_causal=False)
157 with pytest.raises(RuntimeError, match="causal-only"):
158 rocm_aiter_ops.mla_decode_supports_non_causal()
159 with pytest.raises(RuntimeError, match="causal-only"):
160 AiterMLABackend.supports_non_causal()
161
162
163 @pytest.mark.parametrize("has_kernels", [False, True])
164 def test_supports_non_causal_requires_arch_kernels(monkeypatch, has_kernels):
165 """The Python argument alone does not guarantee an arch has the kernel."""
166 from vllm._aiter_ops import rocm_aiter_ops
167 from vllm.v1.attention.backends.mla import rocm_aiter_mla
168 from vllm.v1.attention.backends.mla.rocm_aiter_mla import AiterMLABackend
169
170 _install_fake_aiter_modules(monkeypatch, supports_fp8=True, supports_causal=True)
171 monkeypatch.setattr(
172 rocm_aiter_mla,
173 "_aiter_mla_non_causal_asm_kernels",
174 lambda: has_kernels,
175 )
176 assert rocm_aiter_ops.mla_decode_supports_non_causal() is True
177 assert AiterMLABackend.supports_non_causal() is has_kernels
178
179
180 @pytest.mark.parametrize("causal", [True, False])
181 def test_decode_fwd_forwards_causal_to_aiter(monkeypatch, causal):
182 """The decode op must pass the block mask through; omitting it is causal-only."""
183 seen: list = []
184
185 def mla_decode_fwd(*args, causal=None, **kwargs):
186 seen.append(causal)
187
188 mla_mod: Any = types.ModuleType("aiter.mla")
189 mla_mod.mla_decode_fwd = mla_decode_fwd
190 aiter_mod: Any = types.ModuleType("aiter")
191 aiter_mod.mla = mla_mod
192 monkeypatch.setitem(sys.modules, "aiter", aiter_mod)
193 monkeypatch.setitem(sys.modules, "aiter.mla", mla_mod)
194
195 import torch
196
197 import vllm._aiter_ops as aiter_ops
198
199 aiter_ops._rocm_aiter_mla_decode_fwd_impl(
200 torch.zeros(2, 1, 8),
201 torch.zeros(2, 1, 1, 8),
202 torch.zeros(2, 1, 8),
203 torch.zeros(3, dtype=torch.int32),
204 1,
205 causal=causal,
206 )
207 assert seen == [causal]
208
209
210 def _aiter_mla_validate(
211 monkeypatch,
212 *,
213 dtype,
214 kv_cache_dtype,
215 use_non_causal,
216 ):
217 from vllm.platforms.interface import DeviceCapability
218 from vllm.v1.attention.backend import AttentionType
219 from vllm.v1.attention.backends.mla import rocm_aiter_mla
220 from vllm.v1.attention.backends.mla.rocm_aiter_mla import AiterMLABackend
221
222 _install_fake_aiter_modules(monkeypatch, supports_fp8=True, supports_causal=True)
223 monkeypatch.setattr(
224 rocm_aiter_mla, "_aiter_mla_non_causal_asm_kernels", lambda: True
225 )
226 return AiterMLABackend.validate_configuration(
227 head_size=576,
228 dtype=dtype,
229 kv_cache_dtype=kv_cache_dtype,
230 block_size=1,
231 use_mla=True,
232 has_sink=False,
233 use_sparse=False,
234 use_mm_prefix=False,
235 use_per_head_quant_scales=False,
236 device_capability=DeviceCapability(9, 5),
237 attn_type=AttentionType.DECODER,
238 use_non_causal=use_non_causal,
239 )
240
241
242 def test_non_causal_fp16_auto_cache_is_rejected(monkeypatch):
243 """Pinned AITER aborts fp16 Q; auto-select must fall through to Triton."""
244 import torch
245
246 reasons = _aiter_mla_validate(
247 monkeypatch,
248 dtype=torch.float16,
249 kv_cache_dtype="auto",
250 use_non_causal=True,
251 )
252 assert any("non-causal fp16" in reason for reason in reasons)