SIGN IN SIGN UP

[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