Repository navigation
[ROCm] Enable MXFP4 LoRA support on MI355X and add CDNA4 device IDs - #37268
ChuanLi1101 wants to merge 1 commit into
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request enables MXFP4 LoRA support on ROCm platforms by leveraging the Triton backend, which is a valuable enhancement. The changes are well-structured, with clear logic in get_mxfp4_backend_with_lora to handle the ROCm platform specifically. The addition of new CDNA4 device IDs for MI350X and MI355X is also correctly implemented. The accompanying unit tests are thorough, covering the new backend selection logic, device ID mappings, and capability parsing, which ensures the changes are robust and do not introduce regressions. Overall, this is a high-quality contribution.
|
Hi @ChuanLi1101, the pre-commit checks have failed. Please run: uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-filesThen, commit the changes and push to your branch. For future commits, Tip Is
|
|
This pull request has merge conflicts that must be resolved before it can be |
Enable MXFP4 quantized MoE models to use LoRA adapters on AMD ROCm platforms (MI300X/MI325X/MI355X) via the Triton backend. Previously, select_mxfp4_moe_backend() raised NotImplementedError for all non-CUDA platforms when LoRA is enabled, blocking MXFP4+LoRA on ROCm. The Triton backend already supports LoRA through UnfusedOAITritonExperts and works on ROCm when triton_kernels is available. Also adds MI350X (0x75a0) and MI355X (0x75a3) PCI device IDs to the ROCm device name map for proper GPU identification on CDNA4 hardware. Made-with: Cursor Signed-off-by: Li <chuali@amd.com>
ba36963 to
001eda3
Compare
|
Rebased onto latest main and resolved merge conflicts. The main codebase refactored \Mxfp4Backend\ enum and backend selection logic from \mxfp4.py\ into the new \oracle/mxfp4.py\ module (\select_mxfp4_moe_backend\ / \Mxfp4MoeBackend). Updated the ROCm LoRA support to work with the new architecture — the change now adds ROCm handling in the LoRA path of \select_mxfp4_moe_backend(), using \TRITON_UNFUSED\ backend when \ riton_kernels\ is available. Tests have been rewritten to match the new API. @AndreasKaratzas @tjtanaa @gshtras — could you take a look when you get a chance? This enables MXFP4+LoRA on ROCm (MI300X/MI325X/MI355X) via the Triton backend. Thanks! |
|
Hi @ChuanLi1101, the pre-commit checks have failed. Please run: uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-filesThen, commit the changes and push to your branch. For future commits, Tip Is
|
| """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) |
There was a problem hiding this comment.
What is the purpose of this test? I don't see any mxfp4 testing here with lora. There should be accuracy tests here. But also mxfp4 test for ROCm should include a much broader set of tests, which I think is out of the scope of this PR. So probably best if you find the lora or quantization test file that fits best this new feature and add an accuracy test there instead of this new file.
|
This pull request has merge conflicts that must be resolved before it can be |
|
This pull request has been automatically marked as stale because it has not had any activity within 90 days. It will be automatically closed if no further activity occurs within 30 days. Leave a comment if you feel this pull request should remain open. Thank you! |
|
@ChuanLi1101 I think we can close this as Bowen is taking care of reworking the quant paths here: |
github-actions
Bot
added
unstale
|
This pull request has merge conflicts that must be resolved before it can be |
Summary
get_mxfp4_backend_with_lora()to support AMD ROCm platforms (MI300X/MI325X/MI355X) via the Triton backend. Previously, this function returnedMxfp4Backend.NONEfor all non-CUDA platforms, blocking MXFP4 quantized MoE models from using LoRA adapters on ROCm.0x75a0) and MI355X (PCI ID0x75a3) to_ROCM_DEVICE_ID_NAME_MAPfor proper GPU identification on CDNA4 hardware.Motivation
MXFP4 is a key quantization format for AMD's MI355X (CDNA4/gfx950) GPU, enabling higher throughput for large MoE models like DeepSeek-V3 and MiniMax-M2. However, LoRA fine-tuning support was blocked on ROCm because the MXFP4 backend selection did not consider ROCm platforms when LoRA is enabled.
The Triton backend already supports LoRA through
UnfusedOAITritonExpertsand works on ROCm whentriton_kernelsis available. The CK backend is monolithic and does not support the modular kernel path required for LoRA, so Triton is the correct choice for this use case.Changes
vllm/model_executor/layers/quantization/mxfp4.pyget_mxfp4_backend_with_lora()vllm/platforms/rocm.pytests/quantization/test_mxfp4_rocm_backend.pyTest Plan
get_mxfp4_backend_with_lora()with ROCm platform mockingget_mxfp4_backend()with LoRA flag on ROCm