DataWhale组队学习笔记--llm-algo-leetcode(三)

一起努力啊～2026-07-18 11:30

文章目录

  • [Attention MHA GQA(多头注意力)](#Attention MHA GQA(多头注意力))
    • [一、 为什么传统的 MHA 成了推理瓶颈？](#一、 为什么传统的 MHA 成了推理瓶颈？)
      • [显存暴涨的元凶：KV Cache](#显存暴涨的元凶：KV Cache)
    • [二、 三大注意力机制对比：MHA vs MQA vs GQA](#二、 三大注意力机制对比：MHA vs MQA vs GQA)
      • [1. MHA (Multi-Head Attention) --- 传统多头](#1. MHA (Multi-Head Attention) — 传统多头)
      • [2. MQA (Multi-Query Attention) --- 极端单头](#2. MQA (Multi-Query Attention) — 极端单头)
      • [3. GQA (Grouped-Query Attention) --- 黄金折中方案](#3. GQA (Grouped-Query Attention) — 黄金折中方案)
    • [三、 数学维度与计算对比](#三、 数学维度与计算对比)
  • [FlashAttention 模拟](#FlashAttention 模拟)
    • [一、 核心痛点：GPU 的"计算快，访存慢"](#一、 核心痛点：GPU 的“计算快，访存慢”)
      • [1. GPU 内存层级](#1. GPU 内存层级)
      • [2. 传统 Attention 的"读写噩梦"](#2. 传统 Attention 的“读写噩梦”)
    • [二、 FlashAttention 的两大核心魔法](#二、 FlashAttention 的两大核心魔法)
      • [1. 平铺（Tiling）：分块计算与 Softmax 局部更新](#1. 平铺（Tiling）：分块计算与 Softmax 局部更新)
        • [💡 解决方案：在线 Softmax（Online Softmax）](#💡 解决方案：在线 Softmax（Online Softmax）)
      • [2. 重计算（Recomputation / Gradient Checkpointing）](#2. 重计算（Recomputation / Gradient Checkpointing）)
    • [三、 FlashAttention 的演进之路 (V1 vs V2 vs V3)](#三、 FlashAttention 的演进之路 (V1 vs V2 vs V3))
      • [1. FlashAttention-1 (2022年)](#1. FlashAttention-1 (2022年))
      • [2. FlashAttention-2 (2023年)](#2. FlashAttention-2 (2023年))
      • [3. FlashAttention-3 (2024年及以后)](#3. FlashAttention-3 (2024年及以后))
    • [四、 为什么它是大模型时代的基石？](#四、 为什么它是大模型时代的基石？)

Attention MHA GQA(多头注意力)

可以先把 Attention 记成"当前 token 去回看历史 token"，KV Cache 是把历史信息存起来避免重复算，GQA 则是在效果和显存之间做折中。

在大语言模型（LLM）的演进过程中，如何在高并发、长文本场景下降低显存占用并提升推理速度，一直是研究的重中之重。

在 Self-Attention 机制中，传统的 MHA（Multi-Head Attention，多头注意力） 带来了极大的计算和显存负担。为了解决这个问题，研究者们相继提出了 MQA（Multi-Query Attention） 和目前最主流的 GQA（Grouped-Query Attention，分组查询注意力）。

下面我们由浅入深，对这三种注意力机制进行全方位的对比与深度拆解。

一、 为什么传统的 MHA 成了推理瓶颈？

在 Transformer 的自回归解码（Generation）阶段，模型需要逐字生成 Token。为了避免重复计算，我们需要把之前生成过的所有 Token 的 K K K 和 V V V 向量缓存下来，这被称为 KV Cache。

显存暴涨的元凶：KV Cache

在传统的 MHA 中：

  • 每个 Attention Head（注意力头）都有自己独一份的 Q Q Q、 K K K、 V V V 投影矩阵。
  • 假设模型有 H H H 个头，每个头的维度是 d d d，那么每个 Token 都需要存储 H × d H \times d H×d 大小的 K K K 向量和 V V V 向量。
  • 随着序列长度（Context Length）和并发请求数（Batch Size）的增加，KV Cache 占用的显存会呈线性暴涨，甚至远超模型权重本身，导致 GPU 显存轻易 OOM（Out of Memory）。
  • 访存瓶颈（Memory-Bound）：在 GPU 计算时，从全局显存读取这些庞大 KV Cache 的速度（带宽）跟不上计算单元（Tensor Cores）的速度，导致 GPU 算力利用率极低。

二、 三大注意力机制对比：MHA vs MQA vs GQA

为了解决 KV Cache 的痛点，架构演进经历了以下三个阶段：

1. MHA (Multi-Head Attention) --- 传统多头

  • 结构 ： N N N 个 Query 头、 N N N 个 Key 头、 N N N 个 Value 头（一配一）。
  • 特点 ：每个 Q Q Q 头都有自己专属的 K K K 和 V V V 头。
  • 优缺点：
  • 优点：表达能力最强，模型效果最好。
  • 缺点：KV Cache 极其庞大，推理速度慢，显存开销巨大。

2. MQA (Multi-Query Attention) --- 极端单头

  • 结构 ： N N N 个 Query 头，但只有 1 个 Key 头和 1 个 Value 头（多配一）。
  • 特点 ：所有的 Q Q Q 头共享同一组 K K K 和 V V V 向量。
  • 优缺点：
  • 优点：KV Cache 骤降为原来的 1 / N 1/N 1/N，极大地减少了显存占用和访存带宽，推理速度（Throughput）飞跃式提升。
  • 缺点：由于所有头强行共享同一组 KV，模型理解复杂语义的能力有所下降，容易出现性能退化。

3. GQA (Grouped-Query Attention) --- 黄金折中方案

  • 结构 ：将 N N N 个 Query 头均分成 G G G 个组（Group），每个组内的 Query 头共享 1 个 Key 头和 1 个 Value 头。

  • 特点：介于 MHA 和 MQA 之间的折中方案。

  • 当 G = N G = N G=N 时，GQA 退化为 MHA。

  • 当 G = 1 G = 1 G=1 时，GQA 变为 MQA。

  • 优缺点：

  • 优点：完美平衡了速度与效果。它在保持接近 MHA 的模型精度（Accuracy）的同时，获得了接近 MQA 的极快推理速度和超低显存占用。

三、 数学维度与计算对比

为了更直观地理解，我们来看一下这三种机制在维度 上的差异。

假设：

  • 输入特征维度为 hidden_size = 4096
  • 注意力头数 H Q = 32 H_{Q} = 32 HQ=32，每个头的维度 head_dim = 128
机制 Query 投影输出形状 Key/Value 投影输出形状 KV 显存占用对比 适用场景
MHA [B, S, 32, 128] [B, S, 32, 128] 100% (基准) 早期模型 (GPT-3, LLaMA-1)
MQA [B, S, 32, 128] [B, S, 1, 128] 3.125% ( 1 / 32 1/32 1/32) 极长上下文、轻量级模型 (Falcon)
GQA (以 8 组为例) [B, S, 32, 128] [B, S, 8, 128] 25% ( 8 / 32 8/32 8/32) 现代主流大模型 (LLaMA-3, Mistral, Qwen)

代码实现：

python 复制代码 
import torch
import torch.nn as nn
import math
python 复制代码 
def repeat_kv(hidden_states:torch.Tensor,n_rep:int) -> torch.Tensor:
    #将KV头复制n_rep次，以匹配Query头的数量（GQA/MQA需要）
    batch,num_kv_heads,slen,head_dim = hidden_states.shape
    if n_req == 1:
        return hidden_states
    hidden_states = hidden_states[:,:,Node,:,:].expand(batch,num_kv_heads,n_rep,slen,head_dim)
    return hidden_states.reshape(batch,num_kv_heads * n_rep,slen,head_dim)

class GroupedQueryAttention(nn.Module):
    def __init__(self,hidden_dim:int,num_heads:int,num_kv_heads:int = None):
        super().__init__()
        self.hidden_dim = hidden_dim
        self.num_heads = num_heads
        self.num_kv_heads = num_kv_heads if num_kv_heads is not None else num_heads

        self.num_queries_per_kv = self.num_heads // self.num_kv_heads
        self.head_dim = hidden_dim // num_heads

        #定义投影矩阵
        self.q_proj = nn.Linear(hidden_dim,num_heads * self.head_dim,bias = False)
        self.k_proj = nn.Linear(hidden_dim,self.num_kv_heads * self.head_dim,bias = False)
        self.v_proj = nn.Linear(hidden_dim,self.num_kv_heads * self.head_dim,bias = False)
        self.o_proj = nn.Linear(num_heads * self.head_dim,hidden_dim,bias = False)

    def forward(
        self,
        x:torch.Tensor,
        attention_mask:torch.Tensor = None,
        kv_cache:tuple[torch.Tensor,torch.Tensor] = None
    ):

        batch_size,seq_len,_ = x.shape

        #1、线性投影
        xq,xk,xv = self.q_proj(x),self.k_proj(x),self.v_proj(x)

        #TODO1:Reshape xq,xk,xv 以适配多头注意力计算
        xq = xq.view(batch_size,seq_len,self.num_heads,self.head_dim).transpose(1,2)
        xk = xk.view(batch_size,seq_len,self.num_kv_heads,self.head_dim).transpose(1,2)
        xv = xv.view(batch_size,seq_len,self.num_kv_heads,self.head_dim).transpose(1,2)


        #TODO2:处理KV Cache
        if kv_cache is not None:
            k_cache,v_cache = kv_cache
            xk = torch.cat([k_cache,xk],dim = 2)
            xv = torch.cat([v_cache,xv],dim = 2)

        new_kv_cache = (xk,xv)

        #通过repeat_kv 把GQA的KV头数扩充到和Query数量一致
        xk = repeat_kv(xk,self.num_queries_per_kv)
        xv = repeat_kv(xv,self.num_queries_per_kv)

        #TODO3:计算注意力分数

        scores = torch.matmul(xq,xk.transpose(2,3))/math.sqrt(self.head_dim)

        if attention_mask is not Node:
            scores = scores + attention_mask
        
        probs = torch.nn.functional.softmax(scores,dim=-1)
        output = torch.matmul(probs,xv)

        #TODO4:恢复形状并输出

        output = output.transpose(1,2).contiguous().view(batch_size,seq_len,-1)

        return self.o_proj(output),new_kv_cache

结果：

FlashAttention 模拟

在了解了自注意力机制（Self-Attention）和 GQA 之后，你可能会发现，大模型的计算瓶颈除了参数量和 KV Cache 之外，还有一个非常致命的痛点：内存带宽（Memory Bandwidth）。

传统的 Attention 计算在处理长文本时，速度会急剧下降，甚至导致显存溢出（OOM）。为了彻底解决这一问题，斯坦福大学的 Tri Dao 等人提出了 FlashAttention。

FlashAttention 没有改变 Attention 的数学逻辑（其输出与标准 Attention 完全一致） ，它是一项纯粹的硬件加速与访存优化技术。

一、 核心痛点：GPU 的"计算快，访存慢"

要理解 FlashAttention 的伟大之处，首先需要了解 GPU 的存储架构。

1. GPU 内存层级

在 GPU 内部，存储介质主要分为两种：

  • HBM（High Bandwidth Memory，高带宽显存） ：也就是我们平时说的"显存"（如 A100 的 80GB 显存）。它的容量大，但是读写速度相对较慢（类似于电脑的内存）。
  • SRAM（Static Random-Access Memory，静态随机存取内存） ：也叫 Shared Memory（共享内存） 。它集成在 GPU 芯片内部，容量极小（通常每块只有 100KB 左右），但读写速度极其恐怖（类似于 CPU 的 L1/L2 缓存）。

2. 传统 Attention 的"读写噩梦"

在计算标准 Attention 时，公式为：

Attention ( Q , K , V ) = softmax ( Q K T d k ) V \text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V Attention(Q,K,V)=softmax(dk QKT)V

传统的 PyTorch 实现（通常称为 Baseline）是逐个算子（Kernel）调用 HBM 进行读写的：

  1. 从 HBM 读取 Q Q Q 和 K K K，计算 S = Q K T S = QK^T S=QKT，将大小为 N × N N \times N N×N（ N N N 为序列长度）的中间矩阵 S S S 写回 HBM。
  2. 从 HBM 读取 S S S，计算 P = softmax ( S ) P = \text{softmax}(S) P=softmax(S)，再将 P P P 写回 HBM。
  3. 从 HBM 读取 P P P 和 V V V，计算 O = P V O = PV O=PV，将结果 O O O 写回 HBM。

❌ 致命问题 ：

当序列长度 N N N 很大时（例如 100k）， N × N N \times N N×N 的中间矩阵 S S S 和 P P P 会变得超级巨大，频繁在慢速的 HBM 中进行读写（I/O 交互）消耗了 GPU 80% 以上的时间。这种瓶颈被称为访存受限（Memory-Bound）。

二、 FlashAttention 的两大核心魔法

FlashAttention 的核心思路是：绝不往 HBM 中写入任何 N × N N \times N N×N 的中间矩阵！所有计算尽可能在超高速的 SRAM 中完成。

为了实现这一点，它引入了两个革命性的数学与工程设计：平铺（Tiling） 与 重计算（Recomputation）。

1. 平铺（Tiling）：分块计算与 Softmax 局部更新

因为 SRAM 的容量非常小（仅约 100KB），我们无法一次性把整个 Q , K , V Q, K, V Q,K,V 塞进去。FlashAttention 采用平铺（Tiling）技术，将 Q , K , V Q, K, V Q,K,V 拆分成许多个小分块（Blocks），每次只加载一小块到 SRAM 中计算。

但这里有一个巨大的数学障碍：Softmax 是全局关联的。

softmax ( x i ) = e x i ∑ e x j \text{softmax}(x_i) = \frac{e^{x_i}}{\sum e^{x_j}} softmax(xi)=∑exjexi

为了防止数值溢出，实际计算时需要减去全局最大值 m m m：

softmax ( x i ) = e x i − m ∑ e x j − m \text{softmax}(x_i) = \frac{e^{x_i - m}}{\sum e^{x_j - m}} softmax(xi)=∑exj−mexi−m

如果把矩阵分块，我们怎么在不知道全局最大值 m m m 和全局分母和（归一化因子） d d d 的情况下，计算出正确的 Softmax 呢？

💡 解决方案：在线 Softmax（Online Softmax）

FlashAttention 利用了数学上的递推性质，可以动态更新 Softmax 的值。

假设我们已经计算了前一小块的局部最大值 m ( 1 ) m^{(1)} m(1) 和分母和 d ( 1 ) d^{(1)} d(1)，现在迎来了新的一小块，其局部最大值为 m ( 2 ) m^{(2)} m(2)。

我们可以通过以下公式，无缝地将它们融合成新的全局指标：

m new = max ⁡ ( m ( 1 ) , m ( 2 ) ) m^{\text{new}} = \max(m^{(1)}, m^{(2)}) mnew=max(m(1),m(2))

d new = d ( 1 ) ⋅ e m ( 1 ) − m new + d ( 2 ) ⋅ e m ( 2 ) − m new d^{\text{new}} = d^{(1)} \cdot e^{m^{(1)} - m^{\text{new}}} + d^{(2)} \cdot e^{m^{(2)} - m^{\text{new}}} dnew=d(1)⋅em(1)−mnew+d(2)⋅em(2)−mnew

通过这种缩放更新机制，FlashAttention 能够一边分块读取数据，一边动态修正之前的 Attention 输出 ，从而彻底避免了保存全局 N × N N \times N N×N 矩阵的需要。

2. 重计算（Recomputation / Gradient Checkpointing）

在反向传播（Backward Pass）中，通常需要用到正向传播时产生的激活值（即那两个大矩阵 S S S 和 P P P）。如果我们在正向传播中不保存它们，反向传播该怎么办？

FlashAttention 采用了一种"以计算换空间"的激进策略：

  • 在正向传播时，直接扔掉中间的 Attention 矩阵。
  • 在反向传播需要用到它们时，直接利用 SRAM 里保留的 Q , K , V Q, K, V Q,K,V 分块现场重新计算一次。

虽然这增加了一点点计算量，但由于不用读写慢速的 HBM，省下来的 I/O 时间远远超过了重计算的时间，整体速度反而得到了成倍的提升！

三、 FlashAttention 的演进之路 (V1 vs V2 vs V3)

随着显卡架构（如 Ampere, Hopper）的升级，FlashAttention 也经历了多次迭代：

1. FlashAttention-1 (2022年)

  • 奠定了 Tiling 和 Recomputation 的框架。
  • 相比于标准 Attention，实现了 2~4倍 的速度提升，显存占用从 O ( N 2 ) O(N^2) O(N2) 骤降到 O ( N ) O(N) O(N)。

2. FlashAttention-2 (2023年)

  • 减少了非矩阵乘法（Non-matmul）计算：将不适合 Tensor Core 的运算（如 Softmax 缩放）尽可能减少，让 90% 以上的计算都跑在极快的 Tensor Core 上。
  • 更高效的多头并行 ：在 Thread Block（线程块）维度上进行了更好的切分与调度，GPU 算力利用率（TFLOPs/s）从 V1 的 30~40% 提升到了 50~75%，速度相比 V1 再次翻倍。

3. FlashAttention-3 (2024年及以后)

  • 针对 NVIDIA Hopper 架构（H100 / H200 / B200）进行了底层硬件级榨干。
  • 异步执行（Asynchronous WGMMA） ：利用 H100 的 Tensor Memory Accelerator (TMA)，让"数据从 HBM 搬运到 SRAM"和"Tensor Core 进行矩阵乘法计算"同时异步进行，消除了等待延迟。
  • 支持 FP8（8位浮点数）低精度训练与推理，速度相比 V2 再次提升 1.5~2倍。

四、 为什么它是大模型时代的基石？

特性 标准 Attention FlashAttention
显存占用（存储复杂度） O ( N 2 ) O(N^2) O(N2) （随长度呈平方级暴涨） O ( N ) O(N) O(N) （随长度呈线性增长）
访存特征 频繁读写 HBM（Memory-Bound） 内部高效利用 SRAM（Compute-Bound）
硬件利用率 极低（GPU 经常在"等数据"） 极高（Tensor Core 全速运转）
支持的上下文长度 极难突破 8K 轻松支持 32K / 128K / 1M 级别

一句话总结：

FlashAttention 通过数学公式的巧妙重构（在线 Softmax 递推），把原本需要写在显存里的超大临时表，变成了在 GPU 缓存里闪电般完成的局部计算。它直接拉开了大模型"长文本时代（Long Context）"的序幕！

代码实现：

python 复制代码 
import torch
import math
python 复制代码 
def flash_attention_forward_sim(q,k,v,block_size = 2):
    #纯PyTorch模拟FlashAttention前向传播，假设没有Batch和Head维度，q,k,v的形状都是(seq_len,dim)
    seq_len,dim = q.shape

    #TODO1：初始化输出O，全局最大值m,全局指数和1
    out = torch.zeros((seq_len,dim),device=q.device)
    m = torch.full((seq_len,1),-float('inf'),device = q.device)
    l = torch.zeros((seq_len,1),device=q.device)

    scale = 1.0 / math.sqrt(dim)

    #外层循环：遍历Q的分块
    for i in range(0,seq_len,block_size):
        q_block = q[i:i+block_size]*scale
        m_i = m[i:i+block_size]
        l_i = l[i:i+block_size]
        out_i = out[i:i+block_size]

        #内层循坏：遍历K,V的分块
        for j in range(0,seq_len,block_size):
            k_block = k[j:j+block_size]
            v_block = v[j:j+block_size]

            #TODO2:计算当前块的未归一化分数S_ij
            S_ij = q_block @ k_block.transpose(-2,-1)

            #TODO3:计算当前块的局部最大值m_block,并求出新的全局最大值m_new
            m_block = torch.max(S_ij,dim = -1,keepdim=True)[0]
            m_new = torch.maximum(m_i,m_block)

            #TODO4:计算P_ij = exp(S_ij - m_new)
            P_ij = torch.exp(S_ij - m_new)

            #TODO5:计算当前块的局部指数和l_block,并更新全局指数和l_new
            l_block = torch.sum(P_ij,dim = -1,keepdim=True)
            l_new = l_i * torch.exp(m_i - m_new) + l_block

            #TODO6:更新输出O_i（使用Online Softmax的修正公式）
            out_i = out_i * (l_i * torch.exp(m_i - m_new) / l_new) + (P_ij @ v_block) / l_new

            #更新全局状态
            m_i = m_new
            l_i = l_new
            pass
        #写回全局变量
        out[i:i+block_size] = out_i
        m[i:i+block_size] = m_i
        l[i:i+block_size] = l_i

    return out
python 复制代码 
# 测试你的实现
def test_flash_attention_sim():
    try:
        import math

        def run_case(seq_len, dim, block_size, seed):
            torch.manual_seed(seed)
            q = torch.randn(seq_len, dim)
            k = torch.randn(seq_len, dim)
            v = torch.randn(seq_len, dim)

            scale = 1.0 / math.sqrt(dim)
            scores = (q @ k.transpose(-2, -1)) * scale
            attn = torch.nn.functional.softmax(scores, dim=-1)
            out_ref = attn @ v

            out_sim = flash_attention_forward_sim(q, k, v, block_size=block_size)
            diff = torch.max(torch.abs(out_ref - out_sim))
            print(f"[seq={seq_len}, dim={dim}, block={block_size}] 最大误差: {diff.item():.6e}")
            assert diff < 1e-5, f"计算结果与标准 Attention 不一致！(seq={seq_len}, dim={dim}, block={block_size})"

        run_case(seq_len=8, dim=4, block_size=2, seed=42)
        run_case(seq_len=5, dim=3, block_size=3, seed=7)
        run_case(seq_len=3, dim=2, block_size=1, seed=123)

        print("✅ Online Softmax 与分块计算逻辑正确！")
        print("\n FlashAttention 分块计算逻辑验证通过。")

    except NotImplementedError:
        print("请先完成 TODO 部分的代码！")
        raise
    except (AttributeError, NameError, TypeError, ValueError, AssertionError, RuntimeError) as e:
        if isinstance(e, AttributeError):
            print("代码未完成，无法找到必要的属性")
        elif isinstance(e, NameError):
            print("代码可能未完成，导致了变量未定义")
        elif isinstance(e, TypeError):
            print("代码可能未完成，导致了操作错误")
        elif isinstance(e, ValueError):
            print("代码可能未完成，导致了张量维度错误")
        elif isinstance(e, RuntimeError):
            print("代码可能未完成，导致了运行时错误")
        else:
            print("代码可能未完成，导致了断言失败")
        raise NotImplementedError("请先完成 TODO 部分的代码！") from e
    except Exception as e:
        print(f"❌ 测试失败: {e}")
        raise


test_flash_attention_sim()

结果：

相关推荐
云空1 小时前
《Three.js 可操控3D魔方（键盘转动+打乱+复原）》
javascript·学习·游戏·3d·three.js
kdxiaojie2 小时前
Linux 驱动研究 —— V4L2 （3）
linux·运维·笔记·学习
AINative软件工程2 小时前
LLM 并发工具调用的 5 个生产陷阱：幂等性、竞态与失败补偿
llm·agent·ai编程
Amazing_Cacao2 小时前
CFCA精品可可饮品中级参数推演：突破液态表达的混沌，实现萃取微观结构的绝对物理可控
笔记
十月的皮皮2 小时前
C语言学习笔记20260717-预处理机制
c语言·笔记·学习
吃饱了得干活2 小时前
LangChain 提示词模板：从拼接字符串到组合流水线
langchain·llm
Zzj_tju3 小时前
Alignment Science 与可控生成：偏好学习、宪法式训练和可验证对齐
人工智能·学习·语言模型
eddietao3 小时前
MySQL+SQLServer数据库-笔记
数据库·笔记·mysql
热门推荐
012026年7月AI圈大地震：GPT-5.6被政府限制、Claude入驻Slack、Anthropic自研芯片02GitHub 镜像站点03AI科技热点日报 | 2026年07月01日04幻兽帕鲁 - 服务器管理员权限与 GM 命令完全指南05AI 编程 IDE 全景解析 2026：Agent 全面接管开发链路06国内可直接用、免费额度/永久免费的大模型API清单（含 SiliconFlow、火山、阿里、智谱、百度、Kimi、DeepSeek、DMXAPI 等）072026年AI技术突破与产业落地全景：从GPT-5到多模态智能体的新纪元082026 国产 AI 大模型横评：DeepSeek、通义千问、Kimi、文心一言、星火、豆包谁更能打？09Codex 下载安装指南：Windows 和 macOS 官方版下载102026 年 AI 大模型 & AI 编程工具实战全总结