tensor core实现flash_attn_mma_share_kv源码分析

youzjuer

已于 2025-02-07 10:34:13 修改

阅读量249

点赞数

分类专栏： CUDA和LLM通俗易懂技术站文章标签： cuda 人工智能

于 2025-01-08 10:10:43 首次发布

本文链接：https://blog.youkuaiyun.com/youzjuer/article/details/144985817

版权

CUDA和LLM通俗易懂技术站专栏收录该内容

105 篇文章 ¥59.90 ¥99.00

订阅专栏

超级会员免费看

一源码分析

1.1 函数入口

void flash_attn_mma_stages_split_q_shared_kv(torch::Tensor Q, 
                                             torch::Tensor K, 
                                             torch::Tensor V, 
                                             torch::Tensor O, 
                                             int stages) {
  CHECK_TORCH_TENSOR_DTYPE(Q, torch::kHalf) // Q [B,H,N,D]
  CHECK_TORCH_TENSOR_DTYPE(K, torch::kHalf) // K [B,H,N,D]
  CHECK_TORCH_TENSOR_DTYPE(V, torch::kHalf) // V [B,H,N,D]
  CHECK_TORCH_TENSOR_DTYPE(O, torch::kHalf) // O [B,H,N,D]
  const int d = Q.size(3); // B, H, N, d

  if (stages > 1) {
    switch (d)
    {
    case 32:
      launch_flash_attn_mma_stages_split_q_shared_kv<32,  2>(Q, K, V, O);