@tduncan / Nvidia-Optimization-Stack-VLLM / commits / c92b29a1d4

[Bugfix] Make metadata send non-blocking in GroupCoordinator.isend_tensor_dict (#49274)

Signed-off-by: z00897177 <zhouxinyi6@huawei.com>

AI-Infra Rookie committed Aug 31, 2026 at 07:51 UTC c92b29a1d40644da710209f862b1be0ebd5c2e74

3 files changed +282 -13

tests/distributed/test_comm_ops.py

+136

@@ -5,6 +5,7 @@

5 Run `pytest tests/distributed/test_comm_ops.py`.

6 """

8 +from collections import deque

9 from collections.abc import Callable

10 from typing import Any

11 from unittest.mock import Mock

@@ -208,9 +209,14 @@ def send_recv_tensor_dict_test_worker(

209 class _DummyWork:

210 def __init__(self) -> None:

211 self.wait_calls = 0

212 + self.completed = False

213

214 def wait(self) -> None:

215 self.wait_calls += 1

216 + self.completed = True

217 +

218 + def is_completed(self) -> bool:

219 + return self.completed

220

221

222 class _DummyAllGatherGroup:

@@ -235,6 +241,7 @@ def _make_group_for_unit_test(

241 g.use_cpu_custom_send_recv = False

242 g.device_group = None

243 g.cpu_group = None

244 + g._pending_isends = deque()

245 return g

246

247

@@ -402,6 +409,135 @@ def test_flashinfer_all_reduce_precedes_nccl(monkeypatch: pytest.MonkeyPatch) ->

409 nccl_selector.assert_not_called()

410

411

412 +def test_isend_object_posts_size_then_object_and_releases_on_wait(

413 + monkeypatch: pytest.MonkeyPatch,

414 +) -> None:

415 + posted: list[tuple[torch.Tensor, _DummyWork]] = []

416 +

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

418 + w = _DummyWork()

419 + posted.append((t, w))

420 + return w

421 +

422 + monkeypatch.setattr(torch.distributed, "isend", fake_isend)

423 +

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

425 + handle = g.isend_object({"k": [1, 2, 3]}, dst=1)

426 +

427 + # two sends, in size-then-object order (preserves gloo FIFO).

428 + assert len(posted) == 2

429 + assert posted[0][0].dtype == torch.long

430 + assert posted[0][0].shape == torch.Size([1])

431 + assert posted[0][0].item() == posted[1][0].numel()

432 + assert posted[1][0].dtype == torch.uint8

433 +

434 + # retain holds both source tensors until wait.

435 + assert handle._retained == (posted[0][0], posted[1][0])

436 +

437 + handle.wait()

438 +

439 + # both underlying works drained, retain dropped.

440 + assert all(w.wait_calls == 1 for _, w in posted)

441 + assert handle._retained == ()

442 +

443 + # wait is idempotent: gloo Work.wait() is single-shot for p2p sends (a

444 + # second wait blocks forever), and both the lazy FIFO reap in

445 + # ``_reap_completed_isends`` and any explicit isend caller may wait the

446 + # same metadata handle.

447 + handle.wait()

448 + assert all(w.wait_calls == 1 for _, w in posted)

449 + assert handle._retained == ()

450 +

451 +

452 +def test_isend_tensor_dict_includes_metadata_handle(

453 + monkeypatch: pytest.MonkeyPatch,

454 +) -> None:

455 + posted: list[torch.Tensor] = []

456 +

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

458 + posted.append(t)

459 + return _DummyWork()

460 +

461 + monkeypatch.setattr(torch.distributed, "isend", fake_isend)

462 +

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

464 + td = {"a": torch.arange(4, dtype=torch.float32, device="cpu")}

465 + handles = g.isend_tensor_dict(td, dst=1)

466 +

467 + # size + object (metadata) + one tensor send.

468 + assert len(posted) == 3

469 + assert posted[0].dtype == torch.long

470 + assert posted[1].dtype == torch.uint8

471 + assert posted[2].dtype == torch.float32

472 +

473 + # composite metadata handle + one tensor handle.

474 + assert len(handles) == 2

475 +

476 + for handle in handles:

477 + handle.wait()

478 + assert handles[0]._retained == ()

479 +

480 +

481 +def test_isend_tensor_dict_self_retains_for_fire_and_forget(

482 + monkeypatch: pytest.MonkeyPatch,

483 +) -> None:

484 + works: list[_DummyWork] = []

485 +

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

487 + w = _DummyWork()

488 + works.append(w)

489 + return w

490 +

491 + monkeypatch.setattr(torch.distributed, "isend", fake_isend)

492 +

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

494 + td = {"a": torch.arange(4, dtype=torch.float32, device="cpu")}

495 +

496 + # fire-and-forget: the returned handles are dropped by the caller, so the

497 + # group must self-retain the handles and the source tensor.

498 + g.isend_tensor_dict(td, dst=1)

499 + assert len(g._pending_isends) == 1

500 + handles0, tensors0 = g._pending_isends[0]

501 + assert len(handles0) == 2 # metadata composite + one tensor send

502 + assert tensors0 == [td["a"]]

503 +

504 + # the first send is still in flight, so the next call's reap keeps it.

505 + g.isend_tensor_dict(td, dst=1)

506 + assert len(g._pending_isends) == 2

507 +

508 + # metadata-only send (no tensors) is dropped best-effort on the next

509 + # reap, since it has no reliable completion signal.

510 + g.isend_tensor_dict({"meta": "no-tensors"}, dst=1)

511 + assert len(g._pending_isends) == 3

512 +

513 + # complete every posted work; the next call reaps all three old entries

514 + # (the metadata-only one unconditionally) and keeps only its own.

515 + for w in works:

516 + w.completed = True

517 + g.isend_tensor_dict(td, dst=1)

518 + assert len(g._pending_isends) == 1

519 + # reaped entries had their metadata handles waited as a backstop.

520 + assert handles0[0]._retained == ()

521 +

522 +

523 +def test_send_tensor_dict_sync_path_does_not_self_retain(

524 + monkeypatch: pytest.MonkeyPatch,

525 +) -> None:

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

527 + return _DummyWork()

528 +

529 + monkeypatch.setattr(torch.distributed, "isend", fake_isend)

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

531 + monkeypatch.setattr(GroupCoordinator, "send_object", lambda self, obj, dst: None)

532 +

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

534 + td = {"a": torch.arange(4, dtype=torch.float32, device="cpu")}

535 + g.send_tensor_dict(td, dst=1)

536 +

537 + # the sync path waited every handle, so no retention entry may leak.

538 + assert len(g._pending_isends) == 0

539 +

540 +

541 def test_async_intermediate_tensors_lazy_wait() -> None:

542 work = _DummyWork()

543 post_calls = {"n": 0}

vllm/distributed/parallel_state.py

+140 -3

@@ -27,7 +27,7 @@ import contextlib

27 import gc

28 import pickle

29 import weakref

30 -from collections import namedtuple

30 +from collections import deque, namedtuple

31 from collections.abc import Callable

32 from contextlib import contextmanager, nullcontext

33 from dataclasses import dataclass

@@ -78,6 +78,37 @@ class Handle(Protocol):

78 def wait(self) -> None: ...

81 +class _RetainedHandle:

82 + """Handle that retains source tensors until all posted work completes.

83 +

84 + Used by ``isend_object`` to keep the serialized CPU tensors alive while

85 + gloo copies them in the background; ``wait`` drops the refs once drained.

86 +

87 + ``wait`` must be idempotent: gloo ``Work.wait()`` is single-shot for

88 + point-to-point operations — calling it a second time on an already

89 + completed send blocks forever (verified empirically). Both the worker's

90 + top-of-step drain of ``_pp_send_work`` and ``_reap_completed_isends``

91 + may wait the same entry's metadata handle, so the second caller must

92 + not re-wait the underlying works.

93 + """

94 +

95 + def __init__(self, works: list[Any], retained: tuple[torch.Tensor, ...]) -> None:

96 + self._works = works

97 + self._retained = retained

98 + self._waited = False

99 +

100 + def is_completed(self) -> bool:

101 + return all(w.is_completed() for w in self._works)

102 +

103 + def wait(self) -> None:

104 + if self._waited:

105 + return

106 + for w in self._works:

107 + w.wait()

108 + self._waited = True

109 + self._retained = ()

110 +

111 +

112 def _split_tensor_dict(

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

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

@@ -532,6 +563,15 @@ class GroupCoordinator:

563 and getattr(self.device_communicator, "supports_tensor_dict", False)

564 )

565

566 + # FIFO of in-flight ``isend_tensor_dict`` calls: one entry per call,

567 + # (handles, sent source tensors). Self-retention keeps the send

568 + # buffers alive for fire-and-forget callers that drop the returned

569 + # handles; entries are reaped lazily on the next ``isend_tensor_dict``

570 + # (see ``_reap_completed_isends``). Lazily created instead of plain

571 + # attribute init because some subclasses replicate ``__init__``

572 + # without calling ``super().__init__``.

573 + self._pending_isends: deque[tuple[list[Handle], list[torch.Tensor]]] = deque()

574 +

575 def make_sibling_device_group(self, group_desc: str | None = None) -> ProcessGroup:

576 """Create a new device-side ProcessGroup with the same per-rank membership

577 as this coordinator's `device_group`, but backed by a distinct communicator.

@@ -882,6 +922,33 @@ class GroupCoordinator:

922

923 return obj

924

925 + def isend_object(self, obj: Any, dst: int) -> Handle:

926 + """Non-blocking ``send_object``: isend size + pickled object on the

927 + CPU group. Returns a handle that retains the serialized source

928 + tensors until ``wait`` drains both sends.

929 + """

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

931 + assert dst != self.rank_in_group, (

932 + "Invalid destination rank. Destination rank is the same "

933 + "as the current rank."

934 + )

935 +

936 + object_tensor = torch.frombuffer(pickle.dumps(obj), dtype=torch.uint8)

937 + size_tensor = torch.tensor(

938 + [object_tensor.numel()], dtype=torch.long, device="cpu"

939 + )

940 +

941 + retained = (size_tensor, object_tensor)

942 + works: list[Any] = []

943 + for tensor in retained:

944 + work = torch.distributed.isend(

945 + tensor, dst=self.ranks[dst], group=self.cpu_group

946 + )

947 + if work is not None:

948 + works.append(work)

949 +

950 + return _RetainedHandle(works, retained)

951 +

952 def broadcast_tensor_dict(

953 self,

954 tensor_dict: dict[str, torch.Tensor | Any] | None = None,

@@ -1014,8 +1081,51 @@ class GroupCoordinator:

1081 )

1082 for handle in handles:

1083 handle.wait()

1084 + # The sync path waited every handle, so the self-retention entry

1085 + # appended by ``isend_tensor_dict`` is already drained; drop it

1086 + # instead of leaking the tensor refs until the next async send.

1087 + pending = getattr(self, "_pending_isends", None)

1088 + if pending and pending[-1][0] is handles:

1089 + pending.pop()

1090 return None

1091

1092 + def _reap_completed_isends(self) -> None:

1093 + """Lazily drop self-retained ``isend_tensor_dict`` entries that have

1094 + completed, oldest first (FIFO).

1095 +

1096 + gloo ``Work.is_completed()`` is unreliable (it can report completion

1097 + before the background copy of the source buffer finishes), so the

1098 + metadata handle (``handles[0]``, a gloo-backed ``_RetainedHandle``) is

1099 + never used as the completion gate. Instead we gate on the tensor-send

1100 + handles (``handles[1:]``), which run on the device backend where

1101 + ``is_completed()`` is reliable. Once all tensor sends are done the

1102 + peer must already have posted the matching tensor recvs — and since

1103 + the receiver ingests metadata before tensors, the gloo metadata send

1104 + is then guaranteed complete, so ``wait()`` on it is ~instant and only

1105 + serves to drop the retained refs.

1106 +

1107 + Metadata-only entries (``handles[1:]`` empty, e.g. an empty

1108 + ``IntermediateTensors``) have no reliable completion signal; drop them

1109 + best-effort without blocking on the gloo metadata wait.

1110 + """

1111 + pending = getattr(self, "_pending_isends", None)

1112 + if pending is None:

1113 + # Subclasses that replicate __init__ without calling

1114 + # super().__init__ may not have the attribute yet.

1115 + self._pending_isends = pending = deque()

1116 + while pending:

1117 + handles, _ = pending[0]

1118 + tensor_handles = handles[1:]

1119 + if not tensor_handles:

1120 + pending.popleft()

1121 + continue

1122 + if not all(h.is_completed() for h in tensor_handles):

1123 + # The oldest send is still in flight; FIFO ordering means the

1124 + # rest are no further along, so stop here.

1125 + break

1126 + handles[0].wait()

1127 + pending.popleft()

1128 +

1129 def isend_tensor_dict(

1130 self,

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

@@ -1023,6 +1133,16 @@ class GroupCoordinator:

1133 all_gather_group: "GroupCoordinator | None" = None,

1134 all_gather_tensors: dict[str, bool] | None = None,

1135 ) -> list[Handle]:

1136 + """Send the input tensor dictionary asynchronously.

1137 +

1138 + The returned handles are additionally self-retained in

1139 + ``self._pending_isends`` (together with the sent source tensors) and

1140 + reaped lazily on subsequent calls, so fire-and-forget callers that

1141 + drop the handles cannot free or overwrite the send buffers while the

1142 + sends are still in flight. Callers that do keep the handles may

1143 + ``wait()`` them as before; the retention entry is dropped on reap

1144 + either way.

1145 + """

1146 if self.world_size <= 1:

1147 return []

1148

@@ -1048,12 +1168,16 @@ class GroupCoordinator:

1168 metadata_group = self.cpu_group

1169

1170 metadata_list, tensor_list = _split_tensor_dict(tensor_dict)

1051 - self.send_object(metadata_list, dst=dst)

1171

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

1173 assert len(tensor_keys) == len(tensor_list)

1174

1056 - handles: list[Handle] = []

1175 + # Reap completed sends before posting new ones, bounding the

1176 + # self-retention FIFO.

1177 + self._reap_completed_isends()

1178 + handles: list[Handle] = [self.isend_object(metadata_list, dst=dst)]

1179 +

1180 + sent_tensors: list[torch.Tensor] = []

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

1182 if tensor.numel() == 0:

1183 continue

@@ -1070,6 +1194,19 @@ class GroupCoordinator:

1194 if tensor.is_cuda:

1195 tensor.record_stream(torch.cuda.current_stream(tensor.device))

1196 handles.append(handle)

1197 + sent_tensors.append(tensor)

1198 +

1199 + # Self-retain for fire-and-forget callers: keep the handles and

1200 + # the source tensors alive until the sends complete (reaped

1201 + # lazily by ``_reap_completed_isends``). Without this, dropping

1202 + # the returned handles would GC the ``_RetainedHandle``'s

1203 + # serialized metadata buffers mid-send, and non-CUDA devices

1204 + # have no ``record_stream`` equivalent protecting the tensor

1205 + # sources.

1206 + pending = getattr(self, "_pending_isends", None)

1207 + if pending is None:

1208 + self._pending_isends = pending = deque()

1209 + pending.append((handles, sent_tensors))

1210

1211 return handles

1212

vllm/v1/worker/gpu_worker.py

+6 -10

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

216 self.profiler_config = vllm_config.profiler_config

217

218 self.use_v2_model_runner = vllm_config.use_v2_model_runner

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

220 - self._pp_send_work: list[Handle] = []

219

220 # Resolved lazily on first sleep/wake; persists worker-process state.

221 self._sleep_mode_backend: SleepModeBackend | None = None

@@ -1108,12 +1106,6 @@ class Worker(WorkerBase):

1106 def execute_model(

1107 self, scheduler_output: "SchedulerOutput"

1108 ) -> ModelRunnerOutput | AsyncModelRunnerOutput | None:

1111 - # ensure any previous non-blocking PP sends are complete

1112 - if self._pp_send_work:

1113 - for handle in self._pp_send_work:

1114 - handle.wait()

1115 - self._pp_send_work = []

1116 -

1109 intermediate_tensors = None

1110 forward_pass = scheduler_output.total_num_scheduled_tokens > 0

1111 num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens

@@ -1186,8 +1178,12 @@ class Worker(WorkerBase):

1178 and not get_pp_group().is_last_rank

1179 )

1180

1189 - # launch non-blocking send of intermediate tensors

1190 - self._pp_send_work = get_pp_group().isend_tensor_dict(

1181 + # launch non-blocking send of intermediate tensors; the

1182 + # GroupCoordinator self-retains the source tensors and the

1183 + # metadata handle in ``_pending_isends`` and reaps them lazily on

1184 + # the next ``isend_tensor_dict(is_async=True)`` call, so this

1185 + # step returns without waiting for the receiver to drain.

1186 + get_pp_group().isend_tensor_dict(

1187 output.tensors,

1188 all_gather_group=get_tp_group(),

1189 all_gather_tensors=all_gather_tensors,