混元AI Infra团队 投稿
量子位 | 公众号 QbitAI
大模型竞赛中,算力不再只是堆显卡,更是抢效率。
面对H20等推理卡在主流算子库下难以跑满性能的痛点,腾讯混元AI Infra团队正式开源生产级高性能LLM推理核心算子库HPC-Ops。
该算子库采用CUDA和CuTe从零构建,通过抽象化工程架构、微架构深度适配及指令级极致优化等,降低底层算子开发门槛,将核心算子性能逼近硬件峰值,实现了显著性能突破。
在真实场景下,基于HPC-Ops,混元模型推理QPM提升30%,DeepSeek模型QPM提升17%。
同时,在单算子性能方面,HPC-Ops实现Attention相比FlashInfer/FlashAttention最高提升2.22倍;
GroupGEMM 相比DeepGEMM最高提升1.88倍;FusedMoE相比TensorRT-LLM最高提升1.49倍。
主流算子库亟需更适配的底层支持
在大模型时代,计算效率已成为AI应用及发展的关键瓶颈。
目前主流算子库(如FlashInfer、DeepGEMM)多以NVIDIA H800等高配训练卡为首要优化目标,但限于客观原因,不少大模型的大规模线上推理服务只能采用H20等推理型计算卡。
现有SOTA算子库在这些显卡上往往难以发挥硬件峰值能力。
同时业务侧对极致吞吐、低延迟以及Blockwise FP8等复杂量化策略的需求日益迫切,亟需更适配的底层支持。
总结来看,现有主流算子库主要存在以下痛点。
使用成本高
主流算子库设计复杂,核心Kernel封装深,在其上修改适配成本非常高,除了对代码非常熟悉的开发者,普通的AI研究者很难在其上适配修改。
而大模型的很多加速研究创新,比如量化算法和投机采样等方法都严重依赖于与之匹配的高效算子实现。
比如最开始4bit和8bit的量化算法出来后,虽然理论上加载数据量减少,但由于没有与之匹配的低精度算子实现,低精度量化在很长的一段时间内都是负优化。
目标硬件不匹配
现有的主流算子库都是以H800等显卡为目标优化、NVIDIA提供的CUTLASS等算子更是以Blackwell架构为目标,而目前国内主流的推理显卡则有所不同。
不同硬件间算力带宽的差距导致Kernel的优化方法也会不同,因此现有的算子库在国内主流推理卡上的表现并未发挥出硬件的全部性能。
基于以上问题,腾讯混元使用CUDA和CuTe开发了一套轻量、高效的LLM核心算子库。
用CUDA和CuTe从零构建
该算子库主要包括FusedMoE、Attention、机内/机间通信、Norm、Sampler、以及各类小算子的融合算子模块,整体算子库架构如下图所示。
通过分析任务特性和硬件微架构,将任务的划分逻辑与硬件指令做了更好的对齐,以此获得更好的性能,并且对工程代码进行了适度的抽象,让开发者能聚焦于算法本身,降低维护门槛。
该算子库不仅是高性能生产工具,也可作为开发者深入理解CUTLASS与CuTe工业级开发的实践范本,具体的技术细节如下。
任务特性与硬件能力对齐
针对访存瓶颈的算子,其性能主要受限于数据加载速度。
针对国内的主流推理显卡,通过调整指令发射顺序进行数据预取优化,确保数据传输单元一直处于高利用率。
针对不同的问题规格做了更细致的指令对齐和优化,去除冗余低效指令以减少算力的浪费,如针对Decode Attention和小batch下的GroupGEMM都做了AB矩阵交换的优化;
以此对齐到硬件架构上的wgmma指令,访存带宽可达到硬件峰值能力的80%以上。
△Attention SwapAB示意图精细的任务调度和数据重排
针对每个算子问题,都重新思考了任务数据的划分调度策略,尽可能保证每个SM都任务均衡的同时兼顾cache的连续性。
并且采用了persistent kernel的方式隐藏kernel prologue和epilogue的开销。
另外也通过数据重排减少了额外的操作和显存占用。
比如在FP8 Attention Kernel中创新性采用了Interleave重排技术,解决了指令不匹配的问题,减少线程间数据shuffle,获得了优于业界SOTA的算子性能。
聚焦于计算逻辑本身
GPU编程的复杂度很大程度上来源于操作的复杂性,为了能使用高效指令,一般需要对数据进行多次的重解释和变换等编程技巧,这大大加重了开发者的心智负担。
因此基于CuTe扩展开发vec抽象层统一负责高效数据搬运,利用Layout代数抽象隔离复杂的Tiling与计算逻辑,让开发者能聚焦于算法本身,降低维护门槛。
关键实验结果
通过以上高效算子实现,在混元模型上将QPM端到端提升30%,DeepSeek上QPM提升17%。
同时针对LLM中核算子模块进行了测试,以常用的模型规格(混元、DeepSeek)进行了测试,并对比了目前主流的算子库实现。
实验表明,在LLM的核心模块Attention和FusedMoE上的性能都超越当下SOTA实现。
GroupGEMM
与 DeepGEMM (v2.2.0) 的两种版本进行对比,在Batch<=64的低延迟场景下优势显著,较DeepGEMM最佳表现最高提升1.88x。
且通过流水线掩盖技术Blockwise与PerTensor性能几乎持平;
在大Batch场景下,亦能保持约1.1x的领先优势。该算子同时兼容紧密排布与Token不连续输入,显著减少临时显存用量。
△GroupGEMM性能对比图FusedMoE
完整封装了包括前序数据重排、GroupGEMM及后续Reduce加权平均在内的全流程模块.
并在序列长度取16倍数的均衡分配规格下,对比了vLLM (v0.11.0)与TensorRT-LLM (v1.1.0)的实现。
测试结果显示,该FusedMoE模块在TP场景下相比TensorRT-LLM提升显著,最大性能提升达1.49x;在EP模拟均衡场景下最大提升1.09x。
针对不同输入长度采取的差异化重排策略,进一步确保了整体模块在各类规格下的最优表现。
△FusedMoE性能对比图Attention
针对Prefill场景,测试128~64K的输入长度。
在batch较小时,BF16精度下相比SOTA实现提升1.3x;在大batch时基本与当前SOTA对齐。
针对Decode场景,根据线上SLO约束,搭配一组batch和输入长度的测试用例,BF16精度下提1.35x~2.22x;
FP8精度下,当Sequence Length较小时与SOTA相当,当Sequence Length较大时相比SOTA提升1.09x~2.0x。
△Attention性能对比图算子库当前能力和未来发展方向
作为面向大模型推理场景的高性能算子库,HPC-Ops凭借Attention、FusedMoE、GroupGEMM等核心算子的极致优化,达成最高2.22倍的性能提升,且已在腾讯大规模生产环境中完成验证。
其简洁易用的API可无缝对接vLLM、SGLang等主流推理框架,原生支持BF16、FP8等多精度量化方案。
同时还以CuTe、CUTLASS为基础,提供了数百行代码即可构建SOTA算子的实践范例,为开发者降低了高性能CUDA内核的开发门槛。
在未来的发展规划中,HPC-Ops将持续深耕大模型推理性能的突破方向。
一方面,将重点研发稀疏Attention算子,针对性解决长上下文大模型的内存与算力瓶颈;
另一方面,会拓展更丰富的量化策略,覆盖4bit/8bit混合精度等更多量化方案,进一步平衡推理速度与模型精度。
此外,算子库还将布局计算-通信协同优化的内核,通过融合多GPU间的计算逻辑与通信流程,大幅降低分布式推理场景下的通信开销,为超大规模大模型的高效部署提供底层支撑。
目前,HPC-Ops已在GitHub开放源码供开发者下载使用。
同时腾讯混元Infra团队也表示,欢迎行业内的技术实践者提交高价值PR,参与算子边缘场景优化、教程案例打磨等精准化贡献,共同推动大模型推理技术的边界拓展。
GitHub项目地址:https://github.com/Tencent/hpc-ops