[feature] implement dcp for deepseek_v2 (#14194)

python/sglang/srt/distributed/parallel_state.py

@@ -54,6 +54,7 @@ from sglang.srt.utils import (
     get_current_device_stream_fast,
     get_int_env_var,
     is_cpu,
+    is_cuda,
     is_cuda_alike,
     is_hip,
     is_musa,
@@ -775,6 +776,43 @@ class GroupCoordinator:
         else:
             torch.distributed.all_reduce(input_, group=self.device_group)
 
+    def reduce_scatter_along_dim(
+        self, input_: torch.Tensor, dim: int = -1
+    ) -> torch.Tensor:
+        world_size = self.world_size
+        # Bypass the function if we are using only 1 GPU.
+        if world_size == 1:
+            return input_
+        assert (
+            -input_.dim() <= dim < input_.dim()
+        ), f"Invalid dim ({dim}) for input tensor with shape {input_.size()}"
+
+        if dim < 0:
+            # Convert negative dim to positive.
+            dim += input_.dim()
+
+        with self.use_symmetric_memory(self):
+            # TODO: make sure whether tensor layout affects nccl reduce_scatter
+            # Note: This will produce an incorrect answer if we don't make
+            # the input_tensor contiguous. Possible bug in reduce_scatter_tensor?
+            input_tensor = input_.movedim(dim, 0).contiguous()
+
+        assert input_tensor.shape[0] % world_size == 0
+        chunk_size = input_tensor.shape[0] // world_size
+        output_shape = (chunk_size,) + input_tensor.shape[1:]
+
+        with self.use_symmetric_memory(self):
+            output_tensor = torch.empty(
+                output_shape,
+                dtype=input_tensor.dtype,
+                device=input_tensor.device,
+            )
+
+        self.reduce_scatter_tensor(output_tensor, input_tensor)
+
+        # Reshape before returning
+        return output_tensor.movedim(0, dim)
+
     def _reduce_scatter_tensor(
         self,
         output: torch.Tensor,
@@ -1667,6 +1705,10 @@ def get_attn_cp_group() -> GroupCoordinator:
     return _ATTN_CP
 
 
+def get_dcp_group_no_assert() -> Optional[GroupCoordinator]:
+    return _DCP
+
+
 def get_dcp_group() -> GroupCoordinator:
     assert _DCP is not None, "decode context parallel group is not initialized"
     return _DCP
@@ -1999,12 +2041,12 @@ def initialize_model_parallel(
         raise RuntimeError(
             f"decode_context_parallel_size ({decode_context_parallel_size}) must be >= 1"
         )
-    if decode_context_parallel_size > 1 and not is_hip():
+    if decode_context_parallel_size > 1 and not (is_hip() or is_cuda()):
         raise RuntimeError(
             "Decode context parallel (decode_context_parallel_size > 1) is "
-            "currently only supported on the AMD HIP platform, but got "
+            "currently only supported on the AMD HIP platform or CUDA platform, but got "
             f"decode_context_parallel_size ({decode_context_parallel_size}) "
-            "on a non-HIP platform."
+            "on a non-HIP or non-CUDA platform."
         )
     if tensor_model_parallel_size % decode_context_parallel_size != 0:
         raise RuntimeError(
@@ -2071,6 +2113,15 @@ def initialize_model_parallel(
             group_name="dcp",
             recovered_rank=recovered_rank,
         )
+        if get_tensor_model_parallel_rank() == 0:
+            logger.info(
+                f"DCP enabled, dcp_size={decode_context_parallel_size}, tp_size={tensor_model_parallel_size}"
+            )
+    else:
+        if get_tensor_model_parallel_rank() == 0:
+            logger.info(
+                f"DCP disabled, dcp_size={decode_context_parallel_size}, tp_size={tensor_model_parallel_size}"
+            )
 
     attn_dp_size = attention_data_parallel_size
     attn_cp_size = attention_context_model_parallel_size
@@ -2333,6 +2384,11 @@ def ensure_model_parallel_initialized(
         f"{pp_world_size=} vs. "
         f"{pipeline_model_parallel_size=}"
     )
+    if decode_context_parallel_size > 1:
+        dcp_world_size = get_dcp_group().world_size
+        assert (
+            dcp_world_size == decode_context_parallel_size
+        ), f"decode context parallel group already initialized, but of unexpected size: {dcp_world_size=} {decode_context_parallel_size=}"
 
 
 def model_parallel_is_initialized():
@@ -2383,6 +2439,14 @@ def get_tensor_model_parallel_world_size():
     return get_tp_group().world_size
 
 
+def get_dcp_world_size():
+    return get_dcp_group().world_size
+
+
+def get_dcp_rank():
+    return get_dcp_group().rank_in_group
+
+
 def get_tensor_model_parallel_rank():
     """Return my rank for the tensor model parallel group."""
     return get_tp_group().rank_in_group

python/sglang/srt/entrypoints/engine.py

@@ -1259,6 +1259,11 @@ def _set_envs_and_config(server_args: ServerArgs):
         os.environ["NCCL_NVLS_ENABLE"] = str(
             int(server_args.enable_nccl_nvls or server_args.enable_symm_mem)
         )
+    if "NCCL_GRAPH_MIXING_SUPPORT" not in os.environ or server_args.enable_symm_mem:
+        # Note(wh): NCCL_GRAPH_MIXING_SUPPORT=0 can help improve performance for symmetric kernels.
+        # details in https://github.com/NVIDIA/nccl-tests/issues/333#issuecomment-3103636985
+        if server_args.dcp_size > 1:
+            os.environ["NCCL_GRAPH_MIXING_SUPPORT"] = "0"
     os.environ["CUDA_DEVICE_MAX_CONNECTIONS"] = "8"
     os.environ["CUDA_MODULE_LOADING"] = "AUTO"
 

python/sglang/srt/layers/attention/flashinfer_mla_backend.py

@@ -23,6 +23,13 @@ from sglang.srt.layers.attention.flashinfer_backend import (
     create_flashinfer_kv_indices_triton,
 )
 from sglang.srt.layers.attention.utils import assert_buffer_fits
+from sglang.srt.layers.utils.dcp_utils import (
+    DecodeContextParallelMetadata,
+    dcp_enabled,
+    get_attention_dcp_world_size,
+    plan_dcp_decode_metadata,
+    update_local_kv_lens_for_dcp,
+)
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
 from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import (
     is_in_tc_piecewise_cuda_graph,
@@ -410,6 +417,7 @@ class FlashInferMLAAttnBackend(AttentionBackend):
                 prefix_lens,
                 prefill_wrapper_paged=self.prefill_wrapper_paged,
                 use_ragged=use_ragged,
+                attn_dcp_metadata=forward_batch.attn_dcp_metadata,
             )
             self.forward_metadata = PrefillMetadata(
                 self.prefill_wrapper_paged, use_ragged
@@ -463,6 +471,7 @@ class FlashInferMLAAttnBackend(AttentionBackend):
         if forward_mode.is_decode_or_idle():
             assert seq_lens_cpu is not None
             kv_len_arr_cpu = seq_lens_cpu[:bs].to(torch.int32)
+            update_local_kv_lens_for_dcp(kv_len_arr_cpu)
             self.cuda_graph_kv_indptr_cpu[1 : bs + 1] = torch.cumsum(
                 kv_len_arr_cpu, dim=0
             )
@@ -559,7 +568,13 @@ class FlashInferMLAAttnBackend(AttentionBackend):
             )
         else:
             # mla paged prefill
-            k_buf = self.token_to_kv_pool.get_key_buffer(layer.layer_id).to(q.dtype)
+            if (
+                forward_batch.attn_dcp_metadata is not None
+                and forward_batch.attn_dcp_metadata.dcp_kv_buffer is not None
+            ):
+                k_buf = forward_batch.attn_dcp_metadata.dcp_kv_buffer.to(q.dtype)
+            else:
+                k_buf = self.token_to_kv_pool.get_key_buffer(layer.layer_id).to(q.dtype)
             if q_rope is None:
                 qall = q.view(-1, layer.tp_q_head_num, layer.head_dim)
                 q, q_rope = (
@@ -624,6 +639,7 @@ class FlashInferMLAAttnBackend(AttentionBackend):
         k_buffer = self.token_to_kv_pool.get_key_buffer(layer.layer_id).to(q.dtype)
 
         o = q_nope.new_empty(q_nope.shape)
+        # for decode and dcp_world_size > 1, lse should be returned to compute final attn_out
         # Direct call to run without the wrapper
         o = decode_wrapper.run(
             q_nope,
@@ -631,8 +647,13 @@ class FlashInferMLAAttnBackend(AttentionBackend):
             k_buffer[:, :, : layer.v_head_dim],
             k_buffer[:, :, layer.v_head_dim :],
             out=o,
+            # for decode forward_batch, each dcp rank computes total q and partial kv, thus, we need to return_lse for online softmax to get final attn_output
+            return_lse=forward_batch.forward_mode.is_decode() and dcp_enabled(),
         )
-
+        if isinstance(o, tuple):
+            out, lse = o
+            out = out.view(-1, layer.tp_q_head_num * layer.v_head_dim)
+            return (out, lse)
         return o.view(-1, layer.tp_q_head_num * layer.v_head_dim)
 
 
@@ -640,7 +661,9 @@ class FlashInferMLAIndicesUpdaterDecode:
     def __init__(self, model_runner: ModelRunner, attn_backend: AttentionBackend):
         # Parse Constants
         self.num_local_heads = (
-            model_runner.model_config.num_attention_heads // get_parallel().attn_tp_size
+            model_runner.model_config.num_attention_heads
+            // get_parallel().attn_tp_size
+            * get_attention_dcp_world_size()
         )
         self.kv_lora_rank = model_runner.model_config.kv_lora_rank
         self.qk_nope_head_dim = model_runner.model_config.qk_nope_head_dim
@@ -710,6 +733,16 @@ class FlashInferMLAIndicesUpdaterDecode:
                 kv_indices,
                 self.req_to_token.shape[1],
             )
+
+            if dcp_enabled():
+                plan_dcp_decode_metadata(
+                    kv_lens,
+                    kv_indptr,
+                    kv_indices,
+                    init_metadata_replay,
+                    fast_decode_kwargs,
+                    bs,
+                )
         else:
             kv_indptr, kv_indices = spec_info.kv_indptr, spec_info.kv_indices
 
@@ -775,6 +808,7 @@ class FlashInferMLAIndicesUpdaterPrefill:
         prefill_wrapper_paged: BatchMLAPagedAttentionWrapper,
         use_ragged: bool,
         spec_info: Optional[SpecInput] = None,
+        attn_dcp_metadata: Optional[DecodeContextParallelMetadata] = None,
     ):
         if use_ragged:
             paged_kernel_lens = prefix_lens
@@ -795,6 +829,7 @@ class FlashInferMLAIndicesUpdaterPrefill:
             self.qo_indptr,
             use_ragged,
             spec_info,
+            attn_dcp_metadata=attn_dcp_metadata,
         )
 
     def call_begin_forward(
@@ -810,6 +845,7 @@ class FlashInferMLAIndicesUpdaterPrefill:
         qo_indptr: torch.Tensor,
         use_ragged: bool,
         spec_info: Optional[SpecInput] = None,
+        attn_dcp_metadata: Optional[DecodeContextParallelMetadata] = None,
     ):
         bs = len(seq_lens)
         sm_scale = self.scaling
@@ -861,6 +897,11 @@ class FlashInferMLAIndicesUpdaterPrefill:
             )
         else:
             # mla paged prefill
+            if attn_dcp_metadata is not None:
+                if attn_dcp_metadata.dcp_kv_indptr is not None:
+                    kv_indptr = attn_dcp_metadata.dcp_kv_indptr
+                if attn_dcp_metadata.dcp_kv_indices is not None:
+                    kv_indices = attn_dcp_metadata.dcp_kv_indices
             kv_len_arr = kv_indptr[1:] - kv_indptr[:-1]
             wrapper_paged.plan(
                 qo_indptr,

python/sglang/srt/layers/attention/flashmla_backend.py

@@ -18,9 +18,16 @@ from sglang.srt.layers.attention.utils import (
     get_num_kv_index_blocks_flashmla,
 )
 from sglang.srt.layers.quantization.fp8_kernel import scaled_fp8_quant
+from sglang.srt.layers.utils.dcp_utils import (
+    dcp_enabled,
+    get_attention_dcp_rank,
+    get_attention_dcp_world_size,
+)
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
 from sglang.srt.runtime_context import get_parallel
 
+logger = logging.getLogger(__name__)
+
 if TYPE_CHECKING:
     from sglang.srt.layers.radix_attention import RadixAttention
     from sglang.srt.model_executor.model_runner import ModelRunner
@@ -88,6 +95,10 @@ class FlashMLABackend(FlashInferMLAAttnBackend):
         self.cuda_graph_mla_metadata_view = None
         self.cuda_graph_num_splits_view = None
 
+        # get dcp info
+        self.dcp_world_size = get_attention_dcp_world_size()
+        self.dcp_rank = get_attention_dcp_rank()
+
     def init_forward_metadata_out_graph(
         self,
         forward_batch: ForwardBatch,
@@ -327,6 +338,9 @@ class FlashMLABackend(FlashInferMLAAttnBackend):
 
         reshape_q = q.view(bs, -1, layer.tp_q_head_num, layer.head_dim)
         if self.is_fp8_kvcache:
+            assert (
+                self.dcp_world_size == 1
+            ), "FlashMLA does not support DCP for FP8 kv cache"
             if layer.k_scale is not None:
                 q_scale = layer.k_scale
                 descale_q = layer.k_scale.reshape(1)
@@ -360,7 +374,8 @@ class FlashMLABackend(FlashInferMLAAttnBackend):
 
             return o.view(-1, layer.tp_q_head_num * layer.v_head_dim)
         else:
-            o, _ = flash_mla_with_kvcache(
+            # todo: need check all causal True or False?
+            o, lse = flash_mla_with_kvcache(
                 q=reshape_q,
                 k_cache=k_cache.view(-1, PAGE_SIZE, 1, self.kv_cache_dim),
                 block_table=self.forward_metadata.block_kv_indices[:bs],
@@ -371,8 +386,13 @@ class FlashMLABackend(FlashInferMLAAttnBackend):
                 softmax_scale=layer.scaling,
                 causal=True,
             )
-
-            return o.view(-1, layer.tp_q_head_num * layer.v_head_dim)
+            o = o.view(-1, layer.tp_q_head_num * layer.v_head_dim)
+            # TODO uniform output for forward_decode and forward_extend to
+            # return tuple instead of single output
+            # decode context parallel needs lse to correct attn_output via online softmax
+            if dcp_enabled():
+                return o, lse
+            return o
 
     def forward_extend(
         self,

python/sglang/srt/layers/utils/dcp_utils.py

@@ -0,0 +1,724 @@
+from dataclasses import dataclass
+from typing import Optional
+
+import torch
+import triton
+import triton.language as tl
+
+from sglang.srt.distributed.device_communicators.pynccl_allocator import (
+    use_symmetric_memory,
+)
+from sglang.srt.distributed.parallel_state import (
+    GroupCoordinator,
+    get_dcp_group,
+    get_dcp_group_no_assert,
+    get_dcp_rank,
+    get_dcp_world_size,
+)
+from sglang.srt.server_args import get_global_server_args
+from sglang.srt.utils import is_cuda
+
+
+def dcp_enabled() -> bool:
+    """
+    only checks whether dcp enabled for cuda platform
+    """
+    if get_dcp_group_no_assert() is None:
+        return False
+    if not is_cuda():
+        return False
+    return get_dcp_world_size() > 1
+
+
+def get_attention_dcp_group() -> GroupCoordinator:
+    return get_dcp_group()
+
+
+def get_attention_dcp_world_size() -> int:
+    if not dcp_enabled():
+        return 1
+    return get_dcp_world_size()
+
+
+def get_attention_dcp_rank() -> int:
+    if not dcp_enabled():
+        return 0
+    return get_dcp_rank()
+
+
+@triton.jit
+def _correct_attn_cp_out_kernel(
+    outputs_ptr,
+    new_output_ptr,
+    lses_ptr,
+    vlse_ptr,
+    outputs_stride_B,
+    outputs_stride_H,
+    outputs_stride_D,
+    lses_stride_N,
+    lses_stride_B,
+    lses_stride_H,
+    new_outputs_stride_H,
+    new_outputs_stride_B,
+    new_outputs_stride_D,
+    lse_idx,
+    HEAD_DIM: tl.constexpr,
+    N_ROUNDED: tl.constexpr,
+):
+    """
+    Apply the all-gathered lses to correct each local rank's attention
+    output. we still need perform a cross-rank reduction to obtain the
+    final attention output.
+
+    Args:
+        outputs_ptr (triton.PointerType):
+            Pointer to input tensor of shape [ B, H, D ]
+        lses_ptr (triton.PointerType):
+            Pointer to input tensor of shape [ N, B, H ]
+        new_output_ptr (triton.PointerType):
+            Pointer to output tensor of shape [ H, B, D ]
+        vlse_ptr (triton.PointerType):
+            Pointer to output tensor of shape [ B, H ]
+    """
+    batch_idx = tl.program_id(axis=0).to(tl.int64)
+    head_idx = tl.program_id(axis=1).to(tl.int64)
+
+    # Use int32 for offsets where possible to reduce register pressure
+    b_i32 = batch_idx.to(tl.int32)
+    h_i32 = head_idx.to(tl.int32)
+
+    # Vectorized load of LSE values: shape = [N]
+    num_n_offsets = tl.arange(0, N_ROUNDED)
+    lse_offsets = (
+        num_n_offsets * lses_stride_N + b_i32 * lses_stride_B + h_i32 * lses_stride_H
+    )
+
+    # Compute final LSE using online softmax algorithm (more numerically stable)
+    lse = tl.load(lses_ptr + lse_offsets)
+
+    # Replace NaN and inf with -inf for numerical stability
+    neg_inf = float("-inf")
+    lse = tl.where((lse != lse) | (lse == float("inf")), neg_inf, lse)
+
+    # Online softmax: find max, subtract, exp, sum, log
+    lse_max = tl.max(lse, axis=0)
+    lse_max = tl.where(lse_max == neg_inf, 0.0, lse_max)
+    lse = lse - lse_max
+    lse_exp = tl.exp2(lse)
+    lse_acc = tl.sum(lse_exp, axis=0)
+    final_lse = tl.log2(lse_acc) + lse_max
+
+    # Compute correction factor
+    lse_offset = lse_idx * lses_stride_N + b_i32 * lses_stride_B + h_i32 * lses_stride_H
+    local_lse = tl.load(lses_ptr + lse_offset)
+    lse_diff = local_lse - final_lse
+    lse_diff = tl.where(
+        (lse_diff != lse_diff) | (lse_diff == float("inf")),
+        neg_inf,
+        lse_diff,
+    )
+    factor = tl.exp2(lse_diff)
+
+    # Store final LSE
+    tl.store(vlse_ptr + b_i32 * lses_stride_B + h_i32 * lses_stride_H, final_lse)
+
+    # Load output with vectorized access: shape = [D]
+    d_offsets = tl.arange(0, HEAD_DIM)
+    output_offsets = (
+        batch_idx * outputs_stride_B
+        + head_idx * outputs_stride_H
+        + d_offsets * outputs_stride_D
+    )
+
+    new_output_offsets = (
+        head_idx * new_outputs_stride_H
+        + batch_idx * new_outputs_stride_B
+        + d_offsets * new_outputs_stride_D
+    )
+    # Apply correction and store
+    output = tl.load(outputs_ptr + output_offsets)
+    output = output * factor
+    tl.store(new_output_ptr + new_output_offsets, output)
+
+
+class CPTritonContext:
+    """The CPTritonContext is used to avoid recompilation of the Triton JIT."""
+
+    def __init__(self):
+        self.inner_kernel = None
+
+    def call_kernel(self, kernel, grid, *regular_args, **const_args):
+        if self.inner_kernel is None:
+            self.inner_kernel = kernel[grid](*regular_args, **const_args)
+        else:
+            self.inner_kernel[grid](*regular_args)
+
+
+def correct_attn_out(
+    out: torch.Tensor,
+    lses: torch.Tensor,
+    cp_rank: int,
+    ctx: Optional[CPTritonContext],
+    new_output: torch.Tensor = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Correct the attention output using the all-gathered lses.
+
+    Args:
+        out: Tensor of shape [ B, H, D ]
+        lses: Tensor of shape [ N, B, H ]
+        cp_rank: Current rank in the context-parallel group
+        ctx: Triton context to avoid recompilation
+
+    Returns:
+        Tuple of (out, lse) with corrected attention and final log-sum-exp.
+    """
+    if ctx is None:
+        ctx = CPTritonContext()
+
+    # --- Normalize to 3D views ---
+    if out.ndim == 4 and out.shape[1] == 1:
+        out = out.squeeze(1)
+    assert out.ndim == 3, f"expected out [B,H,D] or [B,1,H,D], got {tuple(out.shape)}"
+
+    if lses.ndim == 4 and lses.shape[-1] == 1:
+        lses = lses.squeeze(-1)
+    if lses.ndim == 4 and lses.shape[1] == 1:
+        lses = lses.squeeze(1)
+    assert lses.ndim == 3, (
+        f"expected lses [N,B,H] (optionally with a 1-sized extra dim), "
+        f"got {tuple(lses.shape)}"
+    )
+
+    B, H, D = out.shape
+    N = lses.shape[0]
+
+    # Strides after we normalized shapes to 3-D views.  The kernel computes
+    # offsets for `vlse_ptr` using lses_stride_B/H, so the output buffer must
+    # have the same B/H stride layout as a slice of `lses`.
+    o_sB, o_sH, o_sD = out.stride()
+    l_sN, l_sB, l_sH = lses.stride()
+    no_sH, no_sB, no_sD = new_output.stride()
+    # Allocate LSE with the same B/H strides as `lses` so writes land correctly
+    # even when `lses` is a non-contiguous view (e.g., 4-D to 3-D squeeze).
+    lse = torch.empty_strided(
+        (B, H), (l_sB, l_sH), device=lses.device, dtype=lses.dtype
+    )
+
+    # Kernel launch config
+    grid = (B, H, 1)
+
+    regular_args = (
+        out,
+        new_output,
+        lses,
+        lse,
+        o_sB,
+        o_sH,
+        o_sD,
+        l_sN,
+        l_sB,
+        l_sH,
+        no_sH,
+        no_sB,
+        no_sD,
+        cp_rank,
+    )
+    const_args = {"HEAD_DIM": D, "N_ROUNDED": N}
+
+    ctx.call_kernel(_correct_attn_cp_out_kernel, grid, *regular_args, **const_args)
+    return new_output, lse
+
+
+def cp_lse_ag_out_rs(
+    cp_attn_out: torch.Tensor,
+    cp_attn_lse: torch.Tensor,
+    cp_group: GroupCoordinator,
+    ctx: Optional[CPTritonContext] = None,
+):
+    """
+    cp_attn_out: [ B, H, D ]
+    cp_attn_lse: [ B, H ]
+    """
+    if cp_group.world_size == 1:
+        return cp_attn_out
+
+    if ctx is None:
+        ctx = CPTritonContext()
+
+    with use_symmetric_memory(cp_group):
+        # cp_attn_out is [B,H,D], we want to transpose it to [H,B,D] for the kernel, and then transpose back after correction.
+        new_output = cp_attn_out.new_empty(
+            cp_attn_out.transpose(0, 1).shape, dtype=torch.float32
+        )
+        cp_attn_lse = cp_attn_lse.to(torch.float32)
+    lses = cp_group.all_gather(cp_attn_lse, dim=0).view(
+        (cp_group.world_size,) + cp_attn_lse.shape
+    )
+    out, _ = correct_attn_out(
+        cp_attn_out, lses, cp_group.rank_in_group, ctx, new_output
+    )
+    out = cp_group.reduce_scatter_along_dim(out, dim=0)
+    return out.to(cp_attn_out.dtype)
+
+
+@triton.jit
+def create_dcp_kv_indices(
+    kv_indptr,
+    extend_lens_ptr,
+    extend_cu_lens_ptr,
+    extend_prefix_lens_ptr,
+    extend_cu_prefix_lens_ptr,
+    kv_indices_ptr,
+    extend_prefix_lens_sum,
+    dcp_world_size: tl.constexpr,
+):
+    BLOCK_SIZE: tl.constexpr = 512
+    pid = tl.program_id(axis=0)
+    prefix_len = tl.load(extend_prefix_lens_ptr + pid)
+    prefix_start = tl.load(extend_cu_prefix_lens_ptr + pid)
+    kv_ind_start = tl.load(kv_indptr + pid)
+    num_loop = tl.cdiv(prefix_len, BLOCK_SIZE)
+    for i in range(num_loop):
+        offset = tl.arange(0, BLOCK_SIZE) + i * BLOCK_SIZE
+        mask = offset < prefix_len
+        data = prefix_start + offset
+        tl.store(kv_indices_ptr + kv_ind_start + offset, data, mask=mask)
+    extend_len = tl.load(extend_lens_ptr + pid)
+    extend_start = tl.load(extend_cu_lens_ptr + pid)
+    num_loop = tl.cdiv(extend_len, BLOCK_SIZE)
+    for i in range(num_loop):
+        offset = tl.arange(0, BLOCK_SIZE) + i * BLOCK_SIZE
+        mask = offset < extend_len
+        data = extend_prefix_lens_sum + extend_start + offset
+        tl.store(
+            kv_indices_ptr + kv_ind_start + prefix_len + offset,
+            data,
+            mask=mask,
+        )
+
+
+@triton.jit
+def update_kv_lens_and_indices(
+    kv_lens: torch.Tensor,
+    kv_lens_cumsum: torch.Tensor,
+    kv_indices: torch.Tensor,
+    local_kv_lens: torch.Tensor,
+    local_kv_lens_cumsum: torch.Tensor,
+    local_kv_indices: torch.Tensor,
+    dcp_rank: tl.constexpr,
+    dcp_world_size: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+):
+    bs_idx = tl.program_id(0)
+    block_idx = tl.program_id(1)
+
+    local_kv_len = tl.load(local_kv_lens + bs_idx)
+    local_kv_indices_start = tl.load(local_kv_lens_cumsum + bs_idx)
+    kv_indices_start = tl.load(kv_lens_cumsum + bs_idx)
+
+    block_start = block_idx * BLOCK_SIZE
+    offsets = block_start + tl.arange(0, BLOCK_SIZE)
+
+    mask = offsets < local_kv_len
+
+    kv_indice_offsets = offsets * dcp_world_size + dcp_rank + kv_indices_start
+    local_kv_indices_offsets = local_kv_indices_start + offsets
+
+    kv_values = tl.load(kv_indices + kv_indice_offsets, mask=mask)
+    tl.store(
+        local_kv_indices + local_kv_indices_offsets,
+        kv_values // dcp_world_size,
+        mask=mask,
+    )
+
+
+@dataclass
+class DecodeContextParallelMetadata:
+    # For decode context parallel
+    dcp_kv_indptr: Optional[torch.Tensor] = None
+    dcp_kv_buffer: Optional[torch.Tensor] = None
+    dcp_kv_indices: Optional[torch.Tensor] = None
+    dcp_local_prefix_kv_indices: Optional[torch.Tensor] = None
+    dcp_extend_prefix_lens_sum: Optional[int] = None
+
+
+def prepare_decode_context_parallel_metadata(
+    seq_lens: torch.Tensor,
+    extend_prefix_lens: torch.Tensor,
+    extend_prefix_lens_cpu: torch.Tensor,
+    extend_seq_lens: torch.Tensor,
+    req_pool_indices: torch.Tensor,
+    req_to_token: torch.Tensor,
+    seq_lens_sum: int,
+    kv_buffer_shape: torch.Size,
+    kv_cache_dtype,
+    kv_cache_device,
+    create_chunked_prefix_cache_kv_indices_fn,
+) -> Optional[DecodeContextParallelMetadata]:
+    if not dcp_enabled():
+        return None
+    # dcp_kv_buffer tokens' layout
+    # [ rank0_r1.prefix_tokens, rank1_r1.prefix_tokens, ..., rank7_r1.prefix_tokens,
+    #   ...,
+    #   rank0_rn.prefix_tokens, rank1_rn.prefix_tokens, ..., rank7_rn.prefix_tokens,
+    #   r1.extend_tokens, r2.extent_tokens, rn.extend_tokens ]
+    extend_prefix_starts = torch.zeros(
+        len(seq_lens),
+        dtype=torch.int32,
+        device=get_global_server_args().device,
+    )
+    extend_cu_prefix_lens = torch.zeros(
+        len(seq_lens) + 1,
+        dtype=torch.int32,
+        device=get_global_server_args().device,
+    )
+    extend_cu_prefix_lens[1:] = torch.cumsum(extend_prefix_lens, dim=0)
+    extend_cu_prefix_lens = extend_cu_prefix_lens[:-1]
+    extend_prefix_lens_sum = sum([i for i in extend_prefix_lens_cpu])
+
+    dcp_prefix_kv_indices = torch.empty(
+        sum(extend_prefix_lens_cpu),
+        dtype=torch.int32,
+        device=get_global_server_args().device,
+    )
+    create_chunked_prefix_cache_kv_indices_fn[(len(seq_lens),)](
+        req_to_token,
+        req_pool_indices,
+        extend_prefix_starts,
+        extend_prefix_lens,
+        extend_cu_prefix_lens,
+        dcp_prefix_kv_indices,
+        req_to_token.shape[1],
+    )
+    dcp_kv_indptr = torch.zeros(
+        len(seq_lens) + 1,
+        dtype=torch.int32,
+        device=get_global_server_args().device,
+    )
+    dcp_kv_indptr[1:] = seq_lens.cumsum(dim=0)
+    dcp_kv_indptr = dcp_kv_indptr[: (len(seq_lens) + 1)]
+    dcp_kv_indices = torch.zeros(
+        seq_lens_sum,
+        dtype=torch.int32,
+        device=get_global_server_args().device,
+    )
+
+    extend_cu_lens = torch.zeros(
+        len(seq_lens) + 1,
+        dtype=torch.int32,
+        device=get_global_server_args().device,
+    )
+    extend_cu_lens[1:] = torch.cumsum(extend_seq_lens, dim=0)
+    extend_cu_lens = extend_cu_lens[:-1]
+
+    create_dcp_kv_indices[(len(seq_lens),)](
+        dcp_kv_indptr,
+        extend_seq_lens,
+        extend_cu_lens,
+        extend_prefix_lens,
+        extend_cu_prefix_lens,
+        dcp_kv_indices,
+        extend_prefix_lens_sum,
+        get_dcp_world_size(),
+    )
+    dcp_local_prefix_kv_indices = (
+        dcp_prefix_kv_indices[
+            dcp_prefix_kv_indices % get_dcp_world_size() == get_dcp_rank()
+        ]
+        // get_dcp_world_size()
+    )
+    dcp_kv_buffer = torch.empty(
+        (
+            seq_lens_sum,
+            *kv_buffer_shape[1:],
+        ),
+        dtype=kv_cache_dtype,
+        device=kv_cache_device,
+    )
+    attn_dcp_metadata = DecodeContextParallelMetadata(
+        dcp_kv_indptr=dcp_kv_indptr,
+        dcp_kv_buffer=dcp_kv_buffer,
+        dcp_kv_indices=dcp_kv_indices,
+        dcp_local_prefix_kv_indices=dcp_local_prefix_kv_indices,
+        dcp_extend_prefix_lens_sum=extend_prefix_lens_sum,
+    )
+    return attn_dcp_metadata
+
+
+def _all_gather_dcp_kv_cache(kv_a: torch.Tensor):
+    dcp_world_size = get_dcp_world_size()
+    # not use symmetric_memory unless torch mem_pool updated, see https://github.com/pytorch/pytorch/issues/178138
+    gathered_kv_a = kv_a.new_empty(
+        (kv_a.shape[0] * dcp_world_size, *kv_a.shape[1:]),
+    )
+    get_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)
+        .reshape(-1, *kv_a.shape[1:])
+    )
+    return gathered_kv_a
+
+
+def all_gather_kv_cache_for_mha_chunk_extend(
+    kv_a: torch.Tensor,
+    k_pe: torch.Tensor,
+    prefix_kv_lens_cpu: torch.Tensor,
+    prefix_starts_cpu: torch.Tensor = None,
+):
+    if dcp_enabled():
+        kv_a = kv_a.unsqueeze(1)
+        gathered_kv = all_gather_kv_cache_for_dcp(
+            kv_a,
+            k_pe,
+            prefix_kv_lens_cpu,
+            prefix_starts_cpu,
+        )
+        kv_a, k_pe = gathered_kv.split([kv_a.shape[-1], k_pe.shape[-1]], dim=-1)
+        kv_a = kv_a.squeeze(1)
+    return kv_a.contiguous(), k_pe.contiguous()
+
+
+def all_gather_kv_cache_for_mha_extend(
+    token_to_kv_pool,
+    attn_mqa,
+    dcp_local_prefix_kv_indices,
+    seq_lens,
+    extend_prefix_lens,
+    extend_prefix_lens_cpu: list[int],
+    extend_seq_lens,
+    kv_a: torch.Tensor,
+    k_pe: torch.Tensor,
+):
+    prefix_kv_a, prefix_k_pe = token_to_kv_pool.get_mla_kv_buffer(
+        attn_mqa, dcp_local_prefix_kv_indices
+    )
+    extend_prefix_lens_cpu = torch.tensor(extend_prefix_lens_cpu)
+    gathered_kv_cache = all_gather_kv_cache_for_dcp(
+        prefix_kv_a,
+        prefix_k_pe,
+        extend_prefix_lens_cpu,
+    )
+    prefix_kv_a, prefix_k_pe = gathered_kv_cache.split(
+        [kv_a.shape[-1], k_pe.shape[-1]], dim=-1
+    )
+    prefix_kv_a = prefix_kv_a.squeeze(1)
+    # re-organize kv with query orders
+    prefix_lens_cu = torch.zeros(
+        len(seq_lens) + 1,
+        dtype=torch.int32,
+        device=kv_a.device,
+    )
+    extend_lens_cu = torch.zeros_like(prefix_lens_cu)
+    prefix_lens_cu[1:] = torch.cumsum(extend_prefix_lens, dim=0)
+    extend_lens_cu[1:] = torch.cumsum(extend_seq_lens, dim=0)
+    kv_a_tuple = ()
+    k_pe_tuple = ()
+    for i in range(len(seq_lens)):
+        kv_a_tuple += (
+            prefix_kv_a[prefix_lens_cu[i] : prefix_lens_cu[i + 1]],
+            kv_a[extend_lens_cu[i] : extend_lens_cu[i + 1]],
+        )
+        k_pe_tuple += (
+            prefix_k_pe[prefix_lens_cu[i] : prefix_lens_cu[i + 1]],
+            k_pe[extend_lens_cu[i] : extend_lens_cu[i + 1]],
+        )
+    kv_a = torch.cat(kv_a_tuple, dim=0)
+    k_pe = torch.cat(k_pe_tuple, dim=0)
+    return kv_a.contiguous(), k_pe.contiguous()
+
+
+def filter_dcp_local_kv_indices(kv_indices: torch.Tensor):
+    if dcp_enabled():
+        kv_indices = (
+            kv_indices[kv_indices % get_dcp_world_size() == get_dcp_rank()]
+            // get_dcp_world_size()
+        )
+    return kv_indices
+
+
+def all_gather_q_for_mla_decode(
+    q_nope_out: torch.Tensor,
+    q_pe: torch.Tensor,
+):
+    with use_symmetric_memory(get_dcp_group()):
+        # transpose q_pe and q_nope_out from [B, H, L] to [H, B, L]
+        combined = torch.cat([q_pe.transpose(0, 1), q_nope_out.transpose(0, 1)], dim=-1)
+    gathered = get_dcp_group().all_gather(combined, dim=0)
+    d_pe = q_pe.size(-1)
+    d_nope = q_nope_out.size(-1)
+    q_pe, q_nope_out = gathered.split([d_pe, d_nope], dim=-1)
+    q_pe = q_pe.transpose(0, 1)
+    q_nope_out = q_nope_out.transpose(0, 1)
+    return q_nope_out, q_pe
+
+
+def all_gather_kv_cache_for_mla_extend(
+    token_to_kv_pool,
+    attn_mqa,
+    extend_prefix_lens_cpu: list[int],
+    dcp_local_prefix_kv_indices,
+    dcp_extend_prefix_lens_sum,
+    dcp_kv_buffer,
+    kv_lora_rank,
+    k_nope,
+    k_pe,
+):
+    cache_k_nope, cache_k_rope = token_to_kv_pool.get_mla_kv_buffer(
+        attn_mqa,
+        dcp_local_prefix_kv_indices,
+    )
+    extend_prefix_lens_cpu = torch.tensor(extend_prefix_lens_cpu)
+    # all gather kv cache into forward_batch.attn_dcp_metadata.dcp_kv_buffer
+    gathered_kv = all_gather_kv_cache_for_dcp(
+        cache_k_nope,
+        cache_k_rope,
+        extend_prefix_lens_cpu,
+        prefix_starts_cpu=torch.zeros_like(extend_prefix_lens_cpu),
+    )
+    dcp_kv_buffer[:dcp_extend_prefix_lens_sum] = gathered_kv
+
+    # copy local kv cache into forward_batch.attn_dcp_metadata.dcp_kv_buffer
+    dcp_kv_buffer[
+        dcp_extend_prefix_lens_sum:,
+        ...,
+        :kv_lora_rank,
+    ] = k_nope
+    dcp_kv_buffer[
+        dcp_extend_prefix_lens_sum:,
+        ...,
+        kv_lora_rank:,
+    ] = k_pe
+
+
+def update_local_kv_lens_for_dcp(kv_len_arr):
+    if not dcp_enabled():
+        return
+    dcp_world_size = get_dcp_world_size()
+    dcp_rank = get_dcp_rank()
+    offset = dcp_rank + 1
+    kv_len_arr.sub_(offset).div_(dcp_world_size, rounding_mode="floor").add_(1)
+
+
+def plan_dcp_decode_metadata(
+    kv_lens: torch.Tensor,
+    kv_indptr: torch.Tensor,
+    kv_indices: torch.Tensor,
+    init_metadata_replay: bool,
+    fast_decode_kwargs: dict,
+    bs: int,
+):
+    local_kv_lens = kv_lens.clone()
+    update_local_kv_lens_for_dcp(local_kv_lens)
+    local_kv_lens.clamp_(min=0)
+
+    if not init_metadata_replay:
+        max_local_len = (
+            int(local_kv_lens.max().item()) if local_kv_lens.numel() > 0 else 0
+        )
+        total_local_len = (
+            int(local_kv_lens.sum().item()) if local_kv_lens.numel() > 0 else 0
+        )
+    else:
+        max_local_len = (
+            int(fast_decode_kwargs["kv_len_arr_cpu"].max().item())
+            if fast_decode_kwargs["kv_len_arr_cpu"].numel() > 0
+            else 0
+        )
+        total_local_len = (
+            int(fast_decode_kwargs["kv_len_arr_cpu"].sum().item())
+            if fast_decode_kwargs["kv_len_arr_cpu"].numel() > 0
+            else 0
+        )
+    local_kv_lens_cumsum = kv_indptr.new_zeros((bs + 1,))
+    local_kv_lens_cumsum[1 : bs + 1] = torch.cumsum(local_kv_lens, dim=0)
+    local_kv_indices = kv_indices.new_empty(total_local_len)
+    BLOCK_SIZE = 128
+    num_blocks = (
+        (max_local_len + BLOCK_SIZE - 1) // BLOCK_SIZE if max_local_len > 0 else 1
+    )
+    grid = (bs, num_blocks)
+    update_kv_lens_and_indices[grid](
+        kv_lens,
+        kv_indptr,
+        kv_indices,
+        local_kv_lens,
+        local_kv_lens_cumsum,
+        local_kv_indices,
+        dcp_rank=get_dcp_rank(),
+        dcp_world_size=get_dcp_world_size(),
+        BLOCK_SIZE=BLOCK_SIZE,
+    )
+    kv_indices[:total_local_len] = local_kv_indices[:total_local_len]
+    kv_lens.copy_(local_kv_lens)
+    kv_indptr[: bs + 1] = local_kv_lens_cumsum[: bs + 1]
+
+
+# all gather kv cache and re-org to query orders
+def all_gather_kv_cache_for_dcp(
+    prefix_kv_a: torch.Tensor,
+    prefix_k_pe: torch.Tensor,
+    prefix_kv_lens_cpu: torch.Tensor,
+    prefix_starts_cpu: torch.Tensor = None,
+):
+    """
+    prefix_kv_a and prefix_k_pe should have same shape, expect for last dim
+    """
+    if not dcp_enabled():
+        return torch.cat([prefix_kv_a, prefix_k_pe], dim=-1)
+    # 1. compute max kv_lens for each seq
+    dcp_world_size = get_dcp_world_size()
+    dcp_rank = get_dcp_rank()
+
+    if prefix_starts_cpu is None:
+        prefix_starts_cpu = torch.zeros_like(prefix_kv_lens_cpu)
+
+    left_pads = prefix_starts_cpu % dcp_world_size > dcp_rank
+    left_pads = left_pads.to(torch.int32)
+    right_pads = (
+        prefix_starts_cpu + prefix_kv_lens_cpu - 1
+    ) % dcp_world_size < dcp_rank
+    right_pads = right_pads.to(torch.int32)
+    padded_lens = (
+        prefix_kv_lens_cpu + (prefix_starts_cpu % dcp_world_size) + dcp_world_size - 1
+    ) // dcp_world_size
+
+    local_kv_lens = padded_lens - left_pads - right_pads
+    local_kv_lens_cu = torch.zeros(
+        len(prefix_kv_lens_cpu) + 1,
+        dtype=torch.int32,
+    )
+    local_kv_lens_cu[1:] = torch.cumsum(local_kv_lens, dim=0)
+
+    padded_kv_cache_arr = []
+    prefix_kv_cache = torch.cat([prefix_kv_a, prefix_k_pe], dim=-1)
+    for req_idx in range(len(prefix_kv_lens_cpu)):
+        padded_tensor = prefix_kv_cache.new_empty(
+            (padded_lens[req_idx].item(),) + prefix_kv_cache.size()[1:]
+        )
+        padded_tensor[
+            left_pads[req_idx] : left_pads[req_idx] + local_kv_lens[req_idx]
+        ] = prefix_kv_cache[local_kv_lens_cu[req_idx] : local_kv_lens_cu[req_idx + 1]]
+        padded_kv_cache_arr.append(padded_tensor)
+
+    padded_kv_cache = torch.cat(padded_kv_cache_arr, dim=0)
+
+    gatherd_kv_cache = _all_gather_dcp_kv_cache(padded_kv_cache)
+
+    # 2. re-org kv cache to query orders
+    padded_lens_cu = torch.zeros(
+        len(prefix_kv_lens_cpu) + 1,
+        dtype=torch.int32,
+    )
+    padded_lens_cu[1:] = torch.cumsum(padded_lens, dim=0)
+    kv_cache_tuple = ()
+    for req_idx in range(len(prefix_kv_lens_cpu)):
+        kv_cache_tuple += (
+            gatherd_kv_cache[
+                padded_lens_cu[req_idx] * dcp_world_size
+                + (prefix_starts_cpu[req_idx] % dcp_world_size) :
+            ][: prefix_kv_lens_cpu[req_idx]],
+        )
+    gatherd_kv_cache = torch.cat(kv_cache_tuple, dim=0)
+
+    return gatherd_kv_cache

python/sglang/srt/managers/scheduler.py

@@ -1765,7 +1765,7 @@ class Scheduler(
             enable_hisparse=self.enable_hisparse,
             full_tokens_per_layer=self.full_tokens_per_layer,
             swa_tokens_per_layer=self.swa_tokens_per_layer,
-            max_total_num_tokens=self.max_total_num_tokens,
+            max_total_num_tokens=self.max_total_num_tokens * self.server_args.dcp_size,
             get_last_batch=lambda: self.last_batch,
             get_running_batch=lambda: self.running_batch,
         )

python/sglang/srt/mem_cache/common.py

@@ -26,13 +26,15 @@ from sglang.srt.mem_cache.triton_ops.common import (
     write_req_to_token_pool_triton,
 )
 from sglang.srt.server_args import ServerArgs, get_global_server_args
-from sglang.srt.utils import is_hip, is_npu, support_triton
+from sglang.srt.utils import is_cuda, is_hip, is_npu, support_triton
 from sglang.srt.utils.common import ceil_align, is_pin_memory_available
 
 _is_npu = is_npu()
 
 _is_hip = is_hip()
 
+_is_cuda = is_cuda()
+
 if TYPE_CHECKING:
     from sglang.srt.managers.schedule_batch import Req, ScheduleBatch
     from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
@@ -440,13 +442,13 @@ def alloc_req_slots(
 
 
 def _alloc_page_size(batch: ScheduleBatch) -> int:
-    # DCP (HIP-only) swaps in a PagedTokenToKVPoolAllocator whose page_size is
-    # server_args.page_size * dcp_size, so it can be > 1 even when
+    # DCP (HIP & CUDA only) swaps in a PagedTokenToKVPoolAllocator whose
+    # page_size is server_args.page_size * dcp_size, so it can be > 1 even when
     # tree_cache.page_size (== server_args.page_size) is 1. Only on the HIP DCP
     # path do we branch on the real allocator's page_size so the paged path is
     # taken; everywhere else tree_cache.page_size is authoritative and the two
     # are equal (dcp_size == 1), so behavior is unchanged.
-    if _is_hip and get_global_server_args().dcp_size > 1:
+    if (_is_hip or _is_cuda) and get_global_server_args().dcp_size > 1:
         return batch.tree_cache.token_to_kv_pool_allocator.page_size
     return batch.tree_cache.page_size
 
@@ -484,6 +486,8 @@ def alloc_for_extend(
     if _alloc_page_size(batch) == 1:
         out_cache_loc = alloc_token_slots(batch.tree_cache, batch.extend_num_tokens)
     else:
+        # Since tree_cache.page_size is (page_size * dcp_world_size), for dcp
+        # on cuda platform, always use alloc_paged_token_slots_extend
         # Paged allocation - build last_loc
         last_loc = [
             (t[-1:] if len(t) > 0 else torch.tensor([-1], device=batch.device))

python/sglang/srt/mem_cache/kv_cache_builder.py

@@ -25,6 +25,9 @@ from typing import TYPE_CHECKING
 
 from sglang.srt.configs.model_config import ModelImpl
 from sglang.srt.environ import envs
+from sglang.srt.layers.utils.dcp_utils import (
+    dcp_enabled,
+)
 from sglang.srt.managers.mm_utils import init_mm_embedding_cache
 from sglang.srt.mem_cache.cache_init_params import CacheInitParams
 from sglang.srt.mem_cache.registry import TreeCacheBuildContext, create_tree_cache
@@ -200,7 +203,12 @@ def build_kv_cache(
         disable=disable_radix_cache,
         req_to_token_pool=req_to_token_pool,
         token_to_kv_pool_allocator=token_to_kv_pool_allocator,
-        page_size=page_size,
+        # When dcp enabled, kv_pool_allocator.page_size is page_size * dcp_size.
+        # TreeCache.page_size should keep the same as allocator.page_size to
+        # avoid kv page eviction conflicts.
+        page_size=(
+            page_size if not dcp_enabled() else token_to_kv_pool_allocator.page_size
+        ),
         is_eagle=spec_algorithm.is_eagle(),
         tp_cache_group=(
             attn_tp_cpu_group if server_args.enable_dp_attention else tp_cpu_group

python/sglang/srt/mem_cache/memory_pool.py

@@ -46,6 +46,11 @@ from sglang.srt.layers.attention.dsa.quant_k_cache import (
 from sglang.srt.layers.attention.dsa.utils import aiter_can_use_preshuffle_paged_mqa
 from sglang.srt.layers.quantization.fp8_kernel import fp8_dtype, is_fp8_fnuz
 from sglang.srt.layers.radix_attention import RadixAttention
+from sglang.srt.layers.utils.dcp_utils import (
+    dcp_enabled,
+    get_attention_dcp_rank,
+    get_attention_dcp_world_size,
+)
 from sglang.srt.mem_cache.allocator.mamba import MambaSlotAllocator
 from sglang.srt.mem_cache.triton_ops.cache_move import (
     copy_all_layer_kv_cache_tiled,
@@ -2246,6 +2251,14 @@ class MLATokenToKVPool(KVCache):
         maybe_detect_oob(loc, 0, self.size + self.page_size, "set_kv_buffer (MLA)")
         layer_id = layer.layer_id
         assert not self.dsa_kv_cache_store_fp8
+        if dcp_enabled():
+            valid_mask = (
+                loc % get_attention_dcp_world_size() == get_attention_dcp_rank()
+            )
+            if not valid_mask.all():
+                loc = loc[valid_mask]
+                cache_k = cache_k[valid_mask]
+
         if cache_k.dtype != self.dtype:
             cache_k = cache_k.to(self.dtype)
 

python/sglang/srt/mem_cache/triton_ops/mla_buffer.py

@@ -5,6 +5,11 @@ import triton
 import triton.language as tl
 
 from sglang.jit_kernel.utils import is_arch_support_pdl
+from sglang.srt.layers.utils.dcp_utils import (
+    dcp_enabled,
+    get_attention_dcp_rank,
+    get_attention_dcp_world_size,
+)
 
 
 @triton.jit
@@ -19,6 +24,8 @@ def set_mla_kv_buffer_kernel(
     nope_dim: tl.constexpr,
     rope_dim: tl.constexpr,
     BLOCK: tl.constexpr,
+    DCP_RANK: tl.constexpr,
+    DCP_WORLD_SIZE: tl.constexpr,
     USE_GDC: tl.constexpr = False,
 ):
     pid_loc = tl.program_id(0)
@@ -33,7 +40,10 @@ def set_mla_kv_buffer_kernel(
         tl.extra.cuda.gdc_wait()
 
     loc = tl.load(loc_ptr + pid_loc).to(tl.int64)
-    dst_ptr = kv_buffer_ptr + loc * buffer_stride + offs
+    is_valid = loc % DCP_WORLD_SIZE == DCP_RANK
+    safe_loc = tl.where(is_valid, loc, 0)
+    safe_loc = safe_loc // DCP_WORLD_SIZE
+    dst_ptr = kv_buffer_ptr + safe_loc * buffer_stride + offs
 
     # Three-way branch to handle boundary correctly while preserving fast path
     if base + BLOCK <= nope_dim:
@@ -68,7 +78,7 @@ def set_mla_kv_buffer_kernel(
 
         src = tl.where(is_nope, src_nope, src_rope)
 
-    tl.store(dst_ptr, src, mask=mask)
+    tl.store(dst_ptr, src, mask=mask & is_valid)
 
     if USE_GDC:
         tl.extra.cuda.gdc_launch_dependents()
@@ -124,6 +134,7 @@ def set_mla_kv_buffer_triton(
         n_loc >= _TMA_BULK_STORE_MIN_LOCS
         and is_arch_support_pdl()
         and can_use_set_mla_kv_buffer(nope_bytes, rope_bytes)
+        and not dcp_enabled()
     ):
         jit_set_mla_kv_buffer(kv_buffer, loc, cache_k_nope, cache_k_rope)
         return
@@ -150,6 +161,8 @@ def set_mla_kv_buffer_triton(
         nope_dim,
         rope_dim,
         BLOCK=BLOCK,
+        DCP_RANK=get_attention_dcp_rank(),
+        DCP_WORLD_SIZE=get_attention_dcp_world_size(),
         **pdl_kwargs,
     )
 

python/sglang/srt/model_executor/forward_batch_deepseek_mha_mixin.py

@@ -28,6 +28,10 @@ class ForwardBatchDeepSeekMHAMixin:
     prefix_chunk_len: Optional[int] = None
     # Start positions of prefix cache for each chunk, (num_prefix_chunks, batch_size)
     prefix_chunk_starts: Optional[torch.Tensor] = None
+    # Start positions of prefix cache for each chunk, (num_prefix_chunks, batch_size), need prefix_chunk_starts_cpu for dcp all gather kv cache
+    prefix_chunk_starts_cpu: Optional[torch.Tensor] = None
+    # length of prefix cache for each chunk, (num_prefix_chunks, batch_size)
+    prefix_chunk_seq_lens_cpu: Optional[torch.Tensor] = None
     # Lengths of prefix cache for each chunk, (num_prefix_chunks, batch_size)
     prefix_chunk_seq_lens: Optional[torch.Tensor] = None
     # Accumulated lengths of prefix cache for each chunk, (num_prefix_chunks, batch_size + 1)
@@ -151,14 +155,19 @@ class ForwardBatchDeepSeekMHAMixin:
                 self.prefix_chunk_len,
             )
         )
-        _, prefix_chunk_seq_lens_cpu = self.get_prefix_chunk_seq_lens(
-            torch.tensor(self.extend_prefix_lens_cpu),
-            self.num_prefix_chunks,
-            self.prefix_chunk_len,
+        prefix_chunk_starts_cpu, prefix_chunk_seq_lens_cpu = (
+            self.get_prefix_chunk_seq_lens(
+                torch.tensor(self.extend_prefix_lens_cpu),
+                self.num_prefix_chunks,
+                self.prefix_chunk_len,
+            )
         )
         self.prefix_chunk_starts = prefix_chunk_starts_cuda
         self.prefix_chunk_seq_lens = prefix_chunk_seq_lens_cuda
 
+        # set prefix_chunk_starts_cpu and prefix_chunk_seq_lens_cpu for dcp to gather chunk kv cache with arbitrary lens
+        self.prefix_chunk_starts_cpu = prefix_chunk_starts_cpu
+        self.prefix_chunk_seq_lens_cpu = prefix_chunk_seq_lens_cpu
         # Metadata for attention backend
         self.prefix_chunk_cu_seq_lens = torch.zeros(
             self.num_prefix_chunks,

python/sglang/srt/model_executor/forward_batch_info.py

@@ -45,6 +45,7 @@ from sglang.srt.layers.dp_attention import (
     set_dp_buffer_len,
     set_is_extend_in_batch,
 )
+from sglang.srt.layers.utils.dcp_utils import DecodeContextParallelMetadata
 from sglang.srt.model_executor.forward_batch_deepseek_mha_mixin import (
     ForwardBatchDeepSeekMHAMixin,
 )
@@ -503,6 +504,9 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
 
     attn_cp_metadata: Optional[ContextParallelMetadata] = None
 
+    # For decode context parallel
+    attn_dcp_metadata: Optional[DecodeContextParallelMetadata] = None
+
     # Decode context parallel KV write mask.
     dcp_kv_mask: Optional[torch.Tensor] = None
 
@@ -860,7 +864,11 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
 
             model_runner.lora_manager.prepare_lora_batch(ret)
 
-        if getattr(model_runner, "dcp_size", 1) > 1 and ret.out_cache_loc is not None:
+        if (
+            getattr(model_runner, "dcp_size", 1) > 1
+            and ret.out_cache_loc is not None
+            and is_hip()
+        ):
             ret.dcp_kv_mask = (
                 ret.positions % model_runner.dcp_size == model_runner.dcp_rank
             )

python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py

@@ -11,7 +11,9 @@ from sglang.srt.configs.model_config import (
     is_deepseek_dsa,
     is_deepseek_v4,
 )
-from sglang.srt.distributed.parallel_state import get_world_group
+from sglang.srt.distributed.parallel_state import (
+    get_world_group,
+)
 from sglang.srt.environ import envs
 from sglang.srt.layers.dp_attention import get_attention_tp_size
 from sglang.srt.mem_cache.allocator import (

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

@@ -34,8 +34,16 @@ from sglang.srt.layers.pooler import EmbeddingPoolerOutput
 from sglang.srt.model_executor.cuda_graph_buffer_registry import (
     build_eager_registry,
 )
+from sglang.srt.model_executor.forward_batch_deepseek_mha_mixin import (
+    create_chunked_prefix_cache_kv_indices,
+)
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors
-from sglang.srt.model_executor.forward_context import ForwardContext, forward_context
+from sglang.srt.model_executor.forward_context import (
+    ForwardContext,
+    forward_context,
+    get_req_to_token_pool,
+    get_token_to_kv_pool,
+)
 from sglang.srt.model_executor.runner.base_runner import BaseRunner
 from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import (
     enable_tc_piecewise_cuda_graph,
@@ -256,6 +264,23 @@ class EagerRunner(BaseRunner):
             forward_batch = self.load_batch(forward_batch, pp_proxy_tensors)
 
         if forward_batch.needs_forward_metadata_init():
+            if hasattr(model_runner.model, "prepare_context_parallel_metadata_for_dcp"):
+                # prepare kv cache buffer for dcp to gather kv cache
+                forward_batch.attn_dcp_metadata = (
+                    model_runner.model.prepare_context_parallel_metadata_for_dcp(
+                        forward_batch.seq_lens,
+                        forward_batch.extend_prefix_lens,
+                        forward_batch.extend_prefix_lens_cpu,
+                        forward_batch.extend_seq_lens,
+                        forward_batch.req_pool_indices,
+                        get_req_to_token_pool().req_to_token,
+                        forward_batch.seq_lens_sum,
+                        get_token_to_kv_pool().get_key_buffer(0).shape,
+                        model_runner.kv_cache_dtype,
+                        model_runner.device,
+                        create_chunked_prefix_cache_kv_indices,
+                    )
+                )
             if hasattr(model_runner.model, "prepare_forward_batch"):
                 # Prepare model-specific attention metadata before planning,
                 # e.g. Moss-VL's prefill cross-attention custom mask.

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

@@ -9,6 +9,12 @@ from sglang.srt.layers.attention.dsa.dequant_k_cache import dequantize_k_cache_p
 from sglang.srt.layers.attention.tbo_backend import TboAttnBackend
 from sglang.srt.layers.attention.utils import concat_and_cast_mha_k_triton
 from sglang.srt.layers.communicator import get_attn_tp_context
+from sglang.srt.layers.utils.dcp_utils import (
+    all_gather_kv_cache_for_mha_chunk_extend,
+    all_gather_kv_cache_for_mha_extend,
+    dcp_enabled,
+    filter_dcp_local_kv_indices,
+)
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 from sglang.srt.model_executor.forward_context import (
     get_attn_backend,
@@ -271,11 +277,24 @@ class DeepseekMHAForwardMixin:
                 kv_a, k_pe = self._get_mla_kv_buffer_from_fp8_for_dsa(forward_batch)
             else:
                 # BF16/FP16 path: directly fetch from cache
-                kv_a, k_pe = self._get_mla_kv_buffer(
-                    forward_batch.fetch_mha_one_shot_kv_indices(),
-                    q.dtype,
-                    forward_batch,
-                )
+                if dcp_enabled():
+                    kv_a, k_pe = all_gather_kv_cache_for_mha_extend(
+                        get_token_to_kv_pool(),
+                        self.attn_mha,
+                        forward_batch.attn_dcp_metadata.dcp_local_prefix_kv_indices,
+                        forward_batch.seq_lens,
+                        forward_batch.extend_prefix_lens,
+                        forward_batch.extend_prefix_lens_cpu,
+                        forward_batch.extend_seq_lens,
+                        kv_a,
+                        k_pe,
+                    )
+                else:
+                    kv_a, k_pe = self._get_mla_kv_buffer(
+                        forward_batch.fetch_mha_one_shot_kv_indices(),
+                        q.dtype,
+                        forward_batch,
+                    )
         if _use_fp8_prefill_attn and self.kv_b_proj.weight.dtype == torch.uint8:
             # MXFP4 weights + FP8 prefill: fuse GEMM, nope/v split, and k_pe cat
             # into a single kernel (fused_gemm_afp4wfp4_split_cat) that writes k and v
@@ -421,6 +440,12 @@ class DeepseekMHAForwardMixin:
             kv_a_normed, k_pe = self._get_mla_kv_buffer(
                 kv_indices, kv_a_dtype, forward_batch
             )
+            kv_a_normed, k_pe = all_gather_kv_cache_for_mha_chunk_extend(
+                kv_a_normed,
+                k_pe,
+                forward_batch.prefix_chunk_seq_lens_cpu[i],
+                forward_batch.prefix_chunk_starts_cpu[i],
+            )
             kv = self.kv_b_proj(kv_a_normed)[0]
             kv = kv.view(
                 -1, self.num_local_heads, self.qk_nope_head_dim + self.v_head_dim
@@ -485,6 +510,7 @@ class DeepseekMHAForwardMixin:
         forward_batch: ForwardBatch,
     ):
         if _is_cuda or _use_aiter_gfx95:
+            kv_indices = filter_dcp_local_kv_indices(kv_indices=kv_indices)
             kv_a, k_pe = get_token_to_kv_pool().get_mla_kv_buffer(
                 self.attn_mha, kv_indices, dst_dtype
             )

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

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import logging
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Optional
 
@@ -20,6 +21,14 @@ from sglang.srt.layers.quantization.fp8_kernel import (
 )
 from sglang.srt.layers.radix_attention import unified_attention_with_output
 from sglang.srt.layers.utils.cp_utils import mla_use_prefill_cp
+from sglang.srt.layers.utils.dcp_utils import (
+    all_gather_kv_cache_for_mla_extend,
+    all_gather_q_for_mla_decode,
+    cp_lse_ag_out_rs,
+    dcp_enabled,
+    get_attention_dcp_group,
+    get_attention_dcp_world_size,
+)
 from sglang.srt.lora.deepseek_mla_correction import (
     apply_q_correction as apply_kv_b_lora_q_correction,
 )
@@ -62,6 +71,7 @@ from sglang.srt.state_capturer.indexer_topk import (
 from sglang.srt.utils import BumpAllocator
 from sglang.srt.utils.custom_op import register_custom_op
 
+logger = logging.getLogger(__name__)
 _SGLANG_EXPERIMENTAL_LORA_OPTI = envs.SGLANG_EXPERIMENTAL_LORA_OPTI.get()
 
 if TYPE_CHECKING:
@@ -529,6 +539,32 @@ class DeepseekMLAForwardMixin:
                 latent_cache, forward_batch, k_nope, k_pe
             )
 
+        # 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 dcp_enabled():
+            if forward_batch.forward_mode.is_decode():
+                # 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,
+                    q_pe=q_pe,
+                )
+            elif forward_batch.forward_mode.is_extend():
+                # for extend, gather kv
+                all_gather_kv_cache_for_mla_extend(
+                    get_token_to_kv_pool(),
+                    self.attn_mqa,
+                    forward_batch.extend_prefix_lens_cpu,
+                    forward_batch.attn_dcp_metadata.dcp_local_prefix_kv_indices,
+                    forward_batch.attn_dcp_metadata.dcp_extend_prefix_lens_sum,
+                    forward_batch.attn_dcp_metadata.dcp_kv_buffer,
+                    self.kv_lora_rank,
+                    k_nope,
+                    k_pe,
+                )
+            else:
+                logger.warning(
+                    f"not supported forward_mode {forward_batch.forward_mode}"
+                )
+
         return (
             q_pe,
             k_pe,
@@ -658,6 +694,22 @@ class DeepseekMLAForwardMixin:
                         topk_indices=topk_indices,
                     )
                     attn_output = fusion_plan.attn_output_buf
+                elif forward_batch.forward_mode.is_decode() and dcp_enabled():
+                    # set return_lse=True to correct attn_output
+                    attn_output, lse = self.attn_mqa_for_dcp_decode(
+                        q_nope_out,
+                        k_nope,
+                        k_nope,
+                        forward_batch,
+                        q_rope=q_pe,
+                        k_rope=k_pe,
+                        **extra_args,
+                        **(
+                            dict(topk_indices=topk_indices)
+                            if topk_indices is not None
+                            else {}
+                        ),
+                    )
                 else:
                     attn_output = self.attn_mqa(
                         q_nope_out,
@@ -714,6 +766,16 @@ class DeepseekMLAForwardMixin:
                 save_kv_cache=save_kv_cache,
                 **(dict(topk_indices=topk_indices) if topk_indices is not None else {}),
             )
+
+        # correct attn_output with respect to lse from other ranks
+        if forward_batch.forward_mode.is_decode() and dcp_enabled():
+            attn_output = attn_output.view(
+                -1,
+                self.num_local_heads * get_attention_dcp_world_size(),
+                self.kv_lora_rank,
+            )
+            attn_output = cp_lse_ag_out_rs(attn_output, lse, get_attention_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

@@ -122,6 +122,11 @@ from sglang.srt.layers.utils.cp_utils import (
     mla_use_prefill_cp,
     prepare_context_parallel_metadata,
 )
+from sglang.srt.layers.utils.dcp_utils import (
+    dcp_enabled,
+    get_attention_dcp_world_size,
+    prepare_decode_context_parallel_metadata,
+)
 from sglang.srt.layers.vocab_parallel_embedding import (
     ParallelLMHead,
     VocabParallelEmbedding,
@@ -1723,6 +1728,18 @@ class DeepseekV2AttentionMLA(
             quant_config=quant_config,
             prefix=add_prefix("attn_mqa", prefix),
         )
+        # use num_local_heads * dcp_world_size because q_nope, q_rope is all gathered from dcp ranks
+        if dcp_enabled():
+            self.attn_mqa_for_dcp_decode = RadixAttention(
+                self.num_local_heads * get_attention_dcp_world_size(),
+                self.kv_lora_rank + self.qk_rope_head_dim,
+                self.scaling,
+                num_kv_heads=1,
+                layer_id=layer_id,
+                v_head_dim=self.kv_lora_rank,
+                quant_config=quant_config,
+                prefix=add_prefix("attn_mqa", prefix),
+            )
 
         self.attn_mha = RadixAttention(
             self.num_local_heads,
@@ -2890,6 +2907,34 @@ class DeepseekV2ForCausalLM(nn.Module, DeepseekV2WeightLoaderMixin):
         self.capture_aux_hidden_states = True
         self.model.layers_to_capture = [val + 1 for val in layer_ids]
 
+    def prepare_context_parallel_metadata_for_dcp(
+        self,
+        seq_lens: torch.Tensor,
+        extend_prefix_lens: torch.Tensor,
+        extend_prefix_lens_cpu: torch.Tensor,
+        extend_seq_lens: torch.Tensor,
+        req_pool_indices: torch.Tensor,
+        req_to_token: torch.Tensor,
+        seq_lens_sum: int,
+        kv_buffer_shape: torch.Size,
+        kv_cache_dtype,
+        kv_cache_device,
+        create_chunked_prefix_cache_kv_indices_fn,
+    ):
+        return prepare_decode_context_parallel_metadata(
+            seq_lens=seq_lens,
+            extend_prefix_lens=extend_prefix_lens,
+            extend_prefix_lens_cpu=extend_prefix_lens_cpu,
+            extend_seq_lens=extend_seq_lens,
+            req_pool_indices=req_pool_indices,
+            req_to_token=req_to_token,
+            seq_lens_sum=seq_lens_sum,
+            kv_buffer_shape=kv_buffer_shape,
+            kv_cache_dtype=kv_cache_dtype,
+            kv_cache_device=kv_cache_device,
+            create_chunked_prefix_cache_kv_indices_fn=create_chunked_prefix_cache_kv_indices_fn,
+        )
+
 
 class DeepseekV3ForCausalLM(DeepseekV2ForCausalLM):
     pass

python/sglang/srt/server_args.py

@@ -911,6 +911,13 @@ class ServerArgs:
             aliases=["--moe-data-parallel-size"],
         ),
     ] = 1
+    dcp_size: A[
+        int,
+        Arg(
+            help="The decode context parallelism size.",
+            aliases=["--decode-context-parallel-size"],
+        ),
+    ] = 1
     enable_prefill_cp: A[
         bool,
         "Enable context parallelism for the prefill phase. Select the layout with --cp-strategy.",
@@ -2703,7 +2710,20 @@ class ServerArgs:
                 "--decode-context-parallel-size) must be >= 1, but got "
                 f"dcp_size={self.dcp_size}."
             )
-        if self.dcp_size > 1 and not is_hip():
+        if not self.dcp_size > 1:
+            return
+        if is_hip():
+            return
+        elif is_cuda():
+            if self.speculative_algorithm is not None:
+                raise ValueError(
+                    "Decode context parallel (--dcp-size / "
+                    "--decode-context-parallel-size > 1) on CUDA platform "
+                    "does not support any speculative algorithm, but got "
+                    f"dcp_size={self.dcp_size} on a CUDA platform with "
+                    "speculative decoding enabled."
+                )
+        else:
             raise ValueError(
                 "Decode context parallel (--dcp-size / "
                 "--decode-context-parallel-size > 1) is currently only "