Skip to content

Commit c9fa898

Browse files
committed
[Compat] Add docs for cross ecosystem custom ops
1 parent 4850894 commit c9fa898

File tree

5 files changed

+128
-2
lines changed

5 files changed

+128
-2
lines changed
Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,46 @@
1+
# 跨生态自定义算子接入
2+
3+
## 概述
4+
5+
随着大模型的兴起,在深度学习框架之上构建自定义算子(Custom Operator)已成为提升模型性能和功能的关键手段。而目前 PyTorch 作为深度学习领域的主流框架之一,拥有大量的自定义算子实现。为了帮助用户更好地将现有的 PyTorch 自定义算子迁移至 PaddlePaddle 框架,我们提供了自定义算子兼容机制,旨在降低迁移成本,提升开发效率。
6+
7+
## 方案介绍
8+
9+
为了方便 PyTorch 自定义算子快速接入 PaddlePaddle 框架,我们提供了如下图所示的兼容机制:
10+
11+
![PyTorch 自定义算子兼容机制示意图](./images/pytorch-op-compatible.drawio.png)
12+
13+
正如图上所示,我们自底向上提供了如下几层支持:
14+
15+
- **C++ API 兼容层**:该层实现了常用 PyTorch C++ API 的兼容接口,用户仍然可以通过调用 PyTorch 风格的 `at::*``torch::*``c10::*` 等命名空间下的函数和类来实现自定义算子逻辑,从而最大限度地复用现有代码,使迁移工作量降至最低。
16+
- **算子注册兼容层**:对于使用 pybind11 进行算子注册的 PyTorch 自定义算子,PaddlePaddle 无需额外修改注册代码;而对于使用 `TORCH_LIBRARY` 宏进行注册并通过 `torch.ops` 调用的算子,我们提供了同名的注册接口,用户无需修改注册代码即可完成迁移。
17+
- **Python 接口兼容层**:对于 Python 端自定义算子封装部分,会不可避免地调用一些 PyTorch 内的 Python 组网 API。为此,我们正在致力于提升 Python 端 API 与 PyTorch 的兼容性,力求让用户在迁移过程中无需修改 Python 端代码。
18+
- **Python API 代理层**:在 Python 端,即便 API 能够完全兼容,用户仍然需要将 `import torch` 替换为 `import paddle`。为此,我们提供了一个轻量级的代理层,用户只需在迁移后的代码开头添加一行 `import paddle.compat.enable_torch_proxy`,后续的 `torch` 下的模块将被重定向至 `paddle` 下的模块,从而实现无缝迁移。
19+
20+
通过以上几层兼容机制,用户可以在最大程度上复用现有的 PyTorch 自定义算子代码,从而大幅降低迁移成本。
21+
22+
此外,对于 TVM FFI 生态的自定义算子,由于我们已经对 TVM FFI 中所需的 DLPack 协议提供了最佳支持,因此用户可以直接将 TVM FFI 生态的自定义算子迁移至 PaddlePaddle 框架,无需额外修改。当然,如果相关算子库在 Python 端调用了 PyTorch 的组网 API,则仍然需要借助上述的 Python API 代理层来完成迁移。
23+
24+
## 迁移步骤
25+
26+
下面我们以一个简单的 PyTorch 自定义算子为例,介绍如何将其迁移至 PaddlePaddle 框架。
27+
28+
## 已迁移算子库
29+
30+
### FlashInfer
31+
32+
### FlashMLA
33+
34+
### DeepGEMM
35+
36+
### TileLang
37+
38+
### Triton
39+
40+
### TorchCodec
41+
42+
### DeepEP
43+
44+
coming soon...
45+
46+
## 参考资料
174 KB
Loading
Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,14 @@
1+
#############
2+
跨生态自定义算子接入
3+
#############
4+
5+
本章节介绍如何将其他深度学习框架算子生态的自定义算子迁移至飞桨框架,主要分为以下内容:
6+
7+
- `使用指南 <./user_guide.html>`_ 介绍跨生态自定义算子的使用方法,以及已经成功接入的自定义算子库列表。
8+
- `原理和迁移 <./design.html>`_ 介绍跨生态自定义算子的设计原理,以及如何将其他深度学习框架的自定义算子迁移至飞桨框架。
9+
10+
.. toctree::
11+
:hidden:
12+
13+
user_guide_cn.md
14+
design_cn.md
Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,64 @@
1+
# 使用指南
2+
3+
## 概述
4+
5+
随着大模型的兴起,在深度学习框架之上构建自定义算子(Custom Operator)已成为提升模型性能和功能的关键手段。而目前 PyTorch 作为深度学习领域的主流框架之一,拥有大量的自定义算子实现。为了帮助用户更好地将现有的 PyTorch 等生态的自定义算子迁移至 PaddlePaddle 框架,我们提供了自定义算子兼容机制,旨在降低迁移成本,提升开发效率。
6+
7+
## 使用步骤
8+
9+
### 一般安装方式
10+
11+
对于使用基于兼容性方案的跨生态自定义算子库,一般情况下只需要 clone 后通过 pip 安装对应的算子库即可使用。下面以 FlashMLA 为例说明安装方式:
12+
13+
```bash
14+
git clone https://github.com/PFCCLab/FlashMLA.git
15+
cd FlashMLA
16+
pip install .
17+
```
18+
19+
对于部分已经发布到 PyPI 的自定义算子库,也可以直接通过 pip 安装。下面以 TorchCodec 为例:
20+
21+
```bash
22+
pip install paddlecodec
23+
```
24+
25+
个别算子库可能会有特殊的安装方式,请参考对应算子库的说明文档进行安装。
26+
27+
### 使用方式
28+
29+
安装完成后,即可在代码中直接导入并使用对应的算子库。为了实现跨生态兼容,用户需要在导入算子库之前,先启用 PaddlePaddle 的 PyTorch 代理层,以确保算子库中 `torch` 模块的调用能够正确映射到 `paddle` 模块。下面以 FlashMLA 为例说明使用方式:
30+
31+
```python
32+
import paddle
33+
34+
paddle.compat.enable_torch_proxy({"flash_mla"})
35+
36+
import flash_mla
37+
# 之后即可使用 flash_mla 下的算子
38+
```
39+
40+
## 已支持的算子库
41+
42+
PaddlePaddle 官方协同社区已经对社区中主流的跨生态自定义算子库进行了适配和测试,用户可以直接使用这些算子库而无需进行额外的修改。
43+
44+
我们将这些算子库统一放在组织 [PFCCLab](https://github.com/PFCCLab) 下,并列在下方。如果下方列表中没有你需要的算子库,可以移步至[原理和迁移](./design_cn.md),了解自定义算子兼容机制的实现原理,以及如何将你需要的算子库进行迁移。
45+
46+
### FlashInfer
47+
48+
#### 安装方式
49+
50+
#### 使用方式
51+
52+
### FlashMLA
53+
54+
### DeepGEMM
55+
56+
### DeepEP
57+
58+
### TorchCodec
59+
60+
## 已支持的 Kernel DSL 库
61+
62+
### TileLang
63+
64+
### Triton

docs/guides/custom_op/index_cn.rst

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -8,11 +8,12 @@
88
2. C++扩展:可将 C++ 类绑定至 Python,支持调用无 Tensor 参数的 C++ 函数。不涉及框架的调度开销
99
3. Python 算子:使用 Python 编写实现前向(forward)和反向(backward)方法,在模型组网中使用的自定义 API
1010

11-
- `自定义 C++算子 <./new_cpp_op_cn.html>`_
11+
此外还会介绍如何将其他深度学习框架算子生态的自定义算子迁移至飞桨框架。
1212

13+
- `自定义 C++算子 <./new_cpp_op_cn.html>`_
1314
- `自定义 C++ 扩展 <./cpp_extension_cn.html>`_
14-
1515
- `自定义 Python 算子 <./new_python_op_cn.html>`_
16+
- `跨生态自定义算子接入 <./cross_ecosystem_custom_op/index_cn.html>`_
1617

1718

1819
.. toctree::
@@ -21,3 +22,4 @@
2122
new_cpp_op_cn.md
2223
cpp_extension_cn.md
2324
new_python_op_cn.md
25+
cross_ecosystem_custom_op/index_cn.rst

0 commit comments

Comments
 (0)