feat: extend DCP to GQA and optimize query communication

docs/en/advance/context_parallel.md

@@ -34,26 +34,38 @@ 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`
 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,

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)

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

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."""

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."""

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

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)

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,20 +71,31 @@ 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:
         return ()
     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

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