【AIGC】大模型面试高频考点-注意力（Attention）篇

最新推荐文章于 2025-05-24 16:52:04 发布

LeeZhao@

最新推荐文章于 2025-05-24 16:52:04 发布

阅读量2.2k

点赞数 19

分类专栏： AIGC高频面试题文章标签： AIGC 面试职场和发展

本文链接：https://blog.youkuaiyun.com/qq_36722887/article/details/144272869

版权

AIGC高频面试题专栏收录该内容

16 篇文章

订阅专栏

【AIGC】大模型面试高频考点-注意力（Attention）篇

（一）手撕单头注意力机制（ScaledDotProductAttention）函数

输入是query和 key-value，注意力机制首先计算query与每个key的关联性（compatibility），每个关联性作为每个value的权重（weight），各个权重与value的乘积相加得到输出。

在这里插入图片描述

class ScaledDotProductAttention(nn.Module):
    """ Scaled Dot-Product Attention """

    def __init__(self, scale):
        super().__init__()

        self.scale = scale
        self.softmax = nn.Softmax(dim=2)

    def forward(self, q, k, v, mask=None):
        u = torch.bmm(q, k.transpose(1, 2)) # 1.Matmul
        u = u / self.scale # 2.Scale

        if mask is not None:
            u = u.masked_fill(mask, -np.inf) # 3.Mask

        attn = self.softmax(u) # 4.Softmax
        output = torch.bmm(attn, v) # 5.Output

        return attn, output


if __name__ == "__main__":
    n_q, n_k, n_v = 2, 4, 4
    d_q, d_k, d_v = 128, 128, 64

    q = torch.randn(batch, n_q, d_q)
    k = torch.randn(batch, n_k, d_k)
    v = torch.randn(batch, n_v, d_v)
    mask = torch.zeros(batch, n_q, n_k).bool()

    attention = ScaledDotProductAttention(scale=np.power(d_k, 0.5))
    attn, output = attention(q, k, v, mask=mask)

    print(attn)
    print(output)

（二）手撕多头注意力（MultiHeadAttention）

class MultiHeadAttention(nn.Module):
    """ Multi-Head Attention """

    def __init__(self, n_head, d_k_, d_v_, d_k, d_v, d_o):
        super().__init__()

        self.n_head = n_head
        self.d_k = d_k
        self.d_v = d_v

        self.fc_q = nn.Linear(d_k_, n_head * d_k)
        self.fc_k = nn.Linear(d_k_, n_head * d_k)
        self.fc_v = nn.Linear(d_v_, n_head * d_v)

        self.attention = ScaledDotProductAttention(scale=np.power(d_k, 0.5))

        self.fc_o = nn.Linear(n_head * d_v, d_o)

    def forward(self, q, k, v, mask=None):

        n_head, d_q, d_k, d_v = self.n_head, self.d_k, self.d_k, self.d_v

        batch, n_q, d_q_ = q.size()
        batch, n_k, d_k_ = k.size()
        batch, n_v, d_v_ = v.size()

        q = self.fc_q(q) # 1.单头变多头
        k = self.fc_k(k)
        v = self.fc_v(v)
        q = q.view(batch, n_q, n_head, d_q).permute(2, 0, 1, 3).contiguous().view(-1, n_q, d_q)
        k = k.view(batch, n_k, n_head, d_k).permute(2, 0, 1, 3).contiguous().view(-1, n_k, d_k)
        v = v.view(batch, n_v, n_head, d_v).permute(2, 0, 1, 3).contiguous().view(-1, n_v, d_v)

        if mask is not None:
            mask = mask.repeat(n_head, 1, 1)
        attn, output = self.attention(q, k, v, mask=mask) # 2.当成单头注意力求输出

        output = output.view(n_head, batch, n_q, d_v).permute(1, 2, 0, 3).contiguous().view(batch, n_q, -1) # 3.Concat
        output = self.fc_o(output) # 4.仿射变换得到最终输出

        return attn, output


if __name__ == "__main__":
    n_q, n_k, n_v = 2, 4, 4
    d_q_, d_k_, d_v_ = 128, 128, 64

    q = torch.randn(batch, n_q, d_q_)
    k = torch.randn(batch, n_k, d_k_)
    v = torch.randn(batch, n_v, d_v_)    
    mask = torch.zeros(batch, n_q, n_k).bool()

    mha = MultiHeadAttention(n_head=8, d_k_=128, d_v_=64, d_k=256, d_v=128, d_o=128)
    attn, output = mha(q, k, v, mask=mask)

    print(attn.size())
    print(output.size())

（三）手撕自注意力机制函数（SelfAttention）

Self-Attention。和Attention类似，他们都是一种注意力机制。不同的是Attention是source对target，输入的source和输出的target内容不同。例如英译中，输入英文，输出中文。而Self-Attention是source对source，是source内部元素之间或者target内部元素之间发生的Attention机制，也可以理解为Target=Source这种特殊情况下的注意力机制。

class SelfAttention(nn.Module):
    """ Self-Attention """

    def __init__(self, n_head, d_k, d_v, d_x, d_o):
        self.wq = nn.Parameter(torch.Tensor(d_x, d_k))
        self.wk = nn.Parameter(torch.Tensor(d_x, d_k))
        self.wv = nn.Parameter(torch.Tensor(d_x, d_v))

        self.mha = MultiHeadAttention(n_head=n_head, d_k_=d_k, d_v_=d_v, d_k=d_k, d_v=d_v, d_o=d_o)

        self.init_parameters()

    def init_parameters(self):
        for param in self.parameters():
            stdv = 1. / np.power(param.size(-1), 0.5)
            param.data.uniform_(-stdv, stdv)

    def forward(self, x, mask=None):
        q = torch.matmul(x, self.wq)   
        k = torch.matmul(x, self.wk)
        v = torch.matmul(x, self.wv)

        attn, output = self.mha(q, k, v, mask=mask)

        return attn, output


if __name__ == "__main__":
    n_x = 4
    d_x = 80

    x = torch.randn(batch, n_x, d_x)
    mask = torch.zeros(batch, n_x, n_x).bool()

    selfattn = SelfAttention(n_head=8, d_k=128, d_v=64, d_x=80, d_o=80)
    attn, output = selfattn(x, mask=mask)

    print(attn.size())
    print(output.size())

（四）GPT2 解码中的KV Cache

无论是Encoder-Decoder结构，还是现在我们最接近AGI的decoder-only的LLM，解码生成时都是自回归auto-regressive的方式。

也就是，解码的时候，先根据当前输入input ，生成下一个 token，然后把新生成的token拼接在input后面，获得新的输入input，再用input生成token，依此迭代，直到生成结束。

我们可以注意到，下一个step的输入其实包含了上一个step的内容，而且只在最后面多了一点点（一个token）。那么下一个step的计算应该也包含了上一个step的计算。

但是模型在推理的时候可不管这些，无论你是不是只要最后一个字的输出，它都把所有输入计算一遍，给出所有输出结果。

也就是说中间有很多我们用不到的计算，这样就造成了浪费。

而且随着生成的结果越来越多，输入的长度也越来越长，上面这个例子里，输入长度就从step0的10个，每步增长1，直到step5的15个。如果输入的instruction是让模型写作文，那可能就有800个step。这个情况下，step0被算了800次，step1被算了799次…这样浪费的计算资源确实不容忽视。

有没有什么办法可以重复利用上一个step里已经计算过的结果，减少浪费呢？

答案就是KV Cache，利用一个缓存，把需要重复利用的中间计算结果存下来，减少重复计算。

而 k 和 v 就是我要缓存的对象。

想象一下，在上面的例子中，假设我们一开始的输入就是3个字，我们第一次预测就是预测第4个字，那么由于一开始没有任何缓存，所有我们每一层还是要老实地计算一遍。然后把 k 、 v 值缓存起来。

则有

在这里插入图片描述

kv cache的下标l表示模型层数。

在进行第二次预测，也就是预测第5个字的时候，在第l层的时候，由于前面我们缓存了每层的ku 值，那本层就只需要算新的 o3，而不用算 o0、o1、o2。

因为第l层的 o0、o1、o2本来会经过FNN层之后进到 l十1 层，再经过新的投影变换，成为 l + 1 层的 k、v值，但是l十 1 层的 k、v 值我们已经缓存过了!

然后我们把本次新增算出来的 k、υ 值也存入缓存。

在这里插入图片描述

这样就节省了attention和FFN的很多重复计算。

transformers中，生成的时候传入use_cache=True就会开启KV Cache。

也可以简单看下GPT2中的实现，中文注释的部分就是使用缓存结果和更新缓存结果

Class GPT2Attention(nn.Module):
    ...
    ...
    def forward(
        self,
        hidden_states: Optional[Tuple[torch.FloatTensor]],
        layer_past: Optional[Tuple[torch.Tensor]] = None,
        attention_mask: Optional[torch.FloatTensor] = None,
        head_mask: Optional[torch.FloatTensor] = None,
        encoder_hidden_states: Optional[torch.Tensor] = None,
        encoder_attention_mask: Optional[torch.FloatTensor] = None,
        use_cache: Optional[bool] = False,
        output_attentions: Optional[bool] = False,
    ) -> Tuple[Union[torch.Tensor, Tuple[torch.Tensor]], ...]:
        if encoder_hidden_states is not None:
            if not hasattr(self, "q_attn"):
                raise ValueError(
                    "If class is used as cross attention, the weights `q_attn` have to be defined. "
                    "Please make sure to instantiate class with `GPT2Attention(..., is_cross_attention=True)`."
                )

            query = self.q_attn(hidden_states)
            key, value = self.c_attn(encoder_hidden_states).split(self.split_size, dim=2)
            attention_mask = encoder_attention_mask
        else:
            query, key, value = self.c_attn(hidden_states).split(self.split_size, dim=2)

        query = self._split_heads(query, self.num_heads, self.head_dim)
        key = self._split_heads(key, self.num_heads, self.head_dim)
        value = self._split_heads(value, self.num_heads, self.head_dim)

        # 过去所存的值
        if layer_past is not None:
            past_key, past_value = layer_past
            key = torch.cat((past_key, key), dim=-2)  # 把当前新的key加入
            value = torch.cat((past_value, value), dim=-2)  # 把当前新的value加入

        if use_cache is True:
            present = (key, value)  # 输出用于保存
        else:
            present = None

        if self.reorder_and_upcast_attn:
            attn_output, attn_weights = self._upcast_and_reordered_attn(query, key, value, attention_mask, head_mask)
        else:
            attn_output, attn_weights = self._attn(query, key, value, attention_mask, head_mask)

        attn_output = self._merge_heads(attn_output, self.num_heads, self.head_dim)
        attn_output = self.c_proj(attn_output)
        attn_output = self.resid_dropout(attn_output)

        outputs = (attn_output, present)
        if output_attentions:
            outputs += (attn_weights,)

        return outputs  # a, present, (attentions)

总的来说，KV Cache是以空间换时间的做法，通过使用快速的缓存存取，减少了重复计算。（注意，只有decoder结构的模型可用，因为有mask attention的存在，使得前面的token可以不用关注后面的token）

（五）手撕 MQA 算法

MQA 让所有的头之间共享同一份 Key 和 Value 矩阵，每个头只单独保留了一份 Query 参数，从而大大减少 Key 和 Value 矩阵的参数量。

class MultiQueryAttention(nn.Module):
    """Multi-Query self attention.

    Using torch or triton attention implemetation enables user to also use
    additive bias.
    """

    def __init__(
        self,
        d_model: int,
        n_heads: int,
        attn_impl: str = 'triton',
        clip_qkv: Optional[float] = None,
        qk_ln: bool = False,
        softmax_scale: Optional[float] = None,
        attn_pdrop: float = 0.0,
        low_precision_layernorm: bool = False,
        verbose: int = 0,
        device: Optional[str] = None,
    ):
        super().__init__()

        self.attn_impl = attn_impl
        self.clip_qkv = clip_qkv
        self.qk_ln = qk_ln

        self.d_model = d_model
        self.n_heads = n_heads
        self.head_dim = d_model // n_heads
        self.softmax_scale = softmax_scale
        if self.softmax_scale is None:
            self.softmax_scale = 1 / math.sqrt(self.head_dim)
        self.attn_dropout_p = attn_pdrop

        self.Wqkv = nn.Linear(
            d_model,
            d_model + 2 * self.head_dim,
            device=device,
        )

        fuse_splits = (d_model, d_model + self.head_dim)
        self.Wqkv._fused = (0, fuse_splits)  # type: ignore

        self.attn_fn = scaled_multihead_dot_product_attention
        self.out_proj = nn.Linear(self.d_model, self.d_model, device=device)
        self.out_proj._is_residual = True  # type: ignore

    def forward(
        self,
        x,
        past_key_value=None,
        attn_bias=None,
        attention_mask=None,
        is_causal=True,
        needs_weights=False,
    ):
        qkv = self.Wqkv(x)                                # (1, 512, 960)

        if self.clip_qkv:
            qkv.clamp_(min=-self.clip_qkv, max=self.clip_qkv)

        query, key, value = qkv.split(                         # query -> (1, 512, 768)
            [self.d_model, self.head_dim, self.head_dim],      # key   -> (1, 512, 96)
            dim=2                                              # value -> (1, 512, 96)
        )

        key_padding_mask = attention_mask

        if self.qk_ln:
            # Applying layernorm to qk
            dtype = query.dtype
            query = self.q_ln(query).to(dtype)
            key = self.k_ln(key).to(dtype)

        context, attn_weights, past_key_value = self.attn_fn(
            query,
            key,
            value,
            self.n_heads,
            past_key_value=past_key_value,
            softmax_scale=self.softmax_scale,
            attn_bias=attn_bias,
            key_padding_mask=key_padding_mask,
            is_causal=is_causal,
            dropout_p=self.attn_dropout_p,
            training=self.training,
            needs_weights=needs_weights,
            multiquery=True,
        )

        return self.out_proj(context), attn_weights, past_key_value

（六）Attention改进版

（1）Flash Attention

Paper:《FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness》

论文链接：https://arxiv.org/abs/2205.1413

FlashAttention是一种加速注意力计算方法，目前已经应用在：GPT-3、Falcon2（阿联酋大模型）、Llama2、Megatron-LM、GPT-4等知名LLM上。
Flash Attention已经集成到了pytorch2.0中，可以很便捷的调用。
FlashAttention旨在加速注意力计算并减少内存占用。FlashAttention利用底层硬件的内存层次知识，例如GPU的内存层次结构，来提高计算速度和减少内存访问开销。 FlashAttention的核心原理是通过将输入分块并在每个块上执行注意力操作，从而减少对高带宽内存（HBM）的读写操作。具体而言，FlashAttention使用平铺和重计算等经典技术，将输入块从HBM加载到SRAM（快速缓存），在SRAM上执行注意力操作，并将结果更新回HBM。FlashAttention减少了内存读写量，从而实现了2-4倍的时钟时间加速。
Timeline: 最新的FlashAttention-2版本进一步优化了FlashAttention算法，使用了更好的并行化和工作分区方法，使得计算速度提高了2倍。FlashAttention-2还支持更高的头维数和多查询注意力等新特性，进一步提升了性能和灵活性。

在这里插入图片描述

详细原理细节请查看：
Flash Attention原理详解(含代码讲解)

（2）Page Attention

论文地址：vLLM: Easy, Fast, and Cheap LLM Serving with PagedAttention | vLLM Blog

原理如下：

步骤1：确定固定块大小：（假设512 token) 等进行内存管理。

步骤2：动态分配内存：对每一个batch 首先各分配一个固定块大小，将该固定块对应的物理地址和当前占用token 个数，在一张block table 表内记录。

步骤3：对于不断增长的token数，每超过512就在runtime阶段重新分配物理内存，并同样将该块的信息记录到blocktable 表中。

步骤4：算子库加载实现，获得每个batch 的逻辑地址即该起始block快的地址，得到block 块，加载kvcache，根据当前总的token个数，分段加载token快进行计算。

步骤5：算子库存储实现，获得每个batch 的逻辑地址即该起始block快的地址，根据当前token索引，获得要存储在哪个逻辑block块，找到对应块的fill个数，偏移后存储到对应位置。

在这里插入图片描述

详细原理细节请查看：PageAttention 论文解析

（3）Flash Attention2

论文地址：https://arxiv.org/pdf/2307.08691

如何扩展Transformer使之能够处理更长的序列一直是一个挑战，**因为其核心注意力层的运行时间和内存占用量随输入序列长度成二次增加。**我们希望能够打破2k序列长度限制，从而能够训练书籍、高分辨率图像和长视频。此外，写作等应用也需要模型能够处理长序列。过去一年中，业界推出了一些远超之前长度的语言模型：GPT-4为32k，MosaicML的MPT为65k，以及Anthropic的Claude为100k。

虽然相比标准Attention，FlashAttention快了2_{4倍，节约了10}20倍内存，但是离设备理论最大throughput和flops还差了很多。本文提出了FlashAttention-2，它具有更好的并行性和工作分区。实验结果显示，FlashAttention-2在正向传递中实现了约2倍的速度提升，达到了理论最大吞吐量的73%，在反向传递中达到了理论最大吞吐量的63%。在每个A100 GPU上的训练速度可达到225 TFLOPs/s。

本文主要贡献和创新点为：

减少了non-matmul FLOPs的数量（消除了原先频繁rescale）。虽然non-matmul FLOPs仅占总FLOPs的一小部分，但它们的执行时间较长，这是因为GPU有专用的矩阵乘法计算单元，其吞吐量高达非矩阵乘法吞吐量的16倍。因此，减少non-matmul FLOPs并尽可能多地执行matmul FLOPs非常重要。
提出了在序列长度维度上并行化。该方法在输入序列很长（此时batch size通常很小）的情况下增加了GPU利用率。即使对于单个head，也在不同的thread block之间进行并行计算。
在一个attention计算块内，将工作分配在一个thread block的不同warp上，以减少通信和共享内存读/写。

详细原理细节请查看：FlashAttention2详解（性能比FlashAttention提升200%）

（4）Flash Attention3

Github地址：https://github.com/Dao-AILab/flash-attention

论文地址：https://tridao.me/publications/flash3/flash3.pdf

FlashAttention、FlashAttention-2开创了一种通过最小化内存读/写来加快 GPU 注意力的方法，现在已经成为了pytorch库的标配了，使用它来加速 Transformer 训练和推理。

使得LLM上下文长度大幅增加，从 2-4K （GPT-3， OPT）到 128K （GPT-4），甚至 1M （Llama 3）

FlashAttention-2 可以在 A100 GPU 上实现高达 70% 的理论最大 FLOPS，但它尚未利用 Hopper GPU 上的新功能来最大限度地提高性能。
FlashAttention-2 在 H100 GPU 上仅实现了 35% 的理论最大 FLOP 利用率。
FlashAttention-3在H100 理论最大 FLOPS 的利用率为 75%，比采用 FP16 的 FlashAttention-2 快 1.5-2.0 倍，最高可达 740 TFLOPS。使用 FP8 时，FlashAttention-3 可达到接近 1.2 PFLOPS，误差比基线 FP8 注意力小 2.6 倍。

详细改进如下：

更高效的 GPU 利用率
- H100 GPU 推出了WGMMA（翘曲矩阵乘法累加）功能，比A100吞吐量高3倍
- H100 GPU 的TMA（张量记忆加速器）功能，可加速全局内存和共享内存之间的数据传输，负责所有索引计算和越界预测。这样可以释放寄存器，增加图块大小和效率的宝贵资源。这导致大型语言模型（LLM）的训练和运行速度比FlashAttention-2快得多（1.5-2 倍）。
以更低的精度获得更好的性能
- FlashAttention-3 可以在保持精度的同时处理较低精度的数字（FP8）。例如，FP16 为 989 TFLOPS，FP8 为 1978 TFLOPS。这允许更快的处理速度并尽可能降低内存使用量，这可能会为运行大规模 AI 操作的客户节省成本并提高效率。
- 具体的做法是：利用QuIP: 2-Bit Quantization of Large Language Models With Guarantees技术，通过非相干处理减少量化误差，即将查询和键与随机正交矩阵相乘，以“分散”异常值并减少量化误差。
能够在 LLM 中使用更长的上下文
- 通过加速注意力机制，FlashAttention-3 使 AI 模型能够更有效地处理更长的文本片段。这可以使应用程序能够理解和生成更长、更复杂的内容，而不会减慢速度。
- 对于 FP16，我们看到比 FlashAttention-2 加速约 1.6-2.0 倍。
- 对于FP8，我们可以达到接近1.2 PFLOPS