feat(dcp): support GPU MLA and DSA decode context parallelism - #1718

Summary

Extend GPU decode context parallelism (DCP) to ordinary MLA, hybrid MLA/KDA, and DSA, and shard the existing DeepSeek V4 Index-K cache. Reuse the existing physical cache pools and scheduler-owned cache groups instead of introducing parallel sharded-pool implementations.

This is a draft for design review. GPU kernel/reference checks pass, but model-level numerical validation is incomplete, especially for DSA. This is not a claim of bitwise equivalence across DCP sizes.

Backend scope

Index scoring uses DeepGEMM for supported NVIDIA SM90+ packed-FP8 configurations, with a portable Triton implementation behind the same candidate API.

Blackwell is backend-dependent, not a blanket GPU support claim. The FlashInfer/TRTLLM sparse attention path currently advertises return_lse=False, so it cannot supply DCP partials. Selection must find a compatible LSE-capable implementation; registering a Hopper+ kernel is not an end-to-end Blackwell model validation. Dense FlashMLA coverage remains SM90-only in the reported tests.

PD cache transfer and KVStore are not supported with DCP. The new FlashMLA/GPU DSA paths reject speculative execution. This PR does not extend Ascend DCP.

Design

Placement and storage

• Recipes declare each cache group's shard_count. The arena remains the source of physical storage geometry; the scheduler continues to allocate virtual blocks through the existing cache-group contract.
• CachePlacement describes ownership and virtual-to-local translation. resolve_cache_slots returns physical slots plus an ownership mask. Writers mask both source loads and destination stores: foreign rows must not write to a dummy/null slot.
• Keep the existing set_mla_kv_buffer / get_mla_kv_buffer pool interfaces and physical pool implementations. Placement and collective reconstruction live above storage, avoiding separate DCP/non-DCP cache-pool classes.
• Shared page-table metadata supplies the layout a kernel needs: position-preserving tables with -1 holes, or compact owned pages with local lengths and original logical positions. Compaction does not change causal positions.
• Keep recurrent KDA state replicated. Give DSA Index-K its own cache group and ownership geometry; do not assume its virtual IDs equal attention-KV IDs. Capacity follows physical demand, so replicated groups can remain the limiting resource.

Prefill and decode

Dense MLA: prefill reconstructs only the requested, bounded history chunk in logical order, then runs the existing TP-local attention path. Decode gathers actual query heads after normalization/RoPE, attends to local KV pages, and merges context partials.

DSA: both prefill and decode select a global sparse candidate set, map it to each rank's owned cache slots, gather query heads, and compute local sparse attention. This is context sharding, not query-context parallelism (QCP).

For local attention outputs O_r and natural-log LSE L_r, the shared merge weights each partial by exp(L_r - logsumexp(L)). It gathers FP32 LSE, weights partials in FP32, and reduce-scatters to the TP head owner before restoring the output dtype. Empty shards contribute zero output and -inf LSE; a sink, when present, is applied once after reconstruction. FP32 merge arithmetic does not undo rounding in a BF16 local output. No experimental FlashMLA FP32-output kernel is included.

Distributed Indexer Top-K

Each rank scores the same queries against its local Index-K pages and retains up to K candidates. Candidate positions are mapped back to global logical offsets; positions and FP32 scores are all-gathered, then one global torch.topk selects K from the union. No full Index-K replication is required for this selection.

Local Top-K is sufficient because an entry excluded from a rank's local K cannot be required for a global top-K set, apart from interchangeable equal-score ties. Equal-score ordering is intentionally unspecified. Invalid/NaN candidates, padded entries, valid counts, causal tails, and forced initial/local windows are handled explicitly. Bounded query tiling limits temporary score memory for both prefill and decode.

Validation

The port's recorded validation for 0140ec34 includes full pre-commit run --all-files, 297 runtime/cache tests, 63 GPU kernel tests, and a two-rank distributed Top-K plus attention/LSE comparison against unsharded references. Eight runtime tests were skipped, including SM90-only dense FlashMLA tests on Blackwell.

The following model measurements are from the predecessor implementation 5d9460cd, before the port to main, not an end-to-end validation of this PR head. H200, TP4, DCP1 vs DCP4, eager, fixed continuation tokens; no FP32-output experiment patches: