#十四点五、手撕代码完整参考实现（新增）

前面列出了手撕代码的题单和评分要点，但面试准备时“知道考什么”和“能手写出来”之间还有一段距离。本轮分成两条线：先用一个最小 NumPy Multi-Head Attention 对齐判题机的输入、权重、mask 和输出；再把它扩展到完整 Transformer。面试主线只使用 NumPy 的数组运算，不调用现成 Attention 层；后面的标准库版本和框架代码作为对照。

#面试主线：NumPy Multi-Head Attention

完整实现：mha_interview.py；测试：test_mha_interview.py。这才是面试现场应该先写的版本：输入和输出统一为 (B, T, D)，W_q/W_k/W_v/W_o 由调用方传入，避免随机初始化掩盖公式，也方便判题机用同一组权重做数值对照。

面试实现的固定契约

Q = linear(X, W_q, b_q)
K = linear(X, W_k, b_k)
V = linear(X, W_v, b_v)

Q = reshape_to_heads(Q)  # (B, H, T, d)
K = reshape_to_heads(K)
V = reshape_to_heads(V)

scores = Q @ transpose(K, -2, -1) / sqrt(d)
scores = apply_causal_and_padding_mask(scores)
weights = stable_softmax(scores, dim=-1)
heads = weights @ V

heads = merge_heads(heads)  # (B, T, D)
output = linear(heads, W_o, b_o)

和参考实现一致的必要条件：必须使用同一组投影权重和 bias、同一 batch-first/layout、同一缩放维度、同一 mask 语义、关闭 dropout，并在相同精度下比较。具体 Hugging Face 模型还可能增加 RoPE、GQA/MQA、RMSNorm 等模块，不能只拿基础 MHA 去对齐整个模型输出。

#参考实现 0：纯 Python Attention 与 Decoder-only Transformer

完整代码：transformer_from_scratch.py；回归测试：test_transformer_from_scratch.py。运行测试：python3 -m unittest discover -s notes/llm-interview-question-bank/code -p 'test_*.py'。实现的核心数值依赖只来自标准库中的 math 和 random；它的目标是让数据流可读，不追求生产级矩阵 kernel 的速度。

先把每一步的 shape 说清楚

Attention 的核心不是一张相似度表，而是一个受 mask 约束的信息路由器

对第 \(i\) 个 query，先计算 \(s_{ij}=q_i\cdot k_j/\sqrt{d}\)，再对允许读取的位置做稳定 softmax，最后计算 \(o_i=\sum_j a_{ij}v_j\)。代码中的 allowed 列表同时承载两种约束：causal attention 要求 \(j\leq i\)，padding mask 要求 key 位置是真实 token。被禁止的位置权重直接为 0，而不是让它们参与归一化。

def one_query_attention(query, keys, values, query_position):
    scores = []
    allowed = []
    for key_position, key in enumerate(keys):
        scores.append(dot(query, key) / sqrt(head_dim))
        allowed.append(key_position <= query_position)  # causal mask

    weights = stable_softmax(scores, allowed)  # 减最大值，避免 exp 溢出
    output = [
        sum(weights[j] * values[j][dimension] for j in range(len(values)))
        for dimension in range(head_dim)
    ]
    return output, weights

这里的代码块只展示逻辑骨架，完整版本把 dot、sqrt、稳定 softmax、矩阵乘法和多头循环全部展开。注意：sqrt(d) 只控制 logits 尺度，不会让 causal mask 自动出现；mask 必须显式构造。

一个 Transformer block 如何串起来

实现采用常见的 pre-norm 路径：x → LayerNorm → Self-Attention → residual add → LayerNorm → GELU FFN → residual add。Attention 在 token 维度混合信息，FFN 对每个 token 独立做非线性变换；两次 residual 让新分支只需学习“如何修正旧表示”。多个 block 叠起来后，最后的 LayerNorm 和词表投影才产生 logits。

代码还提供了一个最小的 greedy generate。它每生成一个 token 都重新跑当前上下文，因此是教学上最直观的 \(O(T^2)\) 路径；生产实现会保存历史 K/V，避免重复计算。这个差异正是“会写 Attention”与“理解 Transformer 推理系统”的连接点。

这份实现的边界

• 权重是随机初始化的，所以输出只用于检查 shape 和信息流，不代表模型已经学会语言。
• 没有反向传播、优化器、dropout、量化、KV cache 或高性能 kernel；下一步如果要继续手撕训练，应该先实现一个小损失和有限差分/反向传播，而不是直接跳到完整训练框架。
• 这里实现的是 decoder-only self-attention。Encoder-decoder 的 cross-attention 只需把 Q 来自 decoder hidden state、K/V 改为 encoder 输出，并使用相应的 padding mask；核心缩放点积不变。

#参考实现 1：手写 Multi-Head Attention（PyTorch API 对照）

题目要求：给定 Q/K/V 三个张量，实现 Scaled Dot-Product Attention + Multi-Head 拆分与合并，支持 padding mask 和 causal mask，并保证数值稳定性。

import torch
import torch.nn as nn
import math

class MultiHeadAttention(nn.Module):
    """
    手写 Multi-Head Attention 完整实现。

    面试评分点对应：
    1. shape 是否正确流转                    -> 每步都有 shape 注释
    2. 是否记得除以 sqrt(d_k)                -> scaled_dot_product_attention 内
    3. causal mask 方向是否正确               -> scaled_dot_product_attention 内，上三角填 -1e9
    4. 数值稳定性（softmax 前减最大值）        -> scaled_dot_product_attention 内
    5. 多头拆分与合并的维度操作是否正确         -> view + transpose / transpose + view
    """

    def __init__(self, d_model: int, num_heads: int, dropout: float = 0.1):
        super().__init__()
        assert d_model % num_heads == 0, "d_model 必须能被 num_heads 整除"

        self.d_model = d_model          # 模型隐藏维度，如 512, 768, 4096
        self.num_heads = num_heads      # 注意力头数，如 8, 12, 32
        self.head_dim = d_model // num_heads   # 每头维度 d_k，如 64

        # 线性投影：把输入映射到 Q/K/V 空间
        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)

        # 输出投影：把多头拼接后的结果映射回 d_model
        self.W_o = nn.Linear(d_model, d_model)

        self.dropout = nn.Dropout(dropout)
        self.scale = math.sqrt(self.head_dim)

    def forward(self, x: torch.Tensor, mask: torch.Tensor = None, is_causal: bool = False):
        B, T, _ = x.shape

        # Step 1: 线性投影
        Q = self.W_q(x)   # (B, T, d_model)
        K = self.W_k(x)   # (B, T, d_model)
        V = self.W_v(x)   # (B, T, d_model)

        # Step 2: 拆分为多头: (B, T, d) -> (B, num_heads, T, head_dim)
        Q = Q.view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
        K = K.view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
        V = V.view(B, T, self.num_heads, self.head_dim).transpose(1, 2)

        # Step 3: Scaled Dot-Product Attention
        attn_output, attn_weights = self.scaled_dot_product_attention(
            Q, K, V, mask=mask, is_causal=is_causal
        )

        # Step 4: 多头合并
        attn_output = attn_output.transpose(1, 2).contiguous()   # (B, T, H, d_k)
        attn_output = attn_output.view(B, T, self.d_model)       # (B, T, d_model)
        # 踩坑点：transpose 后必须 contiguous()，否则 view 可能报错

        # Step 5: 输出投影
        output = self.W_o(attn_output)
        output = self.dropout(output)
        return output, attn_weights

    def scaled_dot_product_attention(self, Q, K, V, mask=None, is_causal=False):
        B, H, T, d_k = Q.shape

        # 注意力分数: (B, H, T, d_k) @ (B, H, d_k, T) -> (B, H, T, T)
        scores = torch.matmul(Q, K.transpose(-2, -1)) / self.scale

        # 应用 mask
        if mask is not None:
            if mask.dim() == 2:
                mask = mask.unsqueeze(1).unsqueeze(2)
            scores = scores.masked_fill(mask == 0, -1e9)

        if is_causal:
            causal_mask = torch.triu(torch.ones(T, T, device=scores.device, dtype=torch.bool), diagonal=1)
            scores = scores.masked_fill(causal_mask.unsqueeze(0).unsqueeze(0), -1e9)

        # 除以 sqrt(d_k) 只控制尺度，不能替代减最大值；torch.softmax 内部已做这一步，
        # 这里显式写出，对应手写 exp / sum 时必须先减每行最大值
        attn_weights = torch.softmax(
            scores - scores.amax(dim=-1, keepdim=True), dim=-1
        )
        attn_weights = self.dropout(attn_weights)

        # 加权求和: (B, H, T, T) @ (B, H, T, d_k) -> (B, H, T, d_k)
        output = torch.matmul(attn_weights, V)
        return output, attn_weights


# 快速测试
def test_mha():
    B, T, d = 2, 10, 512
    num_heads = 8
    mha = MultiHeadAttention(d_model=d, num_heads=num_heads)
    x = torch.randn(B, T, d)

    # Test 1: 无 mask
    out, attn = mha(x)
    assert out.shape == (B, T, d)
    assert attn.shape == (B, num_heads, T, T)
    print("Test 1 通过: 无 mask 前向")

    # Test 2: padding mask
    pad_mask = torch.ones(B, T, dtype=torch.bool)
    pad_mask[:, 8:] = False
    out2, attn2 = mha(x, mask=pad_mask)
    assert torch.allclose(attn2[:, :, :, 8:].sum(), torch.tensor(0.0), atol=1e-6)
    print("Test 2 通过: padding mask 正确")

    # Test 3: causal mask
    out3, attn3 = mha(x, is_causal=True)
    upper_tri = torch.triu(torch.ones(T, T), diagonal=1).bool()
    assert torch.allclose(attn3[:, :, upper_tri].sum(), torch.tensor(0.0), atol=1e-6)
    print("Test 3 通过: causal mask 正确")

    print("\n所有测试通过！")

if __name__ == "__main__":
    test_mha()

#面试官追问要点

1. 为什么除以 sqrt(d_k)？ 假设 Q、K 元素独立同分布 N(0,1)，则点积方差为 d_k。除以 sqrt(d_k) 把方差拉回 1，避免 softmax 输入过大导致梯度消失。
2. 为什么 masked_fill 填 -1e9 而不是 -inf？ 当一整行都被 mask 时，-inf 会让 softmax 在减最大值时出现 -inf - (-inf)，整行变成 NaN，并沿反向传播污染梯度；这与硬件无关。-1e9 只是避免了 NaN：整行都是 -1e9 时 softmax 退化成均匀分布，仍会读到 padding，所以全屏蔽的 query 行还要在输出端清零。另外 fp16 下 -1e9 会溢出成 -inf，半精度实现通常改用 torch.finfo(dtype).min。
3. 为什么 transpose 后要 contiguous()？ PyTorch 的 view 要求张量内存连续。transpose 只改变 strides 不复制数据，所以必须 contiguous() 后再 view。
4. W_q/W_k/W_v 能否合并？ 可以，很多实现用 nn.Linear(d_model, 3*d_model) 一次性投影再切分，效率略高。面试写分开的更清晰。

#参考实现 2：手写简化 PPO 算法（PyTorch）

题目要求：实现简化版 PPO，包含 Actor-Critic、GAE、PPO-Clip 目标、KL early stopping。

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.distributions import Categorical
import numpy as np

class ActorCritic(nn.Module):
    def __init__(self, obs_dim: int, action_dim: int, hidden_dim: int = 256):
        super().__init__()
        self.shared = nn.Sequential(
            nn.Linear(obs_dim, hidden_dim), nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim), nn.ReLU(),
        )
        self.actor_head = nn.Linear(hidden_dim, action_dim)
        self.critic_head = nn.Linear(hidden_dim, 1)

    def forward(self, obs):
        feat = self.shared(obs)
        return self.actor_head(feat), self.critic_head(feat)

    def get_action_and_value(self, obs, action=None):
        logits, value = self.forward(obs)
        dist = Categorical(logits=logits)
        if action is None:
            action = dist.sample()
        return action, dist.log_prob(action), dist.entropy(), value.squeeze(-1)


def compute_gae(rewards, values, dones, gamma=0.99, lam=0.95):
    """GAE: 反向计算，从最后一个时间步往前推。"""
    T = len(rewards)
    advantages = torch.zeros_like(rewards)
    last_gae = 0.0
    for t in reversed(range(T)):
        next_value = 0.0 if t == T - 1 else values[t + 1]
        delta = rewards[t] + gamma * next_value * (1 - dones[t]) - values[t]
        last_gae = delta + gamma * lam * (1 - dones[t]) * last_gae
        advantages[t] = last_gae
    return advantages, advantages + values


class PPOTrainer:
    def __init__(self, ac, lr=3e-4, gamma=0.99, lam=0.95,
                 clip_eps=0.2, value_coef=0.5, entropy_coef=0.01,
                 max_kl=0.015, num_epochs=4, batch_size=64, device="cpu"):
        self.ac = ac.to(device)
        self.optimizer = torch.optim.Adam(self.ac.parameters(), lr=lr)
        self.gamma, self.lam = gamma, lam
        self.clip_eps = clip_eps
        self.value_coef = value_coef
        self.entropy_coef = entropy_coef
        self.max_kl = max_kl
        self.num_epochs = num_epochs
        self.batch_size = batch_size
        self.device = device

    def update(self, traj):
        obs = traj["obs"].to(self.device)
        actions = traj["actions"].to(self.device)
        rewards = traj["rewards"].to(self.device)
        dones = traj["dones"].to(self.device)
        old_log_probs = traj["old_log_probs"].to(self.device)
        old_values = traj["old_values"].to(self.device)
        T = len(obs)

        # Step 1: GAE
        with torch.no_grad():
            advantages, returns = compute_gae(rewards, old_values, dones, self.gamma, self.lam)
            advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)

        # Step 2: PPO epochs
        for epoch in range(self.num_epochs):
            indices = torch.randperm(T)
            for start in range(0, T, self.batch_size):
                idx = indices[start:start + self.batch_size]
                _, new_log_probs, entropy, new_values = self.ac.get_action_and_value(
                    obs[idx], actions[idx]
                )

                # PPO-Clip
                ratio = torch.exp(new_log_probs - old_log_probs[idx])
                surr1 = ratio * advantages[idx]
                surr2 = torch.clamp(ratio, 1 - self.clip_eps, 1 + self.clip_eps) * advantages[idx]
                policy_loss = -torch.min(surr1, surr2).mean()

                value_loss = F.mse_loss(new_values, returns[idx])
                loss = policy_loss + self.value_coef * value_loss - self.entropy_coef * entropy.mean()

                self.optimizer.zero_grad()
                loss.backward()
                nn.utils.clip_grad_norm_(self.ac.parameters(), 0.5)
                self.optimizer.step()

                # KL early stopping
                with torch.no_grad():
                    kl = (old_log_probs[idx] - new_log_probs).mean().abs()
                if kl > self.max_kl:
                    print(f"KL={kl:.4f} > {self.max_kl}, early stopping")
                    return

#手撕 PPO 的面试官评分 checklist

#参考实现 3：极简 Attention（PyTorch API 对照）

import torch, math
def attention(q, k, v, mask=None):
    scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(q.size(-1))
    if mask is not None:
        scores = scores.masked_fill(mask == 0, float('-inf'))
    attn = torch.softmax(scores, dim=-1)
    return torch.matmul(attn, v), attn

追问点：causal mask 怎么加？为什么除以 sqrt(d_k)？多头在哪里拆？如果某一行的 key 全部被屏蔽，softmax 会输出什么？这里用 -inf 填充会得到 NaN；参考实现 1 用 -1e9 覆盖整行则会变成均匀分布，悄悄读到 padding。面试主线的 NumPy 版本因此对全屏蔽行单独返回全 0。