Skip to content

[Bugfix][Distributed] Preserve symmetric collective outputs #58154

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
jeffnvidia wants to merge 7 commits into
base: main
Choose a base branch
from
+84 −28
Open

[Bugfix][Distributed] Preserve symmetric collective outputs #58154

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
d8a8bc6
[Bugfix][Distributed] Preserve symmetric collective outputs
jeffnvidia Sep 22, 2026
cfaf005
Test symmetric collective output ownership under CUDA graphs
jeffnvidia Sep 30, 2026
6d30c7b
Merge branch 'main' into fix/nccl-symm-rs-output-ownership
mergify[bot] Oct 2, 2026
03d6f44
Fix symmetric collective graph test formatting
jeffnvidia Oct 2, 2026
4532149
Merge upstream main into fix/nccl-symm-rs-output-ownership
jeffnvidia Oct 2, 2026
ce072b3
Merge upstream main and preserve symmetric output ownership
jeffnvidia Oct 6, 2026
cfd1205
Merge upstream main to include n-gram speculative decoding fix
jeffnvidia Oct 7, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
98 changes: 77 additions & 21 deletions tests/distributed/test_nccl_symm_mem.py
View file
Open in desktop

@samnordmann samnordmann Sep 28, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

could we add CUDA-graph coverage for both operations? Capture two same-key calls, retain both outputs, then replay with changed inputs and verify both results.

@jeffnvidia jeffnvidia Sep 30, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Added CUDA graph coverage in cfaf005: both AllGather and ReduceScatter capture two same-key calls, retain both outputs, and verify them after two replays with changed inputs. The exact PR head passed both tests on 2xGB200 with no skips. NCCL logs confirmed the symmetric paths.

Original file line number Diff line number Diff line change
Expand Up @@ -182,15 +182,49 @@ def nccl_symm_mem_allgather_worker(local_rank: int, world_size: int):
pytest.skip("NCCL symmetric memory is disabled.")

per_rank_size = test_size_elements // world_size
input_tensor = torch.randint(
1, 23, (per_rank_size,), dtype=dtype, device=device
first_input = torch.randint(1, 23, (per_rank_size,), dtype=dtype, device=device)
second_input = torch.randint(
24, 47, (per_rank_size,), dtype=dtype, device=device
)
first_input_clone = first_input.clone()
second_input_clone = second_input.clone()

# Retain the first result while the second public entry point reuses
# the same scratch key.
first_output = cuda_communicator.all_gather(first_input, dim=0)
second_output = cuda_communicator.all_gatherv(
second_input, dim=0, sizes=[per_rank_size] * world_size
)
output = cuda_communicator.all_gatherv(input_tensor, dim=0)

group = get_tp_group().device_group
expected = torch.empty(test_size_elements, dtype=dtype, device=device)
dist.all_gather_into_tensor(expected, input_tensor, group=group)
torch.testing.assert_close(output, expected, atol=0.0, rtol=0.0)
first_expected = torch.empty(test_size_elements, dtype=dtype, device=device)
second_expected = torch.empty(test_size_elements, dtype=dtype, device=device)
dist.all_gather_into_tensor(first_expected, first_input_clone, group=group)
dist.all_gather_into_tensor(second_expected, second_input_clone, group=group)

assert first_output.data_ptr() != second_output.data_ptr()
torch.testing.assert_close(first_output, first_expected, atol=0.0, rtol=0.0)
torch.testing.assert_close(second_output, second_expected, atol=0.0, rtol=0.0)

graph = torch.cuda.CUDAGraph()
with torch.cuda.graph(graph):
graph_first = cuda_communicator.all_gather(first_input, dim=0)
graph_second = cuda_communicator.all_gatherv(
second_input, dim=0, sizes=[per_rank_size] * world_size
)

assert graph_first.data_ptr() != graph_second.data_ptr()
for iteration in range(2):
first_input.fill_(local_rank + 1 + 10 * iteration)
second_input.fill_(local_rank + 11 + 10 * iteration)
dist.all_gather_into_tensor(first_expected, first_input, group=group)
dist.all_gather_into_tensor(second_expected, second_input, group=group)
graph.replay()
torch.accelerator.synchronize()
torch.testing.assert_close(graph_first, first_expected, atol=0.0, rtol=0.0)
torch.testing.assert_close(
graph_second, second_expected, atol=0.0, rtol=0.0
)


@pytest.mark.skipif(
Expand Down Expand Up @@ -245,29 +279,51 @@ def nccl_symm_mem_reduce_scatter_worker(local_rank: int, world_size: int):
pytest.skip("NCCL symmetric memory is disabled.")

per_rank_size = test_size_elements // world_size
ordinary_input = torch.randint(
first_input = torch.randint(
1, 23, (test_size_elements,), dtype=dtype, device=device
)
ordinary_clone = ordinary_input.clone()
ordinary_output = cuda_communicator.reduce_scatter(ordinary_input, dim=0)
second_input = torch.randint(
1, 23, (test_size_elements,), dtype=dtype, device=device
)
first_input_clone = first_input.clone()
second_input_clone = second_input.clone()

# Retain the first result while the second public entry point reuses
# the same scratch key.
first_output = cuda_communicator.reduce_scatter(first_input, dim=0)
second_output = cuda_communicator.reduce_scatterv(
second_input, dim=0, sizes=[per_rank_size] * world_size
)

group = get_tp_group().device_group
ordinary_expected = torch.empty(per_rank_size, dtype=dtype, device=device)
dist.reduce_scatter_tensor(ordinary_expected, ordinary_clone, group=group)
torch.testing.assert_close(
ordinary_output, ordinary_expected, atol=2.5, rtol=0.1
)
assert is_symmetric_memory_tensor(ordinary_output)
first_expected = torch.empty(per_rank_size, dtype=dtype, device=device)
second_expected = torch.empty(per_rank_size, dtype=dtype, device=device)
dist.reduce_scatter_tensor(first_expected, first_input_clone, group=group)
dist.reduce_scatter_tensor(second_expected, second_input_clone, group=group)

# Both calls reuse rs_out; validate the first result before overwriting it.
for sizes in (None, [per_rank_size] * world_size):
ordinary_v_output = cuda_communicator.reduce_scatterv(
ordinary_input, dim=0, sizes=sizes
assert first_output.data_ptr() != second_output.data_ptr()
torch.testing.assert_close(first_output, first_expected, atol=2.5, rtol=0.1)
torch.testing.assert_close(second_output, second_expected, atol=2.5, rtol=0.1)

graph = torch.cuda.CUDAGraph()
with torch.cuda.graph(graph):
graph_first = cuda_communicator.reduce_scatter(first_input, dim=0)
graph_second = cuda_communicator.reduce_scatterv(
second_input, dim=0, sizes=[per_rank_size] * world_size
)

assert graph_first.data_ptr() != graph_second.data_ptr()
for iteration in range(2):
first_input.fill_(local_rank + 1 + 10 * iteration)
second_input.fill_(local_rank + 11 + 10 * iteration)
dist.reduce_scatter_tensor(first_expected, first_input, group=group)
dist.reduce_scatter_tensor(second_expected, second_input, group=group)
graph.replay()
torch.accelerator.synchronize()
torch.testing.assert_close(graph_first, first_expected, atol=0.0, rtol=0.0)
torch.testing.assert_close(
ordinary_v_output, ordinary_expected, atol=2.5, rtol=0.1
graph_second, second_expected, atol=0.0, rtol=0.0
)
assert is_symmetric_memory_tensor(ordinary_v_output)

pynccl_comm = cuda_communicator.pynccl_comm
assert pynccl_comm is not None
Expand Down
14 changes: 7 additions & 7 deletions vllm/distributed/device_communicators/cuda_communicator.py
View file
Open in desktop

@samnordmann samnordmann Sep 28, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In the PR's description: “does not repeat allocation” should refer specifically to symmetric scratch allocation. clone() still creates fresh result storage. Suggested wording: smth like“Preserves cached symmetric scratch and window registration, adding fresh result storage and one output-sized D2D copy per affected call.”

@jeffnvidia jeffnvidia Sep 30, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Updated the PR description to clarify that cached symmetric scratch and window registration are reused, while each clone() allocates fresh result storage and makes an output-sized device-to-device copy. Thanks for catching the imprecise wording.

Original file line number Diff line number Diff line change
Expand Up @@ -636,11 +636,10 @@ def _get_symm_scratch(
geometric growth bounds their total capacity to less than twice the
current allocation.

Safe across serial MoE layers and eager sequence parallelism: the
producer and collective are ordered on the same stream before the next
same-role operation reuses the buffer. DBO microbatches use distinct
cache entries. Any future cross-layer communication overlap must also
use distinct roles.
These buffers are internal scratch. Returning one would let a later
collective with the same cache key overwrite an earlier live result.
MoE combines may reuse scratch when the producer and collective are
ordered on the same stream. DBO microbatches use distinct cache entries.
"""
from vllm.distributed.device_communicators.pynccl_allocator import (
nccl_symm_mem_context,
Expand Down Expand Up @@ -698,12 +697,13 @@ def _reduce_scatter_symm_mem(
)
symm_input.copy_(input_tensor)

scratch_output = output is None
if output is None:
output = self._get_symm_scratch(
"rs_out", output_shape, input_tensor.dtype, input_tensor.device
)
pynccl_comm.reduce_scatter(output, symm_input)
return output
return output.clone() if scratch_output else output

def get_symmetric_memory_buffer(
self,
Expand Down Expand Up @@ -921,7 +921,7 @@ def _all_gather_symm_mem(self, input_: torch.Tensor) -> torch.Tensor:
"ag_out", out_size, input_.dtype, input_.device
)
pynccl_comm.all_gather(symm_output, input_)
return symm_output
return symm_output.clone()
Comment on lines 923 to +924

@GirasoleY GirasoleY Oct 2, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for catching the issue. However this means we will lose the zero copy behavior.

To allow different use cases, let's add view: bool = False in the argument:

  • return symmem buffer as tensor view when view == true
  • return a clone when view == false

@jeffnvidia jeffnvidia Oct 4, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hi @GirasoleY, thanks.
I agree the clone has a measurable cost. In a two-GPU probe, returning the cached buffer reduced AllGather time from 68 to 36 µs and ReduceScatter from 107 to 88 µs.

My concern with adding view=True to the generic API is that the next same-key collective overwrites any earlier result that is still live. This is a correctness problem, and I don't think it's safe to let this option to users. I haven’t found an in-tree caller with a lifetime guarantee that would make this opt-in safe.

Could we keep this PR focused on the ownership fix and pursue the copy cost separately?
I can work on a separate PR to improve the fix.

@GirasoleY GirasoleY Oct 19, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah ideally we want to either allocate different buffers or be smart about the buffer regions so collective that can race each other do not overlap their symm buffer region.


def _all_gather_batched_symm_mem(
self, inputs: list[torch.Tensor]
Expand Down
Loading