手撕经典算法 #1 Attention 篇整理

推荐的最小实现

import math
import torch


def scaled_dot_product_attention(query, key, value, attn_mask=None):
    """
    query: (batch, heads, q_len, head_dim)
    key:   (batch, heads, k_len, head_dim)
    value: (batch, heads, k_len, head_dim)
    attn_mask: broadcastable to (batch, heads, q_len, k_len);
               True means the position is masked.
    """
    scale = 1.0 / math.sqrt(query.size(-1))
    scores = torch.matmul(query, key.transpose(-1, -2)) * scale

    if attn_mask is not None:
        scores = scores.masked_fill(attn_mask, torch.finfo(scores.dtype).min)

    weights = torch.softmax(scores, dim=-1)
    return torch.matmul(weights, value), weights

Mask 的关键边界

更稳健的 MHA 写法

from torch import nn


class MultiHeadAttention(nn.Module):
    def __init__(self, hidden_size, num_heads, dropout=0.0):
        super().__init__()
        if hidden_size % num_heads != 0:
            raise ValueError("hidden_size must be divisible by num_heads")

        self.hidden_size = hidden_size
        self.num_heads = num_heads
        self.head_dim = hidden_size // num_heads

        self.q_proj = nn.Linear(hidden_size, hidden_size)
        self.k_proj = nn.Linear(hidden_size, hidden_size)
        self.v_proj = nn.Linear(hidden_size, hidden_size)
        self.o_proj = nn.Linear(hidden_size, hidden_size)
        self.dropout = nn.Dropout(dropout)

    def _split_heads(self, x):
        batch, seq_len, _ = x.shape
        x = x.view(batch, seq_len, self.num_heads, self.head_dim)
        return x.transpose(1, 2)

    def _merge_heads(self, x):
        batch, _, seq_len, _ = x.shape
        x = x.transpose(1, 2).contiguous()
        return x.view(batch, seq_len, self.hidden_size)

    def forward(self, hidden_states, attn_mask=None):
        query = self._split_heads(self.q_proj(hidden_states))
        key = self._split_heads(self.k_proj(hidden_states))
        value = self._split_heads(self.v_proj(hidden_states))

        output, weights = scaled_dot_product_attention(query, key, value, attn_mask)
        output = self.o_proj(self._merge_heads(output))
        return output, weights

单层缓存接口的写法

class CachedMultiHeadAttention(MultiHeadAttention):
    def forward(self, hidden_states, attn_mask=None, past_key_value=None, use_cache=False):
        query = self._split_heads(self.q_proj(hidden_states))
        key = self._split_heads(self.k_proj(hidden_states))
        value = self._split_heads(self.v_proj(hidden_states))

        if past_key_value is not None:
            past_key, past_value = past_key_value
            key = torch.cat([past_key, key], dim=2)
            value = torch.cat([past_value, value], dim=2)

        output, weights = scaled_dot_product_attention(query, key, value, attn_mask)
        output = self.o_proj(self._merge_heads(output))
        new_cache = (key, value) if use_cache else None
        return output, new_cache, weights

统一版 GQA/MQA

def repeat_kv(x, repeats):
    # x: (batch, kv_heads, seq_len, head_dim)
    if repeats == 1:
        return x
    batch, kv_heads, seq_len, head_dim = x.shape
    x = x[:, :, None, :, :].expand(batch, kv_heads, repeats, seq_len, head_dim)
    return x.reshape(batch, kv_heads * repeats, seq_len, head_dim)


class GroupedQueryAttention(nn.Module):
    def __init__(self, hidden_size, num_heads, kv_heads):
        super().__init__()
        if hidden_size % num_heads != 0:
            raise ValueError("hidden_size must be divisible by num_heads")
        if num_heads % kv_heads != 0:
            raise ValueError("num_heads must be divisible by kv_heads")

        self.hidden_size = hidden_size
        self.num_heads = num_heads
        self.kv_heads = kv_heads
        self.head_dim = hidden_size // num_heads

        self.q_proj = nn.Linear(hidden_size, hidden_size)
        self.k_proj = nn.Linear(hidden_size, kv_heads * self.head_dim)
        self.v_proj = nn.Linear(hidden_size, kv_heads * self.head_dim)
        self.o_proj = nn.Linear(hidden_size, hidden_size)

    def _split(self, x, heads):
        batch, seq_len, _ = x.shape
        x = x.view(batch, seq_len, heads, self.head_dim)
        return x.transpose(1, 2)

    def forward(self, hidden_states, attn_mask=None):
        query = self._split(self.q_proj(hidden_states), self.num_heads)
        key = self._split(self.k_proj(hidden_states), self.kv_heads)
        value = self._split(self.v_proj(hidden_states), self.kv_heads)

        repeats = self.num_heads // self.kv_heads
        key = repeat_kv(key, repeats)
        value = repeat_kv(value, repeats)

        output, weights = scaled_dot_product_attention(query, key, value, attn_mask)
        batch, _, seq_len, _ = output.shape
        output = output.transpose(1, 2).contiguous().view(batch, seq_len, self.hidden_size)
        return self.o_proj(output), weights

1. 先定义输入：hidden_states 是 (B,T,C)，C = H * D。
2. 说明 Q/K/V：三个线性层把同一个输入投影到 query、key、value 空间。
3. 拆 head：从 (B,T,C) 变为 (B,H,T,D)，便于每个 head 独立做注意力。
4. 计算 score：query @ key.transpose(-1, -2) / sqrt(D)，得到 (B,H,T,T)。
5. 加 mask：causal mask 管未来，padding mask 管无效 token。
6. softmax 后乘 value：得到每个位置的上下文向量。
7. 合并 head：转回 (B,T,C)，再经过输出投影。
8. 如果追问推理优化，再讲 KV Cache、MQA/GQA 为什么减少 decode 阶段的内存带宽。

这篇材料适合作为 Attention 手撕题的入口，因为它覆盖了从 SDPA 到 GQA 的主线。但真正要把它变成可面试、可工程使用的知识，需要把“实现片段”提升成“张量不变量”：每一步都能写出 batch/head/query/key/value 的形状，知道 mask 加在哪个维度，知道 cache 复用的是历史 K/V 而不是历史输出。

最值得沉淀的 insight 是：MQA/GQA 并不是新的注意力公式，而是 K/V 共享策略的改变。这个视角会让实现更 DRY，也更容易解释为什么现代 LLM 在推理服务里偏好 GQA：decode 阶段的瓶颈经常不是算不出矩阵乘，而是读写庞大的 KV cache 太贵。

后续如果继续补全这套“手撕经典算法”笔记，优先补 MLA、RoPE、RMSNorm、SwiGLU、MoE router 和 FlashAttention。它们和 Attention 的关系很紧：一个处理位置，一个处理归一化，一个处理 FFN 表达，一个处理稀疏专家，一个处理 IO-aware attention kernel。