注意力卸载：给大模型请一位“帮厨”，让推理不再卡顿

可视化演示

展示FCFS、SSTF、SCAN三种磁盘调度算法的磁头移动过程，可视化请求队列服务顺序和寻道距离优化。

显存瓶颈

想象你是一位主厨，面前是一张只够放四盘菜的操作台。客人点了一道需要十几个配料的“佛跳墙”——这就是大模型中的注意力机制。注意力计算在一瞬间会生成一张巨大的“配料表”（注意力分数矩阵），这张表的尺寸是序列长度的平方。当序列长度从 2048 涨到 8192，这张表的大小直接膨胀 16 倍。操作台瞬间被铺满，连切菜的地方都没有了。

在技术世界里，这张“配料表”就是注意力机制中的 QK^T 矩阵和 softmax 之后的概率矩阵。它们属于“临时结果”——只在当前层计算时需要，算完下一层就不再使用。但问题在于，这些临时结果必须“同时在场”才能算出最终输出，于是它们会一直占据显存，直到整个解码步完成。显存里还同时住着模型权重、KV Cache 和激活值，四者争抢有限空间，就像四个厨师共用一个操作台。

结果就是：你不得不把每桌客人的批次缩小（降低 batch size），一次只服务一位客人。GPU 虽然算力强大，但因为显存装不下更多的中间结果，大部分计算单元只能空转，利用率可能跌破 30%。这就像米其林大厨被限制一次只炒一个菜，锅铲都闲出火星了，出菜速度却上不去。

卸载原理

帮厨小张想出了一个办法：主厨操作台不够用时，把暂时不用的“半成品配料”端到储物间（CPU 内存）或冷库（SSD）里放着，等主厨喊“需要番茄”时再递回来。这一端一递，就是“注意力卸载”。

注意力卸载的核心理念是：注意力计算的中间结果（QK^T 矩阵、softmax 输出）并非需要在显存中驻留到最后一刻。它们只在当前解码步骤的某一小段计算中被使用。因此，可以在算完某个子模块后，立即将这些中间张量通过 PCIe 总线拷贝到 CPU 内存，甚至写入 SSD；当后续计算（比如加权求和）需要它们时，再拷贝回来。整个过程对上层逻辑透明——模型还是那个模型，只是中间数据在“外置仓库”里周转了一圈。

卸载的粒度可以精细控制。最粗糙的粒度是整个注意力层的输出一起卸，最精细的粒度是每个注意力头（head）单独卸。粒度越细，显存释放越及时，但搬运次数也越多。一个注意力头对应一个 3D 张量，批量卸载时可以用 CUDA 的异步拷贝，让搬运和计算重叠进行。

1. 算力与带宽的权衡：GPU 内计算速度极快，但显存容量有限且昂贵（HBM 每 GB 成本远高于内存和 SSD），因此把不常用的中间数据放到低成本存储是经济的选择。

2. 搬运开销的量化：搬运开销由数据量和传输带宽决定。PCIe 4.0 x16 的理论带宽约 32GB/s，而 HBM2e 的带宽超过 1TB/s，差距约 30 倍。

3. 异步隐藏延迟：真正高效的卸载采用异步执行——把拷贝操作放入独立的 CUDA stream，与主计算流并行，这样搬运时间被“藏”在计算时间背后。

考虑到篇幅，略去具体公式，但其要点是：显存占用 = 权重 + KV Cache + 注意力中间矩阵卸载策略的目标 = 在增加少量搬运开销的前提下，显著降低峰值显存

CPU卸载

把配料端到储物间（CPU 内存）是帮厨最常用的操作。储物间虽然比操作台远（PCIe 带宽远低于显存带宽），但胜在空间几乎无限——一台服务器动辄几百 GB 内存，而一块 GPU 只有 80GB 显存。对于“暂时不用但很快要取回”的注意力中间结果，CPU 是第一个顺手的卸货点。

CPU 卸载的典型实现是：在注意力计算中，QK^T 矩阵算完之后立即通过 cudaMemcpyAsync 拷贝到主机内存（pinned memory），释放显存空间；等到需要计算 softmax 与 V 的加权和时，再拷贝回来。由于拷贝是异步的，GPU 可以继续计算下一层的 LayerNorm，搬运过程不会让计算核心停摆。

但 CPU 卸载有明确的边界：它适合“低频大块”的数据，不适合“高频小块”的数据。如果每个 token 的解码都要搬运几十次，PCIe 的 32GB/s 带宽就会成为瓶颈——就像帮厨频繁往返储物间，累得气喘吁吁，反而拖慢出菜速度。因此，工程实现中通常设定一个显存水位线：只有显存占用超过阈值时，才触发 CPU 卸载。

SSD卸载

当储物间也堆满了，冷库（SSD）就是最后的救兵。SSD 的容量可以达到数 TB，成本仅为内存的十分之一，但延迟高达微秒甚至毫秒级——比内存慢几个数量级。什么数据值得“冻”进冷库？答案是：长序列推理时的历史注意力中间结果，以及不需要频繁访问的模型权重切片。

SSD 卸载的工程实现与 CPU 卸载类似，但有一个关键差异：延迟太高，无法依靠简单的异步拷贝来隐藏。工程师的应对策略是“预取”（prefetch）——在 GPU 还在计算当前层时，就提前把下一层需要的权重或中间结果从 SSD 读到内存，再把内存中的下一批数据拷入显存。这样，冷库的漫长寻址时间被“前置”到计算时间中，用户在体验上几乎感知不到额外延迟。

SSD 卸载最适合的场景是“超长序列 + 超大模型”的组合。当序列长度达到 32K 甚至 128K 时，注意力矩阵的尺寸会突破 1GB 级别，此时任何单卡都无法承载，SSD 成为唯一能存下全部中间结果的“仓库”。虽然单次搬运较慢，但总吞吐量反而因为批次可以加大而提升。

卸载策略

是每次操作台一满就立刻搬运，还是攒一批再搬？帮厨小张刚开始手忙脚乱——一会儿搬太勤累得半死，一会儿搬太晚操作台已经爆了。后来他总结出一套“容量水位线”策略：操作台使用率到达 70% 时开始卸最占地方的半成品，到 90% 时多线程并行搬，低于 50% 时停止搬运。

这套经验对应到工程中，就是“容量阈值触发”策略。系统会动态监控当前显存使用率，当 chạm到预设阈值（如 75%）时，选择最合适的张量卸载。哪类张量“最合适”？优先级分层给出了答案：释放占比大、重新读取频率低的张量（如长序列的 QK^T 矩阵）优先卸载；体量小但频繁访问的 KV Cache 尽量留在显存。

异步重叠执行是让这一切高效运转的引擎。CUDA 提供了 Stream 机制——把卸载操作放进独立的 stream，与主计算 stream 并行运行。主厨继续炒菜（计算），帮厨同时搬运（传输），两者互不阻塞。在高并发场景中，甚至可以同时进行“取回上一层结果”和“卸下当前层结果”，让数据在传输通道上双向流动。这就是所谓的“流水线卸载”。

效果对比

装了帮厨小张的“搬运策略”之后，餐厅的运营报表焕然一新。最显著的变化是翻台率——即模型的吞吐量（每秒钟处理的请求数）。在长序列、大批次的场景下，吞吐量可以提升 50% 到 100%，接近翻倍。这背后的原理很简单：之前因为显存不够而被迫降低的 batch size，现在可以重新调大，GPU 的计算单元终于被喂饱了。

代价是单次请求的延迟略微上升。搬运数据需要时间，哪怕异步隐藏也只能隐藏一部分。长序列场景下，端到端延迟可能增加 10%-20%。这就像帮厨端菜走的路稍微远了点，主厨出单稍慢，但每桌能同时开火的菜变多了。对在线服务而言，只要延迟在 SLA 范围内，吞吐量的提升就是实打实的收益。

此外，卸载对“显存不足”的容错性带来了质变。以前一个 70B 模型配上 40GB 显存，直接报 OOM；现在配合卸载，模型照常运行，只是多了一些访存开销。这使得推理服务可以部署在更小、更便宜的 GPU 上，硬件成本下降明显。这不只是策略的成功，更是“资源调度思维”在大模型时代的一次胜利。

AI 应用场景

Python 演示

import torchimport torch.nn.functional as F class AttentionWithOffload(torch.nn.Module):    """带注意力卸载的解码层：QK^T中间结果可临时转到CPU/SSD"""     def __init__(self, hidden_dim, num_heads, offload_to="cpu"):        super().__init__()        self.num_heads = num_heads        self.head_dim = hidden_dim // num_heads        self.q_proj = torch.nn.Linear(hidden_dim, hidden_dim)        self.k_proj = torch.nn.Linear(hidden_dim, hidden_dim)        self.v_proj = torch.nn.Linear(hidden_dim, hidden_dim)        self.out_proj = torch.nn.Linear(hidden_dim, hidden_dim)        self.offload_to = offload_to     def forward(self, x, kv_cache=None):        batch, seq_len, _ = x.shape        q = self.q_proj(x).reshape(batch, seq_len, self.num_heads, self.head_dim)        k = self.k_proj(x).reshape(batch, seq_len, self.num_heads, self.head_dim)        v = self.v_proj(x).reshape(batch, seq_len, self.num_heads, self.head_dim)         if kv_cache is not None:            k = torch.cat([kv_cache["k"], k], dim=1)            v = torch.cat([kv_cache["v"], v], dim=1)         q = q.transpose(1, 2)        k = k.transpose(1, 2)        v = v.transpose(1, 2)         scores = torch.matmul(q, k.transpose(-2, -1)) * (self.head_dim ** -0.5)         # 关键：注意力分数矩阵暂时不使用时，可以卸载到CPU/SSD        if self.offload_to == "cpu":            scores_offload = scores.cpu()  # 显存释放，数据搬到内存        elif self.offload_to == "ssd":            torch.save(scores, "/tmp/scores_offload.pt")  # 写入SSD            scores_offload = None        else:            scores_offload = scores         # 模拟异步：计算softmax前先取回数据        if scores_offload is not None and scores_offload.device != scores.device:            scores = scores_offload.to(x.device)        elif scores_offload is None:            scores = torch.load("/tmp/scores_offload.pt", map_location=x.device)         attn = F.softmax(scores, dim=-1)        out = torch.matmul(attn, v)        out = out.transpose(1, 2).reshape(batch, seq_len, -1)        return self.out_proj(out), {"k": k.transpose(1, 2), "v": v.transpose(1, 2)}  if __name__ == "__main__":    torch.manual_seed(0)    device = "cuda" if torch.cuda.is_available() else "cpu"    model = AttentionWithOffload(64, 8, offload_to="cpu").to(device)     x = torch.randn(2, 32, 64).to(device)  # batch=2, seq_len=32    out, cache = model(x)    out2, _ = model(x, kv_cache=cache)     print("输出形状:", out.shape)    print("KV Cache长度:", cache["k"].shape[1])    print("验证通过: 推理流程完整，注意力卸载模块工作正常。")

1. scores_offload = scores.cpu()将注意力分数矩阵从显存拷贝到 CPU 内存，立即释放显存空间。
2. torch.save(scores, "/tmp/scores_offload.pt")模拟卸载到 SSD 文件系统，适合超大批次场景。
3. scores = scores_offload.to(x.device)后续计算前把数据取回 GPU（或直接用 CPU 张量继续计算）。
4. head_dim ** -0.5缩放点积注意力，防止数值过大导致 softmax 饱和。

常见误区

面试题