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支持的算子,结合融合场景全景图识别所有算子序列。

第三步:识别融合机会

使用融合场景全景图对照:

  1. MatMul相关融合(重点)
  2. Norm相关融合
  3. MoE相关融合
  4. Attention相关融合
  5. 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 等方式在文件中搜索特定的版本号和特征,验证文件是否真的是目标版本。绝不可以通过文件名直接判定。