diff --git a/tests/quantization/test_mxfp4_rocm_backend.py b/tests/quantization/test_mxfp4_rocm_backend.py new file mode 100644 index 000000000000..203b3fb82a9e --- /dev/null +++ b/tests/quantization/test_mxfp4_rocm_backend.py @@ -0,0 +1,132 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for MXFP4 backend selection on ROCm platforms.""" + +from unittest.mock import MagicMock, patch + +import pytest + +from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import ( + Mxfp4MoeBackend, + select_mxfp4_moe_backend, +) + + +def _make_lora_config(): + """Create a mock FusedMoEConfig with LoRA enabled.""" + config = MagicMock() + config.is_lora_enabled = True + config.moe_backend = "auto" + config.moe_parallel_config.use_batched_activation_format = False + return config + + +def _make_non_lora_config(): + """Create a mock FusedMoEConfig with LoRA disabled.""" + config = MagicMock() + config.is_lora_enabled = False + config.moe_backend = "auto" + config.moe_parallel_config.use_batched_activation_format = False + return config + + +class TestMxfp4LoraBackendROCm: + """Test MXFP4 backend selection for LoRA on ROCm platforms.""" + + @patch( + "vllm.model_executor.layers.fused_moe.oracle.mxfp4.current_platform" + ) + @patch( + "vllm.model_executor.layers.fused_moe.oracle.mxfp4.has_triton_kernels" + ) + def test_rocm_lora_with_triton_kernels( + self, mock_has_triton, mock_platform + ): + mock_platform.is_rocm.return_value = True + mock_platform.is_cuda.return_value = False + mock_platform.get_device_capability.return_value = (9, 5) + mock_has_triton.return_value = True + + config = _make_lora_config() + backend, experts_cls = select_mxfp4_moe_backend(config) + assert backend == Mxfp4MoeBackend.TRITON_UNFUSED + + @patch( + "vllm.model_executor.layers.fused_moe.oracle.mxfp4.current_platform" + ) + @patch( + "vllm.model_executor.layers.fused_moe.oracle.mxfp4.has_triton_kernels" + ) + def test_rocm_lora_without_triton_raises( + self, mock_has_triton, mock_platform + ): + mock_platform.is_rocm.return_value = True + mock_platform.is_cuda.return_value = False + mock_platform.get_device_capability.return_value = (9, 5) + mock_has_triton.return_value = False + + config = _make_lora_config() + with pytest.raises(NotImplementedError, match="triton_kernels"): + select_mxfp4_moe_backend(config) + + @patch( + "vllm.model_executor.layers.fused_moe.oracle.mxfp4.current_platform" + ) + @patch( + "vllm.model_executor.layers.fused_moe.oracle.mxfp4.has_triton_kernels" + ) + def test_non_cuda_non_rocm_lora_raises( + self, mock_has_triton, mock_platform + ): + mock_platform.is_rocm.return_value = False + mock_platform.is_cuda.return_value = False + mock_platform.get_device_capability.return_value = (0, 0) + mock_has_triton.return_value = False + + config = _make_lora_config() + with pytest.raises( + NotImplementedError, match="only supported on CUDA and ROCm" + ): + select_mxfp4_moe_backend(config) + + +class TestROCmDeviceIdNameMap: + """Test that MI350X/MI355X device IDs are properly mapped.""" + + def test_mi350x_device_id(self): + from vllm.platforms.rocm import _ROCM_DEVICE_ID_NAME_MAP + + assert "0x75a0" in _ROCM_DEVICE_ID_NAME_MAP + assert _ROCM_DEVICE_ID_NAME_MAP["0x75a0"] == "AMD_Instinct_MI350X" + + def test_mi355x_device_id(self): + from vllm.platforms.rocm import _ROCM_DEVICE_ID_NAME_MAP + + assert "0x75a3" in _ROCM_DEVICE_ID_NAME_MAP + assert _ROCM_DEVICE_ID_NAME_MAP["0x75a3"] == "AMD_Instinct_MI355X" + + def test_existing_mi300x_unchanged(self): + from vllm.platforms.rocm import _ROCM_DEVICE_ID_NAME_MAP + + assert _ROCM_DEVICE_ID_NAME_MAP["0x74a1"] == "AMD_Instinct_MI300X" + + def test_existing_mi325x_unchanged(self): + from vllm.platforms.rocm import _ROCM_DEVICE_ID_NAME_MAP + + assert _ROCM_DEVICE_ID_NAME_MAP["0x74a5"] == "AMD_Instinct_MI325X" + + +class TestGfx950CapabilityParsing: + """Test that gfx950 arch string is correctly parsed.""" + + def test_gfx950_capability(self): + from vllm.platforms.rocm import _capability_from_gcn_arch + + result = _capability_from_gcn_arch("gfx950") + assert result == (9, 5) + + def test_gfx942_capability(self): + from vllm.platforms.rocm import _capability_from_gcn_arch + + result = _capability_from_gcn_arch("gfx942") + assert result == (9, 4) diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py index 9008bdeeca7e..633e4cbf6922 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py @@ -208,12 +208,19 @@ def select_mxfp4_moe_backend( # LoRA: separate experts backend path if config.is_lora_enabled: + if current_platform.is_rocm(): + if has_triton_kernels(): + logger.info_once("Using Triton backend for mxfp4 lora on ROCm") + return Mxfp4MoeBackend.TRITON_UNFUSED, backend_to_kernel_cls( + Mxfp4MoeBackend.TRITON_UNFUSED + )[0] + raise NotImplementedError( + "Mxfp4 LoRA on ROCm requires triton_kernels." + ) if not current_platform.is_cuda(): - # ROCm: Triton mxfp4 LoRA hits GPU memory faults due to - # triton_kernels.tensor.Tensor / HIP read-only page issues - # during weight swizzle and LoRA forward. Needs work from - # the triton_kernels/aiter side. - raise NotImplementedError("Mxfp4 LoRA is currently only supported on CUDA.") + raise NotImplementedError( + "Mxfp4 LoRA is currently only supported on CUDA and ROCm." + ) if envs.VLLM_MXFP4_USE_MARLIN is False and triton_kernels_supported: logger.info_once("Using Triton backend for mxfp4 lora") return Mxfp4MoeBackend.TRITON_UNFUSED, backend_to_kernel_cls( diff --git a/vllm/platforms/rocm.py b/vllm/platforms/rocm.py index 26b081b47dfc..bde1317ddf4f 100644 --- a/vllm/platforms/rocm.py +++ b/vllm/platforms/rocm.py @@ -55,6 +55,7 @@ # Architecture -> Reason. _ROCM_PARTIALLY_SUPPORTED_MODELS: dict[str, str] = {} _ROCM_DEVICE_ID_NAME_MAP: dict[str, str] = { + # CDNA3 (gfx942) "0x74a0": "AMD_Instinct_MI300A", "0x74a1": "AMD_Instinct_MI300X", "0x74b5": "AMD_Instinct_MI300X", # MI300X VF @@ -63,6 +64,10 @@ "0x74b9": "AMD_Instinct_MI325X", # MI325X VF "0x74a9": "AMD_Instinct_MI300X_HF", "0x74bd": "AMD_Instinct_MI300X_HF", + # CDNA4 (gfx950) + "0x75a0": "AMD_Instinct_MI350X", + "0x75a3": "AMD_Instinct_MI355X", + # RDNA3 "0x744c": "AMD_Radeon_RX7900XTX", }