# 大模型推理加速200%+？这套国产融合算子库把CUDA/Triton全替换了 > 干了几年模型部署，最头疼的事永远是——推理速度上不去，显存不够用，pipeline里一堆小算子来回倒腾数据，bandwidth全浪费在中间结果的读写上了。 如果你也在折腾DeepSeek、Qwen、LLaMA这些大模型的推理加速，那你大概率也踩过同一个坑：框架自带的算子——不管是CUDA原生的还是Triton写的——总有那么几个"慢得离谱"的环节，瓶颈不在计算，而在**访存**和**算子调度的碎片化**。 今天聊的这套东西，就是专门冲着这个痛点来的。 --- ## 目录 1. [这玩意儿到底是什么](#1-这玩意儿到底是什么) 2. [Reduce类算子：Softmax和LayerNorm能有多快](#2-reduce类算子softmax和layernorm能有多快) 3. [大模型融合算子：把七八个小算子揉成一个](#3-大模型融合算子把七八个小算子揉成一个) 4. [GEMM/MOE类算子：量化+融合，显存省一半](#4-gemm-moe类算子量化融合显存省一半) 5. [总结：什么时候该用它](#5-总结什么时候该用它) --- ## 1. 这玩意儿到底是什么 LightOP，一个自研的**融合算子库**。核心思路简单粗暴：把模型推理和训练里那些零散的小算子，该融合的融合，该重写的重写，减少显存读写次数，榨干硬件带宽。 它的覆盖范围，说白了就是针对大模型推理做了**全覆盖**。Attention和GEMM之外的算子，全部来自LightOP，把vLLM框架里原生的CUDA算子和Triton算子全给替换了。 整体可以分成四大块： | 分类 | 包含的算子 | 主要场景 | | ------------------ | ------------------------------------------------------------ | ---------------------------- | | **Reduce类** | Softmax, LayerNorm, RMSNorm, GroupNorm + Elementwise | 基础归一化 & 激活 | | **定制算子** | MSDA, Box Cost, Encode Reg Target, Polyfill | 特定模型定制需求 | | **大模型融合算子** | Moe Fuse Gate, Softmax TopK, Fuse RMS RoPE, MoE Align, MoE Sum, RmsRopeQuantKvcache | DeepSeek/Qwen/LLaMA 推理加速 | | **GEMM/MOE类** | W4A8/W8A8 INT8/FP8 GEMM, W4A16/W8A8量化MOE, MQALogits | 量化推理 & 混合专家 | --- ## 2. Reduce类算子：Softmax和LayerNorm能有多快 ### 2.1 Softmax优化 Softmax看着简单，但在大模型里一跑就是几百万次。瓶颈在哪？不是计算 `exp`，是**访存**——数据从显存搬进搬出，GPU大部分时间在等数据。 LightOP的Softmax做了四件事： ```mermaid flowchart TD A["Input Tensor"] --> B["向量化访存
128-bit宽指令读写"] B --> C["层级规约
Warp Shuffle + 共享内存"] C --> D{"列数是否超阈值?"} D -- "小列数" --> E["Block内Reduce
Max & Sum"] D -- "大列数/超长Sequence" --> F["跨Block全局同步
AtomicAdd"] E --> G["Exp & Normalize & Write"] F --> G G --> H["Output Tensor"] ``` - **向量化访存**：用128-bit宽指令一次性读写，带宽利用率拉满 - **层级规约**：Warp Shuffle + 共享内存配合，尽量减少全局显存读写 - **内核分发**：根据输入列数自动选最佳kernel模板——列少走Block Reduce，列多走Global Atomic - **跨Block全局同步**：超长Sequence场景下，用AtomicAdd做跨Block通信，避免重复访存和重复exp计算 实测显存带宽能干到大几百GB/s级别，对于Reduce类算子来说相当猛了。 ### 2.2 LayerNorm优化 LayerNorm更狠，直接把Input + Residual -> 统计量计算 -> 归一化三个步骤全部**融合成一个kernel**，中间结果全部走寄存器，不写回显存。 ```mermaid flowchart LR subgraph "融合前（传统做法）" X1["Input X"] --> K1["Kernel 1: Add Residual"] R1["Residual R"] --> K1 K1 --> M1["写回显存"] M1 --> K2["Kernel 2: 算Mean/Var"] K2 --> M2["写回显存"] M2 --> K3["Kernel 3: Normalize"] K3 --> Y1["Output"] end subgraph "融合后（LightOP）" X2["Input X"] --> F["Fused LayerNorm Kernel"] R2["Residual R"] --> F F --> Y2["Output"] end ``` 四个关键技术点： 1. **Welford One-Pass**：单次遍历即可算出均值+方差，数值比传统两步法更稳定 2. **深度融合**：Input + Residual → 统计量计算 → 归一化，全程寄存器传递 3. **向量化动态分发**：根据列数自适应选向量化加载方式 + BlockSize配置 4. **Block Reduce**：利用BlockReduce + 共享内存高效规约统计量 实测显存带宽跑到接近1TB/s量级，这个级别的融合效果基本是寄存器级别的数据复用了。 --- ## 3. 大模型融合算子：把七八个小算子揉成一个 这块是LightOP最核心的价值——DeepSeek系列、Qwen、LLaMA这些大模型里，pipeline太长了，小算子一个接一个，每一步都要写回显存再读出来，白白浪费带宽。 ### 3.1 Moe Fuse Gate & Softmax TopK MOE（混合专家）模型里，Gate计算 + Softmax + TopK 是每一层都要跑的固定流程。三个独立的kernel → 三个kernel融合成一个。 ```mermaid flowchart LR subgraph "Moe Fuse Gate" A1["Gate Logits"] --> B1["Fused Kernel"] B1 --> C1["Sigmoid + Bias Add"] C1 --> D1["Cache Top1"] D1 --> E1["Group Select"] E1 --> F1["Shfl Broadcast + Write"] F1 --> G1["Out: Indices/Weights"] end subgraph "Softmax TopK" A2["Scores"] --> B2["Fused Kernel"] B2 --> C2["Cache TopK"] C2 --> D2["TopK Select (Unroll展开)"] D2 --> E2["Partial Exp"] E2 --> F2["Renormalize"] F2 --> G2["Out: Indices/Weights"] end ``` 优化策略拆开看： **Moe Fuse Gate：** - **计算优化**：优先选组，缓存Top1/Top2，减少线程工作负载 - **访存优化**：用Shfl广播结果，指定Lane合并写入，保证全局内存合并访问 - **分支优化**：减少条件跳转，降低线程束分化（warp divergence） **Softmax TopK：** - **计算优化**：不跑全量Softmax，只对TopK做Exp——大部分token的exp根本不用算 - **Argmax与数据擦除**：Unroll展开减少数据擦除开销 - **访存优化**：缓存Top1，优化结果写回逻辑 实测效果：相比vLLM原生实现，fuse_gate平均提升**318%**，softmax_topk平均提升**685%**。 > **踩坑提示**：融合后的kernel对warp调度非常敏感，如果你的batch size或token数不匹配kernel的分发策略，性能可能不升反降。建议先用LightOP自带的benchmark跑一下你的真实shape，别直接上线。 ### 3.2 Fuse RMS RoPE & MoE Align & MoE Sum 这三个是Qwen/LLaMA和DeepSeek MOE场景下高频调用的算子链，融合后的效果同样暴力： | 算子 | 核心优化 | 相对vLLM提升 | | ----------------- | ------------------------------------------------------------ | ------------ | | **Fuse RMS RoPE** | RMS与RoPE数据对齐，免去LDS中转；缓存Cos/Sin供Q/K多头复用 | **542%** | | **MoE Align** | 并行前缀和加速索引计算；分层原子加减少全局原子操作；融合专家map/mask/fill | **227%** | | **MoE Sum** | 128bit向量化读写；多Block并行处理单Token；融合共享专家factor + bias add | **164%** | ```mermaid flowchart TD subgraph "Fuse RMS RoPE" A1["Input Q/K"] --> B1["RMS Norm"] B1 --> C1["Residual Add"] C1 --> D1["RoPE Q/K
缓存Cos/Sin复用"] D1 --> E1["Output"] end subgraph "MoE Align" A2["Input Tokens"] --> B2["Token Reorder
并行前缀和"] B2 --> C2["Expert Map/Mask
分层原子加"] C2 --> D2["Fill Buffer"] D2 --> E2["Sorted Output"] end subgraph "MoE Sum" A3["Input Tokens"] --> B3["Accumulate Outputs
多Block并行"] B3 --> C3["Shared Expert Factor"] C3 --> D3["Add Bias"] D3 --> E3["Output"] end ``` ### 3.3 RmsRopeQuantKvcache：DeepSeek专用融合大杀器 这玩意儿是专门给DeepSeek系列模型做的——把原来**9个访存密集型小算子**揉成一个。 ```mermaid flowchart LR subgraph "融合前：9个小算子各自独立" direction TB K1["RMS Norm Q"] --> M1["写回"] K2["RMS Norm K"] --> M2["写回"] K3["RoPE Q"] --> M3["写回"] K4["RoPE K"] --> M4["写回"] K5["Quant Q"] --> M5["写回"] K6["Quant K"] --> M6["写回"] K7["Cache Q"] --> M7["写回"] K8["Cache K"] --> M8["写回"] K9["Cache V"] --> M9["写回"] M1 --> R1["..."] M2 --> R2["..."] M3 --> R3["..."] M4 --> R4["..."] M5 --> R5["..."] M6 --> R6["..."] M7 --> R7["..."] M8 --> R8["..."] M9 --> R9["..."] end subgraph "融合后：一个Kernel搞定" direction TB F1["RMS (Q)"] --> F2["RMS (K)"] F3["RoPE (Q)"] --> F4["RoPE (K)"] F5["Quant (Q)"] --> F6["Quant (K)"] F7["Cache Q/K/V"] end ``` 核心优化策略： - **Block特化**：RMS、RoPE、Quant三个操作放在同一个kernel的不同线程块里独立计算，省掉三次kernel launch的调度开销 - **Warp分配优化**：原来单个warp只处理一个token，但RoPE的RotDim很小，向量化访存导致7/8的线程在摸鱼。改成**单warp处理多个token**，线程利用率直接拉满 - **全局访存砍半**：融合前每个小算子都要读写全局显存来回倒数据，融合后中间结果全程寄存器/共享内存传递 - **硬件指令加速**：FP32→FP8/BF8的类型转换直接用硬件指令 `__builtin_hcu_cvt_pk_bf8_f32`，两个拼成x4指令，减少指令数同时避免寄存器浪费 实测性能数据，在不同BatchSize下对比融合前后的延迟： ```mermaid xychart-beta title "RmsRopeQuantKvcache 融合前后延迟对比 (us)" x-axis ["BS=1", "BS=2", "BS=4", "BS=8", "BS=16", "BS=24", "BS=32", "BS=48", "BS=64"] y-axis "延迟 (us)" 0 --> 60 bar [48, 45, 42, 40, 33, 30, 28, 26, 26] bar [7, 7, 7, 7, 7, 7, 7, 7, 6] ``` （上图为mermaid近似还原，原始数据来自原文档性能测试页） 平均提升**411%**。注意看——融合后的延迟在不同BatchSize下基本稳定在6-7us，说明瓶颈已经从"算子太多调度不过来"变成了"数据量就这么大，得老老实实搬"。 --- ## 4. GEMM/MOE类算子：量化+融合，显存省一半 ### 4.1 W8A8量化GEMM：Prefill和Decoding分开搞 量化GEMM听起来不新鲜，但LightOP做了一件很聪明的事——**Prefill和Decoding是两种完全不同的场景，不能用一个实现糊弄**。 > Prefill阶段：一次处理几千个token，计算密集。 > Decoding阶段：每次只处理1个token，访存密集。 所以LightOP的W8A8 GEMM**拆成两个实现**： | 阶段 | 特征 | 优化策略 | 关键技术 | | ------------ | ---------- | ---------------- | ------------------------------------------------------------ | | **Decoding** | 访存密集型 | 类CUDA实现 | Input/Weight均不过LDS、16x16x32 TensorCore、WarpSplitK、向量化访存、Smem写回缓冲 | | **Prefill** | 计算密集型 | 手写汇编极致优化 | 多级流水异步拷贝掩盖延迟、WarpSpecial策略（4 Warp读数据 + 8 Warp计算） | SmoothQuant量化流程代码（脱敏版，展示核心逻辑）： ```python def per_token_quant_int8(x): """Per-token量化到int8""" # x: [M, K] absmax = x.abs().amax(dim=-1, keepdim=True) # [M, 1] scale = absmax / 127.0 scale = torch.where(absmax == 0, torch.ones_like(scale), scale) x_quant = torch.round(x / scale).clamp(-128, 127).to(torch.int8) return x_quant, scale def w8a8_gemm_smooth(a, b, scale_a, scale_b, bias=None): """W8A8量化GEMM""" a_quant, scale_a_quant = per_token_quant_int8(a) # 反量化到float做矩阵乘 a_fp = a_quant.to(torch.float32) b_fp = b.to(torch.float32) out = F.linear(a_fp, b_fp.t()) # [M, K] @ [K, N] -> [M, N] # 补偿量化scale out = out * scale_a_quant * scale_b.t() if bias is not None: out = out + bias return out ``` 核心思路是SmoothQuant：通过逐通道的等价缩放变换，把激活值里的离群值平滑迁移到权重上，这样权重和激活值都能用INT8表示，**显存占用直接减半，计算速度翻倍**，而且**不需要重新训练**。 ### 4.2 MQALogits：DeepSeek推理的加速核心 MQALogits是Multi-Query Attention（MQA）模式下快速计算注意力分数的算子。核心思想：用一组共享的K/V状态服务于所有Query头，解码时只需缓存一份KV对。 和GEMM一样，也按照Prefill/Decoding分阶段实现： | 阶段 | 策略 | 关键优化 | | ------------ | -------------------- | ------------------------------------------------------------ | | **Decoding** | 类CUDA（访存密集） | buffer_load预取、BlockSplitK、WarpSplitK、向量化访存、不过LDS | | **Prefill** | 手写汇编（计算密集） | Head间+Head内多级流水、K数据寄存器复用避免重复读取、Clean融合入epilogue | Decoding阶段在token数=16时性能可达同类GPU的290%；Prefill阶段在token数=2048时约为同类GPU的75%（计算密集型场景下这个比例已经相当不错了）。 ```mermaid flowchart LR subgraph "MQA Logits 计算流程" Q["Query Q"] --> E["einsum: Q @ K^T"] K["Shared K"] --> E E --> R["ReLU激活"] R --> W["加权求和
weights @ scores"] W --> L["Logits + Mask"] end ``` ### 4.3 MOE量化方案矩阵 MOE（混合专家）的量化比普通GEMM复杂得多——每个专家的矩阵形状可能不一样，还有两组GEMM叠加。 LightOP支持的量化方案： | 方案 | 权重 (W) | 激活值 (A) | 适用场景 | | ------ | -------- | ---------- | -------------- | | W16A16 | FP16 | FP16 | 基准精度 | | W8A8 | INT8/FP8 | INT8/FP8 | 推理加速 | | W4A8 | INT4 | INT8/FP8 | 极致压缩 | | W4A16 | INT4 | FP16 | 兼顾精度和压缩 | 针对每种量化方案，还按输入规模进一步细分： - **小Size（Decoding场景）**：访存密集型，重点改善访存效率 - **中Size**：最大化显存带宽，访存和计算互相掩盖 - **大Size（Prefill场景）**：计算密集型，最大化计算资源占用 几个值得一提的优化技巧： - **权重离线重排与压缩**：W4A8方案采用**8in1压缩重排+npack4**，让访存合并和向量化成为可能 - **矩阵分块+Tuning**：MOE里有两组GEMM，支持分开tuning搜索最佳分块配置 - **多级流水**：最大化访存和其他并行硬件单元的调用掩盖 - **LDS优化**：Padding + Swizzle策略减少bank冲突；特定大Shape时用多warp warp特化方案手动排布流水 --- ## 5. 总结：什么时候该用它 这一套算子库的核心价值，我觉得可以用三句话概括： 1. **能融合的都融合了**：Reduce类算子3合1、大模型算子9合1——减少的是kernel launch开销和中间结果的显存读写，这些都是"看不见的耗时" 2. **该分开的都分开了**：Prefill和Decoding走不同的实现路径，访存密集型和计算密集型各用各的优化策略——没有"一刀切" 3. **量化不只是精度换速度**：SmoothQuant、W4A8重排压缩、多级流水掩盖——这是一套完整的量化部署方案，不只是把FP16换成INT8 适合什么场景？ - 正在部署DeepSeek/Qwen/LLaMA系列大模型 - 推理瓶颈在算子层面（不是模型本身太大） - 在国产加速卡上做适配（LightOP的架构设计天然支持异构） - 想替换掉vLLM/Triton里的原生算子，获得更好性能 不适合什么场景？ - 你的模型结构特别小众，算子不在覆盖范围内 - 瓶颈在Attention或GEMM本身（这些不在LightOP范畴） - 模型太小，融合收益有限 ---