@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 73e688cb79

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 14, 05:30 UTC

[ROCm][Feature] Enable Pipeline Parallelism with Ray Compiled Graph on ROCm (#24275)

Signed-off-by: charlifu <charlifu@amd.com>

Charlie Fu committed Sep 9, 2025 at 18:27 UTC 73e688cb79786d4403c227b09a39d9a79a201299

3 files changed +16 -3

docker/Dockerfile.rocm

@@ -104,6 +104,7 @@ COPY --from=export_vllm /examples ${COMMON_WORKDIR}/vllm/examples

104 COPY --from=export_vllm /docker ${COMMON_WORKDIR}/vllm/docker

105

106 ENV RAY_EXPERIMENTAL_NOSET_ROCR_VISIBLE_DEVICES=1

107 +ENV RAY_EXPERIMENTAL_NOSET_HIP_VISIBLE_DEVICES=1

108 ENV TOKENIZERS_PARALLELISM=false

109

110 # ENV that can improve safe tensor loading, and end-to-end time

requirements/rocm.txt

+1 -1

@@ -8,7 +8,7 @@ numba == 0.61.2; python_version > '3.9'

8 boto3

9 botocore

10 datasets

11 -ray>=2.10.0,<2.45.0

11 +ray[cgraph]>=2.48.0 # Ray Compiled Graph, required for pipeline parallelism in V1.

12 peft

13 pytest-asyncio

14 tensorizer==2.10.1

vllm/utils/__init__.py

+14 -2

@@ -78,6 +78,7 @@ if TYPE_CHECKING:

78 from argparse import Namespace

80 from vllm.config import ModelConfig, VllmConfig

81 + from vllm.sequence import IntermediateTensors

83 logger = init_logger(__name__)

@@ -1472,7 +1473,8 @@ def current_stream() -> torch.cuda.Stream:

1473 # is hurting performance. Therefore creating a dedicated stream

1474 # per process

1475 if current_platform.is_rocm():

1475 - _current_stream_tls.value = torch.cuda.Stream()

1476 + # torch.cuda.set_stream here is the alias of _pathed_set_stream

1477 + torch.cuda.set_stream(torch.cuda.Stream())

1478 elif current_platform.is_cpu():

1479 _current_stream_tls.value = _StreamPlaceholder()

1480 else:

@@ -2278,7 +2280,8 @@ def weak_ref_tensor(tensor: Any) -> Any:

2280

2281

2282 def weak_ref_tensors(

2281 - tensors: Union[torch.Tensor, list[torch.Tensor], tuple[torch.Tensor]]

2283 + tensors: Union[torch.Tensor, list[torch.Tensor], tuple[torch.Tensor],

2284 + IntermediateTensors]

2285 ) -> Union[torch.Tensor, list[Any], tuple[Any], Any]:

2286 """

2287 Convenience function to create weak references to tensors,

@@ -2290,6 +2293,15 @@ def weak_ref_tensors(

2293 return [weak_ref_tensor(t) for t in tensors]

2294 if isinstance(tensors, tuple):

2295 return tuple(weak_ref_tensor(t) for t in tensors)

2296 +

2297 + # For IntermediateTensors used in pipeline parallelism

2298 + from vllm.sequence import IntermediateTensors

2299 + if isinstance(tensors, IntermediateTensors):

2300 + ret = IntermediateTensors({

2301 + key: weak_ref_tensor(val)

2302 + for key, val in tensors.tensors.items()

2303 + })

2304 + return ret

2305 raise ValueError("Invalid type for tensors")

2306

2307