[Feature] Add fused_sparse_attention_overlap operator for DSA with CPU offloading (#14278)
### What this PR does / why we need it? This PR implement the Sparse KV Cache Offloading for decode part proposed in https://github.com/vllm-project/vllm/issues/48203. Please refer to our RFC for more detail about the motivation and proposed changes of this PR. The code on the framework side will be submitted in a subsequent pull request. This PR adds a new fused AscendC ttentionOverlap` under`csrc/attention/fused_sparse_attention_overlap/`, targeting DeepSeek Sparse Attention (DSA) scenarios where the KV cache is pmory (CPU offloading). In DSA, only the top-k selected tn. With CPU offloading, selectedtokens may hit the on-device KV cache or miss it and reside in host memory. This operator fuses the sparse attention computation with the selection-path KV management (block status / membership map), so that the attention compulaps with and hides the **H2Dloading latency of missed tokens**, instead of stalling the pipeline while waiting for the KV transfer to complete. Main changes: - **op_host**: op definition (aclistered for `ascend910B`, `ascend910C` and `ascend950`), infer-shape/proto, and full tiling implementation supporting `BSND` / `TND` / `PA_BSND` (pageddecode mode. - **op_kernel (arch22)**: For 910C device, kernel implementation for Ascend 910 series, split into cube service, vector service and MLA kernel headers. - **op_kernel (arch35)**: For 950 device, kernel based on the regbase programmingmodel (`util_regbase.h`), with dedicated KV-cache handling (`fused_sparse_attention_overlap_kvcache.h`). - **torch adapter**: `fused_sparst.h` exposing the op to the Pythonside with input checks and output construction (`attention_out`, `selection_kv_actual_seq`). ### Does this PR introduce _any_ user-facing change? Yes. A new custom op `FusedSparseAttentionOverlap` is registered and exposed to the Python side via the torch adapter. No existing APIs or behaviors are changed. ### How was this patch tested? - Added `tests/ut/attention/test_fused_sparse_attention_overlap_precision.py`, which compares the op output against a pure-PyTorch (fp32) reference implementation of paged sparse attention with `rtol=atol=1e-2`, covering both Fn Ascend 950. - End-to-end performance with DSA + CPU offloading was verified on Ascend 910C, showing that the H2D loading of missed tokens is hmputation of hit tokens. - End-to-end integration on Ascend 950 is not enabled yet and will be followed up in a subsequent PR. ### Co-Authors @Misak2333 - vLLM version: v0.27.1 - vLLM main: https://github.com/vllm-project/vllm/commit/ba07e4a48fc951300d97eb506217dd530583dea3 --------- Signed-off-by: lolo-pop <lolopop.rong@gmail.com> Signed-off-by: Misak2333 <167268798+Misak2333@users.noreply.github.com> Co-authored-by: Misak2333 <167268798+Misak2333@users.noreply.github.com>
C
Chase R committed
d1523852b416be34cfea126a9134bb3519771834
Parent: 04da1ce
Committed by GitHub <noreply@github.com>
on 8/25/2026, 9:29:51 AM