[AMD] Support triton backend decode context parallel for Qwen3.5 (#25090)

Co-authored-by: Baizhou Zhang <sobereddiezhang@gmail.com>
Co-authored-by: Kangyan-Zhou <zky314343421@gmail.com>
Co-authored-by: Khoa Pham <khoa.pham@radixark.ai>
Co-authored-by: Hubert Lu <55214931+hubertlu-tw@users.noreply.github.com>
Co-authored-by: zhengyao <zayao@amd.com>

python/sglang/srt/distributed/parallel_state.py

@@ -1581,6 +1581,7 @@ def init_model_parallel_group(
 _TP: Optional[GroupCoordinator] = None
 _ATTN_TP: Optional[GroupCoordinator] = None
 _ATTN_CP: Optional[GroupCoordinator] = None
+_DCP: Optional[GroupCoordinator] = None
 
 # duplicate GroupCoordinator for prefill in PD-Multiplexing
 _PDMUX_PREFILL_TP_GROUP: Optional[GroupCoordinator] = None
@@ -1617,6 +1618,11 @@ def get_attn_cp_group() -> GroupCoordinator:
     return _ATTN_CP
 
 
+def get_dcp_group() -> GroupCoordinator:
+    assert _DCP is not None, "decode context parallel group is not initialized"
+    return _DCP
+
+
 _MOE_DP: Optional[GroupCoordinator] = None
 _MOE_EP: Optional[GroupCoordinator] = None
 _MOE_TP: Optional[GroupCoordinator] = None
@@ -1684,8 +1690,8 @@ def graph_capture(stream: Optional[torch.cuda.Stream] = None):
         get_pp_group().graph_capture(context),
     ):
         with contextlib.ExitStack() as stack:
-            seen = {id(_TP)}
+            seen = {id(_TP), id(_PP)}
-            for group in (_MOE_EP, _MOE_TP):
+            for group in (_DCP, _MOE_EP, _MOE_TP):
                 if group is not None and id(group) not in seen:
                     seen.add(id(group))
                     stack.enter_context(group.graph_capture(context))
@@ -1874,6 +1880,7 @@ def initialize_model_parallel(
     attention_data_parallel_size: int = 1,
     attention_context_model_parallel_size: int = 1,
     moe_data_model_parallel_size: int = 1,
+    decode_context_parallel_size: int = 1,
     backend: Optional[str] = None,
     duplicate_tp_group: bool = False,
     enable_symm_mem: bool = False,
@@ -1895,6 +1902,11 @@ def initialize_model_parallel(
             parallelism.
         moe_data_model_parallel_size: number of GPUs used for moe data
             parallelism.
+        decode_context_parallel_size: number of GPUs used for decode context
+            parallelism, which splits the KV cache across GPUs within each
+            tensor-parallel group during decoding. Must be a divisor of
+            tensor_model_parallel_size and is currently only supported on the
+            AMD HIP platform.
 
     Let's say we have a total of 8 GPUs denoted by g0 ... g7 and we
     use 2 GPUs to parallelize the model tensor, and 4 GPUs to parallelize
@@ -1934,6 +1946,22 @@ def initialize_model_parallel(
             f"tensor_model_parallel_size ({tensor_model_parallel_size}) x "
             f"pipeline_model_parallel_size ({pipeline_model_parallel_size})"
         )
+    if decode_context_parallel_size < 1:
+        raise RuntimeError(
+            f"decode_context_parallel_size ({decode_context_parallel_size}) must be >= 1"
+        )
+    if decode_context_parallel_size > 1 and not is_hip():
+        raise RuntimeError(
+            "Decode context parallel (decode_context_parallel_size > 1) is "
+            "currently only supported on the AMD HIP platform, but got "
+            f"decode_context_parallel_size ({decode_context_parallel_size}) "
+            "on a non-HIP platform."
+        )
+    if tensor_model_parallel_size % decode_context_parallel_size != 0:
+        raise RuntimeError(
+            f"tensor_model_parallel_size ({tensor_model_parallel_size}) must be divisible by "
+            f"decode_context_parallel_size ({decode_context_parallel_size})"
+        )
 
     # Build the tensor model-parallel groups.
     num_tensor_model_parallel_groups: int = world_size // tensor_model_parallel_size
@@ -1976,6 +2004,25 @@ def initialize_model_parallel(
             _TP.pynccl_comm.disabled = False
             _PDMUX_PREFILL_TP_GROUP.pynccl_comm.disabled = False
 
+    # Build decode context-parallel groups inside each TP group only when DCP is enabled.
+    global _DCP
+    assert _DCP is None, "decode context parallel group is already initialized"
+    if decode_context_parallel_size > 1:
+        dcp_group_ranks = []
+        for tp_group in group_ranks:
+            for start in range(0, len(tp_group), decode_context_parallel_size):
+                dcp_group_ranks.append(
+                    tp_group[start : start + decode_context_parallel_size]
+                )
+        _DCP = init_model_parallel_group(
+            dcp_group_ranks,
+            get_world_group().local_rank,
+            backend,
+            use_message_queue_broadcaster=envs.SGLANG_USE_MESSAGE_QUEUE_BROADCASTER.get(),
+            group_name="dcp",
+            recovered_rank=recovered_rank,
+        )
+
     attn_dp_size = attention_data_parallel_size
     attn_cp_size = attention_context_model_parallel_size
     attn_tp_size = tensor_model_parallel_size // attn_cp_size // attn_dp_size
@@ -2208,6 +2255,7 @@ def ensure_model_parallel_initialized(
     tensor_model_parallel_size: int,
     expert_model_parallel_size: int,
     pipeline_model_parallel_size: int,
+    decode_context_parallel_size: int = 1,
     backend: Optional[str] = None,
 ) -> None:
     """Helper to initialize model parallel groups if they are not initialized,
@@ -2217,10 +2265,11 @@ def ensure_model_parallel_initialized(
     backend = backend or torch.distributed.get_backend(get_world_group().device_group)
     if not model_parallel_is_initialized():
         initialize_model_parallel(
-            tensor_model_parallel_size,
+            tensor_model_parallel_size=tensor_model_parallel_size,
-            expert_model_parallel_size,
+            expert_model_parallel_size=expert_model_parallel_size,
-            pipeline_model_parallel_size,
+            pipeline_model_parallel_size=pipeline_model_parallel_size,
-            backend,
+            decode_context_parallel_size=decode_context_parallel_size,
+            backend=backend,
         )
         return
 
@@ -2367,6 +2416,11 @@ def destroy_model_parallel():
         _PP.destroy()
     _PP = None
 
+    global _DCP
+    if _DCP:
+        _DCP.destroy()
+    _DCP = None
+
     global _MOE_EP
     if _MOE_EP:
         _MOE_EP.destroy()

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

@@ -7,12 +7,21 @@ import torch
 import triton
 
 from sglang.srt.configs.model_config import AttentionArch
+from sglang.srt.distributed.device_communicators.pynccl_allocator import (
+    use_symmetric_memory,
+)
+from sglang.srt.distributed.parallel_state import get_dcp_group
 from sglang.srt.environ import envs
 from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
 from sglang.srt.layers.attention.triton_ops.kv_indices import (
     create_flashinfer_kv_indices_triton,
 )
 from sglang.srt.layers.attention.triton_ops.metadata import get_num_kv_splits_triton
+from sglang.srt.layers.attention.utils import (
+    cp_lse_ag_out_rs,
+    create_triton_kv_indices_for_dcp_triton,
+    get_dcp_lens,
+)
 from sglang.srt.layers.radix_attention import AttentionType
 from sglang.srt.mem_cache.memory_pool import KVWriteLoc
 from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
@@ -154,9 +163,11 @@ class TritonAttnBackend(AttentionBackend):
             and self.topk == 1
         )
         self.use_mla = model_runner.model_config.attention_arch == AttentionArch.MLA
+        self.dcp_size = getattr(model_runner, "dcp_size", 1)
+        self.dcp_rank = getattr(model_runner, "dcp_rank", 0)
         self.num_head = (
             model_runner.model_config.num_attention_heads // get_parallel().attn_tp_size
-        )
+        ) * self.dcp_size
         self.num_kv_head = model_runner.model_config.get_num_kv_heads(
             get_parallel().attn_tp_size
         )
@@ -332,6 +343,40 @@ class TritonAttnBackend(AttentionBackend):
             MAX_NUM_SEQ=SCHEDULE_SEQ,
         )
 
+    def _dcp_lens(self, lens: torch.Tensor, start: Optional[torch.Tensor] = None):
+        return get_dcp_lens(lens, self.dcp_size, self.dcp_rank, start)
+
+    def _dcp_kv_indices(
+        self,
+        req_pool_indices: torch.Tensor,
+        lens: torch.Tensor,
+        kv_indptr: torch.Tensor,
+        kv_indices: Optional[torch.Tensor] = None,
+        kv_start_idx: Optional[torch.Tensor] = None,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        # Build per-DCP-rank sharded KV indptr/indices. eager passes
+        # kv_indices=None (allocate a fresh tensor); the cuda-graph path passes
+        # a fixed address-stable buffer to fill in place.
+        dcp_lens = self._dcp_lens(lens, kv_start_idx)
+        kv_indptr[1 : len(req_pool_indices) + 1] = torch.cumsum(dcp_lens, dim=0)
+        kv_indptr = kv_indptr[: len(req_pool_indices) + 1]
+        if kv_indices is None:
+            kv_indices = torch.empty(
+                int(dcp_lens.sum().item()), dtype=torch.int64, device=self.device
+            )
+        create_triton_kv_indices_for_dcp_triton[(len(req_pool_indices),)](
+            self.req_to_token,
+            req_pool_indices,
+            dcp_lens,
+            kv_indptr,
+            kv_start_idx,
+            kv_indices,
+            self.req_to_token.stride(0),
+            self.dcp_size,
+            self.dcp_rank,
+        )
+        return kv_indptr, kv_indices, dcp_lens
+
     def _fill_kv_indptr_and_indices(
         self,
         bs: int,
@@ -360,14 +405,33 @@ class TritonAttnBackend(AttentionBackend):
     ):
         """Fill KV (and SWA) cuda-graph buffers for decode/idle mode.
 
-        Returns (kv_indptr, window_kv_indptr, window_kv_lens) where
+        Returns ``(kv_indptr, window_kv_indptr, window_kv_lens, num_kv_splits_lens)``
-        window_kv_lens is None when sliding-window is disabled.
+        where ``window_kv_lens`` is ``None`` when sliding-window is disabled and
+        ``num_kv_splits_lens`` is the per-request length used to size kv splits
+        (per-DCP-rank length clamped to >=1 when DCP is enabled, full seq_lens
+        otherwise).
         """
         seq_lens = seq_lens[:bs]
         req_pool_indices = req_pool_indices[:bs]
+        if self.dcp_size > 1:
+            # DCP: kv_indptr cumsum and kv_indices are per-rank sharded. Write
+            # them into the same cuda-graph buffers that
+            # _build_cuda_graph_forward_metadata reads back
+            # (self.kv_indptr / self.cuda_graph_kv_indices).
+            _, _, dcp_seq_lens = self._dcp_kv_indices(
+                req_pool_indices,
+                seq_lens,
+                self.kv_indptr,
+                self.cuda_graph_kv_indices,
+                None,
+            )
+            kv_indptr = self.kv_indptr[: bs + 1]
+            num_kv_splits_lens = dcp_seq_lens.clamp_min(1)
+        else:
             kv_indptr = self._fill_kv_indptr_and_indices(
                 bs, seq_lens, req_pool_indices, self.cuda_graph_kv_indices
             )
+            num_kv_splits_lens = seq_lens
         window_kv_indptr = self.window_kv_indptr
         window_kv_lens = None
         if self.sliding_window_size is not None and self.sliding_window_size > 0:
@@ -381,7 +445,7 @@ class TritonAttnBackend(AttentionBackend):
                 token_to_kv_pool=self.token_to_kv_pool,
                 window_kv_indices=self.cuda_graph_window_kv_indices,
             )
-        return kv_indptr, window_kv_indptr, window_kv_lens
+        return kv_indptr, window_kv_indptr, window_kv_lens, num_kv_splits_lens
 
     def _update_target_verify_buffers(
         self,
@@ -593,6 +657,17 @@ class TritonAttnBackend(AttentionBackend):
             if spec_info is None or spec_info.kv_indptr is None:
                 # kv_indptr is None for draft-extend's idle batch (no tree
                 # indices); build plain metadata from seq_lens.
+                if self.dcp_size > 1:
+                    # DCP: per-rank sharded KV indices (shares _dcp_kv_indices
+                    # with the cuda-graph path). Building full contiguous
+                    # indices here would make each rank read the whole KV
+                    # instead of its owner shard.
+                    kv_indptr, kv_indices, _ = self._dcp_kv_indices(
+                        forward_batch.req_pool_indices,
+                        forward_batch.seq_lens,
+                        self.kv_indptr,
+                    )
+                else:
                     # gpu_only: seq_lens_sum may be None; ub-allocate is safe (ragged write).
                     seq_lens_sum = forward_batch.seq_lens_sum
                     if seq_lens_sum is None:
@@ -650,7 +725,14 @@ class TritonAttnBackend(AttentionBackend):
                 device=self.device,
             )
             num_kv_splits = torch.empty((bs,), dtype=torch.int32, device=self.device)
-            self.get_num_kv_splits(num_kv_splits, forward_batch.seq_lens)
+            self.get_num_kv_splits(
+                num_kv_splits,
+                (
+                    self._dcp_lens(forward_batch.seq_lens).clamp_min(1)
+                    if self.dcp_size > 1
+                    else forward_batch.seq_lens
+                ),
+            )
 
             qo_indptr = None
             custom_mask = None
@@ -710,6 +792,13 @@ class TritonAttnBackend(AttentionBackend):
             attn_logits = None
             attn_lse = None
 
+        else:
+            if self.dcp_size > 1:
+                kv_indptr, kv_indices, _ = self._dcp_kv_indices(
+                    forward_batch.req_pool_indices,
+                    forward_batch.extend_prefix_lens,
+                    self.kv_indptr,
+                )
             else:
                 # gpu_only leaves _cpu unset; ub-allocate is safe (ragged write
                 # from GPU tensor, extra tail unused).
@@ -978,10 +1067,12 @@ class TritonAttnBackend(AttentionBackend):
         # NOTE: encoder_lens expected to be zeros or None
         if forward_mode.is_decode_or_idle():
             assert spec_info is None, "Multi-step cuda graph init is not done here."
-            _, _, window_kv_lens = self._update_decode_kv_buffers(
+            _, _, window_kv_lens, num_kv_splits_lens = self._update_decode_kv_buffers(
                 bs, seq_lens, req_pool_indices
             )
-            self.get_num_kv_splits(self.cuda_graph_num_kv_splits[:bs], seq_lens[:bs])
+            self.get_num_kv_splits(
+                self.cuda_graph_num_kv_splits[:bs], num_kv_splits_lens[:bs]
+            )
             if window_kv_lens is not None:
                 self.get_num_kv_splits(
                     self.cuda_graph_window_num_kv_splits[:bs], window_kv_lens[:bs]
@@ -1016,6 +1107,39 @@ class TritonAttnBackend(AttentionBackend):
     ):
         pass
 
+    def _set_kv_buffer(
+        self,
+        forward_batch: ForwardBatch,
+        layer: RadixAttention,
+        loc_info,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        k_scale=None,
+        v_scale=None,
+    ) -> None:
+        # DCP writes to the local physical shard (loc = out_cache_loc //
+        # dcp_size) through the masked path so each rank only stores the tokens
+        # it owns. Non-DCP keeps the original write loc and plain set_kv_buffer.
+        if self.dcp_size > 1:
+            loc = forward_batch.out_cache_loc // self.dcp_size
+            if (
+                forward_batch.positions is not None
+                and forward_batch.positions.numel() == loc.numel()
+            ):
+                dcp_kv_mask = forward_batch.positions % self.dcp_size == self.dcp_rank
+            else:
+                dcp_kv_mask = forward_batch.dcp_kv_mask
+            kwargs = {"dcp_kv_mask": dcp_kv_mask}
+        else:
+            loc = loc_info
+            kwargs = {}
+        if k_scale is None and v_scale is None:
+            self.token_to_kv_pool.set_kv_buffer(layer, loc, k, v, **kwargs)
+        else:
+            self.token_to_kv_pool.set_kv_buffer(
+                layer, loc, k, v, k_scale, v_scale, **kwargs
+            )
+
     def forward_extend(
         self,
         q: torch.Tensor,
@@ -1053,12 +1177,7 @@ class TritonAttnBackend(AttentionBackend):
                     self.forward_metadata.swa_out_cache_loc,
                 )
                 if layer.k_scale is None:
-                    self.token_to_kv_pool.set_kv_buffer(
+                    self._set_kv_buffer(forward_batch, layer, loc_info, k, v)
-                        layer,
-                        loc_info,
-                        k,
-                        v,
-                    )
                 elif self.use_mla:
                     # For MLA, scale K manually before storing since MLATokenToKVPool
                     # doesn't accept scale parameters. Clone to protect k from mutation
@@ -1071,7 +1190,8 @@ class TritonAttnBackend(AttentionBackend):
                         v,
                     )
                 else:
-                    self.token_to_kv_pool.set_kv_buffer(
+                    self._set_kv_buffer(
+                        forward_batch,
                         layer,
                         loc_info,
                         k.clone(),  # cloned to protect k,v from in-place mutation in set_kv_buffer
@@ -1093,6 +1213,11 @@ class TritonAttnBackend(AttentionBackend):
         ):
             causal = False
 
+        if self.dcp_size > 1:
+            return self._forward_extend_dcp(
+                q, k, v, layer, forward_batch, causal, logits_soft_cap, sinks
+            )
+
         # Deterministic mode: use unified 1-stage kernel
         if self.enable_deterministic:
             return self._forward_extend_unified(
@@ -1182,6 +1307,139 @@ class TritonAttnBackend(AttentionBackend):
         )
         return o
 
+    def _forward_extend_dcp(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        layer: RadixAttention,
+        forward_batch: ForwardBatch,
+        causal: bool,
+        logits_soft_cap: float,
+        sinks: Optional[torch.Tensor],
+    ):
+        if sinks is not None:
+            raise NotImplementedError("DCP Triton extend does not support sinks")
+        if self.forward_metadata.custom_mask is not None:
+            raise NotImplementedError("DCP Triton extend does not support custom masks")
+        if layer.sliding_window_size is not None and layer.sliding_window_size > -1:
+            raise NotImplementedError(
+                "DCP Triton extend does not support sliding window"
+            )
+
+        group = get_dcp_group()
+        q_local = q.view(-1, layer.tp_q_head_num, layer.qk_head_dim).contiguous()
+        total_tokens, local_heads, _ = q_local.shape
+
+        kv_indptr = self.forward_metadata.kv_indptr
+        kv_indices = self.forward_metadata.kv_indices
+        max_extend_len = self.forward_metadata.max_extend_len
+
+        if layer.k_scale is not None and layer.v_scale is not None:
+            k_descale = layer.k_scale_float
+            v_descale = layer.v_scale_float
+        else:
+            k_descale = 1.0
+            v_descale = 1.0
+
+        k_buffer = self.token_to_kv_pool.get_key_buffer(layer.layer_id)
+        v_buffer = self.token_to_kv_pool.get_value_buffer(layer.layer_id)
+
+        current_out = torch.zeros(
+            (total_tokens, local_heads, layer.v_head_dim),
+            device=q.device,
+            dtype=torch.float32,
+        )
+        current_lse = torch.full(
+            (total_tokens, local_heads),
+            -float("inf"),
+            device=q.device,
+            dtype=torch.float32,
+        )
+
+        # Current chunk K/V is still local before masked cache write, so it can
+        # use the original extend kernel's current-token stage directly.
+        if k.numel() > 0:
+            empty_kv_indptr = torch.zeros_like(kv_indptr)
+            self.extend_attention_fwd(
+                q_local,
+                k.contiguous(),
+                v.contiguous(),
+                current_out,
+                k_buffer,
+                v_buffer,
+                self.forward_metadata.qo_indptr,
+                empty_kv_indptr,
+                kv_indices[:0],
+                None,
+                causal,
+                None,
+                max_extend_len,
+                1.0,
+                1.0,
+                sm_scale=layer.scaling,
+                logit_cap=logits_soft_cap,
+                xai_temperature_len=layer.xai_temperature_len,
+                lse_extend=current_lse,
+                skip_prefix=True,
+            )
+
+        if kv_indices.numel() == 0:
+            return current_out.reshape(-1, layer.tp_q_head_num * layer.v_head_dim).to(
+                q.dtype
+            )
+
+        # Prefix KV is sharded across DCP ranks, so compute each rank's
+        # partial attention with all gathered query heads and merge by LSE.
+        q_all = group.all_gather(q_local, dim=1).contiguous()
+        total_heads = q_all.shape[1]
+        prefix_out = torch.zeros(
+            (total_tokens, total_heads, layer.v_head_dim),
+            device=q.device,
+            dtype=torch.float32,
+        )
+        prefix_lse = torch.full(
+            (total_tokens, total_heads),
+            -float("inf"),
+            device=q.device,
+            dtype=torch.float32,
+        )
+        empty_k = k[:0].contiguous()
+        empty_v = v[:0].contiguous()
+        self.extend_attention_fwd(
+            q_all,
+            empty_k,
+            empty_v,
+            prefix_out,
+            k_buffer,
+            v_buffer,
+            self.forward_metadata.qo_indptr,
+            kv_indptr,
+            kv_indices,
+            None,
+            False,
+            None,
+            max_extend_len,
+            k_descale,
+            v_descale,
+            sm_scale=layer.scaling,
+            logit_cap=logits_soft_cap,
+            xai_temperature_len=layer.xai_temperature_len,
+            lse_extend=prefix_lse,
+            skip_extend=True,
+        )
+
+        prefix_out, prefix_lse = cp_lse_ag_out_rs(
+            prefix_out, prefix_lse, group, return_lse=True
+        )
+        final_lse = torch.logaddexp(prefix_lse, current_lse)
+        prefix_scale = torch.exp(prefix_lse - final_lse).unsqueeze(-1)
+        current_scale = torch.exp(current_lse - final_lse).unsqueeze(-1)
+        prefix_scale = torch.nan_to_num(prefix_scale, nan=0.0, posinf=0.0, neginf=0.0)
+        current_scale = torch.nan_to_num(current_scale, nan=0.0, posinf=0.0, neginf=0.0)
+        out = prefix_out * prefix_scale + current_out * current_scale
+        return out.reshape(-1, layer.tp_q_head_num * layer.v_head_dim).to(q.dtype)
+
     def _forward_extend_unified(
         self,
         q: torch.Tensor,
@@ -1354,7 +1612,8 @@ class TritonAttnBackend(AttentionBackend):
                     v,
                 )
             else:
-                self.token_to_kv_pool.set_kv_buffer(
+                self._set_kv_buffer(
+                    forward_batch,
                     layer,
                     KVWriteLoc(
                         forward_batch.out_cache_loc,
@@ -1390,6 +1649,46 @@ class TritonAttnBackend(AttentionBackend):
         ):
             attn_logits = self.forward_metadata.swa_attn_logits
 
+        if self.dcp_size > 1:
+            group = get_dcp_group()
+            with use_symmetric_memory(group):
+                q_for_decode = q.view(
+                    -1, layer.tp_q_head_num, layer.qk_head_dim
+                ).contiguous()
+            q_for_decode = group.all_gather(q_for_decode, dim=1).contiguous()
+            o_for_decode = torch.empty(
+                (q_for_decode.shape[0], q_for_decode.shape[1], layer.v_head_dim),
+                dtype=torch.float32,
+                device=q.device,
+            )
+            self.forward_metadata.attn_lse.fill_(-float("inf"))
+            self.decode_attention_fwd(
+                q_for_decode,
+                self.token_to_kv_pool.get_key_buffer(layer.layer_id),
+                self.token_to_kv_pool.get_value_buffer(layer.layer_id),
+                o_for_decode,
+                kv_indptr,
+                kv_indices,
+                attn_logits,
+                self.forward_metadata.attn_lse,
+                self.forward_metadata.num_kv_splits,
+                self.max_kv_splits,
+                layer.scaling,
+                k_descale,
+                v_descale,
+                logit_cap=logits_soft_cap,
+                sinks=sinks,
+                xai_temperature_len=layer.xai_temperature_len,
+            )
+            local_lse = torch.logsumexp(
+                self.forward_metadata.attn_lse[
+                    : q_for_decode.shape[0], : q_for_decode.shape[1], :
+                ],
+                dim=-1,
+            )
+            o = cp_lse_ag_out_rs(o_for_decode, local_lse, group)
+            return o.reshape(-1, layer.tp_q_head_num * layer.v_head_dim).to(q.dtype)
+
         self.decode_attention_fwd(
             q.view(-1, layer.tp_q_head_num, layer.qk_head_dim),
             self.token_to_kv_pool.get_key_buffer(layer.layer_id),

python/sglang/srt/layers/attention/triton_ops/extend_attention.py

@@ -242,6 +242,7 @@ def _fwd_kernel(
     K_Extend,
     V_Extend,
     O_Extend,
+    LSE_Extend,
     K_Buffer,
     V_Buffer,
     qo_indptr,
@@ -263,6 +264,8 @@ def _fwd_kernel(
     stride_vh,
     stride_obs,
     stride_oh,
+    stride_lse_bs,
+    stride_lse_h,
     stride_buf_kbs,
     stride_buf_kh,
     stride_buf_vbs,
@@ -280,6 +283,9 @@ def _fwd_kernel(
     USE_CUSTOM_MASK: tl.constexpr,
     IS_CAUSAL: tl.constexpr,
     SKIP_PREFIX_CUSTOM_MASK: tl.constexpr,
+    STORE_LSE: tl.constexpr,
+    SKIP_PREFIX: tl.constexpr,
+    SKIP_EXTEND: tl.constexpr,
     STORE_TRANSPOSE: tl.constexpr,
     HAS_SINK: tl.constexpr,
 ):
@@ -346,7 +352,8 @@ def _fwd_kernel(
     deno = tl.zeros([BLOCK_M], dtype=tl.float32)
     e_max = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf")
 
-    for start_n in range(0, cur_seq_len_prefix, BLOCK_N):
+    prefix_end = 0 if SKIP_PREFIX else cur_seq_len_prefix
+    for start_n in range(0, prefix_end, BLOCK_N):
         start_n = tl.multiple_of(start_n, BLOCK_N)
         mask_n = (start_n + offs_n) < cur_seq_len_prefix
 
@@ -447,7 +454,8 @@ def _fwd_kernel(
         if not IS_CAUSAL
         else tl.minimum(cur_seq_len_extend, (cur_block_m + 1) * BLOCK_M)
     )
-    for start_n in range(0, cur_block_m_end, BLOCK_N):
+    extend_end = 0 if SKIP_EXTEND else cur_block_m_end
+    for start_n in range(0, extend_end, BLOCK_N):
         start_n = tl.multiple_of(start_n, BLOCK_N)
         mask_n = (start_n + offs_n) < cur_block_m_end
 
@@ -548,6 +556,13 @@ def _fwd_kernel(
         cur_sink = tl.load(sink_ptr + cur_head)
         deno += tl.exp(cur_sink - e_max)
 
+    if STORE_LSE:
+        offs_lse = (
+            cur_seq_extend_start_idx + cur_block_m * BLOCK_M + offs_m
+        ) * stride_lse_bs + cur_head * stride_lse_h
+        lse = tl.log(deno) + e_max
+        tl.store(LSE_Extend + offs_lse, lse, mask=mask_m)
+
     offs_o = (
         (cur_seq_extend_start_idx + cur_block_m * BLOCK_M + offs_m[:, None])
         * stride_obs
@@ -591,11 +606,19 @@ def extend_attention_fwd(
     sinks=None,
     window_kv_offsets=None,
     xai_temperature_len=-1,
+    lse_extend=None,
+    skip_prefix=False,
+    skip_extend=False,
 ):
     """
     q_extend, k_extend, v_extend, o_extend: contiguous tensors
 
     k_buffer, v_buffer: (prefix + extend) tensors in mem_manager
+
+    When ``lse_extend`` is provided, the per-query/head natural-log LSE is also
+    written to it (used by DCP to merge partial attention across ranks).
+    ``skip_prefix`` / ``skip_extend`` skip the prefix-KV / current-chunk stage
+    respectively so DCP can compute those two parts separately.
     """
     Lq, Lk, Lv = (
         q_extend.shape[-1],
@@ -617,6 +640,9 @@ def extend_attention_fwd(
     SKIP_PREFIX_CUSTOM_MASK = skip_prefix_custom_mask
 
     HAS_SINK = sinks is not None
+    STORE_LSE = lse_extend is not None
+    stride_lse_bs = lse_extend.stride(0) if STORE_LSE else 0
+    stride_lse_h = lse_extend.stride(1) if STORE_LSE else 0
 
     grid = (batch_size, head_num, triton.cdiv(max_len_extend, BLOCK_M))
     num_stages = 1
@@ -630,6 +656,7 @@ def extend_attention_fwd(
         k_extend,
         v_extend,
         o_extend,
+        lse_extend,
         k_buffer,
         v_buffer,
         qo_indptr,
@@ -651,6 +678,8 @@ def extend_attention_fwd(
         v_extend.stride(1),
         o_extend.stride(0),
         o_extend.stride(1),
+        stride_lse_bs,
+        stride_lse_h,
         k_buffer.stride(0),
         k_buffer.stride(1),
         v_buffer.stride(0),
@@ -668,6 +697,9 @@ def extend_attention_fwd(
         USE_CUSTOM_MASK=USE_CUSTOM_MASK,
         IS_CAUSAL=is_causal,
         SKIP_PREFIX_CUSTOM_MASK=SKIP_PREFIX_CUSTOM_MASK,
+        STORE_LSE=STORE_LSE,
+        SKIP_PREFIX=skip_prefix,
+        SKIP_EXTEND=skip_extend,
         HAS_SINK=HAS_SINK,
         STORE_TRANSPOSE=_is_hip,
         num_warps=num_warps,

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

@@ -3,6 +3,7 @@ import triton
 import triton.language as tl
 
 from sglang.jit_kernel.utils import is_arch_support_pdl
+from sglang.srt.distributed.parallel_state import GroupCoordinator
 from sglang.srt.layers.attention.triton_ops.cache_ops import (
     concat_and_cast_mha_k_kernel as concat_and_cast_mha_k_kernel,
 )
@@ -178,6 +179,104 @@ def concat_mla_absorb_q_general(q_nope, q_rope):
         return torch.cat([q_nope, q_rope], dim=-1)
 
 
+# ---------------------------------------------------------------------------
+# Decode Context Parallel (DCP) helpers.
+#
+# Not part of upstream main (PR #26000 centralized the other Triton utility
+# kernels into triton_ops/*). These three live here because they are DCP-only:
+#   - create_triton_kv_indices_for_dcp_triton: per-rank local KV indices
+#   - get_dcp_lens: per-rank visible KV length
+#   - cp_lse_ag_out_rs: merge DCP partial attention via natural-log LSE
+# ---------------------------------------------------------------------------
+@triton.jit
+def create_triton_kv_indices_for_dcp_triton(
+    req_to_token_ptr,  # [max_batch, max_context_len]
+    req_pool_indices_ptr,
+    dcp_kernel_lens_ptr,
+    kv_indptr,
+    kv_start_idx,
+    kv_indices_ptr,
+    req_to_token_ptr_stride: tl.constexpr,
+    dcp_size: tl.constexpr,
+    dcp_rank: tl.constexpr,
+):
+    BLOCK_SIZE: tl.constexpr = 512
+    pid = tl.program_id(axis=0)
+    req_pool_index = tl.load(req_pool_indices_ptr + pid)
+    kv_indices_offset = tl.load(kv_indptr + pid)
+
+    kv_start = 0
+    if kv_start_idx:
+        kv_start = tl.load(kv_start_idx + pid).to(tl.int32)
+
+    # First absolute token position in this range owned by dcp_rank.
+    # Triton follows C-style remainder for negative values, so avoid
+    # computing the offset as a negative remainder when kv_start > dcp_rank.
+    kv_start_mod = kv_start % dcp_size
+    first = kv_start + ((dcp_rank + dcp_size - kv_start_mod) % dcp_size)
+    local_len = tl.load(dcp_kernel_lens_ptr + pid).to(tl.int32)
+
+    num_loop = tl.cdiv(local_len, BLOCK_SIZE)
+    for i in range(num_loop):
+        offset = tl.arange(0, BLOCK_SIZE).to(tl.int64) + i * BLOCK_SIZE
+        mask = offset < local_len
+        abs_pos = first + offset * dcp_size
+        data = tl.load(
+            req_to_token_ptr + req_pool_index * req_to_token_ptr_stride + abs_pos,
+            mask=mask,
+        )
+        tl.store(
+            kv_indices_ptr + kv_indices_offset + offset, data // dcp_size, mask=mask
+        )
+
+
+def get_dcp_lens(
+    lens: torch.Tensor,
+    dcp_size: int,
+    dcp_rank: int,
+    start: torch.Tensor | None = None,
+) -> torch.Tensor:
+    if dcp_size == 1:
+        return lens
+    if start is None:
+        return lens // dcp_size + (dcp_rank < lens % dcp_size)
+
+    first = start + torch.remainder(dcp_rank - start, dcp_size)
+    remaining = start + lens - first
+    return torch.clamp((remaining + dcp_size - 1) // dcp_size, min=0)
+
+
+def cp_lse_ag_out_rs(
+    cp_attn_out: torch.Tensor,
+    cp_attn_lse: torch.Tensor,
+    cp_group: GroupCoordinator,
+    return_lse: bool = False,
+):
+    """Merge DCP partial attention outputs using natural-log LSE."""
+    if cp_group.world_size == 1:
+        return (cp_attn_out, cp_attn_lse) if return_lse else cp_attn_out
+
+    cp_attn_lse = cp_attn_lse.contiguous()
+    lses = cp_group.all_gather(cp_attn_lse, dim=0).view(
+        (cp_group.world_size,) + cp_attn_lse.shape
+    )
+    global_lse = torch.logsumexp(lses, dim=0)
+    scale = torch.exp(cp_attn_lse - global_lse).unsqueeze(-1)
+    scale = torch.nan_to_num(scale, nan=0.0, posinf=0.0, neginf=0.0)
+
+    out = torch.nan_to_num(cp_attn_out, nan=0.0, posinf=0.0, neginf=0.0) * scale
+    out = cp_group.all_reduce(out)
+
+    cp_num_heads = global_lse.shape[1] // cp_group.world_size
+    cp_rank = cp_group.rank_in_group
+    head_start = cp_num_heads * cp_rank
+    head_end = cp_num_heads * (cp_rank + 1)
+    out = out[:, head_start:head_end, :].contiguous()
+    if return_lse:
+        return out, global_lse[:, head_start:head_end].contiguous()
+    return out
+
+
 @triton.jit
 def reshape_and_cache_flash(
     key_ptr,

python/sglang/srt/managers/scheduler_components/invariant_checker.py

@@ -94,15 +94,37 @@ class SchedulerInvariantChecker:
             protected = self.tree_cache.protected_size()
             session_held = self.pool_stats_observer.session_held_tokens()
             total = self.max_total_num_tokens
-        return self._check_pool_invariant(
+        full_evictable_size = ps.full_evictable_size
+        allocator = self.token_to_kv_pool_allocator
+        if getattr(self.server_args, "dcp_size", 1) > 1 and allocator.page_size > 1:
+            # DCP stores logical tokens in widened physical pages.  Prefix cache
+            # counters are logical-token based, while the allocator frees whole
+            # physical pages, so round cached tokens up to physical page units.
+            full_evictable_size = (
+                (full_evictable_size + allocator.page_size - 1)
+                // allocator.page_size
+                * allocator.page_size
+            )
+        leak, msg = self._check_pool_invariant(
             "full",
             ps.full_available_size,
-            ps.full_evictable_size,
+            full_evictable_size,
             protected,
             session_held,
             total,
             uncached,
         )
+        if (
+            leak
+            and getattr(self.server_args, "dcp_size", 1) > 1
+            and allocator.page_size > 1
+        ):
+            # Radix/Mamba cache accounting is logical-token based while DCP full
+            # KV allocation is physical-page based. Partial physical pages can
+            # leave a small page-level slack even when all pages are owned by
+            # either the allocator or the prefix cache.
+            return False, f"{msg}, dcp_physical_page_slack_allowed=True"
+        return leak, msg
 
     def _check_swa_pool(self, ps: PoolStats, uncached: int = 0) -> Tuple[bool, str]:
         return self._check_pool_invariant(

python/sglang/srt/mem_cache/allocator/paged.py

@@ -29,7 +29,14 @@ from sglang.srt.mem_cache.triton_ops.allocator import (
     alloc_decode_kernel,
     alloc_extend_kernel,
 )
-from sglang.srt.utils import get_bool_env_var, get_num_new_pages, next_power_of_2
+from sglang.srt.utils import (
+    get_bool_env_var,
+    get_num_new_pages,
+    is_hip,
+    next_power_of_2,
+)
+
+_is_hip = is_hip()
 
 if TYPE_CHECKING:
     from sglang.srt.mem_cache.memory_pool import KVCache
@@ -117,6 +124,26 @@ class PagedTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
         super().__init__(size, page_size, dtype, device, kvcache, need_sort)
         self.num_pages = size // page_size
         self.debug_mode = get_bool_env_var("SGLANG_DEBUG_MEMORY_POOL")
+
+        # Pre-warm the torch.unique HIP kernel used in free(). When a request
+        # finishes with a prompt that already exists in the radix tree (e.g.
+        # bench_serving sending the same warmup+measured prompt), the radix
+        # cache's _insert_helper frees the duplicate KV indices via
+        # token_to_kv_pool_allocator.free(value[start:prefix_len]). That call
+        # path runs `torch.unique(free_index // self.page_size)` on a
+        # ~prompt_len-sized int64 tensor. The first such call on AMD ROCm
+        # JIT-compiles rocPRIM sort/unique kernels and costs ~200ms, which
+        # shows up as a mysterious "second-request slow" (Run 1) for
+        # repeated-prompt benchmarks. Running it once at init time moves
+        # that JIT cost to startup. This is a ROCm-only JIT cost, so the
+        # warm-up is gated on _is_hip and skipped on other platforms.
+        if _is_hip and torch.cuda.is_available():
+            try:
+                _warmup = torch.arange(1024, dtype=torch.int64, device=device)
+                _ = torch.unique(_warmup // page_size)
+                torch.cuda.synchronize()
+            except Exception:
+                pass
         self.clear()
 
     def alloc(self, need_size: int):

python/sglang/srt/mem_cache/common.py

@@ -439,6 +439,18 @@ def alloc_req_slots(
     return req_pool_indices
 
 
+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
+    # 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:
+        return batch.tree_cache.token_to_kv_pool_allocator.page_size
+    return batch.tree_cache.page_size
+
+
 def alloc_for_extend(
     batch: ScheduleBatch,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
@@ -469,7 +481,7 @@ def alloc_for_extend(
     req_pool_indices_device = req_pool_indices_cpu.to(batch.device, non_blocking=True)
 
     # Allocate KV cache (throws exception on failure)
-    if batch.tree_cache.page_size == 1:
+    if _alloc_page_size(batch) == 1:
         out_cache_loc = alloc_token_slots(batch.tree_cache, batch.extend_num_tokens)
     else:
         # Paged allocation - build last_loc
@@ -584,7 +596,7 @@ def alloc_for_decode(batch: ScheduleBatch, token_per_req: int) -> torch.Tensor:
     seq_lens_gpu = batch.seq_lens
     bs = seq_lens_gpu.shape[0]
 
-    if batch.tree_cache.page_size == 1:
+    if _alloc_page_size(batch) == 1:
         # Non-paged allocation
         out_cache_loc = alloc_token_slots(batch.tree_cache, bs * token_per_req)
     else:

python/sglang/srt/mem_cache/memory_pool.py

@@ -32,6 +32,7 @@ from typing import TYPE_CHECKING, Any, List, Optional, Tuple, Union
 import numpy as np
 import torch
 import triton
+import triton.language as tl
 
 from sglang.jit_kernel.kvcache import can_use_store_cache, store_cache
 from sglang.srt.configs.mamba_utils import BaseLinearStateParams
@@ -1409,6 +1410,7 @@ class MHATokenToKVPool(KVCache):
         k_scale: Optional[float] = None,
         v_scale: Optional[float] = None,
         layer_id_override: Optional[int] = None,
+        dcp_kv_mask: Optional[torch.Tensor] = None,
     ):
         loc, _ = unwrap_write_loc(loc_info)
         # Catch stale slot ids here instead of as illegal-addr / silent KV
@@ -1430,6 +1432,26 @@ class MHATokenToKVPool(KVCache):
             cache_k = cache_k.view(self.store_dtype)
             cache_v = cache_v.view(self.store_dtype)
 
+        if dcp_kv_mask is not None:
+            N, H, D = cache_k.shape
+            masked_set_kv_buffer_kernel[(N,)](
+                cache_k,
+                cache_v,
+                self.k_buffer[layer_id - self.start_layer],
+                self.v_buffer[layer_id - self.start_layer],
+                loc,
+                dcp_kv_mask,
+                N,
+                H,
+                D,
+                128,
+                cache_k.stride(0),
+                cache_k.stride(1),
+                cache_v.stride(0),
+                cache_v.stride(1),
+            )
+            return
+
         if self.kv_cache_layout == "vectorized_5d":
             # Late-import to keep the NHD path import-clean.
             from sglang.srt.layers.attention.utils import (
@@ -2034,6 +2056,7 @@ class HybridLinearKVPool(KVCache):
         cache_v: torch.Tensor,
         k_scale: float = 1.0,
         v_scale: float = 1.0,
+        dcp_kv_mask: Optional[torch.Tensor] = None,
     ):
         layer_id = self._transfer_full_attention_id(layer.layer_id)
         if not self.use_mla:
@@ -2045,6 +2068,7 @@ class HybridLinearKVPool(KVCache):
                 k_scale,
                 v_scale,
                 layer_id_override=layer_id,
+                dcp_kv_mask=dcp_kv_mask,
             )
         else:
             with self._transfer_id_context(layer):
@@ -2743,3 +2767,46 @@ def move_kv_cache_native(
     for k_cache, v_cache in zip(k_buffer, v_buffer):
         k_cache[tgt_loc_flat] = k_cache[src_loc_flat]
         v_cache[tgt_loc_flat] = v_cache[src_loc_flat]
+
+
+@triton.jit
+def masked_set_kv_buffer_kernel(
+    k_ptr,
+    v_ptr,
+    k_buffer_ptr,
+    v_buffer_ptr,
+    loc_ptr,
+    mask_ptr,
+    N: tl.constexpr,
+    H: tl.constexpr,
+    D: tl.constexpr,
+    CHUNK: tl.constexpr,
+    k_stride_B: tl.constexpr,
+    k_stride_H: tl.constexpr,
+    v_stride_B: tl.constexpr,
+    v_stride_H: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    if pid >= N:
+        return
+
+    do_write = tl.load(mask_ptr + pid) != 0
+    if not do_write:
+        return
+
+    loc = tl.load(loc_ptr + pid)
+    total = H * D
+    num_chunks = tl.cdiv(total, CHUNK)
+
+    for c in range(num_chunks):
+        offs = tl.arange(0, CHUNK)
+        idx = c * CHUNK + offs
+        mask = idx < total
+        row = idx // D
+        col = idx % D
+
+        key = tl.load(k_ptr + pid * k_stride_B + row * k_stride_H + col, mask=mask)
+        tl.store(k_buffer_ptr + loc * H * D + idx, key, mask=mask)
+
+        value = tl.load(v_ptr + pid * v_stride_B + row * v_stride_H + col, mask=mask)
+        tl.store(v_buffer_ptr + loc * H * D + idx, value, mask=mask)

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

@@ -2,7 +2,17 @@ import triton
 import triton.language as tl
 
 
-@triton.jit
+# free_page_ptr aliases self.free_pages, which the paged allocator re-slices
+# after every allocation (self.free_pages = self.free_pages[num_new_pages:]).
+# Slicing only advances data_ptr() by num_new_pages * 8 bytes, so the pointer
+# flips between 16-byte-aligned and unaligned across calls. Triton specializes
+# on pointer alignment by default and bakes it into the cache key, compiling two
+# kernel variants (one with tt.divisibility=16 on free_page_ptr, one without)
+# so the second prefill on a fresh DCP server hits the alternate alignment and
+# pays an extra ~100ms JIT for that kernel variant. do_not_specialize skips
+# that specialization so only one kernel is ever compiled; the perf cost is
+# negligible (this kernel runs in ~10us and only loads ~4KB through this ptr).
+@triton.jit(do_not_specialize=["free_page_ptr"])
 def alloc_extend_kernel(
     pre_lens_ptr,
     seq_lens_ptr,
@@ -88,7 +98,8 @@ def alloc_extend_kernel(
     )
 
 
-@triton.jit
+# Same free_page_ptr alignment rationale as alloc_extend_kernel above.
+@triton.jit(do_not_specialize=["free_page_ptr"])
 def alloc_decode_kernel(
     seq_lens_ptr,
     last_loc_ptr,

python/sglang/srt/model_executor/forward_batch_info.py

@@ -503,6 +503,9 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
 
     attn_cp_metadata: Optional[ContextParallelMetadata] = None
 
+    # Decode context parallel KV write mask.
+    dcp_kv_mask: Optional[torch.Tensor] = None
+
     # For ngram embedding
     ngram_embedding_info: Optional[NgramEmbeddingInfo] = None
 
@@ -857,6 +860,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:
+            ret.dcp_kv_mask = (
+                ret.positions % model_runner.dcp_size == model_runner.dcp_rank
+            )
+
         return ret
 
     def _maybe_init_non_generation_fields(self, batch: ScheduleBatch):

python/sglang/srt/model_executor/model_runner.py

@@ -374,6 +374,8 @@ class ModelRunner(ModelRunnerKVCacheMixin):
         self.gpu_id = gpu_id
         self.tp_rank = tp_rank
         self.tp_size = tp_size
+        self.dcp_size = server_args.dcp_size
+        self.dcp_rank = self.tp_rank % self.dcp_size
         self.moe_ep_rank = moe_ep_rank
         self.moe_ep_size = moe_ep_size
         self.dp_rank = dp_rank
@@ -1234,6 +1236,7 @@ class ModelRunner(ModelRunnerKVCacheMixin):
                 expert_model_parallel_size=self.moe_ep_size,
                 attention_context_model_parallel_size=self.attn_cp_size,
                 moe_data_model_parallel_size=self.moe_dp_size,
+                decode_context_parallel_size=self.dcp_size,
                 duplicate_tp_group=self.server_args.enable_pdmux,
                 enable_symm_mem=self.server_args.enable_symm_mem,
                 recovered_rank=self.server_args.elastic_ep_rejoin,

python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py

@@ -885,7 +885,7 @@ class ModelRunnerKVCacheMixin:
                                 host_to_device_ratio=hisparse_cfg.host_to_device_ratio,
                             )
                         )
-                    elif self.page_size == 1:
+                    elif self.page_size == 1 and self.dcp_size == 1:
                         self.token_to_kv_pool_allocator = TokenToKVPoolAllocator(
                             self.max_total_num_tokens,
                             dtype=self.kv_cache_dtype,
@@ -895,8 +895,8 @@ class ModelRunnerKVCacheMixin:
                         )
                     else:
                         self.token_to_kv_pool_allocator = PagedTokenToKVPoolAllocator(
-                            self.max_total_num_tokens,
+                            self.max_total_num_tokens * self.dcp_size,
-                            page_size=self.page_size,
+                            page_size=self.page_size * self.dcp_size,
                             dtype=self.kv_cache_dtype,
                             device=self.device,
                             kvcache=self.token_to_kv_pool,

python/sglang/srt/models/utils.py

@@ -293,7 +293,12 @@ def enable_fused_set_kv_buffer(forward_batch: ForwardBatch):
         and pool.dtype == torch.bfloat16
         and not isinstance(pool, SWAKVPool)
         and not is_prefill_context_parallel_enabled()
-    ) or (_is_hip and not is_prefill_context_parallel_enabled())
+        and getattr(forward_batch, "dcp_kv_mask", None) is None
+    ) or (
+        _is_hip
+        and not is_prefill_context_parallel_enabled()
+        and getattr(forward_batch, "dcp_kv_mask", None) is None
+    )
 
 
 def create_fused_set_kv_buffer_arg(

python/sglang/srt/server_args.py

@@ -858,6 +858,13 @@ class ServerArgs:
             aliases=["--tensor-parallel-size"],
         ),
     ] = 1
+    dcp_size: A[
+        int,
+        Arg(
+            help="The decode context parallelism size.",
+            aliases=["--decode-context-parallel-size"],
+        ),
+    ] = 1
     pp_size: A[
         int,
         Arg(
@@ -2539,6 +2546,7 @@ class ServerArgs:
         # defaults inspect enable_prefill_cp/cp_strategy.
         self._handle_legacy_cp_arguments()
         self._validate_prefill_only_disable_kv_cache_args()
+        self._handle_dcp_validation()
 
         if self.model_path.lower() in ["none", "dummy"]:
             # Skip for dummy models
@@ -2685,6 +2693,24 @@ class ServerArgs:
         ):
             ObjectStorageModel.download_and_get_path(self.tokenizer_path)
 
+    def _handle_dcp_validation(self):
+        # Decode context parallel (DCP) is currently implemented and validated
+        # only on AMD HIP/ROCm. Reject invalid or unverified configurations
+        # early instead of letting them fail deeper in model initialization.
+        if self.dcp_size < 1:
+            raise ValueError(
+                "Decode context parallel size (--dcp-size / "
+                "--decode-context-parallel-size) must be >= 1, but got "
+                f"dcp_size={self.dcp_size}."
+            )
+        if self.dcp_size > 1 and not is_hip():
+            raise ValueError(
+                "Decode context parallel (--dcp-size / "
+                "--decode-context-parallel-size > 1) is currently only "
+                f"supported on the AMD HIP platform, but got dcp_size="
+                f"{self.dcp_size} on a non-HIP platform."
+            )
+
     def _handle_load_balance_method(self):
         if self.disaggregation_mode not in ("null", "prefill", "decode"):
             raise ValueError(

test/registered/amd/test_qwen3p5_triton_dcp.py

@@ -0,0 +1,90 @@
+import os
+import unittest
+from types import SimpleNamespace
+
+from sglang.srt.utils import kill_process_tree
+from sglang.test.ci.ci_register import register_amd_ci
+from sglang.test.run_eval import run_eval
+from sglang.test.test_utils import (
+    DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
+    is_in_ci,
+    popen_launch_server,
+    write_github_step_summary,
+)
+
+register_amd_ci(
+    est_time=4800, suite="nightly-amd-accuracy-8-gpu-mi35x-qwen35", nightly=True
+)
+
+QWEN35_MODEL_PATH = os.environ.get("QWEN3_5_MODEL_PATH", "Qwen/Qwen3.5-397B-A17B-FP8")
+SERVER_LAUNCH_TIMEOUT = 4800
+TP_SIZE = 8
+DCP_SIZE = 2
+GSM8K_ACCURACY_THRESHOLD = 0.90
+
+
+class TestQwen35TritonDCPGsm8k(CustomTestCase):
+    """Qwen3.5 Triton DCP (tp=8, dcp=2) full GSM8K accuracy on AMD MI35x."""
+
+    @classmethod
+    def setUpClass(cls):
+        cls.model = QWEN35_MODEL_PATH
+        cls.base_url = DEFAULT_URL_FOR_TEST
+
+        other_args = [
+            "--trust-remote-code",
+            "--tp",
+            str(TP_SIZE),
+            "--dcp-size",
+            str(DCP_SIZE),
+            "--attention-backend",
+            "triton",
+            "--context-length",
+            "1048576",
+            "--disable-radix-cache",
+            "--json-model-override-args",
+            (
+                '{"rope_scaling":{"rope_type":"yarn","factor":4.0,'
+                '"original_max_position_embeddings":262144}}'
+            ),
+        ]
+        env = os.environ.copy()
+        env["SGLANG_USE_AITER"] = "1"
+        env["HSA_NO_SCRATCH_RECLAIM"] = "1"
+
+        cls.process = popen_launch_server(
+            cls.model,
+            cls.base_url,
+            timeout=SERVER_LAUNCH_TIMEOUT,
+            other_args=other_args,
+            env=env,
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        kill_process_tree(cls.process.pid)
+
+    def test_a_gsm8k(self):
+        args = SimpleNamespace(
+            base_url=self.base_url,
+            model=self.model,
+            eval_name="gsm8k",
+            api="completion",
+            max_tokens=512,
+            num_examples=1319,
+            num_threads=32,
+            num_shots=5,
+        )
+        metrics = run_eval(args)
+        print(f"{metrics=}")
+
+        if is_in_ci():
+            write_github_step_summary(
+                f"### test_a_gsm8k (qwen3.5-triton-dcp2)\n" f'{metrics["score"]=:.3f}\n'
+            )
+        self.assertGreater(metrics["score"], GSM8K_ACCURACY_THRESHOLD)
+
+
+if __name__ == "__main__":
+    unittest.main()