Feat[DCP] Support decode context parallelism for DSA models with the flashmla_kv backend

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

@@ -277,6 +277,10 @@ def _correct_attn_cp_out_kernel(
     # Apply correction and store
     output = tl.load(outputs_ptr + output_offsets)
     output = output * factor
+    # A zero-weight rank's raw output can be NaN (e.g. a degenerate DCP shard
+    # with no owned tokens); NaN * 0.0 == NaN, so guard explicitly rather than
+    # trust the multiply.
+    output = tl.where(factor == 0.0, 0.0, output)
     tl.store(new_output_ptr + new_output_offsets, output)
 
 
@@ -553,9 +557,12 @@ def _dcp_lse_combine_kernel(
             + d_offsets * recv_output_stride_D
         )
         partial_out = tl.load(recv_output_ptr + o_offsets).to(tl.float32)
-        acc += partial_out * w
+        # Empty shards may return NaN outputs; NaN * 0 is not zero.
+        acc += tl.where(w == 0.0, 0.0, partial_out * w)
 
-    acc = acc / weight_sum
+    has_valid_weight = weight_sum > 0.0
+    weight_sum = tl.where(has_valid_weight, weight_sum, 1.0)
+    acc = tl.where(has_valid_weight, acc / weight_sum, 0.0)
 
     out_offsets = (
         batch_idx * out_stride_B + head_idx * out_stride_H + d_offsets * out_stride_D
@@ -567,6 +574,7 @@ def _dcp_lse_combine_kernel(
             global_lse = tl.log(weight_sum) + lse_max
         else:
             global_lse = tl.log2(weight_sum) + lse_max
+        global_lse = tl.where(has_valid_weight, global_lse, -float("inf"))
         out_lse_offset = batch_idx * recv_lse_stride_B + head_idx * recv_lse_stride_H
         tl.store(out_lse_ptr + out_lse_offset, global_lse)
 
@@ -660,7 +668,16 @@ def _lse_weighted_combine_cpu(
         weights = torch.pow(2.0, centered)
 
     weight_sum = weights.sum(dim=0, keepdim=True)
-    weights = weights / weight_sum
+    has_valid_weight = weight_sum > 0
+    safe_weight_sum = torch.where(
+        has_valid_weight, weight_sum, torch.ones_like(weight_sum)
+    )
+    weights = torch.where(
+        has_valid_weight,
+        weights / safe_weight_sum,
+        torch.zeros_like(weights),
+    )
 
-    combined = (partial_outputs * weights.unsqueeze(-1)).sum(dim=0)
+    weights = weights.unsqueeze(-1)
+    combined = torch.where(weights == 0, 0.0, partial_outputs * weights).sum(dim=0)
     return combined

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

@@ -49,6 +49,8 @@ def transform_index_page_table_decode_kernel(
     result_ptr: torch.Tensor,
     page_size: tl.constexpr,
     page_table_row_stride: tl.constexpr,
+    dcp_size: tl.constexpr,
+    dcp_rank: tl.constexpr,
 ):
     TOPK: tl.constexpr = 2048
     req_id = tl.program_id(0)
@@ -60,6 +62,10 @@ def transform_index_page_table_decode_kernel(
     loaded_topk_indices = tl.load(topk_indices_ptr + offset)
     mask = loaded_topk_indices >= 0
     loaded_kv_indices = tl.load(page_table_ptr + loaded_topk_indices, mask=mask)
+    if dcp_size > 1:
+        # Keep slots owned by this rank as local rows; others become -1.
+        mask = mask & (loaded_kv_indices % dcp_size == dcp_rank)
+        loaded_kv_indices = loaded_kv_indices // dcp_size
     tl.store(result_ptr + offset, loaded_kv_indices, mask=mask)
     tl.store(result_ptr + offset, -1, mask=~mask)
 
@@ -80,6 +86,8 @@ def transform_index_page_table_prefill_kernel(
     TOPK: tl.constexpr,
     BLOCK_Q: tl.constexpr,
     BLOCK_TOPK: tl.constexpr,
+    dcp_size: tl.constexpr,
+    dcp_rank: tl.constexpr,
 ):
     request_id = tl.program_id(0)
     query_offsets = tl.program_id(1) * BLOCK_Q + tl.arange(0, BLOCK_Q)
@@ -113,6 +121,10 @@ def transform_index_page_table_prefill_kernel(
         mask=valid_topk_mask,
         other=-1,
     )
+    if dcp_size > 1:
+        # DCP owner filter: keep slots on this rank, map global -> local row.
+        owned_mask = valid_topk_mask & (loaded_kv_indices % dcp_size == dcp_rank)
+        loaded_kv_indices = tl.where(owned_mask, loaded_kv_indices // dcp_size, -1)
     tl.store(
         result_ptr
         + token_indices[:, None] * result_stride_0
@@ -127,6 +139,8 @@ def transform_index_page_table_decode_fast(
     topk_indices: torch.Tensor,
     result: Optional[torch.Tensor] = None,
     page_size: int = 1,
+    dcp_size: int = 1,
+    dcp_rank: int = 0,
 ) -> torch.Tensor:
     """
     Transform the page table according to topk indices for sparse topk attention.
@@ -151,6 +165,8 @@ def transform_index_page_table_decode_fast(
         result,
         page_size,
         page_table_row_stride=page_table.stride(0),
+        dcp_size=dcp_size,
+        dcp_rank=dcp_rank,
     )
     return result
 
@@ -163,6 +179,8 @@ def transform_index_page_table_prefill_fast(
     output_num_tokens: Optional[int] = None,
     page_table_is_expanded: bool = False,
     cu_seqlens_q: Optional[torch.Tensor] = None,
+    dcp_size: int = 1,
+    dcp_rank: int = 0,
 ) -> torch.Tensor:
     assert page_size == 1
     assert topk_indices.shape[1] == 2048
@@ -200,6 +218,8 @@ def transform_index_page_table_prefill_fast(
         TOPK=topk_indices.shape[1],
         BLOCK_Q=block_q,
         BLOCK_TOPK=block_topk,
+        dcp_size=dcp_size,
+        dcp_rank=dcp_rank,
         num_warps=4,
     )
     return result