[DeepSeek V4] Correct FlashMLA sparse prefill under DSA context parallel

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

@@ -51,6 +51,9 @@
 from sglang.srt.layers.attention.dsv4.quant_k_cache import (
     quant_to_nope_fp8_rope_bf16_pack_triton,
 )
+from sglang.srt.layers.attention.dsv4.sparse_prefill_gate import (
+    can_use_sparse_prefill,
+)
 from sglang.srt.layers.attention.dsv4.sparse_prefill_utils import (
     SparsePrefillChunkCache,
 )
@@ -74,6 +77,11 @@
 
 logger = logging.getLogger(__name__)
 
+# Warn-once diagnostics for the sparse-prefill CP gate (per compress_ratio for
+# the enabled path; once for the skip path).
+_SPARSE_PREFILL_ENABLED_LOGGED_RATIOS: set[int] = set()
+_SPARSE_PREFILL_SKIPPED_LOGGED = False
+
 SWA_WINDOW = 128
 C4_TOPK = 512
 PAGE_INDEX_ALIGNED_SIZE = 64
@@ -1117,10 +1125,62 @@ def forward(
                     extra_indices.shape[-1] % 64 == 0
                 ), f"{extra_indices.shape=}'s last dimension is not aligned to 64"
 
-            if forward_batch.forward_mode.is_extend_without_speculative() and (
-                q.shape[0] > _LARGE_INDEXER_QUERY_THRESHOLD
-                or envs.SGLANG_OPT_FLASHMLA_SPARSE_PREFILL.get()
-            ):
+            sparse_prefill_requested = (
+                forward_batch.forward_mode.is_extend_without_speculative()
+                and (
+                    q.shape[0] > _LARGE_INDEXER_QUERY_THRESHOLD
+                    or envs.SGLANG_OPT_FLASHMLA_SPARSE_PREFILL.get()
+                )
+            )
+            use_sparse_prefill = sparse_prefill_requested
+            if use_sparse_prefill:
+                from sglang.srt.layers.attention.dsa.utils import (
+                    is_dsa_prefill_cp_round_robin_split,
+                )
+
+                # Under DSA round-robin CP each rank holds a strided 1/cp_size
+                # subset of the chunk's tokens. The sparse path is only correct
+                # when the chunk is a single request whose local row count
+                # matches the reindexed positions; otherwise fall back to the
+                # (already CP-correct) dense flash_mla path.
+                is_cp_round_robin = is_dsa_prefill_cp_round_robin_split()
+                use_sparse_prefill = can_use_sparse_prefill(
+                    q_num_rows=q.shape[0],
+                    batch_size=forward_batch.batch_size,
+                    extend_seq_lens_cpu=forward_batch.extend_seq_lens_cpu,
+                    is_cp_round_robin=is_cp_round_robin,
+                    cp_num_rows=(
+                        core_attn_metadata.positions_casual.shape[0]
+                        if is_cp_round_robin
+                        else None
+                    ),
+                    cp_size=get_attention_cp_size() if is_cp_round_robin else 1,
+                )
+                global _SPARSE_PREFILL_SKIPPED_LOGGED
+                max_extend_seq_len = max(forward_batch.extend_seq_lens_cpu or [0])
+                if (
+                    not use_sparse_prefill
+                    and max_extend_seq_len > SWA_WINDOW
+                    and not _SPARSE_PREFILL_SKIPPED_LOGGED
+                ):
+                    logger.warning(
+                        "DSV4 sparse prefill requested but skipped (dense "
+                        "fallback): q_rows=%s batch_size=%s "
+                        "extend_seq_lens_cpu=%s is_cp_round_robin=%s "
+                        "cp_num_rows=%s",
+                        q.shape[0],
+                        forward_batch.batch_size,
+                        forward_batch.extend_seq_lens_cpu,
+                        is_cp_round_robin,
+                        (
+                            core_attn_metadata.positions_casual.shape[0]
+                            if is_cp_round_robin
+                            else None
+                        ),
+                    )
+                    _SPARSE_PREFILL_SKIPPED_LOGGED = True
+
+            if use_sparse_prefill:
                 return self._forward_prefill_sparse(
                     q=q,
                     layer_id=layer_id,
@@ -1196,19 +1256,57 @@ def _forward_prefill_sparse(
         # q is (b, 1, h_q, d_qk); flash_mla_sparse_fwd takes (s_q, h_q, d_qk).
         q_flat = q.squeeze(1)
 
+        global _SPARSE_PREFILL_ENABLED_LOGGED_RATIOS
+        if compress_ratio not in _SPARSE_PREFILL_ENABLED_LOGGED_RATIOS:
+            logger.warning(
+                "DSV4 sparse prefill enabled: q_rows=%s compress_ratio=%s "
+                "cp_rank=%s cp_size=%s positions=%s",
+                q_flat.shape[0],
+                compress_ratio,
+                get_attention_cp_rank(),
+                get_attention_cp_size(),
+                core_attn_metadata.positions_casual.shape[0],
+            )
+            _SPARSE_PREFILL_ENABLED_LOGGED_RATIOS.add(compress_ratio)
+
         cache = self.forward_metadata.sparse_prefill_cache
         if cache is None:
             # ``swa_window_size`` on the pool is its storage page size, not
             # the model's SWA window — pass both explicitly.
+            seq_lens = forward_batch.seq_lens.to(torch.int32)
+            extend_seq_lens = forward_batch.extend_seq_lens.to(torch.int32)
+            req_pool_indices = forward_batch.req_pool_indices.to(torch.int32)
+
+            # Under DSA round-robin CP this rank holds a strided subset of the
+            # (single, gate-guaranteed) request's tokens. The chunk cache must
+            # build query_start_loc from the local row count and use the
+            # reindexed causal positions for the SWA-window math, not the
+            # implicit contiguous-prefix assumption that holds without CP.
+            local_extend_seq_lens = None
+            positions = None
+            from sglang.srt.layers.attention.dsa.utils import (
+                is_dsa_prefill_cp_round_robin_split,
+            )
+
+            if is_dsa_prefill_cp_round_robin_split():
+                assert seq_lens.numel() == 1
+                assert (
+                    core_attn_metadata.positions_casual.shape[0] == q_flat.shape[0]
+                )
+                local_extend_seq_lens = seq_lens.new_tensor([q_flat.shape[0]])
+                positions = core_attn_metadata.positions_casual.to(torch.int32)
+
             cache = SparsePrefillChunkCache.build(
-                seq_lens=forward_batch.seq_lens.to(torch.int32),
-                extend_seq_lens=forward_batch.extend_seq_lens.to(torch.int32),
-                req_pool_indices=forward_batch.req_pool_indices.to(torch.int32),
+                seq_lens=seq_lens,
+                extend_seq_lens=extend_seq_lens,
+                req_pool_indices=req_pool_indices,
                 req_to_token=self.req_to_token,
                 full_to_swa=token_to_kv_pool.full_to_swa_index_mapping,
                 swa_window_size=SWA_WINDOW,
                 swa_page_size=token_to_kv_pool.swa_window_size,
                 num_qo_tokens=q_flat.shape[0],
+                local_extend_seq_lens=local_extend_seq_lens,
+                positions=positions,
             )
             self.forward_metadata.sparse_prefill_cache = cache
 

python/sglang/srt/layers/attention/dsv4/sparse_prefill_gate.py

@@ -0,0 +1,28 @@
+from typing import List, Optional
+
+
+def can_use_sparse_prefill(
+    *,
+    q_num_rows: int,
+    batch_size: int,
+    extend_seq_lens_cpu: Optional[List[int]],
+    is_cp_round_robin: bool,
+    cp_num_rows: Optional[int] = None,
+    cp_size: int = 1,
+) -> bool:
+    if extend_seq_lens_cpu is None:
+        return False
+    if is_cp_round_robin:
+        # Single request only (the round-robin interleave can't be described
+        # by one query_start_loc across multiple requests), local row count
+        # must match the reindexed positions, AND the chunk must be unpadded:
+        # when the global token count isn't a multiple of cp_size the batch is
+        # ceil-aligned with trailing padding rows (pos=0, all-(-1) page masks)
+        # that poison the c128 mask selection and write OOB SWA combine
+        # indices. Those padded (ragged-tail) chunks fall back to dense.
+        return (
+            batch_size == 1
+            and cp_num_rows == q_num_rows
+            and sum(extend_seq_lens_cpu) % cp_size == 0
+        )
+    return sum(extend_seq_lens_cpu) == q_num_rows

python/sglang/srt/layers/attention/dsv4/sparse_prefill_utils.py

@@ -1,10 +1,7 @@
 """Per-query sparse-index combiner for the FlashMLA sparse prefill path.
 
-Adapts vllm's ``combine_topk_swa_indices`` to sglang's flat-workspace layout.
-Reference:
-https://github.com/vllm-project/vllm/blob/124fac10cb0ea83aee2ffeabac0b413d6b759b26/vllm/models/deepseek_v4/common/ops/cache_utils.py#L476
-
-For each
+Adapts vllm's ``combine_topk_swa_indices`` (vllm/v1/attention/ops/
+deepseek_v4_ops/cache_utils.py) to sglang's flat-workspace layout. For each
 query token in a prefill chunk, emits one row of combined indices into the
 chunk's bf16 KV workspace:
 
@@ -68,6 +65,7 @@ def combine_topk_swa_indices(
     topk: int,
     out_indices: Optional[torch.Tensor] = None,
     out_lens: Optional[torch.Tensor] = None,
+    positions: Optional[torch.Tensor] = None,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """Combine topk + SWA indices into a single ``flash_mla_sparse_fwd`` row.
 
@@ -110,6 +108,9 @@ def combine_topk_swa_indices(
     assert gather_lens.dtype == torch.int32
     assert compressed_base.dtype == torch.int32
     assert swa_base.dtype == torch.int32
+    if positions is not None:
+        assert positions.dtype == torch.int32
+        assert positions.shape == (topk_indices.shape[0],)
     assert compress_ratio >= 1, "COMPRESS_RATIO must be >= 1 (use TOP_K=0 for SWA-only)"
 
     num_tokens = topk_indices.shape[0]
@@ -127,14 +128,17 @@ def combine_topk_swa_indices(
         assert out_indices.dtype == torch.int32
         combined_indices = out_indices
     if out_lens is None:
-        combined_lens = torch.zeros(
+        combined_lens = torch.empty(
             num_tokens, dtype=torch.int32, device=topk_indices.device
         )
     else:
         assert out_lens.shape == (num_tokens,)
         assert out_lens.dtype == torch.int32
         combined_lens = out_lens
 
+    if positions is None:
+        positions = torch.empty(0, dtype=torch.int32, device=topk_indices.device)
+
     NUM_WORKERS = 128
     _combine_topk_swa_indices_kernel[(num_reqs, NUM_WORKERS)](
         combined_indices,
@@ -147,10 +151,12 @@ def combine_topk_swa_indices(
         gather_lens,
         compressed_base,
         swa_base,
+        positions,
         TOP_K=topk,
         COMPRESS_RATIO=compress_ratio,
         WINDOW_SIZE=window_size,
         PADDED_TOP_K=triton.next_power_of_2(topk_indices.shape[-1]),
+        HAS_POSITIONS=positions.numel() != 0,
     )
     return combined_indices, combined_lens
 
@@ -242,7 +248,7 @@ def _build_swa_token_ids_kernel(
 
     first_pos = tl.load(swa_first_pos_ptr + batch_idx)
     gather_len = tl.load(swa_gather_lens_ptr + batch_idx)
-    out_off = tl.load(swa_offsets_ptr + batch_idx).to(tl.int64)
+    out_off = tl.load(swa_offsets_ptr + batch_idx)
     req_pool_idx = tl.load(req_pool_indices_ptr + batch_idx).to(tl.int64)
 
     for i in range(worker_id, gather_len, num_workers):
@@ -266,10 +272,12 @@ def _combine_topk_swa_indices_kernel(
     gather_lens_ptr,
     compressed_base_ptr,
     swa_base_ptr,
+    positions_ptr,
     TOP_K: tl.constexpr,
     COMPRESS_RATIO: tl.constexpr,
     WINDOW_SIZE: tl.constexpr,
     PADDED_TOP_K: tl.constexpr,
+    HAS_POSITIONS: tl.constexpr,
 ):
     batch_idx = tl.program_id(0)
     worker_id = tl.program_id(1)
@@ -293,24 +301,24 @@ def _combine_topk_swa_indices_kernel(
 
     for token_idx in range(query_start + worker_id, query_end, num_workers):
         token_idx_in_query = token_idx - query_start
-        pos = start_pos + token_idx_in_query
+        if HAS_POSITIONS:
+            pos = tl.load(positions_ptr + token_idx)
+        else:
+            pos = start_pos + token_idx_in_query
         # Both the C4 indexer and the C128 metadata builder emit
         # min((pos+1)//compress_ratio, topk_tokens) valid entries. Caller
         # passes TOP_K=0 for SWA-only layers to zero this out.
         topk_len = tl.minimum((pos + 1) // COMPRESS_RATIO, TOP_K)
         swa_len = tl.minimum(pos + 1, WINDOW_SIZE)
 
-        combined_row = token_idx.to(tl.int64) * combined_indices_stride
-        topk_row = token_idx.to(tl.int64) * topk_indices_stride
-
         offset = tl.arange(0, PADDED_TOP_K)
         mask = offset < topk_len
         topk_vals = tl.load(
-            topk_indices_ptr + topk_row + offset,
+            topk_indices_ptr + token_idx * topk_indices_stride + offset,
             mask=mask,
         )
         tl.store(
-            combined_indices_ptr + combined_row + offset,
+            combined_indices_ptr + token_idx * combined_indices_stride + offset,
             topk_vals + compressed_base,
             mask=mask,
         )
@@ -320,7 +328,10 @@ def _combine_topk_swa_indices_kernel(
         # For positions [pos - swa_len + 1, pos], the buffer offsets are
         # [pos - swa_len + 1 - gather_start, pos - gather_start].
         tl.store(
-            combined_indices_ptr + combined_row + topk_len + offset,
+            combined_indices_ptr
+            + token_idx * combined_indices_stride
+            + topk_len
+            + offset,
             swa_base + offset + pos - swa_len + 1 - gather_start,
             mask=offset < swa_len,
         )
@@ -357,6 +368,7 @@ class SparsePrefillChunkCache:
     swa_first_pos: torch.Tensor  # (num_reqs,) int32
     swa_gather_lens: torch.Tensor  # (num_reqs,) int32
     swa_offsets: torch.Tensor  # (num_reqs+1,) int32
+    positions: Optional[torch.Tensor] = None
 
     # c0 pre-computed combine output (entire input set is chunk-invariant).
     c0_combined_indices: torch.Tensor = field(default=None)
@@ -395,12 +407,19 @@ def build(
         swa_window_size: int,
         swa_page_size: int,
         num_qo_tokens: int,
+        local_extend_seq_lens: Optional[torch.Tensor] = None,
+        positions: Optional[torch.Tensor] = None,
     ) -> "SparsePrefillChunkCache":
         device = seq_lens.device
         num_reqs = seq_lens.shape[0]
+        if local_extend_seq_lens is None:
+            local_extend_seq_lens = extend_seq_lens
 
         query_start_loc = torch.zeros(num_reqs + 1, dtype=torch.int32, device=device)
-        query_start_loc[1:] = torch.cumsum(extend_seq_lens, dim=0).to(torch.int32)
+        query_start_loc[1:] = torch.cumsum(local_extend_seq_lens, dim=0).to(torch.int32)
+        assert (
+            int(query_start_loc[-1].item()) == num_qo_tokens
+        ), f"local query rows {int(query_start_loc[-1].item())} != {num_qo_tokens}"
 
         swa_token_ids, swa_first_pos, swa_gather_lens, swa_offsets = (
             build_swa_token_ids(
@@ -424,6 +443,7 @@ def build(
             swa_first_pos=swa_first_pos,
             swa_gather_lens=swa_gather_lens,
             swa_offsets=swa_offsets,
+            positions=positions,
         )
 
         # Pre-compute the c0 combine output: TOPK=0, compressed_base=0,
@@ -441,6 +461,7 @@ def build(
             window_size=swa_window_size,
             compress_ratio=1,
             topk=0,
+            positions=positions,
         )
         cache.c0_workspace = torch.empty(
             (swa_token_ids.shape[0], 1, WORKSPACE_DIM),
@@ -494,6 +515,7 @@ def ensure_c128(self, c128_page_indices: torch.Tensor) -> None:
             window_size=self.swa_window_size,
             compress_ratio=128,
             topk=c128_max,
+            positions=self.positions,
         )
 
         self.c128_flat_token_ids = flat_c128_ids
@@ -569,7 +591,7 @@ def combine_c4_layer(
                 dtype=torch.int32,
                 device=device,
             )
-            self.c4_combined_lens = torch.zeros(
+            self.c4_combined_lens = torch.empty(
                 self.num_qo_tokens, dtype=torch.int32, device=device
             )
         return combine_topk_swa_indices(
@@ -584,4 +606,5 @@ def combine_c4_layer(
             topk=topk,
             out_indices=self.c4_combined_indices,
             out_lens=self.c4_combined_lens,
+            positions=self.positions,
         )