Skip to content

Commit ac9e282

Browse files
chenchris2chenchris2
committed
[Feature] Add FFN op for Ascend950
Signed-off-by: chenchris2 <1349418798@qq.com>
1 parent 7d81170 commit ac9e282

127 files changed

Lines changed: 28104 additions & 0 deletions

File tree

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎csrc/build_aclnn.sh‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -225,6 +225,7 @@ elif [[ "$SOC_VERSION" =~ ^ascend950 ]]; then
225225
"k2q_csr"
226226
"sparse_attention_score"
227227
"mla_prolog_v3"
228+
"ffn"
228229
)
229230

230231
CUSTOM_OPS=$(IFS=';'; echo "${CUSTOM_OPS_ARRAY[*]}")

‎csrc/ffn/CMakeLists.txt‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,18 @@
1+
# -----------------------------------------------------------------------------------------------------------
2+
# Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+
# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+
# CANN Open Software License Agreement Version 2.0 (the "License").
5+
# Please refer to the License for details. You may not use this file except in compliance with the License.
6+
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+
# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY OR FITNESS FOR A PARTICULAR PURPOSE.
8+
# See LICENSE in the root of the software repository for the full text of the License.
9+
# -----------------------------------------------------------------------------------------------------------
10+
11+
file(GLOB SUBDIRECTORIES LIST_DIRECTORIES true RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*)
12+
# 遍历子目录
13+
foreach(SUBDIR ${SUBDIRECTORIES})
14+
# 检查子目录中是否存在 CMakeLists.txt
15+
if(EXISTS ${CMAKE_CURRENT_SOURCE_DIR}/${SUBDIR}/CMakeLists.txt)
16+
add_subdirectory(${SUBDIR})
17+
endif()
18+
endforeach()

‎csrc/ffn/ffn/CMakeLists.txt‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
# -----------------------------------------------------------------------------------------------------------
2+
# Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+
# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+
# CANN Open Software License Agreement Version 2.0 (the "License").
5+
# Please refer to the License for details. You may not use this file except in compliance with the License.
6+
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+
# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+
# See LICENSE in the root of the software repository for the full text of the License.
9+
# -----------------------------------------------------------------------------------------------------------
10+
11+
file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*)
12+
if(NOT ENABLE_TEST AND NOT BENCHMARK)
13+
list(REMOVE_ITEM CURRENT_DIRS tests)
14+
endif()
15+
foreach(SUB_DIR ${CURRENT_DIRS})
16+
if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt")
17+
add_subdirectory(${SUB_DIR})
18+
endif()
19+
endforeach()

‎csrc/ffn/ffn/README.md‎

Lines changed: 188 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,188 @@
1+
# FFN
2+
3+
## 产品支持情况
4+
5+
| 产品 | 是否支持 |
6+
| :----------------------------------------------------------- | :------: |
7+
| <term>Ascend 950PR/Ascend 950DT</term> | √ |
8+
| <term>Atlas A3 训练系列产品/Atlas A3 推理系列产品</term> | √ |
9+
| <term>Atlas A2 训练系列产品/Atlas A2 推理系列产品</term> | × |
10+
| <term>Atlas 200I/500 A2 推理产品</term> | × |
11+
| <term>Atlas 推理系列产品</term> | × |
12+
| <term>Atlas 训练系列产品</term> | × |
13+
| <term>Kirin X90 处理器系列产品</term> | √ |
14+
| <term>Kirin 9030 处理器系列产品</term> | √ |
15+
16+
## 功能说明
17+
18+
- 算子功能：该FFN算子提供MoeFFN和FFN的计算功能。在没有专家分组（expertTokens为空）时是FFN，有专家分组时是MoeFFN，统称为FFN，属于Moe结构。MoE（Mixture-of-Experts，混合专家系统）是一种用于训练万亿参数量级模型的技术。MoE将预测建模任务分解为若干子任务，在每个子任务上训练一个专家模型（Expert Model），开发一个门控模型（Gating Model），该模型会根据输入数据分配一个或多个专家，最终综合多个专家计算结果作为预测结果。Mixture-of-Experts结构的模型是将输入数据分配给最相关的一个或者多个专家，综合涉及的所有专家的计算结果来确定最终结果。
19+
- 计算公式：
20+
21+
- **非量化场景：**
22+
23+
$$
24+
y=activation(x * W1 + b1) * W2 + b2
25+
$$
26+
27+
- **量化场景：**
28+
29+
$$
30+
y=((activation((x * W1 + b1) * deqScale1) * scale + offset) * W2 + b2) * deqScale2
31+
$$
32+
33+
- **伪量化场景：**
34+
35+
$$
36+
y=activation(x * ((W1 + antiquantOffset1) * antiquantScale1) + b1) * ((W2 + antiquantOffset2) * antiquantScale2) + b2
37+
$$
38+
39+
## 参数说明
40+
41+
<table style="undefined;table-layout: fixed; width: 1050px"><colgroup>
42+
<col style="width: 150px">
43+
<col style="width: 300px">
44+
<col style="width: 200px">
45+
<col style="width: 300px">
46+
<col style="width: 100px">
47+
</colgroup>
48+
<thead>
49+
<tr>
50+
<th>参数名</th>
51+
<th>输入/输出/属性</th>
52+
<th>描述</th>
53+
<th>数据类型</th>
54+
<th>数据格式</th>
55+
</tr></thead>
56+
<tbody>
57+
<tr>
58+
<td>x</td>
59+
<td>输入</td>
60+
<td>必选参数，Device侧的aclTensor，公式中的输入x，支持输入的维度最少是2维[M, K1]，最多是8维。</td>
61+
<td>FLOAT16、BFLOAT16、INT8</td>
62+
<td>ND</td>
63+
</tr>
64+
<tr>
65+
<td>weight1</td>
66+
<td>输入</td>
67+
<td>必选参数，Device侧的aclTensor，专家的权重数据，公式中的W1，输入在有/无专家时分别为[E, K1, N1]/[K1, N1]。</td>
68+
<td>FLOAT16、BFLOAT16、INT8、INT4</td>
69+
<td>ND</td>
70+
</tr>
71+
<tr>
72+
<td>weight2</td>
73+
<td>输入</td>
74+
<td>必选参数，Device侧的aclTensor，专家的权重数据，公式中的W2，输入在有/无专家时分别为[E, K2, N2]/[K2, N2]。</td>
75+
<td>FLOAT16、BFLOAT16、INT8、INT4</td>
76+
<td>ND</td>
77+
</tr>
78+
<tr>
79+
<td>expertTokens</td>
80+
<td>输入</td>
81+
<td>可选参数，Host侧的aclIntArray类型，代表各专家的token数，若不为空时可支持的最大长度为256个。</td>
82+
<td>INT64</td>
83+
<td>ND</td>
84+
</tr>
85+
<tr>
86+
<td>bias1</td>
87+
<td>输入</td>
88+
<td>可选参数，Device侧的aclTensor，权重数据修正值，公式中的b1，输入在有/无专家时分别为[E, N1]/[N1]。</td>
89+
<td>FLOAT16、BFLOAT16、FLOAT32、INT32</td>
90+
<td>ND</td>
91+
</tr>
92+
<tr>
93+
<td>bias2</td>
94+
<td>输入</td>
95+
<td>可选参数，Device侧的aclTensor，权重数据修正值，公式中的b2，输入在有/无专家时分别为[E, N2]/[N2]。</td>
96+
<td>FLOAT16、BFLOAT16、FLOAT32、INT32</td>
97+
<td>ND</td>
98+
</tr>
99+
<tr>
100+
<td>scale</td>
101+
<td>输入</td>
102+
<td>可选参数，Device侧的aclTensor，量化参数，量化缩放系数，per-tensor下输入在有/无专家时均为一维向量，输入元素个数在有/无专家时分别为[E]/[1]；per-channel下输入在有/无专家时为二维向量/一维向量，输入元素个数在有/无专家时分别为[E, N1]/[N1]。</td>
103+
<td>FLOAT32</td>
104+
<td>ND</td>
105+
</tr>
106+
<tr>
107+
<td>offset</td>
108+
<td>输入</td>
109+
<td>可选参数，Device侧的aclTensor，量化参数，量化偏移量，一维向量，输入元素个数在有/无专家时分别为[E]/[1]。</td>
110+
<td>FLOAT32</td>
111+
<td>ND</td>
112+
</tr>
113+
<tr>
114+
<td>deqScale1</td>
115+
<td>输入</td>
116+
<td>可选参数，Device侧的aclTensor，量化参数，第一个matmul的反量化缩放系数，输入在有/无专家时分别为[E, N1]/[N1]。</td>
117+
<td>UINT64、INT64、FLOAT32、BFLOAT16</td>
118+
<td>ND</td>
119+
</tr>
120+
<tr>
121+
<td>deqScale2</td>
122+
<td>输入</td>
123+
<td>可选参数，Device侧的aclTensor，量化参数，第二个matmul的反量化缩放系数，输入在有/无专家时分别为[E, N2]/[N2]。</td>
124+
<td>UINT64、INT64、FLOAT32、BFLOAT16</td>
125+
<td>ND</td>
126+
</tr>
127+
<tr>
128+
<td>antiquantScale1</td>
129+
<td>输入</td>
130+
<td>可选参数，Device侧的aclTensor，伪量化参数，第一个matmul的缩放系数，per-channel下输入在有/无专家时分别为[E, N1]/[N1]，per-group下输入在有/无专家时分别为[E, G, N1]/[G, N1]。</td>
131+
<td>FLOAT16、BFLOAT16</td>
132+
<td>ND</td>
133+
</tr>
134+
<tr>
135+
<td>antiquantScale2</td>
136+
<td>输入</td>
137+
<td>可选参数，Device侧的aclTensor，伪量化参数，第二个matmul的缩放系数，per-channel下输入在有/无专家时分别为[E, N2]/[N2]，per-group下输入在有/无专家时分别为[E, G, N2]/[G, N2]。</td>
138+
<td>FLOAT16、BFLOAT16</td>
139+
<td>ND</td>
140+
</tr>
141+
<tr>
142+
<td>antiquantOffset1</td>
143+
<td>输入</td>
144+
<td>可选参数，Device侧的aclTensor，伪量化参数，第一个matmul的偏移量，per-channel下输入在有/无专家时分别为[E, N1]/[N1]，per-group下输入在有/无专家时分别为[E, G, N1]/[G, N1]。</td>
145+
<td>FLOAT16、BFLOAT16</td>
146+
<td>ND</td>
147+
</tr>
148+
<tr>
149+
<td>antiquantOffset2</td>
150+
<td>输入</td>
151+
<td>可选参数，Device侧的aclTensor，伪量化参数，第二个matmul的偏移量，per-channel下输入在有/无专家时分别为[E, N2]/[N2]，per-group下输入在有/无专家时分别为[E, G, N2]/[G, N2]。</td>
152+
<td>FLOAT16、BFLOAT16</td>
153+
<td>ND</td>
154+
</tr>
155+
</tbody>
156+
</table>
157+
158+
## 约束说明
159+
160+
- 有专家时，专家数据的总数需要与x的M保持一致。
161+
- 激活层为geglu/swiglu/reglu时，仅支持无专家分组时的FLOAT16高性能场景（FLOAT16场景指类型为aclTensor的必选参数数据类型都为FLOAT16的场景），且N1=2\*K2。
162+
- 激活层为gelu/fastgelu/relu/silu时，支持有专家或无专家分组的FLOAT16高精度及高性能场景、BFLOAT16场景、量化场景及伪量化场景，且N1=K2。
163+
- 所有场景下需满足K1=N2, K1<65536, K2<65536, M轴在32Byte对齐后小于INT32的最大值。
164+
- 非量化场景不能输入量化参数和伪量化参数，量化场景不能输入伪量化参数，伪量化场景不能输入量化参数。
165+
- 量化场景参数类型：x为INT8、weight为INT8、bias为INT32、scale为FLOAT32、offset为FLOAT32，其余参数类型根据y不同分两种情况：
166+
- y为FLOAT16，deqScale支持数据类型：UINT64、INT64、FLOAT32。
167+
- y为BFLOAT16，deqScale支持数据类型：BFLOAT16。
168+
- 要求deqScale1与deqScale2的数据类型保持一致。
169+
- 量化场景支持scale的per-channel模式参数类型：x为INT8、weight为INT8、bias为INT32、scale为FLOAT32、offset为FLOAT32，其余参数类型根据y不同分两种情况：
170+
- y为FLOAT16，deqScale支持数据类型：UINT64、INT64。
171+
- y为BFLOAT16，deqScale支持数据类型：BFLOAT16。
172+
- 要求deqScale1与deqScale2的数据类型保持一致。
173+
- 伪量化场景支持两种不同参数类型：
174+
- y为FLOAT16、x为FLOAT16、bias为FLOAT16，antiquantScale为FLOAT16、antiquantOffset为FLOAT16，weight支持数据类型INT8和INT4。
175+
- y为BFLOAT16、x为BFLOAT16、bias为FLOAT32，antiquantScale为BFLOAT16、antiquantOffset为BFLOAT16，weight支持数据类型INT8和INT4。
176+
- 当weight1/weight2的数据类型为INT4时，其shape最后一维必须为偶数。
177+
- 伪量化场景，per-group下，antiquantScale1和antiquantOffset1中的K1需要能整除组数G，antiquantScale2和antiquantOffset2中的K2需要能整除组数G。
178+
- 伪量化场景，per-group下目前只支持weight是INT4数据类型的场景。
179+
- innerPrecise参数在BFLOAT16非量化场景，只能配置为0；FLOAT16非量化场景，可以配置为0或者1；量化或者伪量化场景，0和1都可配置，但是配置后不生效。
180+
- Ascend 950PR/950DT仅支持非量化、单专家场景，激活为gelu/silu/swiglu，x/weight数据类型为FLOAT16或BFLOAT16。
181+
- Ascend 950PR/950DT下bias1与bias2需同时存在或同时不存在，支持与x同dtype或FLOAT32的bias。
182+
- Ascend 950PR/950DT下swiglu仅支持linear布局weight1 [2H, K]，gelu/silu支持canonical [K, H]与linear [H, K]双布局。
183+
184+
## 调用说明
185+
186+
| 调用方式 | 调用样例 | 说明 |
187+
|--------------|-------------------------|--------------------------------------------------------------|
188+
| aclnn调用 | [test_aclnn_ffn.cpp](examples/test_aclnn_ffn.cpp) | 通过接口方式调用[FFN](docs/aclnnFFN.md)算子。 |

‎csrc/ffn/ffn/ffn_linear_torch_adpt.h‎

Lines changed: 115 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,115 @@
1+
/*
2+
* Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
#ifndef FFN_LINEAR_TORCH_ADPT_H
17+
#define FFN_LINEAR_TORCH_ADPT_H
18+
19+
#include <cctype>
20+
#include <string>
21+
#include <vector>
22+
23+
#include "op_host/ffn_layout.h"
24+
25+
namespace vllm_ascend {
26+
27+
inline bool ffn_linear_is_swiglu(const std::string &act)
28+
{
29+
static constexpr char kSwiglu[] = "swiglu";
30+
if (act.size() != sizeof(kSwiglu) - 1) {
31+
return false;
32+
}
33+
for (size_t i = 0; i < act.size(); ++i) {
34+
if (std::tolower(static_cast<unsigned char>(act[i])) != kSwiglu[i]) {
35+
return false;
36+
}
37+
}
38+
return true;
39+
}
40+
41+
// 布局识别走公共规则（op_host/ffn_layout.h）：全方阵/消歧失败默认 linear（PyTorch Linear 惯例）。
42+
inline std::vector<int64_t> ffn_linear_npu_output_size(const at::Tensor &x, const at::Tensor &weight1,
43+
const at::Tensor &weight2, const std::string &activation)
44+
{
45+
auto xSizes = x.sizes().vec();
46+
const int64_t xK = xSizes.back();
47+
const int64_t w1d0 = weight1.size(-2);
48+
const int64_t w1d1 = weight1.size(-1);
49+
const int64_t w2d0 = weight2.size(-2);
50+
const int64_t w2d1 = weight2.size(-1);
51+
const bool swiglu = ffn_linear_is_swiglu(activation);
52+
const ffnlayout::FfnLayout layout = ffnlayout::FfnDetectLayout(w1d0, w1d1, w2d0, w2d1, xK, swiglu);
53+
TORCH_CHECK(layout != ffnlayout::FfnLayout::INVALID, "weight1 shape [", w1d0, ", ", w1d1,
54+
"] does not match x K=", xK, " (expect [K,N] canonical or [N,K] linear)");
55+
const bool isLinear = (layout == ffnlayout::FfnLayout::LINEAR);
56+
if (swiglu && !isLinear) {
57+
TORCH_CHECK(false, "swiglu only supports linear layout weight1 [2H,K]");
58+
}
59+
xSizes[xSizes.size() - 1] = isLinear ? w2d0 : w2d1;
60+
return xSizes;
61+
}
62+
63+
// FFN 融合算子（arch35）：y = act(x @ W1^T + b1) @ W2^T + b2，bf16/fp16，gelu/silu/swiglu。
64+
// 直接接收 PyTorch Linear 布局权重 [N,K]（out,in），kernel 内部通过 transB 处理。
65+
inline at::Tensor ffn_linear(const at::Tensor &x, const at::Tensor &weight1, const at::Tensor &weight2,
66+
const c10::optional<at::Tensor> &bias1_opt, const c10::optional<at::Tensor> &bias2_opt,
67+
c10::string_view activation, int64_t inner_precise)
68+
{
69+
TORCH_CHECK(x.dim() >= 2, "x must have at least 2 dimensions");
70+
TORCH_CHECK(weight1.dim() == 2 && weight2.dim() == 2, "weight1/weight2 must be 2D");
71+
72+
at::Tensor x_contiguous = x.contiguous();
73+
at::Tensor w1_contiguous = weight1.contiguous();
74+
at::Tensor w2_contiguous = weight2.contiguous();
75+
76+
c10::optional<at::Tensor> bias1 = (bias1_opt.has_value() && bias1_opt.value().defined())
77+
? c10::optional<at::Tensor>(bias1_opt.value().contiguous())
78+
: c10::nullopt;
79+
c10::optional<at::Tensor> bias2 = (bias2_opt.has_value() && bias2_opt.value().defined())
80+
? c10::optional<at::Tensor>(bias2_opt.value().contiguous())
81+
: c10::nullopt;
82+
83+
const c10::optional<at::Tensor> expertTokens = c10::nullopt;
84+
const c10::optional<at::Tensor> scale = c10::nullopt;
85+
const c10::optional<at::Tensor> offset = c10::nullopt;
86+
const c10::optional<at::Tensor> deqScale1 = c10::nullopt;
87+
const c10::optional<at::Tensor> deqScale2 = c10::nullopt;
88+
const c10::optional<at::Tensor> antiquantScale1 = c10::nullopt;
89+
const c10::optional<at::Tensor> antiquantScale2 = c10::nullopt;
90+
const c10::optional<at::Tensor> antiquantOffset1 = c10::nullopt;
91+
const c10::optional<at::Tensor> antiquantOffset2 = c10::nullopt;
92+
93+
std::string actStr(activation.data(), activation.size());
94+
auto output_size = ffn_linear_npu_output_size(x_contiguous, w1_contiguous, w2_contiguous, actStr);
95+
at::Tensor y = at::empty(output_size, x_contiguous.options());
96+
97+
const char *activationCStr = actStr.c_str();
98+
const bool tokensIndexFlag = false;
99+
100+
EXEC_NPU_CMD(aclnnFFNV2,
101+
x_contiguous, w1_contiguous, w2_contiguous,
102+
expertTokens,
103+
bias1, bias2,
104+
scale, offset,
105+
deqScale1, deqScale2,
106+
antiquantScale1, antiquantScale2,
107+
antiquantOffset1, antiquantOffset2,
108+
activationCStr, inner_precise, tokensIndexFlag,
109+
y);
110+
111+
return y;
112+
}
113+
114+
} // namespace vllm_ascend
115+
#endif

0 commit comments

Comments
 (0)