手撕代码题（模型与 Agent 方向）

本章导读：大模型岗的手撕分两类：一类是通用算法题，另一类是“本岗位特色题”——手写 MHA、RoPE、LoRA、KV Cache、采样、MoE、ReAct。后者近两年权重明显上升，面试官看的不只是 AC，还有维度、mask、数值稳定性和边界。

使用建议：每道题先自己在纸上写，再对照代码；重点记张量维度和易错点，不要背整段实现。

核心要点

• 模型题先写清张量维度，再写计算；维度错了，后面全错。
• attention 的三个易错点：缩放、mask、多头 reshape/transpose。
• RoPE 要能解释“成对旋转”和相对位置性质，代码只需写对 cos/sin 与 rotate_half。
• LoRA 要会算参数量、初始化、merge；注意 scaling 是 \(\alpha/r\)。
• KV Cache 只缓存历史 K/V，当前 token 的 Q 不缓存。
• 采样要区分 top-k（固定数量）与 top-p（累计概率），并处理全 mask/数值稳定。
• MoE 要写对 top-k 路由、权重归一化、负载统计；真实实现用 grouped GEMM。
• ReAct 要写清循环、终止条件、工具异常、最大步数。
• 手撕时先讲思路与复杂度，再写代码，最后跑一个样例。

1. 多头自注意力（MHA）

要求：实现 MultiHeadAttention(d_model, n_heads)，支持 causal mask，返回 [B, S, d_model]。

关键点：

• d_head = d_model // n_heads；
• QKV 投影后拆头：[B, S, H, D] -> [B, H, S, D]；
• 缩放因子 \(1/\sqrt{d_{head}}\)；
• mask 用 -inf 或 dtype 最小值，避免全 mask 行 NaN；
• 输出拼接后过 W_O。

import math
import torch
import torch.nn as nn
import torch.nn.functional as F


class MHA(nn.Module):
    def __init__(self, d_model: int, n_heads: int, dropout: float = 0.0):
        super().__init__()
        assert d_model % n_heads == 0
        self.h = n_heads
        self.dh = d_model // n_heads
        self.qkv = nn.Linear(d_model, 3 * d_model, bias=False)
        self.o = nn.Linear(d_model, d_model, bias=False)
        self.drop = nn.Dropout(dropout)

    def forward(self, x, causal=True):
        B, S, D = x.shape
        q, k, v = self.qkv(x).chunk(3, dim=-1)

        # [B, S, H, Dh] -> [B, H, S, Dh]
        q = q.view(B, S, self.h, self.dh).transpose(1, 2)
        k = k.view(B, S, self.h, self.dh).transpose(1, 2)
        v = v.view(B, S, self.h, self.dh).transpose(1, 2)

        scores = q @ k.transpose(-2, -1) / math.sqrt(self.dh)  # [B,H,S,S]
        if causal:
            mask = torch.triu(torch.ones(S, S, device=x.device, dtype=torch.bool), 1)
            scores = scores.masked_fill(mask, torch.finfo(scores.dtype).min)

        attn = self.drop(F.softmax(scores, dim=-1))
        out = attn @ v                                       # [B,H,S,Dh]
        out = out.transpose(1, 2).contiguous().view(B, S, D)
        return self.o(out)

追问：

• 为什么除以 \(\sqrt{d_h}\)？→ 控制点积方差。
• 为什么用 transpose 后要 contiguous？→ view 需要连续内存。
• 全 mask 行怎么处理？→ 保证至少一个可见位置，或 softmax 后置零。
• 参数量？→ \(4d^2\)。

2. 带 RoPE 的注意力

要求：在 MHA 基础上对 Q/K 应用 RoPE。

关键点：

• 频率 \(\theta_i=\text{base}^{-2i/d}\)；
• 位置 \(m\) 的旋转角 \(m\theta_i\)；
• rotate_half 实现前后半交换；
• 只作用于 Q/K，不作用于 V。

def rope_cache(S, dim, base=10000.0, device="cpu"):
    inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, device=device).float() / dim))
    t = torch.arange(S, device=device).float()
    freqs = torch.outer(t, inv_freq)          # [S, dim/2]
    return torch.cos(freqs), torch.sin(freqs)


def rotate_half(x):
    x1, x2 = x.chunk(2, dim=-1)
    return torch.cat([-x2, x1], dim=-1)


def apply_rope(x, cos, sin):
    # x: [B,H,S,Dh]; cos/sin: [S, Dh/2]
    cos = torch.cat([cos, cos], dim=-1)[None, None, :, :]
    sin = torch.cat([sin, sin], dim=-1)[None, None, :, :]
    return x * cos + rotate_half(x) * sin

追问：

• 为什么内积只依赖相对位置？→ \(R(m)^\top R(n)=R(n-m)\)。
• base 变大有什么影响？→ 波长变长，外推更好，但局部分辨率下降。
• PI/NTK/YaRN 的区别？→ 位置压缩 / 基频缩放 / 分段+温度修正。

3. RMSNorm

要求：实现 RMSNorm，并说明与 LayerNorm 的区别。

class RMSNorm(nn.Module):
    def __init__(self, d: int, eps: float = 1e-6):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(d))
        self.eps = eps

    def forward(self, x):
        # 在最后一维上求均方根
        rms = x.pow(2).mean(dim=-1, keepdim=True).add(self.eps).rsqrt()
        return self.weight * x * rms

追问：

• 为什么去掉均值？→ 特征均值接近 0，缩放是主要作用；省一次归约。
• 为什么用 rsqrt？→ 数值稳定且快。
• 与 LayerNorm 的参数量？→ RMSNorm 少一个 bias。

4. SwiGLU FFN

要求：实现 SwiGLU，并说明 \(d_{ff}=8d/3\) 的来历。

class SwiGLU(nn.Module):
    def __init__(self, d, d_ff):
        super().__init__()
        self.w1 = nn.Linear(d, d_ff, bias=False)  # gate
        self.w3 = nn.Linear(d, d_ff, bias=False)  # value
        self.w2 = nn.Linear(d_ff, d, bias=False)

    def forward(self, x):
        return self.w2(F.silu(self.w1(x)) * self.w3(x))

参数量：\(3d\cdot d_{ff}\)。令其等于 \(8d^2\)，得 \(d_{ff}=8d/3\)。

追问：

• 为什么门控有效？→ 动态选择信息，增加表达力。
• 为什么实际取 128 的倍数？→ 硬件对齐。

5. LoRA 层

要求：实现 LoRA 线性层，支持初始化、前向、merge。

class LoRALinear(nn.Module):
    def __init__(self, base: nn.Linear, r=8, alpha=16, dropout=0.0):
        super().__init__()
        self.base = base
        for p in self.base.parameters():
            p.requires_grad = False

        self.r = r
        self.scaling = alpha / r
        self.A = nn.Parameter(torch.randn(r, base.in_features) * 0.01)
        self.B = nn.Parameter(torch.zeros(base.out_features, r))
        self.drop = nn.Dropout(dropout)
        self.merged = False

    def forward(self, x):
        # B 初始化为 0，保证训练开始时等价于原模型
        if self.merged:
            return self.base(x)
        return self.base(x) + self.drop(x) @ self.A.T @ self.B.T * self.scaling

    @torch.no_grad()
    def merge(self):
        # W' = W + scaling * B @ A；重复调用会重复累加，需先判重
        if self.merged:
            return
        self.base.weight.add_(self.B @ self.A * self.scaling)
        self.merged = True

参数量：\(r(d_{\text{in}}+d_{\text{out}})\)。

追问：

• 为什么 B 初始化为 0？→ 初始时 \(\Delta W=0\)，不破坏基座。
• scaling 为什么是 \(\alpha/r\)？→ 解耦秩与学习率尺度。
• merge 后还能继续训练吗？→ 可以，但需要重新挂 LoRA 或解冻。

6. KV Cache 单步 decode

要求：实现带 KV Cache 的单步注意力：输入当前 token，缓存历史 K/V。

class KVCacheAttention(nn.Module):
    def __init__(self, d_model, n_heads):
        super().__init__()
        self.h = n_heads
        self.dh = d_model // n_heads
        self.qkv = nn.Linear(d_model, 3 * d_model, bias=False)
        self.o = nn.Linear(d_model, d_model, bias=False)

    def forward(self, x, cache=None):
        # x: [B, 1, D]（decode 单步）
        B, S, D = x.shape
        assert S == 1, "KV Cache 版本只演示单步 decode；prefill 请另写带 causal mask 的实现"
        q, k, v = self.qkv(x).chunk(3, dim=-1)
        q = q.view(B, S, self.h, self.dh).transpose(1, 2)
        k = k.view(B, S, self.h, self.dh).transpose(1, 2)
        v = v.view(B, S, self.h, self.dh).transpose(1, 2)

        if cache is not None:
            k_old, v_old = cache
            k = torch.cat([k_old, k], dim=2)   # 拼接历史 K
            v = torch.cat([v_old, v], dim=2)

        new_cache = (k, v)
        scores = q @ k.transpose(-2, -1) / math.sqrt(self.dh)
        # 单步 decode 不需要 causal mask：当前 token 可见全部历史
        attn = F.softmax(scores, dim=-1)
        out = (attn @ v).transpose(1, 2).contiguous().view(B, S, D)
        return self.o(out), new_cache

追问：

• 为什么只缓存 K/V，不缓存 Q？→ Q 只用于当前步，历史 Q 不再需要。
• KV Cache 显存公式？→ \(2L n_{kv}d_h S B\) bytes。
• 为什么 decode 是 memory-bound？→ 每步矩阵-向量，算术强度低。

7. Top-k / Top-p / Temperature 采样

要求：实现一个采样函数，支持 temperature、top-k、top-p。

def sample(logits, temperature=1.0, top_k=None, top_p=None):
    # logits: [V]
    if temperature <= 0:
        return torch.argmax(logits).item()

    logits = logits / temperature
    probs = torch.softmax(logits, dim=-1)

    if top_k is not None and top_k > 0:
        k = min(top_k, probs.numel())
        kth = torch.topk(probs, k).values[-1]
        probs = torch.where(probs < kth, torch.zeros_like(probs), probs)

    if top_p is not None and 0 < top_p < 1:
        sorted_probs, sorted_idx = torch.sort(probs, descending=True)
        cum = torch.cumsum(sorted_probs, dim=-1)
        # 保留累计概率达到 top_p 的最小集合
        mask = cum - sorted_probs >= top_p
        sorted_probs[mask] = 0.0
        probs = torch.zeros_like(probs).scatter(0, sorted_idx, sorted_probs)

    if probs.sum() <= 0:
        return torch.argmax(logits).item()
    probs = probs / probs.sum()
    return torch.multinomial(probs, 1).item()

追问：

• top-k 与 top-p 的区别？→ 固定数量 vs 自适应累计概率。
• temperature 的作用？→ 缩放 logits，改变分布尖锐度。
• 数值稳定？→ 先减 max 再 softmax；全零时回退 argmax。

8. MoE 路由

要求：实现 top-k 门控，返回聚合输出与负载统计。

class TopKRouter(nn.Module):
    def __init__(self, d, n_experts, top_k=2):
        super().__init__()
        self.gate = nn.Linear(d, n_experts, bias=False)
        self.top_k = top_k

    def forward(self, x, experts):
        # x: [N, D]
        logits = self.gate(x)
        probs = torch.softmax(logits, dim=-1)
        top_p, top_i = probs.topk(self.top_k, dim=-1)  # [N,k]

        # 负载统计：每个专家被选中的 token 数
        load = torch.bincount(top_i.flatten(), minlength=len(experts)).float()

        # 简化：逐 token 调专家；生产用 scatter/gather + grouped GEMM
        out = torch.zeros_like(x)
        for i in range(x.size(0)):
            for j in range(self.top_k):
                out[i] += top_p[i, j] * experts[top_i[i, j]](x[i])
        return out, load

追问：

• 为什么需要负载均衡？→ 防止赢家通吃、专家欠训练。
• 辅助损失公式？→ \(\alpha E\sum f_iP_i\)。
• 容量因子？→ 限制每专家 token 数，超出丢弃或走残差。

9. 最小 ReAct Agent

要求：实现一个 ReAct 循环：工具注册、Thought/Action/Observation、最大步数、异常处理。

import json


class Agent:
    def __init__(self, llm, tools, max_steps=8):
        self.llm = llm
        self.tools = tools          # {name: callable}
        self.max_steps = max_steps

    def run(self, question):
        history = [f"Question: {question}"]
        for step in range(self.max_steps):
            prompt = "\n".join(history) + "\nThought:"
            text = self.llm(prompt)

            if "Final Answer:" in text:
                return text.split("Final Answer:", 1)[1].strip()

            if "Action:" not in text:
                history.append(text)
                continue

            try:
                action = text.split("Action:", 1)[1].strip()
                name, arg = action.split("(", 1)
                name = name.strip()
                arg = arg.rsplit(")", 1)[0].strip().strip('"').strip("'")

                if name not in self.tools:
                    raise ValueError(f"unknown tool: {name}")

                obs = self.tools[name](arg)
                history.append(f"{text}\nObservation: {obs}")
            except Exception as e:
                # 工具错误要回填给模型，而不是直接崩溃
                history.append(f"{text}\nObservation: ERROR: {e}")

        return "达到最大步数，未能完成任务。"

追问：

• 怎么终止？→ Final Answer、最大步数、超时、重复动作检测。
• 工具报错怎么办？→ 回填错误、重试、换工具、请求人工。
• 怎么防死循环？→ 预算、动作指纹、状态哈希。
• 安全？→ 工具白名单、参数校验、沙箱、审批。

10. RRF 融合

要求：实现多路召回结果的 Reciprocal Rank Fusion。

def rrf(rankings, k=60):
    # rankings: list[list[doc_id]]，每路按相关性降序
    score = {}
    for docs in rankings:
        for rank, doc in enumerate(docs, start=1):
            score[doc] = score.get(doc, 0.0) + 1.0 / (k + rank)
    return sorted(score.items(), key=lambda x: -x[1])

追问：

• 为什么用 RRF？→ 不需要分数归一化，对异构检索器鲁棒。
• k 的作用？→ 平滑头部优势；常用 60。
• 与加权分数融合对比？→ RRF 更稳，加权需要校准。

11. BM25（简化）

import math
from collections import Counter


def bm25(query, docs, k1=1.5, b=0.75):
    N = len(docs)
    avgdl = sum(len(d) for d in docs) / max(N, 1)
    df = Counter()
    for d in docs:
        for w in set(d):
            df[w] += 1

    scores = []
    q = query.split()
    for d in docs:
        tf = Counter(d)
        s = 0.0
        for w in q:
            if w not in tf:
                continue
            idf = math.log(1 + (N - df[w] + 0.5) / (df[w] + 0.5))
            denom = tf[w] + k1 * (1 - b + b * len(d) / avgdl)
            s += idf * tf[w] * (k1 + 1) / denom
        scores.append(s)
    return scores

追问：

• BM25 的 idf/tf 饱和分别解决什么？→ 稀有词权重、词频边际递减。
• 为什么 RAG 里还要 BM25？→ 精确匹配、编号、人名。
• 与向量检索怎么融合？→ RRF 或加权。

12. 通用算法题清单

题目	关键点	复杂度
反转链表	三指针/递归	O(n)
链表相加	dummy 头、进位	O(n)
排序链表	归并	O(n log n)
二叉树层序	队列	O(n)
路径总和 III	前缀和 + 哈希	O(n)
接雨水	双指针/单调栈	O(n)
无重复最长子串	滑动窗口	O(n)
最大正方形	DP	O(mn)
Top-K	堆/quickselect	O(n log k)/O(n)
字符串解码	栈	O(n)
两正序数组中位数	二分	O(log min(m,n))
最长递增子序列	DP/二分	O(n log n)
只出现一次的数字	异或	O(n)

13. 手撕答题模板

1. 澄清：输入输出、维度、是否 batch、是否 causal、精度要求。
2. 思路：先讲公式/流程，再讲复杂度。
3. 边界：空输入、单 token、全 mask、长度不整除头数。
4. 实现：变量命名清晰，维度注释。
5. 测试：给一个 2×3 的小例子，口算验证。
6. 追问：参数量、显存、数值稳定、怎么并行/加速。

本文档代码用于面试手写与理解，生产实现请使用框架融合算子与成熟库。