分离推理阶段的生成计算,降低大模型响应延迟。
Decode-Branch Transformers: Decoupling the Primary Prefill Path from Additional Decode Computation
- 主路径处理输入,分支路径仅在生成时启用,共享权重但不共享缓存。
- 相同参数量下,验证损失更低,生成质量提升,延迟增长可控。
- 适合对响应速度敏感的部署场景,尤其适用于MoE模型优化。
随着大语言模型服务请求增多,累积推理成本相对于训练成本持续上升。典型服务中,提示词预填充并行执行且受计算限制,而自回归生成是串行的且受内存访问限制。传统宽度或深度扩展会同时增加两阶段成本,因为每层均在两个阶段运行,并增大每次生成步骤的权重读取量。本文提出解码分支变压器(Decode-Branch Transformer),将额外的计算专门分配给后续生成预测,同时保持提示词的完整主路径和单一键值缓存。其主路径独立处理提示词并写入KV缓存;解码分支在预填充阶段被跳过,仅从提示词末尾位置开始激活,添加生成计算但不写入状态,不影响主路径。两条路径共享注意力、MLP和输出矩阵,使用独立的词嵌入并通过轻量耦合连接。分组解码可复用已加载的权重块和主路径的KV缓存,因此新增算力不会成比例增加主导的内存流量或解码延迟。在匹配令牌比较中,解码分支在不同架构和数据设置下均实现更低的验证损失。在MoE模型中,主路径与分支的专家扇出可独立调节,用于权衡提示成本、解码成本与预测质量。研究两种专家分配策略:固定预填充或固定解码计算,揭示了由阶段特异性专家分配带来的预填充-解码质量权衡。
原文摘要 · Abstract (English)
As large language models serve ever more requests, cumulative inference cost is growing relative to the one-time cost of training. In typical serving, prompt prefill runs in parallel and is compute-bound, whereas autoregressive decode is sequential and memory-traffic-bound. Conventional width or depth scaling raises both costs together, since every added layer is evaluated in both phases and enlarges the weights read at each decode step. We instead ask whether additional learned computation can be allocated to continuation prediction while preserving prompt-wide primary computation and a single KV cache. We realize this with the Decode-Branch Transformer. Its primary path alone processes the prompt and writes the KV cache; the decode branch is omitted during prefill and activated only from the final prompt position onward, adding continuation computation without writing state or affecting the primary path. The paths share attention, MLP, and output matrices, using separate token embeddings with lightweight coupling. Grouped decode reuses loaded weight tiles and the primary KV cache across both paths, so the added arithmetic does not proportionally increase dominant memory traffic or decode latency. Across matched-token comparisons, Decode-Branch achieves lower validation loss across architectures and data settings. In MoE models, the primary and branch expert fan-outs become independent knobs for trading prompt cost, decode cost, and predictive quality. We study two expert-allocation regimes, holding prefill or decode computation fixed, and expose a prefill-decode-quality trade-off enabled by phase-specific expert allocation.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。