[DCP] Restore GLM-5.3 Flash decode CP and fix virtual KV addressing

python/sglang/kernels/ops/attention/dsa/tilelang_kernel.py

@@ -271,6 +271,7 @@ def sparse_attention_fwd_kernel_v1(
     block_I=64,
     num_stages=2,
     threads=256,
+    return_lse=False,
 ):
     assert dim == tilelang.math.next_power_of_2(dim) or dim % 64 == 0, (
         f"dim={dim} must be a power of 2 or a multiple of 64"
@@ -318,13 +319,8 @@ def sparse_attention_fwd_kernel_v1(
 
     H_per_block = padded_H if REPLICATE_H == 1 else 64
 
-    @T.prim_func
-    def main(
-        Q: T.Tensor(q_shape, dtype),  # type: ignore
-        KV: T.Tensor(kv_shape, dtype),  # type: ignore
-        Indices: T.Tensor(indices_shape, indices_dtype),  # type: ignore
-        Output: T.Tensor(o_shape, dtype),  # type: ignore
-    ):
+    @T.macro
+    def body(Q, KV, Indices, Output, LSE):
         with T.Kernel(seq_len * REPLICATE_H, batch, kv_group, threads=threads) as (
             bx,
             by,
@@ -419,13 +415,42 @@ def main(
 
             # Rescale
             for h_i, d_i in T.Parallel(H_per_block, D):
-                acc_o[h_i, d_i] /= sumexp[h_i]
+                if return_lse:
+                    # NOTE(kpham-sgl): Empty ranks need zero output and -inf LSE for the merge.
+                    acc_o[h_i, d_i] /= T.if_then_else(sumexp[h_i] > 0, sumexp[h_i], 1)
+                else:
+                    acc_o[h_i, d_i] /= sumexp[h_i]
             for h_i in T.Parallel(H_per_block):
                 sumexp[h_i] = T.log2(sumexp[h_i]) + m_i[h_i] * sm_scale
+            if return_lse:
+                T.copy(sumexp, LSE[b_i, s_i, H0:H1])
 
             T.copy(acc_o, O_shared)
             T.copy(acc_o, Output[b_i, s_i, H0:H1, :])
 
+    if return_lse:
+
+        @T.prim_func
+        def main(
+            Q: T.Tensor(q_shape, dtype),
+            KV: T.Tensor(kv_shape, dtype),
+            Indices: T.Tensor(indices_shape, indices_dtype),
+            LSE: T.Tensor([batch, seq_len, num_heads], accum_dtype),
+            Output: T.Tensor(o_shape, dtype),
+        ):
+            body(Q, KV, Indices, Output, LSE)
+
+    else:
+
+        @T.prim_func
+        def main(
+            Q: T.Tensor(q_shape, dtype),
+            KV: T.Tensor(kv_shape, dtype),
+            Indices: T.Tensor(indices_shape, indices_dtype),
+            Output: T.Tensor(o_shape, dtype),
+        ):
+            body(Q, KV, Indices, Output, Output)
+
     return main
 
 
@@ -1328,14 +1353,26 @@ def tilelang_sparse_fwd(
     indices: torch.Tensor,
     sm_scale: float,
     d_v: int = 512,
-) -> torch.Tensor:
+    return_lse: bool = False,
+) -> torch.Tensor | Tuple[torch.Tensor, torch.Tensor]:
     assert q.dim() == 3 and kv.dim() == 3 and indices.dim() == 3
     num_heads = q.shape[1]
     dim = q.shape[2]
     tail_dim = dim - d_v
     topk = indices.shape[-1]
     assert topk % 64 == 0, "topk must be padded to a multiple of 64"
 
+    if return_lse:
+        assert not _is_hip and tail_dim == 0, "Sparse DCP LSE requires CUDA NoPE MLA"
+        kernel = sparse_attention_fwd_kernel_v1(
+            num_heads, d_v, tail_dim, topk, sm_scale=sm_scale, return_lse=True
+        )
+        lse = torch.empty(
+            (1, q.shape[0], num_heads), dtype=torch.float32, device=q.device
+        )
+        out = kernel(q.unsqueeze(0), kv.unsqueeze(0), indices.unsqueeze(0), lse)
+        return out, lse
+
     if _is_hip:
         is_fp8_kv = kv.dtype in (torch.float8_e4m3fn, torch.float8_e4m3fnuz)
         if is_fp8_kv:

python/sglang/kernels/ops/kvcache/mla_buffer.py

@@ -88,10 +88,13 @@ def set_mla_kv_buffer_kernel_norope(
     kv_buffer_ptr,
     cache_k_nope_ptr,
     loc_ptr,
+    reserved_skip_index,
     buffer_stride: tl.constexpr,
     nope_stride: tl.constexpr,
     nope_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)
@@ -105,13 +108,15 @@ def set_mla_kv_buffer_kernel_norope(
         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 != reserved_skip_index) & (loc % DCP_WORLD_SIZE == DCP_RANK)
+    safe_loc = tl.where(is_valid, loc, 0) // DCP_WORLD_SIZE
+    dst_ptr = kv_buffer_ptr + safe_loc * buffer_stride + offs
 
     src = tl.load(
         cache_k_nope_ptr + pid_loc * nope_stride + offs,
         mask=mask,
     )
-    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()
@@ -177,10 +182,13 @@ def _set_mla_kv_buffer_impl(
             kv_buffer,
             cache_k_nope,
             loc,
+            reserved_skip_index,
             kv_buffer.stride(0),
             cache_k_nope.stride(0),
             nope_dim,
             BLOCK=BLOCK,
+            DCP_RANK=dcp_rank,
+            DCP_WORLD_SIZE=dcp_world_size,
             **pdl_kwargs,
         )
         return

python/sglang/srt/arg_groups/parallel_hook.py

@@ -9,7 +9,6 @@
 
 from sglang.srt.arg_groups.overrides import (
     _data_parallelism_defaults,
-    _dcp_comm_backend_default,
     _dp_lm_head_validation,
     _tp_lm_head_all_to_all_default,
     declare_resolution,
@@ -122,7 +121,8 @@ def handle_context_parallelism(server_args: Any):
 
 
 def handle_decode_context_parallelism(server_args: Any):
-    run_post_process_pass(server_args, _dcp_comm_backend_default)
+    from sglang.srt.configs.model_config import is_deepseek_dsa
+
     cfg = resolving_view(server_args)
     if cfg.dcp_size < 1:
         raise ValueError(
@@ -154,6 +154,29 @@ def handle_decode_context_parallelism(server_args: Any):
                 f"got --dcp-comm-backend={cfg.dcp_comm_backend}."
             )
 
+    if (
+        cfg.dcp_size == 1
+        or parse_connector_type(cfg.model_path) == ConnectorType.INSTANCE
+    ):
+        return
+
+    model_config = model_config_of(server_args)
+    if is_deepseek_dsa(model_config.hf_config) and (
+        not get_platform().is_cuda
+        or model_config.qk_rope_head_dim != 0
+        or cfg.dsa_prefill_backend != "tilelang"
+        or cfg.dsa_decode_backend != "tilelang"
+        or cfg.dsa_topk_backend == "torch"
+        or not envs.SGLANG_DSA_FUSE_TOPK.get()
+        or cfg.enable_hisparse
+        or cfg.enable_prefill_cp
+    ):
+        raise ValueError(
+            "DSA decode context parallelism requires CUDA NoPE MLA, "
+            "tilelang prefill/decode, and fused top-k; "
+            "HiSparse and prefill CP cannot be combined with it."
+        )
+
 
 def handle_data_parallelism(server_args: Any):
     # The dp_size==1 resets moved to the resolution pipeline

python/sglang/srt/arg_groups/pipeline.py

@@ -12,6 +12,7 @@
 
 from sglang.srt.arg_groups.arg_utils import record_fields
 from sglang.srt.arg_groups.overrides import (
+    _dcp_comm_backend_default,
     _page_size_default,
     _pipeline_parallel_overlap_disable,
     _sampling_backend_default,
@@ -179,7 +180,7 @@ def run_resolution_pipeline(server_args: Any) -> None:
     )
 
     run_hook(validate_prefill_only_disable_kv_cache_args, server_args)
-    run_hook(handle_decode_context_parallelism, server_args)
+    run_post_process_pass(server_args, _dcp_comm_backend_default)
 
     # Model-arch prefill CUDA-graph default must land before cuda-graph
     # resolution (the declarative registry materializes too late to affect
@@ -241,6 +242,7 @@ def run_resolution_pipeline(server_args: Any) -> None:
 
     run_hook(handle_model_specific_adjustments, server_args)
     run_hook(default_unset_prefill_decode_interval, server_args)
+    run_hook(handle_decode_context_parallelism, server_args)
     # After the model overrides: Qwen4-Exp declares the PLE offload default there.
     run_hook(handle_offload_compatibility, server_args)
 

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

@@ -348,11 +348,18 @@ def __init__(
         assert model_runner.req_to_token_pool is not None
         self.req_to_token_pool = model_runner.req_to_token_pool
         self.token_to_kv_pool = model_runner.token_to_kv_pool
+        allocator = model_runner.token_to_kv_pool_allocator
+        self.kv_address_space_size = allocator.size_full + allocator.page_size
+        self._dcp_sharded_kv = (
+            get_parallel().dcp_enabled and not model_runner.is_draft_worker
+        )
         self.hisparse_coordinator = model_runner.hisparse_coordinator
         self.req_to_token = model_runner.req_to_token_pool.req_to_token
 
         self.use_mha: bool = False
-        self.supports_mha_one_shot: bool = True
+        # TODO(kpham-sgl): Evaluate whether to enable MHA one-shot with DCP;
+        # handle sharded target and replicated draft prefix KV if enabled.
+        self.supports_mha_one_shot: bool = not get_parallel().dcp_enabled
         self.dsa_prefill_impl: _DSA_IMPL_T = get_exec().kernel.dsa_prefill_backend
         self.dsa_decode_impl: _DSA_IMPL_T = get_exec().kernel.dsa_decode_backend
         self.dsa_topk_backend: DSATopKBackend = DSATopKBackend.resolve(model_runner)
@@ -1045,16 +1052,11 @@ def init_forward_metadata(self, forward_batch: ForwardBatch):
                     f"{page_table_1_flattened.shape[0] = } must be the same as {sum(indexer_seq_lens_cpu) = }"
                 )
 
-                # Validate indices when logical tokens exceed physical capacity
-                # This is likely to be triggered by PP with high kv reuse & parallelism
-                kv_cache_capacity = (
-                    self.token_to_kv_pool.size + self.token_to_kv_pool.page_size
-                )
-                if forward_batch.seq_lens_sum > kv_cache_capacity:
+                if forward_batch.seq_lens_sum > self.kv_address_space_size:
                     max_idx = page_table_1_flattened.max().item()
-                    assert max_idx < kv_cache_capacity, (
+                    assert max_idx < self.kv_address_space_size, (
                         f"Invalid page table index: max={max_idx}, "
-                        f"kv_cache_capacity={kv_cache_capacity}"
+                        f"kv_address_space_size={self.kv_address_space_size}"
                     )
 
             if topk_transform_method == TopkTransformMethod.RAGGED:
@@ -1992,6 +1994,13 @@ def forward_extend(
                     cu_seqlens_q=metadata.cu_seqlens_q,
                 )
 
+        if self._dcp_sharded_kv:
+            if forward_batch.forward_mode.is_extend_without_speculative():
+                assert k is not None
+                kv_cache = self._dcp_gather_extend_kv(layer, forward_batch, k)
+            elif forward_batch.forward_mode.is_target_verify():
+                page_table_1 = self._dcp_global_to_local_kv_indices(page_table_1)
+
         # todo hisparse: to cover more backends
         if self.hisparse_coordinator is not None:
             # flash_mla_sparse_fwd / tilelang require int32 page indices.
@@ -2012,6 +2021,10 @@ def forward_extend(
                 page_table_1=page_table_1,
                 sm_scale=layer.scaling,
                 v_head_dim=layer.v_head_dim,
+                return_lse=(
+                    self._dcp_sharded_kv
+                    and forward_batch.forward_mode.is_target_verify()
+                ),
             )
         elif dsa_impl == "triton":
             from sglang.kernels.ops.attention.dsa.triton_sparse_mla import (
@@ -2272,6 +2285,9 @@ def forward_decode(
                 page_size=1,
             )
 
+        if self._dcp_sharded_kv:
+            page_table_1 = self._dcp_global_to_local_kv_indices(page_table_1)
+
         if dsa_impl == "flashmla_sparse":
             if q_rope is not None:
                 q_all = concat_mla_absorb_q_general(q_nope, q_rope)
@@ -2321,6 +2337,7 @@ def forward_decode(
                 page_table_1=page_table_1,
                 sm_scale=layer.scaling,
                 v_head_dim=layer.v_head_dim,
+                return_lse=self._dcp_sharded_kv,
             )
         elif dsa_impl == "triton":
             return self._forward_triton_decode(
@@ -2958,14 +2975,45 @@ def _forward_standard_mha(
             causal=causal,
         )
 
+    @staticmethod
+    def _dcp_global_to_local_kv_indices(page_table: torch.Tensor) -> torch.Tensor:
+        # TODO(kpham-sgl): Fuse the index conversion and masking into one GPU kernel.
+        parallel = get_parallel()
+        owned = (page_table >= 0) & (
+            page_table % parallel.attn_dcp_size == parallel.attn_dcp_rank
+        )
+        return torch.where(owned, page_table // parallel.attn_dcp_size, -1)
+
+    def _dcp_gather_extend_kv(
+        self, layer: RadixAttention, forward_batch: ForwardBatch, k: torch.Tensor
+    ) -> torch.Tensor:
+        from sglang.srt.layers.dcp.comm import all_gather_kv_cache_for_mla_extend
+
+        metadata = forward_batch.attn_dcp_metadata
+        k_nope = k.view(k.shape[0], 1, self.kv_lora_rank)
+        all_gather_kv_cache_for_mla_extend(
+            token_to_kv_pool=self.token_to_kv_pool,
+            attn_mqa=layer,
+            extend_prefix_lens_cpu=forward_batch.extend_prefix_lens_cpu,
+            dcp_local_prefix_kv_indices=metadata.dcp_local_prefix_kv_indices,
+            dcp_extend_prefix_lens_sum=metadata.dcp_extend_prefix_lens_sum,
+            dcp_kv_buffer=metadata.dcp_kv_buffer,
+            kv_lora_rank=self.kv_lora_rank,
+            k_nope=k_nope,
+            k_pe=k_nope[..., :0],
+        )
+        # NOTE(kpham-sgl): RAGGED top-k requires this per-request [prefix; extend] order.
+        return metadata.dcp_kv_buffer[metadata.dcp_kv_indices]
+
     def _forward_tilelang(
         self,
         q_all: torch.Tensor,
         kv_cache: torch.Tensor,
         v_head_dim: int,
         page_table_1: torch.Tensor,
         sm_scale: float,
-    ) -> torch.Tensor:
+        return_lse: bool = False,
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
         from sglang.kernels.ops.attention.dsa.tilelang_kernel import tilelang_sparse_fwd
 
         # KPool appends up to index_kpool - 1 live tail tokens to the fixed
@@ -2981,13 +3029,18 @@ def _forward_tilelang(
                 dim=-1,
             )
 
-        return tilelang_sparse_fwd(
+        result = tilelang_sparse_fwd(
             q=q_all,
             kv=kv_cache,
             indices=page_table_1.unsqueeze(1),
             sm_scale=sm_scale,
             d_v=v_head_dim,
+            return_lse=return_lse,
         )
+        if return_lse:
+            out, lse = result
+            return out, lse.squeeze(0)
+        return result
 
     def _forward_triton_decode(
         self,
@@ -3548,12 +3601,15 @@ def set_dsa_prefill_impl(self, forward_batch: Optional[ForwardBatch] = None):
                 self.dsa_prefill_impl = "flashmla_sparse"
 
     def get_topk_transform_method(
-        self, forward_mode: Optional[ForwardMode] = None
+        self, forward_mode: ForwardMode
     ) -> TopkTransformMethod:
         """
         SGLANG_DSA_FUSE_TOPK controls whether to fuse the topk transform into the topk kernel.
         This method is used to select the topk transform method which can be fused or unfused.
         """
+        # Note(kpham-sgl): Gathered prefill KV uses sequence offsets, not cache slots.
+        if self._dcp_sharded_kv and forward_mode.is_extend_without_speculative():
+            return TopkTransformMethod.RAGGED
         if (
             # disable for MTP
             self.dsa_kv_cache_store_fp8

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

@@ -210,6 +210,8 @@ def all_gather_kv_cache_for_mha_extend(
     prefix_kv_a, prefix_k_pe = token_to_kv_pool.get_mla_kv_buffer(
         attn_mqa, dcp_local_prefix_kv_indices, dst_dtype=kv_a.dtype
     )
+    if prefix_k_pe is None:
+        prefix_k_pe = prefix_kv_a[..., :0]
     extend_prefix_lens_cpu = torch.tensor(extend_prefix_lens_cpu)
     gathered_kv_cache = all_gather_kv_cache_for_dcp(
         prefix_kv_a,
@@ -285,6 +287,8 @@ def all_gather_kv_cache_for_mla_extend(
             attn_mqa,
             dcp_local_prefix_kv_indices,
         )
+        if cache_k_rope is None:
+            cache_k_rope = cache_k_nope[..., :0]
         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(