diff --git a/docs/en/advance/context_parallel.md b/docs/en/advance/context_parallel.md
index a474ca74b9..2cb741bf5c 100644
--- a/docs/en/advance/context_parallel.md
+++ b/docs/en/advance/context_parallel.md
@@ -34,12 +34,16 @@ lmdeploy serve api_server Qwen/Qwen3-235B-A22B --tp 8 --cp 2
 
 ## PyTorch decode context parallelism
 
-PyTorch DCP distributes MLA KV cache across existing TP ranks, increasing
+PyTorch DCP distributes replicated KV cache across existing TP ranks, increasing
 effective cache capacity without additional GPUs. It supports FlashMLA-backed
-dense MLA and sparse DSA models.
+dense MLA, sparse DSA, and GQA with replicated KV heads.
 
 ```bash
 lmdeploy serve api_server <mla-model> --backend pytorch --tp 4 --dcp 2
+
+# Both Qwen3 models have four KV heads; TP=8 replicates each head twice.
+lmdeploy serve api_server Qwen/Qwen3-30B-A3B --backend pytorch --tp 8 --dcp 2
+lmdeploy serve api_server Qwen/Qwen3-235B-A22B --backend pytorch --tp 8 --dcp 2
 ```
 
 For Python, use `PytorchEngineConfig(tp=4, dcp=2)`. The default `dcp=1`
@@ -47,13 +51,21 @@ disables DCP.
 
 ### Requirements and supported features
 
-- Requires NVIDIA Hopper/SM90 GPUs, FlashMLA, and BF16 activations.
+- MLA requires NVIDIA Hopper/SM90 GPUs, FlashMLA, and BF16 activations.
   Sparse DSA also requires compatible DeepGEMM and TileLang top-k kernels
   (top-k 512 or 2048).
+- GQA uses Triton for decode and FlashAttention-3, when available, for prefill,
+  with a Triton prefill fallback. It supports BF16/FP16 activations and
+  unquantized or per-tensor FP8 KV cache (`--quant-policy fp8` for E4M3,
+  `--quant-policy fp8_e5m2` for E5M2). Cached-prefix chunks are dequantized
+  before gathering. INT8, INT4, TurboQuant, ALiBi, and attention sinks are not supported.
 - `dcp` must divide the attention TP size and replicated KV-head count;
   `dp=1` and `ep=1` are required.
 - Prefix caching and `deepseek_mtp` are supported for compatible MLA models,
   including DeepSeek V3/V3.1, DeepSeek V3.2, and GLM DSA.
+- GQA prefill retains local query heads and gathers cached K/V in bounded
+  chunks. Decode gathers query heads within each replicated-KV group and
+  merges shard outputs using their softmax normalization factors.
 - BF16 KV cache is supported. Sparse MLA also supports FP8 KV cache via
   `--quant-policy fp8`.
 - Sliding-window attention, MemDecode, prefill/decode disaggregation,
diff --git a/lmdeploy/cli/cli.py b/lmdeploy/cli/cli.py
index bf2d1e9a22..10d4fb1504 100644
--- a/lmdeploy/cli/cli.py
+++ b/lmdeploy/cli/cli.py
@@ -58,6 +58,7 @@ def add_parser_chat():
         ArgumentHelper.prefix_cache_decode_state_interval(pt_group)
         ArgumentHelper.prefill_save_alignment(pt_group)
         ArgumentHelper.mooncake_state_save_slots(pt_group)
+        ArgumentHelper.communication_backend(pt_group)
         # common engine args
         dtype_act = ArgumentHelper.dtype(pt_group)
         tp_act = ArgumentHelper.tp(pt_group)
diff --git a/lmdeploy/cli/serve.py b/lmdeploy/cli/serve.py
index 6af5441658..4b878d7cdd 100644
--- a/lmdeploy/cli/serve.py
+++ b/lmdeploy/cli/serve.py
@@ -147,6 +147,7 @@ def add_parser_api_server():
         ArgumentHelper.migration_backend(pt_group)
         ArgumentHelper.cudagraph_capture_batch_sizes(pt_group)
         ArgumentHelper.piecewise_cudagraph_max_tokens(pt_group)
+        ArgumentHelper.communication_backend(pt_group)
         # multi-node serving args
         node_rank_act = ArgumentHelper.node_rank(pt_group)
         num_nodes_act = ArgumentHelper.num_nodes(pt_group)
@@ -291,6 +292,7 @@ def api_server(args):
                 enable_return_routed_experts=args.enable_return_routed_experts,
                 distributed_executor_backend=args.distributed_executor_backend,
                 kv_transfer_config=kv_transfer_config,
+                communication_backend=args.communication_backend,
             )
         else:
             from lmdeploy.messages import TurbomindEngineConfig
diff --git a/lmdeploy/cli/utils.py b/lmdeploy/cli/utils.py
index c026811191..cacc29eba7 100644
--- a/lmdeploy/cli/utils.py
+++ b/lmdeploy/cli/utils.py
@@ -918,6 +918,14 @@ def trust_remote_code(parser):
                                    default=False,
                                    help='Whether to trust remote code from model repositories.')
 
+    @staticmethod
+    def communication_backend(parser):
+        """Select native or automatic PyTorch communication."""
+        return parser.add_argument(
+            '--communication-backend', choices=['nccl', 'auto'], default='nccl',
+            help='PyTorch communication backend. nccl uses native process-group collectives; '
+            'auto opts into eligible optimizations. Default: nccl.')
+
     @staticmethod
     def kv_transfer_config(parser):
         """Add external KV-cache connector configuration."""
diff --git a/lmdeploy/messages.py b/lmdeploy/messages.py
index 3f22316ed7..c3f9f119cb 100644
--- a/lmdeploy/messages.py
+++ b/lmdeploy/messages.py
@@ -548,6 +548,8 @@ class PytorchEngineConfig:
         dllm_denoising_steps: Dllm denoising steps.
         dllm_confidence_threshold: dllm unmasking threshold for
             dynamic unmasking.
+        communication_backend: ``nccl`` (default) uses native process-group collectives.
+            ``auto`` opts into eligible communication optimizations.
         kv_transfer_config: External KV-cache connector configuration. This is
             supported only by the PyTorch engine. ``None`` disables external
             KV-cache transfer.
@@ -609,6 +611,7 @@ class PytorchEngineConfig:
     migration_backend: MigrationBackend = MigrationBackend.DLSlime
     kv_transfer_config: KVTransferConfig | dict[str, Any] | None = None
     piecewise_cudagraph_max_tokens: int | None = None
+    communication_backend: Literal['nccl', 'auto'] = 'nccl'
 
     def __post_init__(self):
         """Check input validation."""
@@ -663,7 +666,8 @@ def __post_init__(self):
             self.kv_transfer_config = KVTransferConfig(**self.kv_transfer_config)
         elif self.kv_transfer_config is not None and not isinstance(self.kv_transfer_config, KVTransferConfig):
             raise TypeError('kv_transfer_config must be a KVTransferConfig, dict, or None')
-
+        if self.communication_backend not in ('nccl', 'auto'):
+            raise ValueError('communication_backend must be nccl or auto')
 
 class ResponseType(enum.Enum):
     """Response type."""
diff --git a/lmdeploy/pytorch/backends/base.py b/lmdeploy/pytorch/backends/base.py
index d03ce96ecc..c435d72814 100644
--- a/lmdeploy/pytorch/backends/base.py
+++ b/lmdeploy/pytorch/backends/base.py
@@ -92,13 +92,14 @@ def model_build_context(ctx_mgr):
         yield
 
     @classmethod
-    def build_communicator(cls, cpu_group, device_group, dist_config):
+    def build_communicator(cls, cpu_group, device_group, dist_config, *, group_name='tp'):
         """Build a device communicator."""
         from .communicator import build_communicator
         return build_communicator(
             cpu_group=cpu_group,
             device_group=device_group,
             dist_config=dist_config,
+            group_name=group_name,
         )
 
     @staticmethod
diff --git a/lmdeploy/pytorch/backends/communicator.py b/lmdeploy/pytorch/backends/communicator.py
index 418458871d..c17e01d451 100644
--- a/lmdeploy/pytorch/backends/communicator.py
+++ b/lmdeploy/pytorch/backends/communicator.py
@@ -1,40 +1,57 @@
 # Copyright (c) OpenMMLab. All rights reserved.
+from typing import Literal
+
 import torch
 from torch import distributed as dist
 
 
 class DeviceCommunicator:
-    """Device-neutral collective operations for one process group."""
+    """Device-neutral collective operations for one process group.
+
+    Workspace creation returns a prepared, caller-owned workspace when available. Collective methods execute with native
+    fallback when an optimized implementation cannot handle the input.
+    """
 
     def __init__(self, device_group: dist.ProcessGroup):
         self.device_group = device_group
 
-    def supports_optimized_all_reduce(self) -> bool:
-        """Whether an optimized all-reduce implementation is available."""
-        return False
-
-    def supports_fused_all_reduce_residual_rms_norm(self) -> bool:
-        """Whether fused all-reduce, residual and RMSNorm is available."""
-        return False
-
-    def try_fused_all_reduce_residual_rms_norm(self,
-                                               input: torch.Tensor,
-                                               residual: torch.Tensor,
-                                               weight: torch.Tensor,
-                                               eps: float):
-        """Run fused all-reduce, residual and RMSNorm when eligible."""
-        return None
-
     def all_reduce_(self, input: torch.Tensor):
         """All-reduce ``input`` in place."""
         dist.all_reduce(input, group=self.device_group)
 
+    def all_gather(self, input: torch.Tensor, *, dim: int = -1,
+                   workspace=None, copy_output: bool = True) -> torch.Tensor:
+        """Gather 2D inputs along dimension 0 or -1 on every group rank.
+
+        With copy_output=False, optimized output may borrow the workspace until its next use. Consume it on the same
+        stream. Native gathering always returns an owned output for multi-rank groups.
+        """
+        if input.dim() != 2 or dim not in (0, -1):
+            raise ValueError('all_gather requires a 2D input and dim=0 or dim=-1')
+        world_size = dist.get_world_size(self.device_group)
+        if world_size == 1:
+            return input
+        input = input.contiguous()
+        gathered = input.new_empty(world_size * input.size(0), input.size(1))
+        dist.all_gather_into_tensor(gathered, input, group=self.device_group)
+        if dim == 0:
+            return gathered
+        return gathered.view(world_size, *input.shape).transpose(0, 1).reshape(input.size(0), -1)
+
+    def create_all_gather_workspace(self, gathered_width: int, device: torch.device, dtype: torch.dtype,
+                                    *, dim: int = -1, reuse_sync: bool = True, capacity_bytes: int = 64 * 1024 * 1024):
+        """Create a caller-owned workspace.
+
+        Disable reuse_sync only when a subsequent same-group collective orders all consumers before the next gather.
+        """
+        return None
+
     def close(self):
-        """Release communicator-owned resources."""
-        pass
+        """Release group-owned resources; gather workspaces belong to
+        callers."""
 
 
 def build_communicator(cpu_group: dist.ProcessGroup, device_group: dist.ProcessGroup,
-                       dist_config):
+                       dist_config, *, group_name: Literal['tp', 'dcp'] = 'tp'):
     """Build the default process-group communicator."""
     return DeviceCommunicator(device_group=device_group)
diff --git a/lmdeploy/pytorch/backends/cp_utils.py b/lmdeploy/pytorch/backends/cp_utils.py
index 4343290097..85f0d21b6f 100644
--- a/lmdeploy/pytorch/backends/cp_utils.py
+++ b/lmdeploy/pytorch/backends/cp_utils.py
@@ -10,29 +10,33 @@ def get_dcp_topk_workspace_size(num_rows: int, topk: int, dcp_size: int) -> int:
     return num_rows * topk * (16 + 8 * dcp_size)
 
 
-def get_dcp_prefill_workspace_size(*, batch_size: int, head_dim: int, block_size: int, dcp_size: int) -> int:
+
+def get_dcp_prefill_workspace_size(*, batch_size: int, kv_width: int, block_size: int, dcp_size: int) -> int:
     """Per-rank prefix KV budget shared by chunk planning and cache sizing.
 
     Target 64 MiB, enlarged if one virtual block per request needs more. FP8 caches also gather dequantized BF16 KV.
     """
     # Simultaneous BF16 buffers: one local block, D gathered, D reordered.
-    minimum = batch_size * block_size * (1 + 2 * dcp_size) * head_dim * 2
+    minimum = batch_size * block_size * (1 + 2 * dcp_size) * kv_width * 2
     return max(64 << 20, minimum)
 
 
 def get_dcp_workspace_size(*, num_prefill_tokens: int, num_decode_tokens: int,
                            batch_size: int, num_heads: int, head_dim: int,
                            dtype: torch.dtype, block_size: int, dcp_size: int,
-                           topk: int | None, score_workspace_bytes: int) -> int:
+                           topk: int | None, score_workspace_bytes: int,
+                           kv_width: int) -> int:
     """Estimate per-rank peak temporary bytes for enabled DCP.
 
     ``num_heads`` is the rank-local query-head count before DCP gathering. Reserve max(indexer, attention) plus final
-    indices shared by both phases. Flattened indexer KV relies on memory headroom.
+    indices shared by both phases.
+    ``kv_width`` counts both K and V elements per cached token for GQA, or the shared
+    latent width for MLA. Flattened indexer KV relies on memory headroom.
     """
     # Prefix KV: local, gathered, and reordered buffers.
     attention_workspace = get_dcp_prefill_workspace_size(
         batch_size=batch_size,
-        head_dim=head_dim,
+        kv_width=kv_width,
         block_size=block_size,
         dcp_size=dcp_size,
     )
@@ -67,8 +71,19 @@ class DCPPrefixChunk:
     local_kv_seqlens: torch.Tensor  # [dcp_size, requests]
     local_cu_seqlens: torch.Tensor  # current rank
 
+    @property
+    def local_capacity(self) -> int:
+        """Padded token capacity per rank, identical across the DCP group."""
+        return self.kv_seqlens.numel() * self.size // self.local_kv_seqlens.size(0)
+
+    def local_block_slice(self, block_size: int) -> slice:
+        """Select local cache blocks for this globally aligned chunk."""
+        virtual_block_size = block_size * self.local_kv_seqlens.size(0)
+        start = self.start // virtual_block_size
+        return slice(start, start + self.size // virtual_block_size)
+
 
-def build_dcp_prefix_chunks(*, prefix_lens: torch.Tensor, prefix_limit: int, block_size: int, head_dim: int,
+def build_dcp_prefix_chunks(*, prefix_lens: torch.Tensor, prefix_limit: int, block_size: int, kv_width: int,
                              dcp_world_rank: tuple[int, int]) -> tuple[DCPPrefixChunk, ...]:
     """Plan identical collective shapes on all ranks without device sync."""
     if prefix_limit <= 0:
@@ -76,11 +91,11 @@ def build_dcp_prefix_chunks(*, prefix_lens: torch.Tensor, prefix_limit: int, blo
     dcp_size, dcp_rank = dcp_world_rank
     batch_size = prefix_lens.numel()
     workspace_bytes = get_dcp_prefill_workspace_size(batch_size=batch_size,
-                                                     head_dim=head_dim,
+                                                     kv_width=kv_width,
                                                      block_size=block_size,
                                                      dcp_size=dcp_size)
     virtual_block_size = block_size * dcp_size
-    bytes_per_block = batch_size * block_size * (1 + 2 * dcp_size) * head_dim * 2
+    bytes_per_block = batch_size * block_size * (1 + 2 * dcp_size) * kv_width * 2
     # Longer prefixes add chunks rather than enlarge the gather buffers.
     chunk_size = (workspace_bytes // bytes_per_block) * virtual_block_size
     ranks = torch.arange(dcp_size, device=prefix_lens.device, dtype=prefix_lens.dtype)[:, None]
@@ -97,6 +112,14 @@ def build_dcp_prefix_chunks(*, prefix_lens: torch.Tensor, prefix_limit: int, blo
     return tuple(chunks)
 
 
+def get_dcp_kv_width(model_config) -> int:
+    """Elements per cached token: shared MLA latent or combined rank-local GQA K/V."""
+    if model_config.use_flash_mla:
+        return model_config.head_dim
+    _, num_kv_heads = model_config.get_num_qkv_head_by_tp()
+    return num_kv_heads * (model_config.k_head_dim + model_config.v_head_dim)
+
+
 def get_dcp_local_seq_lens(seq_lens: torch.Tensor,
                            dcp_world_rank: tuple[int, int]) -> torch.Tensor:
     """Return interleaved per-rank lengths for global sequence lengths."""
@@ -150,7 +173,7 @@ def update_dcp_metadata(attn_metadata, step_context) -> None:
         prefix_lens=attn_metadata.kv_seqlens - attn_metadata.q_seqlens,
         prefix_limit=prefix_limit,
         block_size=step_context.cache_config.block_size,
-        head_dim=step_context.model_config.head_dim,
+        kv_width=get_dcp_kv_width(step_context.model_config),
         dcp_world_rank=dcp_world_rank,
     )
     topk = step_context.model_config.mla_index_topk
diff --git a/lmdeploy/pytorch/backends/cuda/attention/__init__.py b/lmdeploy/pytorch/backends/cuda/attention/__init__.py
index 1c09484b6b..fa4b8a17b1 100644
--- a/lmdeploy/pytorch/backends/cuda/attention/__init__.py
+++ b/lmdeploy/pytorch/backends/cuda/attention/__init__.py
@@ -126,6 +126,11 @@ def _build_paged_attention(spec: PagedAttentionBuildSpec) -> TritonAttentionImpl
         logger.debug('Build FlashMLAImpl Attention')
         from .mla import FlashMLAImpl
         return FlashMLAImpl(use_fa3=use_fa3, **common_args)
+
+    from lmdeploy.pytorch.distributed import get_dcp_world_rank
+    if get_dcp_world_rank()[0] > 1:
+        from .cp import DCPAttentionImpl
+        return DCPAttentionImpl(use_fa3=enable_fa3, **common_args)
     elif enable_fa3:
         logger.debug('Build FA3Impl Attention')
         from .fa3 import FA3Impl
diff --git a/lmdeploy/pytorch/backends/cuda/attention/cp.py b/lmdeploy/pytorch/backends/cuda/attention/cp.py
index fe55818372..15dbbcdd47 100644
--- a/lmdeploy/pytorch/backends/cuda/attention/cp.py
+++ b/lmdeploy/pytorch/backends/cuda/attention/cp.py
@@ -1,63 +1,255 @@
 # Copyright (c) OpenMMLab. All rights reserved.
-"""CUDA DCP query gathering and attention-output merging."""
+"""CUDA DCP attention implementation and shared collectives."""
 
 import torch
 
+from lmdeploy.pytorch.backends.cp_utils import DCPPrefixChunk, get_dcp_local_causal_seq_lens
 
-def gather_dcp_query(query: torch.Tensor, *, dcp_world_size: int) -> torch.Tensor:
-    """Gather [tokens, local_heads, dim] queries along the head axis."""
-    if dcp_world_size == 1:
-        return query
-    from lmdeploy.pytorch.distributed import all_gather_into_tensor
-
-    # Keep the existing packing path for strided queries. Direct gather
-    # saves a copy only when the token-major input is already contiguous.
-    if not query.is_contiguous():
-        transposed = query.transpose(0, 1).contiguous()
-        gathered = transposed.new_empty(
-            dcp_world_size * transposed.size(0), *transposed.shape[1:])
-        all_gather_into_tensor(gathered, transposed, group='dcp')
-        return gathered.transpose(0, 1).contiguous()
-
-    gathered = query.new_empty(dcp_world_size * query.size(0), *query.shape[1:])
-    all_gather_into_tensor(gathered, query, group='dcp')
-    # Gather token-major queries directly, then join rank-local heads.
-    # Contiguous inputs need only this final layout conversion.
-    gathered = gathered.view(dcp_world_size, *query.shape)
-    return gathered.transpose(0, 1).reshape(query.size(0), -1, query.size(2)).contiguous()
-
-
-def merge_dcp_attention(local_output: torch.Tensor,
-                        local_lse: torch.Tensor,
-                        valid_counts: torch.Tensor,
-                        *,
-                        dcp_world_rank: tuple[int, int]) -> torch.Tensor:
-    """Merge normalized CUDA shard outputs and scatter heads back to each rank.
-
-    Inputs are [tokens, gathered_heads, value_dim] outputs and natural-log [tokens, gathered_heads] LSE. Empty rows
-    contribute zero. LSE and correction arithmetic use FP32; reduce-scatter uses the input output dtype.
+from .default import TritonAttentionImpl, TritonAttentionMetaBuilder, TritonAttentionMetadata
+
+
+class DCPManager:
+    """Own shared DCP gathers and the output combine that orders query reuse.
+
+    Layers in one distributed context share this manager. Prepare workspaces during model construction, before CUDA
+    graph capture, and close only after all consumers and graphs have retired.
     """
-    dcp_world_size, dcp_rank = dcp_world_rank
-    if dcp_world_size == 1:
-        return local_output
-    from lmdeploy.pytorch.distributed import all_gather_into_tensor, reduce_scatter_tensor
-    from lmdeploy.pytorch.kernels.cuda.dcp import correct_dcp_attention_output, sanitize_dcp_lse
-
-    local_lse = sanitize_dcp_lse(local_lse, valid_counts)
-    gathered_lse = local_lse.new_empty(
-        dcp_world_size * local_lse.size(0), local_lse.size(1))
-    all_gather_into_tensor(gathered_lse, local_lse, group='dcp')
-    gathered_lse = gathered_lse.view(dcp_world_size, *local_lse.shape)
-    contribution = correct_dcp_attention_output(
-        local_output, gathered_lse, dcp_rank=dcp_rank)
-
-    num_heads = contribution.size(0)
-    # The correction kernel writes [heads, tokens, dim] directly so the
-    # head-sharded reduce-scatter needs no separate transpose/copy.
-    assert num_heads % dcp_world_size == 0
-    local_heads = num_heads // dcp_world_size
-    scattered_output = contribution.new_empty(local_heads,
-                                              contribution.size(1),
-                                              contribution.size(2))
-    reduce_scatter_tensor(scattered_output, contribution, group='dcp')
-    return scattered_output.transpose(0, 1)
+
+    def __init__(self, group):
+        self.group = group.gpu_group
+        self.communicator = group.communicator
+        self.world_size = torch.distributed.get_world_size(self.group)
+        self.rank = group.rank
+        self._query_workspace = None
+        self._lse_workspace = None
+        self._candidate_workspace = None
+        self._prefix_workspaces = {}
+
+    def prepare_attention(self, num_heads: int, head_size: int):
+        """Prepare query and LSE arenas shared across attention layers."""
+        if self._query_workspace is None:
+            width = num_heads * head_size * self.world_size
+            # combine() orders all ranks after attention consumes these queries,
+            # so the next query gather does not need another entry barrier.
+            self._query_workspace = self.communicator.create_all_gather_workspace(
+                width, device=torch.device('cuda'), dtype=torch.bfloat16, reuse_sync=False)
+        if self._lse_workspace is None:
+            # Output reduce-scatter orders LSE readers before the next reuse.
+            # LSE is small; larger inputs fall back to NCCL beyond this budget.
+            self._lse_workspace = self.communicator.create_all_gather_workspace(
+                num_heads * self.world_size, device=torch.device('cuda'), dtype=torch.float32,
+                dim=0, reuse_sync=False, capacity_bytes=1024 * 1024)
+
+    def prepare_candidate_gather(self, topk: int):
+        """Prepare one candidate arena shared across indexer layers."""
+        if self._candidate_workspace is None:
+            self._candidate_workspace = self.communicator.create_all_gather_workspace(
+                topk * 2, device=torch.device('cuda'), dtype=torch.float32, dim=0)
+
+    def prepare_prefix_gather(self, *kv_widths: int):
+        """Share prefix arenas across layers and equal-width K/V payloads."""
+        for width in kv_widths:
+            if width not in self._prefix_workspaces:
+                self._prefix_workspaces[width] = self.communicator.create_all_gather_workspace(
+                    width, device=torch.device('cuda'), dtype=torch.bfloat16, dim=0)
+
+    def gather_prefix(self, local_kv: torch.Tensor, chunk: DCPPrefixChunk,
+                      *, zero_padding: bool = False) -> torch.Tensor:
+        """Gather cached KV and restore sequence order into an owned output.
+
+        The entry barrier protects the previous reorder before arena reuse. Zero padding for consumers such as FA3
+        that may load masked V lanes: 0 * NaN can contaminate valid outputs.
+        """
+        from lmdeploy.pytorch.kernels.cuda.dcp import reorder_dcp_prefill_kv
+
+        flat_kv = local_kv.flatten(1)
+        gathered = self.communicator.all_gather(
+            flat_kv, dim=0, workspace=self._prefix_workspaces.get(flat_kv.size(1)), copy_output=False)
+        gathered = gathered.view(-1, *local_kv.shape[1:])
+        context = torch.zeros_like(gathered) if zero_padding else torch.empty_like(gathered)
+        reorder_dcp_prefill_kv(gathered, context, chunk_kv_seqlens=chunk.kv_seqlens,
+                               kv_start_loc=chunk.cu_seqlens[:-1], local_lens=chunk.local_kv_seqlens)
+        return context
+
+    def gather_query(self, query: torch.Tensor) -> torch.Tensor:
+        """Gather query heads; consume and combine on the same stream before
+        reuse."""
+        if self.world_size == 1:
+            return query
+        output = self.communicator.all_gather(
+            query.flatten(1), workspace=self._query_workspace, copy_output=False)
+        return output.view(query.size(0), -1, query.size(2))
+
+    def gather_candidates(self, packed: torch.Tensor) -> torch.Tensor:
+        """Gather [rows, topk, score/ID] payloads without changing their bits.
+
+        The entry barrier protects the previous top-k consumer before arena reuse. Consume the borrowed output on the
+        same stream before the next candidate gather.
+        """
+        output = self.communicator.all_gather(
+            packed.flatten(1), dim=0, workspace=self._candidate_workspace, copy_output=False)
+        return output.view(self.world_size, *packed.shape)
+
+    def combine(self, local_output: torch.Tensor, local_lse: torch.Tensor,
+                valid_counts: torch.Tensor) -> torch.Tensor:
+        """Merge shard outputs and scatter heads, ordering the next query
+        gather.
+
+        Inputs are [tokens, gathered_heads, value_dim] outputs and natural-log [tokens, gathered_heads] LSE. Empty rows
+        contribute zero. LSE and correction arithmetic use FP32; reduce-scatter uses the input output dtype.
+        """
+        if self.world_size == 1:
+            return local_output
+        from lmdeploy.pytorch.distributed import reduce_scatter_tensor
+        from lmdeploy.pytorch.kernels.cuda.dcp import correct_dcp_attention_output, sanitize_dcp_lse
+
+        local_lse = sanitize_dcp_lse(local_lse, valid_counts)
+        gathered_lse = self.communicator.all_gather(
+            local_lse, dim=0, workspace=self._lse_workspace, copy_output=False)
+        gathered_lse = gathered_lse.view(self.world_size, *local_lse.shape)
+        contribution = correct_dcp_attention_output(
+            local_output, gathered_lse, dcp_rank=self.rank)
+
+        num_heads = contribution.size(0)
+        # The correction kernel writes [heads, tokens, dim] directly so the
+        # head-sharded reduce-scatter needs no separate transpose/copy.
+        assert num_heads % self.world_size == 0
+        local_heads = num_heads // self.world_size
+        scattered_output = contribution.new_empty(local_heads,
+                                                  contribution.size(1),
+                                                  contribution.size(2))
+        reduce_scatter_tensor(scattered_output, contribution, group=self.group)
+        return scattered_output.transpose(0, 1)
+
+    def close(self):
+        """Release shared arenas before their process group is destroyed."""
+        if self._query_workspace is not None:
+            self._query_workspace.close()
+            self._query_workspace = None
+        if self._lse_workspace is not None:
+            self._lse_workspace.close()
+            self._lse_workspace = None
+        if self._candidate_workspace is not None:
+            self._candidate_workspace.close()
+            self._candidate_workspace = None
+        for workspace in self._prefix_workspaces.values():
+            if workspace is not None:
+                workspace.close()
+        self._prefix_workspaces.clear()
+
+
+def get_dcp_manager() -> DCPManager:
+    """Get the context-owned manager during layer initialization."""
+    from lmdeploy.pytorch.distributed import get_dist_manager
+
+    context = get_dist_manager().current_context()
+    if context.dcp_manager is None:
+        context.dcp_manager = DCPManager(context.dcp_group)
+    return context.dcp_manager
+
+
+class DCPAttentionImpl(TritonAttentionImpl):
+    """Keep TP query heads local in prefill and gather them for decode.
+
+    Model configuration ensures that every rank in a DCP group holds the same KV head. Cached prefill gathers bounded
+    K/V chunks; decode evaluates each context shard and reduce-scatters softmax-corrected outputs to the TP heads.
+    """
+
+    def __init__(self, num_heads: int, head_size: int, *, use_fa3: bool = False, **kwargs):
+        super().__init__(num_heads=num_heads, head_size=head_size, **kwargs)
+        self.use_fa3 = use_fa3
+        assert not self.alibi, 'GQA DCP does not support ALiBi'
+        if use_fa3:
+            from lmdeploy.pytorch.third_party.flash_attn_interface import flash_attn_varlen_func
+            self._fa3_prefill = flash_attn_varlen_func
+        self.dcp_manager = get_dcp_manager()
+        self.dcp_manager.prepare_attention(self.num_heads, self.head_size)
+        self.dcp_manager.prepare_prefix_gather(self.num_kv_heads * self.head_size,
+                                               self.num_kv_heads * self.v_head_size)
+
+    def get_step_metadata_provider(self):
+        """Use common DCP lengths without a separate kernel scheduler."""
+        return TritonAttentionMetaBuilder()
+
+    def supports_piecewise_cuda_graph(self) -> bool:
+        return True
+
+    def _decode(self, query: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor,
+                metadata: TritonAttentionMetadata, k_scales_zeros: torch.Tensor = None,
+                v_scales_zeros: torch.Tensor = None) -> torch.Tensor:
+        """Evaluate interleaved shards and restore rank-local query heads."""
+        query = self.dcp_manager.gather_query(query)
+        query_len = query.size(0) // metadata.q_seqlens.numel()
+        local_lens = metadata.dcp_local_kv_seqlens
+        blocks = metadata.block_offsets
+        if query_len > 1:
+            # A global query step need not advance this shard's causal length.
+            local_lens = get_dcp_local_causal_seq_lens(
+                metadata.kv_seqlens, query_len, (self.dcp_world_size, self.dcp_rank))
+            blocks = blocks.repeat_interleave(query_len, dim=0)
+        output, lse = self.paged_attention_fwd(
+            query, k_cache, v_cache, cache_seqlens=local_lens,
+            page_table=blocks, max_seqlen_q=1, softmax_scale=self.scale,
+            softcap=self.logit_softcapping, return_lse=True, quant_policy=metadata.quant_policy,
+            k_scales_zeros=k_scales_zeros, v_scales_zeros=v_scales_zeros)
+        return self.dcp_manager.combine(output, lse, valid_counts=local_lens)
+
+    def _prefill_attention(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, *,
+                           cu_q: torch.Tensor, cu_k: torch.Tensor, max_q: int, max_k: int,
+                           causal: bool, return_lse: bool):
+        """Evaluate one prefill partition with optional natural-log LSE."""
+        kwargs = dict(cu_seqlens_q=cu_q, cu_seqlens_k=cu_k,
+                      max_seqlen_q=max_q, max_seqlen_k=max_k,
+                      softmax_scale=self.scale, softcap=max(0.0, self.logit_softcapping), causal=causal)
+        if self.use_fa3:
+            return self._fa3_prefill(query, key, value, return_lse=return_lse, **kwargs)
+        return self.flash_attention_fwd(query, key, value, kv_layout='shd', return_lse=return_lse, **kwargs)
+
+    def _gather_prefix(self, k_cache: torch.Tensor, v_cache: torch.Tensor,
+                       metadata: TritonAttentionMetadata, chunk: DCPPrefixChunk,
+                       dtype: torch.dtype, k_scales_zeros: torch.Tensor = None,
+                       v_scales_zeros: torch.Tensor = None) -> tuple[torch.Tensor, torch.Tensor]:
+        """Gather and reorder separate K/V caches within the chunk budget."""
+        local_k, local_v = self.flatten_kv_cache(
+            k_cache, v_cache, chunk.local_kv_seqlens[self.dcp_rank],
+            metadata.block_offsets[:, chunk.local_block_slice(k_cache.size(1))],
+            start_loc=chunk.local_cu_seqlens[:-1], out_size=chunk.local_capacity,
+            out_dtype=dtype, flatten_kv_layout='shd', quant_policy=metadata.quant_policy,
+            k_scales_zeros=k_scales_zeros, v_scales_zeros=v_scales_zeros)
+        return (self.dcp_manager.gather_prefix(local_k, chunk, zero_padding=True),
+                self.dcp_manager.gather_prefix(local_v, chunk, zero_padding=True))
+
+    def forward(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor,
+                k_cache: torch.Tensor, v_cache: torch.Tensor, attn_metadata: TritonAttentionMetadata,
+                k_scales_zeros: torch.Tensor = None, v_scales_zeros: torch.Tensor = None,
+                learnable_sink: torch.Tensor = None, inplace: bool = True, **kwargs) -> torch.Tensor:
+        """Write owned KV tokens, then evaluate GQA prefill or decode."""
+        assert learnable_sink is None, 'GQA DCP does not support attention sinks'
+        max_q = (query.size(0) // attn_metadata.q_seqlens.numel()
+                 if attn_metadata.is_decoding else self._get_max_q_seqlen(query, attn_metadata))
+        if key is not None and value is not None:
+            self._fill_kv_cache_impl(key, value, k_cache, v_cache, attn_metadata, max_q,
+                                     k_scales_zeros=k_scales_zeros, v_scales_zeros=v_scales_zeros)
+        if attn_metadata.is_decoding:
+            return self._decode(query, k_cache, v_cache, attn_metadata,
+                                k_scales_zeros=k_scales_zeros, v_scales_zeros=v_scales_zeros)
+
+        assert key is not None and value is not None, 'GQA DCP prefill requires current K/V'
+        chunks = attn_metadata.dcp_prefix_chunks
+        result = self._prefill_attention(
+            query, key, value, cu_q=attn_metadata.cu_seqlens_q, cu_k=attn_metadata.cu_seqlens_q,
+            max_q=max_q, max_k=max_q, causal=True, return_lse=bool(chunks))
+        if not chunks:
+            return result
+        from lmdeploy.pytorch.kernels.cuda.dcp import merge_attention_states
+        output, lse = result
+        del result
+        for chunk in chunks:
+            context_k, context_v = self._gather_prefix(k_cache, v_cache, attn_metadata, chunk, query.dtype,
+                                                      k_scales_zeros=k_scales_zeros, v_scales_zeros=v_scales_zeros)
+            chunk_output, chunk_lse = self._prefill_attention(
+                query, context_k, context_v, cu_q=attn_metadata.cu_seqlens_q, cu_k=chunk.cu_seqlens,
+                max_q=max_q, max_k=chunk.size, causal=False, return_lse=True)
+            output, lse = merge_attention_states(output, lse, chunk_output, chunk_lse)
+            del context_k, context_v, chunk_output, chunk_lse
+        return output
diff --git a/lmdeploy/pytorch/backends/cuda/attention/mla.py b/lmdeploy/pytorch/backends/cuda/attention/mla.py
index e29d686d57..433fa3f6dc 100644
--- a/lmdeploy/pytorch/backends/cuda/attention/mla.py
+++ b/lmdeploy/pytorch/backends/cuda/attention/mla.py
@@ -14,7 +14,7 @@
 from lmdeploy.utils import get_logger
 
 from ..step_metadata import CudaAttentionMetaBuilder
-from .cp import gather_dcp_query, merge_dcp_attention
+from .cp import get_dcp_manager
 from .default import TritonAttentionImpl, TritonAttentionMetadata
 
 logger = get_logger('lmdeploy')
@@ -254,6 +254,11 @@ def __init__(
         assert num_kv_heads == 1, 'MLA requires num kv heads equal to 1'
         self.use_fa3 = use_fa3
 
+        if self.dcp_world_size > 1:
+            self.dcp_manager = get_dcp_manager()
+            self.dcp_manager.prepare_attention(num_heads, head_size)
+            self.dcp_manager.prepare_prefix_gather(head_size)
+
     def get_step_metadata_provider(self):
         """Describe metadata required by this selected implementation."""
         return FlashMLAAttentionMetaBuilder(num_attention_heads=self.num_heads)
@@ -518,23 +523,14 @@ def _gather_dcp_prefix_chunk(
         v_scales_zeros: torch.Tensor = None,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         """Gather one globally ordered cached-prefix chunk."""
-        from lmdeploy.pytorch.distributed import all_gather_into_tensor
-        from lmdeploy.pytorch.kernels.cuda.dcp import reorder_dcp_prefill_kv
-
-        dcp_world_size = self.dcp_world_size
-        block_size = k_cache.size(1)
-        virtual_block_size = block_size * dcp_world_size
-        local_capacity = chunk.kv_seqlens.numel() * chunk.size // dcp_world_size
-        block_start = chunk.start // virtual_block_size
-        num_blocks = chunk.size // virtual_block_size
         local_meta = replace(
             attn_metadata,
-            block_offsets=attn_metadata.block_offsets[:, block_start:block_start + num_blocks],
+            block_offsets=attn_metadata.block_offsets[:, chunk.local_block_slice(k_cache.size(1))],
             kv_start_loc=chunk.local_cu_seqlens[:-1],
             kv_seqlens=chunk.local_kv_seqlens[self.dcp_rank],
             cu_seqlens_k=chunk.local_cu_seqlens,
-            kv_flatten_size=local_capacity,
-            max_kv_seqlen=chunk.size // dcp_world_size,
+            kv_flatten_size=chunk.local_capacity,
+            max_kv_seqlen=chunk.size // self.dcp_world_size,
         )
         local_k, _ = self._flatten_prefill_kv_cache(
             k_cache,
@@ -545,18 +541,7 @@ def _gather_dcp_prefix_chunk(
             k_scales_zeros=k_scales_zeros,
             v_scales_zeros=v_scales_zeros,
         )
-        gathered = local_k.new_empty(dcp_world_size * local_capacity,
-                                     *local_k.shape[1:])
-        all_gather_into_tensor(gathered, local_k, group='dcp')
-
-        context_k = local_k.new_empty(chunk.kv_seqlens.numel() * chunk.size, *local_k.shape[1:])
-        reorder_dcp_prefill_kv(
-            gathered,
-            context_k,
-            chunk_kv_seqlens=chunk.kv_seqlens,
-            kv_start_loc=chunk.cu_seqlens[:-1],
-            local_lens=chunk.local_kv_seqlens,
-        )
+        context_k = self.dcp_manager.gather_prefix(local_k, chunk)
         return context_k, chunk.cu_seqlens
 
     def _prefill_dcp_context(
@@ -772,7 +757,7 @@ def _forward_decoding(
         if self.dcp_world_size == 1:
             return self._decoding_paged(query, k_cache, attn_metadata)
 
-        query = gather_dcp_query(query, dcp_world_size=self.dcp_world_size)
+        query = self.dcp_manager.gather_query(query)
         query_len = query.size(0) // attn_metadata.q_seqlens.numel()
         if query_len > 1:
             # Consecutive global queries do not advance each interleaved
@@ -787,11 +772,10 @@ def _forward_decoding(
             )
         local_output, local_lse = self._decoding_paged(
             query, k_cache, attn_metadata, causal=False, return_lse=True)
-        return merge_dcp_attention(
+        return self.dcp_manager.combine(
             local_output,
             local_lse,
-            valid_counts=attn_metadata.dcp_local_kv_seqlens,
-            dcp_world_rank=(self.dcp_world_size, self.dcp_rank))
+            valid_counts=attn_metadata.dcp_local_kv_seqlens)
 
     def _forward_prefill(
         self,
diff --git a/lmdeploy/pytorch/backends/cuda/attention/sparse_mla.py b/lmdeploy/pytorch/backends/cuda/attention/sparse_mla.py
index bb3ee3d2d7..e6966d7c40 100644
--- a/lmdeploy/pytorch/backends/cuda/attention/sparse_mla.py
+++ b/lmdeploy/pytorch/backends/cuda/attention/sparse_mla.py
@@ -7,7 +7,6 @@
 from lmdeploy.pytorch import envs
 from lmdeploy.utils import get_logger
 
-from .cp import gather_dcp_query, merge_dcp_attention
 from .default import TritonAttentionMetadata
 from .mla import FlashMLAAttentionMetaBuilder, FlashMLAImpl
 
@@ -341,14 +340,11 @@ def _forward_decoding(self,
             dcp_world_rank = self.dcp_world_size, self.dcp_rank
             local_indices, local_counts = filter_and_compact_dcp_indices(
                 nsa_indices, dcp_world_rank=dcp_world_rank)
-            query = gather_dcp_query(query, dcp_world_size=self.dcp_world_size)
+            query = self.dcp_manager.gather_query(query)
             local_output, local_lse = self._decoding_sparse(
                 query, k_cache, local_indices, attn_metadata,
                 return_lse=True, topk_length=local_counts)
-            return merge_dcp_attention(local_output,
-                                       local_lse,
-                                       valid_counts=local_counts,
-                                       dcp_world_rank=dcp_world_rank)
+            return self.dcp_manager.combine(local_output, local_lse, valid_counts=local_counts)
         return self._decoding_sparse(query, k_cache, nsa_indices, attn_metadata)
 
     def _forward_prefill(self,
diff --git a/lmdeploy/pytorch/backends/cuda/comm/communicator.py b/lmdeploy/pytorch/backends/cuda/comm/communicator.py
index cd0a2bddf3..1e2569267d 100644
--- a/lmdeploy/pytorch/backends/cuda/comm/communicator.py
+++ b/lmdeploy/pytorch/backends/cuda/comm/communicator.py
@@ -1,79 +1,104 @@
 # Copyright (c) OpenMMLab. All rights reserved.
+import socket
+from typing import Literal
+
 import torch
 from torch import distributed as dist
 
-from lmdeploy.pytorch import envs as _envs
+from lmdeploy.utils import get_logger
 
 from ...communicator import DeviceCommunicator
-from .flashinfer_allreduce import FlashInferAllReduce
 from .symm_mem_allreduce import SymmetricMemoryAllReduce
 
+logger = get_logger('lmdeploy')
+
 
 class CudaCommunicator(DeviceCommunicator):
-    """Dispatch optional CUDA collectives with a process-group fallback."""
+    """Prepare group-owned providers and dispatch by operation and input
+    shape."""
 
-    def __init__(self, cpu_group: dist.ProcessGroup, device_group: dist.ProcessGroup):
+    def __init__(self, cpu_group: dist.ProcessGroup, device_group: dist.ProcessGroup, *,
+                 group_name: Literal['tp', 'dcp']):
         super().__init__(device_group=device_group)
-        if _envs.enable_flashinfer_allreduce and _envs.enable_symm_mem_allreduce:
-            raise ValueError('FlashInfer and symmetric-memory all-reduce cannot be enabled together.')
-
-        self._flashinfer = (FlashInferAllReduce(cpu_group)
-                            if _envs.enable_flashinfer_allreduce else None)
-        self._symm_mem = (SymmetricMemoryAllReduce(cpu_group)
-                          if _envs.enable_symm_mem_allreduce else None)
-
-    def supports_optimized_all_reduce(self) -> bool:
-        """Whether an optimized all-reduce implementation is available."""
-        return ((self._flashinfer is not None and self._flashinfer.is_available())
-                or (self._symm_mem is not None and self._symm_mem.is_available()))
-
-    def supports_fused_all_reduce_residual_rms_norm(self) -> bool:
-        """Whether fused all-reduce, residual and RMSNorm is available."""
-        return self._flashinfer is not None and self._flashinfer.is_available()
-
-    def try_fused_all_reduce_residual_rms_norm(self,
-                                               input: torch.Tensor,
-                                               residual: torch.Tensor,
-                                               weight: torch.Tensor,
-                                               eps: float):
-        """Run fused all-reduce, residual and RMSNorm when eligible."""
-        if self._flashinfer is None:
-            return None
-        return self._flashinfer.fused_all_reduce_residual_rms_norm(
-            input=input,
-            residual=residual,
-            weight=weight,
-            eps=eps,
-        )
+        self._cpu_group = cpu_group
+        self._all_reduce_provider = SymmetricMemoryAllReduce(cpu_group) if group_name == 'tp' else None
+        if self._all_reduce_provider is not None:
+            self._check_all_reduce_provider()
+        if self._all_reduce_provider is not None:
+            self._all_reduce_provider.prepare()
+            self._check_all_reduce_provider()
+        if dist.get_rank(device_group) == 0:
+            logger.info('Communication backend: auto group=%s symmetric_all_reduce=%s device=%s',
+                        group_name, self._all_reduce_provider is not None, torch.cuda.get_device_name())
+
+    def _check_all_reduce_provider(self):
+        """Keep or retire symmetric all-reduce consistently on every rank."""
+        provider = self._all_reduce_provider
+        flags = [None] * dist.get_world_size(self._cpu_group)
+        dist.all_gather_object(flags, provider.is_available(), group=self._cpu_group)
+        if not all(flags):
+            provider.close()
+            self._all_reduce_provider = None
+            logger.info('Symmetric-memory all-reduce unavailable; using process-group fallback')
 
     def all_reduce_(self, input: torch.Tensor):
-        """Dispatch all-reduce through optimized CUDA backends."""
-        if self._flashinfer is not None and self._flashinfer.all_reduce_(input):
-            return
-        if self._symm_mem is not None and self._symm_mem.all_reduce_(input):
+        if self._all_reduce_provider is not None and self._all_reduce_provider.all_reduce_(input):
             return
         super().all_reduce_(input)
 
-    def close(self):
-        """Release communicator-owned workspaces."""
-        if self._flashinfer is not None:
-            self._flashinfer.close()
-        if self._symm_mem is not None:
-            self._symm_mem.close()
+    def create_all_gather_workspace(self, gathered_width: int, device: torch.device, dtype: torch.dtype,
+                                    *, dim: int = -1, reuse_sync: bool = True, capacity_bytes: int = 64 * 1024 * 1024):
+        from .symm_mem_allgather import SymmetricMemoryAllGather
+
+        workspace = SymmetricMemoryAllGather(
+            self.device_group, dist.get_rank(self.device_group), gathered_width,
+            device=device, dtype=dtype, capacity_bytes=capacity_bytes, dim=dim, reuse_sync=reuse_sync)
+        workspace.prepare()
+        return workspace
 
+    def all_gather(self, input: torch.Tensor, *, dim: int = -1,
+                   workspace=None, copy_output: bool = True) -> torch.Tensor:
+        if workspace is not None:
+            output = workspace.all_gather(input, dim=dim, copy_output=copy_output)
+            if output is not None:
+                return output
+        return super().all_gather(input, dim=dim)
 
-def should_try_symm_mem(dist_config) -> bool:
-    """Whether this configuration is a symmetric-memory candidate."""
-    return (_envs.enable_symm_mem_allreduce and dist_config.dp == 1
-            and dist_config.ep == 1 and dist_config.attn_tp > 1
-            and not dist_config.enable_microbatch)
+    def close(self):
+        if self._all_reduce_provider is not None:
+            self._all_reduce_provider.close()
+            self._all_reduce_provider = None
 
 
 def build_cuda_communicator(cpu_group: dist.ProcessGroup, device_group: dist.ProcessGroup,
-                            dist_config):
-    """Build the optional CUDA communicator for a TP group."""
+                            dist_config, *, group_name: Literal['tp', 'dcp'] = 'tp'):
+    """Agree on configuration and topology before optional collective setup."""
+    backend = dist_config.communication_backend
     compatible = dist_config.dp == 1 and dist_config.ep == 1 and not dist_config.enable_microbatch
-    enabled = _envs.enable_flashinfer_allreduce or should_try_symm_mem(dist_config)
-    if not compatible or not enabled:
+    if not compatible:
+        return None
+
+    device = torch.cuda.current_device()
+    local = (backend, group_name, socket.gethostname(), device,
+             torch.cuda.get_device_name(device), torch.cuda.get_device_capability(device))
+    world_size = dist.get_world_size(device_group)
+    peers = [None] * world_size
+    dist.all_gather_object(peers, local, group=cpu_group)
+    if any(peer[:2] != local[:2] for peer in peers):
+        raise ValueError('communication_backend and group_name must agree across ranks')
+    if backend == 'nccl':
         return None
-    return CudaCommunicator(cpu_group=cpu_group, device_group=device_group)
+
+    same_node = all(peer[2] == local[2] for peer in peers)
+    same_device = all(peer[4:] == local[4:] for peer in peers)
+    peer_access = (same_node and same_device and len({peer[3] for peer in peers}) == world_size
+                   and all(peer[3] == device or torch.cuda.can_device_access_peer(device, peer[3]) for peer in peers))
+    flags = [None] * world_size
+    dist.all_gather_object(flags, peer_access, group=cpu_group)
+    if not all(flags):
+        if dist.get_rank(device_group) == 0:
+            logger.info('Communication policy: process-group fallback (peer topology unavailable)')
+        return None
+
+    return CudaCommunicator(cpu_group=cpu_group, device_group=device_group,
+                            group_name=group_name)
diff --git a/lmdeploy/pytorch/backends/cuda/comm/flashinfer_allreduce.py b/lmdeploy/pytorch/backends/cuda/comm/flashinfer_allreduce.py
deleted file mode 100644
index 9cd5a431b5..0000000000
--- a/lmdeploy/pytorch/backends/cuda/comm/flashinfer_allreduce.py
+++ /dev/null
@@ -1,209 +0,0 @@
-# Copyright (c) OpenMMLab. All rights reserved.
-# Modified from: https://github.com/vllm-project/vllm
-import torch
-from torch import distributed as dist
-
-from lmdeploy.utils import get_logger
-
-logger = get_logger('lmdeploy')
-
-_MiB = 1024 * 1024
-_MAX_TOKEN_NUM = 2048
-
-# Empirical FlashInfer fusion limits used by vLLM:
-# https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/allreduce_rms_fusion.py
-# The separate token cap bounds each process-group workspace.
-_FUSION_MAX_SIZE = {
-    90: {
-        2: 64 * _MiB,
-        4: 2 * _MiB,
-        8: _MiB // 2,
-    },
-    100: {
-        2: 64 * _MiB,
-        4: 32 * _MiB,
-        8: _MiB,
-    },
-    103: {
-        2: 64 * _MiB,
-        4: 64 * _MiB,
-        8: 2 * _MiB,
-    },
-}
-_ONE_SHOT_MAX_SIZE = {
-    90: {
-        2: 32 * _MiB,
-        4: 2 * _MiB,
-        8: _MiB // 2,
-    },
-    100: {
-        2: 32 * _MiB,
-        4: 4 * _MiB,
-        8: _MiB,
-    },
-    103: {
-        2: 32 * _MiB,
-        4: 4 * _MiB,
-        8: 2 * _MiB,
-    },
-}
-
-
-class FlashInferAllReduce:
-    """FlashInfer all-reduce operations for one distributed group."""
-
-    def __init__(self, group: dist.ProcessGroup):
-        """Initialize lazy FlashInfer state for ``group``."""
-        self.group = group
-        self._comm = None
-        self._workspace = None
-        self._hidden_dim = None
-        self._dtype = None
-        self._disabled = False
-        self._world_size = dist.get_world_size(group)
-        major, minor = torch.cuda.get_device_capability()
-        self._device_capability = major * 10 + minor
-        self._max_size = _FUSION_MAX_SIZE.get(self._device_capability, {}).get(self._world_size, 0)
-        self._one_shot_max_size = _ONE_SHOT_MAX_SIZE.get(self._device_capability, {}).get(self._world_size, 0)
-
-    def _disable(self, reason):
-        """Disable this backend after a process-group-wide setup failure."""
-        self._disabled = True
-        logger.warning(
-            f'Disabling FlashInfer all-reduce for this process group: {reason}')
-
-    def _initialize(self) -> bool:
-        """Load and validate the optional FlashInfer communication API."""
-        if self._disabled:
-            return False
-        if self._comm is not None:
-            return True
-        message = 'FlashInfer all-reduce fusion requires flashinfer-python with the unified comm API.'
-        try:
-            import flashinfer.comm as comm
-            from flashinfer.comm.cuda_ipc import cudart
-            if not (hasattr(comm, 'allreduce_fusion') and hasattr(
-                    comm, 'create_allreduce_fusion_workspace')):
-                raise ImportError(message)
-
-            # Resolve FlashInfer's CUDA runtime before TileLang loads its
-            # libcudart shim, which the lazy loader could otherwise select.
-            cudart.cudaSetDevice(torch.cuda.current_device())
-        except Exception as e:
-            self._disable(e)
-            return False
-        self._comm = comm
-        return True
-
-    def is_available(self) -> bool:
-        """Whether the fused collective implementation is available."""
-        return self._max_size > 0 and self._initialize()
-
-    def supports(self, dtype: torch.dtype) -> bool:
-        """Whether this group can handle the given dtype."""
-        return dtype in (torch.float16, torch.bfloat16) and self.is_available()
-
-    def _supports_input(self, input: torch.Tensor) -> bool:
-        """Whether a flattened input satisfies the FlashInfer launch limits."""
-        return (input.nbytes <= self._max_size
-                and input.size(0) <= _MAX_TOKEN_NUM
-                and input.is_contiguous()
-                and self.supports(input.dtype))
-
-    def _get_workspace(self, input: torch.Tensor):
-        """Return the workspace for the input shape and dtype."""
-        hidden_dim = input.size(-1)
-        if self._workspace is not None:
-            assert self._hidden_dim == hidden_dim and self._dtype == input.dtype
-            return self._workspace
-
-        rank = dist.get_rank(self.group)
-        max_token_num = min(_MAX_TOKEN_NUM, self._max_size // input[0].nbytes)
-        try:
-            workspace = self._comm.create_allreduce_fusion_workspace(
-                backend='trtllm',
-                world_size=self._world_size,
-                rank=rank,
-                max_token_num=max_token_num,
-                hidden_dim=hidden_dim,
-                dtype=input.dtype,
-                group=self.group,
-            )
-            if workspace is None:
-                raise RuntimeError('workspace creation returned None')
-        except Exception as e:
-            self._disable(f'workspace initialization failed: {e}')
-            return None
-
-        self._workspace = workspace
-        self._hidden_dim = hidden_dim
-        self._dtype = input.dtype
-        logger.info(
-            f'FlashInfer all-reduce workspace initialized: rank={rank}, world_size={self._world_size}, '
-            f'max_token_num={max_token_num}, hidden_dim={hidden_dim}')
-        return self._workspace
-
-    def all_reduce_(self, input: torch.Tensor) -> bool:
-        """All-reduce ``input`` in place, returning whether it was handled."""
-        if input.dim() < 2 or not input.is_contiguous():
-            return False
-        input_2d = input.flatten(0, -2)
-        if not self._supports_input(input_2d):
-            return False
-
-        workspace = self._get_workspace(input_2d)
-        if workspace is None:
-            return False
-        output = self._comm.allreduce_fusion(
-            input=input_2d,
-            workspace=workspace,
-            pattern=self._comm.AllReduceFusionPattern.kAllReduce,
-            launch_with_pdl=True,
-            trigger_completion_at_end=True,
-            use_oneshot=input_2d.nbytes <= self._one_shot_max_size,
-        )
-        input_2d.copy_(output)
-        return True
-
-    def fused_all_reduce_residual_rms_norm(self,
-                                           input: torch.Tensor,
-                                           residual: torch.Tensor,
-                                           weight: torch.Tensor,
-                                           eps: float):
-        """Fuse all-reduce, residual addition and RMSNorm when supported."""
-        if weight.dtype != input.dtype:
-            return None
-        input_2d = input.flatten(0, -2)
-        residual_2d = residual.flatten(0, -2)
-        if (not self._supports_input(input_2d)
-                or not residual_2d.is_contiguous()
-                or not weight.is_contiguous()):
-            return None
-
-        norm_out = torch.empty_like(input_2d)
-        residual_out = torch.empty_like(residual_2d)
-        workspace = self._get_workspace(input_2d)
-        if workspace is None:
-            return None
-        self._comm.allreduce_fusion(
-            input=input_2d,
-            workspace=workspace,
-            pattern=self._comm.AllReduceFusionPattern.kARResidualRMSNorm,
-            launch_with_pdl=True,
-            # Keep completion conservative; early signaling requires separate
-            # validation for the one-shot and two-shot paths.
-            trigger_completion_at_end=True,
-            use_oneshot=input_2d.nbytes <= self._one_shot_max_size,
-            residual_in=residual_2d,
-            residual_out=residual_out,
-            norm_out=norm_out,
-            rms_gamma=weight,
-            rms_eps=eps,
-        )
-        return norm_out.view_as(input), residual_out.view_as(residual)
-
-    def close(self):
-        """Release the group-bound FlashInfer workspace."""
-        if self._workspace is not None:
-            self._workspace.destroy()
-            self._workspace = None
diff --git a/lmdeploy/pytorch/backends/cuda/comm/symm_mem_allgather.py b/lmdeploy/pytorch/backends/cuda/comm/symm_mem_allgather.py
index cf95b9ed0a..87e3c966b1 100644
--- a/lmdeploy/pytorch/backends/cuda/comm/symm_mem_allgather.py
+++ b/lmdeploy/pytorch/backends/cuda/comm/symm_mem_allgather.py
@@ -1,26 +1,26 @@
 # Copyright (c) OpenMMLab. All rights reserved.
-"""Setup-owned LM-head all-gather with a portable NCCL fallback.
+"""Symmetric-memory workspace for 2D all-gather.
 
 The optional Triton kernels are imported only during collective preparation. Their SGLang source reference is in
 kernels/cuda/symm_mem_allgather.py.
 """
 import importlib
-import logging
 
 import torch
 import torch.distributed as dist
 
-from lmdeploy.pytorch import envs as _envs
+from lmdeploy.utils import get_logger
 
-logger = logging.getLogger(__name__)
+logger = get_logger('lmdeploy')
 
 
-class MultimemAllGatherer:
-    """Own configuration, TP admission and the symmetric arena for one LM-head.
+class SymmetricMemoryAllGather:
+    """Own configuration, collective admission and one symmetric arena.
 
-    Construct on every TP rank after process-group setup. Only preparation performs allocation and consensus; forward
-    never initializes or rebuilds. Call reset_for_weight after coordinated, quiescent device/dtype changes. Runtime rows
-    must match across ranks, as required by NCCL all-gather too.
+    Construct on every group rank after process-group setup. Only preparation performs allocation and consensus; forward
+    never initializes or rebuilds. Re-prepare after coordinated, quiescent device/dtype changes. Runtime rows must match
+    across ranks, as required by NCCL all-gather too. With reuse_sync=False, callers must order all ranks' consumers
+    through a subsequent same-group collective before reusing the arena.
     """
 
     # Keep the conservative single-token NCCL policy internal to the provider.
@@ -28,67 +28,93 @@ class MultimemAllGatherer:
     _WORLD_SIZES = {2, 4, 8}
 
     def __init__(self, group: dist.ProcessGroup, rank: int, gathered_width: int,
-                 device: torch.device, dtype: torch.dtype):
-        self._group = group
+                 *, device: torch.device, dtype: torch.dtype, capacity_bytes: int,
+                 dim: int = -1, reuse_sync: bool = True):
+        self.device = torch.device(device)
+        self.dtype = dtype
+        self._prepared = False
+        self.group = group
         self._rank = rank
         self._gathered_width = gathered_width
-        self._enabled = _envs.enable_symm_mem_lmhead
-        self._capacity_bytes = _envs.symm_mem_lmhead_max_mb * 1024 * 1024
+        self._dim = dim
+        self._capacity_bytes = capacity_bytes
+        self._reuse_sync = reuse_sync
         self._state = None
         self._kernels = None
         self._graph_ready_shapes = set()
-        self._weight_contract = None
-        self.prepare(device, dtype)
 
     def _agree(self, ready: bool, device: torch.device) -> bool:
         flag = torch.tensor(int(ready), dtype=torch.int32, device=device)
-        dist.all_reduce(flag, op=dist.ReduceOp.MIN, group=self._group)
+        dist.all_reduce(flag, op=dist.ReduceOp.MIN, group=self.group)
         return bool(flag.item())
 
     def _same_config(self, values: tuple[int, ...], device: torch.device) -> bool:
         lower = torch.tensor(values, dtype=torch.int64, device=device)
         upper = lower.clone()
-        dist.all_reduce(lower, op=dist.ReduceOp.MIN, group=self._group)
-        dist.all_reduce(upper, op=dist.ReduceOp.MAX, group=self._group)
+        dist.all_reduce(lower, op=dist.ReduceOp.MIN, group=self.group)
+        dist.all_reduce(upper, op=dist.ReduceOp.MAX, group=self.group)
         return bool(torch.equal(lower, upper))
 
     def _disabled(self, reason: str) -> bool:
         if self._rank == 0:
-            logger.warning('symmetric-memory LM-head disabled: %s', reason)
+            logger.warning('symmetric-memory all-gather disabled: %s', reason)
         return False
 
-    def prepare(self, device: torch.device, dtype: torch.dtype) -> bool:
+    def _valid_handle(self, state, kernels) -> bool:
+        """Validate rank metadata and both pointer and signal-pad bounds."""
+        try:
+            handle = state.symm_mem_hdl
+            multicast = int(handle.multicast_ptr or 0)
+            signals = handle.signal_pad_ptrs_dev
+            signal_addr = int(signals.data_ptr() if hasattr(signals, 'data_ptr') else signals or 0)
+            pad_size = getattr(handle, 'signal_pad_size', None)
+            if pad_size is None:
+                pad_size = kernels.symm_mem.get_signal_pad_size()
+            return (int(handle.rank) == self._rank
+                    and int(handle.world_size) == state.world_size
+                    and multicast > 0 and multicast % 16 == 0
+                    and signal_addr > 0 and signal_addr % 8 == 0
+                    and (not hasattr(signals, 'numel') or signals.numel() >= state.world_size)
+                    and int(pad_size) >= kernels._MAX_BLOCKS * state.world_size * 4)
+        except (AttributeError, TypeError, ValueError, RuntimeError, OverflowError):
+            return False
+
+    def _prepare(self, device: torch.device, dtype: torch.dtype) -> bool:
         """Collectively rebuild and admit the arena, outside CUDA graphs."""
         device = torch.device(device)
         if device.type == 'cuda' and device.index is None:
             device = torch.device('cuda', torch.cuda.current_device())
         if device.type == 'cuda' and torch.cuda.is_available() and torch.cuda.is_current_stream_capturing():
             raise RuntimeError('symmetric-memory preparation is not graph capturable')
-        self.release()
-        self._weight_contract = (device, dtype)
+        self.close()
         if device.type != 'cuda' or not torch.cuda.is_available():
             return False
-        if not isinstance(self._group, dist.ProcessGroup):
+        if not isinstance(self.group, dist.ProcessGroup):
             return False
         # Even disabled ranks participate: mixed flags/dtypes must not leave
         # peers entering an optional rendezvous on their own.
-        if not self._agree(self._enabled and dtype == torch.bfloat16, device):
+        supported = dtype == torch.bfloat16 or (self._dim == 0 and dtype == torch.float32)
+        if not self._agree(supported, device):
             return False
-        world_size = dist.get_world_size(self._group)
-        width = self._gathered_width
+        world_size = dist.get_world_size(self.group)
+        # The copy kernel addresses BF16 words, but never converts payloads.
+        width = self._gathered_width * dtype.itemsize // torch.bfloat16.itemsize
+        alignment = 8 if self._dim == 0 else world_size * 8
         try:
             capability_ok = torch.cuda.get_device_capability(device) >= (9, 0)
         except (RuntimeError, AssertionError):
             capability_ok = False
         valid = (world_size in self._WORLD_SIZES
-                 and self._rank == dist.get_rank(self._group)
-                 and width > 0 and width % (world_size * 8) == 0
+                 and self._rank == dist.get_rank(self.group)
+                 and self._dim in (0, -1)
+                 and width > 0 and width % alignment == 0
                  and self._capacity_bytes >= width * torch.bfloat16.itemsize
                  and capability_ok)
         if not self._agree(valid, device):
-            return self._disabled('unsupported TP, dtype, width, capacity or device')
-        if not self._same_config((width, self._capacity_bytes), device):
-            return self._disabled('inconsistent TP arena configuration')
+            return self._disabled('unsupported group size, dtype, width, capacity or device')
+        if not self._same_config((width, dtype.itemsize, self._dim,
+                                  self._capacity_bytes, int(self._reuse_sync)), device):
+            return self._disabled('inconsistent group arena configuration')
 
         kernels = None
         try:
@@ -102,7 +128,7 @@ def prepare(self, device: torch.device, dtype: torch.dtype) -> bool:
         buffer = None
         allocation_ok = False
         try:
-            buffer = kernels._allocate_symmetric_buffer(self._group, max_tokens, width, device)
+            buffer = kernels._allocate_symmetric_buffer(self.group, max_tokens, width, device)
             allocation_ok = (buffer.shape == (max_tokens, width)
                              and buffer.dtype == torch.bfloat16 and buffer.device == device
                              and buffer.is_contiguous() and buffer.storage_offset() == 0
@@ -110,11 +136,11 @@ def prepare(self, device: torch.device, dtype: torch.dtype) -> bool:
         except Exception as exc:
             logger.debug('Symmetric allocation failed: %s', exc)
         if not self._agree(allocation_ok, device):
-            return self._disabled('a TP rank could not allocate its arena')
+            return self._disabled('a rank could not allocate its arena')
 
         # Once all ranks commit, rendezvous errors are fatal, not a rank-local
         # reason to fall back and strand peers inside a collective.
-        state = kernels.create_state(self._group, self._rank, max_tokens, width, device, comm_buff=buffer)
+        state = kernels.create_state(self.group, self._rank, max_tokens, width, device, comm_buff=buffer)
         if not self._agree(self._valid_handle(state, kernels), device):
             return self._disabled('invalid symmetric-memory handle or unsupported multicast')
         self._kernels = kernels
@@ -124,31 +150,56 @@ def prepare(self, device: torch.device, dtype: torch.dtype) -> bool:
                            world_size, width, max_tokens)
         return True
 
-    def _valid_handle(self, state, kernels) -> bool:
-        """Validate rank metadata and both pointer and signal-pad bounds."""
-        try:
-            handle = state.symm_mem_hdl
-            multicast = int(handle.multicast_ptr or 0)
-            signals = handle.signal_pad_ptrs_dev
-            signal_addr = int(signals.data_ptr() if hasattr(signals, 'data_ptr') else signals or 0)
-            pad_size = getattr(handle, 'signal_pad_size', None)
-            if pad_size is None:
-                pad_size = kernels.symm_mem.get_signal_pad_size()
-            return (int(handle.rank) == self._rank
-                    and int(handle.world_size) == state.world_size
-                    and multicast > 0 and multicast % 16 == 0
-                    and signal_addr > 0 and signal_addr % 8 == 0
-                    and (not hasattr(signals, 'numel') or signals.numel() >= state.world_size)
-                    and int(pad_size) >= kernels._MAX_BLOCKS * state.world_size * 4)
-        except (AttributeError, TypeError, ValueError, RuntimeError, OverflowError):
-            return False
+    def is_available(self) -> bool:
+        """Whether the symmetric-memory workspace is ready for gathering."""
+        return self._state is not None
+
+    def prepare(self):
+        """Allocate once during setup, never during forward."""
+        if not self._prepared:
+            self._prepare(self.device, self.dtype)
+            self._prepared = True
 
-    def reset_for_weight(self, weight: torch.Tensor) -> None:
-        """Notify the provider after a quiescent, TP-wide model transition."""
-        if (weight.device, weight.dtype) != self._weight_contract:
-            self.prepare(weight.device, weight.dtype)
+    def reset(self, device: torch.device, dtype: torch.dtype):
+        """Rebuild after a coordinated, quiescent device/dtype transition."""
+        device = torch.device(device)
+        if (device, dtype) == (self.device, self.dtype):
+            return
+        self.device, self.dtype = device, dtype
+        if self._prepared:
+            self._prepared = False
+            self.prepare()
+
+    def all_gather(self, input: torch.Tensor, *, dim: int = -1,
+                   copy_output: bool = True) -> torch.Tensor | None:
+        """Return gathered shards, or None when the input is unsupported.
+
+        With copy_output=False, the output borrows the workspace until its next use.
+        """
+        state = self._state
+        if state is None:
+            return None
+        width = self._gathered_width if dim == 0 else self._gathered_width // state.world_size
+        if not (dim == self._dim and input.dim() == 2 and input.dtype == self.dtype
+                and input.is_contiguous() and input.shape[-1] == width):
+            # Shape/dtype/layout contracts must match across ranks, as for NCCL.
+            return None
+        if input.device != state.device or input.data_ptr() % 16 != 0:
+            raise RuntimeError('multimem all-gather device or alignment changed after group-wide admission')
+        rows = input.shape[0]
+        min_rows = 1 if dim == 0 else self._MIN_TOKENS
+        max_rows = state.max_token_num // state.world_size if dim == 0 else state.max_token_num
+        if not min_rows <= rows <= max_rows:
+            return None
+        if rows not in self._graph_ready_shapes and torch.cuda.is_current_stream_capturing():
+            return None
+        output = self._kernels.all_gather_inner(
+            state, input.view(torch.bfloat16), tp_hidden_dim=state.hidden_dim, dim=dim,
+            skip_entry_sync=not self._reuse_sync, copy_output=copy_output, _validated=True)
+        self._graph_ready_shapes.add(rows)
+        return output.view(self.dtype)
 
-    def release(self) -> None:
+    def close(self) -> None:
         """Release only after all consumers/graphs have been retired."""
         state = self._state
         if state is not None:
@@ -160,24 +211,4 @@ def release(self) -> None:
         self._state = None
         self._kernels = None
         self._graph_ready_shapes.clear()
-
-    def __call__(self, x: torch.Tensor) -> torch.Tensor | None:
-        """Gather admitted logits into an owning output; None means NCCL."""
-        state = self._state
-        if state is None:
-            return None
-        if not (x.dim() == 2 and x.dtype == torch.bfloat16 and x.device == state.device
-                and x.is_contiguous() and x.data_ptr() % 16 == 0
-                and x.shape[-1] * state.world_size == state.hidden_dim):
-            # After static TP admission this is a programming error. A local
-            # NCCL fallback could disagree with peers launching multimem.
-            raise RuntimeError('multimem all-gather input contract changed after TP-wide admission')
-        rows = x.shape[0]
-        if not self._MIN_TOKENS <= rows <= state.max_token_num:
-            return None
-        if rows not in self._graph_ready_shapes and torch.cuda.is_current_stream_capturing():
-            return None
-        output = self._kernels.all_gather_inner(
-            state, x, tp_hidden_dim=self._gathered_width, safe=True, _validated=True)
-        self._graph_ready_shapes.add(rows)
-        return output
+        self._prepared = False
diff --git a/lmdeploy/pytorch/backends/cuda/comm/symm_mem_allreduce.py b/lmdeploy/pytorch/backends/cuda/comm/symm_mem_allreduce.py
index 7f1af61c21..1a9e328e8c 100644
--- a/lmdeploy/pytorch/backends/cuda/comm/symm_mem_allreduce.py
+++ b/lmdeploy/pytorch/backends/cuda/comm/symm_mem_allreduce.py
@@ -43,8 +43,7 @@ class SymmetricMemoryAllReduce:
     def __init__(self, group: dist.ProcessGroup):
         self.group = group
         self._buffer = None
-        self._enabled = False
-        self._max_size = 0
+        self._available = False
 
         world_size = dist.get_world_size(group)
         major, minor = torch.cuda.get_device_capability()
@@ -59,38 +58,53 @@ def __init__(self, group: dist.ProcessGroup):
             op_name = 'multimem_all_reduce_' if self._use_multimem else 'two_shot_all_reduce_'
             if not (hasattr(symm_mem, 'empty') and hasattr(symm_mem, 'rendezvous')
                     and hasattr(torch.ops.symm_mem, op_name)):
-                logger.warning('PyTorch symmetric-memory all-reduce is unavailable in this PyTorch build.')
-                return
+                raise RuntimeError('PyTorch symmetric-memory all-reduce is unavailable in this PyTorch build')
+            self._symm_mem = symm_mem
+            self._available = True
+        except (ImportError, RuntimeError) as e:
+            logger.warning(f'PyTorch symmetric-memory all-reduce is unavailable: {e}')
+
+    def is_available(self) -> bool:
+        """Whether the optimized collective implementation is available."""
+        return self._available
+
+    def prepare(self):
+        """Allocate the group buffer during communicator construction."""
+        if self._buffer is not None or not self._available:
+            return
+
+        group = self.group
+        world_size = dist.get_world_size(group)
+        symm_mem = self._symm_mem
+        try:
             self._buffer = symm_mem.empty(
                 self._max_size // torch.bfloat16.itemsize,
                 dtype=torch.bfloat16,
                 device=torch.device('cuda', torch.cuda.current_device()),
             )
-            handle = symm_mem.rendezvous(self._buffer, group.group_name)
-        except (ImportError, RuntimeError) as e:
-            logger.warning(f'PyTorch symmetric-memory all-reduce is unavailable: {e}')
+        except RuntimeError as e:
+            logger.warning(f'PyTorch symmetric-memory allocation failed: {e}')
             self._buffer = None
+
+        ready = [None] * world_size
+        dist.all_gather_object(ready, self._buffer is not None, group=group)
+        if not all(ready):
+            self.close()
             return
 
+        # Once ranks commit, rendezvous failures are fatal rather than a local fallback.
+        handle = symm_mem.rendezvous(self._buffer, group.group_name)
         if getattr(handle, 'multicast_ptr', 0) == 0:
             logger.warning('PyTorch symmetric-memory all-reduce is unavailable: multicast is not supported.')
-            self._buffer = None
+            self.close()
             return
-        self._enabled = True
+
         if dist.get_rank(group) == 0:
             logger.info(f'Using PyTorch symmetric-memory all-reduce for TP{world_size}.')
 
-    def supports(self, dtype: torch.dtype) -> bool:
-        """Whether this group can all-reduce the given dtype."""
-        return self._enabled and dtype == torch.bfloat16
-
-    def is_available(self) -> bool:
-        """Whether the optimized collective implementation is available."""
-        return self._enabled
-
     def all_reduce_(self, input: torch.Tensor) -> bool:
         """All-reduce ``input`` in place, returning whether it was handled."""
-        if (not self.supports(input.dtype) or not input.is_contiguous()
+        if (self._buffer is None or input.dtype != torch.bfloat16 or not input.is_contiguous()
                 or input.nbytes > self._max_size or input.nbytes % 4 != 0):
             return False
 
@@ -106,4 +120,4 @@ def all_reduce_(self, input: torch.Tensor) -> bool:
     def close(self):
         """Release the symmetric buffer."""
         self._buffer = None
-        self._enabled = False
+        self._available = False
diff --git a/lmdeploy/pytorch/backends/cuda/nsa.py b/lmdeploy/pytorch/backends/cuda/nsa.py
index 1c4dfb9ef7..23423097fd 100644
--- a/lmdeploy/pytorch/backends/cuda/nsa.py
+++ b/lmdeploy/pytorch/backends/cuda/nsa.py
@@ -379,6 +379,9 @@ def __init__(self, topk: int, softmax_scale: float, block_size: int,
                 raise RuntimeError('DCP requires the TileLang sparse index top-k kernel.')
             if _get_deep_gemm() is None:
                 raise RuntimeError('DCP DSA scoring requires a compatible DeepGEMM installation.')
+            from .attention.cp import get_dcp_manager
+            self.dcp_manager = get_dcp_manager()
+            self.dcp_manager.prepare_candidate_gather(topk)
         self._step_meta_group: int | None = None
         self._piecewise_forward: Callable[..., Tensor] | None = None
         self._piecewise_forward_fused: Callable[..., Tensor] | None = None
@@ -490,26 +493,16 @@ def _merge_dcp_topk(self, scores: Tensor, local_indices: Tensor) -> Tensor:
         if self.dcp_world_size == 1:
             return local_indices
 
-        from lmdeploy.pytorch.distributed import all_gather_into_tensor
         from lmdeploy.pytorch.kernels.cuda.sparse_index_dcp_topk import (
             pack_dcp_topk_candidates,
             sparse_dcp_global_topk,
         )
 
-        num_rows = scores.size(0)
         packed = pack_dcp_topk_candidates(
             scores,
             local_indices,
             dcp_world_rank=(self.dcp_world_size, self.dcp_rank))
-        gathered = torch.empty(
-            self.dcp_world_size * num_rows,
-            self.topk,
-            2,
-            dtype=packed.dtype,
-            device=scores.device,
-        )
-        all_gather_into_tensor(gathered, packed, group='dcp')
-        gathered = gathered.view(self.dcp_world_size, num_rows, self.topk, 2)
+        gathered = self.dcp_manager.gather_candidates(packed)
         return sparse_dcp_global_topk(gathered, k=self.topk, fill=self.fill)
 
     def _score_and_select_prefill(
diff --git a/lmdeploy/pytorch/backends/cuda/op_backend.py b/lmdeploy/pytorch/backends/cuda/op_backend.py
index 3294b54f4c..e97ab9532f 100644
--- a/lmdeploy/pytorch/backends/cuda/op_backend.py
+++ b/lmdeploy/pytorch/backends/cuda/op_backend.py
@@ -232,13 +232,14 @@ def get_v4_attention_metadata_cls():
         return CudaV4AttentionMetadata
 
     @classmethod
-    def build_communicator(cls, cpu_group, device_group, dist_config):
+    def build_communicator(cls, cpu_group, device_group, dist_config, *, group_name='tp'):
         """Build a CUDA communicator."""
         from .comm.communicator import build_cuda_communicator
         communicator = build_cuda_communicator(
             cpu_group=cpu_group,
             device_group=device_group,
             dist_config=dist_config,
+            group_name=group_name,
         )
         if communicator is not None:
             return communicator
@@ -246,6 +247,7 @@ def build_communicator(cls, cpu_group, device_group, dist_config):
             cpu_group=cpu_group,
             device_group=device_group,
             dist_config=dist_config,
+            group_name=group_name,
         )
 
     @staticmethod
diff --git a/lmdeploy/pytorch/backends/default/embedding.py b/lmdeploy/pytorch/backends/default/embedding.py
index c5596d71f9..49a94b7e4a 100644
--- a/lmdeploy/pytorch/backends/default/embedding.py
+++ b/lmdeploy/pytorch/backends/default/embedding.py
@@ -1,8 +1,9 @@
 # Copyright (c) OpenMMLab. All rights reserved.
 import torch
-import torch.distributed as dist
 import torch.nn.functional as F
 
+import lmdeploy.pytorch.distributed as dist
+
 from ..embedding import EmbeddingImpl
 
 
diff --git a/lmdeploy/pytorch/backends/default/linear.py b/lmdeploy/pytorch/backends/default/linear.py
index e72fcfe4d6..aea55a1363 100644
--- a/lmdeploy/pytorch/backends/default/linear.py
+++ b/lmdeploy/pytorch/backends/default/linear.py
@@ -1,9 +1,10 @@
 # Copyright (c) OpenMMLab. All rights reserved.
 
 import torch
-import torch.distributed as dist
 import torch.nn.functional as F
 
+import lmdeploy.pytorch.distributed as dist
+
 from ..linear import LinearImpl
 
 
diff --git a/lmdeploy/pytorch/backends/dlinfer/op_backend.py b/lmdeploy/pytorch/backends/dlinfer/op_backend.py
index 44582593be..36d8beed76 100644
--- a/lmdeploy/pytorch/backends/dlinfer/op_backend.py
+++ b/lmdeploy/pytorch/backends/dlinfer/op_backend.py
@@ -128,18 +128,6 @@ def get_cache_backend(cls):
         from .cache import DlinferCacheBackend
         return DlinferCacheBackend
 
-    @classmethod
-    def build_communicator(cls, cpu_group, device_group, dist_config):
-        """Build a DLInfer communicator."""
-        from lmdeploy.pytorch import envs
-        cuda_communicator_enabled = envs.enable_flashinfer_allreduce or envs.enable_symm_mem_allreduce
-        assert not cuda_communicator_enabled, 'CUDA communicators are not supported by DLInfer.'
-        return super().build_communicator(
-            cpu_group=cpu_group,
-            device_group=device_group,
-            dist_config=dist_config,
-        )
-
     @staticmethod
     def get_attention_metadata_cls():
         from .attention import DlinferAttentionMetadata
diff --git a/lmdeploy/pytorch/config.py b/lmdeploy/pytorch/config.py
index ea2a78607f..772eef9551 100644
--- a/lmdeploy/pytorch/config.py
+++ b/lmdeploy/pytorch/config.py
@@ -3,7 +3,7 @@
 import re
 from collections.abc import Callable, Mapping
 from dataclasses import dataclass, field
-from typing import Any
+from typing import Any, Literal
 
 import torch
 
@@ -253,6 +253,9 @@ class DistConfig:
     mlp_tp_mode: TPMode = TPMode.DEFAULT
     moe_tp_mode: TPMode = TPMode.DEFAULT
 
+    # communication
+    communication_backend: Literal['nccl', 'auto'] = 'nccl'
+
     def __post_init__(self):
         """Post init."""
         assert self.dp_rank < self.dp
@@ -326,6 +329,7 @@ def from_engine_config(cls, engine_config: PytorchEngineConfig):
             dp_rank=engine_config.dp_rank,
             enable_microbatch=engine_config.enable_microbatch,
             enable_eplb=engine_config.enable_eplb,
+            communication_backend=engine_config.communication_backend,
             tp=engine_config.tp,
             dcp=engine_config.dcp,
             attn_tp=engine_config.attn_tp_size,
@@ -658,7 +662,11 @@ def from_hf_config(
         model_config.dist_config = dist_config
 
         if dist_config.dcp > 1:
-            assert model_config.use_flash_mla, 'DCP requires FlashMLA attention'
+            if not model_config.use_flash_mla:
+                replicas = model_config.num_replicate_key_value_heads
+                assert replicas % dist_config.dcp == 0, (
+                    'GQA DCP groups must share replicated KV heads: '
+                    f'KV replication factor {replicas} must be divisible by dcp {dist_config.dcp}')
             assert model_config.sliding_window < 0, 'DCP does not support sliding-window attention'
             if model_config.mla_index_topk is not None:
                 from lmdeploy.pytorch import envs
diff --git a/lmdeploy/pytorch/distributed.py b/lmdeploy/pytorch/distributed.py
index 525de1bf58..0a87b87e81 100644
--- a/lmdeploy/pytorch/distributed.py
+++ b/lmdeploy/pytorch/distributed.py
@@ -14,6 +14,7 @@
 
 if TYPE_CHECKING:
     from .backends.communicator import DeviceCommunicator
+    from .backends.cuda.attention.cp import DCPManager
 
 
 @dataclass
@@ -27,32 +28,6 @@ class DistGroup:
     gpu_gather_group: dist.ProcessGroup = None
     communicator: 'DeviceCommunicator' = None
 
-    def supports_optimized_all_reduce(self) -> bool:
-        """Whether this group has an optimized all-reduce implementation."""
-        return self.communicator is not None and self.communicator.supports_optimized_all_reduce()
-
-    def supports_fused_all_reduce_residual_rms_norm(self) -> bool:
-        """Whether this group can fuse all-reduce, residual and RMSNorm."""
-        return (self.communicator is not None
-                and self.communicator.supports_fused_all_reduce_residual_rms_norm())
-
-    def try_fused_all_reduce_residual_rms_norm(self,
-                                               input: torch.Tensor,
-                                               residual: torch.Tensor,
-                                               weight: torch.Tensor,
-                                               eps: float):
-        """Run fused all-reduce, residual and RMSNorm, or return ``None``."""
-        if self.communicator is None:
-            return None
-        return self.communicator.try_fused_all_reduce_residual_rms_norm(
-            input=input, residual=residual, weight=weight, eps=eps)
-
-    def all_reduce_(self, tensor: torch.Tensor):
-        """All-reduce a tensor in place on this group."""
-        if self.communicator is not None:
-            return self.communicator.all_reduce_(tensor)
-        return dist.all_reduce(tensor, group=self.gpu_group)
-
     def close(self):
         """Close groups."""
         if self.communicator is not None:
@@ -237,17 +212,20 @@ def _build_dcp_group(context: 'DistContext', timeout: timedelta,
     )
 
 
-def _build_tp_communicators(context: 'DistContext'):
-    """Attach one communicator to each rank-local, unique TP group."""
+def _build_communicators(context: 'DistContext'):
+    """Attach one communicator to each rank-local, unique TP/DCP group."""
     build_communicator = context.communicator_builder
-    groups = (context.attn_tp_group, context.mlp_tp_group, context.moe_tp_group)
-    for group in {id(group): group for group in groups}.values():
+    tp_groups = (context.attn_tp_group, context.mlp_tp_group, context.moe_tp_group)
+    groups = [('tp', group) for group in {id(group): group for group in tp_groups}.values()]
+    groups.append(('dcp', context.dcp_group))
+    for group_name, group in groups:
         if group.gpu_group is None:
             continue
         group.communicator = build_communicator(
             cpu_group=group.cpu_group,
             device_group=group.gpu_group,
             dist_config=context.dist_config,
+            group_name=group_name,
         )
 
 
@@ -262,6 +240,7 @@ class DistContext:
     mlp_tp_group: DistGroup = None
     moe_tp_group: DistGroup = None
     dcp_group: DistGroup = None
+    dcp_manager: 'DCPManager' = None
 
     cpu_group: dist.ProcessGroup = None
     ep_gpu_group: dist.ProcessGroup = None
@@ -330,18 +309,23 @@ def build(cls,
 
         # tp
         _build_tp_group(context, timeout, cpu_backend=cpu_backend, ccl_backend=ccl_backend)
-        _build_tp_communicators(context)
 
         # cp
         _build_dcp_group(context, timeout, cpu_backend=cpu_backend, ccl_backend=ccl_backend)
 
+        # communicator
+        _build_communicators(context)
+
         # ep
         cls._build_ep_group(context, timeout, ccl_backend=ccl_backend)
 
         return context
 
     def close(self):
-        """Close groups."""
+        """Close DCP resources before their process groups."""
+        if self.dcp_manager is not None:
+            self.dcp_manager.close()
+            self.dcp_manager = None
         if not dist.is_initialized():
             return
         if self.dcp_group is not None:
@@ -479,10 +463,30 @@ def get_group(group_type: str, device: str):
         raise RuntimeError(f'Unknown group type: {group_type}')
 
 
+def _try_optimized_all_reduce(tensor, op, group, async_op):
+    """Run an optimized synchronous CUDA SUM when eligible."""
+    if not (tensor.is_cuda and op == ReduceOp.SUM and not async_op and group is not None):
+        return False
+
+    context = get_dist_manager().current_context()
+    # The public API exposes raw process groups, so find the wrapper that owns this one.
+    for tp_group in (context.attn_tp_group, context.mlp_tp_group, context.moe_tp_group):
+        if tp_group is not None and tp_group.gpu_group is group:
+            if tp_group.communicator is not None:
+                tp_group.communicator.all_reduce_(tensor)
+                return True
+            break
+    return False
+
+
 def all_reduce(tensor, op=ReduceOp.SUM, group='tp', async_op=False):
     """All reduce."""
     if isinstance(group, str):
         group = get_group(group, 'gpu')
+
+    if _try_optimized_all_reduce(tensor, op, group, async_op):
+        return None
+
     return dist.all_reduce(tensor, op, group, async_op)
 
 
diff --git a/lmdeploy/pytorch/engine/executor/__init__.py b/lmdeploy/pytorch/engine/executor/__init__.py
index eec27c22c0..a1e370c6da 100644
--- a/lmdeploy/pytorch/engine/executor/__init__.py
+++ b/lmdeploy/pytorch/engine/executor/__init__.py
@@ -107,6 +107,9 @@ def build_executor(
 
     # Finalize cache policy before any executor copies configs to workers or
     # builds backend operators. Target and memory models share CacheConfig.
+    if (dist_config.dcp > 1 and not model_config.use_flash_mla
+            and cache_config.quant_policy not in (QuantPolicy.NONE, QuantPolicy.FP8, QuantPolicy.FP8_E5M2)):
+        raise ValueError('GQA DCP supports only unquantized or per-tensor FP8 KV cache')
     shared_cache_models = [model_config]
     if memdecode_config := misc_config.memdecode_config:
         shared_cache_models.append(memdecode_config.memory_model_config)
diff --git a/lmdeploy/pytorch/engine/executor/base.py b/lmdeploy/pytorch/engine/executor/base.py
index e23dfef8c5..e534c7aeb7 100644
--- a/lmdeploy/pytorch/engine/executor/base.py
+++ b/lmdeploy/pytorch/engine/executor/base.py
@@ -243,7 +243,7 @@ def _get_dcp_workspace_size(self, num_prefill_tokens: int) -> int:
         if self.cache_config.dcp <= 1:
             return 0
 
-        from lmdeploy.pytorch.backends.cp_utils import get_dcp_workspace_size
+        from lmdeploy.pytorch.backends.cp_utils import get_dcp_kv_width, get_dcp_workspace_size
 
         config = self.cache_config
         model = self.model_config
@@ -262,6 +262,7 @@ def _get_dcp_workspace_size(self, num_prefill_tokens: int) -> int:
             dcp_size=config.dcp,
             topk=model.mla_index_topk,
             score_workspace_bytes=self._get_dsa_score_workspace_size(),
+            kv_width=get_dcp_kv_width(model),
         )
 
     def _get_runtime_size(self, free_mems: list[int], cache_block_sizes: list[_WorkerCachePlanSizes],
diff --git a/lmdeploy/pytorch/engine/executor/ray_executor.py b/lmdeploy/pytorch/engine/executor/ray_executor.py
index 3c77d91590..0462eeab02 100644
--- a/lmdeploy/pytorch/engine/executor/ray_executor.py
+++ b/lmdeploy/pytorch/engine/executor/ray_executor.py
@@ -104,9 +104,9 @@ def _needs_symm_mem_device_setup(dist_config: DistConfig) -> bool:
     memory rendezvous identifies allocations by the process-local device ordinal, so TP actors on one host must instead
     inherit the full visibility and select their assigned GPU.
     """
-    from lmdeploy.pytorch.backends.cuda.comm.communicator import should_try_symm_mem
-    return (should_try_symm_mem(dist_config)
-            or (_envs.enable_symm_mem_lmhead and dist_config.attn_tp > 1))
+    if dist_config.dp != 1 or dist_config.ep != 1 or dist_config.enable_microbatch:
+        return False
+    return dist_config.communication_backend == 'auto' and (dist_config.attn_tp > 1 or dist_config.dcp > 1)
 
 
 class RemoteLogger:
diff --git a/lmdeploy/pytorch/envs.py b/lmdeploy/pytorch/envs.py
index e8cc6eb11d..ec890da70e 100644
--- a/lmdeploy/pytorch/envs.py
+++ b/lmdeploy/pytorch/envs.py
@@ -249,13 +249,6 @@ def _patched_get_env(
     fake_capture = env_to_bool('LMDEPLOY_FAKE_CUDA_GRAPH_CAPTURE', False)
     enable_decode_torch_compile = env_to_bool('LMDEPLOY_ENABLE_DECODE_TORCH_COMPILE', False)
 
-    # cuda communicator
-    enable_flashinfer_allreduce = env_to_bool('LMDEPLOY_ENABLE_FLASHINFER_ALLREDUCE', False)
-    enable_symm_mem_allreduce = env_to_bool('LMDEPLOY_ENABLE_SYMM_MEM_ALLREDUCE', False)
-    # NCCL remains the default; launch and small-token policies are internal.
-    enable_symm_mem_lmhead = env_to_bool('LMDEPLOY_ENABLE_SYMM_MEM_LMHEAD', False)
-    symm_mem_lmhead_max_mb = max(1, env_to_int('LMDEPLOY_SYMM_MEM_LMHEAD_MAX_MB', 64))
-
     # opt-ttft
     opt_ttft_policy = env_to_choice('LMDEPLOY_PT_TTFT_POLICY', 'size', {'fifo', 'size'})
     opt_ttft_short_turns = max(1, env_to_int('LMDEPLOY_PT_TTFT_SHORT_TURNS', 3))
diff --git a/lmdeploy/pytorch/kernels/cuda/dcp.py b/lmdeploy/pytorch/kernels/cuda/dcp.py
index ac005b7d47..9952a6e403 100644
--- a/lmdeploy/pytorch/kernels/cuda/dcp.py
+++ b/lmdeploy/pytorch/kernels/cuda/dcp.py
@@ -263,21 +263,23 @@ def _reorder_dcp_prefill_kv_kernel(
     stride_lr,
     stride_ls,
     DCP_SIZE: tl.constexpr,
+    BLOCK_M: tl.constexpr,
     BLOCK_D: tl.constexpr,
 ):
-    source_row = tl.program_id(0)
+    source_row = tl.program_id(0) * BLOCK_M + tl.arange(0, BLOCK_M)
+    row_mask = source_row < local_capacity * DCP_SIZE
     # Each rank contributes the same capacity, including trailing padding.
     rank = source_row // local_capacity
     local_position = source_row % local_capacity
 
     request = 0
-    request_start = 0
-    request_id = 0
-    position_in_request = 0
-    found = False
+    request_start = tl.full((BLOCK_M, ), 0, tl.int32)
+    request_id = tl.full((BLOCK_M, ), 0, tl.int32)
+    position_in_request = tl.full((BLOCK_M, ), 0, tl.int32)
+    found = tl.full((BLOCK_M, ), False, tl.int1)
     while request < num_sequences:
         request_len = tl.load(LocalLens + rank * stride_lr +
-                              request * stride_ls).to(tl.int32)
+                              request * stride_ls, mask=row_mask, other=0).to(tl.int32)
         owns_position = ((local_position >= request_start)
                          & (local_position < request_start + request_len))
         request_id = tl.where(owns_position, request, request_id)
@@ -291,18 +293,18 @@ def _reorder_dcp_prefill_kv_kernel(
     global_position = position_in_request * DCP_SIZE + rank
     chunk_kv_seqlen = tl.load(ChunkKvSeqLens + request_id, mask=found, other=0)
     output_start = tl.load(KvStartLoc + request_id, mask=found, other=0)
-    valid_row = found & (global_position < chunk_kv_seqlen)
+    valid_row = row_mask & found & (global_position < chunk_kv_seqlen)
 
     dim_offsets = tl.arange(0, BLOCK_D)
     dim_mask = dim_offsets < row_width
-    value = tl.load(Gathered + source_row * stride_gs +
-                    dim_offsets * stride_gd,
-                    mask=valid_row & dim_mask,
+    value = tl.load(Gathered + source_row[:, None] * stride_gs +
+                    dim_offsets[None, :] * stride_gd,
+                    mask=valid_row[:, None] & dim_mask[None, :],
                     other=0.0)
     output_row = output_start + global_position
-    tl.store(Output + output_row * stride_os + dim_offsets * stride_od,
+    tl.store(Output + output_row[:, None] * stride_os + dim_offsets[None, :] * stride_od,
              value,
-             mask=valid_row & dim_mask)
+             mask=valid_row[:, None] & dim_mask[None, :])
 
 
 def sanitize_dcp_lse(local_lse: torch.Tensor,
@@ -429,7 +431,9 @@ def reorder_dcp_prefill_kv(gathered: torch.Tensor, output: torch.Tensor, *,
 
     row_width = gathered_rows.size(1)
     block_d = triton.next_power_of_2(row_width)
-    _reorder_dcp_prefill_kv_kernel[(gathered_rows.size(0), )](
+    # Process four rows per CTA to amortize scheduling and request-scan control overhead.
+    block_m = 4
+    _reorder_dcp_prefill_kv_kernel[(triton.cdiv(gathered_rows.size(0), block_m), )](
         gathered_rows,
         output_rows,
         chunk_kv_seqlens,
@@ -442,6 +446,7 @@ def reorder_dcp_prefill_kv(gathered: torch.Tensor, output: torch.Tensor, *,
         *output_rows.stride(),
         *local_lens.stride(),
         DCP_SIZE=dcp_size,
+        BLOCK_M=block_m,
         BLOCK_D=block_d,
         num_warps=8,
     )
diff --git a/lmdeploy/pytorch/kernels/cuda/fill_kv_cache.py b/lmdeploy/pytorch/kernels/cuda/fill_kv_cache.py
index a931fd2f1c..666ec7c1db 100644
--- a/lmdeploy/pytorch/kernels/cuda/fill_kv_cache.py
+++ b/lmdeploy/pytorch/kernels/cuda/fill_kv_cache.py
@@ -237,6 +237,8 @@ def _fill_kv_cache_fp8_scalar_kernel(
     stride_vch: tl.constexpr,
     stride_vcd: tl.constexpr,
     stride_boff,
+    DCP_SIZE: tl.constexpr,
+    DCP_RANK: tl.constexpr,
     BLOCK: tl.constexpr,
     BLOCK_D: tl.constexpr,
     BLOCK_DV: tl.constexpr,
@@ -250,23 +252,22 @@ def _fill_kv_cache_fp8_scalar_kernel(
     q_seqlen = tl.load(QSeqLens + batch_id)
     kv_seqlen = tl.load(KVSeqLens + batch_id)
     history_seqlen = kv_seqlen - q_seqlen
-    kv_block_id = history_seqlen // BLOCK + block_id
+    kv_block_id = history_seqlen // (BLOCK * DCP_SIZE) + block_id
 
     if kv_seqlen <= 0:
         return
-    if kv_block_id * BLOCK >= kv_seqlen:
+    if kv_block_id * BLOCK * DCP_SIZE >= kv_seqlen:
         return
 
     if is_decoding:
-        page_offs = tl.full((1, ), history_seqlen % BLOCK, dtype=tl.int32)
-        kv_mask = tl.full((1, ), 1, dtype=tl.int1)
+        page_offs = tl.full((1, ), (history_seqlen // DCP_SIZE) % BLOCK, dtype=tl.int32)
+        kv_mask = tl.full((1, ), history_seqlen % DCP_SIZE == DCP_RANK, dtype=tl.int1)
         q_offs = tl.full((1, ), q_startloc, dtype=tl.int32)
     else:
         page_offs = tl.arange(0, BLOCK)
-        kv_offs = kv_block_id * BLOCK + page_offs
+        kv_offs = (kv_block_id * BLOCK + page_offs) * DCP_SIZE + DCP_RANK
         kv_mask = (kv_offs >= history_seqlen) & (kv_offs < kv_seqlen)
-        token_off = q_startloc + kv_block_id * BLOCK - history_seqlen
-        q_offs = token_off + page_offs
+        q_offs = q_startloc + kv_offs - history_seqlen
 
     block_off = tl.load(BlockOffsets + batch_id * stride_boff + kv_block_id)
     block_off = block_off.to(tl.int64)
@@ -899,8 +900,8 @@ def fill_kv_cache(k_states: Tensor,
     if v_states.size(-1) == 0:
         head_dim_v = 0
 
-    if dcp_size > 1 and quant_policy != QuantPolicy.NONE:
-        raise ValueError('DCP generic KV fill currently supports only unquantized cache payloads')
+    if dcp_size > 1 and quant_policy not in (QuantPolicy.NONE, QuantPolicy.FP8, QuantPolicy.FP8_E5M2):
+        raise ValueError('DCP generic KV fill supports only unquantized or per-tensor FP8 cache payloads')
     if max_q_seq_length == 1:
         max_num_blocks = 1
     else:
@@ -1021,6 +1022,8 @@ def fill_kv_cache(k_states: Tensor,
             stride_vch=v_caches.stride(h_dim),
             stride_vcd=v_caches.stride(d_dim),
             stride_boff=block_offsets.stride(0),
+            DCP_SIZE=dcp_size,
+            DCP_RANK=dcp_rank,
             BLOCK=BLOCK,
             BLOCK_D=BLOCK_D,
             BLOCK_DV=BLOCK_DV,
diff --git a/lmdeploy/pytorch/kernels/cuda/pagedattention.py b/lmdeploy/pytorch/kernels/cuda/pagedattention.py
index 1b62c591ca..a9637a0929 100644
--- a/lmdeploy/pytorch/kernels/cuda/pagedattention.py
+++ b/lmdeploy/pytorch/kernels/cuda/pagedattention.py
@@ -576,6 +576,8 @@ def _reduce_split_kernel(
     acc_ptr,
     out_ptr,
     sinks_ptr,
+    lse_ptr,
+    cache_seqlens_ptr,
     stride_ak,
     stride_abs,
     stride_ah,
@@ -587,6 +589,8 @@ def _reduce_split_kernel(
     SPLIT_K: tl.constexpr,
     BLOCK_DV: tl.constexpr,
     USE_PDL: tl.constexpr,
+    NUM_HEADS: tl.constexpr,
+    SEQ_LEN: tl.constexpr,
 ):
     """Second step kernel of split k attention."""
     cur_batch = tl.program_id(1)
@@ -609,6 +613,13 @@ def _reduce_split_kernel(
     if USE_PDL:
         tl.extra.cuda.gdc_wait()
 
+    if lse_ptr is not None:
+        # Empty DCP shards do not write split scratch. Never read it.
+        if tl.load(cache_seqlens_ptr + cur_batch // SEQ_LEN) <= 0:
+            tl.store(out_ptr + offs_dv * stride_od, 0.0, mask=mask_dv)
+            tl.store(lse_ptr + cur_batch * NUM_HEADS + cur_head, -float('inf'))
+            return
+
     m_k = tl.load(acc_ptr + offs_mi)
     l_k = tl.load(acc_ptr + offs_mi + 1)
     # (m_k[:, None] > -float('inf')) produce invalid mask for triton 3.5.1
@@ -625,6 +636,10 @@ def _reduce_split_kernel(
     if sinks_ptr is not None:
         sink = tl.load(sinks_ptr + cur_head).to(l_sum.dtype)
         l_sum = l_sum + tl.exp2(sink * tl_log2(math.e) - m_max)
+    if lse_ptr is not None:
+        # Split accumulators use log2; the DCP merge consumes natural-log LSE.
+        lse = (m_max + tl_log2(l_sum)) / tl_log2(math.e)
+        tl.store(lse_ptr + cur_batch * NUM_HEADS + cur_head, lse)
     acc = acc / (l_sum + 1e-10)
 
     out_offs = offs_dv * stride_od
@@ -794,6 +809,7 @@ def flash_attn_with_kvcache(
     quant_policy: QuantPolicy = QuantPolicy.NONE,
     sinks: Tensor = None,
     kv_layout: str = 'bshd',
+    return_lse: bool = False,
 ):
     """Paged Attention forward.
 
@@ -804,8 +820,12 @@ def flash_attn_with_kvcache(
             per-tensor key scale for per-tensor FP8 KV cache.
         v_scales_zeros: Per-token scale/zero metadata for int KV cache, or
             per-tensor value scale for per-tensor FP8 KV cache.
+        return_lse: Return ``(output, lse)`` with natural-log FP32 LSE shaped
+            [tokens, heads]. Empty cache rows return zero output and -inf LSE.
     """
 
+    if return_lse and quant_policy == QuantPolicy.TURBO_QUANT:
+        raise NotImplementedError('TurboQuant paged attention does not return LSE')
     if kv_layout == 'bshd':
         b_dim, s_dim, h_dim, d_dim = (0, 1, 2, 3)
     elif kv_layout == 'bhsd':
@@ -1032,6 +1052,7 @@ def flash_attn_with_kvcache(
 
     num_warps = 2
     grid = (head, num_tokens)
+    lse = torch.empty((num_tokens, head), device=q.device, dtype=torch.float32) if return_lse else None
     if quant_policy == QuantPolicy.INT4:
         Lv *= 2
         BLOCK_DV *= 2
@@ -1063,6 +1084,8 @@ def flash_attn_with_kvcache(
         _reduce_split_kernel[grid](acc,
                                    o,
                                    sinks,
+                                   lse,
+                                   cache_seqlens,
                                    stride_ak=acc.stride(2),
                                    stride_abs=acc.stride(0),
                                    stride_ah=acc.stride(1),
@@ -1074,7 +1097,9 @@ def flash_attn_with_kvcache(
                                    head_size_v=Lv,
                                    BLOCK_DV=BLOCK_DV,
                                    USE_PDL=use_pdl,
+                                   NUM_HEADS=head,
+                                   SEQ_LEN=seq_len,
                                    num_warps=num_warps,
                                    num_stages=1,
                                    **pdl_launch)
-    return o
+    return (o, lse) if return_lse else o
diff --git a/lmdeploy/pytorch/kernels/cuda/symm_mem_allgather.py b/lmdeploy/pytorch/kernels/cuda/symm_mem_allgather.py
index 486f21396a..8d8af29df3 100644
--- a/lmdeploy/pytorch/kernels/cuda/symm_mem_allgather.py
+++ b/lmdeploy/pytorch/kernels/cuda/symm_mem_allgather.py
@@ -451,13 +451,14 @@ def all_gather_inner(
     hidden_states: torch.Tensor,
     tp_hidden_dim: int,
     skip_entry_sync: bool = False,
-    safe: bool = True,
+    copy_output: bool = True,
     *,
+    dim: int = -1,
     _validated: bool = False,
 ) -> torch.Tensor:
-    """Gather ``[T, H/TP]`` shards into ``[T, H]`` along the hidden dim.
+    """Gather 2D BF16 words along dimension 0 or -1 without arithmetic.
 
-    ``tp_hidden_dim`` is the gathered width ``H``. Returns a clone when ``safe``,
+    ``tp_hidden_dim`` is the output row width. Returns a clone when ``copy_output``,
     else a view into the symmetric buffer (valid until the next collective).
     ``_validated`` is reserved for the admitted provider,
     whose admission check already enforces the immutable dtype/layout/width
@@ -471,10 +472,9 @@ def all_gather_inner(
             f"hidden_states.data_ptr()={hex(hidden_states.data_ptr())} must be "
             f"16-byte aligned for 128-bit multimem.st"
         )
-        assert (
-            tp_hidden_dim % world_size == 0
-        ), f"tp_hidden_dim={tp_hidden_dim} must be divisible by world_size={world_size}"
-    local_hidden = tp_hidden_dim // world_size
+        assert dim in (0, -1)
+        assert dim == 0 or tp_hidden_dim % world_size == 0
+    local_hidden = tp_hidden_dim if dim == 0 else tp_hidden_dim // world_size
     total_tokens, in_hidden = hidden_states.shape
     if not _validated:
         assert local_hidden % _NUMEL_PER_THREAD == 0, (
@@ -489,14 +489,20 @@ def all_gather_inner(
             in_hidden == local_hidden
         ), f"input hidden ({in_hidden}) != this rank's shard ({local_hidden})"
         assert (
-            total_tokens <= state.max_token_num
+            total_tokens * (world_size if dim == 0 else 1) <= state.max_token_num
         ), f"total_tokens={total_tokens} exceeds max_token_num={state.max_token_num}"
 
+    logical_rows = total_tokens
+    if dim == 0:
+        # Copy each rank's entire payload as one contiguous shard. Preserve the
+        # logical row count for launch tuning, without a rank/token transpose.
+        local_hidden *= total_tokens
+        total_tokens = 1
     hidden_offset = local_hidden * state.rank_in_group
     symm_mem_hdl = state.symm_mem_hdl
     num_blocks, block_size, num_warps, numel_per_thread = _launch_config(
         total_tokens * local_hidden,
-        total_tokens=total_tokens,
+        total_tokens=logical_rows,
         world_size=world_size,
     )
     split_barrier = _use_single_barrier(num_blocks, world_size)
@@ -527,8 +533,11 @@ def all_gather_inner(
         # The completion barrier can therefore be a single CTA and still
         # publish the whole grid before the output view is returned.
         barrier_inner(state, slot=1, release=True)
-    output = state.comm_buff[:total_tokens, :tp_hidden_dim]
-    return output.clone() if safe else output
+    if dim == 0:
+        output = state.comm_buff.view(-1)[:world_size * local_hidden].view(world_size * logical_rows, tp_hidden_dim)
+    else:
+        output = state.comm_buff[:total_tokens, :tp_hidden_dim]
+    return output.clone() if copy_output else output
 
 
 def barrier_inner(state: MultimemAllGatherState, *, slot: int,
diff --git a/lmdeploy/pytorch/models/deepseek_v32.py b/lmdeploy/pytorch/models/deepseek_v32.py
index c8dfb15fe9..68a25d6d34 100644
--- a/lmdeploy/pytorch/models/deepseek_v32.py
+++ b/lmdeploy/pytorch/models/deepseek_v32.py
@@ -496,22 +496,12 @@ def __init__(self,
         nn.Module.__init__(self)
         self.layer_idx = layer_idx
         quantization_config = None
-        # Row-wise TP outputs normally reduce inside their projections. An
-        # optimized communicator lets the following RMSNorm consume that
-        # reduction instead. Attention is consumed in this layer.
-        defer_attn_all_reduce = RMSNorm.can_handle_all_reduce('attn')
-        # MLP is consumed by the next target layer, so terminal and MTP blocks
-        # must still reduce their outputs.
-        defer_mlp_all_reduce = (layer_idx < config.num_hidden_layers - 1
-                                and RMSNorm.can_handle_all_reduce('mlp'))
-
         # build attention layer
         self.self_attn = self.attention_cls(
             config,
             layer_idx,
             dtype=dtype,
             device=device,
-            all_reduce=not defer_attn_all_reduce,
             prefix=add_prefix('self_attn', prefix))
 
         # mlp
@@ -521,14 +511,12 @@ def __init__(self,
             layer_idx,
             dtype=dtype,
             device=device,
-            all_reduce=not defer_mlp_all_reduce,
             prefix=mlp_prefix) if
                     (config.n_routed_experts is not None and layer_idx >= config.first_k_dense_replace
                      and layer_idx % config.moe_layer_freq == 0) else DeepseekV2MLP(
                          config,
                          dtype=dtype,
                          device=device,
-                         all_reduce=not defer_mlp_all_reduce,
                          prefix=mlp_prefix))
 
         # build input layer norm
@@ -536,15 +524,13 @@ def __init__(self,
                                        config.rms_norm_eps,
                                        quant_config=quantization_config,
                                        dtype=torch.float32,
-                                       device=device,
-                                       all_reduce_group='mlp')
+                                       device=device)
 
         # build attention layer norm
         self.post_attention_layernorm = RMSNorm(config.hidden_size,
                                                 config.rms_norm_eps,
                                                 dtype=torch.float32,
-                                                device=device,
-                                                all_reduce_group='attn')
+                                                device=device)
 
 class DeepseekV32Model(DeepseekV2Model):
     decoder_layer_cls = DeepseekV32DecoderLayer
diff --git a/lmdeploy/pytorch/nn/embedding.py b/lmdeploy/pytorch/nn/embedding.py
index da85d2dab5..2dc3728a23 100644
--- a/lmdeploy/pytorch/nn/embedding.py
+++ b/lmdeploy/pytorch/nn/embedding.py
@@ -1,10 +1,8 @@
 # Copyright (c) OpenMMLab. All rights reserved.
 import torch
-import torch.distributed as dist
 from torch import nn
 
 from lmdeploy.pytorch.backends import get_backend
-from lmdeploy.pytorch.backends.cuda.comm.symm_mem_allgather import MultimemAllGatherer
 from lmdeploy.pytorch.backends.embedding import EmbeddingBuildSpec
 from lmdeploy.pytorch.backends.linear import LinearBuildSpec
 from lmdeploy.pytorch.distributed import get_dist_group, get_dist_manager, get_tp_world_rank
@@ -161,21 +159,28 @@ def __init__(
             enable_deterministic=get_build_model_context().enable_deterministic,
         )
 
-        self._symm_mem_gatherer = (
-            MultimemAllGatherer(self.tp_group, self.tp_rank,
-                               self.tp * self.vocab_size_padded,
-                               self.weight.device, self.weight.dtype)
-            if self.all_reduce else None)
+        communicator = get_dist_group(layer_type=layer_type).communicator
+        self._logits_gather_workspace = None
+        if self.all_reduce:
+            if communicator is None:
+                from lmdeploy.pytorch.backends.communicator import DeviceCommunicator
+                communicator = DeviceCommunicator(self.tp_group)
+            self._logits_gather_workspace = communicator.create_all_gather_workspace(
+                self.tp * self.vocab_size_padded, device=self.weight.device, dtype=self.weight.dtype)
+
+        self._communicator = communicator
 
     def tie_weights(self, embedding: ParallelEmbedding):
         """Tie the local LM-head shard to a parallel embedding shard."""
         self.weight = embedding.weight
+        if self._logits_gather_workspace is not None:
+            self._logits_gather_workspace.reset(self.weight.device, self.weight.dtype)
 
     def _apply(self, fn, recurse=True):
-        """Notify the provider after coordinated model device/dtype moves."""
+        """Update the workspace after coordinated model device/dtype moves."""
         result = super()._apply(fn, recurse=recurse)
-        if self._symm_mem_gatherer is not None:
-            self._symm_mem_gatherer.reset_for_weight(self.weight)
+        if self._logits_gather_workspace is not None:
+            self._logits_gather_workspace.reset(self.weight.device, self.weight.dtype)
         return result
 
     def get_local_logits(self, hidden_states: torch.Tensor):
@@ -189,22 +194,10 @@ def all_gather_logits(self, local_logits: torch.Tensor) -> torch.Tensor:
         if not self.all_reduce:
             return local_logits[..., :self.vocab_size]
 
-        if self._symm_mem_gatherer is not None:
-            gathered = self._symm_mem_gatherer(local_logits.reshape(-1, local_logits.shape[-1]))
-            if gathered is not None:
-                output_shape = local_logits.shape[:-1] + (self.tp * local_logits.shape[-1], )
-                return gathered.reshape(output_shape)[..., :self.vocab_size]
-
-        input_size = local_logits.size()
-        output_size = (input_size[0] * self.tp, ) + input_size[1:]
-        logits = local_logits.new_empty(output_size)
-        dist.all_gather_into_tensor(logits, local_logits, group=self.tp_group)
-        # The collective concatenates dim 0. Move its rank dimension beside
-        # the vocabulary shard before reconstructing the full last dimension.
-        logits = logits.reshape((self.tp, ) + input_size)
-        logits = logits.movedim(0, local_logits.dim() - 1)
-        logits = logits.reshape(input_size[:-1] + (self.tp * input_size[-1], ))
-        return logits[..., :self.vocab_size]
+        gathered = self._communicator.all_gather(
+            local_logits.reshape(-1, local_logits.shape[-1]), workspace=self._logits_gather_workspace)
+        output_shape = local_logits.shape[:-1] + (self.tp * local_logits.shape[-1], )
+        return gathered.reshape(output_shape)[..., :self.vocab_size]
 
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
         """Compute TP-local logits and all-gather them on every rank."""
diff --git a/lmdeploy/pytorch/nn/linear/base.py b/lmdeploy/pytorch/nn/linear/base.py
index 6dbd7b3ed9..2d938e90a8 100644
--- a/lmdeploy/pytorch/nn/linear/base.py
+++ b/lmdeploy/pytorch/nn/linear/base.py
@@ -7,6 +7,7 @@
 
 from lmdeploy.pytorch.config import TPMode
 from lmdeploy.pytorch.distributed import (
+    all_reduce,
     gather_by_tp_sizes,
     get_dist_group,
     get_dist_manager,
@@ -216,7 +217,7 @@ def _forward_lora(self, x, tp_sizes: list[int] = None):
                     out = out.to(self.tp_reduce_dtype)
                     get_dist_group(self.layer_type).all_reduce_(out)
                 else:
-                    dist.all_reduce(out, group=self.tp_group)
+                    all_reduce(out, group=self.tp_group)
                 out = out.to(output_dtype)
         return out
 
diff --git a/lmdeploy/pytorch/nn/norm.py b/lmdeploy/pytorch/nn/norm.py
index 362c2ef6e5..3a17326d2a 100644
--- a/lmdeploy/pytorch/nn/norm.py
+++ b/lmdeploy/pytorch/nn/norm.py
@@ -3,7 +3,7 @@
 import torch
 from torch import nn
 
-from lmdeploy.pytorch.distributed import get_dist_group, get_tp_world_rank
+from lmdeploy.pytorch.distributed import get_tp_world_rank
 from lmdeploy.pytorch.models.patch import get_build_model_context
 
 from ..backends import get_backend
@@ -40,7 +40,6 @@ def __init__(
         tp: bool = False,
         align: int = 1,
         prefix: str = '',
-        all_reduce_group: str | None = None,
     ):
         super().__init__()
         backend = get_backend()
@@ -69,18 +68,6 @@ def __init__(
         if tp:
             self.weight.weight_loader = self.weight_loader
         self.align = align
-        self.eps = eps
-        self._all_reduce_group = None
-        self._fuse_all_reduce = False
-        if all_reduce_group is not None and self.can_handle_all_reduce(all_reduce_group):
-            group = get_dist_group(all_reduce_group)
-            self._all_reduce_group = group
-            self._fuse_all_reduce = group.supports_fused_all_reduce_residual_rms_norm()
-
-    @staticmethod
-    def can_handle_all_reduce(layer_type: str) -> bool:
-        """Whether this norm can consume a pending all-reduce."""
-        return get_dist_group(layer_type).supports_optimized_all_reduce()
 
     def weight_loader(self, param: nn.Parameter, loaded_weight: torch.Tensor):
         """Weight loader."""
@@ -100,17 +87,6 @@ def create_weight(hidden_size: int, dtype: torch.dtype | None = None, device: to
 
     def forward(self, x: torch.Tensor, residual: torch.Tensor = None):
         """forward."""
-        if residual is not None and self._all_reduce_group is not None:
-            if self._fuse_all_reduce:
-                output = self._all_reduce_group.try_fused_all_reduce_residual_rms_norm(
-                    input=x,
-                    residual=residual,
-                    weight=self.weight,
-                    eps=self.eps,
-                )
-                if output is not None:
-                    return output
-            self._all_reduce_group.all_reduce_(x)
         return self.impl.forward(x, self.weight, residual)
 
 
diff --git a/lmdeploy/pytorch/spec_decode/base.py b/lmdeploy/pytorch/spec_decode/base.py
index fe178e64e6..3fc325a816 100644
--- a/lmdeploy/pytorch/spec_decode/base.py
+++ b/lmdeploy/pytorch/spec_decode/base.py
@@ -18,6 +18,7 @@ def _build_draft_dist_ctx(dist_ctx: DistContext, specdecode_config: SpecDecodeCo
         return None
 
     draft_dist_config = specdecode_config.dist_config
+    draft_dist_config.communication_backend = dist_ctx.dist_config.communication_backend
     if specdecode_config.method in ('qwen3_5_mtp', 'hy3_mtp') or draft_dist_config == dist_ctx.dist_config:
         return dist_ctx
 
diff --git a/lmdeploy/pytorch/spec_decode/proposers/base.py b/lmdeploy/pytorch/spec_decode/proposers/base.py
index 42d3996773..5721938c8a 100644
--- a/lmdeploy/pytorch/spec_decode/proposers/base.py
+++ b/lmdeploy/pytorch/spec_decode/proposers/base.py
@@ -73,6 +73,7 @@ def draft_model_forward(
             model_config=model_config,
             cache_config=cache_engine.cache_config,
             kv_caches=kv_caches,
+            kv_quant_policy=cache_engine.cache_config.quant_policy,
         )
         context.block_caches = cache_engine.block_caches
         with ctx_mgr.context(context):
diff --git a/lmdeploy/pytorch/third_party/flash_attn_interface.py b/lmdeploy/pytorch/third_party/flash_attn_interface.py
index 5dbcc2fbec..3d7b11f4b7 100644
--- a/lmdeploy/pytorch/third_party/flash_attn_interface.py
+++ b/lmdeploy/pytorch/third_party/flash_attn_interface.py
@@ -1,13 +1,30 @@
 # Copyright (c) OpenMMLab. All rights reserved.
 import functools
 
+import torch
 from flash_attn_interface import flash_attn_varlen_func as _flash_attn_varlen_func
 from flash_attn_interface import flash_attn_with_kvcache as _flash_attn_with_kvcache
 
 
 @functools.wraps(_flash_attn_varlen_func)
-def flash_attn_varlen_func(*args, **kwargs):
+def flash_attn_varlen_func(*args, return_lse: bool = False, **kwargs):
+    """Normalize FA3 returns with optional natural-log LSE.
+
+    LSE is FP32 [tokens, heads]. Callers requesting LSE provide cu_seqlens_q and cu_seqlens_k as keyword arguments.
+    Empty request partitions have -inf LSE.
+    """
+    if return_lse:
+        kwargs['return_attn_probs'] = True
     output = _flash_attn_varlen_func(*args, **kwargs)
+    if return_lse:
+        from lmdeploy.pytorch.kernels.cuda.dcp import sanitize_dcp_lse
+
+        output, lse = output
+        cu_q, cu_k = kwargs['cu_seqlens_q'], kwargs['cu_seqlens_k']
+        valid_counts = torch.repeat_interleave(cu_k[1:] - cu_k[:-1], cu_q[1:] - cu_q[:-1],
+                                               output_size=output.size(0))
+        lse = sanitize_dcp_lse(lse.transpose(0, 1), valid_counts)
+        return output, lse
     if isinstance(output, tuple):
         # for old api
         return output[0]
diff --git a/tests/pytorch/engine/test_executor_base.py b/tests/pytorch/engine/test_executor_base.py
index 99aa0fab14..c1b5229c49 100644
--- a/tests/pytorch/engine/test_executor_base.py
+++ b/tests/pytorch/engine/test_executor_base.py
@@ -316,7 +316,7 @@ def test_get_min_num_gpu_blocks_rejects_worker_count_mismatch():
 def test_runtime_size_reserves_dcp_peak_phase(monkeypatch, draft_tokens, topk, score_mb, dtype):
     monkeypatch.setattr(executor_base._envs, 'dsa_indexer_max_logits_mb', score_mb)
     executor = object.__new__(ExecutorBase)
-    executor.model_config = SimpleNamespace(mla_index_topk=topk, head_dim=576, v_head_dim=0,
+    executor.model_config = SimpleNamespace(use_flash_mla=True, mla_index_topk=topk, head_dim=576, v_head_dim=0,
                                             num_attention_heads=64, dtype=dtype)
     executor.specdecode_config = SimpleNamespace(num_speculative_tokens=draft_tokens) if draft_tokens else None
     executor.dist_config = SimpleNamespace(attn_tp=8)
diff --git a/tests/pytorch/kernel/test_fill_kv_cache.py b/tests/pytorch/kernel/test_fill_kv_cache.py
index 2a711f45bb..8c191858b3 100644
--- a/tests/pytorch/kernel/test_fill_kv_cache.py
+++ b/tests/pytorch/kernel/test_fill_kv_cache.py
@@ -855,6 +855,65 @@ def quant_policy(self):
         yield QuantPolicy.FP8_E5M2
 
 
+@pytest.mark.skipif(not torch.cuda.is_available(), reason='CUDA required')
+@pytest.mark.parametrize(('quant_policy', 'cache_dtype'), [
+    (QuantPolicy.FP8, torch.float8_e4m3fn), (QuantPolicy.FP8_E5M2, torch.float8_e5m2)])
+@pytest.mark.parametrize('dcp_size', [2, 4])
+@pytest.mark.parametrize('kv_layout', ['bshd', 'bhsd'])
+@pytest.mark.parametrize('q_lengths', [[1, 1, 1, 1], [0, 1, 5, 35]])
+def test_fill_fp8_dcp_matches_unsharded_bytes(quant_policy, cache_dtype, dcp_size, kv_layout, q_lengths):
+    """Check ownership and untouched slots with unequal K/V widths and strided
+    inputs."""
+    from lmdeploy.pytorch.kernels.cuda.fill_kv_cache import fill_kv_cache
+
+    _skip_unsupported_triton_fp8_dtype(cache_dtype)
+    torch.manual_seed(101)
+    block_size, num_heads = 16, 2
+    history = [0, 31, 32, 65]
+    kv_lengths = [h + q for h, q in zip(history, q_lengths)]
+    q_lens = torch.tensor(q_lengths, device='cuda', dtype=torch.int32)
+    kv_lens = torch.tensor(kv_lengths, device='cuda', dtype=torch.int32)
+    starts = torch.nn.functional.pad(q_lens.cumsum(0, dtype=torch.int32), (1, 0))[:-1]
+    k = torch.randn(sum(q_lengths) * 2, num_heads, 128, device='cuda', dtype=torch.bfloat16)[::2]
+    v = torch.randn(sum(q_lengths) * 2, num_heads, 96, device='cuda', dtype=torch.bfloat16)[::2]
+    # Exercise saturation as well as non-unit, independent K/V scales.
+    k[0, 0, :2] = torch.tensor([1e5, -1e5], device='cuda')
+    scales = dict(k_scales_zeros=torch.tensor(0.25, device='cuda'),
+                  v_scales_zeros=torch.tensor(0.5, device='cuda'))
+    global_blocks, global_count = _dcp_block_offsets(kv_lengths, block_size, 1, k.device)
+    local_blocks, local_count = _dcp_block_offsets(kv_lengths, block_size, dcp_size, k.device)
+
+    def allocate(count, width):
+        cache = torch.full((count, block_size, num_heads, width), 7, device='cuda', dtype=cache_dtype)
+        return cache if kv_layout == 'bshd' else cache.transpose(1, 2).contiguous()
+
+    def canonical(cache):
+        return cache if kv_layout == 'bshd' else cache.transpose(1, 2)
+
+    global_k, global_v = allocate(global_count, 128), allocate(global_count, 96)
+    fill_kv_cache(k, v, global_k, global_v, starts, q_lens, kv_lens, max(q_lengths), global_blocks,
+                  quant_policy=quant_policy, kv_layout=kv_layout, **scales)
+    for rank in range(dcp_size):
+        local_k, local_v = allocate(local_count, 128), allocate(local_count, 96)
+        fill_kv_cache(k, v, local_k, local_v, starts, q_lens, kv_lens, max(q_lengths), local_blocks,
+                      quant_policy=quant_policy, kv_layout=kv_layout, dcp_size=dcp_size, dcp_rank=rank, **scales)
+        for actual, global_cache in [(local_k, global_k), (local_v, global_v)]:
+            expected = torch.full_like(canonical(actual).view(torch.uint8),
+                                       torch.tensor(7, dtype=cache_dtype).view(torch.uint8).item())
+            for request, kv_len in enumerate(kv_lengths):
+                num_global = _div_up(kv_len, block_size)
+                num_local = _div_up(kv_len, block_size * dcp_size)
+                if num_local == 0:
+                    continue
+                global_ids = global_blocks[request, :num_global]
+                local_ids = local_blocks[request, :num_local]
+                owned = canonical(global_cache).view(torch.uint8)[global_ids].flatten(0, 1)[rank::dcp_size]
+                padded = expected[local_ids].flatten(0, 1)
+                padded[:owned.size(0)] = owned
+                expected[local_ids] = padded.view(num_local, block_size, num_heads, -1)
+            assert torch.equal(canonical(actual).view(torch.uint8), expected)
+
+
 def _dcp_owned_current_tokens(states, q_seqlens, history_lens, dcp_size,
                               dcp_rank):
     """Pack current-chunk tokens owned by one interleaved DCP rank."""
diff --git a/tests/pytorch/kernel/test_mla_attention.py b/tests/pytorch/kernel/test_mla_attention.py
index e4700fea4b..b6f1e8e293 100644
--- a/tests/pytorch/kernel/test_mla_attention.py
+++ b/tests/pytorch/kernel/test_mla_attention.py
@@ -9,7 +9,7 @@
 from lmdeploy.pytorch.backends.cuda import attention as attention_module
 from lmdeploy.pytorch.backends.cuda.attention import mla as mla_module
 from lmdeploy.pytorch.backends.cuda.attention import sparse_mla as sparse_mla_module
-from lmdeploy.pytorch.backends.cuda.attention.cp import gather_dcp_query
+from lmdeploy.pytorch.backends.cuda.attention.cp import DCPManager
 from lmdeploy.pytorch.backends.cuda.attention.sparse_mla import (
     FlashMLAIndexMapper,
     FlashMLASparseImpl,
@@ -146,6 +146,7 @@ def test_dcp_query_all_gather_preserves_contiguous_head_order(monkeypatch, devic
     if device == 'cuda' and not torch.cuda.is_available():
         pytest.skip('requires CUDA')
     from lmdeploy.pytorch import distributed
+    from lmdeploy.pytorch.backends.communicator import DeviceCommunicator
 
     queries = torch.arange(dcp_size * num_tokens * 3 * 8, device=device, dtype=torch.float32)
     queries = queries.reshape(dcp_size, num_tokens, 3, 8)
@@ -159,11 +160,14 @@ def fake_all_gather(output, input_tensor, group='tp', async_op=False):
         if not strided:
             assert input_tensor.data_ptr() == rank0_query.data_ptr()
         for rank in range(dcp_size):
-            source = queries[rank].transpose(0, 1) if strided else queries[rank]
+            source = queries[rank].flatten(1)
             output[rank * source.size(0):(rank + 1) * source.size(0)].copy_(source)
 
-    monkeypatch.setattr(distributed, 'all_gather_into_tensor', fake_all_gather)
-    gathered = gather_dcp_query(rank0_query, dcp_world_size=dcp_size)
+    monkeypatch.setattr(torch.distributed, 'all_gather_into_tensor', fake_all_gather)
+    monkeypatch.setattr(torch.distributed, 'get_world_size', lambda group: dcp_size)
+    group = distributed.DistGroup(gpu_group='dcp', communicator=DeviceCommunicator('dcp'))
+    manager = DCPManager(group)
+    gathered = manager.gather_query(rank0_query)
 
     expected = torch.cat(list(queries), dim=1)
     assert torch.equal(gathered, expected)
@@ -175,7 +179,7 @@ def fake_all_gather(output, input_tensor, group='tp', async_op=False):
         torch.cuda.synchronize()
         graph = torch.cuda.CUDAGraph()
         with torch.cuda.graph(graph):
-            graph_output = gather_dcp_query(rank0_query, dcp_world_size=dcp_size)
+            graph_output = manager.gather_query(rank0_query)
         queries.add_(1)
         graph.replay()
         torch.cuda.synchronize()
@@ -495,9 +499,16 @@ def test_dcp_cached_prefill_matches_reference_across_chunks(monkeypatch, sparse,
                            max_q_seqlen=3,
                            input_ids=torch.zeros(1, 5),
                            cache_config=SimpleNamespace(block_size=64),
-                           model_config=SimpleNamespace(head_dim=576, mla_index_topk=512 if sparse else None),
+                           model_config=SimpleNamespace(head_dim=576, use_flash_mla=True,
+                                                        mla_index_topk=512 if sparse else None),
                            kv_quant_policy=0)
     monkeypatch.setattr(distributed, 'get_dcp_world_rank', lambda: (dcp_size, 0))
+    # This prefill test simulates ranks without creating process groups.
+    manager = object.__new__(DCPManager)
+    manager._prefix_workspaces = {}
+    manager.prepare_attention = Mock()
+    manager.prepare_prefix_gather = Mock()
+    monkeypatch.setattr(mla_module, 'get_dcp_manager', lambda: manager)
     # One virtual block per chunk, with an empty second request in every chunk.
     monkeypatch.setattr(cp_utils, 'get_dcp_prefill_workspace_size',
                         lambda **kwargs: 2 * 64 * (1 + 2 * dcp_size) * 576 * 2)
@@ -535,9 +546,10 @@ def test_dcp_cached_prefill_matches_reference_across_chunks(monkeypatch, sparse,
         reference_keys[:prefix_length, :, :512] = latent.flatten(-2).to(torch.bfloat16)
     gather_calls = 0
 
-    def gather(output, local, group='tp'):
+    def gather(local, **kwargs):
         nonlocal gather_calls
-        assert group == 'dcp'
+        local = local.view(-1, 1, 576)
+        output = local.new_empty(dcp_size * local.size(0), 1, 576)
         chunk = metadata.dcp_prefix_chunks[gather_calls]
         prefix = reference_keys[chunk.start:min(chunk.start + chunk.size, prefix_length)]
         torch.testing.assert_close(local[:prefix[::dcp_size].size(0)], prefix[::dcp_size], atol=0, rtol=0)
@@ -547,8 +559,9 @@ def gather(output, local, group='tp'):
             remote.zero_()
             remote[:prefix[rank::dcp_size].size(0)].copy_(prefix[rank::dcp_size])
         gather_calls += 1
+        return output.flatten(1)
 
-    monkeypatch.setattr(distributed, 'all_gather_into_tensor', gather)
+    manager.communicator = SimpleNamespace(all_gather=gather)
     indices = None
     if use_sparse:
         # Non-contiguous selections span cached chunks and current tokens.
@@ -632,7 +645,8 @@ def test_dcp_attention_correction_kernel_matches_torch(dtype):
 
 @pytest.mark.skipif(not torch.cuda.is_available(), reason='requires CUDA')
 @pytest.mark.parametrize('dcp_size', [2, 4])
-def test_reorder_dcp_prefill_kv_handles_uneven_requests(monkeypatch, dcp_size):
+@pytest.mark.parametrize('gather_prefix', [False, True], ids=['kernel', 'gather'])
+def test_reorder_dcp_prefill_kv_handles_uneven_requests(monkeypatch, dcp_size, gather_prefix):
     from lmdeploy.pytorch.backends import cp_utils
     from lmdeploy.pytorch.kernels.cuda.dcp import reorder_dcp_prefill_kv
 
@@ -644,7 +658,7 @@ def test_reorder_dcp_prefill_kv_handles_uneven_requests(monkeypatch, dcp_size):
                         lambda **kwargs: len(lengths) * (1 + 2 * dcp_size) * 2)
     chunks = cp_utils.build_dcp_prefix_chunks(
         prefix_lens=prefix_lens, prefix_limit=max(lengths), block_size=1,
-        head_dim=1, dcp_world_rank=(dcp_size, 0))
+        kv_width=1, dcp_world_rank=(dcp_size, 0))
     for chunk in chunks:
         values = [request * 100 + torch.arange(length, device=device)[chunk.start:chunk.start + chunk.size]
                   for request, length in enumerate(lengths)]
@@ -654,15 +668,25 @@ def test_reorder_dcp_prefill_kv_handles_uneven_requests(monkeypatch, dcp_size):
             owned = torch.cat([value[rank::dcp_size] for value in values])
             gathered[rank, :owned.numel(), 0] = owned
         output = torch.full((len(lengths) * chunk.size, 1), -99, dtype=torch.int32, device=device)
-        expected = torch.full_like(output, -99)
+        expected = torch.full_like(output, 0 if gather_prefix else -99)
         valid_values = torch.cat(values)
         expected[:valid_values.numel(), 0] = valid_values
         gathered = gathered.flatten(0, 1)
-        reorder_dcp_prefill_kv(gathered,
-                               output,
-                               chunk_kv_seqlens=chunk.kv_seqlens,
-                               kv_start_loc=chunk.cu_seqlens[:-1],
-                               local_lens=chunk.local_kv_seqlens)
+        if gather_prefix:
+            def all_gather(local, **kwargs):
+                assert local.size(0) == local_capacity
+                return gathered
+
+            manager = object.__new__(DCPManager)
+            manager.communicator = SimpleNamespace(all_gather=all_gather)
+            manager._prefix_workspaces = {}
+            output = manager.gather_prefix(gathered[:local_capacity], chunk, zero_padding=True)
+        else:
+            reorder_dcp_prefill_kv(gathered,
+                                   output,
+                                   chunk_kv_seqlens=chunk.kv_seqlens,
+                                   kv_start_loc=chunk.cu_seqlens[:-1],
+                                   local_lens=chunk.local_kv_seqlens)
         assert torch.equal(output, expected)
 
 
diff --git a/tests/pytorch/spec_decode/test_spec_agent.py b/tests/pytorch/spec_decode/test_spec_agent.py
index ab0160c9c5..299b2ba00c 100644
--- a/tests/pytorch/spec_decode/test_spec_agent.py
+++ b/tests/pytorch/spec_decode/test_spec_agent.py
@@ -473,13 +473,14 @@ def test_spec_model_agent_method_when_enabled():
     assert agent.method == specdecode_config.method
 
 
-def test_qwen35_mtp_reuses_main_dist_context(monkeypatch):
+@pytest.mark.parametrize('backend', ['auto', 'nccl'])
+def test_qwen35_mtp_reuses_main_dist_context(monkeypatch, backend):
     """Qwen3.5 MTP mirrors the target topology, so it should share groups."""
     from lmdeploy.pytorch.config import DistConfig, SpecDecodeConfig
     from lmdeploy.pytorch.distributed import DistContext
     from lmdeploy.pytorch.spec_decode import base as base_mod
 
-    dist_config = DistConfig(dp=2, ep=2)
+    dist_config = DistConfig(dp=2, ep=2, communication_backend=backend)
     dist_ctx = DistContext(rank=1, dp_rank=1, dist_config=dist_config, ep_gpu_group=object())
     specdecode_config = SpecDecodeConfig(model='draft-model',
                                          method='qwen3_5_mtp',
@@ -492,16 +493,18 @@ def fail_build(*args, **kwargs):
     monkeypatch.setattr(base_mod.DistContext, 'build', staticmethod(fail_build))
 
     assert base_mod._build_draft_dist_ctx(dist_ctx, specdecode_config) is dist_ctx
+    assert specdecode_config.dist_config.communication_backend == backend
 
 
-def test_non_qwen35_mtp_builds_draft_dist_context(monkeypatch):
+@pytest.mark.parametrize('backend', ['auto', 'nccl'])
+def test_non_qwen35_mtp_builds_draft_dist_context(monkeypatch, backend):
     """Other speculative methods keep their separate draft distribution
     path."""
     from lmdeploy.pytorch.config import DistConfig, SpecDecodeConfig
     from lmdeploy.pytorch.distributed import DistContext
     from lmdeploy.pytorch.spec_decode import base as base_mod
 
-    main_dist_config = DistConfig(dp=2, ep=2)
+    main_dist_config = DistConfig(dp=2, ep=2, communication_backend=backend)
     draft_dist_config = DistConfig()
     dist_ctx = DistContext(rank=1, dp_rank=1, dist_config=main_dist_config)
     specdecode_config = SpecDecodeConfig(model='draft-model',
@@ -519,6 +522,7 @@ def fake_build(*, rank, dist_config):
 
     assert base_mod._build_draft_dist_ctx(dist_ctx, specdecode_config) is draft_dist_ctx
     assert build_calls == [(dist_ctx.rank, draft_dist_config)]
+    assert draft_dist_config.communication_backend == backend
 
 
 def test_async_model_forward_dp1_non_last_chunk_skips_remaining_spec_forwards():
diff --git a/tests/pytorch/test_cuda_communicator.py b/tests/pytorch/test_cuda_communicator.py
index 4c98aaa146..31b4e9e0a1 100644
--- a/tests/pytorch/test_cuda_communicator.py
+++ b/tests/pytorch/test_cuda_communicator.py
@@ -1,5 +1,4 @@
 # Copyright (c) OpenMMLab. All rights reserved.
-import sys
 from types import SimpleNamespace
 from unittest.mock import Mock
 
@@ -8,194 +7,322 @@
 
 from lmdeploy.pytorch.backends import communicator as base_communicator_module
 from lmdeploy.pytorch.backends.cuda.comm import communicator as communicator_module
-from lmdeploy.pytorch.backends.dlinfer.op_backend import DlinferOpsBackend
-from lmdeploy.pytorch.nn import norm as norm_module
-
-
-class _Collective:
-
-    def __init__(self, *, result=None, handled=False):
-        self._result = result
-        self._handled = handled
-
-    def is_available(self):
-        return True
-
-    def fused_all_reduce_residual_rms_norm(self, **kwargs):
-        return self._result
-
-    def all_reduce_(self, input):
-        return self._handled
-
-
-def _build_communicator(monkeypatch,
-                        *,
-                        backend,
-                        fused_result=None,
-                        handled=False):
-    collective = _Collective(result=fused_result, handled=handled)
-    use_flashinfer = backend == 'flashinfer'
-    monkeypatch.setattr(communicator_module._envs, 'enable_flashinfer_allreduce', use_flashinfer)
-    monkeypatch.setattr(communicator_module._envs, 'enable_symm_mem_allreduce', not use_flashinfer)
-    flashinfer_cls = Mock(return_value=collective)
-    symm_mem_cls = Mock(return_value=collective)
-    monkeypatch.setattr(communicator_module, 'FlashInferAllReduce', flashinfer_cls)
-    monkeypatch.setattr(communicator_module, 'SymmetricMemoryAllReduce', symm_mem_cls)
-    communicator = communicator_module.CudaCommunicator(cpu_group='cpu', device_group='gpu')
-    enabled_cls, disabled_cls = ((flashinfer_cls, symm_mem_cls)
-                                 if use_flashinfer else (symm_mem_cls, flashinfer_cls))
-    enabled_cls.assert_called_once_with('cpu')
-    disabled_cls.assert_not_called()
-    return communicator, collective
-
-
-def _build_norm(monkeypatch, *, fused=False):
-    impl = Mock()
-    backend = Mock()
-    backend.build_op.return_value = impl
-    group = Mock()
-    group.supports_optimized_all_reduce.return_value = True
-    group.supports_fused_all_reduce_residual_rms_norm.return_value = fused
-    monkeypatch.setattr(norm_module, 'get_backend', lambda: backend)
-    monkeypatch.setattr(norm_module, 'get_dist_group', lambda layer_type: group)
-    norm = norm_module.RMSNorm(4, dtype=torch.float32, device='cpu', all_reduce_group='attn')
-    return norm, impl, group
-
-
-def test_cuda_communicator_dispatch(monkeypatch):
-    fused_result = object()
-    communicator, flashinfer = _build_communicator(
-        monkeypatch, backend='flashinfer', fused_result=fused_result)
-    nccl_all_reduce = Mock()
-    monkeypatch.setattr(communicator_module.dist, 'all_reduce', nccl_all_reduce)
-
-    assert communicator.supports_optimized_all_reduce()
-    assert communicator.supports_fused_all_reduce_residual_rms_norm()
-    assert communicator.try_fused_all_reduce_residual_rms_norm(
-        input=torch.ones(1), residual=torch.ones(1), weight=torch.ones(1), eps=1e-6) is fused_result
-
-    input = torch.ones(1)
-    communicator.all_reduce_(input)
-    nccl_all_reduce.assert_called_once_with(input, group='gpu')
 
-    flashinfer._handled = True
+
+def _run_auto_all_gather(rank, rendezvous):
+    from datetime import timedelta
+
+    from torch import distributed as dist
+
+    from lmdeploy.pytorch.config import DistConfig
+
+    torch.cuda.set_device(rank)
+    dist.init_process_group('nccl', init_method=rendezvous, rank=rank, world_size=2,
+                            timeout=timedelta(seconds=60))
+    cpu_group = dist.new_group(backend='gloo')
+    communicator = communicator_module.build_cuda_communicator(
+        cpu_group,
+        dist.group.WORLD,
+        DistConfig(tp=2, dcp=2, communication_backend='auto'),
+        group_name='dcp')
+    try:
+        workspace = communicator.create_all_gather_workspace(
+            256, device=torch.device('cuda'), dtype=torch.bfloat16)
+        assert workspace.is_available()
+
+        rows = (2, 16, 96, 384)
+        inputs = [torch.empty(row, 128, device='cuda', dtype=torch.bfloat16) for row in rows]
+        graph = torch.cuda.CUDAGraph()
+        outputs = []
+        with torch.cuda.graph(graph):
+            for input in inputs:
+                outputs.append(communicator.all_gather(input, workspace=workspace, copy_output=False))
+
+        for step in range(3):
+            for input in inputs:
+                input.fill_(rank + step)
+            graph.replay()
+            torch.cuda.synchronize()
+            for input, output in zip(inputs, outputs):
+                expected = torch.cat((input, input + 1), dim=-1)
+                torch.testing.assert_close(output, expected, rtol=0, atol=0)
+        graph.reset()
+        workspace.close()
+    finally:
+        communicator.close()
+        dist.destroy_process_group()
+
+
+def _run_auto_all_reduce(rank, rendezvous):
+    from datetime import timedelta
+
+    from torch import distributed as dist
+
+    from lmdeploy.pytorch.config import DistConfig
+
+    torch.cuda.set_device(rank)
+    dist.init_process_group('nccl', init_method=rendezvous, rank=rank, world_size=2,
+                            timeout=timedelta(seconds=60))
+    cpu_group = dist.new_group(backend='gloo')
+    communicator = communicator_module.build_cuda_communicator(
+        cpu_group, dist.group.WORLD, DistConfig(tp=2, communication_backend='auto'))
+    try:
+        provider = communicator._all_reduce_provider
+        assert provider is not None
+        provider.all_reduce_ = Mock(wraps=provider.all_reduce_)
+        for dtype in (torch.bfloat16, torch.float16, torch.float32):
+            tensor = torch.empty(7, 4096, device='cuda', dtype=dtype)
+            tensor.fill_(rank + 1)
+            communicator.all_reduce_(tensor)
+            provider.all_reduce_.assert_called_with(tensor)
+            torch.testing.assert_close(tensor, torch.full_like(tensor, 3), rtol=0, atol=0)
+            graph = torch.cuda.CUDAGraph()
+            with torch.cuda.graph(graph):
+                tensor.fill_(rank + 1)
+                communicator.all_reduce_(tensor)
+            for _ in range(3):
+                graph.replay()
+            torch.cuda.synchronize()
+            torch.testing.assert_close(tensor, torch.full_like(tensor, 3), rtol=0, atol=0)
+            graph.reset()
+        # Strided input retains the native NCCL error.
+        strided = torch.empty(7, 8192, device='cuda', dtype=torch.bfloat16)[:, ::2]
+        assert not provider.all_reduce_(strided)
+        with pytest.raises(ValueError, match='contiguous'):
+            communicator.all_reduce_(strided)
+    finally:
+        communicator.close()
+        dist.destroy_process_group()
+
+
+@pytest.mark.skipif(torch.cuda.device_count() < 2, reason='requires two CUDA GPUs')
+def test_auto_all_reduce_eager_and_graph(tmp_path):
+    from torch.distributed._symmetric_memory import DeviceType, _SymmetricMemory
+
+    if any(torch.cuda.get_device_capability(i) != (9, 0)
+           or not _SymmetricMemory.has_multicast_support(DeviceType.CUDA, i) for i in range(2)):
+        pytest.skip('requires SM90 with multicast support')
+    torch.multiprocessing.spawn(_run_auto_all_reduce,
+                                args=(f'file://{tmp_path}/auto_all_reduce',), nprocs=2, join=True)
+
+
+@pytest.mark.skipif(torch.cuda.device_count() < 2, reason='requires two CUDA GPUs')
+def test_auto_all_gather_reuses_workspace(tmp_path):
+    from torch.distributed._symmetric_memory import DeviceType, _SymmetricMemory
+
+    if any(torch.cuda.get_device_capability(i)[0] < 9
+            or not _SymmetricMemory.has_multicast_support(DeviceType.CUDA, i) for i in range(2)):
+        pytest.skip('requires SM90 or newer with multicast support')
+    if not torch.cuda.can_device_access_peer(0, 1):
+        pytest.skip('requires peer access')
+    torch.multiprocessing.spawn(_run_auto_all_gather,
+                                args=(f'file://{tmp_path}/auto_all_gather',), nprocs=2, join=True)
+
+
+@pytest.fixture
+def comm_env(monkeypatch):
+    """Two ranks with available providers; no CUDA or distributed runtime."""
+    monkeypatch.setattr(torch.cuda, 'get_device_name', lambda *args: 'GPU')
+    monkeypatch.setattr(torch.cuda, 'current_device', lambda: 0)
+    monkeypatch.setattr(communicator_module.dist, 'get_world_size', lambda group: 2)
+    monkeypatch.setattr(communicator_module.dist, 'get_rank', lambda group: 0)
+    monkeypatch.setattr(communicator_module.dist, 'all_gather_object',
+                        lambda output, value, group: output.__setitem__(slice(None), [value, value]))
+    provider = Mock()
+    provider.is_available.return_value = True
+    monkeypatch.setattr(communicator_module, 'SymmetricMemoryAllReduce', Mock(return_value=provider))
+    return provider
+
+
+def test_gather_registration_and_native_fallback(monkeypatch):
+    from lmdeploy.pytorch.backends.cuda.attention.cp import get_dcp_manager
+    from lmdeploy.pytorch.distributed import DistContext, DistGroup, get_dist_manager
+
+    module = base_communicator_module
+    monkeypatch.setattr(module.dist, 'get_world_size', lambda group: 2)
+
+    def gather(output, input, group):
+        output.copy_(torch.cat((input, input + 10), dim=0))
+
+    monkeypatch.setattr(module.dist, 'all_gather_into_tensor', gather)
+    communicator = module.DeviceCommunicator('group')
+    group = DistGroup(gpu_group='group', communicator=communicator)
+    context = DistContext(dcp_group=group)
+    with get_dist_manager().context(context):
+        manager = get_dcp_manager()
+        manager.prepare_attention(2, 4)
+    assert manager._query_workspace is None
+    logits = communicator.create_all_gather_workspace(6, device=torch.device('cpu'), dtype=torch.float32)
+    assert logits is None
+    query = torch.arange(48).view(3, 2, 8)[..., ::2]
+    with get_dist_manager().context(context):
+        torch.testing.assert_close(manager.gather_query(query), torch.cat((query, query + 10), dim=1))
+    input = torch.ones(3, 3)
+    torch.testing.assert_close(communicator.all_gather(input, workspace=logits), torch.cat((input, input + 10), dim=-1))
+    torch.testing.assert_close(communicator.all_gather(input, dim=0), torch.cat((input, input + 10), dim=0))
+
+
+def test_gather_preparation_and_weight_transition(monkeypatch):
+    from lmdeploy.pytorch.backends.cuda.comm.symm_mem_allgather import SymmetricMemoryAllGather
+
+    workspace = SymmetricMemoryAllGather('group', 0, 16, device='cpu', dtype=torch.bfloat16, capacity_bytes=1024)
+    prepare = Mock(return_value=False)
+    monkeypatch.setattr(workspace, '_prepare', prepare)
+    workspace.reset('cpu', torch.float16)
+    prepare.assert_not_called()
+    workspace.prepare()
+    workspace.prepare()
+    prepare.assert_called_once_with(torch.device('cpu'), torch.float16)
+    workspace.reset('cpu', torch.float16)
+    prepare.assert_called_once()
+    workspace.reset('cpu', torch.bfloat16)
+    assert prepare.call_count == 2
+    prepare.assert_called_with(torch.device('cpu'), torch.bfloat16)
+    assert not workspace.is_available()
+    workspace.close()
+    workspace.prepare()
+    assert prepare.call_count == 3
+
+
+def test_cuda_gather_dispatch_and_workspace_ownership(monkeypatch, comm_env):
+    from lmdeploy.pytorch.backends.cuda.attention.cp import get_dcp_manager
+    from lmdeploy.pytorch.backends.cuda.comm import symm_mem_allgather
+    from lmdeploy.pytorch.distributed import DistContext, DistGroup, get_dist_manager
+
+    # The context shares a DCP manager; logits workspaces remain caller-owned.
+    factory = Mock(side_effect=lambda *args, **kwargs: Mock())
+    monkeypatch.setattr(symm_mem_allgather, 'SymmetricMemoryAllGather', factory)
+    communicator = communicator_module.CudaCommunicator('cpu', 'gpu', group_name='dcp')
+    group = DistGroup(gpu_group='gpu', communicator=communicator)
+    context = DistContext(dcp_group=group)
+    with get_dist_manager().context(context):
+        manager = get_dcp_manager()
+        manager.prepare_attention(2, 4)
+        manager = get_dcp_manager()
+        manager.prepare_attention(2, 4)
+    query_workspace = manager._query_workspace
+    logits_workspace = communicator.create_all_gather_workspace(6, device=torch.device('cpu'), dtype=torch.float32)
+    assert factory.call_count == 3
+    lse_workspace = manager._lse_workspace
+    lse_workspace.prepare.assert_called_once()
+    query_workspace.prepare.assert_called_once()
+    logits_workspace.prepare.assert_called_once()
+
+    query = torch.arange(24).view(3, 2, 4)
+    gathered_query = torch.cat((query, query + 10), dim=1)
+    query_workspace.all_gather.return_value = gathered_query.flatten(1)
+    with get_dist_manager().context(context):
+        torch.testing.assert_close(manager.gather_query(query), gathered_query)
+    assert query_workspace.all_gather.call_args.kwargs == {'dim': -1, 'copy_output': False}
+
+    native = Mock(side_effect=lambda output, input, group: output.copy_(torch.cat((input, input + 10))))
+    monkeypatch.setattr(base_communicator_module.dist, 'all_gather_into_tensor', native)
+    input = torch.ones(3, 3)
+    logits_workspace.all_gather.return_value = torch.ones(3, 6)
+    torch.testing.assert_close(communicator.all_gather(input, workspace=logits_workspace), torch.ones(3, 6))
+    logits_workspace.all_gather.assert_called_once_with(input, dim=-1, copy_output=True)
+    native.assert_not_called()
+    logits_workspace.all_gather.return_value = None
+    torch.testing.assert_close(communicator.all_gather(input, workspace=logits_workspace),
+                               torch.cat((input, input + 10), dim=-1))
+    native.assert_called_once()
+
+    manager.prepare_candidate_gather(512)
+    candidate_workspace = manager._candidate_workspace
+    manager.prepare_candidate_gather(512)
+    assert manager._candidate_workspace is candidate_workspace
+    manager.prepare_prefix_gather(576, 128, 576)
+    prefix_workspaces = list(manager._prefix_workspaces.values())
+    assert len(prefix_workspaces) == 2
+    assert context.dcp_manager is manager
+    monkeypatch.setattr(torch.distributed, 'is_initialized', lambda: True)
+    group_close = Mock(wraps=group.close)
+    monkeypatch.setattr(group, 'close', group_close)
+    query_workspace.close.side_effect = lambda: group_close.assert_not_called()
+    lse_workspace.close.side_effect = lambda: group_close.assert_not_called()
+    candidate_workspace.close.side_effect = lambda: group_close.assert_not_called()
+    for workspace in prefix_workspaces:
+        workspace.close.side_effect = lambda: group_close.assert_not_called()
+    context.close()
+    context.close()
+    query_workspace.close.assert_called_once()
+    lse_workspace.close.assert_called_once()
+    candidate_workspace.close.assert_called_once()
+    for workspace in prefix_workspaces:
+        workspace.close.assert_called_once()
+    assert not manager._prefix_workspaces
+    assert context.dcp_manager is None
+    assert manager._query_workspace is manager._lse_workspace is manager._candidate_workspace is None
+    logits_workspace.close.assert_not_called()
+
+
+@pytest.mark.parametrize('native', ['none', 'async', 'max', 'cpu', 'other_group', 'no_communicator'])
+def test_public_all_reduce_routes_only_synchronous_cuda_tp_sum(monkeypatch, native):
+    import lmdeploy.pytorch.distributed as distributed
+
+    communicator = Mock()
+    group = object()
+    context = distributed.DistContext(
+        attn_tp_group=distributed.DistGroup(gpu_group=group, communicator=communicator))
+    if native == 'no_communicator':
+        context.attn_tp_group.communicator = None
+    tensor = SimpleNamespace(is_cuda=native != 'cpu')
+    op = distributed.ReduceOp.MAX if native == 'max' else distributed.ReduceOp.SUM
+    async_op = native == 'async'
+    selected_group = object() if native == 'other_group' else group
+    native_reduce = Mock()
+    monkeypatch.setattr(distributed.dist, 'all_reduce', native_reduce)
+    with distributed.get_dist_manager().context(context):
+        result = distributed.all_reduce(tensor, op=op, group=selected_group, async_op=async_op)
+    if native == 'none':
+        communicator.all_reduce_.assert_called_once_with(tensor)
+        native_reduce.assert_not_called()
+    else:
+        native_reduce.assert_called_once_with(tensor, op, selected_group, async_op)
+        communicator.all_reduce_.assert_not_called()
+        assert result is native_reduce.return_value
+
+
+def test_all_reduce_workspace_failure_agrees_across_ranks(monkeypatch, comm_env):
+    def peer_failure(flags, ready, group):
+        flags[:] = [ready, not comm_env.prepare.called]
+
+    monkeypatch.setattr(communicator_module.dist, 'all_gather_object', peer_failure)
+    communicator = communicator_module.CudaCommunicator('cpu', 'gpu', group_name='tp')
+    comm_env.close.assert_called_once()
+    comm_env.prepare.assert_called_once()
+    assert communicator._all_reduce_provider is None
+    native = Mock()
+    monkeypatch.setattr(communicator_module.dist, 'all_reduce', native)
+    input = torch.ones(2, 4)
     communicator.all_reduce_(input)
-    nccl_all_reduce.assert_called_once()
+    native.assert_called_once_with(input, group='gpu')
+
 
-    communicator, symm_mem = _build_communicator(monkeypatch, backend='symm_mem', handled=True)
+@pytest.mark.parametrize('handled', [True, False])
+def test_auto_all_reduce_falls_back_for_unsupported_inputs(monkeypatch, comm_env, handled):
+    communicator = communicator_module.CudaCommunicator('cpu', 'gpu', group_name='tp')
+    comm_env.all_reduce_.return_value = handled
+    native = Mock()
+    monkeypatch.setattr(communicator_module.dist, 'all_reduce', native)
+    comm_env.prepare.assert_called_once()
+    input = torch.ones(7, 128, dtype=torch.bfloat16)
     communicator.all_reduce_(input)
-    assert communicator.supports_optimized_all_reduce()
-    assert not communicator.supports_fused_all_reduce_residual_rms_norm()
-    nccl_all_reduce.assert_called_once()
-
-
-def test_cuda_communicator_rejects_multiple_backends(monkeypatch):
-    monkeypatch.setattr(communicator_module._envs, 'enable_flashinfer_allreduce', True)
-    monkeypatch.setattr(communicator_module._envs, 'enable_symm_mem_allreduce', True)
-    with pytest.raises(ValueError, match='cannot be enabled together'):
-        communicator_module.CudaCommunicator(cpu_group='cpu', device_group='gpu')
-
-
-def test_rms_norm_fuses_or_falls_back(monkeypatch):
-    norm, impl, group = _build_norm(monkeypatch, fused=True)
-    fused_output = (torch.full((1, 4), 2.0), torch.full((1, 4), 3.0))
-    group.try_fused_all_reduce_residual_rms_norm.return_value = fused_output
-    input = torch.ones((1, 4), dtype=torch.bfloat16)
-    residual = torch.ones_like(input)
-
-    assert norm(input, residual) is fused_output
-    group.try_fused_all_reduce_residual_rms_norm.assert_called_once_with(
-        input=input,
-        residual=residual,
-        weight=norm.weight,
-        eps=norm.eps,
-    )
-    group.all_reduce_.assert_not_called()
-    impl.forward.assert_not_called()
-
-    group.reset_mock()
-    impl.reset_mock()
-    group.try_fused_all_reduce_residual_rms_norm.return_value = None
-    impl.forward.return_value = object()
-    output = norm(input, residual)
-
-    assert output is impl.forward.return_value
-    group.try_fused_all_reduce_residual_rms_norm.assert_called_once()
-    group.all_reduce_.assert_called_once_with(input)
-    impl.forward.assert_called_once_with(input, norm.weight, residual)
-
-
-def test_flashinfer_allreduce_in_place_and_dtype_guards(monkeypatch):
-    from lmdeploy.pytorch.backends.cuda.comm import flashinfer_allreduce as flashinfer_module
-
-    FlashInferAllReduce = flashinfer_module.FlashInferAllReduce
-
-    flashinfer = FlashInferAllReduce.__new__(FlashInferAllReduce)
-    flashinfer.is_available = Mock(return_value=True)
-    assert flashinfer.supports(torch.float16)
-    assert flashinfer.supports(torch.bfloat16)
-    assert not flashinfer.supports(torch.float32)
-
-    flashinfer._max_size = 1024
-    flashinfer._one_shot_max_size = 1024
-    flashinfer._comm = SimpleNamespace(
-        AllReduceFusionPattern=SimpleNamespace(kAllReduce=0),
-        allreduce_fusion=Mock(
-            return_value=torch.full((2, 4), 2.0, dtype=torch.bfloat16)),
-    )
-    flashinfer._get_workspace = Mock(return_value='workspace')
-    flashinfer.supports = Mock(return_value=True)
-
-    input = torch.ones(1, 2, 4, dtype=torch.bfloat16)
-    assert flashinfer.all_reduce_(input)
-    torch.testing.assert_close(input, torch.full_like(input, 2.0))
-    call_kwargs = flashinfer._comm.allreduce_fusion.call_args.kwargs
-    assert call_kwargs['input'].shape == (2, 4)
-    assert call_kwargs['trigger_completion_at_end']
-    assert call_kwargs['use_oneshot']
-
-    fused_calls = flashinfer._comm.allreduce_fusion.call_count
-    bf16_input = torch.ones(2, 4, dtype=torch.bfloat16)
-    output = flashinfer.fused_all_reduce_residual_rms_norm(
-        input=bf16_input,
-        residual=torch.ones_like(bf16_input),
-        weight=torch.ones(4, dtype=torch.float32),
-        eps=1e-6,
-    )
-    assert output is None
-    assert flashinfer._comm.allreduce_fusion.call_count == fused_calls
-
-    unavailable = FlashInferAllReduce.__new__(FlashInferAllReduce)
-    unavailable._comm = None
-    unavailable._disabled = False
-    unavailable._max_size = 1024
-    monkeypatch.setitem(sys.modules, 'flashinfer.comm', None)
-    assert not unavailable.is_available()
-    assert unavailable._disabled
-
-    workspace_error = FlashInferAllReduce.__new__(FlashInferAllReduce)
-    workspace_error.group = 'cpu'
-    workspace_error._world_size = 2
-    workspace_error._max_size = 1024
-    workspace_error._one_shot_max_size = 1024
-    workspace_error._workspace = None
-    workspace_error._hidden_dim = None
-    workspace_error._dtype = None
-    workspace_error._disabled = False
-    create_workspace = Mock(side_effect=RuntimeError('unsupported topology'))
-    workspace_error._comm = SimpleNamespace(
-        AllReduceFusionPattern=SimpleNamespace(kAllReduce=0),
-        create_allreduce_fusion_workspace=create_workspace,
-        allreduce_fusion=Mock(),
-    )
-    monkeypatch.setattr(flashinfer_module.dist, 'get_rank', lambda group: 0)
+    comm_env.all_reduce_.assert_called_once_with(input)
+    if handled:
+        native.assert_not_called()
+    else:
+        native.assert_called_once_with(input, group='gpu')
+
 
-    input = torch.ones(2, 4, dtype=torch.bfloat16)
-    assert not workspace_error.all_reduce_(input)
-    assert workspace_error._disabled
-    assert not workspace_error.all_reduce_(input)
-    create_workspace.assert_called_once()
-    workspace_error._comm.allreduce_fusion.assert_not_called()
+@pytest.mark.parametrize('group_name,available', [('tp', True), ('tp', False), ('dcp', True)])
+def test_auto_initializes_all_reduce_by_availability(comm_env, group_name, available):
+    comm_env.is_available.return_value = available
+    communicator = communicator_module.CudaCommunicator('cpu', 'gpu', group_name=group_name)
+    assert communicator._all_reduce_provider is (comm_env if available and group_name == 'tp' else None)
+    assert communicator_module.SymmetricMemoryAllReduce.call_count == (group_name == 'tp')
+    assert comm_env.close.call_count == (group_name == 'tp' and not available)
+    assert comm_env.prepare.call_count == (group_name == 'tp' and available)
 
 
 def test_symm_mem_allreduce_selects_group_algorithm(monkeypatch):
@@ -215,7 +342,6 @@ def test_symm_mem_allreduce_selects_group_algorithm(monkeypatch):
     communicator = SymmetricMemoryAllReduce.__new__(SymmetricMemoryAllReduce)
     communicator.group = SimpleNamespace(group_name='group')
     communicator._buffer = torch.empty(8, dtype=torch.bfloat16)
-    communicator._enabled = True
     communicator._max_size = communicator._buffer.nbytes
     input = torch.ones(4, dtype=torch.bfloat16)
 
@@ -229,40 +355,26 @@ def test_symm_mem_allreduce_selects_group_algorithm(monkeypatch):
     two_shot.assert_called_once()
 
 
-def test_build_cuda_communicator_gates_unsupported_parallelism(monkeypatch):
-    monkeypatch.setattr(communicator_module._envs, 'enable_flashinfer_allreduce', True)
-    monkeypatch.setattr(communicator_module._envs, 'enable_symm_mem_allreduce', False)
-    communicator_cls = Mock(return_value=object())
-    monkeypatch.setattr(communicator_module, 'CudaCommunicator', communicator_cls)
-
-    for config in (
-            SimpleNamespace(dp=2, ep=1, attn_tp=8, enable_microbatch=False),
-            SimpleNamespace(dp=1, ep=2, attn_tp=8, enable_microbatch=False),
-            SimpleNamespace(dp=1, ep=1, attn_tp=8, enable_microbatch=True),
-    ):
-        assert communicator_module.build_cuda_communicator('cpu', 'device', config) is None
-    communicator_cls.assert_not_called()
-
-    config = SimpleNamespace(dp=1, ep=1, attn_tp=4, enable_microbatch=False)
-    communicator = communicator_module.build_cuda_communicator('cpu', 'device', config)
-    assert communicator is communicator_cls.return_value
-    communicator_cls.assert_called_once_with(cpu_group='cpu', device_group='device')
-
-    monkeypatch.setattr(communicator_module._envs, 'enable_flashinfer_allreduce', False)
-    monkeypatch.setattr(communicator_module._envs, 'enable_symm_mem_allreduce', True)
-    config.attn_tp = 1
-    assert communicator_module.build_cuda_communicator('cpu', 'device', config) is None
-    config.attn_tp = 2
-    assert communicator_module.build_cuda_communicator('cpu', 'device', config) is communicator_cls.return_value
-
-
-def test_dlinfer_communicator_rejects_cuda_options(monkeypatch):
-    monkeypatch.setattr(communicator_module._envs, 'enable_flashinfer_allreduce', True)
-    monkeypatch.setattr(communicator_module._envs, 'enable_symm_mem_allreduce', False)
-    with pytest.raises(AssertionError, match='not supported by DLInfer'):
-        DlinferOpsBackend.build_communicator('cpu', 'device', SimpleNamespace())
-
-    monkeypatch.setattr(communicator_module._envs, 'enable_flashinfer_allreduce', False)
-    communicator = DlinferOpsBackend.build_communicator('cpu', 'device', SimpleNamespace())
-    assert isinstance(communicator, base_communicator_module.DeviceCommunicator)
-    assert communicator.device_group == 'device'
+def test_symm_mem_allreduce_peer_allocation_failure(monkeypatch):
+    import torch.distributed._symmetric_memory as symm_mem
+
+    from lmdeploy.pytorch.backends.cuda.comm import symm_mem_allreduce as module
+
+    monkeypatch.setattr(module.dist, 'get_world_size', lambda group: 2)
+    monkeypatch.setattr(torch.cuda, 'get_device_capability', lambda: (9, 0))
+    monkeypatch.setattr(torch.cuda, 'current_device', lambda: 0)
+    monkeypatch.setattr(torch.ops, 'symm_mem', SimpleNamespace(two_shot_all_reduce_=Mock()))
+    monkeypatch.setattr(symm_mem, 'empty', Mock(return_value=object()))
+    rendezvous = Mock()
+    monkeypatch.setattr(symm_mem, 'rendezvous', rendezvous)
+
+    def peer_failure(flags, ready, group):
+        flags[:] = [ready, False]
+
+    monkeypatch.setattr(module.dist, 'all_gather_object', peer_failure)
+    provider = module.SymmetricMemoryAllReduce(SimpleNamespace(group_name='group'))
+    symm_mem.empty.assert_not_called()
+    provider.prepare()
+    assert not provider.is_available()
+    assert provider._buffer is None
+    rendezvous.assert_not_called()