[Kernel] add custom op GmmSwigluQuantWeightNzTensorList (#3804)
### What this PR does / why we need it? This PR introduces support for adding custom CANN `aclnn` ops to `vllm-ascend`, allowing users to define and use their own custom operators. Key changes include: - Building and installing custom ops into the `vllm-ascend`-specified directory - Binding the `aclnn` op interface to the `torch.ops._C_ascend` module - Enabling invocation of these ops within `vllm-ascend` This PR includes a sample custom op: `aclnnGroupedMatmulSwigluQuantWeightNzTensorList`, which is adapted from the CANN operator [`aclnnGroupedMatmulSwigluQuantWeightNZ`](https://www.hiascend.com/document/detail/zh/canncommercial/83RC1/API/aolapi/context/aclnnGroupedMatmulSwigluQuantWeightNZ.md). Its input parameters `weight` and `weight_scale` now accept `list[torch.Tensor]` (i.e., `at::TensorList`). ### Does this PR introduce _any_ user-facing change? No. - vLLM version: v0.11.2 --------- Signed-off-by: QianChenxi <chenxi.qian.cq@outlook.com>
This commit is contained in:
@@ -0,0 +1,25 @@
|
||||
# Adding a custom aclnn operation
|
||||
|
||||
This document describes how to add a custom aclnn operation to vllm-ascend.
|
||||
|
||||
## How custom aclnn operation works in vllm-ascend?
|
||||
|
||||
Custom aclnn operations are built and installed into `vllm_ascend/cann_ops_custom` directory during the build process of vllm-ascend. Then the aclnn operators are bound to `torch.ops._C_ascend` module, enabling users to invoke them in vllm-ascend python code.
|
||||
|
||||
To enable custom operations, use the following code:
|
||||
|
||||
```python
|
||||
from vllm_ascend.utils import enable_custom_op
|
||||
|
||||
enable_custom_op()
|
||||
```
|
||||
|
||||
## How to add a custom aclnn operation?
|
||||
|
||||
- Create a new operation folder under `csrc` directory
|
||||
- Create `op_host` and `op_kernel` directories for host and kernel source code
|
||||
- Add build options in `csrc/build_aclnn.sh` for supported SOC. Note that multiple ops should be separated with `;`, i.e. `CUSTOM_OPS=op1;op2;op3`
|
||||
- Bind aclnn operators to torch.ops._C_ascend module in `csrc/torch_binding.cpp`
|
||||
- Write a meta implementation in `csrc/torch_binding_meta.cpp` for op being captured into aclgraph
|
||||
|
||||
After a successful build of vllm-ascend, the custom aclnn operation can be invoked in python code.
|
||||
@@ -12,4 +12,5 @@ eplb_swift_balancer.md
|
||||
Multi_Token_Prediction
|
||||
ACL_Graph
|
||||
KV_Cache_Pool_Guide
|
||||
add_custom_aclnn_op
|
||||
:::
|
||||
|
||||
Reference in New Issue
Block a user