[ROCm] Enable Sequence Parallelism for AMD GPUs (MI300X/MI325X/MI355X)

docs/design/fusions.md

@@ -45,7 +45,7 @@ The table below lists the quantization schemes supported by each fusion on each
 | `fuse_attn_quant` (MLA)\*    | FP8 static\*, FP8 per-group\*, NVFP4\*   | FP8 static\*, FP8 per-group\*            | FP8 static\*, FP8 per-group\*            | —             | FP8 static\* (untested)                  |
 | `fuse_rope_kvcache`          | —                                        | —                                        | —                                        | —             | FP16/BF16                                |
 | `enable_qk_norm_rope_fusion` | FP16/BF16                                | FP16/BF16                                | FP16/BF16†                               | FP16/BF16†    | —                                        |
-| `enable_sp`                  | FP16/BF16, FP8 static†                   | FP16/BF16, FP8 static                    | FP16/BF16†                               | FP16/BF16†    | —                                        |
+| `enable_sp`                  | FP16/BF16, FP8 static†                   | FP16/BF16, FP8 static                    | FP16/BF16†                               | FP16/BF16†    | FP16/BF16†                               |
 | `fuse_gemm_comms`            | FP16/BF16, FP8 static†                   | FP16/BF16, FP8 static                    | FP16/BF16†                               | FP16/BF16†    | —                                        |
 | `fuse_norm_quant`            | FP8 static, FP8 per-token, FP8 per-group | FP8 static, FP8 per-token, FP8 per-group | FP8 static, FP8 per-token, FP8 per-group | —             | FP8 static, FP8 per-token, FP8 per-group |
 | `fuse_act_quant`             | FP8 static, NVFP4                        | FP8 static, FP8 per-group (128/64)       | FP8 static, FP8 per-group (128/64)       | —             | FP8 per-group                            |
@@ -56,8 +56,8 @@ The table below lists the quantization schemes supported by each fusion on each
 fused quantization output. See the [`fuse_attn_quant` section](#attention--quantization-fuse_attn_quant)
 for per-backend details.
 
-† `enable_sp` and `fuse_gemm_comms` are only autoconfigured for SM90 today;
-other architectures support requires setting `PassConfig.sp_min_token_num` explicitly.
+† `enable_sp` and `fuse_gemm_comms` are only autoconfigured for SM90 and ROCm MI300X/MI325X/MI355X today;
+other architectures require setting `PassConfig.sp_min_token_num` explicitly.
 SM100 support also requires setting `VLLM_DISABLED_KERNELS=FlashInferFP8ScaledMMLinearKernel`.
 
 ## Enabling / Disabling Fusions
@@ -194,9 +194,9 @@ pass can fuse the reduce-scatter / all-gather with the surrounding GEMMs.
 
 Sequence Parallelism itself does not directly improve performance; it is a prerequisite for the
 AsyncTP pass (`fuse_gemm_comms`). SP is only applied above a minimum token threshold that is
-autoconfigured based on device capability and model `hidden_size`. Currently only active on
-H100/SM90 for models with `hidden_size >= 8192`. The threshold is configurable via
-`PassConfig.sp_min_token_num`.
+autoconfigured based on device capability and model `hidden_size`. Currently autoconfigured on
+H100/SM90 (models with `hidden_size >= 8192`) and ROCm MI300X/MI325X/MI355X (models with
+`hidden_size >= 7168`). The threshold is configurable via `PassConfig.sp_min_token_num`.
 
 The general transformation:
 
@@ -215,7 +215,7 @@ Patterns covered:
 Requires: `use_inductor_graph_partition=True` **or** piecewise compilation with static sizes
 divisible by `tensor_parallel_size`.
 
-Supported hardware: Only tested on NVIDIA CUDA, possibly works on ROCm. FP8 all-gather requires sm90+.
+Supported hardware: NVIDIA CUDA (SM80+) and ROCm (MI300X/MI325X/MI355X). FP8 all-gather requires sm90+.
 
 **Code locations.**
 

tests/compile/conftest.py

@@ -25,6 +25,7 @@ def _mock_platform(is_cuda: bool = True, capability: tuple[int, int] | None = No
         mock_platform = MagicMock()
         mock_platform.is_cuda.return_value = is_cuda
         mock_platform.is_xpu.return_value = False
+        mock_platform.is_rocm.return_value = False
         device_capability = (
             DeviceCapability(*capability) if capability is not None else None
         )
@@ -49,6 +50,33 @@ def is_device_capability_family(
     return _mock_platform
 
 
+@pytest.fixture
+def mock_rocm_platform():
+    """
+    Fixture that returns a factory for creating mocked ROCm platforms.
+
+    Usage:
+        def test_something(mock_rocm_platform):
+            with mock_rocm_platform(capability=(9, 4)):
+                # test code
+    """
+
+    @contextmanager
+    def _mock_platform(capability: tuple[int, int] | None = None):
+        mock_platform = MagicMock()
+        mock_platform.is_cuda.return_value = False
+        mock_platform.is_xpu.return_value = False
+        mock_platform.is_rocm.return_value = True
+        device_capability = (
+            DeviceCapability(*capability) if capability is not None else None
+        )
+        mock_platform.get_device_capability.return_value = device_capability
+        with patch("vllm.platforms.current_platform", mock_platform):
+            yield mock_platform
+
+    return _mock_platform
+
+
 @pytest.fixture
 def mock_xpu_platform():
     """
@@ -65,6 +93,7 @@ def _mock_platform():
         mock_platform = MagicMock()
         mock_platform.is_cuda.return_value = False
         mock_platform.is_xpu.return_value = True
+        mock_platform.is_rocm.return_value = False
         with patch("vllm.platforms.current_platform", mock_platform):
             yield mock_platform
 

tests/compile/test_sequence_parallelism_threshold.py

@@ -110,6 +110,115 @@ def test_hidden_size_boundary(self, mock_cuda_platform):
             assert result is not None
 
 
+class TestGetSequenceParallelismThresholdROCm:
+    """Tests for get_sequence_parallelism_threshold on ROCm platforms."""
+
+    def test_mi300x_gfx942_returns_threshold(self, mock_rocm_platform):
+        """MI300X (gfx942, capability 9.4) should return a calculated threshold."""
+        with mock_rocm_platform(capability=(9, 4)):
+            hidden_size = 7168
+            tp_size = 8
+            element_size = 2
+
+            result = get_sequence_parallelism_threshold(
+                hidden_size=hidden_size,
+                tp_size=tp_size,
+                element_size=element_size,
+            )
+
+            MiB = 1024 * 1024
+            expected = int(
+                (SP_MIN_PER_GPU_SIZE_MB[94] * tp_size * MiB)
+                // (hidden_size * element_size)
+            )
+            assert result == expected
+
+    def test_mi355x_gfx950_returns_threshold(self, mock_rocm_platform):
+        """MI355X (gfx950, capability 9.5) should return a calculated threshold."""
+        with mock_rocm_platform(capability=(9, 5)):
+            hidden_size = 7168
+            tp_size = 4
+            element_size = 2
+
+            result = get_sequence_parallelism_threshold(
+                hidden_size=hidden_size,
+                tp_size=tp_size,
+                element_size=element_size,
+            )
+
+            MiB = 1024 * 1024
+            expected = int(
+                (SP_MIN_PER_GPU_SIZE_MB[95] * tp_size * MiB)
+                // (hidden_size * element_size)
+            )
+            assert result == expected
+
+    def test_mi355x_small_hidden_size_returns_none(self, mock_rocm_platform):
+        """MI355X with hidden_size below threshold should return None."""
+        with mock_rocm_platform(capability=(9, 5)):
+            result = get_sequence_parallelism_threshold(
+                hidden_size=4096,
+                tp_size=4,
+                element_size=2,
+            )
+            assert result is None
+
+    def test_unsupported_rocm_capability_returns_none(self, mock_rocm_platform):
+        """ROCm arch without a configured threshold (e.g. gfx90a) returns None."""
+        with mock_rocm_platform(capability=(9, 0)):
+            result = get_sequence_parallelism_threshold(
+                hidden_size=7168,
+                tp_size=4,
+                element_size=2,
+            )
+            assert result is None
+
+    def test_mi355x_hidden_size_boundary(self, mock_rocm_platform):
+        """Test boundary at the exact hidden_size threshold for MI355X."""
+        with mock_rocm_platform(capability=(9, 5)):
+            # Just below threshold
+            result = get_sequence_parallelism_threshold(
+                hidden_size=SP_MIN_HIDDEN_SIZE[95] - 1,
+                tp_size=4,
+                element_size=2,
+            )
+            assert result is None
+
+            # Exactly at threshold
+            result = get_sequence_parallelism_threshold(
+                hidden_size=SP_MIN_HIDDEN_SIZE[95],
+                tp_size=4,
+                element_size=2,
+            )
+            assert result is not None
+
+    @pytest.mark.parametrize(
+        "capability,cap_int",
+        [
+            ((9, 4), 94),
+            ((9, 5), 95),
+        ],
+    )
+    def test_rocm_tp_scaling(self, mock_rocm_platform, capability, cap_int):
+        """Verify threshold scales with TP size on ROCm."""
+        with mock_rocm_platform(capability=capability):
+            # Use a power-of-two hidden size (>= the 7168 threshold) so the
+            # byte budget divides evenly and floor division scales exactly.
+            assert SP_MIN_HIDDEN_SIZE[cap_int] <= 8192
+            hidden_size = 8192
+            element_size = 2
+
+            result_tp2 = get_sequence_parallelism_threshold(
+                hidden_size=hidden_size, tp_size=2, element_size=element_size
+            )
+            result_tp8 = get_sequence_parallelism_threshold(
+                hidden_size=hidden_size, tp_size=8, element_size=element_size
+            )
+            assert result_tp2 is not None
+            assert result_tp8 is not None
+            assert result_tp8 == 4 * result_tp2
+
+
 # XPU-specific constants (must match sequence_parallelism.py values)
 _XPU_MIN_HIDDEN_SIZE = 4096
 _XPU_MIN_PER_GPU_SIZE_MB = 8.0

vllm/compilation/passes/fusion/sequence_parallelism.py

@@ -37,9 +37,12 @@
 
 # Min hidden size per device capability for sequence parallelism
 # Only apply sequence parallelism for models with hidden_size >= threshold
+# ROCm: capability 94 = MI300X (gfx942), 95 = MI325X/MI355X (gfx950)
 SP_MIN_HIDDEN_SIZE: dict[int, int] = {
     90: 8192,  # H100: only for models with hidden_size >= 8192
     100: 8192,  # Blackwell family: only for models with hidden_size >= 8192
+    94: 7168,  # MI300X (gfx942)
+    95: 7168,  # MI325X/MI355X (gfx950)
 }
 
 # Min size per GPU per device capability for sequence parallelism
@@ -49,8 +52,15 @@
     90: 8,  # 8MB per GPU for H100
     # Use a more conservative threshold on Blackwell so TP8 starts later.
     100: 32,
+    94: 8,  # MI300X (gfx942)
+    95: 8,  # MI325X/MI355X (gfx950)
 }
 
+# ROCm device capabilities that support AITER-fused sequence parallelism
+# (MI3xx only: gfx942 -> 94, gfx950 -> 95). MI250 (gfx90a -> 90) is excluded so
+# it does not alias the CUDA H100 entry in the shared threshold dicts above.
+_ROCM_SP_SUPPORTED_CAPABILITIES = (94, 95)
+
 
 def get_sequence_parallelism_threshold(
     hidden_size: int,
@@ -69,6 +79,9 @@ def get_sequence_parallelism_threshold(
 
     Formula: min_token_num = (min_per_gpu_size_mb * tp_size * MiB) //
              (hidden_size * element_size)
+
+    Supported platforms: CUDA (H100/Blackwell), XPU, and ROCm
+    (MI300X/MI325X/MI355X).
     """
     from vllm.platforms import current_platform
 
@@ -92,6 +105,20 @@ def get_sequence_parallelism_threshold(
         if _hidden is None or _gpu_mb is None:
             return None
         min_hidden_size, min_per_gpu_size_mb = _hidden, _gpu_mb
+    elif current_platform.is_rocm():
+        capability = current_platform.get_device_capability()
+        if capability is None:
+            return None
+        device_capability = capability.to_int()
+
+        # AITER-fused sequence parallelism is only supported on MI3xx
+        # (gfx942 -> 94, gfx950 -> 95). Gate explicitly so MI250 (gfx90a ->
+        # capability 90) is not enabled by aliasing the CUDA H100 entry, which
+        # also lives at capability 90 in the shared threshold dicts.
+        if device_capability not in _ROCM_SP_SUPPORTED_CAPABILITIES:
+            return None
+        min_hidden_size = SP_MIN_HIDDEN_SIZE[device_capability]
+        min_per_gpu_size_mb = SP_MIN_PER_GPU_SIZE_MB[device_capability]
     else:
         return None