Skip to content

feat: support Ascend950 attention and paged KV cache. - #2083

Open
wuyifan2017 wants to merge 2 commits into
xLLM-AI:mainfrom
wuyifan2017:codex/a5-attention-reshape-cache
Open

feat: support Ascend950 attention and paged KV cache.#2083
wuyifan2017 wants to merge 2 commits into
xLLM-AI:mainfrom
wuyifan2017:codex/a5-attention-reshape-cache

Conversation

@wuyifan2017

@wuyifan2017 wuyifan2017 commented Jul 30, 2026

Copy link
Copy Markdown
Contributor

Description

  • Route Ascend950 paged KV cache writes through the ReshapeAndCacheA5 wrapper and register the wrapper in the NPU kernel target.
  • Add Ascend950 PyTorch SDPA fallbacks for packed TND prefill and paged decode attention, including GQA KV-head expansion.
  • Update the xllm-ops submodule to the implementation from xLLM-AI/xllm-ops.
  • Correct the documented paged-cache layout and add Qwen3.5-27B TP4 coverage for full attention and a KV cache crossing the 128-token block boundary.

Related Issues

Depends on xLLM-AI/xllm-ops#31.

Change Type

  • Bug fix
  • New feature
  • Performance improvement
  • Refactor
  • Documentation
  • Test
  • Build or CI

Pull Request Checklist

PR Title and Commit Messages

  • The PR title and each commit message follow the xLLM commit format: <type>: <subject>.

Pre-commit Checks

  • I have installed pre-commit by running pip install pre-commit or an equivalent command.
  • I have installed the hooks with pre-commit install.
  • I have run pre-commit run --all-files and fixed any reported issues.

Self Review

  • I have self-reviewed the code according to .agents/skills/code-review/references/custom-code-style.md, especially code written or assisted by AI.
  • I have rebased this PR onto the latest main branch.

Build and Test Coverage

  • Tests have been added or updated as needed.
  • CUDA: python setup.py build test has passed on a CUDA machine.
  • NPU: python setup.py build test has passed on an NPU machine.
  • MLU: python setup.py build test has passed on an MLU machine.

Reviewer Notes

  • The fallback covers non-chunked packed TND prefill and standard paged decode on Ascend950. ACL graph decode remains unchanged.
  • pre-commit run --all-files passes.


constexpr int64_t kSwaIntMax = 2147483647;

bool is_ascend950() {

@yingxudeng yingxudeng Jul 31, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

is_ascend950() 和 expand_kv_heads 不同文件完全重复定义。可以考虑放到 core/kernels/npu/utils.h 或者其他合适位置

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

已修改

@wuyifan2017
wuyifan2017 force-pushed the codex/a5-attention-reshape-cache branch from afe5924 to 222ddc0 Compare August 3, 2026 03:02
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants