triton结合flash-attn的方案

[zhuchen1109]

想请问下，lmdeploy中triton如果支持flash-attn实现这个是否对推理有提升？我在triton的案例中有看到支持flash-attn。

[lvhan028]

@grimoire

[grimoire]

准确率影响应该不大，只是换了实现方式
性能提升会比较大，主要时针对长文本。triton 的 stage 相关的 pass 对 paged attention 的优化不算特别好（相比 turbomind），换成 flash-attn 可以有可观的提升。

[zhuchen1109]

好的，我想尝试修改这块，这个有什么建议吗？是不是应该主要修改atten kernel部分的代码实现？谢谢

[grimoire]

#2676 目前有一个 PR，仅替换了 prefill 的 flash attention。decoding 通常是 Memory bound 的，所以还是使用 paged attention 实现。有兴趣的话可以协助 review、优化以及测试一下