CoDec: Prefix-Shared Decoding Kernel for LLMs

发表时间: 2026-06 · arXiv:2505.17694 (SIGMOD 2026)

原文: https://arxiv.org/abs/2505.17694

Zhibin Wang, Rui Ning, Chao Fang, Zhonghui Zhang, Xi Lin, Shaobo Ma, Mo Zhou, Xue Li, Zhongfeng Wang, Chengying Huan, Rong Gu, Kun Yang, Guihai Chen, Sheng Zhong, Chen Tian
State Key Laboratory for Novel Software Technology, Nanjing University, China; Alibaba Group, China

速读

一句话结论 本文提出了针对大语言模型解码阶段的前缀共享注意力算子 CoDec，通过在 GPU 共享显存中合并共享前缀的 KV Cache 访存并引入树状规约与负载均衡机制，相比 FlashDecoding 实现了平均 1.9 倍的算子提速与 120.9 倍的全局显存访问量骤降。

要解决什么问题 在大语言模型推理中，随着上下文长度增加，解码阶段的注意力计算成为核心瓶颈。解码阶段需要自回归地逐个生成 token，这种计算模式的算术强度极低，导致其严重受限于 GPU 的全局显存带宽。在文档问答、思维树、少样本提示等真实场景中，多个请求通常会共享相同的前缀 prompt。现有的主流注意力优化算子（如 FlashAttention 和 FlashDecoding）针对的是规则的 4D 张量，当它们处理前缀共享场景时，会将共享相同前缀的查询分配给不同的计算单元独立处理。这种机制导致 GPU 必须从全局显存中为每个请求重复读取同一份共享前缀的 KV Cache，造成极大的访存带宽浪费。此外，前缀共享的 KV Cache 在逻辑上呈现为由多个 3D 张量组成的树状结构，不同节点包含的前缀长度和关联的查询数量差异巨大，如果直接将这种不规则的树状结构映射到硬件计算上，会导致严重的线程块负载不均衡以及复杂的局部结果依赖管理问题。

怎么做的 CoDec 的核心思路是将共享前缀的注意力计算转化为数据访问优化问题，通过在 GPU 的共享显存中合并多个请求对同一前缀 KV 块的读取，彻底消除冗余的全局显存访存。其关键设计由三个部件构成。首先是面向计算的树状 KV Cache 管理，它在全局显存中维护一棵 KV 块构成的树，并建立查询到对应 KV 节点的轻量级索引，使得算子能在内核启动时将共享同一前缀的查询张量直接聚合到共享显存中。其次是块内与块间并行的注意力核，在线程块内部，CoDec 抽象了局部注意力计算（PAC）和局部输出规约（POR）两个原语，PAC 负责在共享显存中计算聚合后的查询与对应 KV 节点的结果，POR 负责利用在线 Softmax 状态合并同一查询的不同局部输出；在线程块间，由于规约操作满足结合律和交换律，CoDec 将树状规约转换为多个独立的并行规约链，最大化并发度。最后是应对不规则树状结构的负载均衡机制，由于单个子任务的耗时既不纯粹受限于 IO 也不纯粹受限于计算，CoDec 引入了一个基于硬件微基准测试的代价估计器 $C_{est}(n_q, n)$ 来预测执行时间。任务划分器利用该估计器，以最小化最慢线程块耗时为目标进行调度：

$$ \arg\min_{b_k, \mathbf{A}} \max_{i=1}^{m} \left( \sum_{j=1}^{t} \mathbf{C}[j] * \mathbf{A}[i, j] \right) $$


其中 $\mathbf{C}[j]$ 是子任务代价，$\mathbf{A}$ 是分配矩阵。系统通过二分查找确定代价上下界，并贪心搜索最优的 KV 维度切分粒度 $b_k$，将切分后的子任务调度到不同线程块中执行，从而在避免细粒度切分开销的同时消除长尾阻塞。

效果如何 实验主要在单张 NVIDIA A100 GPU 上基于 Qwen3-4B 模型展开，测试了包含 LooGLE 真实长文本数据集及多种受控合成前缀树负载，并验证了方法在 MHA、MQA、GQA 等变体下的通用性。在解码阶段注意力算子层面，对比代表长上下文解码 SOTA 路线的 FlashDecoding，CoDec 实现了平均 1.9 倍的执行速度提升，以及高达 120.9 倍的全局显存访问量降低；在端到端耗时层面，对比主流推理框架 vLLM，CoDec 取得了平均 3.8 倍的提速。此外，在共享前缀批处理解码的吞吐量对比中，CoDec 凭借全局视角的负载均衡设计，始终优于代表多级级联注意力路线的 FlashInfer。作者还在 H800、RTX 4090、A6000 等多款 GPU 上进行了测试，证明在显存带宽越受限的硬件上（如 A6000），CoDec 规避访存瓶颈带来的性能优势越显著（提速达 15 倍）。该方法的代价在于任务划分模块需要在 CPU 上运行，随着批处理大小增加会引入几十毫秒的计算开销，但在实际部署中可以通过在多个解码步之间复用同一套划分计划来将此开销摊薄，使其相对于注意力执行时间的占比保持在极低水平。

主要贡献

大型语言模型（LLMs）在问答、规划、代码生成和推荐系统等多种任务中展现了显著的性能，但推理效率仍然是LLM部署的关键，直接影响用户体验和运营成本。提高LLM推理效率的一个极具前景的方向是前缀共享（Prefix-sharing），这基于许多提示（Prompt）共享相同前缀的观察，在文档问答、思维树、投机解码和少样本提示中非常常见。在生成提示词Token的KV缓存的预填充（Prefill）阶段，前缀共享可用于减少不同请求之间共享前缀的KV缓存生成的重复计算和内存消耗。

随着上下文长度的增加，解码（Decode）阶段，特别是注意力计算，成为LLM推理中的关键瓶颈。当在Llama-3.1-8B上运行10万个具有128个输出Token的提示时，解码阶段的注意力计算占总时间的90%。与预填充阶段不同，解码阶段基于生成的KV缓存自回归地生成输出Token，因此需要大量访问KV缓存的内存，表现出并行性不足和内存受限的模式。现有的最先进方法（如FlashAttention和FlashDecoding）通过利用共享内存和增加并行性实现了显著加速。然而，当在前缀共享场景中采用这些技术时，共享相同前缀的查询会被独立的计算单元单独处理，这不可避免地导致了访问共享KV缓存时的重复内存事务。

图1. 在解码阶段的注意力机制中利用前缀共享的动机。
图1. 在解码阶段的注意力机制中利用前缀共享的动机。

本文探索了在解码阶段利用前缀共享的潜力，具体通过优化不同请求间共享KV缓存的内存访问模式，直接解决冗余内存事务问题。将常规的4D张量注意力计算转变为不规则的共享前缀注意力计算面临两大挑战：一是组织共享前缀注意力计算中的复杂依赖关系，需要协调查询张量与KV缓存树的访问，并处理树结构中每个查询的归约操作；二是平衡不规则前缀共享注意力计算的工作负载，因为查询数量和前缀长度的巨大变化导致工作负载高度不规则，且不同的张量形状会导致不同的计算或内存瓶颈。

为应对上述挑战，本文提出了CoDec，这是一个专用的前缀共享注意力算子，用于在解码阶段合并不同请求的共享前缀注意力计算的内存访问。本文的主要贡献如下：
1. 共享前缀注意力内核：提出了一种新颖的共享前缀注意力内核，该内核优化了共享内存和全局内存之间的内存层次结构，并利用了块内和块间并行性。抽象了部分注意力计算（PAC）和部分输出归约（POR）两个基本原语，并在此基础上提出了块间计算任务执行器和专用的基于树的归约机制。
2. 工作负载平衡机制：整合了全面的工作负载平衡机制，包含基于分析的成本估算器、任务划分器和调度器。通过估算成本、智能划分任务和调度执行，在不求助于昂贵的细粒度划分的情况下平衡工作负载。
3. 显著的性能提升：与最先进的FlashDecoding内核相比，CoDec在解码阶段的注意力计算中实现了平均1.9倍的加速和120.9倍的全局内存访问减少，与vLLM相比，每个输出Token的端到端时间（TPOT）实现了3.8倍的加速。

背景知识与设计原则

LLM推理阶段的特征。Transformer架构的推理过程包含预填充（Prefill）和解码（Decode）两个主要阶段。预填充阶段同时处理输入Token序列并缓存对应的KV张量，由于并发处理大量Token，该阶段是计算密集型且受计算限制的。解码阶段每步仅处理一个Token，计算需求显著降低，导致该阶段受内存限制。随着上下文长度增加，解码阶段的注意力计算成为LLM推理的关键瓶颈。

注意力机制的计算逻辑。自注意力算子计算每对Token之间的注意力分数，并据此更新Token的嵌入。多头注意力接收查询（$Q \in \mathbb{R}^{bs \times n_q \times h \times d}$）、键（$K \in \mathbb{R}^{bs \times n \times h \times d}$）和值（$V \in \mathbb{R}^{bs \times n \times h \times d}$）张量。其核心公式为 $O = \operatorname{softmax}\left(\frac{Q K^T}{\sqrt{d}}\right)V$。在解码阶段，每个请求仅处理一个Token（$n_q = 1$），导致并行性不足和算术强度低，使其成为内存受限的操作。

GPU架构与现有内核的局限。GPU逻辑上由多个包含张量核心的块（Block）组成，配备快速但有限的共享内存（片上内存），以及大但慢的全局内存。FlashAttention【7，FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning + 2024 + ICLR + https://arxiv.org/abs/2307.08691】及其后续版本FlashDecoding【8 ，Flash-Decoding for long-context inference + 2023 + Blog + https://crfm.stanford.edu/2023/10/12/flashdecoding.html】通过在批次、头部、查询序列和KV序列维度上分解注意力操作，使其适配共享内存并暴露多路并行性。然而，这些内核专为规则的4D张量设计 。

前缀共享的普遍性与挑战。前缀共享在文档问答、工具使用、少样本提示、自洽性、思维树和投机解码中广泛存在。现有研究在预填充阶段通过树状结构管理KV缓存以实现重用，但在访问KV缓存时，系统仍假定键和值张量为逻辑4D张量结构，从而在解码阶段遭受重复的全局内存访问。这引出了前缀共享解码的两个关键挑战：不规则KV缓存森林上的内核设计，以及高度倾斜的节点工作负载下的工作负载不平衡。

图2. LLM推理的背景知识、GPU架构以及现有的注意力内核。
图2. LLM推理的背景知识、GPU架构以及现有的注意力内核。

CoDec系统架构设计原则。为实现高效的前缀共享解码，内核需要满足IO高效（利用共享内存合并KV缓存访问）和工作负载平衡（在不同块之间平衡不规则的树状工作负载）。CoDec的架构包括：内存管理器（将KV缓存实体化为全局内存中的张量树并维护轻量级索引）、内核执行器（暴露块内和块间并行性，利用部分注意力计算和部分输出归约原语）、成本估算器（预测特定形状的部分注意力计算的执行时间），以及任务划分器和调度器（利用全局视角划分和分配任务以平衡工作负载）。

图3. CoDec系统概览。
图3. CoDec系统概览。

方法细节

KV缓存管理

在FlashAttention/FlashDecoding中，KV张量遵循规则的4D布局，而前缀共享解码中，KV缓存变成了每个前缀3D张量的森林。这使得KV/查询索引复杂化，增加了任务划分的复杂性，并需要树状结构的归约来合并部分结果。

形式化模型与符号定义。令 $\mathcal{R} = \{r_1, \ldots, r_B\}$ 为请求批次。每个请求 $r$ 具有一个当前查询向量 $q_r \in \mathbb{R}^{d_k}$，以及一个（可能共享的）前缀Token序列 $P(r) = (t_1, \dots, t_{\ell(r)})$。系统维护一个KV缓存节点的森林（带有一个虚拟根）$\mathcal{F} = (N, \mathcal{E})$，其中每个节点 $n \in N$ 存储张量 $K_n \in \mathbb{R}^{|n| \times d_k}, V_n \in \mathbb{R}^{|n| \times d_v}$，以及一个有序的Token索引集 $\operatorname{Tok}(n) = \{1, \dots, |n|\}$。边 $(p \to c) \in \mathcal{E}$ 编码了 $p$ 是 $c$ 的前缀。所有查询被堆叠成 $Q \in \mathbb{R}^{B \times d_k}$。系统定义了全局Token地址空间 $\Omega$ 和逻辑展平映射 $\kappa$，但在物理上从不整体实例化逻辑张量 $K, V$。

基于树的KV缓存管理结构。如图4所示，KV缓存被管理为张量树，树中的每个节点代表一个KV缓存块，父子节点之间的边代表两个块之间的前缀关系。所有请求的查询被合并为一个查询张量，每行对应一个单独请求的查询。由于实践中可能有多个不同的前缀被不同请求共享，系统引入了一个虚拟根节点来连接所有请求的前缀。这种虚拟根节点允许将不同前缀的请求批处理在一起，甚至支持非前缀共享的解码。系统还维护了记录KV缓存与查询张量之间双射映射的索引结构。

为部分注意力计算准备张量。共享维护在同一个KV缓存节点中的前缀的查询，实际上构成了一个部分注意力计算。因此，对于每个KV缓存节点，系统维护与该节点共享前缀的请求的查询集合。这使得内核能够聚合这些查询，在由查询形成的张量和对应的KV缓存节点之间进行部分注意力计算。系统不是在全局内存中准备查询张量，而是在部分注意力内核启动时，在线程块共享内存中聚合查询张量 $\mathcal{Q}^{(n)}$，从而减少内存开销。

每个查询的部分结果索引与流式Softmax。Softmax操作本质上在注意力分数张量的每一行（即查询级别）全局运行。系统需要为每个查询维护其前缀由哪些KV缓存节点构成的记录，以便正确索引和检索部分注意力计算结果。为了在不实例化全局分数向量的情况下计算每个查询 $r$ 的注意力，系统在节点（块）上进行流式传输，动态累加Softmax的分子和分母。当访问节点 $n$ 时，计算缩放分数 $s_{r,n} = \frac{1}{\sqrt{d_k}} Q[r, :] K_n^\top \in \mathbb{R}^{|n|}$，并应用可见性掩码得到 $\tilde{s}_{r,n}$。随后在节点内进行局部稳定化，利用节点内最大值 $m_{r,n} = \max_j \tilde{s}_{r,n}[j]$，计算分母贡献 $d_{r,n} = \sum_j e^{\tilde{s}_{r,n}[j] - m_{r,n}}$ 和分子向量贡献 $u_{r,n} = \sum_j e^{\tilde{s}_{r,n}[j] - m_{r,n}} V_n[j, :] \in \mathbb{R}^{d_v}$。为避免合并时的数值溢出/下溢，维护每个查询的全局累加器 $M[r], D[r], U[r, :]$。当合并新节点时，使用共同的对数求和指数参考 $\hat{M} = \max(M[r], m_{r,n})$，并在同一指数框架下更新旧贡献和新贡献。

图4. KV缓存森林。
图4. KV缓存森林。

块内内核原语

在KV缓存管理的基础上，CoDec抽象了两个必要的块内内核原语：部分注意力计算（PAC）内核和部分输出归约（POR）内核。这两个原语被设计在配置有片上共享内存的线程块上执行，显著减少了全局内存访问开销。

部分注意力计算（PAC）内核。该内核负责在查询子张量（$Q \in \mathbb{R}^{n_q \times d}$）及其对应的KV缓存子张量（$K, V \in \mathbb{R}^{n \times d}$）之间执行注意力计算。与普通的注意力操作完全相同，唯一的区别在于PAC内核中的输入查询张量来源于多个请求。内核进一步划分部分注意力计算以利用共享内存，并顺序处理划分的计算，从而在支持更大工作负载的同时保持内存效率。为了进一步减少内存访问开销，PAC内核优化了数据加载模式：不再为每个查询从全局内存加载键和值张量，而是将键和值张量加载到共享内存一次，然后供多个查询重用。这种优化显著减少了内存访问开销，特别是在GQA比率较高时，同时减少了Tensor Core内部因查询张量填充带来的低效。

部分输出归约（POR）内核。由于一个查询及其对应KV缓存之间的操作被划分为多个部分注意力计算，系统需要合并这些结果以获得查询的最终输出。POR内核是一个二元归约操作，合并来自同一查询集的两个不同KV缓存节点的部分注意力计算结果 $O_1$ 和 $O_2$。POR内核接收 $O_1, O_2$ 及其对应的最大注意力分数 $m_1, m_2$ 和指数注意力分数之和 $s_1, s_2$ 作为输入。首先计算两个结果的最大注意力分数 $m = \max(m_1, m_2)$ 和分数之和 $s = s_1 e^{m_1-m} + s_2 e^{m_2-m}$。随后，重新归一化 $O_1$ 和 $O_2$ 并将它们合并为最终输出 $O = \frac{O_1 \cdot s_1 \cdot e^{m_1-m} + O_2 \cdot s_2 \cdot e^{m_2-m}}{s}$。由于输出 $O$ 的大小很容易适配共享内存，POR内核默认在共享内存中执行。

块间启动与树归约

在块内内核原语之上，CoDec开发了块间内核执行器，负责以并行方式执行块内内核原语，为一批请求进行前缀共享注意力计算。

顺序启动与并行执行流程。执行过程主要包含两个步骤：首先是为每个KV缓存节点启动PAC内核，其次是执行树归约以合并每个查询的结果。由于PAC内核的计算是独立的，系统利用极度并行性，并行地为每个KV缓存节点启动PAC内核。随后执行同步操作，确保在进行树归约操作之前准备好所有部分注意力计算结果。

树归约操作的并行化机制。部分注意力计算结果存储在与KV缓存管理系统相同的树结构中。由于归约操作满足结合律（$POR(O_1, O_2) = POR(O_2, O_1)$）和交换律（$POR(POR(O_1, O_2), O_3) = POR(O_1, POR(O_2, O_3))$），这允许系统重组归约操作的顺序。此外，不同查询的归约操作是独立的。基于这两点观察，系统将树归约操作转化为 $bs$ 个独立的归约操作序列（$bs$ 为批次中的查询数）。序列中非相邻边的归约操作可以并行进行。因此，系统在两个维度上加速树归约：一是并行化不同查询的归约操作，二是并行化树中不同节点的归约操作。

复杂度分析对比。CoDec的IO复杂度为 $O(h \cdot d \sum_{i=1}^{node\_num} n[i])$，而FlashAttention的IO复杂度为 $O(h \cdot d \sum_{i=1}^{node\_num} n[i] \times n_q[i])$。给定 $\overline{n_q}$ 为KV缓存共享率的加权平均值，CoDec的IO复杂度比FlashAttention低约 $\overline{n_q}$ 倍。在计算复杂度方面，CoDec与FlashAttention相同。

工作负载平衡

对于一个KV缓存节点，KV长度（前缀长度）和关联查询的数量（共享程度）在不同节点间差异巨大。简单地为每个节点启动一个PAC内核会导致严重的块间负载不平衡；而过度细粒度的拆分会增加调度和归约开销，并可能无法充分利用每个块内的张量核心。

任务划分与调度公式化。部分注意力计算被建模为一组任务 $\mathbf{T}[i] = (n_q[i], n[i])$。任务可以在查询维度（水平切片 $b_q[i]$）和KV缓存维度（垂直切片 $b_k[i]$）上进行划分。任务分配策略表示为张量 $\mathbf{A} \in \mathbb{N}^{m \times t}$，其中 $m$ 是线程块数量。目标是找到最优的任务划分和调度策略，以最小化所有线程块的最大执行时间（即最慢线程块的执行时间）：$\arg\min \max_{i=1}^m (\sum_{j=1}^t \mathbf{C}[j] * \mathbf{A}[i, j])$。

NP-hard问题的启发式求解器。该问题是一个高级并行任务调度问题【16，Bounds for certain multiprocessing anomalies + 1966 + The Bell System Technical Journal】。由于 $n_q \ll n$，在查询维度划分会错过合并KV缓存内存访问的机会，因此系统将水平切片数 $b_q[i]$ 设为1，专注于垂直切片 $b_k[i]$。系统利用两个不等式缩小搜索空间：子任务的总成本不低于原始任务的成本，且每个块的最大成本不低于所有块的平均成本。随着更细粒度的划分，平均成本会因调度开销而增加。利用这种单调性，通过二分查找确定成本的下界 $cost_l$。然后通过设置每个KV缓存节点划分数的上限 $b_k[i] \leq \lceil \frac{C_{est}(n_q[i], n[i])}{cost_l} \rceil$ 来缩小搜索空间。该不等式限制了当成本低于平均成本时的进一步划分。最后，系统对每个KV缓存节点的划分数进行网格搜索，选择最优划分。

基于分析的成本估算。由于理论工作量（计算量 $O(n_q \times n \times d)$，内存访问 $O((n_q + 2n) \times d)$）与实际执行成本之间存在差异，且内核具有固定的启动开销，导致执行成本难以估算。对于小工作量，执行时间受内核启动开销主导；对于大工作量，受计算和内存访问主导。因此，CoDec提出了一种基于分析（profile-based）的方法。在给定硬件和模型配置下，仅有 $n$ 和 $n_q$ 影响PAC内核的执行时间。在部署模型前，系统对不同大小的 $n$ 和 $n_q$ 分析PAC内核并记录执行时间。对于未分析的计算，使用插值函数 $C_{est}(n_q, n)$ 来估算执行时间。

补充细节

CoDec内核模块实现。CoDec内核模块使用 CUDA/C++ 在 NVIDIA CUTLASS 之上实现（约1700行代码）。它接收查询张量、分页KV缓存和解码元数据作为输入，并将注意力输出写入指定的输出张量。它遵循与PagedAttention【22，Efficient Memory Management for Large Language Model Serving with PagedAttention + 2023 + SOSP + https://doi.org/10.1145/3600006.3613165】相同的分页KV缓存布局，并暴露与FlashDecoding兼容的注意力接口，使其能够直接集成到vLLM中。内核利用CUTLASS实现内存和计算之间的细粒度流水线。该内核支持MHA、MQA和GQA。对于使用多头潜在注意力（MLA）的模型，可以通过首先从潜在表示中重建每头KV块，然后应用相同的感知前缀的注意力和归约流水线来扩展CoDec 。

CoDec任务划分模块实现。任务划分模块实现为C++模块（约300行代码），与主解码循环接口。该模块监控解码过程并根据当前工作负载和资源可用性动态划分任务。为了减少开销，任务划分每隔几个解码步骤执行一次，而不是在每个步骤执行。经验表明，划分开销占注意力总时间的1.3%-2.5%，而并行归约在典型的共享前缀工作负载下带来的开销不到部分注意力计算的10%。

实验环境

  • 数据集：实验使用了受控的合成前缀树数据集，以及真实世界长上下文数据集LooGLE。在LooGLE数据集中，通过将共享相同文档上下文的查询分组来实体化共享前缀。
  • 模型架构：默认使用Qwen3-4B模型，其包含32个查询头、8个键/值头，头部维度为128。此外还评估了Llama3.2-1B、Mistral-7B和Llama4-17B等不同大小和注意力变体的模型。
  • 硬件配置：默认实验在单张NVIDIA A100 GPU（40GB，PCIe）上运行。此外，为了评估跨平台性能，还在NVIDIA H800、RTX 4090、A30和RTX A6000上进行了测试。
  • 软件配置：CUDA Toolkit 11.8（运行时版本12.2），vLLM 0.6.6，Python 3.10。基线对比系统包括FlashDecoding（由FlashAttention 2.7.4提供）、vLLM 0.6.6，以及FlashInfer的多级级联注意力。系统原型基于PyTorch 2.6.0实现。

实验结果

1. 与最先进技术（SOTA）的比较
实验评估了在不同序列长度、批次大小、树深度、共享前缀比率和树形状（二叉树2T、三叉树3T、四叉树4T、五叉树5T和退化树DT）下的性能。
* 注意力执行时间：如图5所示，CoDec在所有工作负载下的注意力执行时间均优于FlashDecoding，最高加速3.6倍，平均加速1.9倍。共享前缀越大，加速越显著；不规则工作负载（如退化树）下的加速比规则工作负载更明显，证明了CoDec在平衡不规则工作负载方面的高效性。
* 全局内存访问：如图6所示，CoDec的全局内存访问量显著低于FlashDecoding，降低了14.66倍至409.80倍，平均减少120.85倍。这验证了CoDec在合并内存访问方面的性能增益。
* 端到端延迟（TPOT）：在使用LooGLE数据集的端到端测试中（图7），CoDec相比vLLM平均降低了3.75倍的每个输出Token时间。此外，与FlashInfer的多级级联注意力（图8）相比，CoDec在不同的共享比率下始终实现更低的延迟和更高的吞吐量。这归功于CoDec具备整个前缀树的全局视角来划分工作负载，以及高效的并行树状归约策略。

图5. CoDec与FlashDecoding在执行时间上的对比。
图5. CoDec与FlashDecoding在执行时间上的对比。
图6. CoDec与FlashDecoding在全局内存访问上的对比。
图6. CoDec与FlashDecoding在全局内存访问上的对比。
图7. CoDec与vLLM在端到端时间上的对比。
图7. CoDec与vLLM在端到端时间上的对比。
图8. CoDec与SOTA在真实世界数据集上的对比。
图8. CoDec与SOTA在真实世界数据集上的对比。

2. 消融实验
为了理解每个组件的贡献，实验隔离了三个优化：共享前缀树、工作负载划分和并行执行与归约。如图9所示，在不平衡树（退化树）工作负载下，结合所有优化将延迟从38.0毫秒降至3.5毫秒（加速10.8倍）。在平衡树（二叉树）工作负载下，延迟从578毫秒降至22.2毫秒（加速26.1倍）。结果表明，所有三个优化都对降低延迟有贡献，且工作负载平衡和并行性对计算负载更高的平衡树影响更为显著。

图9. 消融实验。
图9. 消融实验。

3. 划分粒度的影响
如图10所示，实验将CoDec的自适应划分与均匀拆分任务的朴素固定策略进行了比较。结果显示，CoDec的自适应划分和调度策略在所有测试工作负载中超越了最佳的朴素固定策略（1.02-1.04倍）。与无划分的朴素方法相比，CoDec实现了3.20倍至4.39倍的加速（平均3.80倍）。这证明了CoDec策略在平衡子任务工作负载和减少内核启动/同步开销方面的有效性。

图10. 划分粒度的影响。
图10. 划分粒度的影响。

4. 任务划分的开销
如图11所示，生成划分计划的CPU开销随着批次大小的增加而增长，但在批次大小为64时仍保持在几十毫秒内。在实践中，由于一个划分计划可在多个解码步骤中重用，该开销相对于注意力执行时间仍然很小。

图11. 随着批次大小增加，计算任务划分计划的CPU开销。
图11. 随着批次大小增加，计算任务划分计划的CPU开销。

5. 在不同GPU上的性能
如图12所示，在50K Token上下文长度下，CoDec在所有测试的GPU上均优于FlashDecoding。在高端的H800上实现了4.7倍加速（2.094毫秒对比9.900毫秒）。在内存带宽较低的A6000上，CoDec的优势扩大到15倍（2.869毫秒对比43.048毫秒）。这表明CoDec对硬件内存瓶颈不敏感，非常适合跨数据中心和消费级GPU部署。

图12. 在不同GPU上的性能。
图12. 在不同GPU上的性能。

6. 在不同模型上的性能
如图13所示，CoDec在不同的GQA配置（如1:1, 1:2, 1:4）下始终如一地降低了解码延迟，表明其设计不局限于特定的头部布局。同时，在Llama3.2-1B、Mistral-7B和Llama4-17B等不同架构和规模的模型上，CoDec也保持了一致的延迟降低。

图13. 在不同注意力变体和模型下的性能。
图13. 在不同注意力变体和模型下的性能。

结论

本文提出了CoDec，这是一个专用的前缀共享解码算子，旨在通过有效利用跨多个请求的共享KV缓存模式，显著加速占据内存受限的LLM解码阶段主导地位的注意力计算。该方法引入了两项关键创新：一是新颖的共享前缀注意力内核，通过前缀KV缓存树与查询张量之间的复杂索引优化了内存层次结构，同时利用了块内和块间并行性；二是全面的工作负载平衡机制，包含基于分析的成本估算器、智能任务划分和高效的调度算法，以处理不规则的工作负载。实验结果表明，与最先进的FlashDecoding内核相比，CoDec在各种工作负载下实现了高达11.56倍的加速和150.56倍的内存访问减少。未来，CoDec可以很容易地与张量并行等分布式设置集成，探索这些分布式环境下的任务划分是一个有趣的方向。

方法细节中的参考文献汇总

  • 【7】FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning. (2024). ICLR. URL: https://arxiv.org/abs/2307.08691. (引用段落：背景知识-GPU架构与现有内核的局限。原文描述其通过分解张量以适配共享内存来优化注意力计算)。
  • 【8】Flash-Decoding for long-context inference. (2023). URL: https://crfm.stanford.edu/2023/10/12/flashdecoding.html. (引用段落：背景知识-GPU架构与现有内核的局限。原文描述其作为长上下文推理的最先进注意力优化方法)。
  • 【16】Bounds for certain multiprocessing anomalies. (1966). The Bell System Technical Journal. (引用段落：方法细节-工作负载平衡。原文描述任务划分与调度公式化是一个NP-hard的高级并行任务调度问题)。
  • 【22】Efficient Memory Management for Large Language Model Serving with PagedAttention. (2023). SOSP. URL: https://doi.org/10.1145/3600006.3613165. (引用段落：补充细节-CoDec内核模块实现。原文描述CoDec遵循与PagedAttention相同的分页KV缓存布局)。