前言

在昇腾NPU上做自回归推理的解码阶段,每多生成一个 token,KV Cache 的显存开销就膨胀一分,响应延迟也跟着拉长。CANN 社区提供的 Ascend Transformer Boost(ATB)加速库专门针对 Transformer 模型设计,基于昇腾NPU的 AICore、Cube Unit 和 HCCS 链路特性,从算子融合、内存管理和通信调度三个层面提供了端到端的工程基础设施。解决这个问题的工程路径不是单点优化,而是从 KV Cache 量化压缩显存占用量,到 Speculative Decoding 削减解码步数,再到通算融合将通信隐藏到计算背后的一系列系统性手段。ATB库把融合算子、图算子组图、Tiling Cache、Workspace复用和双线程调度等机制封装成可复用的运行时组件,每一项都可以直接服务于推理延迟的压缩,本文逐一分析每个环节在ATB库中的工程实现细节与设计权衡。ATB 库提供了融合算子、图算子组图、Tiling Cache、Workspace 复用和双线程调度等运行时优化机制,每一项都可以直接服务于推理延迟的压缩。本文以 KV Cache 量化、Speculative Decoding 和通算融合三条主线展开,逐一分析每个环节在 ATB 库中的工程实现细节与设计权衡。

KV Cache的显存瓶颈:推理阶段的显存增长曲线

自回归推理的每步解码,注意力层都需要将新计算出的 Key 向量和 Value 向量追加到历史缓存中。对于一个 L 层 Transformer 模型,隐藏维度 d,当前序列长度 s,batch size b,KV Cache 总显存占用为 2 x L x s x b x d x sizeof(dtype)。以 LLaMA-70B 为例,隐藏维度 8192,层数 80,采用 FP16 精度推理且 batch size 为 1 时:序列长度 4096 token 时 KV Cache 需要 2 x 80 x 4096 x 1 x 8192 x 2 个字节,约 10.7 GB;序列长度扩展到 32768 token 时,这一数字膨胀至 85.9 GB。单张 Ascend 910B 的 HBM 容量为 64 GB,意味着长序列推理场景下 KV Cache 独占 HBM 挤占了模型权重的空间,batch size 和序列长度无法同时放大。推理服务为了降低首 token 延迟通常会采用较大的 batch size 并发处理多个请求,KV Cache 的线性膨胀让 batch size 的提升受到严格限制。

增大 batch size 时显存占用按倍数增长。batch size 为 4、序列长度 4096 时 KV Cache 需要 42.8 GB,接近 910B 的 HBM 上限。如果加上模型权重(70B FP16 约 140 GB,需要 4 卡张量并行分摊),每张卡负载超过 50 GB,显存余量所剩无几。KV Cache 的显存增长曲线决定了推理服务的最大并发度和硬件部署成本,成为部署阶段最核心的资源约束。

ATB 库在算子层为 KV Cache 管理提供了两套互补的方案:分页管理和高效缓存写入。分页机制将连续物理显存切分为固定大小的 block,每个请求按需分配 block 页面,页面之间不要求物理连续,通过 block_table 记录虚拟地址到物理 block 的映射。这种设计消除了传统连续显存分配中因请求终止释放而产生的内存碎片。当某个请求完成推理后,其占用的 block 页面立即归还给空闲列表,后续请求可以按任意顺序分配这些页面,不会出现外部碎片。请求的序列长度是动态增长的,每多生成一个 token 就需要一个新的 block 页面。ATB 的 PagedAttention 算子在运行时根据 contextLens 的变化动态追加页面,block_table 随之增长。这种按需分配的机制确保了显存只在真正需要时才被消耗,避免了预分配造成的浪费。

ReshapeAndCacheOperation 负责将解码步产出的 Key 和 Value tensor 写入分页缓存。不同 Ascend 芯片的内存排布存在差异:Ascend 310P 采用 NZ 格式,Ascend 910B 采用 BNSD 格式,Ascend 950 支持 SISO(Single Input Single Output)模式。ATB 库在 src/ops/ops_infer/reshape_and_cache/ 下为每种芯片提供了独立的 ops_runner 实现,通过 reshape_and_cache_ops_runner_310p.cpp、reshape_and_cache_ops_runner_A2_NZ.cpp 和 reshape_and_cache_siso_aclnn_runner.cpp 分别适配。这种芯片级适配保证了 KV Cache 写入操作在不同硬件代际上的 HBM 带宽利用效率。除标准版本外,ATB 还提供了 reshape_and_cache_omni 和 reshape_and_cache_with_stride 等变体,分别支持更灵活的 tensor 排布和跨步访问模式,满足注意力层的多样缓存需求。

ATB 组图机制允许将若干个算子合并为一个 GraphOp。对于 KV Cache 的写入和读取操作,可以将 ReshapeAndCacheOperation 和 PagedAttentionOperation 合并为一张子图,Host 端一次 Setup 完成全部内存分配和 Tiling 计算,一次推理过程中无需重复执行。

atb::Status CreatePagedKVCachePipeline(const atb::PagedAttentionParam &param,
                                         atb::Operation **operation) {
    atb::GraphOpBuilder *graphOpBuilder;
    atb::CreateGraphOpBuilder(&graphOpBuilder);

    graphOpBuilder->Init("KVCachePipeline",
                         inferShapeFunc,
                         {"key_in", "value_in", "block_table", "cache_k", "cache_v"},
                         {"attn_out"});

    auto cacheKV = graphOpBuilder->AddOperation(
        atb::ReshapeAndCacheOperation(param),
        {"key_in", "value_in", "block_table"},
        {"k_cache_block", "v_cache_block"});

    graphOpBuilder->AddOperation(
        atb::PagedAttentionOperation(param),
        {"k_cache_block", "v_cache_block"},
        {"attn_out"});

    *operation = graphOpBuilder->Build();
    atb::DestroyGraphOpBuilder(graphOpBuilder);
    return atb::NO_ERROR;
}

ATB的GraphOpBuilder将ReshapeAndCache和PagedAttention合并为一张子图。如果不组图,两个算子之间NPU Stream上会出现空泡——Host准备第二个算子的Tiling参数时NPU处于空闲。对于长序列推理,每次解码步减少一次Host到NPU的交互,在80层Transformer上累积80倍的节约。

ATB 在图算子 Setup 阶段通过 HBM 内存优化进一步压缩图内算子的总 Workspace 大小。内部实现基于内存 Block 的分裂、合并和尾块优化算法,将图算子中每个子算子的 Workspace 尽可能复用。前一个子算子的临时输出在后一个算子不再需要后立即释放,空间立即分配给后续子算子使用。ATB 仓库的文档说明通过这项优化平均节省 Workspace 50% 以上,提升了批处理推理的 batch size 上限。对于 KV Cache 相关的图算子而言,Workspace 的减少意味着有更多的 HBM 空间用于存储 KV Cache block,间接提高了服务容量。

KV Cache量化:INT8/INT4 KV Cache的精度-延迟权衡

KV Cache 占用显存的核心矛盾来自 FP16 的 16-bit 宽度。如果将其量化为 INT8(8-bit)或 INT4(4-bit),KV Cache 的 HBM 占用可以压缩至原来的 50% 或 25%。代价是量化引入的精度损失必须控制在模型困惑度和下游指标的可接受范围内。INT8 量化对大多数模型几乎没有精度退化,而 INT4 量化需要更精细的 Group-wise 校准。

ATB 库中的 PagedAttentionOperation 在其内部 AttentionFlags 结构体中提供了细粒度的量化控制。useQuant 和 useQuantOffset 控制是否启用 KV Cache 量化。useQKVQuantOffline 和 useQKVQuantOnline 区分了两种量化时机:离线量化在模型加载阶段将权重预量化,不增加推理时延;在线量化在推理过程中对 KV Cache 实时计算 scale factor 并写入,每个 token 生成后都需进行一次量化操作,增加了计算但获得了更大的灵活性。对于大模型推理场景,通常选择离线预量化以最小化在线开销,对于需要动态调整缩放因子的场景才选择在线量化。

在 Per-token 粒度下,每个 token 的 Key 和 Value 使用一个独立的 scale factor。对于 INT8 量化,scale factor 是一个 float32 数值。对于 INT4 量化,需要将每两个 INT4 数值打包为一个 INT8 存储以对齐 HBM 访问粒度。ATB 仓库中 paged_attention 算子的 qscale(quant scale)参数在 paged_attention_ops_runner.h 和 paged_attention_runner_utils.h 中多处出现,该参数默认值为 0 表示不启用量化。设置为非零值后,PageAttention 算子在加载 Cache block 时自动完成反量化,计算注意力输出后重新量化写入。qscale 的值来自 AMCT 工具链在模型校准阶段计算出的量化参数,在推理入口处通过 PagedAttentionParam 传递给算子。对于 Group-wise 量化,qscale 扩展为一个 scale 向量,每个分组对应一个 float32 值,每 32 或 64 个通道共享一个 scale。这种分组的出发点在于注意力头之间的分布差异——不同注意力头关注的语义模式不同,其 Key 和 Value 的动态范围也各异。通过 AMCT 在校准集上统计每个分组的 max 和 min 分布,qscale 向量可以更精确地覆盖全范围,在 INT4 压缩下仍然保持较低的量化误差。qscale 向量在 HBM 中紧邻 Cache block 存放,PageAttention 算子的 tiling 阶段会将其一并加载到 Cube Unit 的 Local Memory 中,避免二次访存。

struct atb::PagedAttentionOperation::AttentionFlags {
    bool useQuantOffset;      
    bool useQuant;            
    bool useBatchRunStatus;   
    bool useMask;             
    bool useQLens;            
    bool useRazorOffset;      
    bool useLogN;             
    bool useQKVQuantOffline;  
    bool useQKVQuantOnline;   
};

AttentionFlags中每一项都是独立的bool开关,允许按需组合。例如只启用useQuant和useQKVQuantOffline,KV Cache在模型加载阶段量化完毕,推理时只做反量化不做重量化,计算代价最低。需要在线调整缩放因子时才打开useQKVQuantOnline。Ascend NPU上INT8的Cube算力是FP16的两倍,量化后Cache数据量减半,HBM带宽压力减半的同时MatMul也跑得更快,延迟收益是乘性的。

ATB 仓库中的 mm_deq_swiglu_quant_mm_deq 算子节点展示了复杂的量化融合模式。这个融合算子将矩阵乘法的反量化、SiLU 门控激活函数、量化以及第二次矩阵乘法的反量化全部融合为一个 kernel。该算子在 src/ops/ops_infer/mm_deq_swiglu_quant_mm_deq/ 下实现,其 mm_deq_swiglu_quant_mm_deq_ops_runner.cpp 负责在 Ascend NPU 上执行完整的量化推理 pipeline。对于 KV Cache 量化的工程实施,这种融合模式提供了直接参考:注意力计算前反量化 KV Cache,利用 Cube Unit 完成 INT8 矩阵乘法,计算后重新量化写入 HBM,中间过程在 NPU 内部完成,不经 Host 参与。这对于减少 Host-Device 交互、缓解 Host Bound 瓶颈有直接作用。融合的思路不止于 KV Cache 量化,ATB 仓库中还包含 gmm_deq_swiglu_quant_gmm_deq 等系列运算,将分组矩阵乘法的量化流程也做了类似的融合优化。

在 INT4 方案中,精度损失控制是决定方案能否落地的关键。Group-wise 量化是缓解 INT4 精度损失的有效手段,将每 32 或 64 个通道分为一组共享同一个 scale factor,相比 per-token 量化在精度和压缩率之间取得了更好的平衡。ATB 的融合算子将反量化和矩阵乘法绑定,使得 scale factor 直接从量化参数中读取,无需额外的数据搬运。实际部署时 INT8 KV Cache 量化几乎不影响模型困惑度,而 INT4 量化需要结合通道级的 scale factor 和敏感的量化校准集,这部分通过 CANN 的 AMCT 工具链预先完成,再以量化参数形式传递给 ATB 算子。ATB 的内存优化机制确保量化 scale factor 存放在 HBM 中与 Cache block 相邻的位置,减少额外的访存开销。选择哪种量化精度不是一个绝对的决策,而是取决于模型对量化的敏感度、目标硬件上的 INT8 与 INT4 计算单元的吞吐差异、以及业务对延迟和精度的容忍程度。

Speculative Decoding:用小模型的"猜测"减少大模型的解码步数

标准自回归解码的串行特点决定了每步只能生成一个 token。对于一个长度为 T 的输出序列,需要 T 次大模型前向计算。大模型通常有 70B 甚至更多参数,单次前向的延迟在毫秒到十毫秒级别。Speculative Decoding 的出发点是用一个轻量级的 Draft Model(参数规模通常为 100M 到 1B)一次生成 K 个候选 token,Target Model 一次性验证这批候选的正确性。验证通过的全部接受,回退处由 Draft Model 重新猜测。理想情况下 T 步解码压缩为 T/K 步,大模型前向次数大幅减少。

ATB 库不直接打包 Speculative Decoding 的端到端实现,但其算子体系和 GraphOpBuilder 组图机制完全支持自行构建。仓库中的 logprobs_sample 采样算子(位于 src/kernels/kernels/logprobs_sample/)可以完成候选 token 的 top-k 采样,toppsample 和 multinomial 算子提供了基于概率分布的多样采样策略。这些采样算子与 PagedAttentionOperation 配合,可以一次性验证 K 个候选 token 的注意力分布。AttentionFlags 中的 useQLens 字段允许传入每个样本的真实序列长度,使得不同候选 token 的验证在同一张计算图中完成。

在构建 Speculative Decoding 的计算图时,Draft Model 的 Linear 层和采样算子的参数规模小,计算开销低。ATB 的 GraphOpBuilder 允许将 Draft 和 Target 的推理图拼接为一体:Draft 部分包含 Linear 或几个轻量注意力层加上采样算子,Target 部分包含完整的 Transformer 层栈加上 PagedAttention。两个部分共用同一组 KV Cache block,Draft 阶段写入的 cache block 在 Target 阶段直接读取。共享 Cache 避免了跨模型的显存拷贝,消除了数据搬移开销。这种共享策略的关键在于 Cache 的写入权限管理:Draft 的推测写入不能污染 Target 已有 Cache,ATB 的 block_table 机制通过拷贝被修改的页面或在写入后做版本检查来保证正确性。对于 J 个并行候选 token 的验证场景,PagedAttentionOperation 支持将所有候选序列拼接到同一个 batch 中一次前向完成,利用其动态 contextLens 和 useMask 标志区分不同位置的有效 token 范围。

atb::Status BuildSpeculativeDecodeGraph(const SpecDecodeParam &param,
                                         atb::Operation **operation) {
    atb::GraphOpBuilder *builder;
    atb::CreateGraphOpBuilder(&builder);

    builder->Init("SpecDecodeGraphOp",
                  inferShapeFunc,
                  {"input_ids", "kv_cache", "draft_weights", "target_weights"},
                  {"accepted_tokens", "new_cache"});

    auto draft_logits = builder->AddOperation(
        atb::Linear(param.draft_params),
        {"input_ids", "draft_weights"}, {"logits_draft"});
    auto draft_tokens = builder->AddOperation(
        atb::LogprobsSample(param.sample_params),
        {draft_logits[0]}, {"candidate_tokens"});
    auto target_logits = builder->AddOperation(
        atb::PagedAttentionOperation(param.attn_params),
        {"candidate_tokens", "kv_cache"}, {"logits_target"});

    *operation = builder->Build();
    atb::DestroyGraphOpBuilder(builder);
    return atb::NO_ERROR;
}

将Draft和Target两次推理合并到一张计算图中,ATB的调度器一次Setup阶段完成全部内存分配和Tiling计算。如果不组图,Draft和Target之间在NPU Stream上会出现多次空泡。Ascend NPU上Kernel发射的Host开销不可忽略,合并为一张图后Draft和Target在NPU上连续执行,Kernel之间几乎无空泡。

在昇腾 NPU 上部署 Speculative Decoding 时,Draft Model 的选择和候选长度 K 的确定需要权衡。ATB 的 Linear 算子和 PagedAttentionOperation 均支持动态 batch size,候选长度 K 可以在每次推理中灵活调整。当 Draft Model 接受率高时增大 K 以进一步减少解码步数,当接受率低时减小 K 以降低无效验证的开销。这种动态调度通过 PagedAttentionParam 中的 contextLens 和 qLens 参数实现,每次前向传入当前批次的实际序列长度,不浪费计算资源。

ATB 仓库中 example/op_demo 目录下的 Demo 展示了单算子的完整调用流程。对于 Speculative Decoding 这种多算子组合场景,可以复用 faupdate_demo.cpp 中的 context 创建、aclrtCreateStream 流管理以及 tensor 生命周期控制模式,在此基础上使用 GraphOpBuilder 组装 Draft 和 Target 推理子图。faupdate_demo.cpp 中先调 CreateContext 建立 ATB 上下文,再调 aclrtCreateStream 创建执行流,通过 SetExecuteStream 将流绑定到 context 上。执行完毕后先释放 Operation 对象,再销毁 Stream 和 Context,这个生命周期管理方式对 Speculative Decoding 的计算组织同样适用。对于更复杂的多步验证场景(如树形 Speculative Decoding),可以利用 ATB 的 Plugin 机制编写自定义的验证逻辑算子,以插件形式注册到 GraphOpBuilder 的算子库中,与原生 Linear、PagedAttention 等算子混合组图。

通算融合:AllReduce与MatMul的流水线重叠

在 Tensor Parallel 推理场景中,Attention 和 MLP 层的输出需要在多卡之间做 AllReduce 同步。Ring AllReduce 包括 ReduceScatter 和 AllGather 两个阶段,每阶段数据量等于单卡上隐藏维度对应的一块数据。在 Ascend 910B 的 8 卡服务器中,HCCS 互联链路提供约 56 GB/s 的跨卡带宽。对于 8192 隐藏维度的 FP16 输出,单次 AllReduce 通信量约 128 KB。虽然单次量不大,但每层 Transformer 需做两次 AllReduce(一次 Attention 输出,一次 MLP 输出),80 层模型累积 160 次通信,每个通信操作都有链路建立和同步等待的开销。

通算融合将通信隐藏到计算背后。做法是:将 Linear 层的输出矩阵在隐藏维度上切分,对第一部分启动 AllReduce 通信,同时让 Cube Unit 计算剩余部分。如果通信和计算耗时相当,通信完全被隐藏,端到端延迟只取决于计算路径。这要求在算子内部实现通信切分和计算切分的精细调度,并且两个路径的数据依赖关系要正确。切分的粒度非常关键:切得太细,通信启动次数增多,链路建立开销占比增大;切得太粗,通信窗口不足,Cube Unit 等待通信完成。通常按照与 HCCS 通信单元一次传输的粒度对齐来切分,使通信循环和计算循环的周期相互匹配。ATB 的线性 ParallelParam 中的 commType 字段提供了一个抽象通信语义,all_reduce、all_gather 和 reduce_scatter 三种通信模式都支持在同一框架下实现与计算的流水化。

ATB 库在 src/kernels/lcal/ 下提供了低层 HAL 通信算子,在 src/ops/ops_infer/ 下封装了高层的 all_reduce、all_gather、reduce_scatter 算子。lcal 目录中的 cmake、include 和 src 子目录共同构成了通信算子的编译和运行时基础。仓库中 linear_parallel 算子(位于 src/ops/ops_infer/linear_parallel/)直接包装了 Linear 和 AllReduce 的并行模式。开发者可以使用这个算子替代 Linear 加独立 AllReduceOperation 的顺序模式,由 linear_parallel 在内部自动编排通算重叠。linear_parallel 的实现将 MatMul 的输出沿隐藏维度切分为两份,通过两个不同的 Stream 或 Event 机制编排执行。除 basic 版本外,ATB 还提供了 grouped_matmul_inplace_add 等融合算子,可以在 AllReduce 前后插入 Element-wise 操作的流水化。

atb::LinearParallelParam parallelParam;
parallelParam.transA = false;
parallelParam.transB = true;
parallelParam.hasBias = false;
parallelParam.commType = atb::CommType::ALL_REDUCE;

atb::Operation *parallelOp = nullptr;
atb::CreateOperation(parallelParam, &parallelOp);

atb::VariantPack variantPack;
variantPack.inTensors = {hiddenStates, weight, allreduceBuf};
variantPack.outTensors = {output};

uint64_t workspaceSize = 0;
parallelOp->Setup(variantPack, workspaceSize, context);
uint8_t *workspacePtr = nullptr;
if (workspaceSize > 0) {
    aclrtMalloc((void**)&workspacePtr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST);
}
parallelOp->Execute(variantPack, workspacePtr, workspaceSize, context);
aclrtSynchronizeStream(stream);

LinearParallelParam中的commType字段决定通信语义。设置为ALL_REDUCE时,算子在Linear计算启动后即刻将切分出的第一份数据通过HCCS发起AllReduce,同时并行推进剩余部分的MatMul计算。Ascend NPU上HCCS引擎和Cube Unit分别占用不同的硬件资源,可以同时运行。但Stream层面需要CANN的aclrtSubscribeEvent机制创建依赖关系,不调用同步会导致AllReduce结果被后续RmsNorm读到不完整的数据。

ATB 仓库中的 mm_deq_swiglu_quant_mm_deq 融合算子也是一个通算重叠的微观案例。它将矩阵乘法后的反量化、激活函数计算、量化编码和第二次矩阵乘法打包为一次 kernel 发射,内部通过 AICore 上的指令级流水线实现计算和访存的深度重叠。这种微观融合和宏观通信重叠的思路一脉相承。在分布式推理场景中,AllReduce 通信可以与后续的 RmsNormAndRopeOperation 或 SwiGLU 激活计算重叠。ATB 的图模式允许将这些操作编排在关联的流上,CANN runtime 负责底层调度。对于 Ascend 910B 的 8 卡配置,通算融合可以将 AllReduce 的通信开销压低至 10% 以下,从软件层面跨越通信墙。类似的思路也体现在 grouped_matmul_with_routing 和 grouped_matmul_inplace_add 两个算子中,它们在 MoE 稀疏计算场景下将多个专家的矩阵乘法与 AllToAll 通信做了流水化编排。

效率对比

ATB 提供的各项优化从不同维度切入推理延迟问题。KV Cache 量化着力于 HBM 带宽和容量瓶颈,Speculative Decoding 针对解码阶段的串行计算瓶颈,通算融合解决分布式推理的通信墙。下表从四个维度概括优化前后的效果(基于 Ascend 910B + LLaMA-70B 场景的工程估算,具体数值因模型规模、batch size、序列长度和硬件配置而异):

维度 优化前 优化后 差异来源
KV Cache 显存占用 FP16 格式占满 HBM,batch 上限受限 INT8 量化后显存减半,batch 上限翻倍 PagedAttention 的 useQuant 标志开启,Cache 以 INT8 存储,读取时反量化,写入时重新量化
单步解码延迟 每 token 需完整大模型 FP16 前向 INT8 量化结合 Speculative Decoding 后延迟降低 40-60% Cube Unit 的 INT8 MatMul 是 FP16 的两倍算力,解码步数因验证压缩而减少
AllReduce 通信开销 每层两次 AllReduce,通信占比 20-30% 通算重叠后通信占比降至 8-12% AllReduce 与 MatMul 通过 HCCS 和 Cube 并行,linear_parallel 算子自动编排
HBM 碎片率 连续分配导致碎片,利用率 50-60% Page block 管理后利用率达 80-90% ReshapeAndCache 和 PagedAttention 的分页机制按 512 token/page 粒度管理

表中每项差异直接关联 ATB 库的具体算子能力。KV Cache 量化依赖 PagedAttention 中 AttentionFlags::useQuant 等标志位和 mm_deq_swiglu_quant_mm_deq 的量化融合模式。Speculative Decoding 的可组合性基于 GraphOpBuilder 的组图机制和 LogprobsSample、PagedAttentionOperation 等算子的动态 batch 支持。通算重叠依赖 linear_parallel 和 all_reduce 算子的协同调度以及 HCCS 链路和 AICore 的物理并行特性。页式缓存碎片优化来自 ReshapeAndCacheOperation 和 PagedAttentionOperation 的 block 管理模式,按固定大小划分 block,空闲回收和按需分配由 CACHE 算子内部管理。每个 block 的大小在模型加载时根据 hidden_dim 计算确定,一般为 512 x hidden_dim 的连续分片。

这些优化之间存在相互影响和协同。KV Cache 量化减小了 HBM 读写量,为 Speculative Decoding 腾出了更大的 batch 空间,可以并行验证更多候选 token。Speculative Decoding 减少了解码步数,降低了 AllReduce 的调用频次,间接减轻了通算重叠的实现压力。通算重叠压缩了每步计算的通信等待时间,使得大模型在更多卡之间做张量并行时扩展效率更高。每一项优化都依托 ATB 库的组图调度、Tiling Cache 复用、Workspace 减量等运行时机制,共同构成了一条从显存管理到计算流水化的完整推理加速链路。在具体的工程实施中,选择从哪一条路径切入取决于当前的瓶颈类型:显存不足时优先 KV Cache 量化,延迟敏感时优先 Speculative Decoding 或通算融合,而这种选择本身也需要结合 ATB 库的运行 profilling 数据来确定。

结尾

KV Cache 的显存瓶颈通过 INT8 和 INT4 量化得到缓解,PagedAttention 中的精细量化控制和融合算子将精度损失约束在可接受范围。Speculative Decoding 借助 GraphOpBuilder 的组图和采样算子的灵活组合,将解码步数压缩至原来的数分之一,直接减少了算子的调用频次。通算融合利用 HCCS 链路和 Cube Unit 的硬件并行能力,将 AllReduce 通信隐藏到 MatMul 计算路径后方。ATB 库提供的 Tiling Cache 复用、HBM 内存优化和组图调度机制,为每一项优化提供了基于 CANN 和昇腾 NPU 硬件特性的工程化落地通路。


仓库地址:https://atomgit.com/cann/ascend-transformer-boost

Logo

鲲鹏昇腾开发者社区是面向全社会开放的“联接全球计算开发者,聚合华为+生态”的社区,内容涵盖鲲鹏、昇腾资源,帮助开发者快速获取所需的知识、经验、软件、工具、算力,支撑开发者易学、好用、成功,成为核心开发者。

更多推荐