diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index 41dbd629bc..25eb3aa063 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -1544,11 +1544,13 @@ steps: - vllm/_aiter_ops.py - tests/kernels/attention/test_rocm_aiter_mla_decode.py - tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py + - tests/kernels/attention/test_rocm_aiter_mla_fp8_support.py - tests/kernels/attention/test_rocm_aiter_mla_op_registration.py commands: - rocm-smi - pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py - pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_decode.py + - pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_fp8_support.py - pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_op_registration.py - label: Kernels Attention Test %N # TBD @@ -1664,6 +1666,7 @@ steps: - tests/kernels/quant_utils.py - tests/kernels/utils.py - vllm/_aiter_ops.py + - vllm/kernels/aiter_ops.py - vllm/_custom_ops.py - vllm/envs.py - vllm/platforms/rocm.py @@ -2452,24 +2455,6 @@ steps: - uv pip install --system conch-triton-kernels - VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py -#----------------------------------------------------------- mi300 · rocm ------------------------------------------------------------# - -- label: ROCm AITER Ops Test # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] - dind: false - agent_pool: mi300_1 - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/_aiter_ops.py - - vllm/envs.py - - vllm/platforms/rocm.py - - tests/rocm/aiter/ - - vllm/v1/attention/backends/mla/rocm_aiter_mla.py - - vllm/v1/attention/selector.py - commands: - - pytest -v -s rocm/aiter/ - #--------------------------------------------------------- mi300 · samplers ----------------------------------------------------------# - label: Samplers Test # TBD @@ -3461,11 +3446,13 @@ steps: - vllm/_aiter_ops.py - tests/kernels/attention/test_rocm_aiter_mla_decode.py - tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py + - tests/kernels/attention/test_rocm_aiter_mla_fp8_support.py - tests/kernels/attention/test_rocm_aiter_mla_op_registration.py commands: - rocm-smi - pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py - pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_decode.py + - pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_fp8_support.py - pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_op_registration.py - label: Kernels Attention Test %N # TBD @@ -3525,6 +3512,7 @@ steps: - tests/kernels/quant_utils.py - tests/kernels/utils.py - vllm/_aiter_ops.py + - vllm/kernels/aiter_ops.py - vllm/_custom_ops.py - vllm/envs.py - vllm/platforms/rocm.py diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index 887ed03e37..4dec86e929 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -183,6 +183,7 @@ steps: - tests/kernels/quantization - tests/kernels/quantization/test_rocm_skinny_gemms.py - vllm/_aiter_ops.py + - vllm/kernels/aiter_ops.py - vllm/platforms/rocm.py - vllm/model_executor/kernels/ depends_on: diff --git a/tests/kernels/attention/test_rocm_aiter_mla_fp8_support.py b/tests/kernels/attention/test_rocm_aiter_mla_fp8_support.py new file mode 100644 index 0000000000..d67df8ee2b --- /dev/null +++ b/tests/kernels/attention/test_rocm_aiter_mla_fp8_support.py @@ -0,0 +1,137 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for ROCm AITER MLA FP8 support detection.""" + +import sys +import types +from typing import Any +from unittest.mock import patch + +import pytest + +from vllm.platforms import current_platform + +_SKIP_UNSUPPORTED_AITER_HARDWARE = True +if current_platform.is_rocm(): + from vllm.platforms.rocm import get_cdna_version + + _SKIP_UNSUPPORTED_AITER_HARDWARE = get_cdna_version() <= 2 + +pytestmark = pytest.mark.skipif( + not current_platform.is_rocm(), reason="ROCm-specific tests" +) + + +@pytest.fixture(autouse=True) +def reset_aiter_mla_support_cache(monkeypatch: pytest.MonkeyPatch) -> None: + import vllm._aiter_ops as aiter_ops + + monkeypatch.setattr(aiter_ops, "_AITER_MLA_SUPPORTS_FP8", None) + + +def _install_fake_aiter_modules( + monkeypatch: pytest.MonkeyPatch, *, supports_fp8: bool +) -> None: + aiter_mod: Any = types.ModuleType("aiter") + mla_mod: Any = types.ModuleType("aiter.mla") + + if supports_fp8: + + def mla_decode_fwd_with_fp8( + q, + kv_buffer, + kv_indptr, + kv_indices, + o, + sm_scale, + q_scale=None, + kv_scale=None, + ): + return None + + mla_decode_fwd: Any = mla_decode_fwd_with_fp8 + + else: + + def mla_decode_fwd_without_fp8( + q, + kv_buffer, + kv_indptr, + kv_indices, + o, + sm_scale, + ): + return None + + mla_decode_fwd = mla_decode_fwd_without_fp8 + + mla_mod.mla_decode_fwd = mla_decode_fwd + aiter_mod.mla = mla_mod + + monkeypatch.setitem(sys.modules, "aiter", aiter_mod) + monkeypatch.setitem(sys.modules, "aiter.mla", mla_mod) + + +def test_aiter_mla_fp8_support_detects_fp8_signature(monkeypatch): + """The support check should detect q_scale and kv_scale parameters.""" + from vllm._aiter_ops import _check_aiter_mla_fp8_support + + _install_fake_aiter_modules(monkeypatch, supports_fp8=True) + + assert _check_aiter_mla_fp8_support() is True + + +def test_aiter_mla_fp8_support_rejects_missing_fp8_signature(monkeypatch): + """The support check should return False when FP8 params are absent.""" + from vllm._aiter_ops import _check_aiter_mla_fp8_support + + _install_fake_aiter_modules(monkeypatch, supports_fp8=False) + + assert _check_aiter_mla_fp8_support() is False + + +@pytest.mark.skipif( + _SKIP_UNSUPPORTED_AITER_HARDWARE, + reason="Installed AITER MLA FP8 check requires CDNA 3 or newer", +) +def test_installed_aiter_mla_supports_fp8(): + """Supported ROCm CI must provide AITER with MLA FP8 scaling.""" + from vllm._aiter_ops import ( + _check_aiter_mla_fp8_support, + is_aiter_found_and_supported, + ) + + assert is_aiter_found_and_supported(), ( + "AITER must be installed on supported ROCm hardware" + ) + assert _check_aiter_mla_fp8_support() is True + + +@pytest.mark.parametrize( + "error_type", + [ImportError, ModuleNotFoundError, AttributeError, ValueError, TypeError], +) +def test_aiter_mla_fp8_support_handles_signature_errors(monkeypatch, error_type): + """The support check should fail closed on import or signature problems.""" + import vllm._aiter_ops as aiter_ops + from vllm._aiter_ops import _check_aiter_mla_fp8_support + + _install_fake_aiter_modules(monkeypatch, supports_fp8=True) + + with patch("inspect.signature", side_effect=error_type("boom")): + assert _check_aiter_mla_fp8_support() is False + assert aiter_ops._AITER_MLA_SUPPORTS_FP8 is False + + +def test_aiter_mla_fp8_support_result_is_cached(monkeypatch): + """The support check should reuse the cached result on later calls.""" + import inspect + + from vllm._aiter_ops import _check_aiter_mla_fp8_support + + _install_fake_aiter_modules(monkeypatch, supports_fp8=True) + + with patch("inspect.signature", wraps=inspect.signature) as signature_mock: + assert _check_aiter_mla_fp8_support() is True + assert _check_aiter_mla_fp8_support() is True + assert signature_mock.call_count == 1 diff --git a/tests/rocm/aiter/test_aiter_hipb_mm_linear_kernel.py b/tests/kernels/quantization/test_aiter_hipb_mm_linear_kernel.py similarity index 100% rename from tests/rocm/aiter/test_aiter_hipb_mm_linear_kernel.py rename to tests/kernels/quantization/test_aiter_hipb_mm_linear_kernel.py diff --git a/tests/rocm/aiter/test_quant_op_schema.py b/tests/kernels/quantization/test_quant_op_schema.py similarity index 100% rename from tests/rocm/aiter/test_quant_op_schema.py rename to tests/kernels/quantization/test_quant_op_schema.py diff --git a/tests/rocm/aiter/test_grouped_quant.py b/tests/kernels/quantization/test_rocm_aiter_grouped_quant.py similarity index 100% rename from tests/rocm/aiter/test_grouped_quant.py rename to tests/kernels/quantization/test_rocm_aiter_grouped_quant.py diff --git a/tests/rocm/aiter/test_mla_fp8_support_check.py b/tests/rocm/aiter/test_mla_fp8_support_check.py deleted file mode 100644 index 28da59a1ae..0000000000 --- a/tests/rocm/aiter/test_mla_fp8_support_check.py +++ /dev/null @@ -1,118 +0,0 @@ -# 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( - "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( - "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( - "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( - "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( - "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"])