Merge branch 'main' into wentao-optimize-dcp-and-add-comm-func

Signed-off-by: yewentao256 <zhyanwentao@126.com>

vllm/distributed/device_communicators/base_device_communicator.py

@@ -181,6 +181,24 @@ class DeviceCommunicatorBase:
         dist.all_reduce(input_, group=self.device_group)
         return input_
 
+    def all_gather_into_tensor(
+        self, output_tensor: torch.Tensor, input_: torch.Tensor
+    ) -> torch.Tensor:
+        input_size = input_.size()
+        expected_output_size = (input_size[0] * self.world_size,) + input_size[1:]
+        assert output_tensor.shape == expected_output_size, (
+            "Invalid output shape for all_gather_into_tensor: "
+            f"expected {expected_output_size}, got {tuple(output_tensor.shape)}"
+        )
+        assert input_.is_contiguous(), (
+            "all_gather_into_tensor requires a contiguous input tensor"
+        )
+        assert output_tensor.is_contiguous(), (
+            "all_gather_into_tensor requires a contiguous output tensor"
+        )
+        dist.all_gather_into_tensor(output_tensor, input_, group=self.device_group)
+        return output_tensor
+
     def all_gather(self, input_: torch.Tensor, dim: int = -1) -> torch.Tensor:
         if dim < 0:
             # Convert negative dim to positive.
@@ -195,7 +213,7 @@ class DeviceCommunicatorBase:
             output_size, dtype=input_.dtype, device=input_.device
         )
         # All-gather.
-        dist.all_gather_into_tensor(output_tensor, input_, group=self.device_group)
+        self.all_gather_into_tensor(output_tensor, input_)
         # Reshape
         output_tensor = output_tensor.reshape((self.world_size,) + input_size)
         output_tensor = output_tensor.movedim(0, dim)
@@ -240,13 +258,32 @@ class DeviceCommunicatorBase:
         )
 
         # Perform reduce-scatter operation
-        torch.distributed.reduce_scatter_tensor(
-            output_tensor, input_tensor, group=self.device_group
-        )
+        self.reduce_scatter_tensor(output_tensor, input_tensor)
 
         # Reshape before returning
         return output_tensor.movedim(0, dim).contiguous()
 
+    def reduce_scatter_tensor(
+        self, output_tensor: torch.Tensor, input_tensor: torch.Tensor
+    ) -> torch.Tensor:
+        expected_input_shape = (output_tensor.shape[0] * self.world_size,) + tuple(
+            output_tensor.shape[1:]
+        )
+        assert input_tensor.shape == expected_input_shape, (
+            "Invalid input shape for reduce_scatter_tensor: "
+            f"expected {expected_input_shape}, got {tuple(input_tensor.shape)}"
+        )
+        assert input_tensor.is_contiguous(), (
+            "reduce_scatter_tensor requires a contiguous input tensor"
+        )
+        assert output_tensor.is_contiguous(), (
+            "reduce_scatter_tensor requires a contiguous output tensor"
+        )
+        torch.distributed.reduce_scatter_tensor(
+            output_tensor, input_tensor, group=self.device_group
+        )
+        return output_tensor
+
     def reduce_scatterv(
         self, input_: torch.Tensor, dim: int = -1, sizes: list[int] | None = None
     ) -> torch.Tensor:

vllm/distributed/device_communicators/cuda_communicator.py

@@ -230,10 +230,24 @@ class CudaCommunicator(DeviceCommunicatorBase):
             torch.distributed.all_reduce(out, group=self.device_group)
         return out
 
+    def all_gather_into_tensor(
+        self, output_tensor: torch.Tensor, input_: torch.Tensor
+    ) -> torch.Tensor:
+        pynccl_comm = self.pynccl_comm
+        assert pynccl_comm is not None and not pynccl_comm.disabled
+        pynccl_comm.all_gather(output_tensor, input_)
+        return output_tensor
+
+    def reduce_scatter_tensor(
+        self, output_tensor: torch.Tensor, input_tensor: torch.Tensor
+    ) -> torch.Tensor:
+        pynccl_comm = self.pynccl_comm
+        assert pynccl_comm is not None and not pynccl_comm.disabled
+        pynccl_comm.reduce_scatter(output_tensor, input_tensor)
+        return output_tensor
+
     def reduce_scatter(self, input_: torch.Tensor, dim: int = -1):
         world_size = self.world_size
-        pynccl_comm = self.pynccl_comm
-        assert pynccl_comm is not None
         if dim < 0:
             # Convert negative dim to positive.
             dim += input_.dim()
@@ -250,7 +264,7 @@ class CudaCommunicator(DeviceCommunicatorBase):
             output_shape, dtype=input_tensor.dtype, device=input_tensor.device
         )
 
-        pynccl_comm.reduce_scatter(output, input_tensor)
+        self.reduce_scatter_tensor(output, input_tensor)
 
         # Reshape before returning
         return output.movedim(0, dim).contiguous()

vllm/distributed/parallel_state.py

@@ -157,6 +157,22 @@ def reduce_scatter_fake(
     return torch.empty(new_shape, dtype=tensor.dtype, device=tensor.device)
 
 
+def reduce_scatter_tensor(
+    output_tensor: torch.Tensor, input_tensor: torch.Tensor, group_name: str
+) -> None:
+    assert group_name in _groups, f"Group {group_name} is not found."
+    group = _groups[group_name]()
+    if group is None:
+        raise ValueError(f"Group {group_name} is destroyed.")
+    group._reduce_scatter_tensor(output_tensor, input_tensor)
+
+
+def reduce_scatter_tensor_fake(
+    output_tensor: torch.Tensor, input_tensor: torch.Tensor, group_name: str
+) -> None:
+    return
+
+
 def all_gather(
     tensor: torch.Tensor, dim: int, world_size: int, group_name: str
 ) -> torch.Tensor:
@@ -175,6 +191,22 @@ def all_gather_fake(
     return torch.empty(new_shape, dtype=tensor.dtype, device=tensor.device)
 
 
+def all_gather_into_tensor(
+    output_tensor: torch.Tensor, input_tensor: torch.Tensor, group_name: str
+) -> None:
+    assert group_name in _groups, f"Group {group_name} is not found."
+    group = _groups[group_name]()
+    if group is None:
+        raise ValueError(f"Group {group_name} is destroyed.")
+    group._all_gather_into_tensor(output_tensor, input_tensor)
+
+
+def all_gather_into_tensor_fake(
+    output_tensor: torch.Tensor, input_tensor: torch.Tensor, group_name: str
+) -> None:
+    return
+
+
 def patched_fused_scaled_matmul_reduce_scatter_fake(
     A: torch.Tensor,
     B: torch.Tensor,
@@ -271,12 +303,26 @@ direct_register_custom_op(
     fake_impl=reduce_scatter_fake,
 )
 
+direct_register_custom_op(
+    op_name="reduce_scatter_tensor",
+    op_func=reduce_scatter_tensor,
+    mutates_args=["output_tensor"],
+    fake_impl=reduce_scatter_tensor_fake,
+)
+
 direct_register_custom_op(
     op_name="all_gather",
     op_func=all_gather,
     fake_impl=all_gather_fake,
 )
 
+direct_register_custom_op(
+    op_name="all_gather_into_tensor",
+    op_func=all_gather_into_tensor,
+    mutates_args=["output_tensor"],
+    fake_impl=all_gather_into_tensor_fake,
+)
+
 # TODO: Remove this once the pytorch fix
 # (https://github.com/pytorch/pytorch/pull/165086) gets released,
 # in either 2.9.1 or 2.10
@@ -528,6 +574,27 @@ class GroupCoordinator:
             raise ValueError("No device communicator found")
         return self.device_communicator.all_reduce(input_)
 
+    def _all_gather_into_tensor(
+        self, output_tensor: torch.Tensor, input_: torch.Tensor
+    ) -> torch.Tensor:
+        if self.device_communicator is None:
+            raise ValueError("No device communicator found")
+        return self.device_communicator.all_gather_into_tensor(output_tensor, input_)
+
+    def all_gather_into_tensor(
+        self, output_tensor: torch.Tensor, input_: torch.Tensor
+    ) -> torch.Tensor:
+        if self.world_size == 1:
+            if output_tensor.data_ptr() != input_.data_ptr():
+                output_tensor.copy_(input_)
+            return output_tensor
+        if self.use_custom_op_call:
+            torch.ops.vllm.all_gather_into_tensor(
+                output_tensor, input_, group_name=self.unique_name
+            )
+            return output_tensor
+        return self._all_gather_into_tensor(output_tensor, input_)
+
     def all_gather(self, input_: torch.Tensor, dim: int = -1) -> torch.Tensor:
         world_size = self.world_size
         # Bypass the function if we are using only 1 GPU.
@@ -582,6 +649,29 @@ class GroupCoordinator:
             raise ValueError("No device communicator found")
         return self.device_communicator.reduce_scatterv(input_, dim, sizes)
 
+    def _reduce_scatter_tensor(
+        self, output_tensor: torch.Tensor, input_tensor: torch.Tensor
+    ) -> torch.Tensor:
+        if self.device_communicator is None:
+            raise ValueError("No device communicator found")
+        return self.device_communicator.reduce_scatter_tensor(
+            output_tensor, input_tensor
+        )
+
+    def reduce_scatter_tensor(
+        self, output_tensor: torch.Tensor, input_tensor: torch.Tensor
+    ) -> torch.Tensor:
+        if self.world_size == 1:
+            if output_tensor.data_ptr() != input_tensor.data_ptr():
+                output_tensor.copy_(input_tensor)
+            return output_tensor
+        if self.use_custom_op_call:
+            torch.ops.vllm.reduce_scatter_tensor(
+                output_tensor, input_tensor, group_name=self.unique_name
+            )
+            return output_tensor
+        return self._reduce_scatter_tensor(output_tensor, input_tensor)
+
     def _reduce_scatter_out_place(self, input_: torch.Tensor, dim: int) -> torch.Tensor:
         if self.device_communicator is None:
             raise ValueError("No device communicator found")

vllm/model_executor/layers/attention/mla_attention.py

@@ -267,7 +267,11 @@ from vllm.v1.attention.backends.utils import (
     get_dcp_local_seq_lens,
     split_decodes_and_prefills,
 )
-from vllm.v1.attention.ops.common import cp_lse_ag_out_rs
+from vllm.v1.attention.ops.common import (
+    cp_all_gather_heads,
+    cp_lse_ag_out_rs,
+    reserve_cp_collective_workspace,
+)
 from vllm.v1.attention.ops.dcp_alltoall import dcp_a2a_lse_reduce
 from vllm.v1.attention.ops.merge_attn_states import merge_attn_states
 from vllm.v1.attention.selector import get_attn_backend
@@ -476,11 +480,35 @@ class MLAAttention(nn.Module, AttentionLayerBase):
         self.use_sparse = use_sparse
 
         _vllm_config = get_current_vllm_config_or_none()
+        dcp_world_size = (
+            1
+            if _vllm_config is None
+            else _vllm_config.parallel_config.decode_context_parallel_size
+        )
         self.dcp_a2a = (
             _vllm_config is not None
-            and _vllm_config.parallel_config.decode_context_parallel_size > 1
+            and dcp_world_size > 1
             and _vllm_config.parallel_config.dcp_comm_backend == "a2a"
         )
+        # reserve collective workspace for dcp
+        if _vllm_config is not None and dcp_world_size > 1:
+            speculative_config = _vllm_config.speculative_config
+            scheduler_config = _vllm_config.scheduler_config
+            model_config = _vllm_config.model_config
+            num_spec_tokens = (
+                speculative_config.num_speculative_tokens
+                if speculative_config is not None
+                else 0
+            )
+            reserve_cp_collective_workspace(
+                max_num_tokens=scheduler_config.max_num_seqs * (1 + num_spec_tokens),
+                total_heads=self.num_heads * dcp_world_size,
+                gather_head_dim=self.kv_lora_rank + self.qk_rope_head_dim,
+                reduce_scatter_head_dim=self.kv_lora_rank,
+                cp_world_size=dcp_world_size,
+                dtype=model_config.dtype,
+                reserve_a2a=self.dcp_a2a,
+            )
 
         # Initialize q/k/v range constants.
         self.q_range = torch.tensor(envs.Q_SCALE_CONSTANT, dtype=torch.float32)
@@ -756,8 +784,8 @@ class MLAAttention(nn.Module, AttentionLayerBase):
                 assert not fp8_attention, "DCP not support fp8 kvcache now."
                 # concatenate mqa_ql_nope and mqa_q_pe -> (B, N, L + P)
                 mqa_q = torch.cat(mqa_q, dim=-1)
-                # mqa_q do allgather in head dim.
-                mqa_q = get_dcp_group().all_gather(mqa_q, dim=1)
+                # mqa_q do allgather in head dim, reuse workspace
+                mqa_q = cp_all_gather_heads(mqa_q, get_dcp_group())
 
             # call decode attn
             if not is_sparse_impl:
@@ -2151,7 +2179,7 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
         q: torch.Tensor,
         kv_c_and_k_pe_cache: torch.Tensor,
         attn_metadata: MLACommonMetadata,
-        k_scale: torch.Tensor,
+        k_scale: torch.Tensor | None,
         dcp_world_size: int,
     ):
         assert k_scale is None, "DCP not support scaled kvcache now."
@@ -2193,8 +2221,8 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
             ]
             assert toks * dcp_world_size <= cur_allgather_workspace.shape[0]
             cur_allgather_kvcache = cur_allgather_workspace[: toks * dcp_world_size]
-            cur_allgather_kvcache.copy_(
-                get_dcp_group().all_gather(local_gathered_kvcache, dim=0)
+            get_dcp_group().all_gather_into_tensor(
+                cur_allgather_kvcache, local_gathered_kvcache
             )
             assert (
                 cur_allgather_kvcache.shape[-1]

vllm/v1/attention/backends/flash_attn.py

@@ -29,7 +29,12 @@ from vllm.v1.attention.backends.fa_utils import (
     is_flash_attn_varlen_func_available,
 )
 from vllm.v1.attention.backends.utils import get_dcp_local_seq_lens
-from vllm.v1.attention.ops.common import cp_lse_ag_out_rs
+from vllm.v1.attention.ops.common import (
+    cp_all_gather_heads,
+    cp_collective_scratch_bytes,
+    cp_lse_ag_out_rs,
+    reserve_cp_collective_workspace,
+)
 from vllm.v1.attention.ops.dcp_alltoall import dcp_a2a_lse_reduce
 from vllm.v1.attention.ops.merge_attn_states import merge_attn_states
 from vllm.v1.worker.workspace import current_workspace_manager
@@ -673,6 +678,45 @@ class FlashAttentionImpl(AttentionImpl):
         self._dcp_dtype: torch.dtype | None = None
         if vllm_config is not None and self.dcp_world_size > 1:
             self._dcp_dtype = vllm_config.model_config.dtype
+            reserve_cp_collective_workspace(
+                max_num_tokens=vllm_config.scheduler_config.max_num_batched_tokens,
+                total_heads=self.num_heads * self.dcp_world_size,
+                gather_head_dim=self.head_size,
+                reduce_scatter_head_dim=self.head_size,
+                cp_world_size=self.dcp_world_size,
+                dtype=self._dcp_dtype,
+                reserve_a2a=dcp_a2a,
+            )
+
+    def _get_dcp_combine_workspace_bytes(self, num_tokens: int) -> int:
+        assert self._dcp_dtype is not None
+        ws, n, h, hd, d = (
+            self.dcp_world_size,
+            num_tokens,
+            self.num_heads,
+            self.head_size,
+            self._dcp_dtype,
+        )
+        if self.dcp_combine is dcp_a2a_lse_reduce:
+            t = (ws, n, h, hd)
+            return cp_collective_scratch_bytes(
+                (t, d),  # send_output
+                (t, d),  # recv_output
+                ((ws, n, h), torch.float32),  # send_lse
+                ((ws, n, h), torch.float32),  # recv_lse
+            )
+        return max(
+            # buffer for all_gather_into_tensor
+            # [ws * n, total_heads]
+            cp_collective_scratch_bytes(
+                ((ws * n, h * ws), torch.float32),
+            ),
+            # output buffer for reduce scatter
+            cp_collective_scratch_bytes(
+                (((h * ws), n, hd), d),
+                ((h, n, hd), d),
+            ),
+        )
 
     def forward(
         self,
@@ -917,16 +961,19 @@ class FlashAttentionImpl(AttentionImpl):
         block_table = attn_metadata.block_table
 
         query = query.contiguous()
-        query_across_dcp = get_dcp_group().all_gather(query, dim=1)
+        query_across_dcp = cp_all_gather_heads(query, get_dcp_group())
         sliding_window_size = (
             list(self.sliding_window) if self.sliding_window is not None else None
         )
         n = query_across_dcp.shape[0]
-        (dcp_context_out,) = current_workspace_manager().get_simultaneous(
-            (
-                (n, self.num_heads * self.dcp_world_size, self.head_size),
-                self._dcp_dtype,
-            ),
+        dcp_context_out, dcp_combine_workspace = (
+            current_workspace_manager().get_simultaneous(
+                (
+                    (n, self.num_heads * self.dcp_world_size, self.head_size),
+                    self._dcp_dtype,
+                ),
+                ((self._get_dcp_combine_workspace_bytes(n),), torch.uint8),
+            )
         )
         context_attn_out, context_lse = flash_attn_varlen_func(
             q=query_across_dcp,
@@ -957,6 +1004,7 @@ class FlashAttentionImpl(AttentionImpl):
             context_lse.transpose(0, 1),
             get_dcp_group(),
             return_lse=True,
+            scratch_workspace=dcp_combine_workspace,
         )
         context_lse_cor = context_lse_cor.transpose(0, 1).contiguous()
 

vllm/v1/attention/backends/flashinfer.py

@@ -66,7 +66,11 @@ from vllm.v1.attention.backends.utils import (
     infer_global_hyperparameters,
     split_decodes_and_prefills,
 )
-from vllm.v1.attention.ops.common import cp_lse_ag_out_rs
+from vllm.v1.attention.ops.common import (
+    cp_all_gather_heads,
+    cp_lse_ag_out_rs,
+    reserve_cp_collective_workspace,
+)
 from vllm.v1.attention.ops.dcp_alltoall import dcp_a2a_lse_reduce
 from vllm.v1.attention.ops.merge_attn_states import merge_attn_states
 from vllm.v1.kv_cache_interface import (
@@ -288,8 +292,8 @@ class BatchDCPPrefillWrapper:
         value: torch.Tensor,
         out: torch.Tensor,
     ):
-        prefill_query_across_dcp = get_dcp_group().all_gather(
-            prefill_query.contiguous(), dim=1
+        prefill_query_across_dcp = cp_all_gather_heads(
+            prefill_query.contiguous(), get_dcp_group()
         )
         output_context_tmp, lse_context_tmp = self._context.run(
             prefill_query_across_dcp,
@@ -1345,6 +1349,17 @@ class FlashInferImpl(AttentionImpl):
         else:
             self.dcp_combine = partial(cp_lse_ag_out_rs, is_lse_base_on_e=False)
 
+        if vllm_config is not None and self.dcp_world_size > 1:
+            reserve_cp_collective_workspace(
+                max_num_tokens=vllm_config.scheduler_config.max_num_batched_tokens,
+                total_heads=self.num_heads * self.dcp_world_size,
+                gather_head_dim=self.head_size,
+                reduce_scatter_head_dim=self.head_size,
+                cp_world_size=self.dcp_world_size,
+                dtype=vllm_config.model_config.dtype,
+                reserve_a2a=dcp_a2a,
+            )
+
     def fused_output_quant_supported(self, quant_key: QuantKey):
         return (
             self.support_trtllm_attn
@@ -1710,8 +1725,8 @@ class FlashInferImpl(AttentionImpl):
                     out_decode = output[:num_decode_tokens]
 
                 if use_dcp:
-                    decode_query = get_dcp_group().all_gather(
-                        decode_query.contiguous(), dim=-2
+                    decode_query = cp_all_gather_heads(
+                        decode_query.contiguous(), get_dcp_group()
                     )
                     output_tmp = torch.empty_like(decode_query)
                     lse = torch.empty(

vllm/v1/attention/ops/common.py

@@ -1,9 +1,16 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from math import prod
+
 import torch
 
 from vllm.distributed.parallel_state import GroupCoordinator
 from vllm.triton_utils import tl, triton
+from vllm.utils.math_utils import round_up
+from vllm.v1.worker.workspace import (
+    current_workspace_manager,
+    is_workspace_manager_initialized,
+)
 
 
 @triton.jit
@@ -107,6 +114,154 @@ class CPTritonContext:
             self.inner_kernel[grid](*regular_args)
 
 
+def cp_collective_scratch_bytes(
+    *shapes_and_dtypes: tuple[tuple[int, ...], torch.dtype],
+) -> int:
+    return sum(
+        round_up(prod(shape) * dtype.itemsize, 256)
+        for shape, dtype in shapes_and_dtypes
+    )
+
+
+def get_cp_collective_scratch_tensors(
+    device: torch.device,
+    *shapes_and_dtypes: tuple[tuple[int, ...], torch.dtype],
+    scratch_workspace: torch.Tensor | None = None,
+) -> list[torch.Tensor]:
+    if scratch_workspace is not None:
+        offset = 0
+        tensors: list[torch.Tensor] = []
+        for shape, dtype in shapes_and_dtypes:
+            actual_bytes = prod(shape) * dtype.itemsize
+            aligned_bytes = round_up(actual_bytes, 256)
+            tensors.append(
+                scratch_workspace[offset : offset + actual_bytes]
+                .view(dtype)
+                .reshape(shape)
+            )
+            offset += aligned_bytes
+        return tensors
+    # get workspace from workspace manager
+    if is_workspace_manager_initialized():
+        workspace_manager = current_workspace_manager()
+        try:
+            return workspace_manager.get_simultaneous(*shapes_and_dtypes)
+        except AssertionError:
+            if not workspace_manager.is_locked():
+                raise
+    return [
+        torch.empty(shape, dtype=dtype, device=device)
+        for shape, dtype in shapes_and_dtypes
+    ]
+
+
+def reserve_cp_collective_workspace(
+    max_num_tokens: int,
+    total_heads: int,
+    gather_head_dim: int,
+    reduce_scatter_head_dim: int,
+    cp_world_size: int,
+    dtype: torch.dtype,
+    lse_dtype: torch.dtype = torch.float32,
+    reserve_a2a: bool = False,
+) -> None:
+    # reserve workspace from workspace manager, call before allgather/reduce_scatter
+    if (
+        cp_world_size <= 1
+        or max_num_tokens <= 0
+        or not is_workspace_manager_initialized()
+    ):
+        return
+
+    assert total_heads % cp_world_size == 0
+    local_heads = total_heads // cp_world_size
+    workspace_manager = current_workspace_manager()
+    workspace_manager.get_simultaneous(
+        ((max_num_tokens * cp_world_size, local_heads, gather_head_dim), dtype),
+    )
+    workspace_manager.get_simultaneous(
+        ((total_heads, max_num_tokens, reduce_scatter_head_dim), dtype),
+        ((local_heads, max_num_tokens, reduce_scatter_head_dim), dtype),
+    )
+    workspace_manager.get_simultaneous(
+        ((cp_world_size * max_num_tokens, total_heads), lse_dtype),
+    )
+    if reserve_a2a:
+        workspace_manager.get_simultaneous(
+            (
+                (cp_world_size, max_num_tokens, local_heads, reduce_scatter_head_dim),
+                dtype,
+            ),
+            (
+                (cp_world_size, max_num_tokens, local_heads, reduce_scatter_head_dim),
+                dtype,
+            ),
+            ((cp_world_size, max_num_tokens, local_heads), lse_dtype),
+            ((cp_world_size, max_num_tokens, local_heads), lse_dtype),
+        )
+
+
+def cp_all_gather_heads(
+    cp_attn_in: torch.Tensor,
+    cp_group: GroupCoordinator,
+) -> torch.Tensor:
+    """All-gather a [B, H_local, D] tensor across ranks on the head axis."""
+    if cp_group.world_size == 1:
+        return cp_attn_in
+
+    cp_attn_in = cp_attn_in.contiguous()
+    batch_size, local_heads, head_dim = cp_attn_in.shape
+    world_size = cp_group.world_size
+
+    (gathered,) = get_cp_collective_scratch_tensors(
+        cp_attn_in.device,
+        ((batch_size * world_size, local_heads, head_dim), cp_attn_in.dtype),
+    )
+    cp_group.all_gather_into_tensor(gathered, cp_attn_in)
+
+    out = torch.empty(
+        (batch_size, local_heads * world_size, head_dim),
+        dtype=cp_attn_in.dtype,
+        device=cp_attn_in.device,
+    )
+    out.view(batch_size, world_size, local_heads, head_dim).copy_(
+        gathered.view(world_size, batch_size, local_heads, head_dim).movedim(0, 1)
+    )
+    return out
+
+
+def cp_reduce_scatter_heads(
+    cp_attn_out: torch.Tensor,
+    cp_group: GroupCoordinator,
+    scratch_workspace: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Reduce-scatter a [B, H_total, D] tensor across ranks on the head axis."""
+    if cp_group.world_size == 1:
+        return cp_attn_out
+
+    batch_size, total_heads, head_dim = cp_attn_out.shape
+    world_size = cp_group.world_size
+    assert total_heads % world_size == 0
+    local_heads = total_heads // world_size
+
+    rs_input, rs_output = get_cp_collective_scratch_tensors(
+        cp_attn_out.device,
+        ((total_heads, batch_size, head_dim), cp_attn_out.dtype),
+        ((local_heads, batch_size, head_dim), cp_attn_out.dtype),
+        scratch_workspace=scratch_workspace,
+    )
+    rs_input.copy_(cp_attn_out.movedim(0, 1))
+    cp_group.reduce_scatter_tensor(rs_output, rs_input)
+
+    out = torch.empty(
+        (batch_size, local_heads, head_dim),
+        dtype=cp_attn_out.dtype,
+        device=cp_attn_out.device,
+    )
+    out.copy_(rs_output.movedim(0, 1))
+    return out
+
+
 def correct_attn_out(
     out: torch.Tensor,
     lses: torch.Tensor,
@@ -184,6 +339,7 @@ def _cp_lse_common(
     cp_group: GroupCoordinator,
     ctx: CPTritonContext | None = None,
     is_lse_base_on_e=True,
+    scratch_workspace: torch.Tensor | None = None,
 ):
     """
     cp_attn_out: [ B, H, D ]
@@ -196,9 +352,14 @@ def _cp_lse_common(
         ctx = CPTritonContext()
 
     cp_attn_lse = cp_attn_lse.contiguous()
-    lses = cp_group.all_gather(cp_attn_lse, dim=0).reshape(
-        (cp_group.world_size,) + cp_attn_lse.shape
+    batch_size, num_heads = cp_attn_lse.shape
+    (lses_flat,) = get_cp_collective_scratch_tensors(
+        cp_attn_lse.device,
+        ((cp_group.world_size * batch_size, num_heads), cp_attn_lse.dtype),
+        scratch_workspace=scratch_workspace,
     )
+    cp_group.all_gather_into_tensor(lses_flat, cp_attn_lse)
+    lses = lses_flat.view((cp_group.world_size,) + cp_attn_lse.shape)
     out, lse = correct_attn_out(
         cp_attn_out,
         lses,
@@ -216,15 +377,21 @@ def cp_lse_ag_out_rs(
     ctx: CPTritonContext | None = None,
     return_lse: bool = False,
     is_lse_base_on_e=True,
+    scratch_workspace: torch.Tensor | None = None,
 ):
     """
     cp_attn_out: [ B, H, D ]
     cp_attn_lse: [ B, H ]
     """
     out, lse = _cp_lse_common(
-        cp_attn_out, cp_attn_lse, cp_group, ctx=ctx, is_lse_base_on_e=is_lse_base_on_e
+        cp_attn_out,
+        cp_attn_lse,
+        cp_group,
+        ctx=ctx,
+        is_lse_base_on_e=is_lse_base_on_e,
+        scratch_workspace=scratch_workspace,
     )
-    out = cp_group.reduce_scatter(out, dim=1)
+    out = cp_reduce_scatter_heads(out, cp_group, scratch_workspace=scratch_workspace)
 
     if return_lse:
         cp_num_heads = lse.shape[1] // cp_group.world_size

vllm/v1/attention/ops/dcp_alltoall.py

@@ -26,10 +26,7 @@ import torch
 import torch.distributed as dist
 
 from vllm.triton_utils import tl, triton
-from vllm.v1.worker.workspace import (
-    current_workspace_manager,
-    is_workspace_manager_initialized,
-)
+from vllm.v1.attention.ops.common import get_cp_collective_scratch_tensors
 
 if TYPE_CHECKING:
     from vllm.distributed.parallel_state import GroupCoordinator
@@ -116,18 +113,15 @@ def _dcp_a2a_send_recv_buffers(
     shape: tuple[int, ...],
     device: torch.device,
     dtype: torch.dtype,
+    scratch_workspace: torch.Tensor | None = None,
 ) -> tuple[torch.Tensor, torch.Tensor]:
-    if is_workspace_manager_initialized():
-        send_buffer, recv_buffer = current_workspace_manager().get_simultaneous(
-            (shape, dtype),
-            (shape, dtype),
-        )
-        return send_buffer, recv_buffer
-
-    return (
-        torch.empty(shape, device=device, dtype=dtype),
-        torch.empty(shape, device=device, dtype=dtype),
+    send_buffer, recv_buffer = get_cp_collective_scratch_tensors(
+        device,
+        (shape, dtype),
+        (shape, dtype),
+        scratch_workspace=scratch_workspace,
     )
+    return send_buffer, recv_buffer
 
 
 @triton.jit
@@ -397,6 +391,7 @@ def dcp_a2a_lse_reduce(
     ctx: CPTritonContext | None = None,
     return_lse: bool = False,
     is_lse_base_on_e: bool = True,
+    scratch_workspace: torch.Tensor | None = None,
 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
     """
     Combine partial attention outputs across DCP ranks using All-to-All.
@@ -433,6 +428,7 @@ def dcp_a2a_lse_reduce(
         (world_size, B, H_per_rank, D + lse_pack_dim),
         device=cp_attn_out.device,
         dtype=cp_attn_out.dtype,
+        scratch_workspace=scratch_workspace,
     )
 
     _dcp_a2a_pack_send(

vllm/v1/worker/workspace.py

@@ -90,7 +90,8 @@ class WorkspaceManager:
         return self._locked
 
     def get_simultaneous(
-        self, *shapes_and_dtypes: tuple[tuple[int, ...], torch.dtype]
+        self,
+        *shapes_and_dtypes: tuple[tuple[int, ...], torch.dtype],
     ) -> list[torch.Tensor]:
         """Get multiple workspace tensors simultaneously from a single allocation.
 
@@ -153,7 +154,8 @@ class WorkspaceManager:
                     )
                 return "unknown"
 
-            if self._locked:
+            can_initialize_locked_ubatch = current_size == 0
+            if self._locked and not can_initialize_locked_ubatch:
                 raise AssertionError(
                     f"Workspace is locked but allocation from '{get_caller_info()}' "
                     f"requires {required_bytes / _MB:.2f} MB, current size is "
@@ -161,17 +163,15 @@ class WorkspaceManager:
                     "Workspace growth is not allowed after locking."
                 )
 
-            # Only resize the requesting ubatch's workspace.  Other
-            # ubatches resize lazily on their next get_simultaneous call.
-            # Resizing all ubatches here would orphan the other ubatch's
-            # old tensor when it still holds views into it (DBO leak).
+            # Only resize the requesting ubatch's workspace. Other ubatches
+            # resize lazily on their next get_simultaneous call. Resizing all
+            # ubatches here would orphan another ubatch's old tensor while it
+            # still holds views into it (DBO leak).
             self._current_workspaces[ubatch_id] = None
             del current_workspace
-            # Release the freed segment back to CUDA so the caching
-            # allocator can reuse the GPU memory for the larger
-            # allocation below. Without this, each resize may leave a
-            # dead segment in reserved memory which can cause higher peak
-            # memory usage.
+            # Release the freed segment back to the accelerator allocator so
+            # the larger allocation below can reuse the memory instead of
+            # leaving dead reserved segments behind.
             torch.accelerator.empty_cache()
             self._current_workspaces[ubatch_id] = torch.empty(
                 (required_bytes,), dtype=torch.uint8, device=self._device
@@ -187,7 +187,6 @@ class WorkspaceManager:
                     required_bytes / _MB,
                     ubatch_id,
                 )
-
         return current_workspace