CATLASS FlashAttention推理示例

FlashAttentionInfer Example Readme【免费下载链接】catlass本项目是CANN的算子模板库提供NPU上高性能矩阵乘及其相关融合类算子模板样例。项目地址: https://gitcode.com/cann/catlass代码组织├── 23_flash_attention_infer │ ├── CMakeLists.txt # CMake构建配置文件 │ ├── gen_data.py │ ├── kernel_common.hpp │ ├── main.cpp │ ├── fai_kernel.cpp │ ├── fai_tiling.cpp │ └── README.md使用示例获取代码之后编译相应的算子可执行文件可参考quickstart接下来先执行gen_data.py生成测试样例测试用例需要从命令行输入, 执行该命令后会在当前路径下生成data目录包含算子的输入数据和用于精度验证的golden数据。然后执行算子这里要注意的是执行算子的输入shape和上面第一步生成数据的shape一致。以下是一个完整的shell脚本示例batch1 qSeqlen177 kvSeqlen512 numHeads1 kvHeads1 headSize128 isVariedLen0 maskType1 dtypebf16 cacheMode1 layout_dtype0 num_blocks2048 inner_prec0 lse_flag0 device0 function build() { rm -rf build rm -rf output bash scripts/build.sh 23_flash_attention_infer } function gen_data() { python3 examples/23_flash_attention_infer/gen_data.py $batch $qSeqlen $kvSeqlen $numHeads $kvHeads $headSize $isVariedLen $maskType $dtype $cacheMode $layout_dtype $num_blocks $inner_prec $lse_flag echo Data gen finished } function run_kernel() { echo Case: B $batch qS $qSeqlen kvS $kvSeqlen qN $numHeads kvN $kvHeads D $headSize mask $maskType cd output/bin/ ./23_flash_attention_infer $batch $qSeqlen $kvSeqlen $numHeads $kvHeads $headSize $isVariedLen $maskType --device $device --dtype $dtype } build gen_data run_kernel执行结果如下说明精度比对成功。Compare success.【免费下载链接】catlass本项目是CANN的算子模板库提供NPU上高性能矩阵乘及其相关融合类算子模板样例。项目地址: https://gitcode.com/cann/catlass创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考