vLLM-Kunlun 框架性能极致优化，充分释放昆仑芯硬件性能

百度智能云技术站
0 阅读8分钟

本文整理自 26 年 3 月 15 日 vLLM-Kunlun Meetup 北京站活动的同名主题演讲。在公众号回复「vLLM-Kunlun」，可以获得此次 Meetup 上半场的 6 个演讲主题材料。

本次分享我会围绕国产芯片适配开源模型的痛点，以及我们在算子、Kernel Launch、框架层、系统吞吐等多维度的优化方案和落地效果展开，带大家了解如何让 vLLM 框架在昆仑芯上发挥出极致的硬件性能。

在做昆仑芯与开源模型的适配过程中，我们遇到了几个核心痛点，也是行业内的共性问题：

  • 第一是 Kernel Launch 开销大，频繁的 Kernel 调用会带来大量的性能损耗；
  • 第二是框架层开销高，原生框架的部分实现无法适配昆仑芯硬件特性，额外消耗算力；
  • 第三是高性能算子适配难，开源模型的各类核心算子不能充分利用昆仑芯的硬件加速能力；
  • 最后是 Decode 吞吐低，推理阶段的解码性能上不去，直接影响业务体验。

这四大痛点也是我们本次优化工作的核心攻坚方向。

图片

针对高性能算子适配的难题，我们的核心解法是推出 vLLM-Kunlun Plugin ，把框架中硬件相关的部分完全独立出来，解耦框架与硬件的依赖，极大简化了昆仑芯的算子适配工作。

重点针对模型中的核心算子模块做定制化适配，包括 FFN、MoE、Attention、PROJ，接下来我们也会逐一介绍每个模块的具体优化策略。

图片

在具体的算子适配中，我们基于昆仑芯的硬件单元特性，为不同类型算子制定了专属的优化策略，做到「算子与硬件精准匹配」：

  • Element-wise 类算子（比如残差加法、激活函数），利用 Cluster 向量指令实现带宽加速；
  • 超越函数（exp、log、sin、cos 等），调用昆仑芯 SFU 硬件单元，实现低延迟、高精度的硬件加速；
  • Matmul 和 Attention 这类强计算、高带宽需求的算子，依托 SDNN 硬件单元做核心加速；
  • FFN 和 PROJ 模块，采用 INT8 量化技术，在保证精度的前提下提升计算和存储效率；
  • Fused_moe 算子做算子融合优化，从内存占用、任务调度、CPU 开销三个维度全面优化；
  • Attention 模块则将 Prefill 和 Decode 阶段分开调度，解决两个阶段负载不均衡的问题，让算力利用更充分。

图片

MoE 算子作为大模型的核心算子之一，我们根据 M * top_k 的数值大小，设计了三级动态优化路径，实现不同规模场景下的性能最优：

  • 当 M * top_k < 400 时，属于极小 Batch 或短序列场景，启用极小规模优化（SMALL PATH） ，极简预处理路径，追求极致的推理延迟；

  • 当 400 ≤ M * top_k ≤ 768 时，为中等规模场景，开启中等规模排序（MEDIUM PATH） ，打开 sort_mode = True，平衡预处理的开销和计算效率；

  • 当 M * top_k > 768 时，是大规模计算场景，采用 大规模高性能（LARGE PATH） ，通过 gen_block_statistic 分块统计 + moe_pre_sorted 预排序，最大化昆仑芯的并行计算吞吐。同时我们还制定了算子融合策略：当 M ≥ 1024 时，moe_fc 会自动融合 SWISH_GLU 激活函数，进一步减少显存的读写次数，降低带宽压力。

图片

针对 Qwen 3.5 / Qwen3-NEXT 模型，我们做了核心的 Kernel 融合优化，推出 split_norm_rope_neox 融合算子，做到「无新增算子」即可跑通 Qwen 3.5，同时实现显著的性能提升。

在 split_norm_rope_neox 融合算子的落地过程中，我们遵循「定位瓶颈 - 针对性开发 - 融合落地 - 效果验证」的步骤：

  • 首先用 Torch Profiler 抓取 Trace，精准定位到 Qwen 3.5 模型推理的性能瓶颈点；
  • 针对 MoE 算子在小 BS（Batch Size）场景下的性能问题，开发专属的小 BS 优化方案；

最终实现 Prefill 吞吐  8%  的提升，同时将 Kernel Launch 次数从 4 次缩减为 1 次，这也是本次优化的核心成果。

这里再给大家展示 split_norm_rope_neox 融合算子的核心 API 调用代码。我们基于 torch.ops.xspeedgate_ops 开发了该算子，集成了 qkv 拆分、归一化、旋转位置编码、门控等所有核心逻辑，入参包含 qkv 张量、归一化权重、位置信息、注意力头参数等，一次性完成多步计算。

对比原生实现，优化前的调用链是 4 步独立 Kernel：q_norm → k_norm → sin_cos → rotary_emb，不仅带来多次 Kernel Launch 的开销，还会产生中间 Tensor 的读写损耗；优化后通过 1 个融合 Kernel，完成 Split + QK Norm + RoPE + Gate 四步合一，既减少了 Kernel Launch 次数，又大幅降低了内存带宽压力，这也是 Prefill 吞吐能提升 8% 的核心原因。

图片

在算子适配的基础上，我们做了 INT8 量化的性能验证，以及昆仑芯 P800 与 NVIDIA A 卡的 INT8 推理延迟对比：

  • 首先看 mimo-v2-flash 模型的 INT8 vs BF16 对比，在 Input = 8192、Output = 1024 的场景下，INT8 量化的 Mean TTFT 远优于 BF16，随着 Batch Size 提升，加速比最高达到 5.45 倍，量化效果显著；

  • 再看 P800 与 NVIDIA A 卡 的 INT8 推理延迟对比，P800 在 Attention 模块的延迟仅 0.102 ms，相比 A 卡 的 0.147 ms，性能提升 34%；qkv_proj 和 o_proj 模块虽略高于 A 卡，但整体在核心计算模块实现了反超，充分验证了昆仑芯在算子优化后的硬件性能。

图片

针对 Kernel Launch 开销大的问题，我们的核心优化方案是基于 CUDA Graph 技术，消除 CPU 侧的调度开销，实现极致的性能提升。具体做了三层优化：

  • 实现 Piecewise CUDAGraph（分段 CUDA 图）和 Full_and_Piecewise CUDAGraph（全量 + 分段 CUDA 图）的适配；
  • 针对性优化 CUDA Graph 的同步开销，减少同步带来的性能损耗；
  • 支持计算与通信同时捕获，让算力和通信资源并行利用，进一步提升效率。

图片

这里我们分享一下 CUDA Graph 的核心性能对比数据，测试场景为 Input = 1024、Output = 1024，以 Piecewise CUDAGraph BS = 1 为基准（1.00 倍）。

从结果能看到，Full & Piecewise CG（全量 + 分段 CUDA 图）在 Output Throughput（输出吞吐） 、Mean TTFT（首令牌生成时间） 、Mean TPOT（每令牌生成时间）  三个核心指标上，均实现了远超 Piecewise CG 的性能提升，且随着 Batch Size 从 1 提升到 32，加速比持续攀升。

其中，full_and_piecewise_cudagraph 比普通的 piecewise_cudagraph  在 Output 吞吐上有 2 倍多的性能提升，Mean TTFT 和 Mean TPOT 也有显著提升。这充分证明，CUDA Graph 技术能有效消除 Kernel Launch 开销，让昆仑芯的并行计算能力得到充分释放。

图片

框架层的优化，我们核心采用「原生实现替换为定制化算子」的思路，针对框架中性能瓶颈的原生实现，开发昆仑芯定制化算子：

比如将 indexer_k_quant_and_cache、flashinfer_rotary_embedding 的原生实现，替换为 torch.ops.xspeedgate_ops 下的定制化算子，充分利用昆仑芯的硬件特性，替代低效的原生框架实现，从底层降低框架层开销。

图片

针对模型中的掩码处理逻辑，我们开发了 get_masked_input_and_mask 定制化算子，替换原生的 Python 实现。

原生实现中存在大量的条件判断和张量类型转换，无法适配昆仑芯硬件。而我们的定制化算子，结合 XPU SIMD 架构和模型自身参数特点，构建了特殊的优化分支，减少了跳转等指令的开销，让掩码处理的计算效率大幅提升，同时通过 torch.ops.xspeedgate_ops 直接调用，实现硬件与算子的深度适配。

图片

针对大模型推理中的随机采样环节，我们发现原生 torch.multinomial 会导致 CPU-GPU 同步，且原生 uniform_ 采用的 Philox4x32-10 随机数算法，因涉及状态跳转，无法在 XPU 上利用 SIMD 算力，实现效率极低。

对此我们做了两大优化：第一，开发 inplace_exponential 定制化算子，替代原生的指数分布生成逻辑，避免 CPU-GPU 同步；第二，将随机数算法从 Philox4x32-10 改为 Philox2x32-10，适配 XPU 的 SIMD 架构，充分利用硬件的并行算力，大幅提升随机采样的效率。

图片

在解决了算子、Kernel Launch、框架层的问题后，我们进一步从系统层面优化推理吞吐，推出三大核心方案：

  • 大规模 EP 并行，充分利用昆仑芯的多 EP 硬件资源，实现大规模并行计算；
  • Dual-Batch Overlap 与 DeepEP，支持 Micro Batch 并行，让计算和通信操作重叠执行，隐藏通信延迟；
  • 支持 MTP（投机解码） ，针对性提升 Decode 阶段的解码性能，这也是提升整体吞吐的关键。

图片

我们对比一下昆仑芯 P800 与 H20 的 Output Throughput（输出吞吐）数据，测试场景为 Input = 1024、Output = 1024，以 H20 BS = 1 为基准（1.00 倍）。

从结果可以看到，P800 在大部分 Batch Size 下，均实现了接近 H20 80%的性能，某些场景下可以做 80%+。

这充分验证了我们多维度优化方案的有效性，也证明了 vLLM-Kunlun 框架能充分释放昆仑芯 P800 的硬件性能，实现远超同类硬件的推理吞吐。

图片

项目地址：github.com/baidu/vLLM-…

avatar
文章
阅读
粉丝