自定义算子接入
当内置算子不满足需求时，你可以利用MindSpore提供的自定义算子功能接入你的算子。
本文档将以 advance_step_flashattn 算子为例，讲解如何在vLLM-MindSpore插件项目中接入一个AscendC自定义算子。
本文重点在于介绍把算子集成进vLLM-MindSpore插件的流程，自定义算子的细节请参考 MindSpore 官方教程：基于CustomOpBuilder的自定义算子。AscendC算子的开发流程请参考昇腾官方文档：Ascend C算子开发。
注：目前vLLM-MindSpore插件的自定义算子仅支持动态图（PyNative Mode）场景。
文件组织结构
接入自定义算子需要在vLLM-MindSpore插件项目的 csrc 目录下添加代码，目录结构如下：
vllm-mindspore/
├── csrc/
│ ├── CMakeLists.txt // 算子编译脚本
│ ├── ascendc/
│ │ ├── CMakeLists.txt // AscendC 算子编译脚本
│ │ ├── adv_step_flash.h // AscendC AdvanceStepFlashattn 算子声明
│ │ ├── adv_step_flash.c // AscendC AdvanceStepFlashattn 算子实现
│ │ └── ...
│ └── module/
│ ├── module.h // 公共模块注册头文件
│ ├── module.cpp // 公共模块注册实现文件
│ ├── adv_step_flash.cpp // 接入层代码，注册 AdvanceStepFlashattn 算子的Python模块接口
│ └── ...
└── vllm_mindspore/
└── _custom_ops.py // 封装自定义算子调用接口
csrc/ascendc/：放置 AscendC 自定义算子的实现代码。csrc/module/：放置算子接入层代码，包括公共模块注册（module.h、module.cpp）和算子接入代码（如adv_step_flash.cpp）。
接入流程
接入一个自定义算子，在算子实现方面，需在csrc目录中，创建算子接口定义，算子实现与算子接入。在完成自定义算子初步的开发与接入后，可添加算子调用接口并进行算子编译和测试。
算子接口声明
在 csrc/ascendc/ 目录下，创建头文件（如 adv_step_flash.h），以声明算子接口，内容参考adv_step_flash.h：
#ifndef VLLM_MINDSPORE_CSRC_ASCENDC_ADV_STEP_FLASH_H
#define VLLM_MINDSPORE_CSRC_ASCENDC_ADV_STEP_FLASH_H
extern void AdvStepFlashKernelEntry(uint32_t blockDims, void *l2ctrl, void *aclStream, uint8_t *sampledTokenIds,
uint8_t *blockTables, uint8_t *seqLensInput, uint8_t *inputTokens,
uint8_t *inputPositions, uint8_t *seqLensOut, uint8_t *slotMapping,
int32_t num_seqs, int32_t block_size, int32_t block_tables_stride);
#endif // VLLM_MINDSPORE_CSRC_ASCENDC_ADV_STEP_FLASH_H
算子实现
在 csrc/ascendc/ 目录下创建实现文件（如 adv_step_flash.c），以实现算子的核心逻辑，内容参考adv_step_flash.c：
#include "kernel_operator.h"
extern "C" __global__ __aicore__ void adv_step_flash_impl(GM_ADDR sampledTokenIds, GM_ADDR blockTables,
GM_ADDR seqLensInput, GM_ADDR inputTokens,
GM_ADDR inputPositions, GM_ADDR seqLensOut,
GM_ADDR slotMapping, int32_t num_seqs, int32_t block_size,
int32_t block_tables_stride) {
// AscendC 算子实现
}
#ifndef __CCE_KT_TEST__
void AdvStepFlashKernelEntry(uint32_t blockDims, void *l2ctrl, void *aclStream, uint8_t *sampledTokenIds,
uint8_t *blockTables, uint8_t *seqLensInput, uint8_t *inputTokens, uint8_t *inputPositions,
uint8_t *seqLensOut, uint8_t *slotMapping, int32_t num_seqs, int32_t block_size,
int32_t block_tables_stride) {
adv_step_flash_impl<<<blockDims, l2ctrl, aclStream>>>(sampledTokenIds, blockTables, seqLensInput, inputTokens,
inputPositions, seqLensOut, slotMapping, num_seqs, block_size,
block_tables_stride);
}
#endif
算子接入
在 csrc/module/ 目录下创建一个新的接入文件（如 adv_step_flash.cpp），内容参考 adv_step_flash.cpp：
#include "ms_extension/api.h"
#include "ascendc/adv_step_flash.h"
#include "module/module.h"
auto pyboost_adv_step_flash(int32_t num_seqs, int32_t num_queries, int32_t block_size,
ms::Tensor input_tokens,
ms::Tensor sampled_token_ids,
ms::Tensor input_positions,
ms::Tensor seq_lens,
ms::Tensor slot_mapping,
ms::Tensor block_tables) {
// 使用 ms::PyboostRunner 调用你的算子
}
VLLM_MS_EXTENSION_MODULE(m) {
m.def("advance_step_flashattn", &pyboost_adv_step_flash, "advance_step_flashattn", pybind11::arg("num_seqs"),
pybind11::arg("num_queries"), pybind11::arg("block_size"), pybind11::arg("input_tokens"),
pybind11::arg("sampled_token_ids"), pybind11::arg("input_positions"), pybind11::arg("seq_lens"),
pybind11::arg("slot_mapping"), pybind11::arg("block_tables"));
}
上面m.def()接口的第一个参数"advance_step_flashattn"就是算子的Python接口名。
module.h 和 module.cpp 文件的作用是基于pybind11创建算子的Python模块。因为一个动态库内只能有一个 PYBIND11_MODULE ，为了让用户可以在一个文件内完成算子接入工作，vLLM-MindSpore插件提供了一个新的注册接口 VLLM_MS_EXTENSION_MODULE 宏。自定义算子动态库加载时，所有算子接口都会被自动注册到同一个Python模块中。
算子调用接口
vLLM-MindSpore插件的自定义算子被编译到了 _C_ops.so 里面，为了方便调用，可以在 vllm_mindspore/_custom_ops.py 添加一个调用接口。如果在算子调用前后需要做额外适配，也可以在这接口内实现。
def advance_step_flashattn(num_seqs: int, num_queries: int, block_size: int,
input_tokens: torch.Tensor,
sampled_token_ids: torch.Tensor,
input_positions: torch.Tensor,
seq_lens: torch.Tensor,
slot_mapping: torch.Tensor,
block_tables: torch.Tensor) -> None:
"""Advance a step on Ascend for existing inputs for a multi-step runner"""
from vllm_mindspore import _C_ops as c_ops
c_ops.advance_step_flashattn(num_seqs=num_seqs,
num_queries=num_queries,
block_size=block_size,
input_tokens=input_tokens,
sampled_token_ids=sampled_token_ids,
input_positions=input_positions,
seq_lens=seq_lens,
slot_mapping=slot_mapping,
block_tables=block_tables)
这里通过 import _C_ops 即可导入自定义算子的Python模块，推荐在调用前再导入，未调用时不需要导入。
算子编译和测试
代码集成：将代码集成至vLLM-MindSpore插件项目。
编译项目：在项目代码根目录下，执行
pip install .，编译安装vLLM-MindSpore插件。测试算子接口：通过
_custom_ops调用算子接口，可以参考测试用例test_custom_advstepflash.py：
from vllm_mindspore import _custom_ops as custom_ops
custom_ops.advance_step_flashattn(...)
自定义算子编译工程
当前MindSpore仅提供了一个 CustomOpBuilder接口 用于在线编译自定义算子，接口内置了默认的编译和链接选项。vLLM-MindSpore插件基于MindSpore的自定义算子功能接入算子，并编译成动态库随包发布。下面是编译流程介绍：
算子扩展库模块
在 setup.py 中，vLLM-MindSpore插件添加了一个 vllm_mindspore._C_ops 扩展，并添加了相应的编译模块：
ext_modules = [Extension("vllm_mindspore._C_ops", sources=[])],
cmdclass = {"build_ext": CustomBuildExt},
这里不需要指定 sources ，是因为vLLM-MindSpore插件通过CMake触发算子编译，自动收集了源文件。
算子编译流程
CustomBuildExt调用CMake执行csrc/CMakeLists.txt，传入必要的环境变量，触发算子编译。通过
ascendc/CMakeLists.txt调用AscendC编译器，编译ascendc目录内的算子算子源码。生成静态库ascendc_kernels_npu.a。递归收集cpp源文件列表
SRC_FILES。生成临时脚本
build_custom_with_ms.py，文件内调用mindspore.CustomOpBuilder编译算子接口。文件里也写入了源文件列表、头文件路径和静态库路径等信息。通过CMake的
add_custom_target命令调用Python脚本编译自定义算子，生成_C_ops.so。将
_C_ops.so重命名成Python包扩展模块的标准名称，如_C_ops.cpython-39-aarch64-linux-gnu.so。