From d8a8bc6cd9371c371e867e6523e4e6af4fc79d5d Mon Sep 17 00:00:00 2001 From: jeffnvidia Date: Tue, 22 Sep 2026 09:43:35 +0300 Subject: [PATCH 1/3] [Bugfix][Distributed] Preserve symmetric collective outputs The symmetric-memory AllGather and ReduceScatter paths returned their cached output scratch directly, so a later same-shape collective could overwrite an earlier live result. Return owned clones while retaining the registered scratch allocations for the collectives. Extend both distributed tests to keep one result live across a same-key call through the companion public entry point, then verify distinct storage and both numerical results. Co-authored-by: OpenAI Codex Signed-off-by: jeffnvidia --- tests/distributed/test_nccl_symm_mem.py | 53 ++++++++++++++----- .../device_communicators/cuda_communicator.py | 14 ++--- 2 files changed, 48 insertions(+), 19 deletions(-) diff --git a/tests/distributed/test_nccl_symm_mem.py b/tests/distributed/test_nccl_symm_mem.py index bd0270fc4849..f31261289b5f 100644 --- a/tests/distributed/test_nccl_symm_mem.py +++ b/tests/distributed/test_nccl_symm_mem.py @@ -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( @@ -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( diff --git a/vllm/distributed/device_communicators/cuda_communicator.py b/vllm/distributed/device_communicators/cuda_communicator.py index fb0c281b4a86..d3c9bc788305 100644 --- a/vllm/distributed/device_communicators/cuda_communicator.py +++ b/vllm/distributed/device_communicators/cuda_communicator.py @@ -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, @@ -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.""" @@ -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() def _all_gather_batched_symm_mem( self, inputs: list[torch.Tensor] From cfaf005b6b5f7477a16f8cbe1ebf972dc4782100 Mon Sep 17 00:00:00 2001 From: jeffnvidia Date: Wed, 30 Sep 2026 11:26:12 +0300 Subject: [PATCH 2/3] Test symmetric collective output ownership under CUDA graphs Capture two same-key AllGather and ReduceScatter calls, then replay with changed inputs and verify both retained results. Co-authored-by: OpenAI Codex Signed-off-by: jeffnvidia --- tests/distributed/test_nccl_symm_mem.py | 36 +++++++++++++++++++++++++ 1 file changed, 36 insertions(+) diff --git a/tests/distributed/test_nccl_symm_mem.py b/tests/distributed/test_nccl_symm_mem.py index f31261289b5f..eba5cbd2cccf 100644 --- a/tests/distributed/test_nccl_symm_mem.py +++ b/tests/distributed/test_nccl_symm_mem.py @@ -154,6 +154,24 @@ def nccl_symm_mem_allgather_worker(local_rank: int, world_size: int): 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( not current_platform.is_cuda(), @@ -233,6 +251,24 @@ def nccl_symm_mem_reduce_scatter_worker(local_rank: int, world_size: int): 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(graph_second, second_expected, atol=0.0, rtol=0.0) + @pytest.mark.skipif( not current_platform.is_cuda(), From 03d6f44f5dcf507de9e6a18cd1df1bb9681097f5 Mon Sep 17 00:00:00 2001 From: jeffnvidia <152798700+jeffnvidia@users.noreply.github.com> Date: Fri, 2 Oct 2026 12:42:01 +0300 Subject: [PATCH 3/3] Fix symmetric collective graph test formatting Assisted-by: OpenAI Codex Signed-off-by: jeffnvidia --- tests/distributed/test_nccl_symm_mem.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/tests/distributed/test_nccl_symm_mem.py b/tests/distributed/test_nccl_symm_mem.py index eba5cbd2cccf..c485d444c85a 100644 --- a/tests/distributed/test_nccl_symm_mem.py +++ b/tests/distributed/test_nccl_symm_mem.py @@ -170,7 +170,9 @@ def nccl_symm_mem_allgather_worker(local_rank: int, world_size: int): 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) + torch.testing.assert_close( + graph_second, second_expected, atol=0.0, rtol=0.0 + ) @pytest.mark.skipif( @@ -267,7 +269,9 @@ def nccl_symm_mem_reduce_scatter_worker(local_rank: int, world_size: int): 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) + torch.testing.assert_close( + graph_second, second_expected, atol=0.0, rtol=0.0 + ) @pytest.mark.skipif(