Attention Residual架构深度解析：如何权衡延迟优化与工程实现？

关键词：Attention Residual、Block AttnRes、推理延迟优化、两阶段计算、软硬协同设计

在月之暗面负责AI Infra工作的这段时间，我常常思考一个问题：一个模型架构的最终形态，真的只关乎算法本身吗？我的看法是，它更像一面镜子，清晰地映照出团队对工程实现、硬件约束，以及算法与系统之间如何协同设计的深刻理解。

今天想聊的Attention Residual架构，便是一个典型的例子。它并非单纯追求“效果更好”，而是从一开始就将模型能力、训练开销、推理延迟和硬件特性纳入统一考量后，诞生的一套面向真实生产环境的系统化设计。

这篇文章，我想从 AI Infra 和推理架构的视角出发（关于训练部分的精彩讨论，我的同事已有详述，此处不再赘述），分享 Attention Residual 背后的一些工程权衡，并尝试将论文中受限于篇幅未能展开的一些细节讲得更透彻一些。

全文将围绕三个核心部分展开：

1. 聚焦推理性能，特别是延迟优化。与成本优化不同，高OTPS（Output Tokens Per Second）场景下的延迟优化会带来额外的技术挑战。Block AttnRes 的设计目标，正是在这种严苛约束下，力求以近乎零增加的延迟和成本开销，显著增强残差结构的表达能力。（论文中提到的推理开销低于2%是一个覆盖多种负载的严谨口径；而在许多常见场景下，这部分额外开销实际上低得多，甚至可视为接近免费。）

2. 回顾架构演进，从 Full AttnRes 到 Block AttnRes。当前的Block版本，是算法效果、硬件限制、训练成本与推理开销共同作用下找到的平衡点。但长远来看，随着硬件演进，我们始终希望向几乎无约束的Full AttnRes迈进。

3. 分享我对跨团队协作的新感悟。这项工作的完成，离不开算法与工程同事之间深入的理解与正向激励。

一、Block AttnRes: 极致延迟和成本优化

本节将先阐明大模型推理中延迟优化与成本优化的差异，以及这种差异如何影响我们对AttnRes架构的设计选择。进而说明，当前的Block AttnRes是如何在相比标准残差连接显著增强表达能力的同时，将额外的延迟与成本开销压制到极低水平的。

1.1 延迟和成本优化的差异

通常，优化延迟与优化成本是高度相关的。例如通过算子融合优化I/O，或优化Flash Attention这类核心计算，往往能同时降低延迟和成本。

但两者并非完全等价，尤其在当下的大模型推理场景中，差异愈发明显。典型的两种情况如下：

1. 延迟低，但成本不一定低：有些方案主要通过增加并行度、缩短关键路径来降低端到端延迟。这对用户体验有益，但模型的总访存量和计算量未必真正下降，因此成本改善可能有限。ScMoE这类方案便接近此方向。
2. 成本低，但延迟不一定低：另一类典型是latency-bound算子。例如topk（稀疏注意力或MoE路由中的topk计算）、一些小尺寸矩阵乘法（如HyperConnection中的 tf32_hc_prenorm_gemm 或路由门计算）等。这些算子在大批次（batch）下可能有不错的带宽和算力利用率，从吞吐和成本角度看问题不大；但在小批次解码（decode）场景下，它们对端到端延迟的影响却相当显著。

因此，若仅从成本或平均吞吐角度评估，很容易低估某些结构在真实推理中的延迟代价。尤其是在解码场景下，更需要警惕关键路径上是否引入了新的latency-bound开销，因为这些开销虽然占比未必高，却常对端到端延迟产生明显影响。对于成本优化，已有诸多成熟方法，如通过增大批次、优化GPU流水线来提升计算与带宽利用率。但对于延迟优化，尤其是解码小批次场景，问题则更为棘手。

正如我之前文章提到的，Latency-bound算子通常带来两大挑战：

1. 难以通过增加并行度解决：像FFN或Attention中的许多memory-bound计算，尚可通过增加并行度换取延迟；但topk和小矩阵乘法这类计算本身并行度不高。虽然可采用Split-K等方法增加并行性，但由于合并开销大，最终收益有限。
2. 硬件升级不一定带来收益：新一代硬件如Blackwell，算力和带宽可能倍增，但这不意味着这些算子的延迟会同步下降。其瓶颈往往不在理论算力，而在启动、同步或访存上。我们甚至观察到某些topk实现在更新的GPU上反而更慢。

正因如此，这类开销的优化思路与常规的compute-bound/memory-bound算子不同。对于latency-bound问题，我们通常更关心以下几点：

1. 能否减少I/O：最直接的方式是通过融合消除不必要的HBM访存，因为很多时候昂贵的不是计算，而是访存延迟。
2. 能否与其他模块重叠：如果一个小算子本身无法占满GPU，最理想的情况不是单独加速它，而是让它与其他计算并行，从而不落在关键路径上。
3. 能否摊薄开销：例如配合投机采样（speculative decoding），某些小算子的延迟在批次增大时增长不明显，其边际成本便有机会被摊薄。

后续我们将看到，这些认知直接塑造了AttnRes的最终架构。Block AttnRes的设计初衷并非单纯“节省计算”，而是围绕一个明确目标展开：在尽量不引入额外计算开销的前提下，最大化残差结构的表达能力。其后的Block Residual与两阶段计算（two-phase computation），本质都是为实现此目标服务的。

1.2 基于两阶段计算的性能优化

如果朴素地（naively）实现Block AttnRes，每一层都需要读取所有先前的块表示（block representation），然后计算一次注意力残差。这会导致两个问题：

1. 对于预填充（prefill）负载：虽然通过Block AttnRes将注意力计算的key和value压缩到了8个块，但query Q 的访存仍有较大开销。
2. 对于解码（decode）负载：除了访存开销，由于注意力计算本身相对复杂，受限于共享内存（shared memory）空间等因素，较难与前后的All-Reduce、RMSNorm等计算有效融合，从而难以降低高OTPS场景下的延迟。

Block AttnRes的一个关键设计是：每一层的注意力查询（attention query）是一个与当前隐藏状态（hidden state）解耦的可学习参数。这一设计从算法上看很轻量，但从系统角度却至关重要，因为它意味着同一块（block）内的所有查询都可以提前提取，并统一执行一次批处理（batched）的块间注意力（inter-block attention）。

基于此，我们最终将计算拆分为两个阶段：

1. 阶段一：批处理块间注意力（Phase 1: batched inter-block attention） 对一个块中的所有层，统一与之前所有的块表示进行一次批量注意力计算，得到每一层对应的块间部分结果及softmax统计量（如lse、max等）。
2. 阶段二：顺序块内注意力 + 在线softmax合并（Phase 2: sequential intra-block attention + online softmax merge） 在块内仍按层顺序推进，用不断更新的部分和（partial sum）计算块内部分，然后通过在线softmax（online softmax）将两部分结果精确合并。

现在我们来分析两阶段计算带来的好处：

1. 优化I/O：通过阶段一的批量计算，原本每一层都需重复读取的历史块表示被平摊了，本质上将“每层读一次”变为“每个块读一次”。
2. 更易融合：阶段二中的在线softmax合并是逐元素（elementwise）计算，这使其能很自然地与All-Reduce、RMSNorm等算子融合，进一步减少额外I/O。
3. 提升并行度：阶段一可以与块内第一个层的部分计算重叠，因此即使它本身引入了额外工作，也并非完全落在关键路径上。
4. 保持数值精确：这一点非常重要：两阶段计算并非近似算法，而是通过在线softmax实现了与原始注意力完全等价的精确合并。这决定了它不仅仅是一个“工程技巧”，而是可以作为正式架构的一部分。需要说明的是，这个两阶段方法已在真实模型上通过了正确性验证。因此，Block AttnRes的推理延迟开销并非以牺牲精度为代价，而是通过更合理的计算组织方式实现的。这种对性能与精确度的极致追求，正是现代人工智能与系统设计中软硬协同的核心体现。

1.3 显存空间优化

关于Attention Residual的讨论中，许多人的第一印象是：由于需要存储更多层的输出结果，显存会不会爆炸？这里我们解释在推理阶段为何显存开销可以忽略不计；甚至在后续章节，我们会说明即使对于Full Attention，这个显存开销也完全可接受。

对于Block AttnRes，如果直接存储完整的 hidden_states 表示，那么在长上下文下显存压力会非常大。例如128K token、8个块、hidden_dim=7168的配置下，仅块表示就可能需要约15GB显存。这显然是不可忽略的开销。

整体优化思路是沿序列（sequence）维度进行分片（shard），将块表示分摊到张量并行（Tensor Parallel, TP）的各个设备上。这样每张卡只需保留本地序列分片对应的块缓存，而非整段序列的完整缓存。

1. 显存占用线性下降：如果有 N_TP 张TP卡，则单卡上的块缓存从 O(N) 下降到 O(N/N_TP)。对于上述128K场景，单卡显存可从约15GB降至1.9GB左右。
2. 能融入现有TP通信路径：阶段二的在线softmax合并本质上是逐元素的，因此可以自然地嵌入标准的TP All-Reduce路径中，具体而言是在reduce-scatter之后进行本地合并，完成RMSNorm后再通过all-gather恢复结果。这使其更容易与已有的融合内核（fused kernel）体系衔接。

更进一步，如果结合分块预填充（chunked prefill），例如每次处理32K的块大小（已足以打满稀疏模型MoE的计算），那么这部分显存开销还可以继续下降到远低于1GB的量级。

1.4 性能分析

此处我们对AttnRes的推理开销进行更具体的分析，并主要与标准残差连接（standard residual connection）对比。为避免 layer、block 等术语混淆，先做如下约定：

1. Transformer Decoder Block：指一个完整的Transformer块，包含一次Attention和一次FFN。
2. layer：沿用论文记法，特指一次AttnRes执行的位置。由于AttnRes在Attention前和FFN前各执行一次，因此 1个Transformer Decoder Block = 2个layer；一个拥有64个 Transformer Decoder Block 的模型，在论文记法中对应128个 layer。
3. AttnRes block：两阶段计算中的分块单位，由若干个 layer 组成。
4. D：隐藏层维度（hidden dimension），此处默认使用7168进行计算。

先看基线（baseline）。对于标准残差连接，在实现中通常直接与后续的PreNorm融合，即常见的 FusedAddRMSNorm。该算子的输入是 hidden_states 和 residual，输出是归一化的新 hidden_states 以及更新后的 residual，因此总访存量大致可记为 4D。

再看Block AttnRes。其计算分为两个阶段。

• 阶段二：在线softmax合并 此步骤的输入是 hidden_states、prefix_sum 和 partial_attention，输出是新的 hidden_states 和 prefix_sum，因此总访存为 5D。也就是说，仅看这一步，相比标准残差连接只多了 D 的访存。
• 阶段一：块间注意力 记AttnRes的块数为 N，每个AttnRes块内包含的 layer 数为 S。在一个典型配置中，若模型有64个 Transformer Decoder Block，则对应128个 layer；若将AttnRes按每16个 layer 一组分块，则有 S=16，整个模型对应 N=8 个AttnRes块。对于最后一个AttnRes块，阶段一需要读取前面8个块表示。若将key和value的读取都计入，此阶段最坏情况下的读取开销近似为 16D。但这仅是最后一个块的情况，并非平均成本，平均约为 8D。更重要的是平摊后的开销：因为这 8D 的读取会被当前块内的16个layer共同分摊，所以平均到每个layer上，额外读取仅增加约 0.5D。此外，阶段一还需为当前块内的每个layer写出一个 partial_attention 结果，因此总共会有 16D 的写回，平摊到每个layer即 D。所以从平摊角度看，阶段一为每个layer带来的额外访存约为 1.5D。

将两部分合起来看，Block AttnRes每个layer相比标准残差连接增加的主要访存开销大致是：
(Phase2: 1D) + (Phase1 per layer: 1.5D) = 2.5D

这个增量相对较小，大致可理解为仅多了一次很轻的激活（activation）计算。这也是为何我们说，在许多实际负载下，Block AttnRes带来的推理开销非常低。实际测试中，其带来的额外开销类似于RMSNorm：即使在大于50 OTPS的输出场景下，这部分解码延迟通常也小于2%。

例如在 batch=128、64个 Transformer Decoder Block 的配置下，端到端解码延迟增加不足0.5ms；若再考虑多租户处理（MTP），此开销还会被进一步平摊。对于32K预填充这类场景，由于注意力等计算占比很高，这部分额外开销通常可忽略不计，甚至在性能测试的抖动中可能出现基线延迟略差的情况。

二、从 Full AttnRes 到 Block AttnRes 的迭代历程

若仅从算法表达能力角度看，Full AttnRes 其实是最直接、也最“正确”的版本。它让每一层都能直接对所有历史层输出进行选择性聚合，形式最完整，也最贴近“将注意力迁移到深度维度”的原始直觉。

本节将简述我们对Full Attention的优化尝试，并解释为何最终综合考虑硬件限制、训练与推理开销、算法效果后，选择了Block AttnRes。

以下分析历程基于我个人所知信息总结，更多是从自身视角出发，不一定代表全貌。

2.1 算法的提出与初步验证

首先是苏剑林老师提出了非常通用的数学形式，即通用的全注意力（general full attention），并给出了完善的理论分析，包括各种数学推导。随后，苏老师及我们算法同事的初步实验也真实证明了其在效果上的优势。这给了我们Infra同事极强的信心，其重要性我将在后续的算法工程协作部分详述。

2.2 Full Attention 显存问题的解决

与大家一样，我看到这个架构后的第一个顾虑就是显存占用。64层模型需存储128份隐藏状态。以32K上下文为例，就需约60GB存储空间。但我们很快发现此问题可以解决，因为完全可以按照序列维度进行切分。例如8卡并行，则仅需约7.5GB显存。

2.3 Full Attention 访存问题的解决

在优化显存的过程中，我反而意识到真正棘手的问题并非显存，而是访存（memory access）。因为访存开销是层数的平方量级。一个64层的大模型，总访存量约为 O(L^2 * D) 字节。假设32K上下文，这对应约7.7 TB的访存量。即在H200这类GPU上，将引入约2秒的额外计算时间；即使按上述将序列切分到8卡，也需要约0.3秒的计算时间，这已相当于激活参数量为40B模型的计算时间。

为解决此问题，才构思出论文中的两阶段计算算法（Full版本的两阶段计算详见论文附录）。其核心是通过批处理（batching）方式来平摊访存开销。通过批处理，可将每层的访存量从 O(L * D) 优化到 O(D * sqrt(L))，即从 O(128D) 优化到 O(22.6D)，将访存量降至原先的约六分之一。结合一些I/O融合及计算并行，此延迟变得可以接受。因此我们从推理Infra角度判定：Full AttnRes是可行的。

2.4 训练开销与 Block AttnRes 算法的确定

同时，训练框架的同事也在极致优化显存、通信等。但后续发现，训练中跨流水线并行（PP）的通信问题仍较难解决（论文中有提及，此处不展开细节）。因此，大家综合考虑后确定了Block AttnRes结构。

最终的块大小设置，综合权衡了训练效率、算法效果与推理效率。块大小=8的配置，既能满足训练效率，又能获得大部分算法收益，同时也与上文 O(D * sqrt(L)) 对应的最佳块大小（约11.3）较为接近。而从推理效率角度，当时新增了一个关键考量因素：公司流量激增，使得延迟优化与成本优化的压力空前巨大。因此，在最终定版时，我们的设计目标就是做到几乎无性能开销。

说了这么多，主要是想阐明我们在设计Block AttnRes过程中的整体思考历程，以及最终方案的权衡取舍。但我们从未放弃表达能力更强的Full AttnRes。或许未来某天，拥有更强大的硬件时，我们将能看到Full AttnRes的上线。

三、团队协作

这个项目历时颇久，最终技术报告只体现了结果，但中间大家付出了诸多尝试与努力，甚至有些尝试并不顺利。项目的推进极度依赖整个团队的协作，既有算法同事的理论分析、方案制定、扩展（scaling）实验与效果排查，也有工程同事在优化显存与通信、保障基建稳定性、优化推理性能与支持各类评测等方面的巨大贡献，包括我自己也花费了大量时间进行性能分析与对分验证。

在此过程中，我感受最深的是团队成员间的正向激励。算法与工程同事之间会积极给予对方正反馈与信心，这对项目的最终完成至关重要。

剑林老师多次提及，当时我设计完两阶段算法后，直接拍板“Full AttnRes推理能搞定”，这给了算法同事极大的鼓舞和推动作用。但反过来，正是他最初对整个方案的完备理论分析，以及算法同事（张宇 @yzhangcs，我们的天才少年Nathan，剑林老师，以及其他多位我不知其知乎账号的同事）的实验结果，给了我们Infra同事极大的信心。我正是在认真读完剑林老师的理论分析后深感信服，才开始非常认真地投入性能优化的。因此，这是双向的激励。

公司内部有一个梗图，在遇到算法或架构问题时，有同事会拿出来，其中两句话我认为很好地体现了我们这种互相理解与鼓励：

1. 第一句是从我的Infra角度说的。背景是剑林老师常有各种算法想法，但他总担心系统开销过大、难以实现，常来与我讨论。我有一次很认真地回复：“数学上合理的架构，工程上没道理实现不了；如果实现不了，大概率是工程这边哪里理解不对。”
2. 第二句是张宇从算法角度说的：“如果系统上实现不了，大概率是算法没真正想清楚。”

我颇为欣赏这种大家互相理解、支持的氛围，也坚信正向激励能更好地推动项目完成。更多此类高质量的技术讨论与协作，也欢迎在 云栈社区 这样的开发者社区中持续发生。

事实上，最近我并非处于项目完成后开香槟的快乐状态，而是始终有一种压力感。这种压力并非源于人或工作本身（我们团队氛围很好，我工作得非常开心），更多的是来自当前面临的技术挑战：大家已逐步感受到一个AI临界点（tipping point）的到来，而我们面临的技术战役也越来越艰巨。

• 例如，流量的指数级增长、算法层面探索的日益深入、参数量与模型架构持续探索带来的未知性等，都带来了巨大的技术挑战。
• 再如，我负责模型架构时，常会担心自己哪里稍有算错，就可能导致巨大的成本开销。

每当与同事聊起这些，彼此间的鼓励对我都颇有助益。