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..a5493f5c3c07 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,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