A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 10, 16:30 UTC
[Bugfix][Hardware][AMD] Fix exception types in AITER MLA FP8 check (#31177)
Signed-off-by: c0de128 <kevin.mckay@outlook.com> Co-authored-by: Claude Opus 4.5 <noreply@anthropic.com>
Kevin McKay committed
Jan 6, 2026 at 00:10 UTC
1fb0209bbc476df3ddbe1dbbaf22fa57248cc8b1
2 files changed
+129
-1
tests/rocm/aiter/test_mla_fp8_support_check.py
new
+118
@@ -0,0 +1,118 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""
4
+Unit tests for AITER MLA FP8 support detection.
5
+
6
+These tests verify that the _check_aiter_mla_fp8_support() function
7
+correctly handles various error conditions without crashing.
8
+"""
9
+
10
+from unittest.mock import patch
11
+
12
+import pytest
13
+
14
+
15
+class TestAiterMlaFp8SupportCheck:
16
+ """Test cases for _check_aiter_mla_fp8_support() function."""
17
+
18
+ def setup_method(self):
19
+ """Reset the global cache before each test."""
20
+ import vllm._aiter_ops as aiter_ops
21
+
22
+ aiter_ops._AITER_MLA_SUPPORTS_FP8 = None
23
+
24
+ @patch("vllm._aiter_ops.is_aiter_found_and_supported", return_value=True)
25
+ def test_import_error_handling(self, mock_supported):
26
+ """Test that ImportError is handled gracefully."""
27
+ import vllm._aiter_ops as aiter_ops
28
+ from vllm._aiter_ops import _check_aiter_mla_fp8_support
29
+
30
+ aiter_ops._AITER_MLA_SUPPORTS_FP8 = None
31
+
32
+ # Should return False without raising
33
+ with patch(
34
+ "vllm._aiter_ops.inspect.signature",
35
+ side_effect=ImportError("No module"),
36
+ ):
37
+ result = _check_aiter_mla_fp8_support()
38
+ assert result is False
39
+
40
+ @patch("vllm._aiter_ops.is_aiter_found_and_supported", return_value=True)
41
+ def test_module_not_found_error_handling(self, mock_supported):
42
+ """Test that ModuleNotFoundError is handled gracefully."""
43
+ import vllm._aiter_ops as aiter_ops
44
+ from vllm._aiter_ops import _check_aiter_mla_fp8_support
45
+
46
+ aiter_ops._AITER_MLA_SUPPORTS_FP8 = None
47
+
48
+ with patch(
49
+ "vllm._aiter_ops.inspect.signature",
50
+ side_effect=ModuleNotFoundError("Module not found"),
51
+ ):
52
+ # Should return False without raising
53
+ assert _check_aiter_mla_fp8_support() is False
54
+ # Cache should be set to False
55
+ assert aiter_ops._AITER_MLA_SUPPORTS_FP8 is False
56
+
57
+ @patch("vllm._aiter_ops.is_aiter_found_and_supported", return_value=True)
58
+ def test_attribute_error_handling(self, mock_supported):
59
+ """Test that AttributeError is handled gracefully."""
60
+ import vllm._aiter_ops as aiter_ops
61
+ from vllm._aiter_ops import _check_aiter_mla_fp8_support
62
+
63
+ aiter_ops._AITER_MLA_SUPPORTS_FP8 = None
64
+
65
+ with patch(
66
+ "vllm._aiter_ops.inspect.signature",
67
+ side_effect=AttributeError("No attribute"),
68
+ ):
69
+ assert _check_aiter_mla_fp8_support() is False
70
+ assert aiter_ops._AITER_MLA_SUPPORTS_FP8 is False
71
+
72
+ @patch("vllm._aiter_ops.is_aiter_found_and_supported", return_value=True)
73
+ def test_value_error_handling(self, mock_supported):
74
+ """Test that ValueError is handled gracefully (no signature)."""
75
+ import vllm._aiter_ops as aiter_ops
76
+ from vllm._aiter_ops import _check_aiter_mla_fp8_support
77
+
78
+ aiter_ops._AITER_MLA_SUPPORTS_FP8 = None
79
+
80
+ with patch(
81
+ "vllm._aiter_ops.inspect.signature",
82
+ side_effect=ValueError("No signature"),
83
+ ):
84
+ assert _check_aiter_mla_fp8_support() is False
85
+ assert aiter_ops._AITER_MLA_SUPPORTS_FP8 is False
86
+
87
+ @patch("vllm._aiter_ops.is_aiter_found_and_supported", return_value=True)
88
+ def test_type_error_handling(self, mock_supported):
89
+ """Test that TypeError is handled gracefully (not callable)."""
90
+ import vllm._aiter_ops as aiter_ops
91
+ from vllm._aiter_ops import _check_aiter_mla_fp8_support
92
+
93
+ aiter_ops._AITER_MLA_SUPPORTS_FP8 = None
94
+
95
+ with patch(
96
+ "vllm._aiter_ops.inspect.signature",
97
+ side_effect=TypeError("Not a callable"),
98
+ ):
99
+ assert _check_aiter_mla_fp8_support() is False
100
+ assert aiter_ops._AITER_MLA_SUPPORTS_FP8 is False
101
+
102
+ @patch("vllm._aiter_ops.is_aiter_found_and_supported", return_value=True)
103
+ def test_result_caching(self, mock_supported):
104
+ """Test that the result is cached after first check."""
105
+ import vllm._aiter_ops as aiter_ops
106
+
107
+ # Set cache to True
108
+ aiter_ops._AITER_MLA_SUPPORTS_FP8 = True
109
+
110
+ from vllm._aiter_ops import _check_aiter_mla_fp8_support
111
+
112
+ # Should return cached value without re-checking
113
+ result = _check_aiter_mla_fp8_support()
114
+ assert result is True
115
+
116
+
117
+if __name__ == "__main__":
118
+ pytest.main([__file__, "-v"])
vllm/_aiter_ops.py
+11
-1
@@ -288,7 +288,17 @@ def _check_aiter_mla_fp8_support() -> bool:
288
_AITER_MLA_SUPPORTS_FP8 = (
289
"q_scale" in sig.parameters and "kv_scale" in sig.parameters
290
)
291
- except Exception:
291
+ except (
292
+ ImportError,
293
+ ModuleNotFoundError,
294
+ AttributeError,
295
+ ValueError,
296
+ TypeError,
297
+ ):
298
+ # ImportError/ModuleNotFoundError: aiter.mla module not available
299
+ # AttributeError: mla_decode_fwd doesn't exist
300
+ # ValueError: mla_decode_fwd has no signature (e.g., built-in)
301
+ # TypeError: mla_decode_fwd is not a callable
302
_AITER_MLA_SUPPORTS_FP8 = False
303
return _AITER_MLA_SUPPORTS_FP8
304