Skip to content

add random_sample to adapt ascend BF16 format#1191

Open
Jingbo-gao wants to merge 1 commit into
InfiniTensor:mainfrom
Jingbo-gao:randomsample-ascend-adapt
Open

add random_sample to adapt ascend BF16 format#1191
Jingbo-gao wants to merge 1 commit into
InfiniTensor:mainfrom
Jingbo-gao:randomsample-ascend-adapt

Conversation

@Jingbo-gao
Copy link
Copy Markdown
Contributor

PR内容:修改random_sample算子在昇腾ascend上对BF16数据格式进行适配
修改内容:增加src\infiniop\ops\random_sample\ascend\randomsample_aclnn.cc中BF16数据格式的处理分支,使用“aclrtMallocHost” 分配 CPU 内存,“aclrtMemcpy” 将数据从 Device 拷贝到 Host,在 CPU 端使用 “_bf16_to_f32” 函数将 BF16 转换为 FP32,通过“aclrtMemcpy”将 FP32 数据拷贝回 Device,调用 TopK 和 random_sample_kernel 时强制使用 FP32
现状:infiniop算子测试全部通过
注意:infiniop算子测试需要修改I32为I64进行验证
测试样例:
randomsample2
测试结果:
randomsample

@Jingbo-gao Jingbo-gao requested a review from a team June 2, 2026 08:34
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.

1 participant