[Feature] DCP: A2A + FlashInfer-MNNVL comm backends and q-replicate (Helix) (#21637)

Co-authored-by: Hao Phan <htphan@nvidia.com>
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
Co-authored-by: kpham-sgl <khoa.pham@radixark.ai>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Baizhou Zhang <sobereddiezhang@gmail.com>

python/sglang/kernels/ops/attention/dcp_kernels.py

@@ -20,7 +20,7 @@ Consolidated from the two merged DCP implementations:
   - _correct_attn_cp_out_kernel / correct_attn_out / CPTritonContext  (PR #14194)
 """
 
-from typing import Optional
+from typing import Optional, Tuple
 
 import torch
 import triton
@@ -331,3 +331,198 @@ def correct_attn_out(
 
     ctx.call_kernel(_correct_attn_cp_out_kernel, grid, *regular_args, **const_args)
     return new_output, lse
+
+
+# A2A DCP reduce: LSE-weighted combine of N partial attention outputs
+# (used by the a2a / fi_a2a communication backends, see comm.py).
+
+
+def _lse_pack_dim(output_dtype: torch.dtype) -> int:
+    """Number of output-dtype elements needed to store one fp32 LSE value."""
+    return torch.finfo(torch.float32).bits // torch.finfo(output_dtype).bits
+
+
+@triton.jit
+def _dcp_lse_combine_kernel(
+    recv_output_ptr,
+    recv_lse_ptr,
+    out_ptr,
+    out_lse_ptr,
+    recv_output_stride_N,
+    recv_output_stride_B,
+    recv_output_stride_H,
+    recv_output_stride_D,
+    recv_lse_stride_N,
+    recv_lse_stride_B,
+    recv_lse_stride_H,
+    out_stride_B,
+    out_stride_H,
+    out_stride_D,
+    N: tl.constexpr,
+    HEAD_DIM: tl.constexpr,
+    IS_BASE_E: tl.constexpr,
+    RETURN_LSE: tl.constexpr,
+):
+    """Combine N partial attention outputs weighted by their LSE values.
+
+    Grid: (B, H_local).
+    Each program handles one (batch, head) position across all N shards.
+
+    Two-pass approach:
+    Pass 1: find max LSE and weight sum across shards
+    Pass 2: accumulate weighted outputs
+    """
+    batch_idx = tl.program_id(0).to(tl.int64)
+    head_idx = tl.program_id(1).to(tl.int64)
+    d_offsets = tl.arange(0, HEAD_DIM)
+
+    lse_base = batch_idx * recv_lse_stride_B + head_idx * recv_lse_stride_H
+
+    # Pass 1: find max LSE across N shards
+    lse_max = tl.load(recv_lse_ptr + lse_base).to(tl.float32)
+    lse_max = tl.where(
+        (lse_max != lse_max) | (lse_max == float("inf")), -float("inf"), lse_max
+    )
+    for i in tl.static_range(1, N):
+        lse_i = tl.load(recv_lse_ptr + lse_base + i * recv_lse_stride_N).to(tl.float32)
+        lse_i = tl.where(
+            (lse_i != lse_i) | (lse_i == float("inf")), -float("inf"), lse_i
+        )
+        lse_max = tl.where(lse_i > lse_max, lse_i, lse_max)
+
+    lse_max = tl.where(lse_max == -float("inf"), 0.0, lse_max)
+
+    # Pass 2: accumulate weighted outputs
+    weight_sum = tl.zeros([], dtype=tl.float32)
+    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)
+
+    for i in tl.static_range(N):
+        lse_i = tl.load(recv_lse_ptr + lse_base + i * recv_lse_stride_N).to(tl.float32)
+        lse_i = tl.where(
+            (lse_i != lse_i) | (lse_i == float("inf")), -float("inf"), lse_i
+        )
+        centered = lse_i - lse_max
+        if IS_BASE_E:
+            w = tl.exp(centered)
+        else:
+            w = tl.exp2(centered)
+        weight_sum += w
+
+        o_offsets = (
+            i * recv_output_stride_N
+            + batch_idx * recv_output_stride_B
+            + head_idx * recv_output_stride_H
+            + d_offsets * recv_output_stride_D
+        )
+        partial_out = tl.load(recv_output_ptr + o_offsets).to(tl.float32)
+        acc += partial_out * w
+
+    acc = acc / weight_sum
+
+    out_offsets = (
+        batch_idx * out_stride_B + head_idx * out_stride_H + d_offsets * out_stride_D
+    )
+    tl.store(out_ptr + out_offsets, acc.to(out_ptr.dtype.element_ty))
+
+    if RETURN_LSE:
+        if IS_BASE_E:
+            global_lse = tl.log(weight_sum) + lse_max
+        else:
+            global_lse = tl.log2(weight_sum) + lse_max
+        out_lse_offset = batch_idx * recv_lse_stride_B + head_idx * recv_lse_stride_H
+        tl.store(out_lse_ptr + out_lse_offset, global_lse)
+
+
+def dcp_lse_combine_triton(
+    recv_output: torch.Tensor,
+    recv_lse: torch.Tensor,
+    is_lse_base_on_e: bool = True,
+    return_lse: bool = False,
+) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+    """Launch the Triton LSE-combine kernel.
+
+    Args:
+        recv_output: [N, B, H_local, D] partial outputs from each DCP rank.
+        recv_lse:    [N, B, H_local]    log-sum-exp from each DCP rank.
+        is_lse_base_on_e: True if LSE uses base-e (FlashAttention),
+                          False if base-2 (FlashInfer).
+        return_lse: If True, also return the combined global LSE.
+
+    Returns:
+        (combined_output [B, H_local, D], combined_lse [B, H_local] or None)
+    """
+    N, B, H_local, D = recv_output.shape
+    out = torch.empty(
+        (B, H_local, D), device=recv_output.device, dtype=recv_output.dtype
+    )
+    out_lse = (
+        torch.empty((B, H_local), device=recv_lse.device, dtype=recv_lse.dtype)
+        if return_lse
+        else recv_lse.new_empty(0)
+    )
+
+    grid = (B, H_local)
+    _dcp_lse_combine_kernel[grid](
+        recv_output,
+        recv_lse,
+        out,
+        out_lse,
+        recv_output.stride(0),
+        recv_output.stride(1),
+        recv_output.stride(2),
+        recv_output.stride(3),
+        recv_lse.stride(0),
+        recv_lse.stride(1),
+        recv_lse.stride(2),
+        out.stride(0),
+        out.stride(1),
+        out.stride(2),
+        N=N,
+        HEAD_DIM=D,
+        IS_BASE_E=is_lse_base_on_e,
+        RETURN_LSE=return_lse,
+    )
+    return out, (out_lse if return_lse else None)
+
+
+def _lse_weighted_combine_cpu(
+    partial_outputs: torch.Tensor,
+    partial_lses: torch.Tensor,
+    is_lse_base_on_e: bool = True,
+) -> torch.Tensor:
+    """CPU reference: combine N partial attention outputs using LSE weights.
+
+    Args:
+        partial_outputs: [N, B, H_local, D]
+        partial_lses:    [N, B, H_local]
+        is_lse_base_on_e: base-e (True) or base-2 (False)
+
+    Returns:
+        [B, H_local, D] combined output
+    """
+    N, B, H_local, D = partial_outputs.shape
+    partial_outputs = partial_outputs.float()
+    partial_lses = partial_lses.float()
+
+    # Sanitize
+    partial_lses = torch.where(
+        torch.isnan(partial_lses) | torch.isinf(partial_lses),
+        torch.full_like(partial_lses, float("-inf")),
+        partial_lses,
+    )
+
+    # max LSE for numerical stability
+    lse_max, _ = partial_lses.max(dim=0)
+    lse_max = torch.where(lse_max == float("-inf"), torch.zeros_like(lse_max), lse_max)
+
+    centered = partial_lses - lse_max.unsqueeze(0)
+    if is_lse_base_on_e:
+        weights = torch.exp(centered)
+    else:
+        weights = torch.pow(2.0, centered)
+
+    weight_sum = weights.sum(dim=0, keepdim=True)
+    weights = weights / weight_sum
+
+    combined = (partial_outputs * weights.unsqueeze(-1)).sum(dim=0)
+    return combined

python/sglang/srt/distributed/device_communicators/pynccl.py

@@ -303,6 +303,60 @@ class PyNcclCommunicator:
             cudaStream_t(stream.cuda_stream),
         )
 
+    def all_to_all_single(
+        self,
+        output_tensor: torch.Tensor,
+        input_tensor: torch.Tensor,
+    ):
+        """All-to-All over the flattened leading dim: each rank sends the i-th
+        equal-sized chunk to rank i and receives rank i's chunk into output
+        position i. Uses ncclGroupStart/End to fuse the sends/recvs into a
+        single NCCL operation, which is CUDA-graph-capturable (used by the DCP
+        a2a communication backend)."""
+        if self.disabled:
+            return
+        assert input_tensor.device == self.device, (
+            f"this nccl communicator is created to work on {self.device}, "
+            f"but the input tensor is on {input_tensor.device}"
+        )
+        assert output_tensor.device == self.device, (
+            f"this nccl communicator is created to work on {self.device}, "
+            f"but the output tensor is on {output_tensor.device}"
+        )
+        stream = self._resolve_stream()
+        # Equal-split all-to-all: fail loudly instead of silently truncating the tail.
+        assert input_tensor.numel() == output_tensor.numel(), (
+            f"all_to_all_single: input numel ({input_tensor.numel()}) != output "
+            f"numel ({output_tensor.numel()})"
+        )
+        assert input_tensor.numel() % self.world_size == 0, (
+            f"all_to_all_single: input numel ({input_tensor.numel()}) not "
+            f"divisible by world_size ({self.world_size})"
+        )
+        chunk_size = input_tensor.numel() // self.world_size
+        dtype = ncclDataTypeEnum.from_torch(input_tensor.dtype)
+        self.nccl.ncclGroupStart()
+        for i in range(self.world_size):
+            send_buf = input_tensor.narrow(0, i * chunk_size, chunk_size)
+            self.nccl.ncclSend(
+                buffer_type(send_buf.data_ptr()),
+                chunk_size,
+                dtype,
+                i,
+                self.comm,
+                cudaStream_t(stream.cuda_stream),
+            )
+            recv_buf = output_tensor.narrow(0, i * chunk_size, chunk_size)
+            self.nccl.ncclRecv(
+                buffer_type(recv_buf.data_ptr()),
+                chunk_size,
+                dtype,
+                i,
+                self.comm,
+                cudaStream_t(stream.cuda_stream),
+            )
+        self.nccl.ncclGroupEnd()
+
     def broadcast(self, tensor: torch.Tensor, src: int):
         if self.disabled:
             return

python/sglang/srt/distributed/parallel_state.py

@@ -1072,7 +1072,12 @@ class GroupCoordinator:
         return True
 
     def _all_to_all_single(self, output: torch.Tensor, input: torch.Tensor) -> None:
-        torch.distributed.all_to_all_single(output, input, group=self.device_group)
+        # pynccl path keeps the a2a exchange CUDA-graph-capturable (DCP a2a backend).
+        pynccl_comm = self.pynccl_comm
+        if pynccl_comm is not None and not pynccl_comm.disabled:
+            pynccl_comm.all_to_all_single(output, input)
+        else:
+            torch.distributed.all_to_all_single(output, input, group=self.device_group)
 
     def all_to_all_single(self, output: torch.Tensor, input: torch.Tensor):
         if self.world_size == 1:

python/sglang/srt/layers/activation.py

@@ -57,12 +57,53 @@ _is_xpu = is_xpu()
 _use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip
 
 if _is_cuda:
+    from sgl_kernel import gelu_and_mul as _sgl_gelu_and_mul
+    from sgl_kernel import gelu_tanh_and_mul as _sgl_gelu_tanh_and_mul
+    from sgl_kernel import silu_and_mul as _sgl_silu_and_mul
+
     from sglang.kernels.ops.activation.activation import (
-        gelu_and_mul,
-        gelu_tanh_and_mul,
-        relu2,
-        silu_and_mul,
+        gelu_and_mul as _jit_gelu_and_mul,
     )
+    from sglang.kernels.ops.activation.activation import (
+        gelu_tanh_and_mul as _jit_gelu_tanh_and_mul,
+    )
+    from sglang.kernels.ops.activation.activation import (
+        relu2,
+    )
+    from sglang.kernels.ops.activation.activation import (
+        silu_and_mul as _jit_silu_and_mul,
+    )
+
+    # The jit act-and-mul kernel requires the per-rank hidden size to be a
+    # multiple of the vector width (kMaxVecBytes/dtype: 32B on SM100+, else 16B --
+    # RuntimeCheck "hidden size must be divisible by vector size" in
+    # kernels/jit/csrc/elementwise/activation.cuh). Route unaligned shapes to the
+    # sgl_kernel implementation (e.g. DeepSeek-V2-Lite dense 10944/8 = 1368 at
+    # tp8, 1368 % 16 != 0).
+    _jit_act_max_vec_bytes: Optional[int] = None
+
+    def _jit_act_supported(out: torch.Tensor) -> bool:
+        global _jit_act_max_vec_bytes
+        if _jit_act_max_vec_bytes is None:
+            major, _ = torch.cuda.get_device_capability()
+            _jit_act_max_vec_bytes = 32 if major >= 10 else 16
+        return out.shape[-1] % (_jit_act_max_vec_bytes // out.dtype.itemsize) == 0
+
+    def _act_and_mul(jit_fn, sgl_fn, input: torch.Tensor, out=None) -> torch.Tensor:
+        if out is None:
+            out = input.new_empty(*input.shape[:-1], input.shape[-1] // 2)
+        (jit_fn if _jit_act_supported(out) else sgl_fn)(input, out)
+        return out
+
+    def silu_and_mul(input: torch.Tensor, out=None) -> torch.Tensor:
+        return _act_and_mul(_jit_silu_and_mul, _sgl_silu_and_mul, input, out)
+
+    def gelu_and_mul(input: torch.Tensor, out=None) -> torch.Tensor:
+        return _act_and_mul(_jit_gelu_and_mul, _sgl_gelu_and_mul, input, out)
+
+    def gelu_tanh_and_mul(input: torch.Tensor, out=None) -> torch.Tensor:
+        return _act_and_mul(_jit_gelu_tanh_and_mul, _sgl_gelu_tanh_and_mul, input, out)
+
 elif _is_xpu:
     from sgl_kernel import gelu_and_mul, gelu_tanh_and_mul, silu_and_mul
 elif _is_hip:

python/sglang/srt/layers/dcp/__init__.py

@@ -40,9 +40,11 @@ from sglang.srt.layers.dcp.comm import (
     all_gather_q_for_mla_decode,
     cp_lse_ag_out_rs_mha,
     cp_lse_ag_out_rs_mla,
+    dcp_a2a_lse_reduce,
     dcp_enabled,
     get_attention_dcp_rank,
     get_attention_dcp_world_size,
+    init_fi_a2a_workspace,
 )
 from sglang.srt.layers.dcp.layout import (
     filter_dcp_local_kv_indices,
@@ -61,6 +63,8 @@ from sglang.srt.layers.dcp.metadata import DecodeContextParallelMetadata
 
 __all__ = [
     "DecodeContextParallelMetadata",
+    "dcp_a2a_lse_reduce",
+    "init_fi_a2a_workspace",
     "all_gather_kv_cache_for_dcp",
     "all_gather_kv_cache_for_mha_chunk_extend",
     "all_gather_kv_cache_for_mha_extend",

python/sglang/srt/layers/dcp/comm.py

@@ -25,7 +25,12 @@ from typing import Optional
 
 import torch
 
-from sglang.kernels.ops.attention.dcp_kernels import CPTritonContext, correct_attn_out
+from sglang.kernels.ops.attention.dcp_kernels import (
+    CPTritonContext,
+    _lse_pack_dim,
+    correct_attn_out,
+    dcp_lse_combine_triton,
+)
 from sglang.srt.distributed.device_communicators.pynccl_allocator import (
     use_symmetric_memory,
 )
@@ -141,7 +146,14 @@ def _all_gather_dcp_kv_cache(kv_a: torch.Tensor):
     gathered_kv_a = kv_a.new_empty(
         (kv_a.shape[0] * dcp_world_size, *kv_a.shape[1:]),
     )
-    parallel.dcp_group.all_gather_into_tensor(gathered_kv_a, kv_a)
+    # pynccl has no fp8 dtype; all-gather is a byte copy, so transport an fp8 KV
+    # cache as raw bytes via a uint8 view (works with --kv-cache-dtype fp8_*).
+    if kv_a.dtype in (torch.float8_e4m3fn, torch.float8_e5m2):
+        parallel.dcp_group.all_gather_into_tensor(
+            gathered_kv_a.view(torch.uint8), kv_a.contiguous().view(torch.uint8)
+        )
+    else:
+        parallel.dcp_group.all_gather_into_tensor(gathered_kv_a, kv_a)
     gathered_kv_a = (
         gathered_kv_a.reshape((dcp_world_size,) + kv_a.shape)
         .transpose(0, 1)
@@ -193,6 +205,12 @@ def all_gather_kv_cache_for_mha_extend(
         [kv_a.shape[-1], k_pe.shape[-1]], dim=-1
     )
     prefix_kv_a = prefix_kv_a.squeeze(1)
+    # torch.cat can't promote fp8 (gathered prefix) + bf16 (current extend), so
+    # align dtypes first (dequant the fp8 prefix; exact for the scale=1.0 default).
+    if prefix_kv_a.dtype != kv_a.dtype:
+        prefix_kv_a = prefix_kv_a.to(kv_a.dtype)
+    if prefix_k_pe.dtype != k_pe.dtype:
+        prefix_k_pe = prefix_k_pe.to(k_pe.dtype)
     # re-organize kv with query orders
     prefix_lens_cu = torch.zeros(
         len(seq_lens) + 1,
@@ -342,3 +360,229 @@ def all_gather_kv_cache_for_dcp(
     gatherd_kv_cache = torch.cat(kv_cache_tuple, dim=0)
 
     return gatherd_kv_cache
+
+
+# ---------------------------------------------------------------------------
+# A2A communication backend for DCP decode (alternative to AG+RS above): exchange
+# per-head partial outputs + LSEs across DCP ranks, then combine locally with the
+# Triton LSE kernel. fi_a2a delegates the exchange to FlashInfer MNNVL (#2951).
+# ---------------------------------------------------------------------------
+
+# Per-process singleton: MNNVL workspace + this rank's cp position. Populated
+# once, pre-CUDA-graph-capture, by init_fi_a2a_workspace().
+_FI_A2A_STATE: Optional[dict] = None
+
+
+def init_fi_a2a_workspace(cp_group: "GroupCoordinator") -> None:
+    # Call once per process BEFORE CUDA-graph capture: the FlashInfer init syncs
+    # the stream and barriers cross-rank, neither of which is capturable.
+    global _FI_A2A_STATE
+    if _FI_A2A_STATE is not None:
+        return
+    if cp_group.world_size == 1:
+        return
+
+    import torch.distributed as dist
+
+    try:
+        from flashinfer.comm.dcp_alltoall import (
+            decode_cp_a2a_allocate_mnnvl_workspace,
+            decode_cp_a2a_init_workspace,
+        )
+        from flashinfer.comm.mapping import Mapping
+        from flashinfer.comm.mnnvl import MnnvlConfig, is_mnnvl_fabric_supported
+    except ImportError as e:
+        raise ImportError(
+            "--dcp-comm-backend fi_a2a requires FlashInfer with the DCP "
+            "all-to-all kernel (flashinfer #2951); could not import "
+            "flashinfer.comm.dcp_alltoall."
+        ) from e
+
+    # Reuse the MoE adapter: its Split() returns a CommBackend (what FlashInfer's
+    # Mapping expects); the flashinfer_comm_fusion copy has drifted to return a
+    # raw ProcessGroup, so don't swap without re-checking the Split() contract.
+    from sglang.srt.layers.moe.token_dispatcher.flashinfer_utils import (
+        TorchDistributedCommBackend,
+    )
+
+    if not is_mnnvl_fabric_supported(torch.cuda.current_device()):
+        raise RuntimeError(
+            "--dcp-comm-backend fi_a2a requires MNNVL fabric memory (e.g. "
+            "GB200 NVL72); is_mnnvl_fabric_supported() returned False. Use "
+            "--dcp-comm-backend a2a or ag_rs on clusters without MNNVL."
+        )
+
+    cp_size = cp_group.world_size
+    cp_rank = cp_group.rank_in_group
+    mapping = Mapping(
+        world_size=cp_size,
+        rank=cp_rank,
+        gpus_per_node=torch.cuda.device_count(),
+        cp_size=cp_size,
+        tp_size=1,
+        pp_size=1,
+    )
+    workspace = decode_cp_a2a_allocate_mnnvl_workspace(
+        mapping,
+        mnnvl_config=MnnvlConfig(
+            comm_backend=TorchDistributedCommBackend(cp_group.device_group)
+        ),
+    )
+    decode_cp_a2a_init_workspace(workspace, cp_rank, cp_size)
+    # REQUIRED barrier before the first alltoall: every rank must finish init,
+    # else a rank writes a peer's FIFO before it is ready -> deadlock.
+    dist.barrier(group=cp_group.device_group)
+    _FI_A2A_STATE = {
+        "workspace": workspace,
+        "cp_rank": cp_rank,
+    }
+
+
+def dcp_a2a_lse_reduce(
+    cp_attn_out: torch.Tensor,
+    cp_attn_lse: torch.Tensor,
+    cp_group: "GroupCoordinator",
+    is_lse_base_on_e: bool = True,
+    cuda_graph_buffers: Optional[dict] = None,
+    comm_backend: str = "a2a",
+) -> torch.Tensor:
+    """A2A DCP reduce: all-to-all exchange of head partials, then local Triton
+    combine. Output + fp32 LSE are packed into ONE all_to_all (LSE reinterpreted
+    as output-dtype columns along D) -> 1 NCCL call/layer instead of 2.
+    is_lse_base_on_e: True=base-e (FlashAttention), False=base-2 (FlashInfer-MLA).
+    """
+    if cp_group.world_size == 1:
+        return cp_attn_out
+
+    if comm_backend == "fi_a2a":
+        return _dcp_fi_a2a_lse_reduce(
+            cp_attn_out, cp_attn_lse, cp_group, is_lse_base_on_e
+        )
+
+    N = cp_group.world_size
+    B, H, D = cp_attn_out.shape
+    assert H % N == 0, f"num_heads ({H}) must be divisible by dcp_size ({N})"
+    H_per_rank = H // N
+    out_dtype = cp_attn_out.dtype
+    lpd = _lse_pack_dim(out_dtype)  # 2 for bf16/fp16
+
+    # Reshape [B, H, D] -> [N, B, H/N, D] — split heads across ranks
+    reshaped_out = cp_attn_out.view(B, N, H_per_rank, D).permute(1, 0, 2, 3)
+    reshaped_lse = cp_attn_lse.view(B, N, H_per_rank).permute(1, 0, 2)
+
+    if cuda_graph_buffers is not None:
+        # CUDA graph path with pre-allocated fused buffers.
+        send_combined = cuda_graph_buffers["send_combined"]
+        recv_combined = cuda_graph_buffers["recv_combined"]
+        send_lse_stg = cuda_graph_buffers["send_lse"]
+        recv_lse_stg = cuda_graph_buffers["recv_lse"]
+
+        send_combined[:, :B, :, :D].copy_(reshaped_out)
+        send_lse_stg[:, :B, :].copy_(reshaped_lse)
+        send_combined[:, :, :, D:].copy_(
+            send_lse_stg.view(out_dtype).view(N, -1, H_per_rank, lpd)
+        )
+
+        cp_group.all_to_all_single(
+            recv_combined.reshape(-1).view(torch.uint8),
+            send_combined.reshape(-1).view(torch.uint8),
+        )
+        recv_output = recv_combined[:, :B, :, :D]
+        recv_lse_stg.view(out_dtype).view(N, -1, H_per_rank, lpd).copy_(
+            recv_combined[:, :, :, D:]
+        )
+        recv_lse = recv_lse_stg[:, :B, :]
+    else:
+        send_lse_contig = reshaped_lse.contiguous()  # [N, B, H_per_rank] fp32
+        send_combined = torch.empty(
+            N,
+            B,
+            H_per_rank,
+            D + lpd,
+            dtype=out_dtype,
+            device=cp_attn_out.device,
+        )
+        recv_combined = torch.empty_like(send_combined)
+
+        send_combined[:, :, :, :D].copy_(reshaped_out)
+        send_combined[:, :, :, D:].copy_(
+            send_lse_contig.view(out_dtype).view(N, B, H_per_rank, lpd)
+        )
+
+        # Transport as raw bytes (uint8): the output may be fp8 (fp8 KV cache),
+        # which pynccl's dtype enum can't send; byte a2a is exact for equal chunks.
+        cp_group.all_to_all_single(
+            recv_combined.reshape(-1).view(torch.uint8),
+            send_combined.reshape(-1).view(torch.uint8),
+        )
+
+        recv_output = recv_combined[:, :, :, :D]
+        recv_lse_stg = torch.empty(
+            N,
+            B,
+            H_per_rank,
+            dtype=torch.float32,
+            device=cp_attn_out.device,
+        )
+        recv_lse_stg.view(out_dtype).view(N, B, H_per_rank, lpd).copy_(
+            recv_combined[:, :, :, D:]
+        )
+        recv_lse = recv_lse_stg
+
+    combined, _ = dcp_lse_combine_triton(
+        recv_output, recv_lse, is_lse_base_on_e=is_lse_base_on_e
+    )
+    return combined
+
+
+def _dcp_fi_a2a_lse_reduce(
+    cp_attn_out: torch.Tensor,
+    cp_attn_lse: torch.Tensor,
+    cp_group: "GroupCoordinator",
+    is_lse_base_on_e: bool = True,
+) -> torch.Tensor:
+    """fi_a2a: delegate only the cross-rank exchange to FlashInfer's MNNVL kernel,
+    then reuse the local Triton LSE combine. FlashInfer takes output + LSE as
+    separate tensors: partial_o [B, H_per_rank, cp_size, D] (peer axis 2nd-to-last),
+    softmax_stats [B, H_per_rank, cp_size, 2] fp32 (S padded 1->2).
+    """
+    from flashinfer.comm.dcp_alltoall import decode_cp_a2a_alltoall
+
+    state = _FI_A2A_STATE
+    assert state is not None, (
+        "fi_a2a workspace not initialized — call init_fi_a2a_workspace(dcp_group) "
+        "at model-runner init (before CUDA graph capture)."
+    )
+
+    N = cp_group.world_size
+    B, H, D = cp_attn_out.shape
+    assert H % N == 0, f"num_heads ({H}) must be divisible by dcp_size ({N})"
+    H_per_rank = H // N
+
+    # FlashInfer sends partial_o[..., peer, :] to `peer`; head h -> peer h//H_per_rank,
+    # so the peer axis is the outer head split: [B,N,H_pr,D] -> [B,H_pr,N,D].
+    partial_o = cp_attn_out.view(B, N, H_per_rank, D).permute(0, 2, 1, 3).contiguous()
+    # softmax_stats: fp32 [B, H_per_rank, N, S=2] (FI requires S>=2 & even);
+    # carry the LSE in lane 0, lane 1 is ignored by the combine.
+    lse_view = cp_attn_lse.view(B, N, H_per_rank).permute(0, 2, 1)  # [B,H_pr,N]
+    softmax_stats = torch.zeros(
+        B, H_per_rank, N, 2, dtype=torch.float32, device=cp_attn_out.device
+    )
+    softmax_stats[..., 0] = lse_view
+
+    o_out, stats_out = decode_cp_a2a_alltoall(
+        partial_o,
+        softmax_stats,
+        state["workspace"],
+        state["cp_rank"],
+        N,
+    )
+
+    # o_out[b,hpr,src] = rank src's partial for local head hpr -> combine layout.
+    recv_output = o_out.permute(2, 0, 1, 3).contiguous()  # [N, B, H_per_rank, D]
+    recv_lse = stats_out[..., 0].permute(2, 0, 1).contiguous()  # [N, B, H_per_rank]
+
+    combined, _ = dcp_lse_combine_triton(
+        recv_output, recv_lse, is_lse_base_on_e=is_lse_base_on_e
+    )
+    return combined

python/sglang/srt/model_executor/model_runner.py

@@ -162,6 +162,7 @@ from sglang.srt.model_executor.runner import (
 from sglang.srt.platforms import current_platform
 from sglang.srt.runtime_context import (
     get_global_dwdp_manager,
+    get_parallel,
     get_server_args,
     set_global_dwdp_manager,
 )
@@ -856,6 +857,49 @@ class ModelRunner:
         self.prefill_attention_backend_str = backends.prefill_attention_backend_str
         self.decode_attention_backend_str = backends.decode_attention_backend_str
 
+        if self.server_args.dcp_size > 1 and self.server_args.dcp_replicate_q_proj:
+            self._prepare_replicated_q_proj()
+
+    def _prepare_replicated_q_proj(self) -> None:
+        # --dcp-replicate-q-proj: gather each rank's attn_tp head-shard of
+        # q_b_proj / w_kc into full-head buffers once here (pre-capture) so the
+        # MLA decode path can skip the per-layer Q all-gather. bf16/fp16 only.
+        from sglang.srt.layers.quantization.unquant import UnquantizedLinearMethod
+        from sglang.srt.models.deepseek_v2 import DeepseekV2AttentionMLA
+
+        dcp_group = get_parallel().dcp_group
+        if dcp_group.world_size <= 1:
+            return
+        n_prepared = 0
+        for m in self.model.modules():
+            if not isinstance(m, DeepseekV2AttentionMLA):
+                continue
+            if m.w_kc is None:
+                continue
+            qp = m.q_b_proj if m.has_q_b_proj else m.q_proj
+            # q-replicate only supports the unquantized bf16/fp16 absorb path;
+            # quantized q-proj (packed weights) and non-16-bit w_kc keep the
+            # per-layer Q all-gather.
+            if (
+                m.w_kc.dtype not in (torch.bfloat16, torch.float16)
+                or not isinstance(qp.quant_method, UnquantizedLinearMethod)
+                or qp.weight.dtype not in (torch.bfloat16, torch.float16)
+            ):
+                logger.warning(
+                    "dcp_replicate_q_proj: skipping quantized q-proj/w_kc "
+                    "(bf16/fp16 only); this layer keeps the Q all-gather."
+                )
+                continue
+            m.w_kc_qrep = dcp_group.all_gather(m.w_kc.contiguous(), dim=0)
+            m.q_b_proj_qrep_weight = dcp_group.all_gather(
+                qp.weight.data.contiguous(), dim=0
+            )
+            n_prepared += 1
+        logger.info(
+            "dcp_replicate_q_proj: prepared full-head Q weights for %d MLA layers",
+            n_prepared,
+        )
+
     def init_cuda_graphs(self, capture_decode_cuda_graph: bool = True):
         capture = capture_cuda_graphs(
             model_runner=self, capture_decode_cuda_graph=capture_decode_cuda_graph

python/sglang/srt/model_executor/runner/base_runner.py

@@ -218,6 +218,7 @@ class BaseRunner(ABC):
             return
 
         self._pre_initialize_flashinfer_allreduce_workspace()
+        self._pre_initialize_fi_a2a_workspace()
 
         if should_run_flashinfer_autotune(self.model_runner):
             buffers, batch_size = self._autotune_buffers()
@@ -256,6 +257,19 @@ class BaseRunner(ABC):
             dtype=mr.dtype,
         )
 
+    def _pre_initialize_fi_a2a_workspace(self):
+        """Allocate the FlashInfer MNNVL all-to-all workspace for the fi_a2a DCP
+        comm backend; must run before CG capture (it syncs the stream + barriers
+        cross-rank, uncapturable) and raises early on non-MNNVL platforms.
+        """
+        mr = self.model_runner
+        if mr.server_args.dcp_size <= 1 or mr.server_args.dcp_comm_backend != "fi_a2a":
+            return
+
+        from sglang.srt.layers.dcp import init_fi_a2a_workspace
+
+        init_fi_a2a_workspace(get_parallel().dcp_group)
+
     def _flashinfer_autotune(self, *, buffers, batch_size):
         """Run flashinfer autotune.
 

python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py

@@ -24,6 +24,7 @@ from sglang.srt.layers.dcp import (
     all_gather_kv_cache_for_mla_extend,
     all_gather_q_for_mla_decode,
     cp_lse_ag_out_rs_mla,
+    dcp_a2a_lse_reduce,
 )
 from sglang.srt.layers.quantization.fp8_utils import (
     materialize_bpreshuffle_fp8_scale_tuple,
@@ -249,6 +250,19 @@ class DeepseekMLAForwardMixin:
             self.q_lora_rank is not None
             and self._can_fuse_bmm_into_attention(forward_batch)
         )
+        # --dcp-replicate-q-proj: project full-head Q locally from pre-gathered
+        # weights and skip the per-layer Q all-gather (bf16 decode absorb only).
+        q_replicate_active = (
+            get_server_args().dcp_replicate_q_proj
+            and get_parallel().dcp_enabled
+            and forward_batch.forward_mode.is_decode()
+            and not self.use_deep_gemm_bmm
+            and self.w_kc_qrep is not None
+            and self.q_b_proj_qrep_weight is not None
+        )
+        if q_replicate_active:
+            # force standard absorb so the full-head w_kc bmm runs
+            fuse_bmm_attention = False
         q_lora = None
         topk_indices = None
         q_nope = None
@@ -350,6 +364,7 @@ class DeepseekMLAForwardMixin:
                 and get_is_capture_mode()
                 and forward_batch.forward_mode.is_decode_or_idle()
                 and q_lora is not None
+                and not q_replicate_active
             ):
                 current_stream = torch.cuda.current_stream()
                 self.alt_stream.wait_stream(current_stream)
@@ -373,7 +388,15 @@ class DeepseekMLAForwardMixin:
                 current_stream.wait_stream(self.alt_stream)
             else:
                 k_nope = k_nope.unsqueeze(1)
-                q = self.q_b_proj_forward(q)
+                if q_replicate_active:
+                    # full-head Q from the gathered weight (skips Q all-gather)
+                    q = torch.nn.functional.linear(q, self.q_b_proj_qrep_weight).view(
+                        -1,
+                        self.num_local_heads * get_parallel().attn_dcp_size,
+                        self.qk_head_dim,
+                    )
+                else:
+                    q = self.q_b_proj_forward(q)
 
                 # Hoist these above the DSA indexer split op so the indexer
                 # and the composite bmm+attention split op are adjacent in FX.
@@ -395,9 +418,18 @@ class DeepseekMLAForwardMixin:
                             self.layer_id, prev_topk_indices
                         )
         else:
-            q = self.q_proj(hidden_states)[0].view(
-                -1, self.num_local_heads, self.qk_head_dim
-            )
+            if q_replicate_active:
+                q = torch.nn.functional.linear(
+                    hidden_states, self.q_b_proj_qrep_weight
+                ).view(
+                    -1,
+                    self.num_local_heads * get_parallel().attn_dcp_size,
+                    self.qk_head_dim,
+                )
+            else:
+                q = self.q_proj(hidden_states)[0].view(
+                    -1, self.num_local_heads, self.qk_head_dim
+                )
             latent_cache = self.kv_a_proj_with_mqa(hidden_states)[0]
             k_nope = latent_cache[..., : self.kv_lora_rank]
             k_nope = self.kv_a_layernorm(k_nope).unsqueeze(1)
@@ -406,7 +438,14 @@ class DeepseekMLAForwardMixin:
             q_nope, q_pe, k_pe = self._split_q_nope_pe(q, latent_cache)
 
         _kvb_q = None
-        if fusion_plan is not None:
+        if q_replicate_active:
+            # full-head absorb with the pre-gathered w_kc (q_nope already full-head)
+            q_nope_out = (
+                torch.bmm(q_nope.transpose(0, 1), self.w_kc_qrep)
+                .transpose(0, 1)
+                .contiguous()
+            )
+        elif fusion_plan is not None:
             # The composite split op fills q_nope_out_buf and attention reads
             # this transposed alias directly.
             q_nope_out = fusion_plan.q_nope_out_view
@@ -556,7 +595,7 @@ class DeepseekMLAForwardMixin:
 
         # all_gather q_pe, q_nope_out,take tp8 as an example， q_pe [B, H, ROPE_DIM], q_nope_out [B, H, NOPE_DIM] gathered to [B, H * dcp_world_size, ROPE_DIM] [B, H * dcp_world_size, NOPE_DIM] for decode batch, and all gather k_pe, k_nope for extend batch.
         if get_parallel().dcp_enabled:
-            if forward_batch.forward_mode.is_decode():
+            if forward_batch.forward_mode.is_decode() and not q_replicate_active:
                 # if forward_batch.forward_mode is decode, gather q
                 q_nope_out, q_pe = all_gather_q_for_mla_decode(
                     q_nope_out=q_nope_out,
@@ -792,10 +831,22 @@ class DeepseekMLAForwardMixin:
                 self.num_local_heads * get_parallel().attn_dcp_size,
                 self.kv_lora_rank,
             )
-            attn_output = cp_lse_ag_out_rs_mla(
-                attn_output, lse, get_parallel().dcp_group
-            )
-            attn_output = attn_output.transpose(0, 1)
+            dcp_comm_backend = get_server_args().dcp_comm_backend
+            if dcp_comm_backend in ("a2a", "fi_a2a"):
+                # A2A exchange of head partials + LSE, then local Triton combine.
+                # MLA decode LSE is base-2 (FlashInfer-MLA/FlashMLA) -> base_on_e=False.
+                attn_output = dcp_a2a_lse_reduce(
+                    attn_output.contiguous(),
+                    lse.contiguous(),
+                    get_parallel().dcp_group,
+                    is_lse_base_on_e=False,
+                    comm_backend=dcp_comm_backend,
+                )
+            else:
+                attn_output = cp_lse_ag_out_rs_mla(
+                    attn_output, lse, get_parallel().dcp_group
+                )
+                attn_output = attn_output.transpose(0, 1)
         attn_output = attn_output.view(-1, self.num_local_heads, self.kv_lora_rank)
 
         _kvb_v = None

python/sglang/srt/models/deepseek_v2.py

@@ -1761,6 +1761,11 @@ class DeepseekV2AttentionMLA(
         self.w_vc = None
         self.w_scale = 1.0
 
+        # Full-head Q/absorb weights for --dcp-replicate-q-proj, gathered once
+        # pre-CUDA-graph-capture by the model runner; None unless replicate is on.
+        self.w_kc_qrep = None
+        self.q_b_proj_qrep_weight = None
+
         self.w_scale_k = None
         self.w_scale_v = None
         self.use_deep_gemm_bmm = False

python/sglang/srt/server_args.py

@@ -1038,6 +1038,29 @@ class ServerArgs:
         ),
         NS("parallel"),
     ] = 1
+    dcp_comm_backend: A[
+        str,
+        Arg(
+            help="Communication backend for the decode context-parallel (DCP) "
+            "attention reduction: 'ag_rs' (AllGather + ReduceScatter), 'a2a' "
+            "(fused NCCL All-to-All exchange of output+LSE + local Triton LSE "
+            "combine), or 'fi_a2a' (FlashInfer MNNVL All-to-All kernel; requires "
+            "SM90+ and MNNVL fabric memory, e.g. GB200 NVL72).",
+            choices=["ag_rs", "a2a", "fi_a2a"],
+        ),
+        NS("parallel"),
+    ] = "ag_rs"
+    dcp_replicate_q_proj: A[
+        bool,
+        Arg(
+            help="For MLA decode context parallelism with the a2a/fi_a2a "
+            "backend: replicate the Q projection so each DCP rank computes the "
+            "full-head query locally (redundant projection compute), eliminating "
+            "the per-layer head-dim all-gather of Q. Trades a small amount of "
+            "extra GEMM for one fewer collective per layer.",
+        ),
+        NS("parallel"),
+    ] = False
     enable_prefill_cp: A[
         bool,
         "Enable context parallelism for the prefill phase. Select the layout with --cp-strategy.",
@@ -3587,6 +3610,30 @@ class ServerArgs:
                 "--decode-context-parallel-size) must be >= 1, but got "
                 f"dcp_size={self.dcp_size}."
             )
+        if self.dcp_comm_backend in ("a2a", "fi_a2a") and self.dcp_size <= 1:
+            raise ValueError(
+                f"--dcp-comm-backend {self.dcp_comm_backend} only affects the "
+                "decode context-parallel attention reduction and therefore "
+                "requires --dcp-size / --decode-context-parallel-size > 1, but "
+                f"got dcp_size={self.dcp_size}."
+            )
+        if self.dcp_comm_backend == "fi_a2a" and not is_cuda():
+            raise ValueError(
+                "--dcp-comm-backend fi_a2a delegates the exchange to FlashInfer's "
+                "MNNVL All-to-All kernel, which requires an NVIDIA CUDA platform "
+                "with SM90+ and MNNVL fabric memory (e.g. GB200 NVL72). The "
+                "authoritative fabric probe runs at model-runner init; use 'a2a' "
+                "or 'ag_rs' on clusters without MNNVL."
+            )
+        if self.dcp_replicate_q_proj:
+            if self.dcp_size <= 1:
+                raise ValueError("--dcp-replicate-q-proj requires --dcp-size > 1.")
+            if self.dcp_comm_backend not in ("a2a", "fi_a2a"):
+                raise ValueError(
+                    "--dcp-replicate-q-proj only applies to the a2a/fi_a2a DCP "
+                    "communication backend (it removes the head-dim Q all-gather); "
+                    f"got --dcp-comm-backend={self.dcp_comm_backend}."
+                )
         if not self.dcp_size > 1:
             return
         if is_hip():