[AMD][CI] Fix Pipeline + Context Parallelism test group (#46650)
Signed-off-by: aarushjain29 <Aarushi.Jain2@amd.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>
Aarushi Jain committed
Jun 24, 2026 at 22:23 UTC
dc55936f6477406f430a42778774af4ee28d89c0
1 file changed
+7
-6
tests/distributed/test_pp_cudagraph.py
+7
-6
@@ -1,7 +1,8 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import pytest
4
-from typing_extensions import LiteralString
4
+
5
+from vllm.platforms import current_platform
6
7
from ..utils import compare_two_settings, create_new_process_for_each_test
8
@@ -14,15 +15,13 @@ from ..utils import compare_two_settings, create_new_process_for_each_test
15
)
16
@pytest.mark.parametrize(
17
"ATTN_BACKEND",
17
- [
18
- "FLASH_ATTN",
19
- ],
18
+ [None] if current_platform.is_rocm() else ["FLASH_ATTN"],
19
)
20
@create_new_process_for_each_test()
21
def test_pp_cudagraph(
22
PP_SIZE: int,
23
MODEL_NAME: str,
25
- ATTN_BACKEND: LiteralString,
24
+ ATTN_BACKEND: str | None,
25
):
26
cudagraph_args = [
27
# use half precision for speed and memory savings in CI environment
@@ -32,8 +31,10 @@ def test_pp_cudagraph(
31
str(PP_SIZE),
32
"--distributed-executor-backend",
33
"mp",
35
- f"--attention-backend={ATTN_BACKEND}",
34
]
35
+ # On ROCm, defer to the platform attention selector instead of forcing a backend.
36
+ if ATTN_BACKEND is not None:
37
+ cudagraph_args.append(f"--attention-backend={ATTN_BACKEND}")
38
39
eager_args = cudagraph_args + ["--enforce-eager"]
40