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

Merged
wzhao18 merged 8 commits into from
Oct 20, 2026
Merged

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

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
8 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
228a9a7
Fix symmetric collective graph test formatting
jeffnvidia Oct 2, 2026
93efe47
Merge upstream main into fix/nccl-symm-rs-output-ownership
jeffnvidia Oct 2, 2026
69556ef
Merge upstream main and preserve symmetric output ownership
jeffnvidia Oct 6, 2026
b547c07
Merge upstream main to include n-gram speculative decoding fix
jeffnvidia Oct 7, 2026
ed5c669
Merge branch 'main' into fix/nccl-symm-rs-output-ownership
wzhao18 Oct 20, 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
53 changes: 41 additions & 12 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
Contributor 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 @@ -130,15 +130,29 @@ 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)


@pytest.mark.skipif(
Expand Down Expand Up @@ -193,16 +207,31 @@ 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
input_tensor = torch.randint(
first_input = torch.randint(
1, 23, (test_size_elements,), dtype=dtype, device=device
)
input_clone = input_tensor.clone()
output = cuda_communicator.reduce_scatter(input_tensor, 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
expected = torch.empty(per_rank_size, dtype=dtype, device=device)
dist.reduce_scatter_tensor(expected, input_clone, group=group)
torch.testing.assert_close(output, expected, atol=2.5, rtol=0.1)
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)

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)


@pytest.mark.skipif(
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
Contributor 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 @@ -559,11 +559,11 @@ def _get_symm_scratch(
call (~0.5 ms/RS+AG pair, dwarfing the NVLS transfer itself). Instead we
allocate once per ``(role, shape, dtype)``, register once, and reuse.

Safe for serial (eager) sequence parallelism: each collective's result
is consumed on the same stream before the next same-role collective
reuses the buffer. Distinct roles (e.g. ``rs_in`` vs ``ag_out``, both
full-size) get distinct buffers so a reduce-scatter input copy never
clobbers a still-live all-gather output.
These buffers are internal scratch. Returning one would let a later
collective with the same cache key overwrite an earlier live result.
Distinct roles (e.g. ``rs_in`` vs ``ag_out``, both full-size) get
distinct buffers so a reduce-scatter input copy never clobbers a
still-live all-gather output.
"""
from vllm.distributed.device_communicators.pynccl_allocator import (
nccl_symm_mem_context,
Expand Down Expand Up @@ -616,7 +616,7 @@ def _reduce_scatter_symm_mem(
symm_input.copy_(input_tensor)

pynccl_comm.reduce_scatter(symm_output, symm_input)
return symm_output
return symm_output.clone()

def send(self, tensor: torch.Tensor, dst: int | None = None) -> None:
"""Sends a tensor to the destination rank in a blocking way."""
Expand Down Expand Up @@ -815,7 +815,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 969 to +970

@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
Contributor 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.

@jeffnvidia jeffnvidia Oct 20, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

Agreed, that sounds like a good longer-term direction. Could we merge this PR as an initial correctness fix and address avoiding the extra copy in a follow-up? That would require ensuring buffers aren't reused while earlier results are still live.


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