From 1fb0209bbc476df3ddbe1dbbaf22fa57248cc8b1 Mon Sep 17 00:00:00 2001 From: Kevin McKay Date: Tue, 6 Jan 2026 00:10:59 -0600 Subject: [PATCH] [Bugfix][Hardware][AMD] Fix exception types in AITER MLA FP8 check (#31177) Signed-off-by: c0de128 Co-authored-by: Claude Opus 4.5 --- .../rocm/aiter/test_mla_fp8_support_check.py | 118 ++++++++++++++++++ vllm/_aiter_ops.py | 12 +- 2 files changed, 129 insertions(+), 1 deletion(-) create mode 100644 tests/rocm/aiter/test_mla_fp8_support_check.py diff --git a/tests/rocm/aiter/test_mla_fp8_support_check.py b/tests/rocm/aiter/test_mla_fp8_support_check.py new file mode 100644 index 0000000000..e3dc0f8ea1 --- /dev/null +++ b/tests/rocm/aiter/test_mla_fp8_support_check.py @@ -0,0 +1,118 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Unit tests for AITER MLA FP8 support detection. + +These tests verify that the _check_aiter_mla_fp8_support() function +correctly handles various error conditions without crashing. +""" + +from unittest.mock import patch + +import pytest + + +class TestAiterMlaFp8SupportCheck: + """Test cases for _check_aiter_mla_fp8_support() function.""" + + def setup_method(self): + """Reset the global cache before each test.""" + import vllm._aiter_ops as aiter_ops + + aiter_ops._AITER_MLA_SUPPORTS_FP8 = None + + @patch("vllm._aiter_ops.is_aiter_found_and_supported", return_value=True) + def test_import_error_handling(self, mock_supported): + """Test that ImportError is handled gracefully.""" + import vllm._aiter_ops as aiter_ops + from vllm._aiter_ops import _check_aiter_mla_fp8_support + + aiter_ops._AITER_MLA_SUPPORTS_FP8 = None + + # Should return False without raising + with patch( + "vllm._aiter_ops.inspect.signature", + side_effect=ImportError("No module"), + ): + result = _check_aiter_mla_fp8_support() + assert result is False + + @patch("vllm._aiter_ops.is_aiter_found_and_supported", return_value=True) + def test_module_not_found_error_handling(self, mock_supported): + """Test that ModuleNotFoundError is handled gracefully.""" + import vllm._aiter_ops as aiter_ops + from vllm._aiter_ops import _check_aiter_mla_fp8_support + + aiter_ops._AITER_MLA_SUPPORTS_FP8 = None + + with patch( + "vllm._aiter_ops.inspect.signature", + side_effect=ModuleNotFoundError("Module not found"), + ): + # Should return False without raising + assert _check_aiter_mla_fp8_support() is False + # Cache should be set to False + assert aiter_ops._AITER_MLA_SUPPORTS_FP8 is False + + @patch("vllm._aiter_ops.is_aiter_found_and_supported", return_value=True) + def test_attribute_error_handling(self, mock_supported): + """Test that AttributeError is handled gracefully.""" + import vllm._aiter_ops as aiter_ops + from vllm._aiter_ops import _check_aiter_mla_fp8_support + + aiter_ops._AITER_MLA_SUPPORTS_FP8 = None + + with patch( + "vllm._aiter_ops.inspect.signature", + side_effect=AttributeError("No attribute"), + ): + assert _check_aiter_mla_fp8_support() is False + assert aiter_ops._AITER_MLA_SUPPORTS_FP8 is False + + @patch("vllm._aiter_ops.is_aiter_found_and_supported", return_value=True) + def test_value_error_handling(self, mock_supported): + """Test that ValueError is handled gracefully (no signature).""" + import vllm._aiter_ops as aiter_ops + from vllm._aiter_ops import _check_aiter_mla_fp8_support + + aiter_ops._AITER_MLA_SUPPORTS_FP8 = None + + with patch( + "vllm._aiter_ops.inspect.signature", + side_effect=ValueError("No signature"), + ): + assert _check_aiter_mla_fp8_support() is False + assert aiter_ops._AITER_MLA_SUPPORTS_FP8 is False + + @patch("vllm._aiter_ops.is_aiter_found_and_supported", return_value=True) + def test_type_error_handling(self, mock_supported): + """Test that TypeError is handled gracefully (not callable).""" + import vllm._aiter_ops as aiter_ops + from vllm._aiter_ops import _check_aiter_mla_fp8_support + + aiter_ops._AITER_MLA_SUPPORTS_FP8 = None + + with patch( + "vllm._aiter_ops.inspect.signature", + side_effect=TypeError("Not a callable"), + ): + assert _check_aiter_mla_fp8_support() is False + assert aiter_ops._AITER_MLA_SUPPORTS_FP8 is False + + @patch("vllm._aiter_ops.is_aiter_found_and_supported", return_value=True) + def test_result_caching(self, mock_supported): + """Test that the result is cached after first check.""" + import vllm._aiter_ops as aiter_ops + + # Set cache to True + aiter_ops._AITER_MLA_SUPPORTS_FP8 = True + + from vllm._aiter_ops import _check_aiter_mla_fp8_support + + # Should return cached value without re-checking + result = _check_aiter_mla_fp8_support() + assert result is True + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index 3c156333bd..b443f77352 100644 --- a/vllm/_aiter_ops.py +++ b/vllm/_aiter_ops.py @@ -288,7 +288,17 @@ def _check_aiter_mla_fp8_support() -> bool: _AITER_MLA_SUPPORTS_FP8 = ( "q_scale" in sig.parameters and "kv_scale" in sig.parameters ) - except Exception: + except ( + ImportError, + ModuleNotFoundError, + AttributeError, + ValueError, + TypeError, + ): + # ImportError/ModuleNotFoundError: aiter.mla module not available + # AttributeError: mla_decode_fwd doesn't exist + # ValueError: mla_decode_fwd has no signature (e.g., built-in) + # TypeError: mla_decode_fwd is not a callable _AITER_MLA_SUPPORTS_FP8 = False return _AITER_MLA_SUPPORTS_FP8