model_optimizer支持用cuteDSL实现自定义fmha算子了
背景
期望在model_optimizer中使用cuteDSL来优化pi0.5 fmha性能,但是pi05 llm fmha的dim=256,tensorrt-edge-llm还不支持,只好自己实现了
下面是详细的实现步骤
1. AOT 编译(需 SM100/SM110 GPU + cutlass-dsl 4.4.1)
pip install -r kernelSrc/requirements-cutedsl.txt
python kernelSrc/build_cutedsl.py--kernelsfmha--gpu_archsm_110-j4--output_dircpp/kernels/cuteDSLArtifact