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

[ChuanLi1101]

Summary

This PR enables the Sequence Parallelism compilation pass on ROCm platforms (AMD MI300X, MI325X, MI355X), which was previously restricted to CUDA only.

Changes

Core: vllm/compilation/passes/fusion/sequence_parallelism.py

• Relaxed the platform check in get_sequence_parallelism_threshold() from is_cuda() to is_cuda_alike(), enabling ROCm support
• Added SP threshold configurations for ROCm device capabilities:
  • 94 (gfx942 / MI300X): min_hidden_size=7168, 8MB/GPU
  • 95 (gfx950 / MI325X/MI355X): min_hidden_size=7168, 8MB/GPU

Docs: docs/design/fusions.md

• Updated support matrix to show ROCm FP16/BF16 support for enable_sp
• Updated SP section description to reflect ROCm autoconfiguration

Tests:

• tests/compile/conftest.py: Extended mock_cuda_platform fixture to support is_cuda_alike parameter for ROCm platform mocking
• tests/compile/test_sequence_parallelism_threshold.py: Added TestGetSequenceParallelismThresholdROCm class with tests for MI300X (gfx942) and MI355X (gfx950) threshold calculation, boundary conditions, and TP scaling

Motivation

The Sequence Parallelism pass converts AllReduce → RMSNorm patterns into ReduceScatter → RMSNorm → AllGather, reducing memory pressure on the residual tensor and laying groundwork for GEMM+comms fusion (AsyncTP). This optimization was previously CUDA-only. ROCm users with multi-GPU setups (DeepSeek-V3, Kimi-K2, etc. on MI300X/MI355X) can now benefit from this pass.

Note: AITER-specific RMSNorm patterns were intentionally excluded from this PR since they will be superseded by upcoming vLLM IR changes. The existing C++ RMSNorm patterns work on ROCm when rms_norm custom op is enabled. Once vLLM IR merges, the patterns will be universal.

Related: Central fusion work tracker #36066 (Sequence Parallelism was marked ❓ for ROCm)

Test plan

• ruff check and ruff format pass on all changed files
• Unit tests: pytest tests/compile/test_sequence_parallelism_threshold.py (includes new ROCm tests)
• ROCm CI: Verify on MI325 2-GPU pool (test-amd.yaml already runs test_sequence_parallelism.py)

[gemini-code-assist]

Code Review

This pull request successfully enables Sequence Parallelism for ROCm platforms by generalizing the platform check and adding ROCm-specific configurations and pattern matchers for AITER operations. The changes are well-structured and accompanied by a comprehensive set of unit and distributed tests.

My main feedback is to improve the robustness of the new distributed test by removing its dependency on an external Hugging Face model configuration. Mocking this configuration would make the test more reliable and self-contained. Overall, this is a great contribution that extends an important optimization to a new hardware platform.

[gemini-code-assist]

This test relies on fetching a model configuration from the Hugging Face Hub ("RedHatAI/Llama-3.2-1B-Instruct-FP8"), which introduces an external dependency. This can make the test flaky if the model is unavailable or due to network issues. The test's behavior depends on the hidden_size from this external config to determine if sequence parallelism should be applied.

To make the test more robust and self-contained, I recommend mocking the configuration fetching to remove this external dependency. This also makes the test's preconditions (like hidden_size being large enough to trigger sequence parallelism) explicit.

You can use unittest.mock.patch to replace vllm.config.model.get_config and return a mock object with a fixed hidden_size. For example:

from unittest.mock import MagicMock, patch

# ... inside _rocm_aiter_sp_worker_fn
    mock_hf_config = MagicMock()
    # Ensure hidden_size is large enough to trigger sequence parallelism
    mock_hf_config.hidden_size = 8192
    with patch("vllm.config.model.get_config", return_value=mock_hf_config):
        model_config = ModelConfig(
            model="mock_model", trust_remote_code=True, dtype=dtype, seed=42
        )
        vllm_config = VllmConfig(
            model_config=model_config,
            device_config=device_config,
            compilation_config=compilation_config,
        )

        with set_current_vllm_config(vllm_config):
            # ... rest of the test

This change would require adding the necessary imports at the top of the file.

[AndreasKaratzas]

This is the same model used for the CUDA test. Other than that, using HF model is the easiest way to perform E2E test.

[ChuanLi1101]

@zejun-chen @AndreasKaratzas Could you please review this PR? It enables Sequence Parallelism for ROCm platforms (MI300X/MI325X/MI355X). Thanks!

[ChuanLi1101]

Correction: @zejun-chen @AndreasKaratzas @wuhuikx Could you please review this PR? It enables Sequence Parallelism for ROCm platforms (MI300X/MI325X/MI355X). Thanks!

[AndreasKaratzas]

Sequence parallelism is something that @tjtanaa took a look recently in his fusions PR. @tjtanaa can you helps review this?

[AndreasKaratzas]

I think it's on the right track, apart from few points.

[AndreasKaratzas]

Minor one: you can drop that comment from there, the class name has all the information.

[AndreasKaratzas]

The proper way here is to use current_platform.is_rocm(). However, in this case we are not only trying to run this on ROCm, but also mi3xx, because AITER is not supported on MI250. So you will need to import from vllm.platforms.rocm import on_mi3xx and then use on_mi3xx().

[ProExpertProg]

Thank you for this work!

Can we get e2e perf & accuracy eval, and can you check e2e_correctness tests pass, and add them to AMD CI?

The new patterns will soon be redundant due to vLLM IR. To reduce churn, I would prefer to first merge vLLM IR PRs and reap immediate benefits.

Instead, could you start working on the gemm+comms fusion (AsyncTP), get it to work, and see how much performance benefit we can achieve? By the time that's done, vLLM IR RMSNorm will be in so we can merge this PR and then the AsyncTP one right after. How does that sound?

[mergify]

Documentation preview: https://vllm--36851.org.readthedocs.build/en/36851/

[tjtanaa]

@ProExpertProg @AndreasKaratzas @ChuanLi1101 iirc, the unit test passes, but it could not mean it is functionally usable. Fusion Passes unit tests only attempted to fuse/split the nodes.

@ChuanLi1101 please follow @ProExpertProg feedback. Provide vllm serve accuracy results, in addition to unit tests results.

[ChuanLi1101]

Thanks for the feedback @ProExpertProg!

vLLM IR & pattern churn: Agreed — I've dropped the AITER-specific patterns (AiterFirstAllReduceRMSNormPattern / AiterMiddleAllReduceRMSNormPattern) and their ROCm-specific distributed test in the latest push. The PR now only contains:

• Platform check relaxation (is_cuda() → is_cuda_alike())
• SP threshold configs for gfx942 (MI300X) and gfx950 (MI325X/MI355X)
• ROCm threshold unit tests
• Updated fusions.md support matrix

Once vLLM IR RMSNorm merges, the existing patterns will work on ROCm without needing AITER-specific ones.

AsyncTP (GEMM+comms fusion): I'd love to work on this. I did some investigation — the current AsyncTPPass relies on torch.ops.symm_mem primitives (fused_matmul_reduce_scatter, fused_all_gather_matmul, and their scaled variants). ROCm SymmetricMemory support was merged into PyTorch, but there are known limitations: multicast/multimem is not available on ROCm hardware, and 7 AsyncTP tests are currently skipped on ROCm including the fused GEMM+collective ops. I'll investigate whether these symm_mem ops actually work on MI300X/MI325X with RCCL and report back.

E2E eval: I'll run e2e perf & accuracy benchmarks on MI325X and share numbers. I'll also add the correctness_e2e/test_sequence_parallel.py tests to AMD CI.

Proposed plan:

1. ✅ Slim down this PR (done — dropped AITER patterns, kept only minimal SP enablement)
2. Investigate torch.ops.symm_mem on MI300X/MI325X, share findings
3. Work on AsyncTP for ROCm based on investigation results
4. Run e2e perf benchmarks for both SP-only and SP+AsyncTP once ready

Does this look good?

[ProExpertProg]

Yep, keep us posted, thank you!

[ProExpertProg]

Oh also, could you create a tracking issue with the work and I can assign you and add it to the central tracker?

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @ChuanLi1101.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[ChuanLi1101]

@AndreasKaratzas — thanks for the review! Addressed your points. The test suite was also reworked since your review (the HF-model distributed E2E test was replaced with mockable capability-based unit tests), so responding against the current state:

1. MI3xx gating (was: use on_mi3xx() instead of is_rocm()) — Done, with a twist. Since the tests now drive behavior through a mocked current_platform.get_device_capability(), on_mi3xx() (which reads the import-time _ON_MI3XX constant and can't be mocked) wouldn't be testable on non-MI3xx CI. So I gate explicitly on the MI3xx capabilities instead:

_ROCM_SP_SUPPORTED_CAPABILITIES = (94, 95)  # gfx942 (MI300X), gfx950 (MI325X/MI355X)
...
elif current_platform.is_rocm():
    ...
    if device_capability not in _ROCM_SP_SUPPORTED_CAPABILITIES:
        return None

This is functionally equivalent to on_mi3xx() (MI250/gfx90a stays disabled) but mockable. It also fixes a latent bug: MI250 (gfx90a) reports capability 90, which aliased the CUDA H100 entry (also 90) in the shared threshold dicts and would previously have enabled SP on MI250. There's a unit test for this (test_unsupported_rocm_capability_returns_none, capability (9, 0) -> None).

2. Redundant comment (class name says it all) — Resolved by the test rewrite; the old comment is gone. The ROCm test class now just carries a one-line docstring.

3. HF model in the E2E test — The HF-model distributed test was dropped entirely in favor of self-contained mocked unit tests (TestGetSequenceParallelismThresholdROCm), which also resolves gemini's flakiness concern about the external HF config dependency.

PR is MERGEABLE against latest main. Note: I couldn't run the pytest suite in my local env (no ROCm build / missing deps) — leaving that to CI and on-hardware validation. Could you take another look / re-approve when you get a chance? Thanks!

(AI-assisted change + comment; on-hardware validation by the submitter.)

[Rohan138]

Closing due to #60517