基于算子优化的大模型高效推理落地实践

01

背景：推理优化的分层栈与算子地位

大模型推理性能优化的全栈视角下，优化层次可划分为四层：应用层负责请求分发与调度；框架层依托 vLLM、TensorRT-LLM 等推理框架完成批处理与 KV Cache 显存管理；图优化层在计算图层面执行常量折叠、死代码消除与算子融合；算子优化层则进一步打开每个算子内部细节，直接对接硬件微架构，是性能挖潜的最后一道关口。

算子优化层之所以关键，在于其直接面向硬件执行单元。GPU 算子的编写需明确每个线程的计算任务——将完整张量计算拆分至 Grid、Block、Thread 各层级，这决定了计算资源能否被充分利用。当前算子优化的主要瓶颈已非计算单元吞吐，而是“内存墙”——计算单元等待数据的空闲周期占据了可观比例。因此，访存优化往往成为算子调优的核心战场。

算子优化的核心技术可归为三类：

• 算法创新：从数学原理层面重新设计计算方案。FlashAttention 通过在线 Softmax 算法将中间数据始终保留在高速片上缓存，同时将 O(N²)显存占用降至 O(N)，兼具算法重设计与访存优化的双重收益。

• 算子融合：将多个独立算子合并为单一内核，消除中间结果的全局内存读写。FlashAttention 本质上亦是一次算子融合，只是需借助在线 Softmax 才能达成融合效果。

• 常规编写优化：涵盖计算任务拆分、访存指令优化（合并访存、向量化加载）、计算指令优化（利用 Tensor Core 及 SFU 等专用硬件单元）以及软件/硬件流水线优化（数据预取重叠访存与计算、利用 Warp 调度器隐藏延迟）。

中兴通讯的推理业务横跨 NLP、CV、语音等多方向，硬件覆盖云端（NVIDIA A100/H20）、边缘（NVIDIA L20）与端侧（Jetson Orin、瑞芯微 NPU、海光 GPU 等），算子供应的异构性使得优化策略需因地制宜。

02

挑战：开发工作量爆炸与厂商软件栈封闭

算子优化的第一个挑战来自开发工作量的膨胀。一方面，不同硬件微架构与输入形状的组合形成了庞大的最优实现空间。芯片代际更迭带来的架构差异（如英伟达 GPU 从 Volta 到 Blackwell，SM、Tensor Core、共享内存容量均有巨变），以及推理中批量大小、序列长度的动态变化，使得同一算子的最优配置不断漂移。以 CUTLASS 矩阵乘为例，需同时协调线程块分块、指令形状、流水线级数等数十个参数，组合空间极其庞大。另一方面，算子开发本身是多领域知识的交汇——需通晓算法、硬件微架构、驱动运行时、内存层次乃至 PTX 指令级调试，一个 CUDA GEMM 从朴素实现到接近 cuBLAS 性能，往往要经历十余步精细调优，周期漫长。

第二个挑战来自厂商软件栈的封闭性。许多芯片厂商仅提供高层 API，用户无法干预内部算法选择与优化决策：无法将相邻算子融合为单一内核；无法针对特定输入形状调整实现（厂商内置启发式策略在非典型形状下可能选择次优内核）；当出现精度异常或回退 CPU 等问题时，用户无法定位根因并修复，排查成本极高。

03

落地实践：五条路径应对双重挑战

1. 算子迁移：跨硬件复用的快速通道

面对为海光 K100 GPU 优化 GroundingDINO 2 模型中 MSDeformAttn 算子的需求，团队并未选择从零编写 HIP 代码。由于海光软件栈兼容 AMD ROCm 生态，而 Deformable DETR 算子库中已存在 CUDA 实现，团队通过 hipify 工具自动将 CUDA Runtime API 转换为 HIP Runtime API，在海光 K100 上编译验证后，通过 Custom OP 接口注册至 ONNX Runtime 的 ROCm 后端，随后针对 K100 的 SM 数量等微架构特征做深度调优。

这一迁移路径的效果是阶梯式的：简单迁移与集成后，算子推理延迟从分解为基本算子组合时的 4.84ms 降至 0.9ms；针对海光微架构完成适配调优后，进一步降至 0.17ms。算子迁移的本质，是利用不同硬件软件栈之间的同构性（如 CUDA 与 HIP 的 API 相似度），将已验证的高质量实现跨硬件复用，避免重复造轮子。

2. 预编译搜索+在线适配：为每种输入形状匹配最优配置

CUTLASS 等开源算子模板库虽提供了高度参数化的 GEMM 实现，但其默认策略是对所有输入形状使用固定线程块与 Warp 形状。当矩阵 M 维度较小时（对应小 batch size 或小序列长度），固定的大 Tile 会引入大量 Padding 导致无效计算。在 Qwen3 模型 W4A4 量化推理场景中，这一缺陷使得小输入形状下的 INT4 GEMM 较 W8A8 方案劣化高达 57%。

团队的解决方案分三阶段：准备阶段，基于 CUTLASS 模板参数化算子配置，用一份代码生成多种配置的算子实例；离线阶段，为目标硬件遍历所有配置与输入形状的组合，构建“甜点配置表”（即每种输入形状对应的最优参数组合）；在线阶段，推理时根据实际输入形状查表选择最优配置，零开销切换至对应内核。

在 L20 硬件上，优化后的 W4A4 方案不仅解决了小输入劣化问题，更实现了全面优于 W8A8 的性能，batch=1 时从 0.473ms 降至 0.131ms，提速 58%。

同样的策略被应用于 Qwen3.5-397B 模型在 H20 集群上的 FP8 GEMM 算子优化。Prefill 阶段序列长度大、M 维尺寸大，属计算密集型，需大 Tile 配置；Decode 阶段序列长度小、M 维尺寸小，属访存密集型，需小 Tile 配置。预编译搜索为两阶段分别匹配不同内核，最终 TPOT（Time Per Output Token）延迟优化 8.9%。

3. 常规优化：访存、流水线与硬件指令的三管齐下

当算子迁移与搜索调优仍不足以拉满性能时，常规优化手段便成为必选项。以 BigVGAN 语音生成模型中的抗混叠激活算子为例，其原生 CUDA 实现较为粗糙，团队在 L20 硬件上进行了五个层面的优化。

共享内存分担寄存器压力：原始实现将中间结果全部存储在寄存器中，寄存器溢出后数据被换出至 HBM 的本地内存，引入高访存开销。优化后将中间结果存入共享内存，释放寄存器资源。

消除 Bank 冲突：引入共享内存后，张量列访问导致同一 Warp 内不同线程访问同一 Bank，产生冲突。通过每行 Padding 一个空元素的方式，使得线程访问分散至不同 Bank。

向量化访存指令：原始实现逐元素写回全局内存（STG.32 指令），优化后将四个 float 打包为 float4，一次写入128bit（STG.128 指令），指令发射次数减少，访存效率提升。

硬件流水线优化：原始实现每线程处理 32 个元素，负载过重导致每个 SM 上可驻留的活跃 Warp 数量减少，不利于 Warp 调度器通过切换线程束来隐藏访存延迟。优化后每线程仅处理 4 个元素，寄存器用量下降，SM 上可驻留 Warp 数增加，硬件流水线的延迟隐藏能力得以充分发挥。

计算指令优化：原始实现中三角函数与指数运算由多条 CUDA 指令组合完成，优化后改用 SM 中 SFU（Special Function Unit）硬件单元的原生指令。SFU 专为三角函数与超越函数设计，以可接受的精度牺牲换取三倍吞吐。最终，在五重优化的作用下，端到端推理延迟优化 15%，单算子延迟优化 38%。

4. 等效替换与厂商协作：封闭生态下的治理闭环

面对瑞芯微 NPU 等厂商底层 AI 软件栈（编译器、运行时、算子库）闭源的现实，团队摸索出两条治理路径。

等效算子替换：当定位到问题算子后，在 PyTorch 层面将其替换为语义等价的替代实现，间接影响下游软件栈的行为。在 distill_anything_v2 模型中，GeLU 算子和特定属性的 Resize 算子被软件栈回退至 CPU 执行，引入大量 CPU-NPU 数据搬运开销。团队将 GeLU 替换为 Tanh 近似实现，将特定属性的 Resize 算子替换为其他属性组合，成功消除回退，实现了大幅推理加速。在 NebulaVoice TTS 模型中，掩码与 Softmax 算子存在精度误差导致输出语音异常，团队将稀疏矩阵乘实现的掩码替换为 torch.where 实现，将 Softmax 分解为 topk+exp+sum 手动实现，虽带来轻微性能牺牲，但确保了结果正确性。

推动厂商升级：当问题根源深埋于厂商运行时库实现缺陷时，团队选择与厂商协作。在 StepAudio TTS 模型中，Pad 算子全部回退至 CPU 执行，而厂商文档却标明支持 NPU 调度。团队提供可复现材料，推动厂商定位根因——通道方向的 Pad 存在对齐限制，当前运行时库未能将其调度至 NPU。厂商完成 V1 版本修复后，Reshape 算子又出现类似问题，团队继续推动 V2 版本修复。这一协作路径虽然周期较长，但能从根本上解决软件栈缺陷，惠及后续所有模型落地。

04

未来规划：算子优化的自动化演进

团队预判算子编写将必然从手工迈向自动化，这一演进在两条技术路线上同时推进。

基于算子级编译器的自动编译优化。Triton 作为生态最完善的代表，将类 Python 的 Tile DSL 编译为 LLVM IR 直至机器码，自动处理共享内存管理与软件流水线，已集成至 PyTorch 的 torch.compile 中。TileLang（复用部分 TVM 编译基础设施）分离数据流与调度，支持三级编程粒度，多后端覆盖面广。FlagTree 作为 Fork 自 Triton 的统一编译器，通过 TritonGPU 与 Linalg 两条编译路线兼容 13 个硬件后端，插件化设计支持芯片厂商自维护后端模块。cuTile 则是英伟达官方编译器，面向 Blackwell 架构的 TMA 与 Tensor Core 提供 30+MLIR 优化 Pass。

现状是，现有算子级编译器仍主要服务于 GPU 生态，对 DSA 架构芯片的支持尚处早期甚至空白。团队的目标是基于开源编译器适配异构芯片，开发硬件相关优化 Pass，显著降低算子开发与迁移成本。

基于 LLM Agent 的算子自动生成。LLM 天然适合压缩难以形式化的硬件专家知识——内存层次、指令选择、并行策略——这些知识恰好难以用传统规则表达。结合 Agent 闭环（生成→编译→测试→反馈→再优化），可将算子开发转化为自动化搜索问题。

当前相关研究已形成三个层次：基础设施层（KernelBench、TritonBench 等数据集与基准）、专用 LLM 层（KernelLLM、AutoTriton、Kevin 等）、Agent 系统层（CUDA-LLM、EvoEngineer、KernelGen 等）。但同样，这些工作主要服务于 GPU/Triton 生态。团队计划构建面向异构芯片的算子编写 Agent 系统，填补这一生态空缺。

05

总结

大模型推理的降本增效诉求将算子优化推至前台。面对开发工作量大与厂商软件栈封闭两大核心挑战，团队通过七个真实落地案例验证了五条差异化策略的有效性：算子迁移实现跨硬件复用，预编译搜索与在线适配应对输入形状动态变化，常规优化挖掘硬件极限性能，等效替换绕过闭源软件栈缺陷，推动厂商升级根除底层问题。展望未来，算子级编译器与 LLM Agent 的结合将为异构硬件生态带来自动化编写能力，从根本上重塑算子优化的开发范式。

分享嘉宾

INTRODUCTION

施霄航博士

中兴通讯

AI 算法工程师

博士毕业于南京大学计算机科学与技术学院，研究方向为边缘计算，读博期间在国际高水平会议及期刊上发表多篇学术论文。2025 年通过中兴通讯"蓝剑计划"加入端侧推理团队，主要从事编译优化与算子优化相关工作，致力于推动大模型推理优化技术的创新、研发和在异构设备上的工程落地。

点个在看你最好看

SPRING HAS ARRIVED