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

[andrewdoro]

Summary

Makes the FlashMLA sparse prefill path (#25418) correct under DSA prefill context parallelism (--dsa-prefill-cp-mode round-robin-split). Today the sparse path is selected by a CP-blind gate and then runs with global, contiguous query geometry, which is wrong once round-robin CP strides each rank's tokens — producing incorrect long-context output (and, on padded chunks, an out-of-bounds index). This is a focused, minimal replacement for the (closed) draft #26215: 3 files, +172/-23, no kernel changes.

Background / problem

#25418 added _forward_prefill_sparse, gated only by:

is_extend_without_speculative() and (q.shape[0] > _LARGE_INDEXER_QUERY_THRESHOLD
                                     or SGLANG_OPT_FLASHMLA_SPARSE_PREFILL.get())

There is no cp_size/cp_rank term. Under round-robin-split, rank r physically holds the strided subset {r, r+cp_size, ...} of the chunk's tokens (the dense path is corrected for this by apply_cp_reindex), but the sparse path is not:

• SparsePrefillChunkCache.build derives query_start_loc from the full (global) extend_seq_lens, while q_flat only carries this rank's strided rows.
• the combine kernel derives each query's SWA position from a contiguous start_pos + token_idx_in_query, which is false under striding.

So a CP round-robin prefill chunk that hits the sparse gate computes attention against the wrong positions.

Two concrete failure modes from ceil-align padding

When the global token count isn't a multiple of cp_size, the batch is ceil-aligned with trailing padding rows (positions = 0, all--1 page masks) that survive the round-robin split into each rank's q_flat:

1. c128 mask poisoning (silent wrong output). ensure_c128 selects last_q_per_req = query_start_loc[-1] - 1 to source the per-request c128 page-index mask. On ranks holding a trailing padding row, that row's mask is all -1; clamp_min(0) collapses it to slot 0, so the request's entire c128 compressed workspace is dequantized from a polluted slot. (The method's own docstring already warns about this all--1 hazard for the first-query case; padding reintroduces it via the last row.)
2. SWA combine OOB read (crash). A padding row has pos = 0, so the combine kernel emits exactly one SWA index swa_base - gather_start; for any continuation/chunked prefill (gather_start large) this is a large negative value other than -1, which violates flash_mla_sparse_fwd's contract (indices must be -1 or >= s_kv) and reads out of bounds on the exactly-sized KV workspace.

The fix

Two mechanisms, both confined to the sparse path and gated behind the existing trigger so non-CP behavior is unchanged:

1. Striding correctness — thread local_extend_seq_lens (the per-rank local row count) and positions (the CP-reindexed positions_casual) into SparsePrefillChunkCache.build, so query_start_loc and the SWA-window position math use per-rank-local query geometry while the KV gather stays global.

2. Padding safety — a new can_use_sparse_prefill gate runs sparse under CP only for a single, unpadded request:

   is_cp_round_robin -> batch_size == 1
                        and cp_num_rows == q_num_rows
                        and sum(extend_seq_lens_cpu) % cp_size == 0

   Multi-request and ceil-align-padded chunks fall back to the already-CP-correct dense flash_mla_with_kvcache. Because chunked_prefill_size is normally a multiple of cp_size, full chunks still run sparse; only a ragged final chunk falls back, so the perf impact is negligible.

This relies on the request-local c4_sparse_raw_indices that #25418 already produces via the v1 topk path, so no indexer/kernel changes are needed.

Validation

GPU token-exact comparison on 4×B200, tp=4, nightly-dev-cu13-20260604-14ed9b44 + this patch, greedy (temperature=0):

arm	config
reference	dense + CP (--enable-dsa-prefill-context-parallel, sparse env off)
under test	sparse + CP (sparse env on)

10/10 prompts produced byte-identical output between the two arms, including two unpadded ~10.6k-token prompts (%cp_size == 0) that exercise the real sparse path over ~83 c128 blocks, padded prompts that take the dense fallback (no crash), and short coherence prompts. No CUDA errors during inference.

Relation to #26215

#26215 bundled this CP support together with the (now-merged) #25418 integration, a separate mixed-chunk feature, and a stale-base divergence, and is closed. This PR is the isolated CP-correctness delta on top of merged #25418.

Test plan

• python -m py_compile on all touched files
• GPU token-exact sparse+CP == dense+CP (4×B200, tp=4, as above)
• CI

---

CI States

Latest PR Test (Base): ❌ Run #26953815469
Latest PR Test (Extra): ❌ Run #26953814552

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[github-actions]

Thanks @andrewdoro. Closing this because it has had no updates in 99 days.

Reopen it if the work is still relevant.

Some directories moved recently, so an older branch may need retargeting:
sgl-kernel/ -> python/sglang/kernels/aot/, python/sglang/jit_kernel/
-> python/sglang/kernels/jit/, docs/ -> docs/docs/ (.mdx),
bench_serving.py -> benchmark/serving.py, test/srt/ -> test/registered/.