From cbd458de86edf502d198d98b730c3129853fe789 Mon Sep 17 00:00:00 2001 From: Chuan Li Date: Wed, 15 Jul 2026 10:16:44 -0700 Subject: [PATCH 1/2] [ROCm] Enable Sequence Parallelism for AMD GPUs (MI300X/MI325X/MI355X) Enable the SequenceParallelismPass on ROCm by adding an is_rocm() branch to get_sequence_parallelism_threshold, alongside the existing CUDA and XPU branches. Threshold configs are added for gfx942 (MI300X, capability 94) and gfx950 (MI325X/MI355X, capability 95), autoconfigured for models with hidden_size >= 7168. - vllm/compilation/passes/fusion/sequence_parallelism.py: ROCm branch + SP_MIN_HIDDEN_SIZE / SP_MIN_PER_GPU_SIZE_MB entries for 94/95 - tests/compile/conftest.py: add mock_rocm_platform fixture; pin is_rocm on the CUDA/XPU mocks - tests/compile/test_sequence_parallelism_threshold.py: ROCm threshold tests - docs/design/fusions.md: mark enable_sp supported on ROCm MI3xx Signed-off-by: Chuan Li Made-with: Cursor --- docs/design/fusions.md | 14 +-- tests/compile/conftest.py | 29 +++++ .../test_sequence_parallelism_threshold.py | 109 ++++++++++++++++++ .../passes/fusion/sequence_parallelism.py | 20 ++++ 4 files changed, 165 insertions(+), 7 deletions(-) diff --git a/docs/design/fusions.md b/docs/design/fusions.md index c9991f75cdb8..98ba9bde98f2 100644 --- a/docs/design/fusions.md +++ b/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.** diff --git a/tests/compile/conftest.py b/tests/compile/conftest.py index 7d15b5c47e55..330d65c28325 100644 --- a/tests/compile/conftest.py +++ b/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 diff --git a/tests/compile/test_sequence_parallelism_threshold.py b/tests/compile/test_sequence_parallelism_threshold.py index 090b77b330ab..a402fb26f110 100644 --- a/tests/compile/test_sequence_parallelism_threshold.py +++ b/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 diff --git a/vllm/compilation/passes/fusion/sequence_parallelism.py b/vllm/compilation/passes/fusion/sequence_parallelism.py index c4caaaedec22..6a955b928599 100644 --- a/vllm/compilation/passes/fusion/sequence_parallelism.py +++ b/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,6 +52,8 @@ 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) } @@ -69,6 +74,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 +100,18 @@ 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() + + # Check if device has configured thresholds (gfx942/gfx950) + _hidden = SP_MIN_HIDDEN_SIZE.get(device_capability) + _gpu_mb = SP_MIN_PER_GPU_SIZE_MB.get(device_capability) + if _hidden is None or _gpu_mb is None: + return None + min_hidden_size, min_per_gpu_size_mb = _hidden, _gpu_mb else: return None From 1d6dd9c0b5cf6e5c313b6c76c87491267bdc1cba Mon Sep 17 00:00:00 2001 From: Chuan Li Date: Wed, 15 Jul 2026 21:34:03 -0700 Subject: [PATCH 2/2] [ROCm] Gate sequence parallelism to MI3xx capabilities Explicitly restrict the ROCm sequence-parallelism threshold path to MI3xx capabilities (gfx942 -> 94, gfx950 -> 95) instead of looking up the shared threshold dicts directly. MI250 (gfx90a) reports capability 90, which aliases the CUDA H100 entry (also 90) in the shared dicts and would have enabled sequence parallelism unintentionally. AITER is not supported on MI250, so this addresses @AndreasKaratzas's MI3xx-gating review comment. Signed-off-by: Chuan Li Co-authored-by: Cursor --- .../passes/fusion/sequence_parallelism.py | 17 ++++++++++++----- 1 file changed, 12 insertions(+), 5 deletions(-) diff --git a/vllm/compilation/passes/fusion/sequence_parallelism.py b/vllm/compilation/passes/fusion/sequence_parallelism.py index 6a955b928599..a5493f5c3c07 100644 --- a/vllm/compilation/passes/fusion/sequence_parallelism.py +++ b/vllm/compilation/passes/fusion/sequence_parallelism.py @@ -56,6 +56,11 @@ 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, @@ -106,12 +111,14 @@ def get_sequence_parallelism_threshold( return None device_capability = capability.to_int() - # Check if device has configured thresholds (gfx942/gfx950) - _hidden = SP_MIN_HIDDEN_SIZE.get(device_capability) - _gpu_mb = SP_MIN_PER_GPU_SIZE_MB.get(device_capability) - if _hidden is None or _gpu_mb is None: + # 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, min_per_gpu_size_mb = _hidden, _gpu_mb + 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