@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 3d2a026fd0

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 8, 08:45 UTC

[Feature] Pipeline Parallel Async send/recv, 2.9% E2E throughput improvement (#33368)

Signed-off-by: yewentao256 <zhyanwentao@126.com> Signed-off-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Co-authored-by: Nick Hill <nhill@redhat.com>

Wentao Ye committed Feb 13, 2026 at 03:38 UTC 3d2a026fd0317204752d7933408aff19aaa80cfd

3 files changed +298 -81

tests/distributed/test_comm_ops.py

+107

@@ -19,6 +19,8 @@ from vllm.distributed import (

19 tensor_model_parallel_all_reduce,

20 tensor_model_parallel_reduce_scatter,

21 )

22 +from vllm.distributed.parallel_state import GroupCoordinator, TensorMetadata

23 +from vllm.v1.worker.gpu_worker import AsyncIntermediateTensors

25 from ..utils import (

26 init_test_distributed_environment,

@@ -200,6 +202,111 @@ def send_recv_tensor_dict_test_worker(

202 torch.testing.assert_close(recv_dict["f"], test_dict["f"])

203

204

205 +class _DummyWork:

206 + def __init__(self) -> None:

207 + self.wait_calls = 0

208 +

209 + def wait(self) -> None:

210 + self.wait_calls += 1

211 +

212 +

213 +class _DummyAllGatherGroup:

214 + def __init__(self, world_size: int, rank_in_group: int) -> None:

215 + self.world_size = world_size

216 + self.rank_in_group = rank_in_group

217 +

218 + def all_gather(self, t: torch.Tensor, dim: int = 0) -> torch.Tensor:

219 + # duplicate local slice across ranks.

220 + assert dim == 0

221 + return torch.cat([t for _ in range(self.world_size)], dim=0)

222 +

223 +

224 +def _make_group_for_unit_test(

225 + rank_in_group: int = 0, world_size: int = 2

226 +) -> GroupCoordinator:

227 + # avoid running GroupCoordinator.__init__ (it wires up real process groups).

228 + g = GroupCoordinator.__new__(GroupCoordinator)

229 + g.world_size = world_size

230 + g.rank_in_group = rank_in_group

231 + g.ranks = list(range(world_size))

232 + g.use_cpu_custom_send_recv = False

233 + g.device_group = None

234 + g.cpu_group = None

235 + return g

236 +

237 +

238 +def test_irecv_tensor_dict_send_allgather_postprocess_binds_keys(

239 + monkeypatch: pytest.MonkeyPatch,

240 +) -> None:

241 + def fake_irecv(t: torch.Tensor, *args: Any, **kwargs: Any) -> _DummyWork:

242 + t.fill_(1)

243 + return _DummyWork()

244 +

245 + monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True)

246 + monkeypatch.setattr(torch.distributed, "irecv", fake_irecv)

247 +

248 + g = _make_group_for_unit_test(rank_in_group=0, world_size=2)

249 + # 2 tensors so we can catch late-binding bugs in postprocess closures.

250 + metadata_list = [

251 + ("a", TensorMetadata("cpu", torch.int32, torch.Size([4]))),

252 + ("b", TensorMetadata("cpu", torch.int32, torch.Size([4]))),

253 + ]

254 + g.recv_object = lambda src=None: metadata_list # type: ignore[method-assign]

255 +

256 + ag = _DummyAllGatherGroup(world_size=2, rank_in_group=0)

257 + td, handles, postprocess = g.irecv_tensor_dict(all_gather_group=ag)

258 +

259 + assert td is not None

260 + assert len(handles) == 2

261 + assert len(postprocess) == 2

262 +

263 + # before postprocess, dict holds the TP slice (shape 2).

264 + assert td["a"].shape == torch.Size([2])

265 + assert td["b"].shape == torch.Size([2])

266 +

267 + # simulate worker-side "defer wait": wait + postprocess later.

268 + for handle in handles:

269 + handle.wait()

270 + for fn in postprocess:

271 + fn()

272 +

273 + # after postprocess, dict values are reconstructed to full shape (shape 4),

274 + # and each key should be updated independently

275 + assert td["a"].shape == torch.Size([4])

276 + assert td["b"].shape == torch.Size([4])

277 + torch.testing.assert_close(td["a"], torch.ones(4, dtype=torch.int32))

278 + torch.testing.assert_close(td["b"], torch.ones(4, dtype=torch.int32))

279 +

280 +

281 +def test_async_intermediate_tensors_lazy_wait() -> None:

282 + work = _DummyWork()

283 + post_calls = {"n": 0}

284 +

285 + def post() -> None:

286 + post_calls["n"] += 1

287 +

288 + it = AsyncIntermediateTensors(

289 + {"x": torch.tensor([1])},

290 + comm_handles=[work],

291 + comm_postprocess=[post],

292 + )

293 +

294 + # accessing non-tensor attributes should not trigger wait.

295 + assert it.kv_connector_output is None

296 + assert work.wait_calls == 0

297 + assert post_calls["n"] == 0

298 +

299 + # first access of `.tensors` triggers wait + postprocess.

300 + _ = it.tensors

301 + assert work.wait_calls == 1

302 + assert post_calls["n"] == 1

303 +

304 + # subsequent access should not re-wait.

305 + _ = it.tensors

306 + assert work.wait_calls == 1

307 + assert post_calls["n"] == 1

308 +

309 +

310 @ray.remote(num_gpus=1, max_calls=1)

311 def send_recv_test_worker(

312 monkeypatch: pytest.MonkeyPatch,

vllm/distributed/parallel_state.py

+137 -76

@@ -33,7 +33,7 @@ from contextlib import contextmanager, nullcontext

33 from dataclasses import dataclass

34 from datetime import timedelta

35 from multiprocessing import shared_memory

36 -from typing import Any

36 +from typing import Any, Protocol

37 from unittest.mock import patch

39 import torch

@@ -64,6 +64,14 @@ class GraphCaptureContext:

64 TensorMetadata = namedtuple("TensorMetadata", ["device", "dtype", "size"])

67 +class Handle(Protocol):

68 + """Minimal async work handle used by P2P send/recv methods."""

69 +

70 + def is_completed(self) -> bool: ...

71 +

72 + def wait(self) -> None: ...

73 +

74 +

75 def _split_tensor_dict(

76 tensor_dict: dict[str, torch.Tensor | Any],

77 ) -> tuple[list[tuple[str, Any]], list[torch.Tensor]]:

@@ -780,6 +788,20 @@ class GroupCoordinator:

788 async_handle.wait()

789 return tensor_dict

790

791 + def _should_use_all_gather(

792 + self,

793 + key: str,

794 + numel: int,

795 + all_gather_group: "GroupCoordinator | None",

796 + all_gather_tensors: dict[str, bool] | None,

797 + ) -> bool:

798 + if all_gather_group is None:

799 + return False

800 + use_all_gather = numel % all_gather_group.world_size == 0

801 + if all_gather_tensors is not None:

802 + use_all_gather = all_gather_tensors.get(key, use_all_gather)

803 + return use_all_gather

804 +

805 def send_tensor_dict(

806 self,

807 tensor_dict: dict[str, torch.Tensor | Any],

@@ -808,6 +830,35 @@ class GroupCoordinator:

830 # Bypass the function if we are using only 1 GPU.

831 if not torch.distributed.is_initialized() or self.world_size == 1:

832 return tensor_dict

833 + handles = self.isend_tensor_dict(

834 + tensor_dict,

835 + dst=dst,

836 + all_gather_group=all_gather_group,

837 + all_gather_tensors=all_gather_tensors,

838 + )

839 + for handle in handles:

840 + handle.wait()

841 + return None

842 +

843 + def isend_tensor_dict(

844 + self,

845 + tensor_dict: dict[str, torch.Tensor | Any],

846 + dst: int | None = None,

847 + all_gather_group: "GroupCoordinator | None" = None,

848 + all_gather_tensors: dict[str, bool] | None = None,

849 + ) -> list[Handle]:

850 + if self.world_size <= 1:

851 + return []

852 +

853 + if self.use_cpu_custom_send_recv:

854 + if self.device_communicator is None:

855 + raise ValueError("No device communicator found")

856 + # custom device communicator path is synchronous

857 + self.device_communicator.send_tensor_dict( # type: ignore

858 + tensor_dict, dst

859 + )

860 + return []

861 +

862 all_gather_size = 1 if all_gather_group is None else all_gather_group.world_size

863 all_gather_rank = (

864 0 if all_gather_group is None else all_gather_group.rank_in_group

@@ -820,53 +871,31 @@ class GroupCoordinator:

871 dst = (self.rank_in_group + 1) % self.world_size

872 assert dst < self.world_size, f"Invalid dst rank ({dst})"

873

823 - if self.use_cpu_custom_send_recv:

824 - if self.device_communicator is None:

825 - raise ValueError("No device communicator found")

826 - self.device_communicator.send_tensor_dict( # type: ignore

827 - tensor_dict, dst

828 - )

829 - return None

830 -

831 - metadata_list: list[tuple[Any, Any]] = []

832 - assert isinstance(tensor_dict, dict), (

833 - f"Expecting a dictionary, got {type(tensor_dict)}"

834 - )

874 metadata_list, tensor_list = _split_tensor_dict(tensor_dict)

836 - # `metadata_list` lives in CPU memory.

837 - # `send_object_list` has serialization & deserialization,

838 - # all happening on CPU. Therefore, we can use the CPU group.

875 self.send_object(metadata_list, dst=dst)

876

877 tensor_keys = [k for k, v in tensor_dict.items() if isinstance(v, torch.Tensor)]

878 assert len(tensor_keys) == len(tensor_list)

879

880 + handles: list[Handle] = []

881 for key, tensor in zip(tensor_keys, tensor_list):

882 if tensor.numel() == 0:

846 - # Skip sending empty tensors.

883 continue

884

849 - # send-allgather: send only a slice, then do allgather.

850 - use_all_gather = (

851 - all_gather_group is not None and tensor.numel() % all_gather_size == 0

852 - )

853 - use_all_gather = (

854 - all_gather_tensors.get(key, use_all_gather)

855 - if all_gather_tensors

856 - else use_all_gather

857 - )

858 - if use_all_gather:

885 + if self._should_use_all_gather(

886 + key, tensor.numel(), all_gather_group, all_gather_tensors

887 + ):

888 tensor = tensor.reshape(all_gather_size, -1)[all_gather_rank]

889

861 - if tensor.is_cpu:

862 - # use metadata_group for CPU tensors

863 - torch.distributed.send(

864 - tensor, dst=self.ranks[dst], group=metadata_group

865 - )

866 - else:

867 - # use group for GPU tensors

868 - torch.distributed.send(tensor, dst=self.ranks[dst], group=group)

869 - return None

890 + comm_group = metadata_group if tensor.is_cpu else group

891 + handle = torch.distributed.isend(

892 + tensor, dst=self.ranks[dst], group=comm_group

893 + )

894 + if tensor.is_cuda:

895 + tensor.record_stream(torch.cuda.current_stream(tensor.device))

896 + handles.append(handle)

897 +

898 + return handles

899

900 def recv_tensor_dict(

901 self,

@@ -895,6 +924,38 @@ class GroupCoordinator:

924 # Bypass the function if we are using only 1 GPU.

925 if not torch.distributed.is_initialized() or self.world_size == 1:

926 return None

927 + tensor_dict, handles, postprocess = self.irecv_tensor_dict(

928 + src=src,

929 + all_gather_group=all_gather_group,

930 + all_gather_tensors=all_gather_tensors,

931 + )

932 + for handle in handles:

933 + handle.wait()

934 + for fn in postprocess:

935 + fn()

936 + return tensor_dict

937 +

938 + def irecv_tensor_dict(

939 + self,

940 + src: int | None = None,

941 + all_gather_group: "GroupCoordinator | None" = None,

942 + all_gather_tensors: dict[str, bool] | None = None,

943 + ) -> tuple[

944 + dict[str, torch.Tensor | Any] | None,

945 + list[Handle],

946 + list[Callable[[], None]],

947 + ]:

948 + if not torch.distributed.is_initialized() or self.world_size == 1:

949 + return None, [], []

950 + if self.use_cpu_custom_send_recv:

951 + if self.device_communicator is None:

952 + raise ValueError("No device communicator found")

953 + # custom device communicator path is synchronous

954 + sync_tensor_dict = self.device_communicator.recv_tensor_dict( # type: ignore

955 + src

956 + )

957 + return sync_tensor_dict, [], []

958 +

959 all_gather_size = 1 if all_gather_group is None else all_gather_group.world_size

960 all_gather_rank = (

961 0 if all_gather_group is None else all_gather_group.rank_in_group

@@ -907,57 +968,57 @@ class GroupCoordinator:

968 src = (self.rank_in_group - 1) % self.world_size

969 assert src < self.world_size, f"Invalid src rank ({src})"

970

910 - if self.use_cpu_custom_send_recv:

911 - if self.device_communicator is None:

912 - raise ValueError("No device communicator found")

913 - return self.device_communicator.recv_tensor_dict( # type: ignore

914 - src

915 - )

916 -

971 recv_metadata_list = self.recv_object(src=src)

972 tensor_dict: dict[str, Any] = {}

973 + handles: list[Handle] = []

974 + postprocess: list[Callable[[], None]] = []

975 +

976 for key, value in recv_metadata_list:

977 if isinstance(value, TensorMetadata):

921 - tensor = torch.empty(value.size, dtype=value.dtype, device=value.device)

922 - if tensor.numel() == 0:

923 - # Skip broadcasting empty tensors.

924 - tensor_dict[key] = tensor

925 - continue

926 -

927 - # send-allgather: send only a slice, then do allgather.

928 - use_all_gather = (

929 - all_gather_group is not None

930 - and tensor.numel() % all_gather_size == 0

931 - )

932 - use_all_gather = (

933 - all_gather_tensors.get(key, use_all_gather)

934 - if all_gather_tensors

935 - else use_all_gather

978 + full_tensor = torch.empty(

979 + value.size, dtype=value.dtype, device=value.device

980 )

981 + if full_tensor.numel() == 0:

982 + tensor_dict[key] = full_tensor

983 + continue

984

938 - if use_all_gather:

939 - orig_shape = tensor.shape

940 - tensor = tensor.reshape(all_gather_size, -1)[all_gather_rank]

941 -

942 - if tensor.is_cpu:

943 - # use metadata_group for CPU tensors

944 - torch.distributed.recv(

945 - tensor, src=self.ranks[src], group=metadata_group

985 + if self._should_use_all_gather(

986 + key, full_tensor.numel(), all_gather_group, all_gather_tensors

987 + ):

988 + orig_shape = full_tensor.shape

989 + slice_tensor = full_tensor.reshape(all_gather_size, -1)[

990 + all_gather_rank

991 + ]

992 + comm_group = metadata_group if slice_tensor.is_cpu else group

993 + handle = torch.distributed.irecv(

994 + slice_tensor, src=self.ranks[src], group=comm_group

995 )

996 + handles.append(handle)

997 +

998 + def _postprocess(

999 + key: str = key,

1000 + slice_tensor: torch.Tensor = slice_tensor,

1001 + orig_shape: tuple[int, ...] = tuple(orig_shape),

1002 + all_gather_group=all_gather_group,

1003 + ) -> None:

1004 + assert all_gather_group is not None

1005 + tensor_dict[key] = all_gather_group.all_gather(

1006 + slice_tensor, dim=0

1007 + ).reshape(orig_shape)

1008 +

1009 + postprocess.append(_postprocess)

1010 + tensor_dict[key] = slice_tensor

1011 else:

948 - # use group for GPU tensors

949 - torch.distributed.recv(tensor, src=self.ranks[src], group=group)

950 - if use_all_gather:

951 - # do the allgather

952 - tensor = all_gather_group.all_gather( # type: ignore

953 - tensor, dim=0

1012 + comm_group = metadata_group if full_tensor.is_cpu else group

1013 + handle = torch.distributed.irecv(

1014 + full_tensor, src=self.ranks[src], group=comm_group

1015 )

955 - tensor = tensor.reshape(orig_shape)

956 -

957 - tensor_dict[key] = tensor

1016 + handles.append(handle)

1017 + tensor_dict[key] = full_tensor

1018 else:

1019 tensor_dict[key] = value

960 - return tensor_dict

1020 +

1021 + return tensor_dict, handles, postprocess

1022

1023 def barrier(self):

1024 """Barrier synchronization among the group.

vllm/v1/worker/gpu_worker.py

+54 -5

@@ -4,6 +4,7 @@

5 import gc

6 import os

7 +from collections.abc import Callable

8 from contextlib import AbstractContextManager, nullcontext

9 from types import NoneType

10 from typing import TYPE_CHECKING, Any, cast

@@ -30,6 +31,7 @@ from vllm.distributed.kv_transfer import (

31 has_kv_transfer_group,

32 )

33 from vllm.distributed.parallel_state import (

34 + Handle,

35 get_pcp_group,

36 get_pp_group,

37 get_tp_group,

@@ -68,6 +70,38 @@ if TYPE_CHECKING:

70 from vllm.v1.worker.gpu_model_runner import GPUModelRunner

73 +class AsyncIntermediateTensors(IntermediateTensors):

74 + """IntermediateTensors with lazy comm synchronization"""

75 +

76 + def __init__(

77 + self,

78 + tensors: dict[str, torch.Tensor],

79 + comm_handles: list[Handle] | None = None,

80 + comm_postprocess: list[Callable[[], None]] | None = None,

81 + ) -> None:

82 + super().__init__(tensors)

83 + self._comm_handles = comm_handles

84 + self._comm_postprocess = comm_postprocess

85 + self._comm_waited = False

86 +

87 + def wait_for_comm(self) -> None:

88 + if self._comm_waited:

89 + return

90 + if self._comm_handles:

91 + for handle in self._comm_handles:

92 + handle.wait()

93 + if self._comm_postprocess:

94 + for fn in self._comm_postprocess:

95 + fn()

96 + self._comm_waited = True

97 +

98 + def __getattribute__(self, name: str):

99 + # ensure `.tensors` is ready before use

100 + if name == "tensors" and not object.__getattribute__(self, "_comm_waited"):

101 + object.__getattribute__(self, "wait_for_comm")()

102 + return object.__getattribute__(self, name)

103 +

104 +

105 class Worker(WorkerBase):

106 def __init__(

107 self,

@@ -113,6 +147,8 @@ class Worker(WorkerBase):

147 raise ValueError(f"Unknown profiler type: {self.profiler_config.profiler}")

148

149 self.use_v2_model_runner = envs.VLLM_USE_V2_MODEL_RUNNER

150 + # pending non-blocking PP send work from the previous iteration

151 + self._pp_send_work: list[Handle] = []

152

153 def sleep(self, level: int = 1) -> None:

154 from vllm.device_allocator.cumem import CuMemAllocator

@@ -600,6 +636,12 @@ class Worker(WorkerBase):

636 def execute_model(

637 self, scheduler_output: "SchedulerOutput"

638 ) -> ModelRunnerOutput | AsyncModelRunnerOutput | None:

639 + # ensure any previous non-blocking PP sends are complete

640 + if self._pp_send_work:

641 + for handle in self._pp_send_work:

642 + handle.wait()

643 + self._pp_send_work = []

644 +

645 intermediate_tensors = None

646 forward_pass = scheduler_output.total_num_scheduled_tokens > 0

647 num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens

@@ -637,12 +679,18 @@ class Worker(WorkerBase):

679 }

680

681 if forward_pass and not get_pp_group().is_first_rank:

640 - tensor_dict = get_pp_group().recv_tensor_dict(

641 - all_gather_group=get_tp_group(),

642 - all_gather_tensors=all_gather_tensors,

682 + tensor_dict, comm_handles, comm_postprocess = (

683 + get_pp_group().irecv_tensor_dict(

684 + all_gather_group=get_tp_group(),

685 + all_gather_tensors=all_gather_tensors,

686 + )

687 )

688 assert tensor_dict is not None

645 - intermediate_tensors = IntermediateTensors(tensor_dict)

689 + intermediate_tensors = AsyncIntermediateTensors(

690 + tensor_dict,

691 + comm_handles=comm_handles,

692 + comm_postprocess=comm_postprocess,

693 + )

694

695 with self.annotate_profile(scheduler_output):

696 output = self.model_runner.execute_model(

@@ -660,7 +708,8 @@ class Worker(WorkerBase):

708 and not get_pp_group().is_last_rank

709 )

710

663 - get_pp_group().send_tensor_dict(

711 + # launch non-blocking send of intermediate tensors

712 + self._pp_send_work = get_pp_group().isend_tensor_dict(

713 output.tensors,

714 all_gather_group=get_tp_group(),

715 all_gather_tensors=all_gather_tensors,