Skip to content

feat: adapt flash attention to Ascend backend#1401

Open
Jingbo-gao wants to merge 7 commits into
InfiniTensor:mainfrom
Jingbo-gao:feat_flash_atten
Open

feat: adapt flash attention to Ascend backend#1401
Jingbo-gao wants to merge 7 commits into
InfiniTensor:mainfrom
Jingbo-gao:feat_flash_atten

Conversation

@Jingbo-gao

Copy link
Copy Markdown
Contributor

PR内容:添加昇腾flash attention特性适配

  1. 主要文件修改:
    python/infinicore/ops/ascend_flash_attn.py:新增基于 torch_npu 的 Ascend paged FlashAttention prefill/decode Python 封装;
    src/infinicore/ops/mha_kvcache/mha_kvcache_flashattn.cc:将 Ascend decode 请求隔离出通用 CUDA FlashAttention 路径并交由专用实现处理;
    src/infinicore/ops/mha_kvcache/mha_kvcache_flashattn_ascend.cc:新增基于 aclnnFusedInferAttentionScoreV4 的 Ascend paged decode 实现,支持 BNSD Query 和 BnBsH KV cache;
    src/infinicore/ops/multi_head_attention_varlen/mha_varlen_flashattn.cc:将 Ascend varlen 请求路由到专用实现,同时保留其他后端原有 ATen、FlashAttention 和 SDPA 行为;
    src/infinicore/ops/multi_head_attention_varlen/mha_varlen_flashattn_ascend.cc:新增基于 FIA V4 的 Ascend TND varlen prefill 实现,支持 paged KV cache、累计序列长度和因果掩码;
    src/infiniop/ops/flash_attention/operator.cc:新增 Ascend FlashAttention descriptor 的设备分发入口;
  2. 推理精度正常:
    ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 python examples/bench.py --device ascend --model=/workspace/models/FM9G_70B_SFT_MHA/ --tp=8 --input-len=32 --output-len=256 --batch-size=16 --enable-paged-attn --attn=flash-attn --enable-graph
图片_137_08

@Jingbo-gao
Jingbo-gao requested a review from a team July 16, 2026 13:41
wooway777
wooway777 previously approved these changes Jul 22, 2026

@wooway777 wooway777 left a comment

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.

Image

其他测试和推理都跑不了了,报错一致。

另外建议把昇腾的实现放在单独文件夹里,参考Hygon。

可以把commits squash一下

@wooway777
wooway777 dismissed their stale review July 22, 2026 03:11

误点

@wooway777 wooway777 left a comment

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.

参考上一条评论

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