国产AI芯片上的MoE大模型部署与性能调优实战

上个月我们团队接到一个有点特殊的任务把刚开源的新一代MoE大模型V4完整部署到新进场的那批国产AI加速卡上推理吞吐目标是不低于同价位GPU平台的七成。说实话接到任务那一刻我心里并没有底。过去几年我们团队在CUDA生态里做模型推理优化做得太顺手了算子库要什么有什么调试工具链齐全坑都被人踩平了。换成国产AI芯片之后算子覆盖到什么程度、图编译器顺不顺手、显存管理能不能扛住千亿级参数的MoE架构这些问题只有真机跑过才会浮出水面。这篇文章不打算讲概念就记录我们这次从拿卡、装环境、跑通、再到把性能一点点抠出来的完整过程。如果你也准备把大型模型迁到国产AI芯片上那这里的坑大概率你可以提前避开一大半。1. 适配前必须盘清楚大模型搬到国产芯片难的不只是算子很多人以为模型迁移就是把权重拷过去、推理引擎换一下跑起来就完事。真上手之后才知道这个链条远比想象中长。大模型推理不是把整个模型当成一个巨大的黑盒算子去执行而是几百个小算子按计算图串起来协同工作。任何一个算子在目标芯片上缺实现、有精度bug或者性能劣化都会直接影响端到端的效果和延迟。所以我拿到适配任务后做的第一件事不是急着写代码而是先把这个模型在计算图上彻底拆明白。1.1 算子覆盖度一半的坑都埋在这里MoE模型本身相比传统稠密模型多出不少特殊计算。路由计算负责决定每个token分发到哪些专家GroupedGemm负责对不同专家做批量矩阵乘还有token级别的MoEScatter和MoECombine这类带索引的动态调度算子。这些算子不是标准矩阵乘法那种规整操作非常依赖目标平台的算子库有没有成熟实现。我们当时写了一个算子扫描脚本把PyTorch计算图导出ONNX之后用图遍历把所有算子节点拆出来再和芯片厂商提供的算子支持列表逐条比对。最终发现两项算子在官方算子库里没有高效率实现其中一项在首次迭代时被迫走了CPU回退。CPU回退这个词看起来轻描淡写实际上后果很严重——回退算子在每个step都会被调用一次回退带来的Host-Device同步开销就有可能把整体性能拖慢一个数量级。如果不提前做扫描这种问题会隐藏在庞大的日志里排查成本极高。注意算子扫描这一步做得好后面少熬三个通宵。别嫌麻烦一定要把模型里每个自定义算子都和芯片厂商的支持列表对一遍并记录实现方式、精度要求、是否有后端kernel。1.2 图编译的翻译再润色问题模型计算图进到芯片平台之后不是直接执行的。芯片厂商一般会有一层中间IR把ONNX或者PyTorch导出的图转换成自己的中间表示再做算子融合、layout转换、内存规划最后生成可执行文件。这个过程很像先翻译再润色中间任何一步不满足约束都得打回重来。我们碰到最头疼的就是动态shape约束。ONNX导出的图里如果sequence_length这个维度是动态的——这是大模型推理的常态prompt长度和生成长度都不可预知——很多图编译期的优化就无法做了。常量折叠算不了内存池分配没法精确规划有些算子融合规则也会因为shape不确定而直接失效。我们的解决办法是牺牲一点通用性把推理切成prefill和decode两个阶段prefill阶段按上限长度固定batch内最大token数decode阶段每次只处理当前beam的增量shape。虽然多了一些分支逻辑但换来了图编译优化的大幅提升。1.3 显存与KV Cache管理千亿参数的胃口大模型显存占用有三大块权重参数、中间激活、KV Cache。MoE模型权重虽然分散在几百个专家里但总量依旧可观激活值在prefill阶段会随序列长度急剧膨胀而KV Cache在长对话场景下是最大的显存吞金兽。芯片平台的显存管理SDK和成熟GPU生态相比差距不小。GPU生态里显存分配、释放、池化、碎片整理都做得非常成熟开箱就能用。但我们这次用的国产卡SDK动态分配和释放的API存在几个比较隐晦的问题频繁alloc/free会增大碎片率显存池的复用策略有些呆板多线程并发分配时偶尔还会出现段错误。最后我们干脆不走常规路线自己写了一个预分配的显存池管理器按块把KV Cache规划成固定大小的环形buffer彻底绕开SDK动态分配的问题。这个改动虽然前期花了两天时间但把后续调优阶段的显存问题一次性地解决了。1.4 工具链成熟度缺什么就自己造什么以前在GPU生态里排查问题性能分析器一开算子耗时、内存带宽、卡利用率全都清清楚楚。迁移到国产芯片后厂商虽然提供了一些性能分析API但粒度粗、维度少很多关键指标拿不到比如逐算子时间戳、SM级别的流水线stall、内存写回冲突等。试过之后发现不顶用我们转而自己造了一套轻量级微基准工具针对每个耗时算子单独构造典型shape的输入端到端地测三遍取中位数再把结果和理论峰值带宽做对比。这套工具看起来简陋但实际帮我们定位了好几个性能瓶颈——比如某个融合算子虽然功能正确但显存访问模式是AoS而非SoA带宽利用率不到三成。这些数据如果只靠厂商profiler根本看不出来。2. 环境搭建与首次跑通从零到一的完整路径芯片平台迁移里有一个特别容易被低估的环节环境搭建。GPU生态经过这么多年的沉淀驱动、容器运行时、通信库、编译器的版本彼此兼容得很好装一个镜像基本开箱即用。国产芯片生态这点上还在追赶版本矩阵踩坑概率高很多。我们一开始就吃过版本不匹配的亏所以后来立了一条规矩任何一台机器动手装环境之前先锁定一套经过验证的版本基线。2.1 版本矩阵先行别用最新版赌运气我们最终稳定下来的版本组合大概是这样的驱动版本固定在某大版本的小版本号固件升级到对应配套版本推理框架锁定在一个发布版本图编译器使用与该发布版本配套的版本。这个组合看着平淡无奇但它是我们反复重装了几次系统之后才试出来的。版本不匹配的典型症状很有迷惑性有时候是随机段错误有时候是HBM读写报错但板卡状态正常有时候模型加载到一半进程直接退出日志里连个error级别都打不出来。更坑的是这些症状不一定每次都能复现排查难度极大。后来我们意识到这大概率是固件和驱动之间的同步问题升级固件之后症状消失。自那以后版本矩阵就成了我们所有机器上线的第一道检查项。实操心得国产AI芯片的软件栈迭代速度很快尤其是固件、驱动、算子库这三个组件经常同步更新。上线之前一定要确认三者版本配套最好跟厂商技术支持确认一份推荐基线不要想当然选最新版本。2.2 推理引擎选型官方闭源、图编译自研、还是开源框架原生支持适配到国产芯片上推理引擎的选择基本上是三条路。第一条是直接用厂商官方提供的闭源推理引擎开箱即用对自家芯片优化最激进但自定义程度低想改内部调度策略很难。第二条是走图编译路线把模型转成芯片IR之后自己封装执行器控制力最强但是工作量大而且需要对硬件底层有相当深的理解。第三条是看看开源推理框架是否已经原生适配了这款芯片如果适配顺利生态和社区支持都比较好但适配质量取决于厂商和开源社区的合作深度。当时我们的判断是千亿级MoE模型对推理灵活性要求很高——需要精细控制prefill和decode的资源占比、自定义KV Cache管理、调整专家并行策略——这些用官方闭源引擎很难完全满足。最后选择了第二条路线在图编译之上自己包了一层控制逻辑。代价是初期开发量增加但回报是后续调优的每一分性能都握在自己手里。2.3 模型转换实操ONNX导出、静态化、算子替换模型转换是整个链路里最繁琐的一步。我们先用PyTorch导出ONNX然后做静态化处理固定推理时的部分维度import torch # 导出ONNX固定batch1和max_seq_len2048 dummy_input ( torch.ones(1, 2048, dtypetorch.int64), torch.ones(1, 2048, dtypetorch.int64), ) torch.onnx.export( model, dummy_input, v4_static.onnx, opset_version17, input_names[input_ids, attention_mask], output_names[logits, hidden_states], dynamic_axesNone, # 全部固定换取图编译优化空间 )导出静态图只是第一步。接下来需要用图转换工具把ONNX转成芯片IR这一步最容易出问题——有些算子图转换器不支持有些算子被拆成了若干子操作导致性能下降有些算子精度匹配不上。我们的做法是维护一个算子映射表逐一把模型里的算子映射到芯片IR的原生算子或者组合实现。GroupedGemm这类特殊算子是重灾区映射表里总是优先选择原生支持路径只有原生路径缺失时才落到通用分解实现。KV Cache的layout也在这个阶段改掉了。原模型里的KV Cache布局是PyTorch默认的[batch, num_heads, seq_len, head_dim]但芯片IR上跑得快的布局往往是[batch, seq_len, num_heads, head_dim]或者分块后的[batch, block_idx, num_heads, block_size, head_dim]。layout转换在GPU生态里经常是性能优化的起点在国产芯片上同样有效只是API使用方式不同需要自己踩一遍。2.4 第一次跑通的那一刻性能还很惨经过两天多的折腾我们终于看到了模型生成的第一个中文句子。那一刻的喜悦持续了大概三秒钟——因为紧接着跑出来的基准数据让人心凉单batch、序列长度32的条件下首token延迟比目标差了大约12倍每秒生成token数低得难以接受。不过逻辑正确性是好的生成内容、概率分布和GPU平台上的输出基本对齐说明算子和精度没有根本性问题。当时我其实松了口气。干这行的人都明白性能差不可怕可怕的是方向压根不对。只要能出正确结果说明整条链路已经打通剩下的就是性能调优——往哪个方向挖、能挖多深那是下一个阶段的事。我们把这套首次跑通的配置命名为baseline后面的所有优化都以它作为基准。3. 性能调优实录从能跑到跑得快性能调优阶段是整个适配过程里最考验工程经验的环节。和之前做GPU推理优化类似思路是有章法的先看精度模式对不对再做算子级融合然后压显存和KV Cache最后调并行策略和并发参数。每一步都要拿数据说话不能靠猜。3.1 精度模式FP16、BF16还是混合精度大模型推理阶段不像训练那样有梯度累积但精度选择依然会显著影响显存和计算速度。FP16能省一半显存但MoE模型里有些参数分布较广直接用FP16会出现溢出生成质量肉眼可见地变差。BF16的指数范围和FP32一致更稳但一开始我们发现芯片对BF16的原生计算支持没有FP16充分某些算子在BF16下甚至会退化为低效率实现。最后的方案是混合精度绝大多数GEMM算子走BF16数值敏感的归一化和部分注意力计算保留FP32同时给某些前向算子加动态缩放做保护。这个策略说起来简单但落实起来需要在算子映射表里逐项标注精度偏好再写脚本自动生成对应的IR配置。折腾完这轮生成质量和baseline对齐显存占用比纯FP32版本下降了约四成。3.2 图编译进阶算子融合和内存复用算子融合在国产芯片上的收益比在GPU上更明显。原因很简单芯片显存带宽相对有限每减少一次中间结果的落盘和重读收益就大一分。我们做的最有价值的一个融合是把RMSNorm、RoPE旋转位置编码和随后的QKV投影GEMM合并成一个复合算子。原来这三个算子分别执行中间结果要写回显存再读出来融合之后三个操作在一遍数据流里完成显存访问量直接降了将近三成。内存复用是另一个大头。计算图里多个算子的中间结果生命周期不重叠理论上可以共享同一块显存。图编译器虽然自带内存规划但对MoE模型里动态索引相关的算子往往比较保守宁可多分配也不复用。我们手工分析了一遍计算图把几个大块中间缓冲区的生命周期重叠关系理出来给图编译器提供了显式复用提示峰值显存又往下降了十几个GB。注意算子融合不能盲目追求效果而不顾实现稳定性。融合算子在边界条件下更容易暴露精度问题每次融合改动后都要跑一遍回归测试对比生成结果和基线的一致性。别为了性能把正确性丢了。3.3 KV Cache量化显存和并发量的直接杠杆KV Cache量化是这次调优里性价比最高的一项改动。MoE模型虽然参数多但实际显存瓶颈往往不在权重而在KV Cache——尤其是并发拉高、序列变长之后KV Cache占用的显存会急速膨胀。量化方案我们选了8bit KV CacheK和V都按per-channel方式做对称量化。每token KV Cache的显存占用公式是这样的总显存 层数 × 2(K和V) × 注意力头数 × 头维度 × 字节数 × 序列长度 × batch数换成8bit之后字节数从4降到1KV Cache直接缩水到原来的四分之一。实际测试中生成质量几乎无损困惑度指标变化在零点几个点以内。这个改动让batch size从8提到了32在线业务的并发承载量立刻完全不一样了。3.4 并行策略与并发配置多卡收益不一定线性张量并行在国产芯片上同样适用但卡间通信拓扑需要重新评估。之前GPU平台上TP8跑得很溜换到国产芯片集群后我们发现TP4时通信效率就开始大幅下降。原因涉及卡间互联带宽和拓扑结构差异高TP配置下每步通信等待时间太长计算资源空转严重。最后我们把TP调回2配合数据并行加多实例部署整体吞吐反而更高。另一个重要结构是把prefill和decode分开部署。prefill阶段计算密集decode阶段访存密集混跑在同一批卡上会互相干扰。我们拆成两组卡分别执行后prefill吞吐提高了约45%decode阶段的P95首token延迟也降到了可接受范围。这个结构在GPU生态里也很常用但国产芯片的资源分配比GPU更敏感拆分的收益更明显。3.5 实测数据对照每次优化都要留下证据调优不能凭感觉走每一次变更都要记录前后数据。这是我们从第一天就定下的规矩后来发现它救了我们很多次——因为有些改动表面上变快了但换个并发粒度就劣化有些改动看着没提升但叠加别的优化后产生了协同效应。配置prefill吞吐(tokens/s)decode吞吐(tokens/s)P95首token延迟显存占用初始基线32085980ms96GB图编译融合内存复用560140620ms78GB叠加KV Cache 8bit量化590300480ms52GB叠加并行策略调整620335390ms49GB上面是我们在某测试环境里记录的相对结果数值本身不具普适性但趋势很有参考价值融合和量化带来的提升在不同芯片平台上方向是一致的只是幅度有差异。品一下就能看出从初始基线到最终配置decode吞吐翻了差不多四倍这才是适配工作的真实含金量。4. 工程化落地多卡通信、故障处理与监控模型能跑、性能达标离一个可以长期稳定运行的服务还有很长的路。这一节说的才是真正上线前最磨人的部分——多卡通信、故障处理、监控告警、版本兼容。说白了能跑不算什么能一直跑才是本事。4.1 集合通信库多卡训练的隐形瓶颈多卡并行离不开集合通信库。在GPU生态里通信库是默认标配API稳定、行为可预期。国产芯片的通信库起步晚一些行为差异比较大。我们遇到一个典型案例跨卡通信初始化时默认拓扑探测策略会把部分流量路由到慢速链路上导致多机训练时通信时间占比从预期的10%暴涨到35%。后来通过手动指定卡间拓扑索引和端口映射才把通信耗时压回去。还有一个细节是集合通信的超时设置。大模型训练里all-reduce操作动辄几百MB数据默认超时阈值经常不够用一旦超时任务整条链路崩溃从checkpoint恢复又要浪费大量时间。建议在通信库初始化时显式配置充足的重试和超时参数并且加一层任务级看门狗出现卡死时自动拉现场日志、重启节点而不是等值班人员半夜被报警电话吵醒。4.2 高频故障实录过热降频、驱动崩溃、OOM国产AI芯片的故障表现有自己的脾气。最典型的是长期满载后的降频问题。我们有一块卡在持续跑推理负载半小时后性能悄悄掉了三成卡利用率看起来却还是100%。排查了半天才发现是芯片温度触发了自动降频保护散热方案没有兜住高密度计算的热量。后来在控制面加了温度阈值告警同时调整了负载均衡策略让热度高的卡自动减少分配的任务量。驱动崩溃这类问题也不能忽视。某一次我们开多线程并发预填充显存池时驱动直接挂掉整机需要reboot。这个问题查了几天最后靠抓取崩溃堆栈定位到是SDK中某个内存管理接口的线程安全问题。我们的应对是绕开这个接口改成单线程统一管理显存分配再分发使用问题就再没出现过。OOM的问题前面提过靠显存池和量化已经基本解决这里不再展开。4.3 监控指标与日志分析监控这件事做得多细都不为过。芯片平台下值得盯的核心指标比GPU要多几项卡利用率、HBM显存占用、芯片温度、通信带宽、队列深度、任务耗时分布。我们当时自研了一个轮询脚本每五秒采集一次这些指标写入时序数据库做成看板设置多级告警。有一个看起来很不起眼但实际非常有用的指标是任务耗时分布里的排队时延。如果计算任务本身很短但排队时延很高说明调度策略有问题如果计算时延很高说明算子实现有瓶颈如果通信时延高说明并行配置不太合理。数据结构化了之后很多问题不用猜看一眼分位数曲线就能定位方向。4.4 版本与兼容性管理最后说版本管理。国产芯片的软件栈迭代频率快固件、驱动、算子库、图编译器任何一方升级都可能改变算子行为。我们遇到过固件升级后某个之前跑得好好的算子性能突然崩掉的情况也遇到过图编译器升级后生成的IR结构变化导致显存规划完全错乱的情况。解决方案没什么黑科技就是立规矩生产环境固定版本基线升级先在一个隔离的验证环境做全量回归对比精度、性能、稳定性三项指标全部通过之后才允许上生产。同时每台机器的版本信息全部登记在案出问题时能快速判断是否和版本变更相关。适配工作可以激进生产环境必须保守。5. 常见问题排查速查表适配过程中我们遇到的问题五花八门我把最典型的整理成速查表给后面接手的人当手册用。现象可能原因解决思路模型加载到一半进程退出权重文件格式不兼容、显存不足检查权重格式转换脚本确认SDK显存接口调用无残留句柄首token延迟极长算子走了CPU回退扫描算子映射表定位回退算子替换成原生实现生成结果乱码或概率分布异常精度模式配置不当特定算子溢出逐算子检查精度混合精度下对敏感算子提升到FP32峰值显存远超预期图编译内存规划过于保守、KV Cache未量化手工指定内存复用提示开启KV Cache量化满载半小时后性能下滑芯片温度过高触发降频保护优化散热增加温度告警和负载均衡策略多卡并行吞吐不随卡数线性提升通信拓扑配置不当、并行策略过激进调低TP维度检查通信库拓扑探测结果随机段错误或死锁固件与驱动版本不匹配按版本矩阵重新升级确保驱动固件配套训练任务中途卡死无响应集合通信超时配置过短增加超时阈值通信库显式配置重试机制部署容器无法识别板卡容器运行时与驱动版本不匹配重建容器镜像确认包含对应版本的驱动用户态组件长序列生成时显存暴涨KV Cache按需分配导致碎片化改用预分配显存池用环形buffer管理KV Cache推理吞吐高但延迟毛刺多调度策略导致任务排队不均统计任务排队时延调整多实例间的负载均衡权重升级固件后算子性能下降新版本算子行为变化回退固件版本或改用兼容的算子映射配置这些坑看起来琐碎但每一个都真实地消耗过我们不少时间。表里写的都是解决思路具体落地时不同芯片平台的API差异很大但排查路径基本可以参考。写在后面的一点体会这次适配做完我最大的感受是国产AI芯片这几年真的进步了但它和成熟生态之间的差距不是硬件算力的问题而是软件生态时间积累的问题。算子库覆盖度、编译器的成熟度、调试工具链的完整度、社区解决方案的沉淀这些都需要时间不是砸钱和堆人就能一夜追平的。对我们搞工程的人来说这反而意味着一个窗口期——谁能早一步把模型的适配路径摸透谁就能在后续的项目里建立真正的护城河。最后再分享一个小技巧适配过程中一定要养成每次修改都留基线数据的习惯。我们这次能快速定位问题很大程度上得益于每轮调优都有前后对比数据。这比任何技术亮点都重要——它是整个调优过程的坐标系没有坐标系后面的优化都是在黑暗中瞎摸。如果你正要开始类似的芯片迁移项目我的建议是先花两天时间把模型的算子彻底盘清楚再决定架构方案不要急着跑代码方向对了速度才有意义。