前言

大语言模型推理与训练存在本质差异。训练阶段以批量数据吞吐为核心优化目标，而推理阶段则面临截然不同的压力组合：长序列输入带来的显存非线性增长、动态请求长度导致的批处理效率塌陷、以及Decoding阶段算力利用率低下等结构性矛盾。在昇腾NPU上运行LLM推理，如果沿用传统GPU时代的老路子，往往会在显存墙和调度效率上同时碰壁。

ascend-transformer-boost仓库（下文简称ATB加速库）是CANN软件栈面向Transformer推理场景推出的专用加速组件，提供了从底层算子到图算子再到插件机制的全套能力。该仓库定位明确：基于华为Ascend AI处理器，专门为Transformer模型的推理场景做深度性能优化。与通用PyTorch推理路径相比，ATB加速库在显存占用、调度开销和算子融合三个维度上均做了定向改进。

本文从LLM推理的特殊性出发，深入剖析ATB加速库在KV Cache管理、Continuous Batching调度和算子融合三个关键环节的设计决策和调优思路，结合仓库中实际的代码结构与API设计，给出一套可落地的昇腾NPU推理加速实践路径。

KV Cache管理：显存是LLM推理的第一瓶颈

问题根源：注意力机制的显存膨胀

在Transformer的Decode过程中，每个新生成的Token都需要与完整的历史序列重新计算注意力。假设模型hidden dimension为H、上下文长度为S、batch size为B，则单层注意力机制中KV Cache的显存占用可以表示为2乘以B乘以S乘以H（分别对应K和V两个方向）。对于一个70B参数的模型，H通常为8192，当序列长度扩展到8192时，单层KV Cache的显存需求就已经超过百GB。多层堆叠后，显存消耗成为制约推理吞吐量的第一约束。

传统方案在显存管理上采用连续内存分配策略，即预先为每个请求分配固定长度的KV Cache缓冲区。这种方式的缺陷在于：请求长度不可预测，提前分配过短则溢出，分配过长则大量显存被无效占用。在长序列场景下，这种静态分配机制带来的显存碎片化和浪费尤为严重。

分块管理与PagedAttention的工程实现

ATB加速库在KV Cache管理上引入了分块（Block）管理策略，将连续的KV Cache划分为固定大小的Block进行管理。每个Block包含若干个Token对应的Key和Value向量，通过Block级分配替代Token级分配，有效减少内存碎片。在ATB的KVCacheTilingData结构中可以看到关键配置参数的设计：

namespace AtbOps {
struct KVCacheTilingData {
    uint32_t batch{0};
    uint32_t maxSeqLen{0};
    uint32_t hiddenSize{0};
    uint32_t batchPerCore{0};
};
}

KVCacheTilingData中定义了batch（并发请求数）、maxSeqLen（最大序列长度）、hiddenSize（隐藏层维度）和batchPerCore（每核处理的batch数）。这四个参数共同决定了KV Cache在HBM上的物理布局和Tiling切分策略。其中maxSeqLen对应预分配的最大Block数量，而hiddenSize决定了每个Block的存储粒度。

TilingData结构将KV Cache的切分参数显式暴露给开发者，而非在库内部黑盒处理。在昇腾NPU上，数据从Global Memory（HBM）到Local Memory（UB）的搬运是性能瓶颈之一，通过提前指定Tiling策略，ATB可以在算子下发前计算出最优的数据分块大小，避免运行时反复推导。对比连续内存分配，Block级管理使得多个请求可以共享同一块物理内存的不同Block，当某个请求结束时，其占用的Block可以立即被其他请求复用，无需等待整个预分配缓冲区释放。

在mixkernels/kvcache目录下，KVCache算子提供了按Block粒度读写KV Cache的能力。对比连续分配方案，分块管理的显存利用率在长序列场景下通常可提升数倍，其差异来源正是Block级复用消除了为每个请求预分配最大长度缓存的开销。

显存优化对推理吞吐的影响

在昇腾NPU上运行LLM推理时，显存约束直接决定了可支持的Batch Size上限。ATB加速库通过分块KV Cache管理，使得系统在相同显存条件下可以容纳更多并发请求。以下是使用分块KV Cache管理前后的关键指标对比：

维度	连续内存分配	分块KV Cache（ATB）	差异来源
显存占用（8192 seq, 70B模型, batch=4）	预分配4份最大序列缓存，固定约64GB	按实际使用Block动态分配，约28GB	Block级复用减少预分配冗余
显存利用率	受最长相符预分配约束，平均约35%	动态分配，实际使用率可达70%以上	无固定长度预分配
Batch Size上限（单卡，80GB HBM）	受预分配策略限制，batch=2~3	分块复用支持batch=6~8	显存碎片减少
长序列支持（显存固定）	超过预分配长度需拒绝请求	按Block扩展，支持到硬件上限	无固定预分配约束

Continuous Batching：动态请求调度打破静态批处理壁垒

静态Batching的致命缺陷

传统推理服务采用固定Batch的批处理方式：将多个请求打包为一个Batch统一处理，等待该Batch中最长请求完成后释放。这种方式在请求长度差异较大的LLM推理场景中存在严重问题。当Batch中混入短请求时，短请求在生成少量Token后就已达到终止条件，但必须等待同Batch中长请求全部完成才能释放计算资源。这种"木桶效应"导致大量计算资源浪费在等待短请求的无效迭代上。

更关键的是，静态Batching在prefill阶段和decode阶段的处理逻辑不同。Prefill阶段处理输入Prompt，计算量与Token数成正比；Decode阶段逐Token生成，每步计算量固定但需要反复调度。两种阶段混合Batch时，调度复杂度进一步上升。

Continuous Batching的调度机制

Continuous Batching（又称Iteration-level Batching或Dynamic Batching）的核心思想是：在每个生成步（Generation Step）结束后，立即检查已完成的请求并释放其占用的Batch Slot，同时将新到达的请求动态插入到释放的Slot中。ATB加速库通过图算子（GraphOp）机制支持了这种动态调度能力。

在仓库的ops_infer目录中可以看到多种注意力相关算子（faupdate、block_copy、gather等），它们共同构成了Continuous Batching的底层支撑。以faupdate算子为例，它负责在每个Step结束后更新FlashAttention的输出状态，配合KV Cache的分块管理，可以在请求粒度上精确控制显存占用。

ATB提供的图算子机制允许开发者将多个基础算子组合为一个大粒度的图算子进行统一调度。对于Continuous Batching场景，这意味着Prefill和Decode阶段的注意力计算可以封装为统一的图算子接口，由ATB的运行时负责在每次Step结束时进行批量调度优化：

atb::Status CreateLlamaMlpOperationByGraphOpBuilder(const LlamaMlpParamGb &param, atb::Operation **operation)
{
    atb::GraphOpBuilder* graphOpBuilder;
    CreateGraphOpBuilder(&graphOpBuilder);
    graphOpBuilder->Init(
        "LlamaMlpGraphOp",
        inferShapeFunc,
        {"hidden_states", "weight"},
        {"mlp_out"}
    );
    graphOpBuilder->Reshape("hidden_states", reshape_01_2, "hidden_states_");
    graphOpBuilder->AddOperation(Linear(param), {"hidden_states_", "weight"}, {"linear_out"});
    graphOpBuilder->Reshape("linear_out", unsqueueze_0, "linear_out_");
    graphOpBuilder->AddOperation(Split(param), {"linear_out_"}, {"gate_out", "up_out"});
    graphOpBuilder->AddOperation(Swish(param), {"gate_out"}, {"swish_out"});
    graphOpBuilder->AddOperation(Mul(param), {"swish_out", "up_out"}, {"mlp_out"});
    *operation = graphOpBuilder->Build();
    DestroyGraphOpBuilder(graphOpBuilder);
    return atb::NO_ERROR;
}

GraphOpBuilder将多个独立算子组合为一个逻辑单元后，ATB的运行时可以在Setup阶段统一规划中间张量的显存复用策略。连续内存分配下一个算子的Workspace在算子完成后才能释放，而图算子内部可以通过Block级内存复用让前序算子的Workspace直接被后续算子覆盖使用。这种设计使得Continuous Batching在请求动态加入和退出时，中间缓冲区的分配和释放开销大幅降低。运行时优化层面的Tiling Cache机制缓存了推理过程中重复的Tiling计算结果，在动态Shape场景下有效减少每次Setup的重复推导。

ATB的运行时还提供了双线程下发优化模式：一个线程负责批量执行Setup，另一个线程负责批量下发Execute，两者在流水线上并行推进。这一机制直接针对Host Bound场景设计，在请求级Batch动态切换时，Host侧的调度开销往往是NPU利用率的瓶颈，双线程流水化可以有效消除该瓶颈。

算子融合：减少Kernel下发开销的核心手段

Host Bound问题的根源

深度学习模型的推理过程可以抽象为算子下发与设备执行的交替过程：Host（CPU）准备算子的上下文参数、Tiling信息和输入输出地址，然后通知Device（NPU）执行。在算子数量众多且每个算子独立下发的场景下，Host下发速度跟不上NPU执行速度，就会产生Host Bound问题——NPU算力闲置等待算子下发，数据流上出现大量空泡。

在LLM推理中，单层Transformer需要执行的算子包括：RotaryEmbedding（位置编码）、QKV投影（Matmul）、Softmax（注意力归一化）、残差LayerNorm、RMSNorm等多个独立Kernel。以RotaryEmbedding为例，它对Query和Key向量施加与位置相关的旋转变换，在Attention计算前完成。如果这些算子各自独立下发，Prefill阶段几十层堆叠下来，Host Bound的开销会严重拖累整体延迟。

融合算子的设计与实现

ATB加速库的核心优化策略之一就是提供经过精心设计的融合算子，将多个相邻计算步骤合并为一次NPU内核调用。融合算子减少了下发次数和数据搬运次数，是解决Host Bound问题的根本手段。

以RMSNorm为例，ATB在src/kernels/kernels/norm/rmsnorm目录下提供了高性能实现。RMSNorm的计算公式为输出等于输入乘以归一化因子除以RMS加epsilon的平方根再乘以gamma权重。在LLM中，RMSNorm通常出现在每个Transformer Block的输入和输出位置，调用频率极高。将RMSNorm与其他邻近算子（如残差加法）融合，可以将原本需要三次独立内存访问（读取输入、写入归一化结果、写入残差）的操作合并为一次NPU内核调用。

ATB仓库中example/op_demo/rms_norm目录提供了RmsNormOperation的C++调用示例，展示了标准的参数配置方式：

// 创建RmsNorm算子参数
atb::infer::RmsNormParam rmsNormParam;
rmsNormParam.layerType = atb::infer::RmsNormParam::RmsNormType::RMS_NORM_NORM;
rmsNormParam.normParam.quantType = atb::infer::QuantType::QUANT_UNQUANT;
rmsNormParam.epsilon = 1e-6;

// 输入Tensor: x [batch, seq_len, hidden_size] 格式为BF16或FP16
// gamma权重: [hidden_size] 格式为BF16或FP16
// 输出Tensor: output [batch, seq_len, hidden_size]

atb::Operation* rmsNormOp = nullptr;
atb::infer::CreateRmsNormOperation(rmsNormParam, &rmsNormOp);

RmsNormOperation的参数设计将计算精度（epsilon）、量化类型（quantType）和归一化类型（layerType）三个维度解耦。在昇腾NPU的硬件约束下，不同hidden_size和seq_len组合对应不同的Tiling策略。通过将epsilon和quantType作为独立参数暴露，算子内部的Tiling计算可以在Setup阶段针对具体Shape做定制化推导，而不是在Execute阶段实时计算。归一化类型参数（RMS_NORM_NORM）区分了标准RMSNorm与其他变体，确保算子内核选择最匹配的硬件指令路径。

融合算子的价值在Attention相关操作上体现得最为充分。ATB仓库中ops_infer目录下包含faupdate（FlashAttention更新）、block_copy（块拷贝）等算子，mixkernels目录下包含kvcache、fusion、fastsoftmax等融合算子。在实际LLM推理中，Attention层的典型融合模式是将RotaryEmbedding、QKV投影、Attention Score计算和Softmax合并为一个融合Kernel。ATB提供的gmm_add、mm_deq_swiglu_quant_mm_deq等融合算子就是针对MoE架构中Gate和MLP融合的工程实现。

逐算子执行与融合执行的性能对比

算子融合对推理性能的影响主要体现在两个维度：减少Host下发次数和消除设备侧内存带宽瓶颈。以下对比表格展示了逐算子执行与融合执行在LLM推理中的核心指标差异：

维度	逐算子执行（原始PyTorch路径）	融合算子执行（ATB加速库）	差异来源
单层Transformer Kernel下发次数	8~12次（QKV投影+Rotary+Softmax+LayerNorm等独立下发）	2~3次（融合后的大粒度算子）	算子融合减少独立下发
NPU流空泡比例（Host Bound占比）	高（每次独立下发间存在调度空隙）	低（融合算子内部调度无间隙）	批量下发优化
HBM带宽占用	多次读写中间Tensor，冗余数据传输	中间结果在UB内复用，减少HBM访问	内存复用优化
单次Prefill延迟（70B, seq=2048）	基准值	约为逐算子执行的0.5~0.7倍	融合减少整体调度开销
推理吞吐（tokens/s，单卡）	基准值	约为逐算子执行的1.5~2.5倍	减少下发+融合计算效率提升
Workspace显存占用（图算子级）	各算子独立分配，总和较大	中间Tensor复用，节省约50%	ATB HBM内存Block级复用

在ATB仓库中，融合算子的落地通过两条路径实现。第一条是ATB官方提供的预置融合算子（如kvcache、fusion目录下已实现的融合Kernel），开发者直接调用即可。第二条是GraphOpBuilder机制，开发者可以将自定义算子与ATB原生算子组合为图算子，由ATB负责内部的调度优化和内存复用。

对于追求极致性能的团队，ATB还提供了Plugin机制（ops_customize目录），允许开发者编写自定义融合算子并注册到ATB算子库中。这一机制为前沿模型架构（如新的Attention变体、特殊位置编码）提供了灵活的扩展能力，同时保持了对ATB原有调度优化管线的兼容性。

结尾

LLM推理在昇腾NPU上的性能优化是一个多维度交织的系统工程。KV Cache的分块管理从根本上解决了长序列场景下的显存利用率问题，使单卡能够支持更大的并发Batch；Continuous Batching的动态调度打破了静态Batch的壁垒，让计算资源在任何时刻都尽可能被有效请求所占用；算子融合则从Host下发效率和NPU计算效率两个方向同时发力，减少了Kernel间的空泡和冗余内存访问。

ascend-transformer-boost仓库提供的不仅仅是单点优化，而是一套覆盖底层算子、高层API、运行时优化和扩展机制的综合能力。开发者在实践中需要根据具体模型结构、请求分布和硬件条件，在这三个优化维度上做合理的取舍与组合。理解ATB的设计意图和约束条件，是将其能力充分发挥的前提。

https://atomgit.com/cann/ascend-transformer-boost