从推理架构的角度，谈谈 Attention Residual 架构一些背后的想法

📌 一句话摘要

本文从 AI Infra 推理架构角度，深入分析 Attention Residual 架构的设计思考，阐述如何通过 Block AttnRes 和 two-phase computation 在几乎不增加延迟和成本的前提下显著增强模型表达能力，并分享了从 Full AttnRes 到 Block AttnRes 的迭代历程与团队协作感悟。

📝 详细摘要

作为月之暗面 AI Infra 团队成员，作者从推理架构角度深入分析 Attention Residual 架构的设计逻辑。文章首先阐述延迟优化与成本优化的差异，指出 latency bound 算子带来的独特挑战；接着详细解析 Block AttnRes 的核心设计——通过 two-phase computation（batched inter-block attention + sequential intra-block attention）实现 IO 优化、容易做 fusion、提升并行度并保持数值精确；然后说明显存空间优化思路（沿 sequence 维度 shard 融入现有 TP 通信路径）；并给出具体性能数据（额外开销小于 2%）。文章还回顾了从 Full AttnRes 到 Block AttnRes 的迭代历程，综合考虑算法表达能力、训练开销、推理延迟和硬件限制最终选择 block_num=8 的设计方案。最终分享了算法与工程团队协作的正向激励如何推动项目完成。

💡 主要观点

- 模型架构设计不仅是算法问题，更是工程实现、硬件约束和算法系统协同设计的综合体现 Attention Residual 架构从一开始就把模型能力、训练开销、推理延迟和硬件特点放在一起思考，是面向真实系统的架构设计典范。

延迟优化和成本优化不等价，latency bound 算子带来独特挑战 topk、小矩阵 gemm 等算子在小 batch decode 场景下对端到端延迟影响明显，且难以通过增加并行度或硬件升级来解决，需要特殊优化思路。

Block AttnRes 通过 two-phase computation 实现近乎无开销的架构增强 Phase 1 的 batched inter-block attention 把'每层读一次'变成'每个 block 读一次'平摊访存；Phase 2 的 online softmax merge 可与 all-reduce、RMSNorm 融合，且保持数值精确等价。

显存优化通过沿 sequence 维度 shard 融入现有 TP 通信路径 单卡 block cache 从 NTd 下降到 N(T/P)d，128K 场景单卡从 15GB 降到 1.9GB；online softmax merge 可自然嵌入 reduce-scatter 路径。

从 Full AttnRes 到 Block AttnRes 是算法效果、硬件限制、训练和推理开销的平衡 Full AttnRes 显存可解但访存开销大（64 层模型 32K 上下文需 7.7TB 访存），结合 two-phase 可优化但训练跨 PP 通信难解决，最终选择 block_num=8 的 Block AttnRes。

算法与工程团队的互相正向激励是项目成功的关键 算法同事的完备理论分析给了工程团队信心，工程团队对方案可行性的积极反馈也推进了算法优化，这种双向激励对复杂系统项目尤为重要。

💬 文章金句

- 一个模型架构最后设计成什么样子，往往不只是算法问题，它同时也反映了团队对工程实现、硬件约束，以及算法和系统如何协同设计的理解。

• 对于 latency bound 问题，我们通常更关心：能不能减少 IO、能不能和其他模块 overlap、能不能把开销摊薄。
• two-phase 并不是一个近似算法，而是通过 online softmax 实现了和原始 attention 完全等价的精确合并。
• 数学上合理的架构，工程上没道理实现不了，如果实现不了，大概率是工程这边哪个地方理解不对。
• 如果系统上实现不了，大概率是算法没真正想清楚。

📊 文章信息

AI 评分：92

来源：青稞AI

作者：青稞AI

分类：人工智能

语言：中文

阅读时间：28 分钟

字数：6993

标签： Attention Residual, 推理架构, 模型优化, two-phase computation, Block AttnRes

阅读完整文章