自定义算子 (ops/)

一、模块概述

ops目录包含针对昇腾NPU优化的自定义算子，涵盖激活函数、归一化、线性层、位置编码、注意力算子、MoE融合算子等核心操作。共计90+个Python文件，包括大量Triton内核实现。

graph LR
    A[ops/] --> B[基础算子]
    A --> C[模型专用算子]
    A --> D[Triton内核]
    A --> E[MoE融合算子]
    
    B --> B1[activation.py]
    B --> B2[layernorm.py]
    B --> B3[linear.py]
    B --> B4[rotary_embedding.py]
    
    C --> C1[mla.py]
    C --> C2[dsa.py]
    C --> C3[gdn.py]
    C --> C4[kimi_kda.py]
    
    D --> D1[triton/rms_norm.py]
    D --> D2[triton/rope.py]
    D --> D3[triton/kda/]
    D --> D4[triton/fla/]
    
    E --> E1[fused_moe/fused_moe.py]
    E --> E2[fused_moe/moe_mlp.py]
    E --> E3[fused_moe/routed_experts.py]
                

二、文件结构

ops/
├── __init__.py
├── activation.py             # AscendQuickGELU, AscendSiluAndMul, AscendSwigluOAIAndMul
├── layernorm.py              # AscendRMSNorm, AscendGemmaRMSNorm, AscendRMSNormGated
├── linear.py                 # AscendQKVParallelLinear, AscendRowParallelLinear
├── linear_op.py              # CustomLinearOp, MLPColumnParallelOp, DSV4OProjColumnParallelOp
├── rotary_embedding.py       # AscendRotaryEmbedding, AscendYaRNRotaryEmbedding
├── rope_dsv4.py              # ComplexExpRotaryEmbedding, RopeGlobalState
├── mla.py                    # AscendMultiHeadLatentAttention, mla_forward()
├── dsa.py                    # AscendDeepseekSparseAttention, dsa_forward()
├── gdn.py                    # AscendGatedDeltaNetAttention
├── gdn_attn_builder.py       # AscendGDNAttentionMetadataBuilder
├── kimi_kda.py               # AscendKimiK3DeltaAttention
├── conv.py                   # AscendConv3dLayer
├── cv_linear.py              # CVLinearWrapper
├── rel_pos_attention.py      # AscendRelPosAttention
├── vocab_parallel_embedding.py
├── bailing_moe_linear_attn.py # AscendBailingMoELinearAttention
├── mm_encoder_attention.py    # 多模态编码器注意力
├── qwen2_decoder.py          # Qwen2解码器算子
├── register_custom_ops.py    # 算子注册入口
├── triton/                   # Triton内核目录（30+文件）
│   ├── rms_norm.py           # RMSNorm Triton内核
│   ├── rope.py               # RoPE Triton内核
│   ├── batch_memcpy.py       # 批量内存拷贝
│   ├── bincount.py           # 计数内核
│   ├── penalty.py            # 惩罚采样内核
│   ├── eplb.py               # EPLB Triton内核
│   ├── sfa_cp.py             # SFA上下文并行内核
│   ├── dsa_cp.py             # DSA上下文并行内核
│   ├── kda/                  # Kimi KDA内核（kda.py, fused_recurrent_kda.py等）
│   ├── fla/                  # FLA闪存线性注意力内核（chunk.py, wy_fast.py等）
│   ├── linearnorm/           # 融合线性归一化内核
│   ├── batch_invariant/      # 批次不变性内核
│   ├── activation/           # 激活函数内核（swiglu_quant.py等）
│   └── mamba/                # Mamba模型内核
└── fused_moe/                # MoE融合算子（10+文件）
    ├── fused_moe.py          # 融合MoE入口
    ├── moe_mlp.py            # MoE MLP计算
    ├── routed_experts.py     # 路由专家
    ├── shared_experts.py     # 共享专家
    ├── token_dispatcher.py   # Token分发器
    ├── gate_linear.py        # 门控线性层
    ├── eplb.py               # 专家并行负载均衡
    └── router/               # 路由器

三、基础算子详解

3.1 RMSNorm算子示例

graph TD
    A[输入X] --> B[计算均方根]
    B --> C[归一化 X/RMS]
    C --> D[乘以权重 weight]
    D --> E[输出]
                

class AscendRMSNorm(torch.nn.Module):
    def __init__(self, hidden_size, eps=1e-6):
        super().__init__()
        self.weight = torch.nn.Parameter(torch.ones(hidden_size))
        self.eps = eps

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # 使用NPU专用内核计算RMSNorm
        return rms_norm_npu(x, self.weight, self.eps)

四、模型专用算子

五、MoE融合算子 (fused_moe/)

MoE（Mixture of Experts）融合算子是大型模型推理的关键组件，负责专家路由、门控计算和负载均衡。

graph LR
    A[fused_moe.py] --> B[gate_linear.py 门控计算]
    A --> C[router/ 路由器]
    A --> D[routed_experts.py 路由专家]
    A --> E[shared_experts.py 共享专家]
    A --> F[token_dispatcher.py Token分发]
    A --> G[eplb.py 负载均衡]
    
    C -->|选择专家| D
    F -->|分发Token| D
    D -->|计算结果| F
                

六、Triton内核 (triton/)

Triton是内核编译器，用于编写高性能GPU/NPU内核。triton-ascend是Triton的昇腾版本，支持在NPU上编写高性能内核。ops/triton/目录包含30+个Triton内核文件。

triton/
├── rms_norm.py           # RMSNorm Triton内核
├── rope.py               # RoPE Triton内核
├── batch_memcpy.py       # 批量内存拷贝内核
├── bincount.py           # 计数内核
├── penalty.py            # 惩罚采样内核
├── eplb.py               # EPLB Triton内核
├── sfa_cp.py             # SFA上下文并行内核
├── dsa_cp.py             # DSA上下文并行内核
├── kda/                  # Kimi KDA内核
│   ├── kda.py            # KDA主内核
│   ├── fused_recurrent_kda.py  # 融合递归KDA
│   ├── chunk_delta_h.py  # 块增量计算
│   ├── cumsum.py         # 累积和
│   ├── solve_tril.py     # 下三角求解
│   └── l2norm.py         # L2归一化
├── fla/                  # FLA闪存线性注意力内核
│   ├── chunk.py          # 块处理
│   ├── chunk_o.py        # 块输出
│   ├── chunk_delta_h.py  # 块增量
│   ├── wy_fast.py        # WY快速算法
│   └── l2norm.py         # L2归一化
├── linearnorm/           # 融合线性归一化内核
│   ├── split_qkv_rmsnorm_rope.py  # 融合QKV+RMSNorm+RoPE
│   ├── split_qkv_rmsnorm_mrope.py  # 融合QKV+RMSNorm+MRoPE
│   └── split_qkv_tp_rmsnorm_rope.py  # 融合QKV+TP+RMSNorm+RoPE
├── batch_invariant/      # 批次不变性内核
│   ├── softmax.py        # 批次不变Softmax
│   ├── rmsnorm.py        # 批次不变RMSNorm
│   ├── mean.py           # 批次不变均值
│   └── matmul.py         # 批次不变矩阵乘法
├── activation/           # 激活函数内核
│   ├── swiglu_quant.py   # 融合SwiGLU+量化
│   └── swiglustep.py     # SwiGLU Step
└── mamba/                # Mamba模型内核
    ├── causal_conv1d.py  # 因果卷积
    ├── lightning_attn.py  # Lightning注意力
    └── postprocess.py    # 后处理

七、算子注册流程

sequenceDiagram
    participant Dev as 开发者
    participant Ops as ops/算子
    participant Reg as register_custom_ops
    participant Model as models/模型
    participant NPU as NPU设备
    
    Dev->>Ops: 定义算子类
    Ops->>Reg: 注册到vLLM torch library
    Model->>Reg: 查找已注册算子
    Reg-->>Model: 返回算子实现
    Model->>NPU: 调用算子执行
    NPU-->>Model: 返回计算结果
                

# 注册自定义算子到vLLM的torch library
def _maybe_all_gather_and_maybe_unpad_impl(...):
    # AllGather + Unpad融合算子
    ...

def _maybe_pad_and_reduce_impl(...):
    # Pad + Reduce融合算子
    ...

def _quantize_impl(...):
    # 量化算子
    ...