Skip to content

fix rope op to adapt with ascend - #1217

Merged
wooway777 merged 2 commits into
InfiniTensor:mainfrom
Jingbo-gao:rope-ascend-adapt
Jun 9, 2026
Merged

wooway777 merged 2 commits into
InfiniTensor:mainfrom
Jingbo-gao:rope-ascend-adapt

Conversation

@Jingbo-gao

@Jingbo-gao Jingbo-gao commented Jun 9, 2026 •

Copy link
Copy Markdown
Contributor

PR内容:修改rope算子在昇腾ascend上对BF16数据格式和GPT_NEOX方法进行适配
修改内容:增加InfiniCore\src\infiniop\ops\rope\ascend\rope_ascend.h中对rope_kernel_neox_launch的声明;修改InfiniCore\src\infiniop\ops\rope\ascend\rope_ascend.cc,增加对GPT_NEOX方法的路由;修改InfiniCore\src\infiniop\ops\rope\ascend\rope_ascend_kernel.cpp,增加GPT_NEOX的实现方法,并且新增BF16实现逻辑,将BF16 cast成FP16,计算完成后再cast成BF16;取消InfiniCore/test/infiniop/rope.py对于昇腾的特殊性注释。
现状:infiniop算子测试全部通过
测试样例:
rope2
测试结果:
rope1

@Jingbo-gao
Jingbo-gao requested a review from a team June 9, 2026 03:41
@wooway777
wooway777 merged commit 099cc3e into InfiniTensor:main Jun 9, 2026
5 checks passed
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