前言

ascend-transformer-boost（ATB）是昇腾CANN生态中面向大模型推理场景的核心加速库，运行在昇腾NPU之上，为Transformer架构的推理流程提供了动态批处理、投机解码和KV缓存复用三项关键能力。大模型推理的瓶颈并不单一：请求的输入长度参差不齐会导致GPU/NPU计算资源的闲置，自回归解码的逐token生成特性限制了单请求吞吐，而重复计算的KV缓存则浪费了宝贵的显存带宽。ATB针对这三个维度的瓶颈分别给出了对应方案，这三项技术并非独立运作——动态批处理为投机解码提供了批量验证的基础，KV缓存复用则为动态批处理中的请求调度提供了显存管理的底层支撑。本文从机制原理到部署实践，对ATB的这三项核心技术进行完整拆解，并结合端到端推理场景中的调优方法，给出在昇腾NPU上实际部署ATB时需要关注的配置项与权衡策略。

动态批处理机制

静态批处理的工作方式是：把若干请求凑成一批，等这一批全部完成后再处理下一批。问题在于，如果同一批请求中有的输入只有几十个token，有的输入长达数千token，短请求完成后只能空等长请求，NPU的计算单元处于闲置状态。这种等待在自回归解码阶段更为严重——不同请求的生成长度差异更大，短请求可能只需要几十步解码，长请求可能需要数百步，静态批处理下短请求的等待时间完全被长请求拖住。

动态批处理的核心改动在于：允许新请求在当前批次尚未全部完成时就加入批次。当某个请求完成解码后，它的槽位立刻被释放，新的等待请求可以填充进来，NPU始终处于接近满载的计算状态。ATB的动态批处理调度器维护一个请求队列和一个正在执行的批次，每当有请求完成或新请求到来时，调度器会根据当前批次的空闲槽位和显存预算决定是否将新请求加入当前批次。

填充策略是动态批处理中一个不可回避的问题。同一批次中不同请求的输入长度不同，需要将它们对齐到相同长度才能进行批量矩阵运算。Padding策略的做法是将所有请求的输入填充到批次中最长请求的长度，短请求会填充大量无效token。当批次中请求长度差异较大时，Padding带来的计算浪费相当可观。Packing策略则不同：它将多个短请求的token拼接在一起，填满一个"虚拟"的长序列，使得NPU的每个计算单元都在处理有效token。ATB在Prefill阶段支持Packing模式，通过将不同请求的token紧凑排列，减少了无效计算。在Decode阶段，由于每个请求每步只生成一个token，批次中各请求的token数已经对齐为1，Padding和Packing的差异消失。

动态批处理下的调度算法需要考虑公平性问题。如果不加干预，调度器可能倾向于优先处理短请求（因为它们完成得快，能快速释放槽位），导致长请求在队列中等待时间过长。ATB的调度器引入了基于等待时间的优先级机制：请求在队列中等待的时间越长，其调度优先级越高。这种策略避免了长请求被持续饥饿的情况。调度器还会根据当前批次的负载情况动态调整新请求的接纳速率——当批次接近满载时，即使队列中有等待的请求也不会强行加入，因为过大的批次会导致单步解码时间增长，反而拖慢所有请求的响应速度。

# ATB动态批处理调度配置示例
atb_config = AtbConfig()
atb_config.batch_schedule = "dynamic"
atb_config.max_batch_size = 32
atb_config.max_queue_latency_ms = 50
#  max_queue_latency_ms caps the wait time before a request enters
# the running batch; without this cap, long-queue scenarios cause tail
# latency spikes because requests pile up waiting for sparse slots.

动态批处理引入了一个微妙的开销来源：每次批次组成发生变化（请求加入或离开）时，注意力计算的掩码矩阵需要重新构建。ATB通过预分配固定大小的掩码缓冲区并在每次调度时增量更新，避免了反复分配显存的开销。这种增量更新策略的代价是掩码缓冲区始终占用着最大批次大小对应的显存，即便当前批次远未达到最大批次。

投机解码技术

自回归解码的固有瓶颈在于每步只能生成一个token——模型必须等待当前token生成完毕才能开始下一个token的计算，这是一个严格的串行过程。投机解码的思路是：用一个参数量远小于目标模型的小模型（Draft Model）快速生成若干候选token，再用目标模型（Target Model）一次性验证这些候选token，保留被目标模型接受的token，拒绝的token及之后的候选token全部丢弃，并从拒绝位置重新开始Draft-Verify循环。

ATB中投机解码的执行流程分为两个阶段。Draft阶段：小模型以自回归方式连续生成K个候选token（K称为Draft长度），这一步的计算量远小于目标模型的单步前向计算。Verify阶段：目标模型对这K个候选token进行一次前向计算，同时产出每个位置的概率分布，将候选token与概率分布中的最高概率token进行比对，若一致则接受，不一致则拒绝。被接受的token数量占Draft总长度的比例称为接受率。接受率直接决定了投机解码的加速比——如果接受率为0.8、Draft长度为5，则平均每个Draft-Verify循环能产出4个有效token，相比自回归解码的1个token有了明显提升。

接受率受多个因素影响。Draft模型与Target模型的分布差异是最核心的因素：如果两个模型对同一上下文的输出分布越接近，接受率越高。Draft长度K的选择也影响接受率——K越大，后面的候选token偏离Target模型分布的概率越高，接受率会随位置递减。ATB允许用户配置Draft长度，并在运行时根据实测的接受率动态调整K值：当接受率低于设定阈值时自动缩短Draft长度，当接受率高于阈值时适度增加Draft长度。这种自适应策略在输入分布变化较大的场景下比固定Draft长度更稳健。

投机解码的显存开销是一个容易被低估的问题。Draft模型本身需要占用额外的显存来存储权重，Verify阶段需要为K个token的KV缓存预分配空间，即使这些token可能被拒绝。ATB的显存预算管理策略是在初始化时为投机解码预留一个固定的显存池，Draft模型的权重和Verify阶段的临时KV缓存都从这个池中分配。当可用显存不足以同时容纳Draft模型和Target模型时，ATB会回退到标准的自回归解码模式，不会因为显存不足而导致推理中断。

# ATB投机解码配置示例
atb_config.speculative_decoding = True
atb_config.draft_model_path = "/models/draft-7b"
atb_config.draft_length = 5
atb_config.acceptance_threshold = 0.6
atb_config.spec_memory_budget_mb = 4096
#  spec_memory_budget_mb reserves a fixed VRAM pool so that the
# draft model weights and speculative KV cache never encroach on the
# target model's working set; without this isolation, OOM can occur
# mid-inference when both models compete for the same HBM region.

投机解码在对话场景中的表现与续写场景有差异。对话场景中，用户输入的新问题与之前的多轮对话拼接后作为Prefill输入，Prefill阶段不使用投机解码（Prefill本身是并行计算，不需要投机）。投机解码只在Decode阶段生效，而对话场景中Decode阶段的生成长度通常较短（几十到上百token），Draft-Verify循环的启动开销占比相对更高。续写或长文本生成场景中，Decode阶段持续数百甚至数千步，投机解码的加速效果更加明显。ATB在实现上针对短Decode场景做了优化：当预测剩余生成长度小于Draft长度的两倍时，自动关闭投机解码，避免启动Draft模型的收益无法覆盖其开销。

KV缓存复用

Transformer推理中的KV缓存是指：在自回归解码过程中，已生成token对应的Key和Value矩阵需要被保存下来，供后续每一步解码的注意力计算使用。如果不缓存这些Key和Value，每一步都需要重新计算所有已生成token的KV表示，计算量随生成长度呈二次增长。KV缓存将这个二次计算降为线性——每步只需要计算当前新token的KV表示，再与缓存中的KV拼接后做注意力计算。

传统的KV缓存实现方式是为每个请求预分配一块连续的显存，大小等于最大序列长度乘以每token的KV缓存大小。这种预分配方式有两个问题：一是在序列长度未达到最大值时，预分配的显存大部分处于闲置状态；二是不同请求的实际生成长度差异很大，有的请求可能只用了预分配空间的十分之一，浪费了大量显存。当并发请求数增多时，这种浪费会直接限制最大并发数。

PagedAttention的核心思想是将KV缓存按固定大小的Block进行管理，每个Block存储若干token的KV数据，请求的KV缓存由一组Block组成，这些Block在物理显存上不必连续，通过Block表进行映射。当请求需要更多KV缓存空间时，只需从Block池中分配新的Block并加入Block表，不需要预分配最大长度的连续空间。这种按需分配的方式将KV缓存的显存利用率大幅提高，因为Block只有在被实际使用时才分配。

ATB对PagedAttention的昇腾适配需要处理昇腾NPU的硬件特性。昇腾NPU的Cube计算单元在执行矩阵乘法时对数据的内存布局有特定要求，PagedAttention的实现需要确保Block中的KV数据在送入Cube单元时的内存访问模式是高效的。ATB将Block的大小设计为与昇腾NPU的内存访问粒度对齐，使得单个Block的读取可以在一次内存事务中完成，避免了跨Block边界的不连续访问带来的带宽损失。

KV Cache的预分配策略在ATB中表现为：启动时根据配置的最大并发请求数和最大序列长度，预先分配一个Block池，池中的Block数量等于最大并发请求数乘以每个请求预估需要的Block数。这个预估基于历史统计或用户配置的平均序列长度。当实际请求的KV缓存需求超出预估时，ATB的动态扩展机制会尝试从系统可用显存中分配新的Block。动态扩展有上限——当可用显存低于安全阈值时，ATB会拒绝新请求入队而不是冒着OOM风险继续扩展。这种拒绝策略虽然可能导致部分请求排队等待，但保证了正在执行的请求不会因为显存不足而中断。

多请求间KV Cache的共享场景出现在多轮对话中。当用户在同一对话会话中发送新消息时，之前轮次的KV缓存仍然有效，只需要在新消息的Prefill阶段计算新输入token的KV表示并追加到已有缓存中。ATB通过会话ID将同一对话的KV缓存Block关联在一起，新轮次复用已有Block，只分配新token对应的Block。不同会话的KV缓存是严格隔离的——一个请求的注意力计算只能访问自己的Block表，不会读取到其他请求的KV数据，这种隔离通过Block表的访问控制实现，而非依赖物理内存的隔离。

# ATB KV缓存配置示例
atb_config.kv_cache_type = "paged"
atb_config.kv_block_size = 16  # tokens per block
atb_config.max_num_seqs = 64
atb_config.kv_cache_gpu_mem_utilization = 0.85
#  kv_cache_gpu_mem_utilization caps the fraction of HBM that
# the block pool may consume; on Ascend NPU the Cube unit requires
# contiguous workspace for matmul, so leaving 15% free prevents
# fragmentation-triggered allocation failures that would crash the
# inference loop mid-request.

Block大小（kv_block_size）的选择是一个需要权衡的参数。Block越大，Block表越短，映射管理的开销越小，但每个Block内未被使用的token空间浪费可能更大（类似于文件系统中的内部碎片）。Block越小，空间利用率越高，但Block表更长，映射查找的次数增加，且不连续内存访问的概率上升。ATB默认的Block大小经过针对昇腾NPU内存访问特性的调优，在常见的序列长度分布下提供了较好的折中。在序列长度普遍较短（几百token以内）的场景中，可以适当减小Block大小以提高空间利用率；在序列长度普遍较长的场景中，增大Block大小可以降低映射管理的开销。

端到端推理性能调优

Transformer推理的端到端流程分为Prefill阶段和Decode阶段，这两个阶段的计算负载特性截然不同。Prefill阶段需要一次性计算整个输入序列所有token的KV表示，计算量与输入长度的平方成正比，属于计算密集型操作，NPU的Cube单元在这个阶段处于高利用率状态。Decode阶段每步只处理一个token，计算量远小于Prefill，但需要从KV缓存中读取所有已生成token的KV数据，内存访问量随已生成长度线性增长，属于内存带宽密集型操作，NPU的计算单元在这个阶段往往利用率不足。

这种计算负载差异意味着两个阶段的调优方向不同。Prefill阶段的瓶颈在于计算能力，增大批量大小可以让Cube单元更充分地并行处理多个请求的Prefill计算。Decode阶段的瓶颈在于显存带宽，增大批量大小虽然可以让多个请求的KV缓存读取合并为更大的内存事务，但当批量大小超过一定阈值后，显存带宽会成为硬瓶颈，继续增大批量只会增加单步解码时间而不会提高吞吐。

ATB提供了Profile接口用于采集推理过程中的硬件性能数据，包括每步解码的耗时、Cube和Vector单元的利用率、显存带宽利用率、HBM使用量等。通过分析Profile数据可以判断当前推理流程的瓶颈所在：如果Cube利用率在Decode阶段持续偏低，说明瓶颈在显存带宽；如果显存带宽利用率已经接近上限而吞吐仍然不够，说明需要考虑减少单个请求的KV缓存占用来容纳更多并发请求。

批量大小（max_batch_size）和最大并发token数（max_num_tokens）是ATB中两个相互关联的配置项。max_batch_size限制了同一批次中最大请求个数，max_num_tokens限制了同一批次中最大token总数。在Prefill阶段，由于Packing的存在，max_num_tokens是实际起约束作用的参数——即使max_batch_size允许32个请求同时执行，如果这些请求的总token数超过了max_num_tokens，调度器只会接纳部分请求。在Decode阶段，每个请求每步只贡献1个token，max_batch_size成为主要约束。ATB允许用户分别配置Prefill和Decode阶段的max_batch_size和max_num_tokens，以适应两个阶段的不同负载特性。

在实际部署中，max_batch_size的设置需要权衡吞吐和延迟。较大的max_batch_size可以提高吞吐（更多请求并行），但每个请求的响应延迟也会增大（因为每步解码需要处理更多请求的KV缓存读取）。对延迟敏感的场景（如实时对话）需要较小的max_batch_size，对吞吐敏感的场景（如离线批量推理）可以设置较大的max_batch_size。ATB的调度器支持在运行时根据队列深度动态调整实际批次大小：当队列中的请求数远小于max_batch_size时，不会为了填满批次而等待，而是立即开始处理已有的请求，这种策略在低负载场景下避免了不必要的等待延迟。

Prefill阶段的另一个调优维度是Chunked Prefill。当输入序列非常长时，一次性计算整个序列的Prefill可能导致单次计算时间过长，延迟了同一批次中其他请求的Decode步。ATB支持将长序列的Prefill拆分为多个Chunk，每个Chunk的大小不超过配置的max_num_tokens，这样长序列的Prefill可以与其他请求的Decode交替执行，避免长Prefill独占NPU导致的Decode停顿。Chunked Prefill的代价是增加了调度次数和KV缓存的拼接操作，但在混合了长短序列的在线推理场景中，它对降低尾延迟有明显帮助。

效率对比

维度	使用前	使用后	差异来源
推理吞吐	静态批处理下短请求空等长请求，NPU计算单元利用率低	动态批处理持续填满批次槽位，NPU利用率接近饱和	动态批处理的请求级调度消除了空等闲置
显存占用	每请求预分配最大序列长度的连续KV缓存，大量空间闲置	PagedAttention按Block按需分配，未使用的Block可供其他请求复用	PagedAttention将预分配的固定开销变为按需分配的弹性开销
端到端延迟	自回归解码每步仅生成1个token，串行瓶颈明显	投机解码在Draft接受率较高时每步平均产出多个token，有效缩短Decode阶段总步数	投机解码用小模型的计算量换取了目标模型的跳步验证，但Draft模型额外占用显存和算力
批量公平性	静态批处理下所有请求等待最慢请求完成，短请求被长请求拖慢	动态批处理下请求完成后立即释放槽位，短请求的等待时间不再受长请求制约，使用后批量公平性并未降低	调度器的等待时间优先级机制防止了长请求饥饿，同时短请求的完成时间不再与长请求耦合

上表中的"使用后批量公平性并未降低"是一个需要说明的约束：动态批处理在提升短请求响应速度的同时，并没有以牺牲长请求的等待时间为代价。调度器的等待时间优先级确保了长请求在队列中不会无限等待，而短请求的提前完成释放的资源反而为长请求提供了更多计算机会。但动态批处理并非在所有维度上都有收益——调度器的运行本身引入了额外的CPU开销，每步调度决策的耗时在极高并发场景下会变得不可忽略；投机解码在Draft接受率较低时（例如Draft模型与Target模型分布差异较大），Verify阶段的计算量相比直接自回归解码反而更高，此时投机解码不仅没有加速，还增加了延迟。

结尾

ATB通过动态批处理、投机解码和KV缓存复用三项技术，分别从计算利用率、解码速度和显存效率三个维度为昇腾NPU上的大模型推理提供了优化方案。动态批处理解决了静态批处理下请求长度差异导致的计算闲置问题，Packing策略在Prefill阶段进一步消除了Padding的无效计算。投机解码利用Draft-Verify两阶段流程突破了自回归解码的串行瓶颈，自适应Draft长度和显存预算管理使得投机解码在不同负载条件下都能安全运行。PagedAttention的昇腾适配实现了KV缓存的按需分配和动态扩展，Block粒度的显存管理与NPU的内存访问特性对齐，多请求间KV缓存的共享和隔离策略兼顾了效率与安全。

---

仓库链接：https://atomgit.com/cann/ascend-transformer-boost