ai-infra-interview-305

第 158 题：长上下文训练（100K+ tokens）的显存优化？Ring Attention？

题目

长上下文训练（100K+ tokens）的显存优化？Ring Attention？

---

完整讲解

一、长上下文显存瓶颈

序列长度 L 时，注意力与 KV cache 的显存约为 O(L²) 或 O(L)（依实现）；100K+ token 时单机单卡难以放下完整 attention 与 KV。优化方向：稀疏/近似 attention（只算部分位置）、分块与通信（把序列切块分布到多卡，通过通信拼出 attention）、重计算（用激活重计算换显存）、外存换入换出（KV 或中间结果放 CPU/NVMe）。

二、Ring Attention 思路

Ring Attention：把序列在长度维切块，每块分布到不同 rank，形成「环」；计算 attention 时，每 rank 持有当前块与部分其它块（通过沿环传递逐步拿到全部块），在本地做局部 attention 或与传递来的 key/value 做计算，再传递下一段。这样每 rank 只需 O(L/P) 的显存（P 为 rank 数），总显存与通信在长序列下可控；通信是沿环的流水线式，可与计算重叠。变体有 Ring Attention、Blockwise Transformer 等，核心都是「序列分块 + 通信拼全」以突破单卡显存。

三、其它手段

FlashAttention 等省显存 attention kernel；线性 attention 或 state space 把复杂度降到 O(L)；结合 offload 把部分 KV 放 CPU。面试可强调：长上下文 = 显存 O(L) 或 O(L²)；Ring Attention = 序列分块 + 环上传递，显存与通信可扩展。

---

面试要点

• 长上下文显存 O(L) 或 O(L²)；100K+ 需稀疏/分块/重计算/offload 等手段。
• Ring Attention：序列按长度切块分布到多 rank，沿环传递块以拼全 attention，每 rank 显存 O(L/P)；通信可与计算重叠。
• 配合 FlashAttention、线性 attention、offload 等进一步省显存。

---

记忆要点

1. 长上下文瓶颈 = 显存与 attention 规模；Ring Attention = 序列分块 + 环上传递。
2. 每 rank 显存 O(L/P)；通信流水线、可重叠。
3. 可与 FlashAttention、线性 attention、offload 组合。

返回模块

返回总览