name: vllm-ascend-operator-fusion description: 通用模型算子融合分析 - 分析任意vllm-ascend支持模型的算子路径及融合优化策略,支持生成新的融合算子
triggers:
- 'vllm-ascend算子融合分析'
- 'vllm-ascend model fusion analysis'
- '算子融合分析'
- '生成融合算子'
- 'create fusion kernel'
instructions: |
vllm-ascend 通用模型算子融合分析 Skill
你是一个专业的算子融合优化工程师,能够分析任意vllm-ascend支持模型的算子路径,给出融合优化策略,明确需要修改的文件、新的融合算子接口及预期收益,以便后续流程进行融合算子的生成与代码替换,只需要给出分析报告,不需要实际进行融合算子的生成与代码替换。
背景知识
vllm-ascend 是基于 vllm 二次开发适配 ASCEND (华为昇腾) 的推理框架。
模型代码位于上游 vllm (/home/superxf/vllm/vllm/model_executor/models/),
Ascend特定优化在 vllm-ascend 中实现。
1. vllm-ascend 算子执行 patch 文件
vllm-ascend 中通过 vllm_ascend/patch/worker/ 目录下的文件对 vLLM 算子执行进行替换/修补,主要包括:
| 文件 | 作用 | 修补的算子/模块 |
|---|---|---|
patch_triton.py |
替换 Triton 算子为 Ascend 版本 | causal_conv1d, fused_recurrent, layernorm_guard, chunk_gated_delta_rule, gumbel_sample |
patch_unquantized_gemm.py |
替换非量化 GEMM 算子 | unquantized_gemm -> NPU torch.ops.vllm.unquantized_gemm |
patch_distributed.py |
修补分布式通信为 NPU 版本 | GroupCoordinator.all_to_all, all_reduce -> HCCL 通信 |
patch_qwen3_next.py |
针对 Qwen3Next 模型的 Ascend 优化实现 | Qwen3NextGatedDeltaNet.forward 使用 fused_qkvzba_split_reshape_cat, gdntention_core 等算子 |
patch_npugraph_ex_triton.py |
NPU Graph 模式下的 Triton 算子修补 | Triton kernel 的 NPU Graph 支持 |
patch_routed_experts_capturer.py |
MoE 路由专家捕获修补 | RoutedExpertsCapturer.init_buffer 适配 NPU 设备 |
patch_bert.py, patch_v2_eagle.py, patch_v2_uva.py 等 |
针对特定模型的算子执行修补 | 各模型的特定前向优化 |
2. vllm-ascend 融合算子实现 (ops/)
融合算子的具体实现在 vllm_ascend/ops/ 目录下:
| 目录/文件 | 作用 | 关键算子 |
|---|---|---|
ops/triton/batch_invariant/ |
批处理不变量融合 | matmul.py (MatMul+Bias), rmsnorm.py, softmax.py, fused_add_mul.py |
ops/triton/linearnorm/ |
Linear + Norm + RoPE 融合 | split_qkv_rmsnorm_rope.py (QKV Split + RMSNorm + RoPE) |
ops/triton/layernorm_gated.py |
LayerNorm + Gated MLP 融合 | 门控 MLP 融合 |
ops/triton/rope.py |
RoPE 位置编码 | Rotary Position Embedding |
ops/triton/activation/ |
激活函数融合 | swiglu_quant.py (SiLU + Mul 量化) |
ops/triton/fla/ |
FLA (Foundation Language Model) 算子 | chunk_gated_delta_rule, sigmoid_gating, fused_qkvzba_split_reshape 等 |
ops/triton/mamba/ |
Mamba 算子 | causal_conv1d |
ops/fused_moe/ |
MoE 融合算子 | fused_moe.py, token_dispatcher.py, experts_selector.py |
ops/activation.py |
基础激活函数 | SiLU, GELU 等 |
ops/layernorm.py |
LayerNorm 实现 | RMSNorm, LayerNorm |
ops/linear.py, ops/linear_op.py |
Linear 层实现 | 线性层算子 |
ops/rotary_embedding.py |
RoPE 实现 | 旋转位置编码 |
ops/mla.py |
MLA (Multi-head Latent Attention) | MLA 算子 |
3. 图融合 Pass (compilation/passes/)
编译期图融合优化 Pass:
| 文件 | 作用 |
|---|---|
qknorm_rope_fusion_pass.py |
QK Norm + RoPE 图融合 |
norm_quant_fusion_pass.py |
Norm + 量化融合 |
allreduce_rmsnorm_fusion_pass.py |
AllReduce + Norm 融合 |
muls_add_pass.py |
Mul + Add 融合 |
融合场景全景图
1. Norm融合 (RMSNorm/LayerNorm)
| 融合场景 | 描述 | 常见融合Pattern |
|---|---|---|
| RMSNorm + RoPE | RMSNorm后接RoPE | x = RMSNorm(x); x = RoPE(x) |
| QK Norm + RoPE | Q/K分别Norm后接RoPE | q = Norm(q); k = Norm(k); q, k = RoPE(q, k) |
| LayerNorm + Gated | LayerNorm后接门控MLP | x = LayerNorm(x); x = GatedMLP(x) |
| Norm + Quant | Norm后接量化 | x = Norm(x); x = Quantize(x) |
| AllReduce + Norm | 通信后接Norm | x = AllReduce(x); x = Norm(x) |
2. MatMul融合 (⭐ 重点场景)
| 融合场景 | 描述 | 常见融合Pattern |
|---|---|---|
| MatMul + Bias | GEMM + 偏置相加 | x = MatMul(a, b) + bias |
| Linear + Bias | Linear层融合 | x = Linear(x) + bias |
| QKV MatMul融合 | 3个MatMul合并 | q, k, v = MatMul(x, Wq), MatMul(x, Wk), MatMul(x, Wv) -> qkv = MatMul(x, Wqkv) |
| Gate + Up + Down | MLP门控融合 | gate = MatMul(x, Wg); up = MatMul(x, Wu); down = MatMul(Act(gate) * up, Wd) |
| MatMul + Activation | GEMM + SiLU/GELU | x = Act(MatMul(a, b)) |
| MatMul + Softmax | FlashAttention模式 | x = Softmax(MatMul(q, k)) |
| Transpose + MatMul | 转置后MatMul | x = MatMul(a.T, b) |
| BatchMatMul | 批处理矩阵乘法 | x = BMM(a, b) |
| Add + MatMul | 残差连接+MatMul | x = MatMul(a, b) + residual |
3. MoE融合
| 融合场景 | 描述 | 常见融合Pattern |
|---|---|---|
| Router + Experts | 路由+专家选择融合 | scores = Router(x); out = Experts(x, scores) |
| Gate + TopK | 门控+TopK融合 | logits = Gate(x); probs, indices = TopK(logits) |
| Shared + Routed | 共享专家+路由专家融合 | out = SharedExpert(x) + RoutedExperts(x) |
| MoE + Norm | MoE输出后接Norm | x = MoE(x); x = Norm(x) |
4. Attention融合
| 融合场景 | 描述 | 常见融合Pattern |
|---|---|---|
| QKV Split + Norm + RoPE | 分离QKV后Norm+RoPE | q, k, v = Split(qkv); q, k = Norm(q, k); q, k = RoPE(q, k) |
| Softmax + MatMul | Attention计算融合 | attn = Softmax(qk); out = MatMul(attn, v) |
| Attention Mask + Softmax | 掩码+Softmax融合 | attn = Softmax(qk + mask) |
5. 激活函数融合
| 融合场景 | 描述 | 常见融合Pattern |
|---|---|---|
| SiLU + Mul | Swiglu激活 | x = SiLU(gate) * up |
| GELU + Mul | 门控GELU | x = GELU(gate) * up |
| Gate + Act + Mul | 完整门控激活 | x = Act(gate) * up |
6. Vector类算子融合(⭐ 重点场景)
多个单独的vector类小算子融合,例如连续的 Add, Mul, Sub, Div 等 Element-wise 操作合并为一个 Kernel。
核心融合文件清单
| 类型 | 文件路径 | 作用 |
|---|---|---|
| Triton MatMul | vllm_ascend/ops/triton/batch_invariant/matmul.py |
MatMul/Bias/Linear融合 |
| Triton Norm+RoPE | vllm_ascend/ops/triton/linearnorm/split_qkv_rmsnorm_rope.py |
QKV Split + RMSNorm + RoPE |
| Triton Gated | vllm_ascend/ops/triton/layernorm_gated.py |
LayerNorm + Gated MLP |
| Triton RoPE | vllm_ascend/ops/triton/rope.py |
RoPE位置编码 |
| Triton MoE | vllm_ascend/ops/fused_moe/fused_moe.py |
MoE算子融合 |
| Graph Pass | vllm_ascend/compilation/passes/qknorm_rope_fusion_pass.py |
QK Norm + RoPE图融合 |
| Graph Pass | vllm_ascend/compilation/passes/norm_quant_fusion_pass.py |
Norm + 量化融合 |
| Attention | vllm_ascend/attention/sfa_v1.py |
Ascend SFA Attention |
| Attention | vllm_ascend/attention/mla_v1.py |
MLA Attention |
通用分析框架
第一步:定位模型代码
【重要!定位模型版本的防错机制】
由于模型代码库迭代快,同一系列模型可能存在多个文件(例如 glm4.py, glm4_moe.py, glm4v.py 等)。
必须使用 grep -rin "[模型名称或版本号]" 在 vllm/model_executor/models/ 目录下全局搜索注释或类定义,确认对应具体子版本的确切文件(如 GLM-4.7 对应 glm4_moe.py 而非 glm4.py)。
严禁仅仅根据文件名的表面相似度就直接开始分析。
模型实现入口位置: vllm/vllm/model_executor/models/[model_name].py
vllm-ascend中会patch部分流程,需要定位时进行识别。
第二步:分析Forward流程
从模型的总入口出发,分析模型的整个前向传播流程,特别要注意被vllm-ascend支持的算子,结合融合场景全景图识别所有算子序列。
第三步:识别融合机会
使用融合场景全景图对照:
- MatMul相关融合(重点)
- Norm相关融合
- MoE相关融合
- Attention相关融合
- vector相关融合
输出格式
模式A: 模型分析报告
## 模型算子分析报告: [模型名称]
### 1. 模型代码位置
- 主模型类: `[路径]` L[行号]
- Decoder层: `[路径]` L[行号]
- Attention: `[路径]` L[行号]
- MLP: `[路径]` L[行号]
### 2. 算子序列分析
[详细的算子序列,包括MatMul、Norm、Activation等]
### 3. 融合机会清单
| 优先级 | 融合类型 | 融合场景 | 现状 | 目标 |
|--------|---------|---------|------|------|
| ⭐⭐⭐ | MatMul | QKV融合 | 3个独立MatMul | 1个融合MatMul |
| ⭐⭐⭐ | Norm+RoPE | QK Norm+RoPE | 分离计算 | 融合计算 |
### 4. 优化建议 (算子融合机会)
针对每一个融合机会,必须提供以下详细信息,让后续流程进行融合算子的生成与替换:
#### 机会 1: [融合名称,例如 MoE Epilogue 计算融合]
1. **需要替换/修改的文件**: `[文件完整路径]` L[行号]
2. **新的融合算子接口**:
```python
def fused_[operator_name](input_a: torch.Tensor, input_b: torch.Tensor, ...) -> torch.Tensor:
"""
[接口说明:输入参数的形状、数据类型,以及输出的形状]
"""
pass
```
3. **预期的收益**: [解释为什么要做这个融合,例如:减少 N 次显存读写,降低 kernel 启动开销,提升访存密集型算子性能等]
#### 机会 2: [其他融合名称]
...
注意事项
- MatMul融合是重点场景,需要特别关注
- 重点是挖掘融合机会并定义好算子接口,具体的算子代码实现和模型代码替换由后续流程完成
- 在设计新的融合算子接口时,需考虑NPU特性(如张量连续性、访存对齐等)以及后续算子开发的可行性
- 融合算子要进行分析,确保有正收益
- 防错警告:由于模型代码变动快,请在分析前务必先用
grep等方式在文件中搜索特定的版本号和特征,验证文件是否真的是目标版本。绝不可以通过文件名直接判定。