梳理下Flash Attention的dispatch逻辑

文章预览

↑ 点击 蓝字  关注极市平台 作者丨GiantPandaCV 来源丨GiantPandaCV 编辑丨极市平台 极市导读   本文分析了Flash Attention在不同场景下的内核调度逻辑，特别关注了在解码阶段何时会使用split_kv实现，并探讨了影响这一决策的因素，如K序列的最大长度、注意力头数和GPU的流处理器数量。   >> 加入极市CV技术交流群，走在计算机视觉的最前沿 0x1. 前言 这篇文章来源是当运行下面的对HuggingFace Qwen2.5-7B-Instruct模型使用Flash Attention的代码时，使用Nsight System工具抓取的kernel trace会发现在prefill和decode阶段，Flash Attention调用了不同的kernel并且decoding的Flash Attention kernel使用了split_kv的实现。然后如果把下面代码中max_new_tokens改成64，我发现在Nsight System工具抓取的kernel trace中，decode阶段的Flash Attention kernel又变成了和prefill阶段一样的kernel，并没有使用split_kv的实现。这篇文 ………………………………