18
csrc/gmm/grouped_matmul_swiglu_quant/CMakeLists.txt
Normal file
18
csrc/gmm/grouped_matmul_swiglu_quant/CMakeLists.txt
Normal file
@@ -0,0 +1,18 @@
|
||||
# -----------------------------------------------------------------------------------------------------------
|
||||
# Copyright (c) 2025 Huawei Technologies Co., Ltd.
|
||||
# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
# CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
# Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
# See LICENSE in the root of the software repository for the full text of the License.
|
||||
# -----------------------------------------------------------------------------------------------------------
|
||||
|
||||
file(GLOB SUBDIRECTORIES LIST_DIRECTORIES true RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*)
|
||||
# 遍历子目录
|
||||
foreach(SUBDIR ${SUBDIRECTORIES})
|
||||
# 检查子目录中是否存在 CMakeLists.txt
|
||||
if(EXISTS ${CMAKE_CURRENT_SOURCE_DIR}/${SUBDIR}/CMakeLists.txt)
|
||||
add_subdirectory(${SUBDIR})
|
||||
endif()
|
||||
endforeach()
|
||||
185
csrc/gmm/grouped_matmul_swiglu_quant/README.md
Normal file
185
csrc/gmm/grouped_matmul_swiglu_quant/README.md
Normal file
@@ -0,0 +1,185 @@
|
||||
# GroupedMatmulSwigluQuant
|
||||
|
||||
## 产品支持情况
|
||||
|
||||
| 产品 | 是否支持 |
|
||||
| :----------------------------------------------------------- | :------: |
|
||||
| <term>Atlas A3 训练系列产品/Atlas A3 推理系列产品</term> | √ |
|
||||
| <term>Atlas A2 训练系列产品/Atlas A2 推理系列产品</term> | √ |
|
||||
| <term>Kirin X90 处理器系列产品</term> | √ |
|
||||
| <term>Kirin 9030 处理器系列产品</term> | √ |
|
||||
|
||||
## 功能说明
|
||||
|
||||
- 接口功能：融合GroupedMatmul 、dquant、swiglu和quant，详细解释见计算公式。
|
||||
- 计算公式：
|
||||
|
||||
- **定义**：
|
||||
|
||||
- **⋅** 表示矩阵乘法。
|
||||
- **⊙** 表示逐元素乘法。
|
||||
- $\left \lfloor x\right \rceil$ 表示将x四舍五入到最近的整数。
|
||||
- $\mathbb{Z_8} = \{ x \in \mathbb{Z} | −128≤x≤127 \}$
|
||||
- $\mathbb{Z_{32}} = \{ x \in \mathbb{Z} | -2147483648≤x≤2147483647 \}$
|
||||
- **输入**：
|
||||
|
||||
- $X∈\mathbb{Z_8}^{M \times K}$：输入矩阵（左矩阵），M是总token 数，K是特征维度。
|
||||
- $W∈\mathbb{Z_8}^{E \times K \times N}$：分组权重矩阵（右矩阵），E是专家个数，K是特征维度，N是输出维度。
|
||||
- $bias∈\mathbb{Z_{32}}^{E \times N}$：矩阵乘计算的偏移值，E是专家个数，N是输出维度。
|
||||
- $offset∈\mathbb{R}^{E \times N}$：per-channel非对称反量化的偏移，E是专家个数，N是输出维度。
|
||||
- $w\_scale∈\mathbb{R}^{E \times N}$：分组权重矩阵（右矩阵）的逐通道缩放因子，E是专家个数，N是输出维度。
|
||||
- $x\_scale∈\mathbb{R}^{M}$：输入矩阵（左矩阵）的逐 token缩放因子，M是总token 数。
|
||||
- $groupList∈\mathbb{N}^{E}$：前缀和的分组索引列表。
|
||||
- **输出**：
|
||||
|
||||
- $Q∈\mathbb{Z_8}^{M \times N / 2}$：量化后的输出矩阵。
|
||||
- $Q\_scale∈\mathbb{R}^{M}$：量化缩放因子。
|
||||
- $Q\_offset∈\mathbb{R}^{M}$：量化偏移因子。
|
||||
- **计算过程**
|
||||
|
||||
- 1.根据groupList[i]确定当前分组的 token ，$i \in [0,Len(groupList)]$。
|
||||
|
||||
>例子：假设groupList=[3,4,4,6]，从0开始计数。
|
||||
>
|
||||
>第0个右矩阵`W[0,:,:]`，对应索引位置[0,3)的token`x[0:3]`（共3-0=3个token），对应`x_scale[0:3]`、`w_scale[0]`、`bias[0]`、`offset[0]`、`Q[0:3]`、`Q_scale[0:3]`、`Q_offset[0:3]`；
|
||||
>
|
||||
>第1个右矩阵`W[1,:,:]`，对应索引位置[3,4)的token`x[3:4]`（共4-3=1个token），对应`x_scale[3:4]`、`w_scale[1]`、`bias[1]`、`offset[1]`、`Q[3:4]`、`Q_scale[3:4]`、`Q_offset[3:4]`；
|
||||
>
|
||||
>第2个右矩阵`W[2,:,:]`，对应索引位置[4,4)的token`x[4:4]`（共4-4=0个token），对应`x_scale[4:4]`、`w_scale[2]`、`bias[2]`、`offset[2]`、`Q[4:4]`、`Q_scale[4:4]`、`Q_offset[4:4]`；
|
||||
>
|
||||
>第3个右矩阵`W[3,:,:]`，对应索引位置[4,6)的token`x[4:6]`（共6-4=2个token），对应`x_scale[4:6]`、`w_scale[3]`、`bias[3]`、`offset[3]`、`Q[4:6]`、`Q_scale[4:6]`、`Q_offset[4:6]`；
|
||||
>
|
||||
>请注意：groupList中未指定的部分将不会参与更新。
|
||||
>例如groupList=[12,14,18]，X的shape为[30，:]。
|
||||
>
|
||||
>则第一个输出Q的shape为[30，:]，其中Q[18:，：]的部分不会进行更新和初始化，其中数据为显存空间申请时的原数据。
|
||||
>
|
||||
>同理，第二个输出Q的shape为[30]，其中Q\_scale[18:]的部分不会进行更新或初始化，其中数据为显存空间申请时的原数据。
|
||||
>
|
||||
>即输出的Q[:groupList[-1],:]和Q\_scale[:groupList[-1]]为有效数据部分。
|
||||
|
||||
- 2.根据分组确定的入参进行如下计算：
|
||||
|
||||
$C_{i} = (X_{i}\cdot W_{i} )\odot x\_scale_{i\ BroadCast} \odot w\_scale_{i\ BroadCast}$
|
||||
|
||||
$C_{i,act}, gate_{i} = split(C_{i})$
|
||||
|
||||
$S_{i}=Swish(C_{i,act})\odot gate_{i}$ 其中$Swish(x)=\frac{x}{1+e^{-x}}$
|
||||
|
||||
>注：当前版本不支持$bias_{i}$、$offset_{i}$，未来版本将支持的计算公式如下：
|
||||
>$C_{i} =(X_{i}\cdot W_{i} + bias_{i\ BroadCast})\odot x\_scale_{i\ BroadCast} \odot w\_scale_{i\ BroadCast}+offset_{i\ BroadCast}$
|
||||
|
||||
- 3.确定量化方式
|
||||
|
||||
- 当量化方式为对称量化时：
|
||||
|
||||
$Q\_scale_{i} = \frac{max(|S_{i}|)}{127}$
|
||||
|
||||
$Q_{i} = \left \lfloor \frac{S_{i}}{Q\_scale_{i}}\right \rceil $
|
||||
|
||||
- 当量化方式为非对称量化时：(暂不支持)
|
||||
|
||||
$Q\_scale_{i} = \frac{max(S_{i})-min(S_{i})}{255}$
|
||||
|
||||
$Q\_offset_{i} = -128 - \left \lfloor \frac{min(S_{i})}{Q\_scale_{i}}\right \rceil$
|
||||
|
||||
$Q_{i} = \left \lfloor \frac{S_{i}}{ Q\_scale_{i} } + Q\_offset_{i}\right \rceil $
|
||||
|
||||
## 参数说明
|
||||
|
||||
<table style="table-layout: auto; width: 100%">
|
||||
<thead>
|
||||
<tr>
|
||||
<th style="white-space: nowrap">参数名</th>
|
||||
<th style="white-space: nowrap">输入/输出/属性</th>
|
||||
<th style="white-space: nowrap">描述</th>
|
||||
<th style="white-space: nowrap">数据类型</th>
|
||||
<th style="white-space: nowrap">数据格式</th>
|
||||
</tr>
|
||||
</thead>
|
||||
<tbody>
|
||||
<tr>
|
||||
<td style="white-space: nowrap">x</td>
|
||||
<td style="white-space: nowrap">输入</td>
|
||||
<td style="white-space: nowrap">左矩阵，公式中的X。</td>
|
||||
<td style="white-space: nowrap">INT8</td>
|
||||
<td style="white-space: nowrap">ND</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style="white-space: nowrap">weight</td>
|
||||
<td style="white-space: nowrap">输入</td>
|
||||
<td style="white-space: nowrap">权重矩阵，公式中的W。</td>
|
||||
<td style="white-space: nowrap">INT8</td>
|
||||
<td style="white-space: nowrap">ND / NZ</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style="white-space: nowrap">bias</td>
|
||||
<td style="white-space: nowrap">输入</td>
|
||||
<td style="white-space: nowrap">矩阵乘计算的偏移值，公式中的bias。</td>
|
||||
<td style="white-space: nowrap">INT32</td>
|
||||
<td style="white-space: nowrap">ND</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style="white-space: nowrap">offset</td>
|
||||
<td style="white-space: nowrap">输入</td>
|
||||
<td style="white-space: nowrap">per-channel非对称反量化的偏移，公式中的offset。</td>
|
||||
<td style="white-space: nowrap">FLOAT32</td>
|
||||
<td style="white-space: nowrap">ND</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style="white-space: nowrap">weightScale</td>
|
||||
<td style="white-space: nowrap">输入</td>
|
||||
<td style="white-space: nowrap">右矩阵的量化因子，公式中的w_scale。</td>
|
||||
<td style="white-space: nowrap">FLOAT、FLOAT16、BFLOAT16</td>
|
||||
<td style="white-space: nowrap">ND</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style="white-space: nowrap">xScale</td>
|
||||
<td style="white-space: nowrap">输入</td>
|
||||
<td style="white-space: nowrap">左矩阵的量化因子，公式中的x_scale。</td>
|
||||
<td style="white-space: nowrap">FLOAT32</td>
|
||||
<td style="white-space: nowrap">ND</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style="white-space: nowrap">groupList</td>
|
||||
<td style="white-space: nowrap">输入</td>
|
||||
<td style="white-space: nowrap">指示每个分组参与计算的Token个数，公式中的groupList。</td>
|
||||
<td style="white-space: nowrap">INT64</td>
|
||||
<td style="white-space: nowrap">ND</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style="white-space: nowrap">output</td>
|
||||
<td style="white-space: nowrap">输出</td>
|
||||
<td style="white-space: nowrap">输出的量化因子，公式中的Q。</td>
|
||||
<td style="white-space: nowrap">FLOAT</td>
|
||||
<td style="white-space: nowrap">ND</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style="white-space: nowrap">outputScale</td>
|
||||
<td style="white-space: nowrap">输出</td>
|
||||
<td style="white-space: nowrap">输出的量化因子，公式中的Q_scale。</td>
|
||||
<td style="white-space: nowrap">FLOAT</td>
|
||||
<td style="white-space: nowrap">ND </td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style="white-space: nowrap">outputOffset</td>
|
||||
<td style="white-space: nowrap">输出</td>
|
||||
<td style="white-space: nowrap">输出的非对称量化的偏移，公式中的Q_offset。</td>
|
||||
<td style="white-space: nowrap">FLOAT</td>
|
||||
<td style="white-space: nowrap">ND</td>
|
||||
</tr>
|
||||
</tbody>
|
||||
</table>
|
||||
|
||||
- Kirin X90/Kirin 9030 处理器系列产品: 不支持BFLOAT16。
|
||||
|
||||
## 约束说明
|
||||
|
||||
- N轴长度不能超过10240。
|
||||
- K轴长度不能超过65536。
|
||||
|
||||
## 调用说明
|
||||
|
||||
| 调用方式 | 调用样例 | 说明 |
|
||||
|--------------|-------------------------|--------------------------------------------------------------|
|
||||
| aclnn调用 | [test_aclnn_grouped_matmul_swiglu_quant](examples/test_aclnn_grouped_matmul_swiglu_quant.cpp) | 通过接口方式调用[GroupedMatmulSwigluQuant](docs/aclnnGroupedMatmulSwigluQuant.md)算子。 |
|
||||
28
csrc/gmm/grouped_matmul_swiglu_quant/op_host/CMakeLists.txt
Normal file
28
csrc/gmm/grouped_matmul_swiglu_quant/op_host/CMakeLists.txt
Normal file
@@ -0,0 +1,28 @@
|
||||
# -----------------------------------------------------------------------------------------------------------
|
||||
# Copyright (c) 2025 Huawei Technologies Co., Ltd.
|
||||
# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
# CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
# Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
# See LICENSE in the root of the software repository for the full text of the License.
|
||||
# -----------------------------------------------------------------------------------------------------------
|
||||
add_op_to_compiled_list()
|
||||
if (BUILD_OPEN_PROJECT)
|
||||
target_sources(op_host_aclnnExc PRIVATE
|
||||
grouped_matmul_swiglu_quant_def.cpp
|
||||
)
|
||||
add_ops_compile_options(
|
||||
OP_NAME GroupedMatmulSwigluQuant
|
||||
OPTIONS --cce-auto-sync=on
|
||||
-Wno-deprecated-declarations
|
||||
)
|
||||
endif()
|
||||
|
||||
if(NOT BUILD_OPS_RTY_KERNEL)
|
||||
add_modules_sources(OPTYPE grouped_matmul_swiglu_quant ACLNNTYPE aclnn_exclude)
|
||||
target_include_directories(${OPHOST_NAME}_tiling_obj PRIVATE
|
||||
${CMAKE_CURRENT_SOURCE_DIR}
|
||||
)
|
||||
endif()
|
||||
|
||||
@@ -0,0 +1,122 @@
|
||||
/**
|
||||
* Copyright (c) 2025 Huawei Technologies Co., Ltd.
|
||||
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
* CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
* See LICENSE in the root of the software repository for the full text of the License.
|
||||
*/
|
||||
|
||||
/*!
|
||||
* \file grouped_matmul_swiglu_quant_def.cpp
|
||||
* \brief
|
||||
*/
|
||||
|
||||
#include "register/op_def_registry.h"
|
||||
namespace ops {
|
||||
class GroupedMatmulSwigluQuant : public OpDef {
|
||||
public:
|
||||
explicit GroupedMatmulSwigluQuant(const char *name) : OpDef(name)
|
||||
{
|
||||
this->Input("x")
|
||||
.ParamType(REQUIRED)
|
||||
.DataType({ge::DT_INT8, ge::DT_INT8, ge::DT_INT8, ge::DT_INT8, ge::DT_INT8})
|
||||
.Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
|
||||
this->Input("weight")
|
||||
.ParamType(REQUIRED)
|
||||
.DataType({ge::DT_INT8, ge::DT_INT8, ge::DT_INT8, ge::DT_INT4, ge::DT_INT4})
|
||||
.Format({ge::FORMAT_FRACTAL_NZ, ge::FORMAT_FRACTAL_NZ, ge::FORMAT_FRACTAL_NZ, ge::FORMAT_ND,
|
||||
ge::FORMAT_FRACTAL_NZ});
|
||||
this->Input("weight_scale")
|
||||
.ParamType(REQUIRED)
|
||||
.DataType({ge::DT_FLOAT, ge::DT_BF16, ge::DT_FLOAT16, ge::DT_UINT64, ge::DT_UINT64})
|
||||
.Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
|
||||
this->Input("x_scale")
|
||||
.ParamType(REQUIRED)
|
||||
.DataType({ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT})
|
||||
.Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
|
||||
this->Input("weight_assistance_matrix")
|
||||
.ParamType(OPTIONAL)
|
||||
.DataType({ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT})
|
||||
.Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
|
||||
this->Input("group_list")
|
||||
.ParamType(REQUIRED)
|
||||
.DataType({ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64})
|
||||
.Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
|
||||
this->Output("y")
|
||||
.ParamType(REQUIRED)
|
||||
.DataType({ge::DT_INT8, ge::DT_INT8, ge::DT_INT8, ge::DT_INT8, ge::DT_INT8})
|
||||
.Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
|
||||
this->Output("y_scale")
|
||||
.ParamType(REQUIRED)
|
||||
.DataType({ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT})
|
||||
.Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
|
||||
this->Attr("is_enable_weight_assistance_matrix").AttrType(OPTIONAL).Bool(true);
|
||||
this->Attr("dequant_mode").AttrType(OPTIONAL).Int(0);
|
||||
this->Attr("limited").AttrType(OPTIONAL).Float(0.0f);
|
||||
|
||||
OpAICoreConfig aicore_config;
|
||||
aicore_config.DynamicCompileStaticFlag(true)
|
||||
.DynamicFormatFlag(true)
|
||||
.DynamicRankSupportFlag(true)
|
||||
.DynamicShapeSupportFlag(true)
|
||||
.NeedCheckSupportFlag(false)
|
||||
.PrecisionReduceFlag(true);
|
||||
|
||||
this->AICore().AddConfig("ascend910b", aicore_config);
|
||||
this->AICore().AddConfig("ascend910_93", aicore_config);
|
||||
|
||||
OpAICoreConfig config_kirin = GetKirinCoreConfig();
|
||||
this->AICore().AddConfig("kirinx90", config_kirin);
|
||||
this->AICore().AddConfig("kirin9030", config_kirin);
|
||||
}
|
||||
|
||||
private:
|
||||
OpAICoreConfig GetKirinCoreConfig() const
|
||||
{
|
||||
OpAICoreConfig config_kirin;
|
||||
config_kirin.DynamicCompileStaticFlag(true)
|
||||
.DynamicFormatFlag(true)
|
||||
.DynamicRankSupportFlag(true)
|
||||
.DynamicShapeSupportFlag(true)
|
||||
.NeedCheckSupportFlag(false)
|
||||
.PrecisionReduceFlag(true);
|
||||
config_kirin.Input("x")
|
||||
.ParamType(REQUIRED)
|
||||
.DataType({ge::DT_INT8, ge::DT_INT8})
|
||||
.Format({ge::FORMAT_ND, ge::FORMAT_ND});
|
||||
config_kirin.Input("weight")
|
||||
.ParamType(REQUIRED)
|
||||
.DataType({ge::DT_INT8, ge::DT_INT8})
|
||||
.Format({ge::FORMAT_FRACTAL_NZ, ge::FORMAT_FRACTAL_NZ});
|
||||
config_kirin.Input("weight_scale")
|
||||
.ParamType(REQUIRED)
|
||||
.DataType({ge::DT_FLOAT, ge::DT_FLOAT16})
|
||||
.Format({ge::FORMAT_ND, ge::FORMAT_ND});
|
||||
config_kirin.Input("x_scale")
|
||||
.ParamType(REQUIRED)
|
||||
.DataType({ge::DT_FLOAT, ge::DT_FLOAT})
|
||||
.Format({ge::FORMAT_ND, ge::FORMAT_ND});
|
||||
config_kirin.Input("weight_assistance_matrix")
|
||||
.ParamType(OPTIONAL)
|
||||
.DataType({ge::DT_FLOAT, ge::DT_FLOAT})
|
||||
.Format({ge::FORMAT_ND, ge::FORMAT_ND});
|
||||
config_kirin.Input("group_list")
|
||||
.ParamType(REQUIRED)
|
||||
.DataType({ge::DT_INT64, ge::DT_INT64})
|
||||
.Format({ge::FORMAT_ND, ge::FORMAT_ND});
|
||||
config_kirin.Output("y")
|
||||
.ParamType(REQUIRED)
|
||||
.DataType({ge::DT_INT8, ge::DT_INT8})
|
||||
.Format({ge::FORMAT_ND, ge::FORMAT_ND});
|
||||
config_kirin.Output("y_scale")
|
||||
.ParamType(REQUIRED)
|
||||
.DataType({ge::DT_FLOAT, ge::DT_FLOAT})
|
||||
.Format({ge::FORMAT_ND, ge::FORMAT_ND});
|
||||
return config_kirin;
|
||||
}
|
||||
};
|
||||
|
||||
OP_ADD(GroupedMatmulSwigluQuant);
|
||||
} // namespace ops
|
||||
@@ -0,0 +1,53 @@
|
||||
/**
|
||||
* Copyright (c) 2025 Huawei Technologies Co., Ltd.
|
||||
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
* CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
* See LICENSE in the root of the software repository for the full text of the License.
|
||||
*/
|
||||
|
||||
/*!
|
||||
* \file grouped_matmul_swiglu_quant_proto.cpp
|
||||
* \brief
|
||||
*/
|
||||
#include "register/op_impl_registry.h"
|
||||
#include "log/log.h"
|
||||
#include "platform/platform_info.h"
|
||||
|
||||
using namespace ge;
|
||||
namespace ops {
|
||||
const int64_t X_INDEX = 0;
|
||||
const int64_t WEIGHTSCALE_INDEX = 2;
|
||||
const int64_t M_DIM_INDEX = 0;
|
||||
const int64_t DIM_LEN = 2;
|
||||
const int64_t SPLIT_RATIO = 2;
|
||||
static ge::graphStatus InferShape4GroupedMatmulSwigluQuant(gert::InferShapeContext *context)
|
||||
{
|
||||
const gert::Shape *xShape = context->GetInputShape(X_INDEX);
|
||||
const gert::Shape *weightScaleShape = context->GetInputShape(WEIGHTSCALE_INDEX);
|
||||
int64_t m = xShape->GetDim(M_DIM_INDEX);
|
||||
int64_t N_DIM_INDEX = weightScaleShape->GetDimNum() - 1;
|
||||
int64_t n = static_cast<int64_t>(weightScaleShape->GetDim(N_DIM_INDEX) / SPLIT_RATIO);
|
||||
auto outShape = context->GetOutputShape(0);
|
||||
outShape->SetDimNum(DIM_LEN);
|
||||
outShape->SetDim(0, m);
|
||||
outShape->SetDim(1, n);
|
||||
auto outScaleShape = context->GetOutputShape(1);
|
||||
outScaleShape->SetDimNum(1);
|
||||
outScaleShape->SetDim(0, m);
|
||||
return GRAPH_SUCCESS;
|
||||
}
|
||||
|
||||
static graphStatus InferDataType4GroupedMatmulSwigluQuant(gert::InferDataTypeContext *context)
|
||||
{
|
||||
context->SetOutputDataType(0, DataType::DT_INT8);
|
||||
context->SetOutputDataType(1, DataType::DT_FLOAT);
|
||||
return GRAPH_SUCCESS;
|
||||
}
|
||||
|
||||
IMPL_OP_INFERSHAPE(GroupedMatmulSwigluQuant)
|
||||
.InferShape(InferShape4GroupedMatmulSwigluQuant)
|
||||
.InferDataType(InferDataType4GroupedMatmulSwigluQuant);
|
||||
} // namespace ops
|
||||
@@ -0,0 +1,288 @@
|
||||
/**
|
||||
* Copyright (c) 2025 Huawei Technologies Co., Ltd.
|
||||
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
* CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
* See LICENSE in the root of the software repository for the full text of the License.
|
||||
*/
|
||||
|
||||
/*!
|
||||
* \file grouped_matmul_swiglu_quant_tiling.cpp
|
||||
* \brief
|
||||
*/
|
||||
#include <climits>
|
||||
#include <graph/utils/type_utils.h>
|
||||
#include "register/op_impl_registry.h"
|
||||
#include "log/log.h"
|
||||
#include "err/ops_err.h"
|
||||
#include "tiling_base/tiling_base.h"
|
||||
#include "grouped_matmul_swiglu_quant_tiling.h"
|
||||
using namespace ge;
|
||||
using namespace AscendC;
|
||||
using namespace GroupedMatmulSwigluQuantTiling;
|
||||
using namespace Ops::Transformer::OpTiling;
|
||||
namespace {
|
||||
template <typename T>
|
||||
static inline auto AlignUp(T a, T base) -> T
|
||||
{
|
||||
if (base == 0) {
|
||||
return 0;
|
||||
}
|
||||
return (a + base - 1) / base * base;
|
||||
}
|
||||
} // namespace
|
||||
|
||||
namespace optiling {
|
||||
|
||||
struct GMMSwigluCompileInfo {
|
||||
uint64_t ubSize_ = 0;
|
||||
uint32_t aicNum_ = 0;
|
||||
uint32_t baseM_ = 128;
|
||||
uint32_t baseN_ = 256;
|
||||
};
|
||||
|
||||
static int64_t CalMaxRowInUb_A8W4(const gert::TilingContext *context, const uint64_t ubSize, const uint64_t n)
|
||||
{
|
||||
const uint64_t ALIGNMENT = 8;
|
||||
const float WEIGHT_FACTOR = 8.5;
|
||||
const uint64_t ALIGNMENT_TERM_FACTOR = 4;
|
||||
const uint64_t LINEAR_TERM_FACTOR = 6;
|
||||
const uint64_t CONSTANT_TERM = 64;
|
||||
const uint64_t MIN_ROW_THRESHOLD = 1;
|
||||
|
||||
// 表达式：8.5 * row * n + 4 * alignUp(row, 8) + 6n + 64 <= ubSize
|
||||
|
||||
// 忽略对齐项的初始估计
|
||||
int64_t maxRowEstimate =
|
||||
(ubSize - CONSTANT_TERM - LINEAR_TERM_FACTOR * n) / static_cast<int64_t>(WEIGHT_FACTOR * n);
|
||||
|
||||
// 考虑对齐影响
|
||||
uint64_t alignedRow = (maxRowEstimate + ALIGNMENT - 1) / ALIGNMENT * ALIGNMENT;
|
||||
uint64_t totalSize = static_cast<uint64_t>(WEIGHT_FACTOR * maxRowEstimate * n) +
|
||||
ALIGNMENT_TERM_FACTOR * alignedRow + LINEAR_TERM_FACTOR * n + CONSTANT_TERM;
|
||||
|
||||
// 如果超过UB大小，逐步减少row直到满足条件
|
||||
while (totalSize > ubSize && maxRowEstimate > 0) {
|
||||
maxRowEstimate--;
|
||||
alignedRow = (maxRowEstimate + ALIGNMENT - 1) / ALIGNMENT * ALIGNMENT;
|
||||
totalSize = static_cast<uint64_t>(WEIGHT_FACTOR * maxRowEstimate * n) + ALIGNMENT_TERM_FACTOR * alignedRow +
|
||||
LINEAR_TERM_FACTOR * n + CONSTANT_TERM;
|
||||
}
|
||||
|
||||
if (maxRowEstimate < MIN_ROW_THRESHOLD) {
|
||||
OP_LOGE(context->GetNodeName(), "GMM_SWIGLU_QUANT TILING: No valid row found for n = %lu, ubSize = %lu\n", n,
|
||||
ubSize);
|
||||
return 0;
|
||||
}
|
||||
return maxRowEstimate;
|
||||
}
|
||||
|
||||
static int64_t CalMaxRowInUb(const gert::TilingContext *context, const uint64_t ubSize, const uint64_t n)
|
||||
{
|
||||
uint64_t tmpBufSize = (n / SWIGLU_REDUCE_FACTOR) * FP32_DTYPE_SIZE;
|
||||
uint64_t perchannleBufSize = n * FP32_DTYPE_SIZE * DOUBLE_BUFFER;
|
||||
uint64_t reduceMaxResBufSize = BLOCK_BYTE;
|
||||
uint64_t reduceMaxTmpBufSize = BLOCK_BYTE;
|
||||
const uint64_t CONSTANT_TERM = 64;
|
||||
int64_t remainUbSize = ubSize - tmpBufSize - perchannleBufSize - reduceMaxResBufSize - reduceMaxTmpBufSize;
|
||||
int64_t maxRowInUb =
|
||||
remainUbSize / (n * INT32_DTYPE_SIZE + n / SWIGLU_REDUCE_FACTOR + FP32_DTYPE_SIZE) / DOUBLE_BUFFER;
|
||||
int64_t curUb = DOUBLE_BUFFER * (maxRowInUb * (INT32_DTYPE_SIZE * n + n / SWIGLU_REDUCE_FACTOR) +
|
||||
AlignUp(maxRowInUb, FP32_BLOCK_SIZE) * FP32_DTYPE_SIZE);
|
||||
if (curUb > remainUbSize) {
|
||||
// 64 : make sure ub does not excceed maxUbSize after align up to 8
|
||||
maxRowInUb = (remainUbSize - CONSTANT_TERM) /
|
||||
(n * INT32_DTYPE_SIZE + n / SWIGLU_REDUCE_FACTOR + FP32_DTYPE_SIZE) / DOUBLE_BUFFER;
|
||||
}
|
||||
if (maxRowInUb < 1) {
|
||||
// when n > (ubSize - 72) / 19 = 10330, maxRowInUb < 1
|
||||
OP_LOGE(context->GetNodeName(), "GMM_SWIGLU_QUANT TILING: n should not be greater than 10240, now is %lu\n", n);
|
||||
}
|
||||
return maxRowInUb;
|
||||
}
|
||||
|
||||
static void SetTilingKey(gert::TilingContext *context, bool isSplitWorkSpace, bool isA8W4MSD)
|
||||
{
|
||||
if (isA8W4MSD) { // A8W4 MSD tiling_key使用4
|
||||
context->SetTilingKey(A8W4_MSD_TILING_KEY_MODE);
|
||||
context->SetScheduleMode(BATCH_MODE_SCHEDULE);
|
||||
} else if (isSplitWorkSpace) {
|
||||
context->SetTilingKey(SPLITWORKSPACE_TILING_KEY_MODE);
|
||||
context->SetScheduleMode(BATCH_MODE_SCHEDULE);
|
||||
} else {
|
||||
context->SetTilingKey(COMMON_TILING_KEY_MODE);
|
||||
context->SetScheduleMode(BATCH_MODE_SCHEDULE);
|
||||
}
|
||||
}
|
||||
|
||||
ASCENDC_EXTERN_C graphStatus TilingGMMSwigluQuant(gert::TilingContext *context)
|
||||
{
|
||||
// set info
|
||||
OP_LOGD(context->GetNodeName(), "Begin Run GMM Swiglu Tiling .");
|
||||
auto xDesc = context->GetInputDesc(X_INDEX);
|
||||
OP_CHECK_NULL_WITH_CONTEXT(context, xDesc);
|
||||
auto weightDesc = context->GetInputDesc(WEIGHT_INDEX);
|
||||
OP_CHECK_NULL_WITH_CONTEXT(context, weightDesc);
|
||||
ge::DataType xDType = xDesc->GetDataType();
|
||||
ge::DataType weightDType = weightDesc->GetDataType();
|
||||
|
||||
bool isA8W4MSD = (xDType == ge::DataType::DT_INT8 && weightDType == ge::DataType::DT_INT4);
|
||||
auto compileInfoPtr = context->GetCompileInfo<GMMSwigluCompileInfo>();
|
||||
auto xTensor = context->GetInputTensor(X_INDEX);
|
||||
OP_CHECK_NULL_WITH_CONTEXT(context, xTensor);
|
||||
const int64_t m = xTensor->GetStorageShape().GetDim(0);
|
||||
const int64_t k = xTensor->GetStorageShape().GetDim(1);
|
||||
auto wTensor = context->GetInputTensor(WEIGHT_INDEX);
|
||||
OP_CHECK_NULL_WITH_CONTEXT(context, wTensor);
|
||||
// swiglu limit 0 means clamp is disabled.
|
||||
auto attrs = context->GetAttrs();
|
||||
float limited = 0.0f;
|
||||
if (attrs != nullptr) {
|
||||
if (const double *limitedPtr = attrs->GetAttrPointer<double>(ATTR_INDEX_LIMITED)) {
|
||||
limited = static_cast<float>(*limitedPtr);
|
||||
}
|
||||
}
|
||||
OP_CHECK_IF(!(limited >= 0.0f),
|
||||
OPS_REPORT_VECTOR_INNER_ERR(context->GetNodeName(), "limited should be non-negative"),
|
||||
return GRAPH_FAILED);
|
||||
int64_t n = 0;
|
||||
if (wTensor->GetStorageShape().GetDimNum() == ND_WEIGHT_DIM_LIMIT) { // ND
|
||||
n = wTensor->GetStorageShape().GetDim(DIM_2);
|
||||
} else if (wTensor->GetStorageShape().GetDimNum() == NZ_WEIGHT_DIM_LIMIT) { // NZ
|
||||
n = wTensor->GetStorageShape().GetDim(DIM_1) * wTensor->GetStorageShape().GetDim(DIM_4);
|
||||
}
|
||||
auto wScaleTensor = context->GetInputTensor(WEIGHT_SCALE_INDEX);
|
||||
OP_CHECK_NULL_WITH_CONTEXT(context, wScaleTensor);
|
||||
int64_t quantGroupNum = 0;
|
||||
if (wScaleTensor->GetStorageShape().GetDimNum() == PERCHANNEL_WSCALE_DIM_LIMIT) { // perChannel
|
||||
quantGroupNum = 1;
|
||||
} else if (wScaleTensor->GetStorageShape().GetDimNum() == PERGROUP_WSCALE_DIM_LIMIT) { // perGroup
|
||||
quantGroupNum = wScaleTensor->GetStorageShape().GetDim(1);
|
||||
}
|
||||
auto groupListTensor = context->GetDynamicInputTensor(GROUPLIST_INDEX, 0);
|
||||
OP_CHECK_NULL_WITH_CONTEXT(context, groupListTensor);
|
||||
const int64_t groupNum = groupListTensor->GetStorageShape().GetDim(0);
|
||||
GMMSwigluQuantTilingData tilingData;
|
||||
int64_t row = 0;
|
||||
if (isA8W4MSD) {
|
||||
row = CalMaxRowInUb_A8W4(context, compileInfoPtr->ubSize_, n);
|
||||
} else {
|
||||
row = CalMaxRowInUb(context, compileInfoPtr->ubSize_, n);
|
||||
}
|
||||
|
||||
tilingData.gmmSwigluBaseParams.set_groupNum(groupNum);
|
||||
tilingData.gmmSwigluBaseParams.set_coreNum(compileInfoPtr->aicNum_);
|
||||
tilingData.gmmSwigluBaseParams.set_K(k);
|
||||
tilingData.gmmSwigluBaseParams.set_N(n);
|
||||
tilingData.gmmSwigluBaseParams.set_M(m);
|
||||
tilingData.gmmSwigluBaseParams.set_baseM(A8W4_BASEM);
|
||||
tilingData.gmmSwigluBaseParams.set_baseN(A8W4_BASEN);
|
||||
tilingData.gmmSwigluBaseParams.set_limited(limited);
|
||||
tilingData.gmmSwiglu.set_maxProcessRowNum(row);
|
||||
tilingData.gmmSwiglu.set_groupListLen(groupNum);
|
||||
tilingData.gmmSwiglu.set_tokenLen(n);
|
||||
|
||||
tilingData.gmmSwigluBaseParams.set_quantGroupNum(quantGroupNum);
|
||||
|
||||
auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
|
||||
using namespace matmul_tiling;
|
||||
|
||||
MatmulApiTiling tiling(ascendcPlatform);
|
||||
tiling.SetAType(TPosition::GM, CubeFormat::ND, matmul_tiling::DataType::DT_INT4);
|
||||
tiling.SetBType(TPosition::GM, CubeFormat::NZ, matmul_tiling::DataType::DT_INT4);
|
||||
tiling.SetCType(TPosition::GM, CubeFormat::ND, matmul_tiling::DataType::DT_FLOAT16);
|
||||
tiling.SetBias(false);
|
||||
tiling.SetShape(A8W4_BASEM, A8W4_BASEN, k);
|
||||
tiling.SetFixSplit(A8W4_BASEM, A8W4_BASEN, A8W4_BASEK);
|
||||
tiling.SetOrgShape(m, n, k);
|
||||
tiling.SetBufferSpace(-1, -1, -1);
|
||||
OP_CHECK_IF(
|
||||
tiling.GetTiling(tilingData.mmTilingData) == -1,
|
||||
OPS_REPORT_VECTOR_INNER_ERR(context->GetNodeName(), "grouped_matmul_swiglu_quant_tiling, get tiling failed"),
|
||||
return GRAPH_FAILED);
|
||||
if (isA8W4MSD) {
|
||||
tilingData.mmTilingData.set_baseM(A8W4_BASEM);
|
||||
tilingData.mmTilingData.set_baseN(A8W4_BASEN);
|
||||
tilingData.mmTilingData.set_baseK(A8W4_BASEK);
|
||||
tilingData.mmTilingData.set_dbL0B(DOUBLE_BUFFER);
|
||||
tilingData.mmTilingData.set_stepKa(NUM_FOUR);
|
||||
tilingData.mmTilingData.set_stepKb(NUM_FOUR);
|
||||
tilingData.mmTilingData.set_depthA1(NUM_EIGHT);
|
||||
tilingData.mmTilingData.set_depthB1(NUM_EIGHT);
|
||||
tilingData.mmTilingData.set_stepM(1);
|
||||
tilingData.mmTilingData.set_stepN(1);
|
||||
}
|
||||
auto workspaceSizes = context->GetWorkspaceSizes(1);
|
||||
int64_t usrWorkspaceLimit = USER_WORKSPACE_LIMIT;
|
||||
int64_t mLimit = 0;
|
||||
if (isA8W4MSD) {
|
||||
mLimit = ((usrWorkspaceLimit / DOUBLE_WORKSPACE_SPLIT) / (k * sizeof(int8_t) + DOUBLE_ROW * n * sizeof(half)));
|
||||
} else {
|
||||
mLimit = ((usrWorkspaceLimit / DOUBLE_WORKSPACE_SPLIT) / INT32_DTYPE_SIZE) / n;
|
||||
}
|
||||
OP_CHECK_IF(mLimit <= 0,
|
||||
OPS_REPORT_VECTOR_INNER_ERR(context->GetNodeName(), "mLimit is %ld must over then 0.", mLimit),
|
||||
return GRAPH_FAILED);
|
||||
tilingData.gmmSwigluBaseParams.set_mLimit(mLimit);
|
||||
if (isA8W4MSD) {
|
||||
int workSpaceMTemp = mLimit * DOUBLE_WORKSPACE_SPLIT;
|
||||
tilingData.gmmSwigluBaseParams.set_workSpaceOffset1(workSpaceMTemp * k * sizeof(int8_t));
|
||||
tilingData.gmmSwigluBaseParams.set_workSpaceOffset2(2 * workSpaceMTemp * n * sizeof(half));
|
||||
workspaceSizes[0] =
|
||||
SYS_WORKSPACE_SIZE + // 系统预留16MB
|
||||
(workSpaceMTemp * k * sizeof(int8_t)) + // 第一阶段 预处理左矩阵 (mLimit, K) * int8 * 2(double WorkSpace)
|
||||
(DOUBLE_ROW * workSpaceMTemp * n *
|
||||
sizeof(half)); // 第二阶段 矩阵乘结果 (2 * mLimit, N) * fp16 * 2(double WorkSpace)
|
||||
} else {
|
||||
int workSpaceMTemp = (mLimit * DOUBLE_WORKSPACE_SPLIT > m ? m : mLimit * DOUBLE_WORKSPACE_SPLIT);
|
||||
tilingData.gmmSwigluBaseParams.set_workSpaceOffset1(0);
|
||||
tilingData.gmmSwigluBaseParams.set_workSpaceOffset2(0);
|
||||
workspaceSizes[0] = SYS_WORKSPACE_SIZE + (workSpaceMTemp * n * sizeof(int32_t));
|
||||
}
|
||||
bool isSplitWorkSpace = m > mLimit * DOUBLE_WORKSPACE_SPLIT;
|
||||
OP_LOGD(context->GetNodeName(), "grouped_matmul_swiglu_quant_tiling.");
|
||||
OP_LOGD(context->GetNodeName(), "gmmSwigluBaseParams.groupNum: %ld", groupNum);
|
||||
OP_LOGD(context->GetNodeName(), "gmmSwigluBaseParams.coreNum: %u ", compileInfoPtr->aicNum_);
|
||||
OP_LOGD(context->GetNodeName(), "gmmSwigluBaseParams.M: %ld", m);
|
||||
OP_LOGD(context->GetNodeName(), "gmmSwigluBaseParams.K: %ld", k);
|
||||
OP_LOGD(context->GetNodeName(), "gmmSwigluBaseParams.N: %ld", n);
|
||||
OP_LOGD(context->GetNodeName(), "gmmSwigluBaseParams.baseM: %ld", A8W4_BASEM);
|
||||
OP_LOGD(context->GetNodeName(), "gmmSwigluBaseParams.baseN: %ld", A8W4_BASEN);
|
||||
OP_LOGD(context->GetNodeName(), "gmmSwigluBaseParams.mLimit: %ld", mLimit);
|
||||
OP_LOGD(context->GetNodeName(), "gmmSwigluBaseParams.quantGroupNum: %ld", quantGroupNum);
|
||||
OP_LOGD(context->GetNodeName(), "gmmSwiglu.maxProcessRowNum: %ld", row);
|
||||
OP_LOGD(context->GetNodeName(), "gmmSwiglu.groupListLen: %ld", groupNum);
|
||||
OP_LOGD(context->GetNodeName(), "gmmSwiglu.tokenLen: %ld", n);
|
||||
OP_LOGD(context->GetNodeName(), "USER_WORKSPACE_LIMIT: %ld", usrWorkspaceLimit);
|
||||
OP_LOGD(context->GetNodeName(), "workspaceSizes: %lu", workspaceSizes[0]);
|
||||
OP_LOGD(context->GetNodeName(), "isSplitWorkSpace: %s", isSplitWorkSpace ? "true" : "false");
|
||||
OP_LOGD(context->GetNodeName(), "GMMSWIGLUQUANT_TILING: baseM is %u, baseK is %u, baseN is %u.", A8W4_BASEM, A8W4_BASEK, A8W4_BASEN);
|
||||
SetTilingKey(context, isSplitWorkSpace, isA8W4MSD);
|
||||
tilingData.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
|
||||
context->SetBlockDim(compileInfoPtr->aicNum_); // block dim is the number of aicube
|
||||
context->GetRawTilingData()->SetDataSize(tilingData.GetDataSize());
|
||||
OP_LOGD(context->GetNodeName(), "End Run GMM Swiglu Tiling.");
|
||||
return GRAPH_SUCCESS;
|
||||
}
|
||||
|
||||
ASCENDC_EXTERN_C graphStatus TilingPrepareForGMMSwigluQuant(gert::TilingParseContext *context)
|
||||
{
|
||||
// get info
|
||||
fe::PlatFormInfos *platformInfoPtr = context->GetPlatformInfo();
|
||||
OP_CHECK_NULL_WITH_CONTEXT(context, platformInfoPtr);
|
||||
auto compileInfoPtr = context->GetCompiledInfo<GMMSwigluCompileInfo>();
|
||||
OP_CHECK_NULL_WITH_CONTEXT(context, compileInfoPtr);
|
||||
|
||||
auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfoPtr);
|
||||
compileInfoPtr->aicNum_ = ascendcPlatform.GetCoreNumAic();
|
||||
ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, compileInfoPtr->ubSize_);
|
||||
OP_LOGD(context->GetNodeName(), "ubSize is %lu, aicNum is %u.", compileInfoPtr->ubSize_, compileInfoPtr->aicNum_);
|
||||
return GRAPH_SUCCESS;
|
||||
}
|
||||
|
||||
IMPL_OP_OPTILING(GroupedMatmulSwigluQuant)
|
||||
.Tiling(TilingGMMSwigluQuant)
|
||||
.TilingParse<GMMSwigluCompileInfo>(TilingPrepareForGMMSwigluQuant);
|
||||
} // namespace optiling
|
||||
@@ -0,0 +1,95 @@
|
||||
/**
|
||||
* Copyright (c) 2025 Huawei Technologies Co., Ltd.
|
||||
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
* CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
* See LICENSE in the root of the software repository for the full text of the License.
|
||||
*/
|
||||
|
||||
/*!
|
||||
* \file grouped_matmul_swiglu_quant_tiling.h
|
||||
* \brief
|
||||
*/
|
||||
#ifndef AIR_CXX_RUNTIME_V2_OP_IMPL_GROUPED_MATMUL_SWIGLU_QUANT_H
|
||||
#define AIR_CXX_RUNTIME_V2_OP_IMPL_GROUPED_MATMUL_SWIGLU_QUANT_H
|
||||
|
||||
#include <set>
|
||||
#include "register/tilingdata_base.h"
|
||||
#include "tiling/tiling_api.h"
|
||||
|
||||
namespace optiling {
|
||||
// GMM 基本信息
|
||||
BEGIN_TILING_DATA_DEF(GMMSwigluBaseParams)
|
||||
TILING_DATA_FIELD_DEF(uint32_t, groupNum);
|
||||
TILING_DATA_FIELD_DEF(uint32_t, coreNum);
|
||||
TILING_DATA_FIELD_DEF(uint32_t, K);
|
||||
TILING_DATA_FIELD_DEF(uint32_t, N);
|
||||
TILING_DATA_FIELD_DEF(uint32_t, M);
|
||||
TILING_DATA_FIELD_DEF(uint32_t, baseM);
|
||||
TILING_DATA_FIELD_DEF(uint32_t, baseN);
|
||||
TILING_DATA_FIELD_DEF(uint32_t, mLimit);
|
||||
TILING_DATA_FIELD_DEF(uint32_t, workSpaceOffset1);
|
||||
TILING_DATA_FIELD_DEF(uint32_t, workSpaceOffset2);
|
||||
TILING_DATA_FIELD_DEF(uint32_t, quantGroupNum);
|
||||
TILING_DATA_FIELD_DEF(float, limited);
|
||||
END_TILING_DATA_DEF;
|
||||
REGISTER_TILING_DATA_CLASS(GMMSwigluBaseParamsOp, GMMSwigluBaseParams)
|
||||
|
||||
// SwigluQuant部分tiling 基本信息
|
||||
BEGIN_TILING_DATA_DEF(GMMSwiglu)
|
||||
TILING_DATA_FIELD_DEF(uint32_t, maxProcessRowNum);
|
||||
TILING_DATA_FIELD_DEF(uint32_t, groupListLen);
|
||||
TILING_DATA_FIELD_DEF(uint32_t, tokenLen);
|
||||
END_TILING_DATA_DEF;
|
||||
REGISTER_TILING_DATA_CLASS(GMMSwigluOp, GMMSwiglu)
|
||||
|
||||
// 结构体集合
|
||||
BEGIN_TILING_DATA_DEF(GMMSwigluQuantTilingData)
|
||||
TILING_DATA_FIELD_DEF_STRUCT(GMMSwigluBaseParams, gmmSwigluBaseParams);
|
||||
TILING_DATA_FIELD_DEF_STRUCT(GMMSwiglu, gmmSwiglu);
|
||||
TILING_DATA_FIELD_DEF_STRUCT(TCubeTiling, mmTilingData);
|
||||
END_TILING_DATA_DEF;
|
||||
|
||||
REGISTER_TILING_DATA_CLASS(GroupedMatmulSwigluQuant, GMMSwigluQuantTilingData)
|
||||
} // namespace optiling
|
||||
|
||||
namespace GroupedMatmulSwigluQuantTiling {
|
||||
constexpr uint32_t X_INDEX = 0;
|
||||
constexpr uint32_t WEIGHT_INDEX = 1;
|
||||
constexpr uint32_t WEIGHT_SCALE_INDEX = 2;
|
||||
constexpr uint32_t GROUPLIST_INDEX = 5;
|
||||
constexpr uint32_t BATCH_MODE_SCHEDULE = 1;
|
||||
constexpr uint32_t DIM_0 = 0;
|
||||
constexpr uint32_t DIM_1 = 1;
|
||||
constexpr uint32_t DIM_2 = 2;
|
||||
constexpr uint32_t DIM_3 = 3;
|
||||
constexpr uint32_t DIM_4 = 4;
|
||||
constexpr uint32_t NUM_FOUR = 4;
|
||||
constexpr uint32_t NUM_EIGHT = 8;
|
||||
constexpr uint32_t SYS_WORKSPACE_SIZE = 16 * 1024 * 1024;
|
||||
constexpr int64_t USER_WORKSPACE_LIMIT = 64 * 1024 * 1024;
|
||||
constexpr int64_t DOUBLE_WORKSPACE_SPLIT = 2;
|
||||
constexpr uint32_t ATTR_INDEX_LIMITED = 2;
|
||||
constexpr int64_t INT32_DTYPE_SIZE = 4;
|
||||
constexpr int64_t FP32_DTYPE_SIZE = 4;
|
||||
constexpr int64_t FP32_BLOCK_SIZE = 8;
|
||||
constexpr int64_t BLOCK_BYTE = 32;
|
||||
constexpr int64_t SWIGLU_REDUCE_FACTOR = 2;
|
||||
constexpr int64_t DOUBLE_BUFFER = 2;
|
||||
constexpr int64_t ND_WEIGHT_DIM_LIMIT = 3;
|
||||
constexpr int64_t NZ_WEIGHT_DIM_LIMIT = 5;
|
||||
constexpr int64_t DOUBLE_ROW = 2;
|
||||
constexpr int64_t PERCHANNEL_WSCALE_DIM_LIMIT = 2;
|
||||
constexpr int64_t PERGROUP_WSCALE_DIM_LIMIT = 3;
|
||||
constexpr int64_t A8W4_MSD_TILING_KEY_MODE = 2;
|
||||
constexpr int64_t SPLITWORKSPACE_TILING_KEY_MODE = 1;
|
||||
constexpr int64_t COMMON_TILING_KEY_MODE = 0;
|
||||
constexpr int64_t A8W4_TOKEN_THRESHOLD = 32;
|
||||
constexpr int64_t A8W4_BASEM = 128;
|
||||
constexpr int64_t A8W4_BASEK = 256;
|
||||
constexpr int64_t A8W4_BASEN = 256;
|
||||
} // namespace GroupedMatmulSwigluQuantTiling
|
||||
|
||||
#endif // AIR_CXX_RUNTIME_V2_OP_IMPL_GROUPED_MATMUL_SWIGLU_QUANT_H
|
||||
@@ -0,0 +1,540 @@
|
||||
/**
|
||||
* Copyright (c) 2025 Huawei Technologies Co., Ltd.
|
||||
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
* CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
* See LICENSE in the root of the software repository for the full text of the License.
|
||||
*/
|
||||
#include <dlfcn.h>
|
||||
#include <new>
|
||||
#include "aclnn_kernels/contiguous.h"
|
||||
#include "acl/acl.h"
|
||||
#include "aclnn/aclnn_base.h"
|
||||
#include "aclnn_kernels/common/op_error_check.h"
|
||||
#include "opdev/common_types.h"
|
||||
#include "opdev/data_type_utils.h"
|
||||
#include "opdev/format_utils.h"
|
||||
#include "opdev/op_dfx.h"
|
||||
#include "opdev/op_executor.h"
|
||||
#include "opdev/op_log.h"
|
||||
#include "opdev/platform.h"
|
||||
#include "opdev/shape_utils.h"
|
||||
#include "opdev/tensor_view_utils.h"
|
||||
#include "opdev/make_op_executor.h"
|
||||
#include "grouped_matmul_swiglu_quant.h"
|
||||
#include "aclnn_grouped_matmul_swiglu_quant_weight_nz.h"
|
||||
#include "aclnn_grouped_matmul_swiglu_quant.h"
|
||||
|
||||
using namespace op;
|
||||
|
||||
#ifdef __cplusplus
|
||||
extern "C" {
|
||||
#endif
|
||||
|
||||
static constexpr int64_t SPLIT = 2L;
|
||||
static constexpr int64_t K_LIMIT_A8W8 = 65536L;
|
||||
static constexpr int64_t K_LIMIT_A8W4 = 20000L;
|
||||
static constexpr int64_t N_LIMIT = 10240L;
|
||||
static constexpr int64_t NZ_DIM_4_INT8 = 32L;
|
||||
static constexpr int64_t NZ_DIM_4_INT4 = 64L;
|
||||
static constexpr int64_t NZ_DIM_3 = 16L;
|
||||
static constexpr int64_t OUTPUT_IDX_0 = 0L;
|
||||
static constexpr int64_t OUTPUT_IDX_1 = 1L;
|
||||
static constexpr int64_t DIM_IDX_0 = 0L;
|
||||
static constexpr int64_t DIM_IDX_1 = 1L;
|
||||
static constexpr int64_t DIM_IDX_2 = 2L;
|
||||
static constexpr int64_t DIM_IDX_3 = 4L;
|
||||
static constexpr size_t X_DIM_LIMIT = 2UL;
|
||||
static constexpr size_t WEIGHT_ND_DIM_LIMIT = 3UL;
|
||||
static constexpr size_t WEIGHT_NZ_DIM_LIMIT = 5UL;
|
||||
static constexpr size_t WEIGHT_SCALE_DIM_LIMIT = 2UL;
|
||||
static constexpr size_t WEIGHT_SCALE_PERGROUP_DIM_LIMIT = 3UL;
|
||||
static constexpr size_t WEIGHT_SCALE_PERCHANNEL_DIM_LIMIT = 2UL;
|
||||
static constexpr size_t TOKEN_SCALE_DIM_LIMIT = 1UL;
|
||||
static constexpr size_t BIAS_DIM_LIMIT = 2UL;
|
||||
static constexpr size_t GROUP_LIST_DIM_LIMIT = 1UL;
|
||||
static constexpr size_t QUANTOUT_DIM_LIMIT = 2UL;
|
||||
static constexpr size_t QUANTSCALEOUT_DIM_LIMIT = 1UL;
|
||||
static constexpr size_t INT4_PER_INT32 = 8UL;
|
||||
bool isEnableWeightAssistanceMatrix = false;
|
||||
int dequantMode = 0;
|
||||
|
||||
static const std::initializer_list<DataType> X_DTYPE_SUPPORT_LIST = {DataType::DT_INT8};
|
||||
static const std::initializer_list<DataType> WEIGHT_DTYPE_SUPPORT_LIST = {DataType::DT_INT8, DataType::DT_INT4};
|
||||
static const std::initializer_list<DataType> WEIGHT_SCALE_DTYPE_SUPPORT_LIST = {
|
||||
DataType::DT_FLOAT, DataType::DT_FLOAT16, DataType::DT_BF16};
|
||||
static const std::initializer_list<DataType> WEIGHT_SCALE_A8W4_DTYPE_SUPPORT_LIST = {DataType::DT_UINT64};
|
||||
static const std::initializer_list<DataType> X_SCALE_DTYPE_SUPPORT_LIST = {DataType::DT_FLOAT, DataType::DT_FLOAT16,
|
||||
DataType::DT_BF16};
|
||||
static const std::initializer_list<DataType> GROUP_LIST_DTYPE_SUPPORT_LIST = {DataType::DT_INT64};
|
||||
static const std::initializer_list<DataType> QUANTOUT_DTYPE_SUPPORT_LIST = {DataType::DT_INT8};
|
||||
static const std::initializer_list<DataType> QUANTSCALEOUT_DTYPE_SUPPORT_LIST = {DataType::DT_FLOAT};
|
||||
static const std::initializer_list<DataType> BIAS_DTYPE_SUPPORT_LIST = {DataType::DT_FLOAT};
|
||||
static bool CheckNotNull(const aclTensor *x, const aclTensor *weight, const aclTensor *bias, const aclTensor *offset,
|
||||
const aclTensor *weightScale, const aclTensor *xScale, const aclTensor *groupList,
|
||||
const aclTensor *output, const aclTensor *outputScale, const aclTensor *outputOffset)
|
||||
{
|
||||
OP_CHECK_NULL(x, return false);
|
||||
OP_CHECK_NULL(weight, return false);
|
||||
OP_CHECK_NULL(weightScale, return false);
|
||||
OP_CHECK_NULL(xScale, return false);
|
||||
OP_CHECK_NULL(groupList, return false);
|
||||
OP_CHECK_NULL(output, return false);
|
||||
OP_CHECK_NULL(outputScale, return false);
|
||||
if (x->GetDataType() == DataType::DT_INT8 && weight->GetDataType() == DataType::DT_INT8 && bias != nullptr) {
|
||||
OP_LOGW("aclnnGroupedMatmulSwiGluQuant, The current version does not support the scenario that bias is not 0. "
|
||||
"Features and accuracy are not guaranteed if inputting bias with values other than 0.");
|
||||
} else if (x->GetDataType() == DataType::DT_INT8 && weight->GetDataType() == DataType::DT_INT4 && bias == nullptr) {
|
||||
OP_LOGE(ACLNN_ERR_PARAM_INVALID,
|
||||
"aclnnGroupedMatmulSwiGluQuant, The current version does not support the scenario that without bias. "
|
||||
"When x is Int8 and weight is int4, bias serves as an auxiliary matrix to weight, and this parameter "
|
||||
"cannot be nullptr.");
|
||||
return false;
|
||||
}
|
||||
if (offset != nullptr) {
|
||||
OP_LOGW(
|
||||
"aclnnGroupedMatmulSwiGluQuant, The current version does not support the scenario where offset is not 0. "
|
||||
"Features and accuracy are not guaranteed if inputting bias with values other than 0s.");
|
||||
}
|
||||
if (outputOffset != nullptr) {
|
||||
OP_LOGW("aclnnGroupedMatmulSwiGluQuant, The current version does not support the scenario where outputOffset "
|
||||
"is not 0. Features and accuracy are not guaranteed if inputting bias with values other than 0s.");
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
static bool CheckInputOutDims_A8W8(const aclTensor *x, const aclTensor *weight, const aclTensor *weightScale,
|
||||
const aclTensor *xScale, const aclTensor *groupList, const aclTensor *output,
|
||||
const aclTensor *outputScale)
|
||||
{
|
||||
OP_CHECK_WRONG_DIMENSION(x, X_DIM_LIMIT, return false);
|
||||
op::Format weightViewFormat = weight->GetViewFormat();
|
||||
if (IsPrivateFormat(weightViewFormat)) {
|
||||
OP_CHECK_WRONG_DIMENSION(weight, WEIGHT_NZ_DIM_LIMIT, return false);
|
||||
} else {
|
||||
OP_CHECK_WRONG_DIMENSION(weight, WEIGHT_ND_DIM_LIMIT, return false);
|
||||
}
|
||||
OP_CHECK_WRONG_DIMENSION(weightScale, WEIGHT_SCALE_DIM_LIMIT, return false);
|
||||
OP_CHECK_WRONG_DIMENSION(xScale, TOKEN_SCALE_DIM_LIMIT, return false);
|
||||
OP_CHECK_WRONG_DIMENSION(groupList, GROUP_LIST_DIM_LIMIT, return false);
|
||||
OP_CHECK_WRONG_DIMENSION(output, QUANTOUT_DIM_LIMIT, return false);
|
||||
OP_CHECK_WRONG_DIMENSION(outputScale, QUANTSCALEOUT_DIM_LIMIT, return false);
|
||||
return true;
|
||||
}
|
||||
|
||||
static bool CheckInputOutDims_A8W4(const aclTensor *x, const aclTensor *weight, const aclTensor *bias,
|
||||
const aclTensor *weightScale, const aclTensor *xScale, const aclTensor *groupList,
|
||||
const aclTensor *output, const aclTensor *outputScale)
|
||||
{
|
||||
OP_CHECK_WRONG_DIMENSION(x, X_DIM_LIMIT, return false);
|
||||
op::Format weightViewFormat = weight->GetViewFormat();
|
||||
if (IsPrivateFormat(weightViewFormat)) {
|
||||
OP_CHECK_WRONG_DIMENSION(weight, WEIGHT_NZ_DIM_LIMIT, return false);
|
||||
} else {
|
||||
OP_CHECK_WRONG_DIMENSION(weight, WEIGHT_ND_DIM_LIMIT, return false);
|
||||
}
|
||||
// 支持pergroup、perchannel量化weightScale分别为2维和3维
|
||||
OP_CHECK_MAX_DIM(weightScale, WEIGHT_SCALE_PERGROUP_DIM_LIMIT, return false);
|
||||
OP_CHECK_MIN_DIM(weightScale, WEIGHT_SCALE_PERCHANNEL_DIM_LIMIT, return false);
|
||||
OP_CHECK_WRONG_DIMENSION(bias, BIAS_DIM_LIMIT, return false);
|
||||
OP_CHECK_WRONG_DIMENSION(xScale, TOKEN_SCALE_DIM_LIMIT, return false);
|
||||
OP_CHECK_WRONG_DIMENSION(groupList, GROUP_LIST_DIM_LIMIT, return false);
|
||||
OP_CHECK_WRONG_DIMENSION(output, QUANTOUT_DIM_LIMIT, return false);
|
||||
OP_CHECK_WRONG_DIMENSION(outputScale, QUANTSCALEOUT_DIM_LIMIT, return false);
|
||||
return true;
|
||||
}
|
||||
|
||||
static bool CheckInputOutShape_A8W8(const aclTensor *x, const aclTensor *weight, const aclTensor *weightScale,
|
||||
const aclTensor *xScale, const aclTensor *groupList, const aclTensor *output,
|
||||
const aclTensor *outputScale)
|
||||
{
|
||||
int64_t m = x->GetViewShape().GetDim(0);
|
||||
int64_t k = x->GetViewShape().GetDim(1);
|
||||
int64_t n = weightScale->GetViewShape().GetDim(1);
|
||||
int64_t e = weight->GetViewShape().GetDim(0);
|
||||
if (n % SPLIT != 0) {
|
||||
OP_LOGE(ACLNN_ERR_PARAM_INVALID, "aclnnGroupedMatmulSwiGluQuant, N is %ld , not an even number.", n);
|
||||
return false;
|
||||
}
|
||||
int64_t nAfterHalve = static_cast<int64_t>(n / SPLIT);
|
||||
// x的shape期望为[M, K]
|
||||
op::Shape xExpectShape = {m, k};
|
||||
// weight的NDshape期望为[E, K, N]
|
||||
op::Shape weightNDExpectShape = {e, k, n};
|
||||
// weight的NZshape期望为[E, N // 32, K // 16, 16, 32]
|
||||
op::Shape weightNZExpectShape = {e, static_cast<int64_t>(n / NZ_DIM_4_INT8), static_cast<int64_t>(k / NZ_DIM_3),
|
||||
NZ_DIM_3, NZ_DIM_4_INT8};
|
||||
// weightScale的shape期望为[E, N]
|
||||
op::Shape weightScaleExpectShape = {e, n};
|
||||
// xScale的shape期望为[E, N]
|
||||
op::Shape xScaleExpectShape = {m};
|
||||
// output的shape期望为[M, N]
|
||||
op::Shape outputExpectShape = {m, nAfterHalve};
|
||||
// outputScale的shape期望为[M]
|
||||
op::Shape outputScaleExpectShape = {m};
|
||||
op::Format weightViewFormat = weight->GetViewFormat();
|
||||
OP_CHECK_SHAPE_NOT_EQUAL_WITH_EXPECTED_SIZE(x, xExpectShape, return false);
|
||||
if (IsPrivateFormat(weightViewFormat)) {
|
||||
OP_CHECK_SHAPE_NOT_EQUAL_WITH_EXPECTED_SIZE(weight, weightNZExpectShape, return false);
|
||||
} else {
|
||||
OP_CHECK_SHAPE_NOT_EQUAL_WITH_EXPECTED_SIZE(weight, weightNDExpectShape, return false);
|
||||
}
|
||||
OP_CHECK_SHAPE_NOT_EQUAL_WITH_EXPECTED_SIZE(weightScale, weightScaleExpectShape, return false);
|
||||
OP_CHECK_SHAPE_NOT_EQUAL_WITH_EXPECTED_SIZE(xScale, xScaleExpectShape, return false);
|
||||
|
||||
OP_CHECK_SHAPE_NOT_EQUAL_WITH_EXPECTED_SIZE(output, outputExpectShape, return false);
|
||||
OP_CHECK_SHAPE_NOT_EQUAL_WITH_EXPECTED_SIZE(outputScale, outputScaleExpectShape, return false);
|
||||
// groupList的长度应小于等于weight的专家数
|
||||
int64_t groupListLen = groupList->GetViewShape().GetDim(0);
|
||||
if (groupListLen > e) {
|
||||
OP_LOGE(ACLNN_ERR_PARAM_INVALID,
|
||||
"aclnnGroupedMatmulSwiGluQuant A8W8, Length of 'groupList' out of range"
|
||||
" (expected to be in range of [1, %ld], but got %ld)",
|
||||
e, groupListLen);
|
||||
return false;
|
||||
}
|
||||
if (n > N_LIMIT) {
|
||||
OP_LOGE(ACLNN_ERR_PARAM_INVALID,
|
||||
"aclnnGroupedMatmulSwiGluQuant A8W8: The current version does not support the scenario that "
|
||||
"N(%ld) is greater than %ld.",
|
||||
n, N_LIMIT);
|
||||
return false;
|
||||
}
|
||||
if (k >= K_LIMIT_A8W8) {
|
||||
OP_LOGE(ACLNN_ERR_PARAM_INVALID,
|
||||
"aclnnGroupedMatmulSwiGluQuant A8W8, The current version does not support the scenario."
|
||||
"The tail axis dimension of input0(x) is %ld, which need lower than %ld.",
|
||||
k, K_LIMIT_A8W8);
|
||||
return false;
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
static bool CheckInputOutShape_A8W4(const aclTensor *x, const aclTensor *weight, const aclTensor *bias,
|
||||
const aclTensor *weightScale, const aclTensor *xScale, const aclTensor *groupList,
|
||||
const aclTensor *output, const aclTensor *outputScale)
|
||||
{
|
||||
int64_t e = weight->GetViewShape().GetDim(0);
|
||||
int64_t m = x->GetViewShape().GetDim(0);
|
||||
int64_t k = x->GetViewShape().GetDim(1);
|
||||
int64_t n = 1;
|
||||
int64_t KGroupCount = 1; // K轴的组数，perchannel场景相当于pergroup场景中的组数为1
|
||||
int64_t KGroupSize = k; // K轴每组的元素个数
|
||||
op::Shape weightScaleExpectShape;
|
||||
// 通过weightScale的维度判断是否为perchannel 或 pergroup量化模式
|
||||
if (weightScale->GetViewShape().GetDimNum() == WEIGHT_SCALE_PERCHANNEL_DIM_LIMIT) {
|
||||
// weightScale入参在perchannel场景期望shape [E, N]
|
||||
n = weightScale->GetViewShape().GetDim(DIM_IDX_1);
|
||||
weightScaleExpectShape = {e, n};
|
||||
} else if (weightScale->GetViewShape().GetDimNum() == WEIGHT_SCALE_PERGROUP_DIM_LIMIT) {
|
||||
// weightScale入参在pergroup场景期望shape [E, KGroupCount, N]
|
||||
n = weightScale->GetViewShape().GetDim(DIM_IDX_2);
|
||||
KGroupCount = weightScale->GetViewShape().GetDim(DIM_IDX_1);
|
||||
KGroupSize = KGroupCount > 0 ? k / KGroupCount : k;
|
||||
weightScaleExpectShape = {e, KGroupCount, n};
|
||||
}
|
||||
if (KGroupCount == 0 || k % KGroupCount != 0) {
|
||||
OP_LOGE(ACLNN_ERR_PARAM_INVALID,
|
||||
"aclnnGroupedMatmulSwiGluQuant, "
|
||||
"The number of groups along the k-axis is %ld, and the length of the k-axis is %ld, which is illegal. "
|
||||
"The number of groups must be greater than 0, and k-axis length %% number of groups == 0 must be true.",
|
||||
KGroupCount, k);
|
||||
return false;
|
||||
}
|
||||
if (n % SPLIT != 0) {
|
||||
OP_LOGE(ACLNN_ERR_PARAM_INVALID, "aclnnGroupedMatmulSwiGluQuant, N is %ld , which must even number.", n);
|
||||
return false;
|
||||
}
|
||||
int64_t nAfterHalve = static_cast<int64_t>(n / SPLIT);
|
||||
// x的shape期望为[M, K]
|
||||
op::Shape xExpectShape = {m, k};
|
||||
// weight的NDshape期望为[E, K, N]
|
||||
op::Shape weightNDExpectShape = {e, k, n};
|
||||
op::Shape biasExpectShape = {e, n};
|
||||
// weight的NZshape期望为[E, N // 64, K // 16, 16, 64]
|
||||
op::Shape weightNZExpectShape = {e, static_cast<int64_t>(n / NZ_DIM_4_INT4), static_cast<int64_t>(k / NZ_DIM_3),
|
||||
NZ_DIM_3, NZ_DIM_4_INT4};
|
||||
// xScale的shape期望为[E, N]
|
||||
op::Shape xScaleExpectShape = {m};
|
||||
// output的shape期望为[M, N]
|
||||
op::Shape outputExpectShape = {m, nAfterHalve};
|
||||
// outputScale的shape期望为[M]
|
||||
op::Shape outputScaleExpectShape = {m};
|
||||
op::Format weightViewFormat = weight->GetViewFormat();
|
||||
OP_CHECK_SHAPE_NOT_EQUAL_WITH_EXPECTED_SIZE(x, xExpectShape, return false);
|
||||
if (IsPrivateFormat(weightViewFormat)) {
|
||||
OP_CHECK_SHAPE_NOT_EQUAL_WITH_EXPECTED_SIZE(weight, weightNZExpectShape, return false);
|
||||
} else {
|
||||
OP_CHECK_SHAPE_NOT_EQUAL_WITH_EXPECTED_SIZE(weight, weightNDExpectShape, return false);
|
||||
}
|
||||
OP_CHECK_SHAPE_NOT_EQUAL_WITH_EXPECTED_SIZE(bias, biasExpectShape, return false);
|
||||
OP_CHECK_SHAPE_NOT_EQUAL_WITH_EXPECTED_SIZE(weightScale, weightScaleExpectShape, return false);
|
||||
OP_CHECK_SHAPE_NOT_EQUAL_WITH_EXPECTED_SIZE(xScale, xScaleExpectShape, return false);
|
||||
OP_CHECK_SHAPE_NOT_EQUAL_WITH_EXPECTED_SIZE(output, outputExpectShape, return false);
|
||||
OP_CHECK_SHAPE_NOT_EQUAL_WITH_EXPECTED_SIZE(outputScale, outputScaleExpectShape, return false);
|
||||
// groupList的长度应小于等于weight的专家数
|
||||
int64_t groupListLen = groupList->GetViewShape().GetDim(0);
|
||||
if (groupListLen > e) {
|
||||
OP_LOGE(ACLNN_ERR_PARAM_INVALID,
|
||||
"aclnnGroupedMatmulSwiGluQuant A8W4, Length of 'groupList' out of range"
|
||||
" (expected to be in range of [1, %ld], but got %ld)",
|
||||
e, groupListLen);
|
||||
return false;
|
||||
}
|
||||
if (n > N_LIMIT) {
|
||||
OP_LOGE(ACLNN_ERR_PARAM_INVALID,
|
||||
"aclnnGroupedMatmulSwiGluQuant A8W4, The current version does not support the scenario."
|
||||
"where N after halve is %ld greater than %ld.",
|
||||
n, N_LIMIT);
|
||||
return false;
|
||||
}
|
||||
if (k >= K_LIMIT_A8W4) {
|
||||
OP_LOGE(ACLNN_ERR_PARAM_INVALID,
|
||||
"aclnnGroupedMatmulSwiGluQuant A8W4, The current version does not support the scenario."
|
||||
"The tail axis dimension of input0(x) is %ld, which need lower than %ld.",
|
||||
k, K_LIMIT_A8W4);
|
||||
return false;
|
||||
}
|
||||
(void)KGroupSize;
|
||||
return true;
|
||||
}
|
||||
|
||||
static bool CheckDtypeValid(const aclTensor *x, const aclTensor *weight, const aclTensor *bias,
|
||||
const aclTensor *weightScale, const aclTensor *xScale, const aclTensor *groupList,
|
||||
const aclTensor *output, const aclTensor *outputScale)
|
||||
{
|
||||
OP_CHECK_DTYPE_NOT_SUPPORT(x, X_DTYPE_SUPPORT_LIST, return false);
|
||||
OP_CHECK_DTYPE_NOT_SUPPORT(weight, WEIGHT_DTYPE_SUPPORT_LIST, return false);
|
||||
if (weight->GetDataType() == DataType::DT_INT4) {
|
||||
OP_CHECK_DTYPE_NOT_SUPPORT(bias, BIAS_DTYPE_SUPPORT_LIST, return false);
|
||||
}
|
||||
if (weight->GetDataType() == DataType::DT_INT4) {
|
||||
OP_CHECK_DTYPE_NOT_SUPPORT(weightScale, WEIGHT_SCALE_A8W4_DTYPE_SUPPORT_LIST, return false);
|
||||
} else if (weight->GetDataType() == DataType::DT_INT8) {
|
||||
OP_CHECK_DTYPE_NOT_SUPPORT(weightScale, WEIGHT_SCALE_DTYPE_SUPPORT_LIST, return false);
|
||||
}
|
||||
OP_CHECK_DTYPE_NOT_SUPPORT(xScale, X_SCALE_DTYPE_SUPPORT_LIST, return false);
|
||||
OP_CHECK_DTYPE_NOT_SUPPORT(groupList, GROUP_LIST_DTYPE_SUPPORT_LIST, return false);
|
||||
OP_CHECK_DTYPE_NOT_SUPPORT(output, QUANTOUT_DTYPE_SUPPORT_LIST, return false);
|
||||
OP_CHECK_DTYPE_NOT_SUPPORT(outputScale, QUANTSCALEOUT_DTYPE_SUPPORT_LIST, return false);
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
static bool CheckFormat(const aclTensor *x, const aclTensor *weight, const aclTensor *output)
|
||||
{
|
||||
bool isNZ = weight->GetStorageFormat() == op::Format::FORMAT_FRACTAL_NZ;
|
||||
if ((x->GetDataType() == DataType::DT_INT8 && weight->GetDataType() == DataType::DT_INT8) && !isNZ) {
|
||||
// fp16 in fp32 out that is split k template, not precision-advanced now
|
||||
OP_LOGE(ACLNN_ERR_PARAM_INVALID,
|
||||
"aclnnGroupedMatmulSwiGluQuant, The current version does not support the scenario."
|
||||
"weight Format expect is FRACTAL_NZ, but got [%s].",
|
||||
op::ToString(weight->GetStorageFormat()).GetString());
|
||||
return false;
|
||||
}
|
||||
if (IsPrivateFormat(x->GetStorageFormat())) {
|
||||
OP_LOGE(ACLNN_ERR_PARAM_INVALID,
|
||||
"aclnnGroupedMatmulSwiGluQuant, The current version does not support the scenario."
|
||||
"x Format Not support Private Format.");
|
||||
return false;
|
||||
}
|
||||
if (IsPrivateFormat(output->GetStorageFormat())) {
|
||||
OP_LOGE(ACLNN_ERR_PARAM_INVALID,
|
||||
"aclnnGroupedMatmulSwiGluQuant, The current version does not support the scenario."
|
||||
"output Format Not support Private Format.");
|
||||
return false;
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
static void UnpackInt32ToInt4(const aclTensor *&tensorS32, const std::string &tensorType)
|
||||
{
|
||||
OP_LOGD("Unpack %s from int32 to int4 start.", tensorType.c_str());
|
||||
auto tensorS4 = const_cast<aclTensor *>(tensorS32);
|
||||
op::Shape tensorShape = tensorS4->GetViewShape();
|
||||
auto viewShapeDim = tensorShape.GetDimNum();
|
||||
tensorShape[viewShapeDim - 1] = tensorShape[viewShapeDim - 1] * INT4_PER_INT32;
|
||||
tensorS4->SetViewShape(tensorShape);
|
||||
tensorS4->SetStorageShape(tensorShape);
|
||||
tensorS4->SetDataType(DataType::DT_INT4);
|
||||
OP_LOGD("Unpack %s from int32 to int4 finished.", tensorType.c_str());
|
||||
}
|
||||
|
||||
static aclnnStatus CheckParams(const aclTensor *x, const aclTensor *weight, const aclTensor *bias,
|
||||
const aclTensor *offset, const aclTensor *weightScale, const aclTensor *xScale,
|
||||
const aclTensor *groupList, const aclTensor *output, const aclTensor *outputScale,
|
||||
const aclTensor *outputOffset)
|
||||
{
|
||||
// 1. 检查参数是否为空指针
|
||||
CHECK_RET(CheckNotNull(x, weight, bias, offset, weightScale, xScale, groupList, output, outputScale, outputOffset),
|
||||
ACLNN_ERR_PARAM_NULLPTR);
|
||||
// A8W8场景
|
||||
if (x->GetDataType() == DataType::DT_INT8 && weight->GetDataType() == DataType::DT_INT8) {
|
||||
// 2. 校验输入、输出参数维度
|
||||
CHECK_RET(CheckInputOutDims_A8W8(x, weight, weightScale, xScale, groupList, output, outputScale),
|
||||
ACLNN_ERR_PARAM_INVALID);
|
||||
|
||||
// 3. 校验输入、输出shape参数
|
||||
CHECK_RET(CheckInputOutShape_A8W8(x, weight, weightScale, xScale, groupList, output, outputScale),
|
||||
ACLNN_ERR_PARAM_INVALID);
|
||||
}
|
||||
// A8W4场景 INT32为兼容torch_npu考虑，实际计算时，1个INT32数据会被视为8个INT4数据
|
||||
if ((x->GetDataType() == DataType::DT_INT8 && weight->GetDataType() == DataType::DT_INT4) ||
|
||||
(x->GetDataType() == DataType::DT_INT8 && weight->GetDataType() == DataType::DT_INT32)) {
|
||||
// 将INT32视为8个Int4数据，调整viewShape和dtype便于后续统一校验
|
||||
if (weight->GetDataType() == DataType::DT_INT32) {
|
||||
UnpackInt32ToInt4(weight, "weight");
|
||||
}
|
||||
if (weightScale->GetDataType() == DataType::DT_INT64) {
|
||||
auto weightScale_fix = const_cast<aclTensor *>(weightScale);
|
||||
weightScale_fix->SetDataType(DataType::DT_UINT64);
|
||||
}
|
||||
// 2. 校验输入、输出参数维度
|
||||
CHECK_RET(CheckInputOutDims_A8W4(x, weight, bias, weightScale, xScale, groupList, output, outputScale),
|
||||
ACLNN_ERR_PARAM_INVALID);
|
||||
|
||||
// 3. 校验输入、输出shape参数
|
||||
CHECK_RET(CheckInputOutShape_A8W4(x, weight, bias, weightScale, xScale, groupList, output, outputScale),
|
||||
ACLNN_ERR_PARAM_INVALID);
|
||||
}
|
||||
// 4. 检查输入的数据类型是否在支持的数据类型范围之内
|
||||
CHECK_RET(CheckDtypeValid(x, weight, bias, weightScale, xScale, groupList, output, outputScale),
|
||||
ACLNN_ERR_PARAM_INVALID);
|
||||
|
||||
// 5. 检查数据形状是否支持
|
||||
CHECK_RET(CheckFormat(x, weight, output), ACLNN_ERR_PARAM_INVALID);
|
||||
return ACLNN_SUCCESS;
|
||||
}
|
||||
|
||||
static aclnnStatus aclnnGroupedMatmulSwigluQuantGetWorkspaceSizeCommon(
|
||||
const aclTensor *x, const aclTensor *weight, const aclTensor *bias, const aclTensor *offset,
|
||||
const aclTensor *weightScale, const aclTensor *xScale, const aclTensor *groupList, double limited, aclTensor *output,
|
||||
aclTensor *outputScale, aclTensor *outputOffset, uint64_t *workspaceSize, aclOpExecutor **executor)
|
||||
{
|
||||
// 固定写法，创建OpExecutor
|
||||
auto uniqueExecutor = CREATE_EXECUTOR();
|
||||
CHECK_RET(uniqueExecutor.get() != nullptr, ACLNN_ERR_INNER_CREATE_EXECUTOR);
|
||||
// 固定写法，参数检查
|
||||
|
||||
auto ret = CheckParams(x, weight, bias, offset, weightScale, xScale, groupList, output, outputScale, outputOffset);
|
||||
|
||||
CHECK_RET(ret == ACLNN_SUCCESS, ret);
|
||||
// 空Tensor场景
|
||||
if (output->IsEmpty() || groupList->IsEmpty() || outputScale->IsEmpty()) {
|
||||
*workspaceSize = 0;
|
||||
uniqueExecutor.ReleaseTo(executor);
|
||||
return ACLNN_SUCCESS;
|
||||
}
|
||||
// 转连续
|
||||
x = l0op::Contiguous(x, uniqueExecutor.get());
|
||||
CHECK_RET(x != nullptr, ACLNN_ERR_INNER_NULLPTR);
|
||||
// 若weight为私有格式，则不应该做连续性转换 （l0op::Contiguous接口会把viewShape赋值给storageShape）
|
||||
if (IsPrivateFormat(weight->GetStorageFormat())) {
|
||||
weight->SetOriginalShape(weight->GetViewShape());
|
||||
} else {
|
||||
weight = l0op::Contiguous(weight, uniqueExecutor.get());
|
||||
}
|
||||
CHECK_RET(weight != nullptr, ACLNN_ERR_INNER_NULLPTR);
|
||||
weightScale = l0op::Contiguous(weightScale, uniqueExecutor.get());
|
||||
CHECK_RET(weightScale != nullptr, ACLNN_ERR_INNER_NULLPTR);
|
||||
xScale = l0op::Contiguous(xScale, uniqueExecutor.get());
|
||||
CHECK_RET(xScale != nullptr, ACLNN_ERR_INNER_NULLPTR);
|
||||
groupList = l0op::Contiguous(groupList, uniqueExecutor.get());
|
||||
CHECK_RET(groupList != nullptr, ACLNN_ERR_INNER_NULLPTR);
|
||||
// 调用L0算子能力
|
||||
if (bias != nullptr) {
|
||||
isEnableWeightAssistanceMatrix = true;
|
||||
bias = l0op::Contiguous(bias, uniqueExecutor.get());
|
||||
CHECK_RET(bias != nullptr, ACLNN_ERR_INNER_NULLPTR);
|
||||
}
|
||||
if (isEnableWeightAssistanceMatrix && weightScale->GetViewShape().GetDimNum() == WEIGHT_SCALE_PERGROUP_DIM_LIMIT) {
|
||||
dequantMode = 1;
|
||||
}
|
||||
auto ret_0 = l0op::GroupedMatmulSwigluQuant(x, weight, weightScale, xScale, groupList, limited, bias,
|
||||
isEnableWeightAssistanceMatrix, dequantMode, uniqueExecutor.get());
|
||||
CHECK_RET(ret_0 != std::tuple(nullptr, nullptr), ACLNN_ERR_INNER_NULLPTR);
|
||||
auto out0 = std::get<OUTPUT_IDX_0>(ret_0);
|
||||
auto ret_1 = l0op::ViewCopy(out0, output, uniqueExecutor.get());
|
||||
CHECK_RET(ret_1 != nullptr, ACLNN_ERR_INNER_NULLPTR);
|
||||
auto out1 = std::get<OUTPUT_IDX_1>(ret_0);
|
||||
auto ret_2 = l0op::ViewCopy(out1, outputScale, uniqueExecutor.get());
|
||||
CHECK_RET(ret_2 != nullptr, ACLNN_ERR_INNER_NULLPTR);
|
||||
*workspaceSize = uniqueExecutor->GetWorkspaceSize();
|
||||
uniqueExecutor.ReleaseTo(executor);
|
||||
return ACLNN_SUCCESS;
|
||||
}
|
||||
|
||||
aclnnStatus aclnnGroupedMatmulSwigluQuantGetWorkspaceSize(const aclTensor *x, const aclTensor *weight,
|
||||
const aclTensor *bias, const aclTensor *offset,
|
||||
const aclTensor *weightScale, const aclTensor *xScale,
|
||||
const aclTensor *groupList, double limited, aclTensor *output,
|
||||
aclTensor *outputScale, aclTensor *outputOffset,
|
||||
uint64_t *workspaceSize, aclOpExecutor **executor)
|
||||
{
|
||||
OP_CHECK_COMM_INPUT(workspaceSize, executor);
|
||||
L2_DFX_PHASE_1(aclnnGroupedMatmulSwigluQuant, DFX_IN(x, weight, bias, offset, weightScale, xScale, groupList, limited),
|
||||
DFX_OUT(output, outputScale, outputOffset));
|
||||
// 固定写法，创建OpExecutor
|
||||
return aclnnGroupedMatmulSwigluQuantGetWorkspaceSizeCommon(x, weight, bias, offset, weightScale, xScale, groupList, limited,
|
||||
output, outputScale, outputOffset, workspaceSize,
|
||||
executor);
|
||||
}
|
||||
|
||||
aclnnStatus aclnnGroupedMatmulSwigluQuantWeightNZGetWorkspaceSize(const aclTensor *x, const aclTensor *weight,
|
||||
const aclTensor *bias, const aclTensor *offset,
|
||||
const aclTensor *weightScale, const aclTensor *xScale,
|
||||
const aclTensor *groupList, double limited, aclTensor *output,
|
||||
aclTensor *outputScale, aclTensor *outputOffset,
|
||||
uint64_t *workspaceSize, aclOpExecutor **executor)
|
||||
{
|
||||
OP_CHECK_COMM_INPUT(workspaceSize, executor);
|
||||
L2_DFX_PHASE_1(aclnnGroupedMatmulSwigluQuantWeightNZ,
|
||||
DFX_IN(x, weight, bias, offset, weightScale, xScale, groupList),
|
||||
DFX_OUT(output, outputScale, outputOffset));
|
||||
// weight在该场景下强制绑定StorageFormat 和 ViewFormat 为NZ
|
||||
CHECK_RET(weight != nullptr, ACLNN_ERR_PARAM_NULLPTR);
|
||||
auto storgeShape = weight->GetStorageShape();
|
||||
auto viewShape = weight->GetViewShape();
|
||||
aclTensor *weightNZ = const_cast<aclTensor *>(weight);
|
||||
CHECK_COND((storgeShape.GetDimNum() == WEIGHT_NZ_DIM_LIMIT), ACLNN_ERR_PARAM_INVALID,
|
||||
"aclnnGroupedMatmulSwigluQuantWeightNZ, The dimnum of storageShape for second input (weight)"
|
||||
"must be 5. \n But StorageShape got %s , and dimNum is %lu.",
|
||||
op::ToString(storgeShape).GetString(), storgeShape.GetDimNum());
|
||||
// weight的StorageFormat无条件视为NZ
|
||||
weightNZ->SetStorageFormat(op::Format::FORMAT_FRACTAL_NZ);
|
||||
if (viewShape.GetDimNum() == WEIGHT_NZ_DIM_LIMIT) {
|
||||
// 若weight的viewShape为5维则视为NZ
|
||||
weightNZ->SetViewFormat(op::Format::FORMAT_FRACTAL_NZ);
|
||||
} else if (viewShape.GetDimNum() == WEIGHT_ND_DIM_LIMIT) {
|
||||
// 若weight的viewShape为3维则视为ND
|
||||
weightNZ->SetViewFormat(op::Format::FORMAT_ND);
|
||||
}
|
||||
// 调用公共接口
|
||||
return aclnnGroupedMatmulSwigluQuantGetWorkspaceSizeCommon(x, weight, bias, offset, weightScale, xScale, groupList, limited,
|
||||
output, outputScale, outputOffset, workspaceSize,
|
||||
executor);
|
||||
}
|
||||
|
||||
aclnnStatus aclnnGroupedMatmulSwigluQuant(void *workspace, uint64_t workspaceSize, aclOpExecutor *executor,
|
||||
aclrtStream stream)
|
||||
{
|
||||
L2_DFX_PHASE_2(aclnnGroupedMatmulSwigluQuant);
|
||||
CHECK_COND(CommonOpExecutorRun(workspace, workspaceSize, executor, stream) == ACLNN_SUCCESS, ACLNN_ERR_INNER,
|
||||
"This is an error in GroupedMatmulSwigluQuant launch aicore");
|
||||
return ACLNN_SUCCESS;
|
||||
}
|
||||
|
||||
aclnnStatus aclnnGroupedMatmulSwigluQuantWeightNZ(void *workspace, uint64_t workspaceSize, aclOpExecutor *executor,
|
||||
aclrtStream stream)
|
||||
{
|
||||
L2_DFX_PHASE_2(aclnnGroupedMatmulSwigluQuantWeightNZ);
|
||||
CHECK_COND(CommonOpExecutorRun(workspace, workspaceSize, executor, stream) == ACLNN_SUCCESS, ACLNN_ERR_INNER,
|
||||
"This is an error in GroupedMatmulSwigluQuantWeightNZ launch aicore");
|
||||
return ACLNN_SUCCESS;
|
||||
}
|
||||
|
||||
#ifdef __cplusplus
|
||||
}
|
||||
#endif
|
||||
@@ -0,0 +1,58 @@
|
||||
/**
|
||||
* Copyright (c) 2025 Huawei Technologies Co., Ltd.
|
||||
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
* CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
* See LICENSE in the root of the software repository for the full text of the License.
|
||||
*/
|
||||
#ifndef OP_API_INC_GROUPED_MATMUL_SWIGLU_QUANT_H
|
||||
#define OP_API_INC_GROUPED_MATMUL_SWIGLU_QUANT_H
|
||||
#include "aclnn/aclnn_base.h"
|
||||
|
||||
#ifdef __cplusplus
|
||||
extern "C" {
|
||||
#endif
|
||||
|
||||
/**
|
||||
* @brief aclnnGroupedMatmulSwigluQuant的第一段接口，根据具体的计算流程，计算workspace大小。
|
||||
* @domain aclnn_ops_infer
|
||||
*
|
||||
* @param [in] x: 表示公式中的x，数据类型支持INT8数据类型，数据格式支持ND。
|
||||
* @param [in] weight:
|
||||
* 表示公式中的weight，数据类型支持INT8数据类型，数据格式支持NZ。
|
||||
* @param [in] weightScale:
|
||||
* 表示量化参数，数据类型支持FLOAT16、BFLOAT16、FLOAT32数据类型，数据格式支持ND，支持的最大长度为128个。 表示per
|
||||
* Channel参数，数据类型支持FLOAT16，BFLOAT16数据类型，数据格式支持ND。
|
||||
* @param [in] xScale:
|
||||
* 表示per Token量化参数，数据类型支持FLOAT32数据类型，数据格式支持ND。
|
||||
* @param [in] groupList: 必选参数，代表输入和输出分组轴上的索引情况，数据类型支持INT64。
|
||||
* @param [out] quantOutput: 表示公式中的out，数据类型支持INT8数据类型，数据格式支持ND。
|
||||
* @param [out] quantScaleOutput: 表示公式中的outQuantScale，数据类型支持Float32数据类型。
|
||||
* @param [out] workspaceSize: 返回用户需要在npu device侧申请的workspace大小。
|
||||
* @param [out] executor: 返回op执行器，包含算子计算流程。
|
||||
* @return aclnnStatus: 返回状态码。
|
||||
*/
|
||||
__attribute__((visibility("default"))) aclnnStatus aclnnGroupedMatmulSwigluQuantGetWorkspaceSize(
|
||||
const aclTensor *x, const aclTensor *weight, const aclTensor *bias, const aclTensor *offset,
|
||||
const aclTensor *weightScale, const aclTensor *xScale, const aclTensor *groupList, double limited, aclTensor *output,
|
||||
aclTensor *outputScale, aclTensor *outputOffset, uint64_t *workspaceSize, aclOpExecutor **executor);
|
||||
|
||||
/**
|
||||
* @brief aclnnGroupedMatmulSwigluQuant的第二段接口，用于执行计算。
|
||||
* @param [in] workspace: 在npu device侧申请的workspace内存起址。
|
||||
* @param [in] workspaceSize: 在npu
|
||||
* device侧申请的workspace大小，由第一段接口aclnnGroupedMatmulSwigluQuantGetWorkspaceSize获取。
|
||||
* @param [in] stream: acl stream流。
|
||||
* @param [in] executor: op执行器，包含了算子计算流程。
|
||||
* @return aclnnStatus: 返回状态码。
|
||||
*/
|
||||
__attribute__((visibility("default"))) aclnnStatus
|
||||
aclnnGroupedMatmulSwigluQuant(void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream);
|
||||
|
||||
#ifdef __cplusplus
|
||||
}
|
||||
#endif
|
||||
|
||||
#endif
|
||||
@@ -0,0 +1,60 @@
|
||||
/**
|
||||
* Copyright (c) 2025 Huawei Technologies Co., Ltd.
|
||||
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
* CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
* See LICENSE in the root of the software repository for the full text of the License.
|
||||
*/
|
||||
#ifndef OP_API_INC_GROUPED_MATMUL_SWIGLU_QUANT_WEIGHT_NZ_H
|
||||
#define OP_API_INC_GROUPED_MATMUL_SWIGLU_QUANT_WEIGHT_NZ_H
|
||||
#include "aclnn/aclnn_base.h"
|
||||
|
||||
#ifdef __cplusplus
|
||||
extern "C" {
|
||||
#endif
|
||||
|
||||
/**
|
||||
* @brief aclnnGroupedMatmulSwigluQuantWeightNZ的第一段接口，根据具体的计算流程，计算workspace大小。
|
||||
* @domain aclnn_ops_infer
|
||||
*
|
||||
* @param [in] x: 表示公式中的x，数据类型支持INT8数据类型，数据格式支持ND。
|
||||
* @param [in] weight:
|
||||
* 表示公式中的weight，数据类型支持INT8数据类型，数据格式支持NZ。
|
||||
* @param [in] weightScale:
|
||||
* 表示量化参数，数据类型支持FLOAT16、BFLOAT16、FLOAT32数据类型，数据格式支持ND，支持的最大长度为128个。 表示per
|
||||
* Channel参数，数据类型支持FLOAT16，BFLOAT16数据类型，数据格式支持ND。
|
||||
* @param [in] xScale:
|
||||
* 表示per Token量化参数，数据类型支持FLOAT32数据类型，数据格式支持ND。
|
||||
* @param [in] groupList: 必选参数，代表输入和输出分组轴上的索引情况，数据类型支持INT64。
|
||||
* @param [out] quantOutput: 表示公式中的out，数据类型支持INT8数据类型，数据格式支持ND。
|
||||
* @param [out] quantScaleOutput: 表示公式中的outQuantScale，数据类型支持Float32数据类型。
|
||||
* @param [out] workspaceSize: 返回用户需要在npu device侧申请的workspace大小。
|
||||
* @param [out] executor: 返回op执行器，包含算子计算流程。
|
||||
* @return aclnnStatus: 返回状态码。
|
||||
*/
|
||||
__attribute__((visibility("default"))) aclnnStatus aclnnGroupedMatmulSwigluQuantWeightNZGetWorkspaceSize(
|
||||
const aclTensor *x, const aclTensor *weight, const aclTensor *bias, const aclTensor *offset,
|
||||
const aclTensor *weightScale, const aclTensor *xScale, const aclTensor *groupList, double limited, aclTensor *output,
|
||||
aclTensor *outputScale, aclTensor *outputOffset, uint64_t *workspaceSize, aclOpExecutor **executor);
|
||||
|
||||
/**
|
||||
* @brief aclnnGroupedMatmulSwigluQuantWeightNZ的第二段接口，用于执行计算。
|
||||
* @param [in] workspace: 在npu device侧申请的workspace内存起址。
|
||||
* @param [in] workspaceSize: 在npu
|
||||
* device侧申请的workspace大小，由第一段接口aclnnGroupedMatmulSwigluQuantWeightNZGetWorkspaceSize获取。
|
||||
* @param [in] stream: acl stream流。
|
||||
* @param [in] executor: op执行器，包含了算子计算流程。
|
||||
* @return aclnnStatus: 返回状态码。
|
||||
*/
|
||||
__attribute__((visibility("default"))) aclnnStatus aclnnGroupedMatmulSwigluQuantWeightNZ(void *workspace,
|
||||
uint64_t workspaceSize,
|
||||
aclOpExecutor *executor,
|
||||
aclrtStream stream);
|
||||
|
||||
#ifdef __cplusplus
|
||||
}
|
||||
#endif
|
||||
|
||||
#endif
|
||||
@@ -0,0 +1,58 @@
|
||||
/**
|
||||
* Copyright (c) 2025 Huawei Technologies Co., Ltd.
|
||||
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
* CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
* See LICENSE in the root of the software repository for the full text of the License.
|
||||
*/
|
||||
|
||||
#include "opdev/op_log.h"
|
||||
#include "opdev/op_dfx.h"
|
||||
#include "opdev/make_op_executor.h"
|
||||
#include "grouped_matmul_swiglu_quant.h"
|
||||
|
||||
using namespace op;
|
||||
|
||||
namespace l0op {
|
||||
OP_TYPE_REGISTER(GroupedMatmulSwigluQuant);
|
||||
|
||||
const std::tuple<aclTensor *, aclTensor *>
|
||||
GroupedMatmulSwigluQuant(const aclTensor *x, const aclTensor *weight, const aclTensor *perChannelScale,
|
||||
const aclTensor *perTokenScale, const aclTensor *groupList, double limited,
|
||||
const aclTensor *weightAssistanceMatrix, bool isEnableWeightAssistanceMatrix, int dequantMode,
|
||||
aclOpExecutor *executor)
|
||||
{
|
||||
L0_DFX(GroupedMatmulSwigluQuant, x, weight, perChannelScale, perTokenScale, weightAssistanceMatrix, groupList, limited,
|
||||
isEnableWeightAssistanceMatrix, dequantMode);
|
||||
if (x == nullptr) {
|
||||
OP_LOGE(ACLNN_ERR_PARAM_INVALID, "x is nullptr.");
|
||||
return std::tuple(nullptr, nullptr);
|
||||
}
|
||||
int64_t m = perTokenScale->GetViewShape().GetDim(0);
|
||||
int64_t n = perChannelScale->GetViewShape().GetDim(1);
|
||||
int64_t nAfterHalve = static_cast<int64_t>(n / 2);
|
||||
gert::Shape outShape({m, nAfterHalve});
|
||||
gert::Shape scaleOutShape({m});
|
||||
auto out = executor->AllocTensor(outShape, DataType::DT_INT8, ge::FORMAT_ND);
|
||||
auto scaleOut = executor->AllocTensor(scaleOutShape, DataType::DT_FLOAT, ge::FORMAT_ND);
|
||||
auto ret = INFER_SHAPE(GroupedMatmulSwigluQuant,
|
||||
OP_INPUT(x, weight, perChannelScale, perTokenScale, weightAssistanceMatrix, groupList),
|
||||
OP_OUTPUT(out, scaleOut), OP_ATTR(isEnableWeightAssistanceMatrix, dequantMode, limited));
|
||||
if (ret != ACLNN_SUCCESS) {
|
||||
OP_LOGE(ACLNN_ERR_PARAM_INVALID, "InferShape failed.");
|
||||
return std::tuple(nullptr, nullptr);
|
||||
}
|
||||
ret = ADD_TO_LAUNCHER_LIST_AICORE(
|
||||
GroupedMatmulSwigluQuant,
|
||||
OP_INPUT(x, weight, perChannelScale, perTokenScale, weightAssistanceMatrix, groupList),
|
||||
OP_OUTPUT(out, scaleOut), OP_ATTR(isEnableWeightAssistanceMatrix, dequantMode, limited));
|
||||
if (ret != ACLNN_SUCCESS) {
|
||||
OP_LOGE(ACLNN_ERR_PARAM_INVALID, "ADD_TO_LAUNCHER_LIST_AICORE failed.");
|
||||
return std::tuple(nullptr, nullptr);
|
||||
}
|
||||
return std::tie(out, scaleOut);
|
||||
}
|
||||
|
||||
} // namespace l0op
|
||||
@@ -0,0 +1,23 @@
|
||||
/**
|
||||
* Copyright (c) 2025 Huawei Technologies Co., Ltd.
|
||||
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
* CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
* See LICENSE in the root of the software repository for the full text of the License.
|
||||
*/
|
||||
#ifndef OP_API_INC_LEVEL0_OP_GROUPED_MATMUL_SWIGLU_QUANT_OP_H
|
||||
#define OP_API_INC_LEVEL0_OP_GROUPED_MATMUL_SWIGLU_QUANT_OP_H
|
||||
|
||||
#include "opdev/op_executor.h"
|
||||
|
||||
namespace l0op {
|
||||
const std::tuple<aclTensor *, aclTensor *>
|
||||
GroupedMatmulSwigluQuant(const aclTensor *x, const aclTensor *weight, const aclTensor *perChannelScale,
|
||||
const aclTensor *perTokenScale, const aclTensor *groupList, double limited,
|
||||
const aclTensor *weightAssistanceMatrix, bool isEnableWeightAssistanceMatrix, int dequantMode,
|
||||
aclOpExecutor *executor);
|
||||
}
|
||||
|
||||
#endif
|
||||
@@ -0,0 +1,197 @@
|
||||
/**
|
||||
* Copyright (c) 2025 Huawei Technologies Co., Ltd.
|
||||
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
* CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
* See LICENSE in the root of the software repository for the full text of the License.
|
||||
*/
|
||||
|
||||
/*!
|
||||
* \file grouped_matmul_swiglu_quant.cpp
|
||||
* \brief
|
||||
*/
|
||||
#ifndef ASCENDC_GROUPED_MATMUL_SWIGLU_QUANT_PIPELINE_H
|
||||
#define ASCENDC_GROUPED_MATMUL_SWIGLU_QUANT_PIPELINE_H
|
||||
#include "grouped_matmul_swiglu_quant.h"
|
||||
#include <typeinfo>
|
||||
#include "grouped_matmul_swiglu_quant_a8w4_msd_pre.h"
|
||||
#include "grouped_matmul_swiglu_quant_a8w4_msd_mid.h"
|
||||
#include "grouped_matmul_swiglu_quant_a8w4_msd_post.h"
|
||||
#include "grouped_matmul_swiglu_quant_utils.h"
|
||||
using namespace AscendC;
|
||||
using namespace matmul;
|
||||
#ifdef GMM_SWIGLU_QUANT_A8W4_MSD
|
||||
|
||||
namespace GROUPED_MATMUL_SWIGLU_QUANT {
|
||||
|
||||
template <class mmType>
|
||||
class GMMSwigluQuantPipelineSchedule {
|
||||
private:
|
||||
typename mmType::MT &mm;
|
||||
TPipe *pipe;
|
||||
const GMMSwigluBaseParams *__restrict gmmBaseParams;
|
||||
const GMMSwiglu *__restrict gmmSwiglu;
|
||||
// WorkSpaceSplitConfig控制Workspace切割方式的结构体;
|
||||
WorkSpaceSplitConfig workspaceSplitConfig;
|
||||
WorkSpaceSplitConfig tempWorkspaceSplitConfig;
|
||||
// 记录GM_ADDR的结构体
|
||||
GMAddrParams gmAddrParams;
|
||||
// 前处理GMMA8W4PreProcess类
|
||||
GMMA8W4PreProcess preProcess;
|
||||
// 中间处理GMMA8W4MidProcess类
|
||||
GMMA8W4MidProcess<mmType> midProcess;
|
||||
// 后处理GMMA8W4PostProcess类
|
||||
GMMA8W4PostProcess postProcess;
|
||||
GlobalTensor<int64_t> groupListGM;
|
||||
__aicore__ inline void InitWorkSpaceSplitConfig(WorkSpaceSplitConfig &workspaceSplitConfig);
|
||||
|
||||
__aicore__ inline void UpdateWorkSpaceSplitConfig(WorkSpaceSplitConfig &workspaceSplitConfig,
|
||||
int32_t workspaceSplitLoopIdx);
|
||||
|
||||
public:
|
||||
__aicore__ inline GMMSwigluQuantPipelineSchedule(typename mmType::MT &mm_,
|
||||
const GMMSwigluBaseParams *__restrict gmmBaseParamsIN,
|
||||
const GMMSwiglu *__restrict gmmSwigluIN, TPipe *tPipeIN)
|
||||
: mm(mm_), midProcess(mm), gmmBaseParams(gmmBaseParamsIN), gmmSwiglu(gmmSwigluIN), pipe(tPipeIN)
|
||||
{
|
||||
}
|
||||
__aicore__ inline void Init(GM_ADDR x, GM_ADDR weight, GM_ADDR weightScale, GM_ADDR xScale,
|
||||
GM_ADDR weightAssistanceMatrix, GM_ADDR groupList, GM_ADDR y, GM_ADDR yScale,
|
||||
GM_ADDR workspace);
|
||||
__aicore__ inline void Process();
|
||||
};
|
||||
|
||||
template <class mmType>
|
||||
__aicore__ inline void GMMSwigluQuantPipelineSchedule<mmType>::Init(GM_ADDR x, GM_ADDR weight, GM_ADDR weightScale,
|
||||
GM_ADDR xScale, GM_ADDR weightAssistanceMatrix,
|
||||
GM_ADDR groupList, GM_ADDR y, GM_ADDR yScale,
|
||||
GM_ADDR workspace)
|
||||
{
|
||||
gmAddrParams.xGM = x;
|
||||
gmAddrParams.weightGM = weight;
|
||||
gmAddrParams.weightScaleGM = weightScale;
|
||||
gmAddrParams.xScaleGM = xScale;
|
||||
gmAddrParams.weightAuxiliaryMatrixGM = weightAssistanceMatrix;
|
||||
gmAddrParams.groupListGM = groupList;
|
||||
gmAddrParams.yGM = y;
|
||||
gmAddrParams.yScaleGM = yScale;
|
||||
gmAddrParams.workSpaceGM = workspace;
|
||||
gmAddrParams.workSpaceOffset1 = gmmBaseParams->workSpaceOffset1 / 2;
|
||||
gmAddrParams.workSpaceOffset2 = gmmBaseParams->workSpaceOffset1;
|
||||
gmAddrParams.workSpaceOffset3 = gmmBaseParams->workSpaceOffset1 + gmmBaseParams->workSpaceOffset2 / 2;
|
||||
groupListGM.SetGlobalBuffer((__gm__ int64_t *)gmAddrParams.groupListGM);
|
||||
InitWorkSpaceSplitConfig(workspaceSplitConfig);
|
||||
}
|
||||
|
||||
template <class mmType>
|
||||
__aicore__ inline void GMMSwigluQuantPipelineSchedule<mmType>::Process()
|
||||
{
|
||||
// 1.对每次workspace切分做大循环。
|
||||
preProcess.Init(gmAddrParams, gmmBaseParams);
|
||||
midProcess.Init(gmAddrParams, gmmBaseParams);
|
||||
postProcess.Init(gmAddrParams, gmmBaseParams, gmmSwiglu);
|
||||
|
||||
// 1.前处理提前下发一次
|
||||
preProcess.Process(workspaceSplitConfig, 0, pipe);
|
||||
for (int64_t workspaceSplitLoopIdx = 0; workspaceSplitLoopIdx < workspaceSplitConfig.loopCount;
|
||||
workspaceSplitLoopIdx++) {
|
||||
// 更新workspaceSplitConfig
|
||||
UpdateWorkSpaceSplitConfig(workspaceSplitConfig, workspaceSplitLoopIdx);
|
||||
if ASCEND_IS_AIV {
|
||||
pipe->Reset();
|
||||
}
|
||||
|
||||
SyncAll<false>();
|
||||
// 2.第n次中处理 && 第n+1次前处理 && 第n-1次后处理 并行
|
||||
midProcess.Process(workspaceSplitConfig, workspaceSplitLoopIdx);
|
||||
|
||||
preProcess.Process(workspaceSplitConfig, workspaceSplitLoopIdx + 1, pipe);
|
||||
if ASCEND_IS_AIV {
|
||||
pipe->Reset();
|
||||
SyncAll<true>();
|
||||
}
|
||||
postProcess.Process(tempWorkspaceSplitConfig, workspaceSplitLoopIdx - 1, pipe);
|
||||
// 3.第n-1次后处理需要保留第n次的切分数据
|
||||
tempWorkspaceSplitConfig = workspaceSplitConfig;
|
||||
// reset
|
||||
if ASCEND_IS_AIV {
|
||||
pipe->Reset();
|
||||
}
|
||||
SyncAll<false>();
|
||||
// 3.前一次后处理 && 后一次MM 并行
|
||||
}
|
||||
// reset
|
||||
if ASCEND_IS_AIV {
|
||||
pipe->Reset();
|
||||
}
|
||||
SyncAll<false>();
|
||||
// // 4.最后一次后处理
|
||||
postProcess.Process(workspaceSplitConfig, workspaceSplitConfig.loopCount - 1, pipe);
|
||||
if ASCEND_IS_AIV {
|
||||
pipe->Destroy();
|
||||
}
|
||||
}
|
||||
|
||||
template <class mmType>
|
||||
__aicore__ inline void
|
||||
GMMSwigluQuantPipelineSchedule<mmType>::InitWorkSpaceSplitConfig(WorkSpaceSplitConfig &workspaceSplitConfig)
|
||||
{
|
||||
workspaceSplitConfig.M = groupListGM.GetValue(gmmSwiglu->groupListLen - 1);
|
||||
workspaceSplitConfig.loopCount = Ceil(workspaceSplitConfig.M, gmmBaseParams->mLimit);
|
||||
workspaceSplitConfig.notLastTaskSize = gmmBaseParams->mLimit;
|
||||
workspaceSplitConfig.lastLoopTaskSize =
|
||||
workspaceSplitConfig.M - (workspaceSplitConfig.loopCount - 1) * gmmBaseParams->mLimit;
|
||||
workspaceSplitConfig.leftMatrixStartIndex = 0;
|
||||
workspaceSplitConfig.rightMatrixExpertStartIndex = 0;
|
||||
workspaceSplitConfig.rightMatrixExpertNextStartIndex = 0;
|
||||
workspaceSplitConfig.isLastLoop = false;
|
||||
}
|
||||
|
||||
template <class mmType>
|
||||
__aicore__ inline void
|
||||
GMMSwigluQuantPipelineSchedule<mmType>::UpdateWorkSpaceSplitConfig(WorkSpaceSplitConfig &workspaceSplitConfig,
|
||||
int32_t workspaceSplitLoopIdx)
|
||||
{
|
||||
if (workspaceSplitLoopIdx < 0)
|
||||
return;
|
||||
workspaceSplitConfig.leftMatrixStartIndex = workspaceSplitLoopIdx * gmmBaseParams->mLimit;
|
||||
workspaceSplitConfig.rightMatrixExpertStartIndex = workspaceSplitConfig.rightMatrixExpertNextStartIndex;
|
||||
workspaceSplitConfig.rightMatrixExpertEndIndex = workspaceSplitConfig.rightMatrixExpertStartIndex;
|
||||
// 计算右专家矩阵的终止索引(rightMatrixExpertEndIndex) 和下一次的起始索引(rightMatrixExpertNextStartIndex)
|
||||
int32_t curTaskNum = 0;
|
||||
int32_t nextTaskNum = 0;
|
||||
while (workspaceSplitConfig.rightMatrixExpertEndIndex < gmmSwiglu->groupListLen) {
|
||||
curTaskNum = groupListGM.GetValue(workspaceSplitConfig.rightMatrixExpertEndIndex) -
|
||||
workspaceSplitConfig.leftMatrixStartIndex;
|
||||
int32_t nextTaskIdx = workspaceSplitConfig.rightMatrixExpertEndIndex >= gmmSwiglu->groupListLen - 1 ?
|
||||
gmmSwiglu->groupListLen - 1 :
|
||||
workspaceSplitConfig.rightMatrixExpertEndIndex + 1;
|
||||
nextTaskNum = groupListGM.GetValue(nextTaskIdx) - workspaceSplitConfig.leftMatrixStartIndex;
|
||||
if (curTaskNum > gmmBaseParams->mLimit) {
|
||||
workspaceSplitConfig.rightMatrixExpertNextStartIndex = workspaceSplitConfig.rightMatrixExpertEndIndex;
|
||||
break;
|
||||
} else if (curTaskNum == gmmBaseParams->mLimit && nextTaskNum > gmmBaseParams->mLimit) {
|
||||
workspaceSplitConfig.rightMatrixExpertNextStartIndex = workspaceSplitConfig.rightMatrixExpertEndIndex + 1;
|
||||
break;
|
||||
} else if (nextTaskNum > gmmBaseParams->mLimit) {
|
||||
workspaceSplitConfig.rightMatrixExpertEndIndex++;
|
||||
workspaceSplitConfig.rightMatrixExpertNextStartIndex = workspaceSplitConfig.rightMatrixExpertEndIndex;
|
||||
break;
|
||||
}
|
||||
workspaceSplitConfig.rightMatrixExpertEndIndex++;
|
||||
}
|
||||
workspaceSplitConfig.isLastLoop = workspaceSplitLoopIdx == workspaceSplitConfig.loopCount - 1 ? true : false;
|
||||
|
||||
if (workspaceSplitConfig.isLastLoop) {
|
||||
workspaceSplitConfig.rightMatrixExpertEndIndex =
|
||||
workspaceSplitConfig.rightMatrixExpertEndIndex >= gmmSwiglu->groupListLen ?
|
||||
gmmSwiglu->groupListLen - 1 :
|
||||
workspaceSplitConfig.rightMatrixExpertEndIndex;
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace GROUPED_MATMUL_SWIGLU_QUANT
|
||||
#endif // GMM_SWIGLU_QUANT_A8W4_MSD
|
||||
#endif // ASCENDC_GROUPED_MATMUL_SWIGLU_QUANT_PIPELINE_H
|
||||
@@ -0,0 +1,102 @@
|
||||
/**
|
||||
* Copyright (c) 2025 Huawei Technologies Co., Ltd.
|
||||
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
* CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
* See LICENSE in the root of the software repository for the full text of the License.
|
||||
*/
|
||||
|
||||
/*!
|
||||
* \file grouped_matmul_swiglu_quant.cpp
|
||||
* \brief
|
||||
*/
|
||||
#include "grouped_matmul_swiglu_quant.h"
|
||||
#include "grouped_matmul_swiglu_pipeline.h"
|
||||
#include "grouped_matmul_swiglu_quant_utils.h"
|
||||
#include <typeinfo>
|
||||
#include "grouped_matmul_swiglu_quant_split_ws.h"
|
||||
using namespace AscendC;
|
||||
using namespace matmul;
|
||||
using namespace GROUPED_MATMUL_SWIGLU_QUANT;
|
||||
|
||||
#define GMM_CV_SPLIT_IMP(computeClass, dtypeWeightScale, transA, transB, sync) \
|
||||
do { \
|
||||
using xType = MatmulType<AscendC::TPosition::GM, CubeFormat::ND, DTYPE_X, false>; \
|
||||
using weightType = MatmulType<AscendC::TPosition::GM, CubeFormat::NZ, DTYPE_WEIGHT, false>; \
|
||||
using yType = MatmulType<AscendC::TPosition::GM, CubeFormat::ND, int32_t>; \
|
||||
using matmulType = MMImplTypeStatic<xType, weightType, yType>; \
|
||||
matmulType::MT mm; \
|
||||
GET_TILING_DATA_MEMBER(GMMSwigluQuantTilingData, gmmSwigluBaseParams, gmmSwigluBaseParams_, tiling); \
|
||||
GET_TILING_DATA_MEMBER(GMMSwigluQuantTilingData, mmTilingData, mmTilingData_, tiling); \
|
||||
GET_TILING_DATA_MEMBER(GMMSwigluQuantTilingData, gmmSwiglu, gmmSwiglu_, tiling); \
|
||||
if ASCEND_IS_AIC { \
|
||||
mm.SetSubBlockIdx(0); \
|
||||
mm.Init(&mmTilingData_, &tPipe); \
|
||||
} \
|
||||
computeClass<matmulType, sync, dtypeWeightScale> computeOp(mm); \
|
||||
computeOp.Init(x, weight, weightScale, xScale, groupList, y, yScale, user1, &gmmSwigluBaseParams_, \
|
||||
&mmTilingData_, &gmmSwiglu_, &tPipe); \
|
||||
computeOp.Process(); \
|
||||
} while (0)
|
||||
|
||||
#define GMM_CV_SPLIT_IMP_A8W4_MSD(computeClass, dtypeWeightScale, transA, transB, sync) \
|
||||
do { \
|
||||
GET_TILING_DATA_MEMBER(GMMSwigluQuantTilingData, gmmSwigluBaseParams, gmmSwigluBaseParams_, tiling); \
|
||||
GET_TILING_DATA_MEMBER(GMMSwigluQuantTilingData, mmTilingData, mmTilingData_, tiling); \
|
||||
GET_TILING_DATA_MEMBER(GMMSwigluQuantTilingData, gmmSwiglu, gmmSwiglu_, tiling); \
|
||||
using xType = MatmulType<TPosition::GM, CubeFormat::ND, int4b_t, false>; \
|
||||
using weightType = MatmulType<TPosition::GM, wFormat, int4b_t, false>; \
|
||||
using yType = MatmulType<TPosition::GM, CubeFormat::ND, half, false>; \
|
||||
using matmulType = MMImplType<xType, weightType, yType>; \
|
||||
matmulType::MT mm; \
|
||||
if ASCEND_IS_AIC { \
|
||||
mm.SetSubBlockIdx(0); \
|
||||
mm.Init(&mmTilingData_); \
|
||||
} \
|
||||
computeClass<matmulType> op(mm, &gmmSwigluBaseParams_, &gmmSwiglu_, &tPipe); \
|
||||
op.Init(x, weight, weightScale, xScale, weightAssistanceMatrix, groupList, y, yScale, user1); \
|
||||
\
|
||||
op.Process(); \
|
||||
} while (0)
|
||||
|
||||
extern "C" __global__ __aicore__ void grouped_matmul_swiglu_quant(GM_ADDR x, GM_ADDR weight, GM_ADDR weightScale,
|
||||
GM_ADDR xScale, GM_ADDR weightAssistanceMatrix,
|
||||
GM_ADDR groupList, GM_ADDR y, GM_ADDR yScale,
|
||||
GM_ADDR workspace, GM_ADDR tiling)
|
||||
{
|
||||
TPipe tPipe;
|
||||
AscendCUtils::SetOverflow(1);
|
||||
KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_MIX_AIC_1_2);
|
||||
GM_ADDR user1 = GetUserWorkspace(workspace);
|
||||
#if defined(GMM_SWIGLU_QUANT_A8W8)
|
||||
if (TILING_KEY_IS(0)) { // antiquant msd
|
||||
KERNEL_TASK_TYPE(0, KERNEL_TYPE_MIX_AIC_1_2);
|
||||
GMM_CV_SPLIT_IMP(GMMSwigluCompute, // computeClass
|
||||
DTYPE_WEIGHT_SCALE,
|
||||
false, // transA
|
||||
false, // transB
|
||||
false // sync
|
||||
);
|
||||
} else if (TILING_KEY_IS(1)) {
|
||||
KERNEL_TASK_TYPE(1, KERNEL_TYPE_MIX_AIC_1_2);
|
||||
GMM_CV_SPLIT_IMP(GMMSwigluSplitWorkSpaceCompute, // computeClass
|
||||
DTYPE_WEIGHT_SCALE,
|
||||
false, // transA
|
||||
false, // transB
|
||||
false // sync
|
||||
);
|
||||
}
|
||||
#elif defined(GMM_SWIGLU_QUANT_A8W4_MSD)
|
||||
if (TILING_KEY_IS(2)) {
|
||||
KERNEL_TASK_TYPE(2, KERNEL_TYPE_MIX_AIC_1_2);
|
||||
GMM_CV_SPLIT_IMP_A8W4_MSD(GMMSwigluQuantPipelineSchedule, // computeClass
|
||||
DTYPE_WEIGHT_SCALE,
|
||||
false, // transA
|
||||
false, // transB
|
||||
false // sync
|
||||
);
|
||||
}
|
||||
#endif
|
||||
}
|
||||
@@ -0,0 +1,547 @@
|
||||
/**
|
||||
* Copyright (c) 2025 Huawei Technologies Co., Ltd.
|
||||
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
* CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
* See LICENSE in the root of the software repository for the full text of the License.
|
||||
*/
|
||||
|
||||
/*!
|
||||
* \file grouped_matmul_swiglu_quant.h
|
||||
* \brief
|
||||
*/
|
||||
#ifndef ASCENDC_GROUPED_MATMUL_SWIGLU_QUANT_H
|
||||
#define ASCENDC_GROUPED_MATMUL_SWIGLU_QUANT_H
|
||||
|
||||
#include "grouped_matmul_swiglu_quant_utils.h"
|
||||
namespace GROUPED_MATMUL_SWIGLU_QUANT {
|
||||
/** @brief internal computation class
|
||||
*/
|
||||
template <class mmType, bool sync = false, typename CHANNELDTYPE = float>
|
||||
class GMMSwigluCompute {
|
||||
public:
|
||||
using AT = typename mmType::AT::T;
|
||||
using BT = typename mmType::BT::T;
|
||||
using B = typename mmType::BT;
|
||||
using CT = typename mmType::CT::T;
|
||||
using BiasT = typename mmType::BiasT::T;
|
||||
using WT = int8_t;
|
||||
constexpr static bool transposeX = mmType::AT::isTrans;
|
||||
constexpr static bool transposeW = mmType::BT::isTrans;
|
||||
static constexpr float FLOAT_INF = 3e+99;
|
||||
/** @brief constructor */
|
||||
__aicore__ inline GMMSwigluCompute(typename mmType::MT &mm_) : mm(mm_)
|
||||
{
|
||||
}
|
||||
|
||||
__aicore__ inline void Init(GM_ADDR x, GM_ADDR weight, GM_ADDR perChannelScale, GM_ADDR perTokenScale,
|
||||
GM_ADDR groupList, GM_ADDR quantOutput, GM_ADDR quantScaleOutput, GM_ADDR workspace,
|
||||
const GMMSwigluBaseParams *__restrict gmmBaseParamsIN,
|
||||
const TCubeTiling *__restrict mmTilingDataIN, const GMMSwiglu *__restrict gmmSwigluIN,
|
||||
TPipe *tPipeIN);
|
||||
__aicore__ inline void Process();
|
||||
|
||||
private:
|
||||
__aicore__ inline void MMCompute(uint32_t groupIdx, MNConfig &mnConfig, uint32_t coreIdx);
|
||||
|
||||
__aicore__ inline void UpdateMnConfig(MNConfig &mnConfig);
|
||||
|
||||
__aicore__ inline void SetMNConfig(const int32_t splitValue, const uint32_t groupIdx, MNConfig &mnConfig);
|
||||
|
||||
__aicore__ inline void SetMKN(const int32_t splitValue, const uint32_t groupIdx, MNConfig &mnConfig);
|
||||
|
||||
__aicore__ inline uint64_t GetWOffset(uint32_t tailN, uint32_t k);
|
||||
|
||||
__aicore__ inline void CubeProcess(MNConfig &mnConfig);
|
||||
__aicore__ inline void VecProcess(VecConfig &vecConfig);
|
||||
__aicore__ inline void MNBlockIdxCompute(MNConfig &mnConfig, const uint32_t curBlock, const uint32_t count,
|
||||
const uint32_t thresholdM_dimN);
|
||||
template <typename DTYPE_CS>
|
||||
__aicore__ inline void UpdateChannelScale(uint32_t loopidx, VecConfig &vecConfig);
|
||||
__aicore__ inline void VectorCompute(uint32_t loopidx, VecConfig &vecConfig);
|
||||
template <typename DTYPE_CS>
|
||||
__aicore__ inline void PreLoadTokenAndChannel(LocalTensor<float> &channelScaleLocal, VecConfig &vecConfig);
|
||||
__aicore__ inline void UpdateVecConfig(uint32_t blockIdx, VecConfig &vecConfig);
|
||||
__aicore__ inline void customDataCopyIn(uint32_t outLoopIdx, VecConfig &vecConfig);
|
||||
__aicore__ inline void customDataCopyOut(VecConfig &vecConfig);
|
||||
__aicore__ inline void Dequant(uint32_t loopidx, VecConfig &vecConfig);
|
||||
__aicore__ inline void Quant(uint32_t loopidx);
|
||||
__aicore__ inline void Swiglu(uint32_t loopidx);
|
||||
|
||||
private:
|
||||
typename mmType::MT &mm;
|
||||
const GMMSwigluBaseParams *__restrict gmmBaseParams;
|
||||
const GMMSwiglu *__restrict gmmSwiglu;
|
||||
const TCubeTiling *__restrict mmTilingData;
|
||||
TPipe *pipe;
|
||||
GlobalTensor<int8_t> xGM;
|
||||
GlobalTensor<int8_t> weightGM;
|
||||
GlobalTensor<CHANNELDTYPE> perChannelScaleGM;
|
||||
GlobalTensor<float> perTokenScaleGM;
|
||||
GlobalTensor<int64_t> groupListGM;
|
||||
GlobalTensor<int8_t> quantOutputGM;
|
||||
GlobalTensor<float> quantScaleOutputGM;
|
||||
GlobalTensor<int32_t> mmOutGM;
|
||||
// define the que
|
||||
TQue<QuePosition::VECIN, 1> mmOutQueue;
|
||||
TQue<QuePosition::VECIN, 1> perChannelScaleInQueue;
|
||||
TQue<QuePosition::VECOUT, 1> quantOutQueue;
|
||||
TQue<QuePosition::VECOUT, 1> quantScaleOutQueue;
|
||||
TBuf<TPosition::VECCALC> reduceWorkspace;
|
||||
uint32_t blockIdx = 0;
|
||||
int32_t preOffset = 0;
|
||||
int64_t aicCoreNum = 0;
|
||||
int64_t aivCoreNum = 0;
|
||||
float limited = FLOAT_INF;
|
||||
GM_ADDR xTensorPtr;
|
||||
GM_ADDR weightTensorPtr;
|
||||
};
|
||||
|
||||
template <typename mmType, bool sync, typename CHANNELDTYPE>
|
||||
__aicore__ inline void GMMSwigluCompute<mmType, sync, CHANNELDTYPE>::Init(
|
||||
GM_ADDR x, GM_ADDR weight, GM_ADDR perChannelScale, GM_ADDR perTokenScale, GM_ADDR groupList, GM_ADDR quantOutput,
|
||||
GM_ADDR quantScaleOutput, GM_ADDR workspace, const GMMSwigluBaseParams *__restrict gmmSwigluBaseParamsIn,
|
||||
const TCubeTiling *__restrict mmTilingDataIN, const GMMSwiglu *__restrict gmmSwigluIN, TPipe *tPipeIN)
|
||||
{
|
||||
aicCoreNum = GetBlockNum();
|
||||
aivCoreNum = aicCoreNum * 2;
|
||||
blockIdx = GetBlockIdx();
|
||||
mmTilingData = mmTilingDataIN;
|
||||
gmmBaseParams = gmmSwigluBaseParamsIn;
|
||||
gmmSwiglu = gmmSwigluIN;
|
||||
pipe = tPipeIN;
|
||||
xTensorPtr = x;
|
||||
limited = gmmBaseParams->limited;
|
||||
weightTensorPtr = weight;
|
||||
groupListGM.SetGlobalBuffer((__gm__ int64_t *)groupList, gmmSwiglu->groupListLen);
|
||||
mmOutGM.SetGlobalBuffer((__gm__ int32_t *)workspace, gmmBaseParams->M * gmmSwiglu->tokenLen);
|
||||
if ASCEND_IS_AIV {
|
||||
perChannelScaleGM.SetGlobalBuffer((__gm__ CHANNELDTYPE *)perChannelScale,
|
||||
gmmSwiglu->groupListLen * gmmSwiglu->tokenLen);
|
||||
perTokenScaleGM.SetGlobalBuffer((__gm__ float *)perTokenScale, gmmBaseParams->M);
|
||||
quantOutputGM.SetGlobalBuffer((__gm__ int8_t *)quantOutput,
|
||||
gmmBaseParams->M * gmmSwiglu->tokenLen / SWIGLU_REDUCE_FACTOR);
|
||||
quantScaleOutputGM.SetGlobalBuffer((__gm__ float *)quantScaleOutput, gmmBaseParams->M);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename mmType, bool sync, typename CHANNELDTYPE>
|
||||
__aicore__ inline void GMMSwigluCompute<mmType, sync, CHANNELDTYPE>::Process()
|
||||
{
|
||||
MNConfig mnConfig;
|
||||
VecConfig vecConfig;
|
||||
CubeProcess(mnConfig);
|
||||
VecProcess(vecConfig);
|
||||
}
|
||||
|
||||
template <typename mmType, bool sync, typename CHANNELDTYPE>
|
||||
template <typename DTYPE_CS>
|
||||
__aicore__ inline void
|
||||
GMMSwigluCompute<mmType, sync, CHANNELDTYPE>::PreLoadTokenAndChannel(LocalTensor<float> &channelScaleLocal,
|
||||
VecConfig &vecConfig)
|
||||
{
|
||||
DataCopyExtParams copyChannelParams{1, static_cast<uint32_t>(gmmSwiglu->tokenLen * sizeof(DTYPE_CS)), 0, 0, 0};
|
||||
DataCopyPadExtParams<DTYPE_CS> padParams{false, 0, 0, 0};
|
||||
if constexpr (!IsSameType<DTYPE_CS, float>::value) {
|
||||
LocalTensor<DTYPE_CS> dstLocalT = channelScaleLocal.template ReinterpretCast<DTYPE_CS>();
|
||||
DataCopyPad(dstLocalT[gmmSwiglu->tokenLen], perChannelScaleGM[vecConfig.curGroupIdx * gmmSwiglu->tokenLen],
|
||||
copyChannelParams, padParams);
|
||||
PipeBarrier<PIPE_ALL>();
|
||||
Cast(channelScaleLocal, dstLocalT[gmmSwiglu->tokenLen], RoundMode::CAST_NONE, gmmSwiglu->tokenLen);
|
||||
} else {
|
||||
DataCopyPad(channelScaleLocal, perChannelScaleGM[vecConfig.curGroupIdx * gmmSwiglu->tokenLen],
|
||||
copyChannelParams, padParams);
|
||||
}
|
||||
perChannelScaleInQueue.EnQue(channelScaleLocal);
|
||||
}
|
||||
|
||||
template <typename mmType, bool sync, typename CHANNELDTYPE>
|
||||
__aicore__ inline void GMMSwigluCompute<mmType, sync, CHANNELDTYPE>::MMCompute(uint32_t groupIdx, MNConfig &mnConfig,
|
||||
uint32_t coreIdx)
|
||||
{
|
||||
uint32_t tailN = mnConfig.nIdx * mnConfig.singleN;
|
||||
uint32_t curSingleN = mnConfig.nIdx < mnConfig.blockDimN - 1 ? mnConfig.singleN : mnConfig.n - tailN;
|
||||
uint32_t curSingleM =
|
||||
mnConfig.mIdx < mnConfig.blockDimM - 1 ? mnConfig.singleM : mnConfig.m - mnConfig.mIdx * mnConfig.singleM;
|
||||
uint64_t xOffset = mnConfig.mIdx * mnConfig.singleM * mnConfig.k;
|
||||
if constexpr (transposeX) {
|
||||
xOffset = mnConfig.mIdx * mnConfig.singleM;
|
||||
}
|
||||
uint64_t outOffset = mnConfig.mIdx * mnConfig.singleM * mnConfig.n + tailN;
|
||||
xGM.SetGlobalBuffer((__gm__ int8_t *)xTensorPtr + mnConfig.xBaseOffset);
|
||||
weightGM.SetGlobalBuffer((__gm__ int8_t *)weightTensorPtr + mnConfig.wBaseOffset + GetWOffset(tailN, mnConfig.k));
|
||||
if (mnConfig.blockDimM == 1) {
|
||||
weightGM.SetL2CacheHint(CacheMode::CACHE_MODE_DISABLE);
|
||||
}
|
||||
mnConfig.workSpaceOffset = outOffset + mnConfig.yBaseOffset;
|
||||
mm.SetOrgShape(mnConfig.m, mnConfig.n, mnConfig.k);
|
||||
mm.SetSingleShape(curSingleM, curSingleN, mnConfig.k);
|
||||
mm.SetTensorA(xGM[xOffset], transposeX);
|
||||
mm.SetTensorB(weightGM, transposeW);
|
||||
mm.template IterateAll<sync>(mmOutGM[mnConfig.workSpaceOffset], 0);
|
||||
}
|
||||
|
||||
template <typename mmType, bool sync, typename CHANNELDTYPE>
|
||||
__aicore__ inline void GMMSwigluCompute<mmType, sync, CHANNELDTYPE>::UpdateMnConfig(MNConfig &mnConfig)
|
||||
{
|
||||
if constexpr (B::format == CubeFormat::NZ) {
|
||||
mnConfig.wBaseOffset += AlignUp<16>(mnConfig.k) * AlignUp<32>(mnConfig.n); // 16: nz format last two dim size
|
||||
} else {
|
||||
mnConfig.wBaseOffset += mnConfig.k * mnConfig.n;
|
||||
}
|
||||
mnConfig.nAxisBaseOffset += mnConfig.n;
|
||||
mnConfig.mAxisBaseOffset += mnConfig.m;
|
||||
mnConfig.xBaseOffset += mnConfig.m * mnConfig.k;
|
||||
mnConfig.yBaseOffset += mnConfig.m * mnConfig.n;
|
||||
}
|
||||
|
||||
template <typename mmType, bool sync, typename CHANNELDTYPE>
|
||||
__aicore__ inline void GMMSwigluCompute<mmType, sync, CHANNELDTYPE>::SetMNConfig(const int32_t splitValue,
|
||||
const uint32_t groupIdx,
|
||||
MNConfig &mnConfig)
|
||||
{
|
||||
SetMKN(splitValue, groupIdx, mnConfig);
|
||||
mnConfig.baseM = BASIC_M;
|
||||
mnConfig.baseN = BASIC_N;
|
||||
mnConfig.singleM = SINGLE_CORE_M;
|
||||
mnConfig.singleN = SINGLE_CORE_N;
|
||||
}
|
||||
|
||||
template <typename mmType, bool sync, typename CHANNELDTYPE>
|
||||
__aicore__ inline void GMMSwigluCompute<mmType, sync, CHANNELDTYPE>::SetMKN(const int32_t splitValue,
|
||||
const uint32_t groupIdx, MNConfig &mnConfig)
|
||||
{
|
||||
mnConfig.m = static_cast<uint32_t>(splitValue);
|
||||
mnConfig.k = gmmBaseParams->K; // tilingData
|
||||
mnConfig.n = gmmBaseParams->N; // tilingData
|
||||
}
|
||||
|
||||
template <typename mmType, bool sync, typename CHANNELDTYPE>
|
||||
__aicore__ inline uint64_t GMMSwigluCompute<mmType, sync, CHANNELDTYPE>::GetWOffset(uint32_t tailN, uint32_t k)
|
||||
{
|
||||
uint64_t wOffset = 0;
|
||||
if constexpr (mmType::BT::format == CubeFormat::NZ) {
|
||||
wOffset = tailN * AlignUp<16>(k); // 16: nz format last two dim size
|
||||
} else {
|
||||
wOffset = tailN;
|
||||
}
|
||||
return wOffset;
|
||||
}
|
||||
|
||||
template <typename mmType, bool sync, typename CHANNELDTYPE>
|
||||
__aicore__ inline void GMMSwigluCompute<mmType, sync, CHANNELDTYPE>::CubeProcess(MNConfig &mnConfig)
|
||||
{
|
||||
if ASCEND_IS_AIC {
|
||||
preOffset = 0;
|
||||
int32_t prevSplitValue = 0;
|
||||
for (uint32_t groupIdx = 0, count = 0; groupIdx < gmmSwiglu->groupListLen; ++groupIdx) {
|
||||
UpdateMnConfig(mnConfig);
|
||||
int32_t currSplitValue = static_cast<int32_t>(groupListGM.GetValue(groupIdx));
|
||||
int32_t splitValue = currSplitValue - prevSplitValue;
|
||||
prevSplitValue = currSplitValue;
|
||||
SetMNConfig(splitValue, groupIdx, mnConfig);
|
||||
if (mnConfig.m <= 0 || mnConfig.k <= 0 || mnConfig.n <= 0) {
|
||||
continue;
|
||||
}
|
||||
mnConfig.blockDimM = Ceil(mnConfig.m, mnConfig.singleM);
|
||||
mnConfig.blockDimN = Ceil(mnConfig.n, mnConfig.singleN);
|
||||
|
||||
uint32_t curCount = count + mnConfig.blockDimM * mnConfig.blockDimN;
|
||||
uint32_t curBlock = blockIdx >= count ? blockIdx : blockIdx + gmmBaseParams->coreNum;
|
||||
uint32_t thresholdM_dimN = THRESHOLD_BLOCK_NUM * mnConfig.blockDimN;
|
||||
|
||||
while (curBlock < curCount) {
|
||||
MNBlockIdxCompute(mnConfig, curBlock, count, thresholdM_dimN);
|
||||
MMCompute(groupIdx, mnConfig, blockIdx);
|
||||
curBlock += aicCoreNum;
|
||||
}
|
||||
count = curCount % gmmBaseParams->coreNum;
|
||||
}
|
||||
SyncAll<false>();
|
||||
}
|
||||
}
|
||||
|
||||
template <typename mmType, bool sync, typename CHANNELDTYPE>
|
||||
__aicore__ inline void GMMSwigluCompute<mmType, sync, CHANNELDTYPE>::VecProcess(VecConfig &vecConfig)
|
||||
{
|
||||
if ASCEND_IS_AIV {
|
||||
UpdateVecConfig(blockIdx, vecConfig);
|
||||
if (blockIdx < vecConfig.usedCoreNum) {
|
||||
LocalTensor<float> channelScaleLocal = perChannelScaleInQueue.AllocTensor<float>();
|
||||
LocalTensor<int32_t> mmLocal = mmOutQueue.AllocTensor<int32_t>();
|
||||
LocalTensor<int8_t> quantLocal = quantOutQueue.AllocTensor<int8_t>();
|
||||
LocalTensor<float> quantScaleLocal = quantScaleOutQueue.AllocTensor<float>();
|
||||
mmOutQueue.EnQue(mmLocal);
|
||||
quantScaleOutQueue.EnQue(quantScaleLocal);
|
||||
quantOutQueue.EnQue(quantLocal);
|
||||
PreLoadTokenAndChannel<CHANNELDTYPE>(channelScaleLocal, vecConfig);
|
||||
}
|
||||
SyncAll<false>();
|
||||
if (blockIdx < vecConfig.usedCoreNum) {
|
||||
for (uint32_t outLoopIdx = 0; outLoopIdx < vecConfig.outLoopNum; outLoopIdx++) {
|
||||
vecConfig.innerLoopNum =
|
||||
outLoopIdx == (vecConfig.outLoopNum - 1) ? vecConfig.tailLoopNum : gmmSwiglu->maxProcessRowNum;
|
||||
customDataCopyIn(outLoopIdx, vecConfig);
|
||||
for (uint32_t innerLoopIdx = 0; innerLoopIdx < vecConfig.innerLoopNum; innerLoopIdx++) {
|
||||
UpdateChannelScale<CHANNELDTYPE>(innerLoopIdx, vecConfig);
|
||||
VectorCompute(innerLoopIdx, vecConfig);
|
||||
}
|
||||
customDataCopyOut(vecConfig);
|
||||
}
|
||||
|
||||
LocalTensor<float> channelScaleLocal = perChannelScaleInQueue.DeQue<float>();
|
||||
LocalTensor<int32_t> mmLocal = mmOutQueue.DeQue<int32_t>();
|
||||
LocalTensor<int8_t> quantLocal = quantOutQueue.DeQue<int8_t>();
|
||||
LocalTensor<float> quantScaleLocal = quantScaleOutQueue.DeQue<float>();
|
||||
perChannelScaleInQueue.FreeTensor(channelScaleLocal);
|
||||
mmOutQueue.FreeTensor(mmLocal);
|
||||
quantScaleOutQueue.FreeTensor(quantScaleLocal);
|
||||
quantOutQueue.FreeTensor(quantLocal);
|
||||
} else {
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename mmType, bool sync, typename CHANNELDTYPE>
|
||||
__aicore__ inline void
|
||||
GMMSwigluCompute<mmType, sync, CHANNELDTYPE>::MNBlockIdxCompute(MNConfig &mnConfig, const uint32_t curBlock,
|
||||
const uint32_t count, const uint32_t thresholdM_dimN)
|
||||
{
|
||||
mnConfig.mIdx = (curBlock - count) / mnConfig.blockDimN;
|
||||
mnConfig.nIdx = (curBlock - count) % mnConfig.blockDimN;
|
||||
}
|
||||
|
||||
template <typename mmType, bool sync, typename CHANNELDTYPE>
|
||||
__aicore__ inline void GMMSwigluCompute<mmType, sync, CHANNELDTYPE>::UpdateVecConfig(uint32_t blockIdx,
|
||||
VecConfig &vecConfig)
|
||||
{
|
||||
// 第一步 读取grouplist reduceSum 计算总数据个数
|
||||
int64_t prevM = 0;
|
||||
for (uint32_t groupIdx = 0; groupIdx < gmmSwiglu->groupListLen; groupIdx++) {
|
||||
int64_t currM = groupListGM.GetValue(groupIdx);
|
||||
int64_t tempM = currM - prevM;
|
||||
prevM = currM;
|
||||
vecConfig.M += tempM;
|
||||
}
|
||||
// 第二步 计算分核
|
||||
uint32_t eachCoreTaskNum = (vecConfig.M + aivCoreNum - 1) / aivCoreNum;
|
||||
vecConfig.usedCoreNum = vecConfig.M >= aivCoreNum ? aivCoreNum : vecConfig.M;
|
||||
uint32_t tailCoreIdx = vecConfig.M - (eachCoreTaskNum - 1) * vecConfig.usedCoreNum;
|
||||
vecConfig.taskNum = blockIdx < tailCoreIdx ? eachCoreTaskNum : eachCoreTaskNum - 1;
|
||||
vecConfig.startIdx =
|
||||
blockIdx < tailCoreIdx ? eachCoreTaskNum * blockIdx : ((eachCoreTaskNum - 1) * blockIdx + tailCoreIdx);
|
||||
vecConfig.curIdx = vecConfig.startIdx;
|
||||
vecConfig.startOffset = vecConfig.startIdx * gmmSwiglu->tokenLen;
|
||||
vecConfig.curOffset = vecConfig.startOffset;
|
||||
int64_t curStartIdx = vecConfig.startIdx;
|
||||
prevM = 0;
|
||||
for (uint32_t groupIdx = 0; groupIdx < gmmSwiglu->groupListLen; groupIdx++) {
|
||||
int64_t currM = groupListGM.GetValue(groupIdx);
|
||||
int64_t tempM = currM - prevM;
|
||||
prevM = currM;
|
||||
if (curStartIdx >= 0 && curStartIdx - tempM < 0) {
|
||||
vecConfig.curGroupIdx = groupIdx;
|
||||
vecConfig.nextUpadteInterVal = tempM - curStartIdx;
|
||||
}
|
||||
curStartIdx -= tempM;
|
||||
}
|
||||
// 第三步 计算总数据量
|
||||
vecConfig.outLoopNum = (vecConfig.taskNum + gmmSwiglu->maxProcessRowNum - 1) / gmmSwiglu->maxProcessRowNum;
|
||||
vecConfig.tailLoopNum = vecConfig.taskNum % gmmSwiglu->maxProcessRowNum ?
|
||||
vecConfig.taskNum % gmmSwiglu->maxProcessRowNum :
|
||||
gmmSwiglu->maxProcessRowNum;
|
||||
pipe->Reset();
|
||||
// 第四步 申请空间
|
||||
pipe->InitBuffer(mmOutQueue, DOUBLE_BUFFER, gmmSwiglu->maxProcessRowNum * gmmSwiglu->tokenLen * sizeof(int32_t));
|
||||
pipe->InitBuffer(perChannelScaleInQueue, DOUBLE_BUFFER, gmmSwiglu->tokenLen * sizeof(float));
|
||||
pipe->InitBuffer(quantOutQueue, DOUBLE_BUFFER,
|
||||
gmmSwiglu->maxProcessRowNum * gmmSwiglu->tokenLen / SWIGLU_REDUCE_FACTOR * sizeof(int8_t));
|
||||
pipe->InitBuffer(quantScaleOutQueue, DOUBLE_BUFFER,
|
||||
AlignUp<int32_t>(gmmSwiglu->maxProcessRowNum, ALIGN_8_ELE) * sizeof(float));
|
||||
// two 32 byte buffer for reduceMax calculation in Quant.
|
||||
pipe->InitBuffer(reduceWorkspace, gmmSwiglu->tokenLen / SWIGLU_REDUCE_FACTOR * sizeof(float) + UB_BLOCK_UNIT_SIZE +
|
||||
UB_BLOCK_UNIT_SIZE);
|
||||
}
|
||||
|
||||
template <typename mmType, bool sync, typename CHANNELDTYPE>
|
||||
__aicore__ inline void GMMSwigluCompute<mmType, sync, CHANNELDTYPE>::customDataCopyIn(uint32_t outLoopIdx,
|
||||
VecConfig &vecConfig)
|
||||
{
|
||||
LocalTensor<int32_t> _inMMLocal_0 = mmOutQueue.DeQue<int32_t>();
|
||||
DataCopyExtParams copyParams_0{
|
||||
1, static_cast<uint32_t>(vecConfig.innerLoopNum * gmmSwiglu->tokenLen * sizeof(int32_t)), 0, 0, 0};
|
||||
DataCopyPadExtParams<int32_t> padParams_0{false, 0, 0, 0};
|
||||
DataCopyPad(_inMMLocal_0, mmOutGM[vecConfig.curOffset], copyParams_0, padParams_0);
|
||||
|
||||
mmOutQueue.EnQue(_inMMLocal_0);
|
||||
|
||||
LocalTensor<int32_t> _inMMLocal_1 = mmOutQueue.DeQue<int32_t>();
|
||||
|
||||
Cast(_inMMLocal_1.ReinterpretCast<float>(), _inMMLocal_1, RoundMode::CAST_NONE,
|
||||
vecConfig.innerLoopNum * gmmSwiglu->tokenLen);
|
||||
|
||||
mmOutQueue.EnQue(_inMMLocal_1);
|
||||
LocalTensor<float> _inMMLocal_2 = mmOutQueue.DeQue<float>();
|
||||
SetFlag<HardEvent::S_V>(EVENT_ID0);
|
||||
for (uint32_t i = 0; i < vecConfig.innerLoopNum; i++) {
|
||||
WaitFlag<HardEvent::S_V>(EVENT_ID0);
|
||||
float scale = perTokenScaleGM.GetValue(vecConfig.curIdx);
|
||||
SetFlag<HardEvent::S_V>(EVENT_ID0);
|
||||
WaitFlag<HardEvent::S_V>(EVENT_ID0);
|
||||
Muls(_inMMLocal_2[i * gmmSwiglu->tokenLen], _inMMLocal_2[i * gmmSwiglu->tokenLen], scale, gmmSwiglu->tokenLen);
|
||||
SetFlag<HardEvent::S_V>(EVENT_ID0);
|
||||
vecConfig.curIdx++;
|
||||
}
|
||||
WaitFlag<HardEvent::S_V>(EVENT_ID0);
|
||||
vecConfig.curOffset = vecConfig.curIdx * gmmSwiglu->tokenLen;
|
||||
mmOutQueue.EnQue(_inMMLocal_2);
|
||||
}
|
||||
|
||||
template <typename mmType, bool sync, typename CHANNELDTYPE>
|
||||
template <typename DTYPE_CS>
|
||||
__aicore__ inline void GMMSwigluCompute<mmType, sync, CHANNELDTYPE>::UpdateChannelScale(uint32_t loopIdx,
|
||||
VecConfig &vecConfig)
|
||||
{
|
||||
// 更新perChannel
|
||||
if (unlikely(vecConfig.nextUpadteInterVal == 0)) {
|
||||
int64_t loop = gmmSwiglu->groupListLen - vecConfig.curGroupIdx;
|
||||
while (loop--) {
|
||||
int64_t curTemp = groupListGM.GetValue(vecConfig.curGroupIdx);
|
||||
vecConfig.curGroupIdx++;
|
||||
int64_t nextTemp = groupListGM.GetValue(vecConfig.curGroupIdx);
|
||||
if (nextTemp != curTemp) {
|
||||
vecConfig.nextUpadteInterVal = nextTemp - curTemp;
|
||||
break;
|
||||
}
|
||||
}
|
||||
LocalTensor<float> _inChannel = perChannelScaleInQueue.DeQue<float>();
|
||||
DataCopyExtParams copyParams{1, static_cast<uint32_t>(gmmSwiglu->tokenLen * sizeof(DTYPE_CS)), 0, 0, 0};
|
||||
DataCopyPadExtParams<DTYPE_CS> padParams{false, 0, 0, 0};
|
||||
if constexpr (!IsSameType<DTYPE_CS, float>::value) {
|
||||
LocalTensor<DTYPE_CS> dstLocalT = _inChannel.template ReinterpretCast<DTYPE_CS>();
|
||||
DataCopyPad(dstLocalT[gmmSwiglu->tokenLen], perChannelScaleGM[vecConfig.curGroupIdx * gmmSwiglu->tokenLen],
|
||||
copyParams, padParams);
|
||||
PipeBarrier<PIPE_ALL>();
|
||||
Cast(_inChannel, dstLocalT[gmmSwiglu->tokenLen], RoundMode::CAST_NONE, gmmSwiglu->tokenLen);
|
||||
} else {
|
||||
DataCopyPad(_inChannel, perChannelScaleGM[vecConfig.curGroupIdx * gmmSwiglu->tokenLen], copyParams,
|
||||
padParams);
|
||||
}
|
||||
PipeBarrier<PIPE_ALL>();
|
||||
perChannelScaleInQueue.EnQue(_inChannel);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename mmType, bool sync, typename CHANNELDTYPE>
|
||||
__aicore__ inline void GMMSwigluCompute<mmType, sync, CHANNELDTYPE>::VectorCompute(uint32_t loopIdx,
|
||||
VecConfig &vecConfig)
|
||||
{
|
||||
Dequant(loopIdx, vecConfig);
|
||||
Swiglu(loopIdx);
|
||||
Quant(loopIdx);
|
||||
}
|
||||
|
||||
template <typename mmType, bool sync, typename CHANNELDTYPE>
|
||||
__aicore__ inline void GMMSwigluCompute<mmType, sync, CHANNELDTYPE>::Dequant(uint32_t loopIdx, VecConfig &vecConfig)
|
||||
{
|
||||
// perChanelScale * perTokenScale
|
||||
LocalTensor<float> mmLocal = mmOutQueue.DeQue<float>();
|
||||
LocalTensor<float> perChannelLocal = perChannelScaleInQueue.DeQue<float>();
|
||||
Mul(mmLocal[loopIdx * gmmSwiglu->tokenLen], mmLocal[loopIdx * gmmSwiglu->tokenLen], perChannelLocal,
|
||||
gmmSwiglu->tokenLen);
|
||||
vecConfig.nextUpadteInterVal--;
|
||||
mmOutQueue.EnQue(mmLocal);
|
||||
perChannelScaleInQueue.EnQue(perChannelLocal);
|
||||
}
|
||||
|
||||
template <typename mmType, bool sync, typename CHANNELDTYPE>
|
||||
__aicore__ inline void GMMSwigluCompute<mmType, sync, CHANNELDTYPE>::Swiglu(uint32_t loopIdx)
|
||||
{
|
||||
// 高阶API swiglu
|
||||
LocalTensor<float> _inMMLocal = mmOutQueue.DeQue<float>();
|
||||
float beta = 1.0f;
|
||||
LocalTensor<float> workspaceLocal = reduceWorkspace.Get<float>();
|
||||
LocalTensor<float> src0Local =
|
||||
_inMMLocal[loopIdx * gmmSwiglu->tokenLen + gmmSwiglu->tokenLen / SWIGLU_REDUCE_FACTOR];
|
||||
LocalTensor<float> src1Local = _inMMLocal[loopIdx * gmmSwiglu->tokenLen];
|
||||
if (limited > 0.0f) {
|
||||
Mins(src0Local, src0Local, limited, gmmSwiglu->tokenLen / 2);
|
||||
PipeBarrier<PIPE_V>();
|
||||
Maxs(src0Local, src0Local, (-1.0f * limited), gmmSwiglu->tokenLen / 2);
|
||||
PipeBarrier<PIPE_V>();
|
||||
Mins(src1Local, src1Local, limited, gmmSwiglu->tokenLen / 2);
|
||||
PipeBarrier<PIPE_V>();
|
||||
}
|
||||
SwiGLU<float, false>(workspaceLocal, src0Local, src1Local, beta, gmmSwiglu->tokenLen / SWIGLU_REDUCE_FACTOR);
|
||||
PipeBarrier<PIPE_ALL>();
|
||||
DataCopyParams repeatParams{1, static_cast<uint16_t>((gmmSwiglu->tokenLen / SWIGLU_REDUCE_FACTOR) / ALIGN_8_ELE), 0,
|
||||
0};
|
||||
DataCopy(_inMMLocal[loopIdx * gmmSwiglu->tokenLen], workspaceLocal, repeatParams);
|
||||
mmOutQueue.EnQue(_inMMLocal);
|
||||
}
|
||||
|
||||
template <typename mmType, bool sync, typename CHANNELDTYPE>
|
||||
__aicore__ inline void GMMSwigluCompute<mmType, sync, CHANNELDTYPE>::Quant(uint32_t loopIdx)
|
||||
{
|
||||
LocalTensor<float> _inMMLocal = mmOutQueue.DeQue<float>();
|
||||
uint64_t preOffset = loopIdx * gmmSwiglu->tokenLen;
|
||||
uint64_t halfTokenLen = gmmSwiglu->tokenLen / BISECT;
|
||||
Abs(_inMMLocal[preOffset + gmmSwiglu->tokenLen / BISECT], _inMMLocal[preOffset], halfTokenLen);
|
||||
PipeBarrier<PIPE_V>();
|
||||
// reduceMax
|
||||
LocalTensor<float> workLocal = reduceWorkspace.Get<float>(halfTokenLen);
|
||||
LocalTensor<float> reduceResLocal =
|
||||
reduceWorkspace.GetWithOffset<float>(FLOAT_UB_BLOCK_UNIT_SIZE, halfTokenLen * sizeof(float));
|
||||
LocalTensor<float> reduceTmpLocal = reduceWorkspace.GetWithOffset<float>(
|
||||
FLOAT_UB_BLOCK_UNIT_SIZE, halfTokenLen * sizeof(float) + UB_BLOCK_UNIT_SIZE);
|
||||
ReduceMaxTemplate(reduceResLocal, workLocal, _inMMLocal[preOffset + gmmSwiglu->tokenLen / BISECT], reduceTmpLocal,
|
||||
static_cast<uint32_t>(halfTokenLen));
|
||||
|
||||
int32_t eventIdVToS = static_cast<int32_t>(GetTPipePtr()->FetchEventID(HardEvent::V_S));
|
||||
SetFlag<HardEvent::V_S>(eventIdVToS);
|
||||
WaitFlag<HardEvent::V_S>(eventIdVToS);
|
||||
float quantScale = reduceResLocal.GetValue(0) / QUANT_SCALE_INT8;
|
||||
LocalTensor<float> quantScaleLocal = quantScaleOutQueue.DeQue<float>();
|
||||
quantScaleLocal.SetValue(loopIdx, quantScale);
|
||||
quantScale = 1 / quantScale;
|
||||
int32_t eventIdSToV = static_cast<int32_t>(GetTPipePtr()->FetchEventID(HardEvent::S_V));
|
||||
SetFlag<HardEvent::S_V>(eventIdSToV);
|
||||
WaitFlag<HardEvent::S_V>(eventIdSToV);
|
||||
Muls(_inMMLocal[preOffset], _inMMLocal[preOffset], quantScale, halfTokenLen);
|
||||
PipeBarrier<PIPE_V>();
|
||||
LocalTensor<int8_t> quantLocal = quantOutQueue.DeQue<int8_t>();
|
||||
int32_t dstTempOffset = static_cast<int32_t>(preOffset / BISECT);
|
||||
int32_t srcTempOffset = static_cast<int32_t>(preOffset);
|
||||
int32_t tempCount = static_cast<int32_t>(halfTokenLen);
|
||||
LocalTensor<int8_t> castSpace = reduceWorkspace.Get<int8_t>(UB_BLOCK_UNIT_SIZE);
|
||||
CastFp32ToInt8Template(quantLocal, _inMMLocal, castSpace, dstTempOffset, srcTempOffset, tempCount);
|
||||
mmOutQueue.EnQue(_inMMLocal);
|
||||
quantOutQueue.EnQue(quantLocal);
|
||||
}
|
||||
|
||||
template <typename mmType, bool sync, typename CHANNELDTYPE>
|
||||
__aicore__ inline void GMMSwigluCompute<mmType, sync, CHANNELDTYPE>::customDataCopyOut(VecConfig &vecConfig)
|
||||
{
|
||||
// perChanelScale * perTokenScale
|
||||
LocalTensor<float> quantScaleLocal = quantScaleOutQueue.DeQue<float>();
|
||||
DataCopyParams copyParams_0{1, (uint16_t)(vecConfig.innerLoopNum * sizeof(float)), 0, 0};
|
||||
PipeBarrier<PIPE_ALL>();
|
||||
DataCopyPad(quantScaleOutputGM[vecConfig.startIdx], quantScaleLocal, copyParams_0);
|
||||
LocalTensor<int8_t> quantLocal = quantOutQueue.DeQue<int8_t>();
|
||||
DataCopyParams copyParams_1{
|
||||
1, (uint16_t)(vecConfig.innerLoopNum * gmmSwiglu->tokenLen / SWIGLU_REDUCE_FACTOR * sizeof(int8_t)), 0, 0};
|
||||
PipeBarrier<PIPE_ALL>();
|
||||
DataCopyPad(quantOutputGM[vecConfig.startIdx * gmmSwiglu->tokenLen / SWIGLU_REDUCE_FACTOR], quantLocal,
|
||||
copyParams_1);
|
||||
PipeBarrier<PIPE_ALL>();
|
||||
vecConfig.startIdx += vecConfig.innerLoopNum;
|
||||
vecConfig.startOffset = vecConfig.startIdx * gmmSwiglu->tokenLen;
|
||||
quantOutQueue.EnQue(quantLocal);
|
||||
quantScaleOutQueue.EnQue(quantScaleLocal);
|
||||
}
|
||||
|
||||
} // namespace GROUPED_MATMUL_SWIGLU_QUANT
|
||||
#endif // ASCENDC_GROUPED_MATMUL_QUANT_MIXCORE_H
|
||||
@@ -0,0 +1,251 @@
|
||||
/**
|
||||
* Copyright (c) 2025 Huawei Technologies Co., Ltd.
|
||||
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
* CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
* See LICENSE in the root of the software repository for the full text of the License.
|
||||
*/
|
||||
|
||||
/*!
|
||||
* \file grouped_matmul_antiquant_a8w4_msd.h
|
||||
* \brief
|
||||
*/
|
||||
|
||||
#ifndef ASCENDC_GROUPED_MATMUL_SWIGLU_QUANT_A8W4_MSD_MID_H
|
||||
#define ASCENDC_GROUPED_MATMUL_SWIGLU_QUANT_A8W4_MSD_MID_H
|
||||
|
||||
#include "grouped_matmul_swiglu_quant_utils.h"
|
||||
#include "grouped_matmul_swiglu_quant.h"
|
||||
|
||||
#ifdef GMM_SWIGLU_QUANT_A8W4_MSD
|
||||
namespace GROUPED_MATMUL_SWIGLU_QUANT {
|
||||
using namespace matmul;
|
||||
using namespace AscendC;
|
||||
|
||||
constexpr uint32_t BUFFER_NUM = 1;
|
||||
|
||||
template <typename T>
|
||||
__aicore__ inline void DataCopyPad2DA8W4(const LocalTensor<T> dst, const GlobalTensor<T> src, uint32_t dim1,
|
||||
uint32_t dim0, uint32_t srcDim0)
|
||||
{
|
||||
DataCopyExtParams params;
|
||||
params.blockCount = dim1;
|
||||
params.blockLen = dim0 * sizeof(T);
|
||||
params.srcStride = (srcDim0 - dim0) * sizeof(T);
|
||||
// 32: int32 -> float16, 为防止跨行数据进入同一32B block，提前每行按偶数block对齐
|
||||
params.dstStride = Ceil(dim0 * sizeof(T), 32) % 2;
|
||||
|
||||
DataCopyPadExtParams<T> padParams{true, 0, 0, 0};
|
||||
DataCopyPad(dst, src, params, padParams);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__aicore__ inline void DataCopyPad2DA8W4ND(const LocalTensor<T> dst, const GlobalTensor<T> src, uint32_t dim1,
|
||||
uint32_t dim0, uint32_t srcDim0)
|
||||
{
|
||||
DataCopyExtParams params;
|
||||
params.blockCount = dim1;
|
||||
params.blockLen = dim0 * sizeof(T);
|
||||
params.srcStride = (srcDim0 - dim0) * sizeof(T);
|
||||
params.dstStride = 0;
|
||||
|
||||
DataCopyPadExtParams<T> padParams{true, 0, 0, 0};
|
||||
DataCopyPad(dst, src, params, padParams);
|
||||
return;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__aicore__ inline void DataCopyPad2DA8W4(const GlobalTensor<T> dst, const LocalTensor<T> src, uint32_t dim1,
|
||||
uint32_t dim0, uint32_t srcDim0, uint32_t dstDim0)
|
||||
{
|
||||
DataCopyExtParams params;
|
||||
params.blockCount = dim1;
|
||||
params.blockLen = dim0 * sizeof(T);
|
||||
// 32: ub访问粒度为32B
|
||||
params.srcStride = (srcDim0 - dim0) * sizeof(T) / 32;
|
||||
params.dstStride = (dstDim0 - dim0) * sizeof(T);
|
||||
DataCopyPad(dst, src, params);
|
||||
}
|
||||
|
||||
template <class mmType>
|
||||
class GMMA8W4MidProcess {
|
||||
public:
|
||||
using bT = typename mmType::BT;
|
||||
|
||||
public:
|
||||
__aicore__ inline GMMA8W4MidProcess(typename mmType::MT &matmul) : mm(matmul)
|
||||
{
|
||||
}
|
||||
__aicore__ inline void Init(const GMAddrParams gmAddrParams,
|
||||
const GMMSwigluBaseParams *__restrict gmmSwigluBaseParamsIN);
|
||||
__aicore__ inline void Process(WorkSpaceSplitConfig &workspaceSplitConfig, int64_t workspaceSplitLoopIdx);
|
||||
|
||||
private:
|
||||
__aicore__ inline void MMCompute(uint32_t groupIdx, MNConfig &mnConfig, WorkSpaceSplitConfig &workspaceSplitConfig);
|
||||
__aicore__ inline void SetMNConfig(const int32_t splitValue, MNConfig &mnConfig);
|
||||
__aicore__ inline void UpdateMnConfig(MNConfig &mnConfig);
|
||||
|
||||
private:
|
||||
typename mmType::MT &mm;
|
||||
const uint32_t HALF_ALIGN = 16;
|
||||
GlobalTensor<int4b_t> xGM;
|
||||
GlobalTensor<int4b_t> xGM1;
|
||||
GlobalTensor<int4b_t> xGM2;
|
||||
GlobalTensor<int4b_t> weightGM;
|
||||
|
||||
GlobalTensor<half> mmOutGM;
|
||||
GlobalTensor<half> mmOutGM1;
|
||||
GlobalTensor<half> mmOutGM2;
|
||||
GlobalTensor<int64_t> groupListGM;
|
||||
GlobalTensor<uint64_t> weightScaleGM;
|
||||
// define the que
|
||||
uint32_t subBlockIdx = 0;
|
||||
uint32_t coreIdx = 0;
|
||||
uint32_t quantGroupSize = 0;
|
||||
uint32_t vecCount = 0;
|
||||
uint32_t xRowSumCount = 0;
|
||||
const GMMSwigluBaseParams *__restrict gmmBaseParams;
|
||||
};
|
||||
|
||||
template <typename mmType>
|
||||
__aicore__ inline void GMMA8W4MidProcess<mmType>::Init(const GMAddrParams gmAddrParams,
|
||||
const GMMSwigluBaseParams *__restrict gmmSwigluBaseParamsIN)
|
||||
{
|
||||
if ASCEND_IS_AIC {
|
||||
gmmBaseParams = gmmSwigluBaseParamsIN;
|
||||
xRowSumCount = gmmBaseParams->M;
|
||||
xGM1.SetGlobalBuffer((__gm__ int4b_t *)gmAddrParams.workSpaceGM); // 从前处理中获得的结果
|
||||
xGM2.SetGlobalBuffer(
|
||||
(__gm__ int4b_t *)((__gm__ int8_t *)gmAddrParams.workSpaceGM + gmAddrParams.workSpaceOffset1));
|
||||
weightGM.SetGlobalBuffer((__gm__ int4b_t *)gmAddrParams.weightGM);
|
||||
weightScaleGM.SetGlobalBuffer((__gm__ uint64_t *)gmAddrParams.weightScaleGM);
|
||||
groupListGM.SetGlobalBuffer((__gm__ int64_t *)gmAddrParams.groupListGM);
|
||||
mmOutGM1.SetGlobalBuffer(
|
||||
(__gm__ half *)((__gm__ int8_t *)gmAddrParams.workSpaceGM + gmAddrParams.workSpaceOffset2));
|
||||
mmOutGM2.SetGlobalBuffer(
|
||||
(__gm__ half *)((__gm__ int8_t *)gmAddrParams.workSpaceGM + gmAddrParams.workSpaceOffset3));
|
||||
quantGroupSize = gmmBaseParams->K / gmmBaseParams->quantGroupNum; // 约束为整除关系
|
||||
subBlockIdx = GetSubBlockIdx();
|
||||
coreIdx = GetBlockIdx();
|
||||
}
|
||||
}
|
||||
|
||||
template <typename mmType>
|
||||
__aicore__ inline void GMMA8W4MidProcess<mmType>::UpdateMnConfig(MNConfig &mnConfig)
|
||||
{
|
||||
if constexpr (bT::format == CubeFormat::NZ) {
|
||||
mnConfig.wBaseOffset += AlignUp<16>(mnConfig.k) * AlignUp<32>(mnConfig.n); // 16: nz format last two dim size
|
||||
} else {
|
||||
mnConfig.wBaseOffset += mnConfig.k * mnConfig.n;
|
||||
}
|
||||
mnConfig.nAxisBaseOffset += mnConfig.n;
|
||||
mnConfig.mAxisBaseOffset += mnConfig.m;
|
||||
mnConfig.xBaseOffset += mnConfig.m * mnConfig.k;
|
||||
mnConfig.yBaseOffset += mnConfig.m * mnConfig.n;
|
||||
}
|
||||
|
||||
template <typename mmType>
|
||||
__aicore__ inline void GMMA8W4MidProcess<mmType>::SetMNConfig(const int32_t splitValue, MNConfig &mnConfig)
|
||||
{
|
||||
mnConfig.m = static_cast<int64_t>(splitValue);
|
||||
mnConfig.baseM = gmmBaseParams->baseM;
|
||||
mnConfig.baseN = gmmBaseParams->baseN;
|
||||
mnConfig.singleM = gmmBaseParams->baseM;
|
||||
mnConfig.singleN = gmmBaseParams->baseN;
|
||||
}
|
||||
|
||||
template <typename mmType>
|
||||
__aicore__ inline void GMMA8W4MidProcess<mmType>::Process(WorkSpaceSplitConfig &workspaceSplitConfig,
|
||||
int64_t workspaceSplitLoopIdx)
|
||||
{
|
||||
if ASCEND_IS_AIC {
|
||||
if (workspaceSplitLoopIdx >= workspaceSplitConfig.loopCount || workspaceSplitLoopIdx < 0) {
|
||||
return;
|
||||
}
|
||||
xGM = (workspaceSplitLoopIdx % 2 == 0 ? xGM1 : xGM2);
|
||||
mmOutGM = (workspaceSplitLoopIdx % 2 == 0 ? mmOutGM1 : mmOutGM2);
|
||||
MNConfig mnConfig;
|
||||
mnConfig.baseM = gmmBaseParams->baseM;
|
||||
mnConfig.baseN = gmmBaseParams->baseN;
|
||||
mnConfig.singleM = gmmBaseParams->baseM;
|
||||
mnConfig.singleN = gmmBaseParams->baseN;
|
||||
mnConfig.k = gmmBaseParams->K; // tilingData
|
||||
mnConfig.n = gmmBaseParams->N; // tilingData
|
||||
mnConfig.blockDimN = Ceil(mnConfig.n, mnConfig.singleN);
|
||||
int32_t prevSplitValue = workspaceSplitLoopIdx * workspaceSplitConfig.notLastTaskSize;
|
||||
for (uint32_t groupIdx = workspaceSplitConfig.rightMatrixExpertStartIndex, preCount = 0;
|
||||
groupIdx <= workspaceSplitConfig.rightMatrixExpertEndIndex; ++groupIdx) {
|
||||
UpdateMnConfig(mnConfig);
|
||||
int32_t currSplitValue = static_cast<int32_t>(groupListGM.GetValue(groupIdx));
|
||||
currSplitValue = currSplitValue > (workspaceSplitLoopIdx + 1) * gmmBaseParams->mLimit ?
|
||||
(workspaceSplitLoopIdx + 1) * gmmBaseParams->mLimit :
|
||||
currSplitValue;
|
||||
|
||||
int32_t splitValue = (currSplitValue - prevSplitValue) * 2; // 2: int8 has been split in 2 int4
|
||||
prevSplitValue = currSplitValue;
|
||||
|
||||
SetMNConfig(splitValue, mnConfig);
|
||||
if (mnConfig.m <= 0 || mnConfig.k <= 0 || mnConfig.n <= 0) {
|
||||
continue;
|
||||
}
|
||||
mnConfig.blockDimM = Ceil(mnConfig.m, mnConfig.singleM);
|
||||
mm.SetOrgShape(mnConfig.m, mnConfig.n, mnConfig.k);
|
||||
uint32_t curCount = preCount + mnConfig.blockDimN * mnConfig.blockDimM;
|
||||
uint32_t curBlock = coreIdx >= preCount ? coreIdx : coreIdx + gmmBaseParams->coreNum;
|
||||
while (curBlock < curCount) {
|
||||
mnConfig.mIdx = (curBlock - preCount) / mnConfig.blockDimN;
|
||||
mnConfig.nIdx = (curBlock - preCount) % mnConfig.blockDimN;
|
||||
MMCompute(groupIdx, mnConfig, workspaceSplitConfig);
|
||||
curBlock += gmmBaseParams->coreNum;
|
||||
}
|
||||
preCount = curCount % gmmBaseParams->coreNum;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename mmType>
|
||||
__aicore__ inline void GMMA8W4MidProcess<mmType>::MMCompute(uint32_t groupIdx, MNConfig &mnConfig,
|
||||
WorkSpaceSplitConfig &workspaceSplitConfig)
|
||||
{
|
||||
uint32_t tailN = mnConfig.nIdx * mnConfig.singleN;
|
||||
uint32_t curSingleN = mnConfig.singleN;
|
||||
if (unlikely(mnConfig.nIdx == mnConfig.blockDimN - 1)) {
|
||||
curSingleN = gmmBaseParams->N - tailN;
|
||||
}
|
||||
uint32_t curSingleM = mnConfig.singleM;
|
||||
if (unlikely(mnConfig.mIdx == mnConfig.blockDimM - 1)) {
|
||||
curSingleM = mnConfig.m - mnConfig.mIdx * mnConfig.singleM;
|
||||
}
|
||||
uint64_t weightOffset = 0;
|
||||
if constexpr (mmType::BT::format == CubeFormat::NZ) {
|
||||
weightOffset = static_cast<uint64_t>(groupIdx) * gmmBaseParams->N * gmmBaseParams->K + tailN * gmmBaseParams->K;
|
||||
} else {
|
||||
weightOffset = static_cast<uint64_t>(groupIdx) * gmmBaseParams->N * gmmBaseParams->K + tailN;
|
||||
}
|
||||
mm.SetSingleShape(curSingleM, curSingleN, quantGroupSize); // 8, 256, 512 --> 514us
|
||||
GlobalTensor<int4b_t> weightSlice;
|
||||
uint64_t outOffset = mnConfig.mIdx * mnConfig.singleM * mnConfig.n + tailN;
|
||||
mnConfig.workSpaceOffset = outOffset + mnConfig.yBaseOffset;
|
||||
for (uint32_t loopK = 0; loopK < gmmBaseParams->quantGroupNum; loopK++) {
|
||||
mm.SetTensorA(
|
||||
xGM[mnConfig.xBaseOffset + mnConfig.mIdx * mnConfig.k * mnConfig.singleM + loopK * quantGroupSize]);
|
||||
if constexpr (mmType::BT::format == CubeFormat::NZ) {
|
||||
weightSlice = weightGM[weightOffset + loopK * quantGroupSize * 64];
|
||||
} else {
|
||||
weightSlice = weightGM[weightOffset + loopK * quantGroupSize * gmmBaseParams->N];
|
||||
}
|
||||
if (mnConfig.blockDimM == 1) {
|
||||
weightSlice.SetL2CacheHint(CacheMode::CACHE_MODE_DISABLE);
|
||||
}
|
||||
mm.SetTensorB(weightSlice);
|
||||
mm.SetQuantVector(weightScaleGM[groupIdx * gmmBaseParams->N * gmmBaseParams->quantGroupNum +
|
||||
loopK * gmmBaseParams->N + tailN]);
|
||||
mm.Iterate();
|
||||
mm.GetTensorC(mmOutGM[mnConfig.workSpaceOffset], loopK == 0 ? 0 : 1);
|
||||
}
|
||||
}
|
||||
} // namespace GROUPED_MATMUL_SWIGLU_QUANT
|
||||
#endif // GMM_SWIGLU_QUANT_A8W4_MSD
|
||||
#endif // ASCENDC_GROUPED_MATMUL_SWIGLU_QUANT_A8W4_MSD_MID_H
|
||||
@@ -0,0 +1,385 @@
|
||||
/**
|
||||
* Copyright (c) 2025 Huawei Technologies Co., Ltd.
|
||||
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
* CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
* See LICENSE in the root of the software repository for the full text of the License.
|
||||
*/
|
||||
|
||||
/*!
|
||||
* \file grouped_matmul_swiglu_quant_a8w4_msd_post.h
|
||||
* \brief
|
||||
*/
|
||||
|
||||
#ifndef ASCENDC_GROUPED_MATMUL_SWIGLU_QUANT_A8W4_MSD_POST_H
|
||||
#define ASCENDC_GROUPED_MATMUL_SWIGLU_QUANT_A8W4_MSD_POST_H
|
||||
#include "grouped_matmul_swiglu_quant_utils.h"
|
||||
#include "kernel_operator.h"
|
||||
#ifdef GMM_SWIGLU_QUANT_A8W4_MSD
|
||||
namespace GROUPED_MATMUL_SWIGLU_QUANT {
|
||||
using namespace AscendC;
|
||||
#define DOUBLE_BUFFER 2
|
||||
constexpr float DEFAULT_MUL_SCALE = 16.0f;
|
||||
class GMMA8W4PostProcess {
|
||||
public:
|
||||
__aicore__ inline GMMA8W4PostProcess(){};
|
||||
__aicore__ inline void Init(const GMAddrParams gmAddrParams,
|
||||
const GMMSwigluBaseParams *__restrict gmmSwigluBaseParamsIN,
|
||||
const GMMSwiglu *__restrict gmmSwigluIN);
|
||||
|
||||
__aicore__ inline void Process(WorkSpaceSplitConfig &workspaceSplitConfig, int64_t workspaceSplitLoopIdx,
|
||||
TPipe *pipe);
|
||||
|
||||
private:
|
||||
__aicore__ inline void UpdateVecConfig(uint32_t blockIdx, VecConfig &vecConfig,
|
||||
WorkSpaceSplitConfig &workspaceSplitConfig, int64_t workspaceSplitLoopIdx,
|
||||
TPipe *pipe);
|
||||
|
||||
__aicore__ inline void UpdateAuxiliaryMatrix(uint32_t loopIdx, VecConfig &vecConfig);
|
||||
|
||||
__aicore__ inline void VectorCompute(uint32_t loopIdx, VecConfig &vecConfig,
|
||||
WorkSpaceSplitConfig &workspaceSplitConfig);
|
||||
|
||||
__aicore__ inline void customDataCopyIn(uint32_t outLoopIdx, GlobalTensor<half> &mmOutGM, VecConfig &vecConfig,
|
||||
WorkSpaceSplitConfig &workspaceSplitConfig);
|
||||
|
||||
__aicore__ inline void customDataCopyOut(VecConfig &vecConfig, WorkSpaceSplitConfig &workspaceSplitConfig);
|
||||
|
||||
__aicore__ inline void PreLoadAuxiliaryMatrix(VecConfig &vecConfig);
|
||||
|
||||
__aicore__ inline void Quant(uint32_t loopIdx, VecConfig &vecConfig);
|
||||
|
||||
__aicore__ inline void Swiglu(uint32_t loopIdx, VecConfig &vecConfig);
|
||||
|
||||
__aicore__ inline void MergeAuxiliaryMatrix(uint32_t loopIdx, VecConfig &vecConfig);
|
||||
|
||||
__aicore__ inline void MulPertokenScale(uint32_t loopIdx, VecConfig &vecConfig,
|
||||
WorkSpaceSplitConfig &workspaceSplitConfig);
|
||||
const GMMSwiglu *__restrict gmmSwiglu;
|
||||
const GMMSwigluBaseParams *__restrict gmmBaseParams;
|
||||
GlobalTensor<float> perTokenScaleGM;
|
||||
GlobalTensor<int64_t> groupListGM;
|
||||
GlobalTensor<int8_t> quantOutputGM;
|
||||
GlobalTensor<float> weightAuxiliaryMatrixGM;
|
||||
GlobalTensor<float> quantScaleOutputGM;
|
||||
GlobalTensor<half> mmOutGM1;
|
||||
GlobalTensor<half> mmOutGM2;
|
||||
GlobalTensor<half> mmOutGM;
|
||||
LocalTensor<float> mmLocal_fp32;
|
||||
LocalTensor<half> mmLocal_fp16;
|
||||
TQue<QuePosition::VECIN, 1> weightAuxiliaryMatrixInQueue;
|
||||
TQue<QuePosition::VECIN, 1> mmOutQueue;
|
||||
TQue<QuePosition::VECOUT, 1> quantOutQueue;
|
||||
TQue<QuePosition::VECOUT, 1> quantScaleOutQueue;
|
||||
TBuf<TPosition::VECCALC> reduceWorkspace;
|
||||
uint32_t blockIdx = 0;
|
||||
int64_t aicCoreNum = 0;
|
||||
int64_t aivCoreNum = 0;
|
||||
};
|
||||
|
||||
__aicore__ inline void GMMA8W4PostProcess::Init(const GMAddrParams gmAddrParams,
|
||||
const GMMSwigluBaseParams *__restrict gmmSwigluBaseParamsIN,
|
||||
const GMMSwiglu *__restrict gmmSwigluIN)
|
||||
{
|
||||
if ASCEND_IS_AIV {
|
||||
aicCoreNum = GetBlockNum();
|
||||
aivCoreNum = aicCoreNum * 2;
|
||||
blockIdx = GetBlockIdx();
|
||||
gmmBaseParams = gmmSwigluBaseParamsIN;
|
||||
gmmSwiglu = gmmSwigluIN;
|
||||
weightAuxiliaryMatrixGM.SetGlobalBuffer((__gm__ float *)gmAddrParams.weightAuxiliaryMatrixGM); // E, N
|
||||
groupListGM.SetGlobalBuffer((__gm__ int64_t *)gmAddrParams.groupListGM, gmmSwiglu->groupListLen);
|
||||
mmOutGM1.SetGlobalBuffer(
|
||||
(__gm__ half *)((__gm__ int8_t *)gmAddrParams.workSpaceGM + gmAddrParams.workSpaceOffset2));
|
||||
mmOutGM2.SetGlobalBuffer(
|
||||
(__gm__ half *)((__gm__ int8_t *)gmAddrParams.workSpaceGM + gmAddrParams.workSpaceOffset3));
|
||||
perTokenScaleGM.SetGlobalBuffer((__gm__ float *)gmAddrParams.xScaleGM, gmmBaseParams->M);
|
||||
quantOutputGM.SetGlobalBuffer((__gm__ int8_t *)gmAddrParams.yGM,
|
||||
gmmBaseParams->M * gmmSwiglu->tokenLen / SWIGLU_REDUCE_FACTOR);
|
||||
quantScaleOutputGM.SetGlobalBuffer((__gm__ float *)gmAddrParams.yScaleGM, gmmBaseParams->M);
|
||||
}
|
||||
}
|
||||
|
||||
__aicore__ inline void GMMA8W4PostProcess::customDataCopyIn(uint32_t outLoopIdx, GlobalTensor<half> &mmOutGM,
|
||||
VecConfig &vecConfig,
|
||||
WorkSpaceSplitConfig &workspaceSplitConfig)
|
||||
{
|
||||
mmLocal_fp16 = mmOutQueue.DeQue<half>();
|
||||
mmLocal_fp32 = mmLocal_fp16.ReinterpretCast<float>();
|
||||
const int64_t processNum = 2 * vecConfig.innerLoopNum * gmmSwiglu->tokenLen;
|
||||
DataCopyExtParams copyParams_0{1, static_cast<uint32_t>(processNum * sizeof(half)), 0, 0, 0};
|
||||
DataCopyPadExtParams<half> padParams_0{false, 0, 0, 0};
|
||||
DataCopyPad(mmLocal_fp16[processNum], mmOutGM[vecConfig.curOffset * DOUBLE_ROW], copyParams_0, padParams_0);
|
||||
|
||||
mmOutQueue.EnQue(mmLocal_fp16);
|
||||
mmLocal_fp16 = mmOutQueue.DeQue<half>();
|
||||
// 1. fp16 -> fp32
|
||||
Cast(mmLocal_fp32, mmLocal_fp16[processNum], RoundMode::CAST_NONE, processNum);
|
||||
PipeBarrier<PIPE_V>();
|
||||
int32_t eventIdSToV = static_cast<int32_t>(GetTPipePtr()->FetchEventID(HardEvent::S_V));
|
||||
// 2. high_4bit * 16 + low_4bit
|
||||
for (uint32_t i = 0; i < vecConfig.innerLoopNum; i++) {
|
||||
Muls(mmLocal_fp32[(DOUBLE_ROW * i) * gmmSwiglu->tokenLen], mmLocal_fp32[(DOUBLE_ROW * i) * gmmSwiglu->tokenLen],
|
||||
DEFAULT_MUL_SCALE, gmmSwiglu->tokenLen);
|
||||
PipeBarrier<PIPE_V>();
|
||||
Add(mmLocal_fp32[i * gmmSwiglu->tokenLen], mmLocal_fp32[(DOUBLE_ROW * i) * gmmSwiglu->tokenLen],
|
||||
mmLocal_fp32[(DOUBLE_ROW * i + 1) * gmmSwiglu->tokenLen], gmmSwiglu->tokenLen);
|
||||
PipeBarrier<PIPE_V>();
|
||||
vecConfig.curIdx++;
|
||||
}
|
||||
vecConfig.curOffset = vecConfig.curIdx * gmmSwiglu->tokenLen;
|
||||
PipeBarrier<PIPE_V>();
|
||||
}
|
||||
|
||||
__aicore__ inline void GMMA8W4PostProcess::VectorCompute(uint32_t loopIdx, VecConfig &vecConfig,
|
||||
WorkSpaceSplitConfig &workspaceSplitConfig)
|
||||
{
|
||||
// 1.辅助矩阵加回
|
||||
MergeAuxiliaryMatrix(loopIdx, vecConfig);
|
||||
// 2.perToken反量化
|
||||
MulPertokenScale(loopIdx, vecConfig, workspaceSplitConfig);
|
||||
// 3.Swiglu
|
||||
Swiglu(loopIdx, vecConfig);
|
||||
// 4.Quant
|
||||
Quant(loopIdx, vecConfig);
|
||||
}
|
||||
|
||||
__aicore__ inline void GMMA8W4PostProcess::MergeAuxiliaryMatrix(uint32_t loopIdx, VecConfig &vecConfig)
|
||||
{
|
||||
// perChanelScale * perTokenScale
|
||||
mmLocal_fp32 = mmOutQueue.DeQue<float>();
|
||||
LocalTensor<float> weightAuxiliaryMatrixLocal = weightAuxiliaryMatrixInQueue.DeQue<float>();
|
||||
Add(mmLocal_fp32[loopIdx * gmmSwiglu->tokenLen], mmLocal_fp32[loopIdx * gmmSwiglu->tokenLen], weightAuxiliaryMatrixLocal,
|
||||
gmmSwiglu->tokenLen);
|
||||
vecConfig.nextUpadteInterVal--;
|
||||
PipeBarrier<PIPE_V>();
|
||||
weightAuxiliaryMatrixInQueue.EnQue(weightAuxiliaryMatrixLocal);
|
||||
}
|
||||
|
||||
__aicore__ inline void GMMA8W4PostProcess::MulPertokenScale(uint32_t loopIdx, VecConfig &vecConfig,
|
||||
WorkSpaceSplitConfig &workspaceSplitConfig)
|
||||
{
|
||||
float scale = perTokenScaleGM.GetValue(loopIdx + workspaceSplitConfig.leftMatrixStartIndex + vecConfig.startIdx);
|
||||
int32_t eventIdSToV = static_cast<int32_t>(GetTPipePtr()->FetchEventID(HardEvent::S_V));
|
||||
SetFlag<HardEvent::S_V>(eventIdSToV);
|
||||
WaitFlag<HardEvent::S_V>(eventIdSToV);
|
||||
Muls(mmLocal_fp32[loopIdx * gmmSwiglu->tokenLen], mmLocal_fp32[loopIdx * gmmSwiglu->tokenLen], scale, gmmSwiglu->tokenLen);
|
||||
PipeBarrier<PIPE_V>();
|
||||
}
|
||||
|
||||
__aicore__ inline void GMMA8W4PostProcess::Swiglu(uint32_t loopIdx, VecConfig &vecConfig)
|
||||
{
|
||||
// 高阶API swiglu
|
||||
float beta = 1.0f;
|
||||
LocalTensor<float> workspaceLocal = reduceWorkspace.Get<float>();
|
||||
LocalTensor<float> src0Local =
|
||||
mmLocal_fp32[loopIdx * gmmSwiglu->tokenLen + gmmSwiglu->tokenLen / SWIGLU_REDUCE_FACTOR];
|
||||
LocalTensor<float> src1Local = mmLocal_fp32[loopIdx * gmmSwiglu->tokenLen];
|
||||
|
||||
SwiGLU<float, false>(workspaceLocal, src0Local, src1Local, beta, gmmSwiglu->tokenLen / SWIGLU_REDUCE_FACTOR);
|
||||
PipeBarrier<PIPE_V>();
|
||||
DataCopyParams repeatParams{1, static_cast<uint16_t>((gmmSwiglu->tokenLen / SWIGLU_REDUCE_FACTOR) / ALIGN_8_ELE), 0,
|
||||
0};
|
||||
DataCopy(mmLocal_fp32[loopIdx * gmmSwiglu->tokenLen], workspaceLocal, repeatParams);
|
||||
|
||||
PipeBarrier<PIPE_V>();
|
||||
}
|
||||
|
||||
__aicore__ inline void GMMA8W4PostProcess::Quant(uint32_t loopIdx, VecConfig &vecConfig)
|
||||
{
|
||||
uint64_t preOffset = loopIdx * gmmSwiglu->tokenLen;
|
||||
uint64_t halfTokenLen = gmmSwiglu->tokenLen / BISECT;
|
||||
Abs(mmLocal_fp32[preOffset + gmmSwiglu->tokenLen / BISECT], mmLocal_fp32[preOffset], halfTokenLen);
|
||||
PipeBarrier<PIPE_V>();
|
||||
// reduceMax
|
||||
LocalTensor<float> workLocal = reduceWorkspace.Get<float>(halfTokenLen);
|
||||
LocalTensor<float> reduceResLocal =
|
||||
reduceWorkspace.GetWithOffset<float>(FLOAT_UB_BLOCK_UNIT_SIZE, halfTokenLen * sizeof(float));
|
||||
LocalTensor<float> reduceTmpLocal = reduceWorkspace.GetWithOffset<float>(
|
||||
FLOAT_UB_BLOCK_UNIT_SIZE, halfTokenLen * sizeof(float) + UB_BLOCK_UNIT_SIZE);
|
||||
ReduceMaxTemplate(reduceResLocal, workLocal, mmLocal_fp32[preOffset + gmmSwiglu->tokenLen / BISECT], reduceTmpLocal,
|
||||
static_cast<uint32_t>(halfTokenLen));
|
||||
int32_t eventIdVToS = static_cast<int32_t>(GetTPipePtr()->FetchEventID(HardEvent::V_S));
|
||||
SetFlag<HardEvent::V_S>(eventIdVToS);
|
||||
WaitFlag<HardEvent::V_S>(eventIdVToS);
|
||||
float quantScale = reduceResLocal.GetValue(0) / QUANT_SCALE_INT8;
|
||||
LocalTensor<float> quantScaleLocal = quantScaleOutQueue.DeQue<float>();
|
||||
quantScaleLocal.SetValue(loopIdx, quantScale);
|
||||
quantScale = 1 / quantScale;
|
||||
int32_t eventIdSToV = static_cast<int32_t>(GetTPipePtr()->FetchEventID(HardEvent::S_V));
|
||||
SetFlag<HardEvent::S_V>(eventIdSToV);
|
||||
WaitFlag<HardEvent::S_V>(eventIdSToV);
|
||||
Muls(mmLocal_fp32[preOffset], mmLocal_fp32[preOffset], quantScale, halfTokenLen);
|
||||
PipeBarrier<PIPE_V>();
|
||||
LocalTensor<int8_t> quantLocal = quantOutQueue.DeQue<int8_t>();
|
||||
int32_t dstTempOffset = static_cast<int32_t>(preOffset / BISECT);
|
||||
int32_t srcTempOffset = static_cast<int32_t>(preOffset);
|
||||
int32_t tempCount = static_cast<int32_t>(halfTokenLen);
|
||||
LocalTensor<int8_t> castSpace = reduceWorkspace.Get<int8_t>(UB_BLOCK_UNIT_SIZE);
|
||||
CastFp32ToInt8Template(quantLocal, mmLocal_fp32, castSpace, dstTempOffset, srcTempOffset, tempCount);
|
||||
mmOutQueue.EnQue(mmLocal_fp32);
|
||||
quantOutQueue.EnQue(quantLocal);
|
||||
}
|
||||
|
||||
__aicore__ inline void GMMA8W4PostProcess::UpdateVecConfig(uint32_t blockIdx, VecConfig &vecConfig,
|
||||
WorkSpaceSplitConfig &workspaceSplitConfig,
|
||||
int64_t workspaceSplitLoopIdx, TPipe *pipe)
|
||||
{
|
||||
// 第一步 读取grouplist reduceSum 计算总数据个数
|
||||
vecConfig.M = workspaceSplitLoopIdx < workspaceSplitConfig.loopCount - 1 ? workspaceSplitConfig.notLastTaskSize :
|
||||
workspaceSplitConfig.lastLoopTaskSize;
|
||||
// 第二步 计算分核
|
||||
uint32_t eachCoreTaskNum = (vecConfig.M + aivCoreNum - 1) / aivCoreNum;
|
||||
vecConfig.usedCoreNum = vecConfig.M >= aivCoreNum ? aivCoreNum : vecConfig.M;
|
||||
uint32_t tailCoreIdx = vecConfig.M - (eachCoreTaskNum - 1) * vecConfig.usedCoreNum;
|
||||
vecConfig.taskNum = blockIdx < tailCoreIdx ? eachCoreTaskNum : eachCoreTaskNum - 1;
|
||||
vecConfig.startIdx =
|
||||
blockIdx < tailCoreIdx ? eachCoreTaskNum * blockIdx : ((eachCoreTaskNum - 1) * blockIdx + tailCoreIdx);
|
||||
vecConfig.curIdx = vecConfig.startIdx;
|
||||
vecConfig.startOffset = vecConfig.startIdx * gmmSwiglu->tokenLen;
|
||||
vecConfig.curOffset = vecConfig.startOffset;
|
||||
int64_t curStartIdx = vecConfig.startIdx;
|
||||
int64_t prevM = workspaceSplitLoopIdx * workspaceSplitConfig.notLastTaskSize;
|
||||
for (uint32_t groupIdx = workspaceSplitConfig.rightMatrixExpertStartIndex;
|
||||
groupIdx <= workspaceSplitConfig.rightMatrixExpertEndIndex; groupIdx++) {
|
||||
int64_t currM = groupListGM.GetValue(groupIdx);
|
||||
int64_t tempM = currM - prevM;
|
||||
prevM = currM;
|
||||
if (curStartIdx >= 0 && curStartIdx - tempM < 0) {
|
||||
vecConfig.curGroupIdx = groupIdx;
|
||||
vecConfig.nextUpadteInterVal = tempM - curStartIdx;
|
||||
}
|
||||
curStartIdx -= tempM;
|
||||
}
|
||||
// 第三步 计算总数据量
|
||||
vecConfig.outLoopNum = (vecConfig.taskNum + gmmSwiglu->maxProcessRowNum - 1) / gmmSwiglu->maxProcessRowNum;
|
||||
vecConfig.tailLoopNum = vecConfig.taskNum % gmmSwiglu->maxProcessRowNum ?
|
||||
vecConfig.taskNum % gmmSwiglu->maxProcessRowNum :
|
||||
gmmSwiglu->maxProcessRowNum;
|
||||
|
||||
// 第四步 申请空间
|
||||
// 2 * row * n * sizeof(float) + row * n / 2 * sizeof(int8) + alignUp<row, 8> * sizeof(float) + n * sizeof(float) +
|
||||
// n / 2 *sizeof(float) + 64 < 191 * 1024
|
||||
pipe->InitBuffer(mmOutQueue, 1, 2 * gmmSwiglu->maxProcessRowNum * gmmSwiglu->tokenLen * sizeof(float));
|
||||
pipe->InitBuffer(quantOutQueue, 1,
|
||||
gmmSwiglu->maxProcessRowNum * gmmSwiglu->tokenLen / SWIGLU_REDUCE_FACTOR * sizeof(int8_t));
|
||||
pipe->InitBuffer(quantScaleOutQueue, 1, AlignUp<int32_t>(gmmSwiglu->maxProcessRowNum, ALIGN_8_ELE) * sizeof(float));
|
||||
pipe->InitBuffer(weightAuxiliaryMatrixInQueue, 1, gmmSwiglu->tokenLen * sizeof(float));
|
||||
// two 32 byte buffer for reduceMax calculation in Quant.
|
||||
pipe->InitBuffer(reduceWorkspace, gmmSwiglu->tokenLen / SWIGLU_REDUCE_FACTOR * sizeof(float) + UB_BLOCK_UNIT_SIZE +
|
||||
UB_BLOCK_UNIT_SIZE);
|
||||
}
|
||||
|
||||
__aicore__ inline void GMMA8W4PostProcess::PreLoadAuxiliaryMatrix(VecConfig &vecConfig)
|
||||
{
|
||||
LocalTensor<float> weightAuxiliaryMatrixLocal = weightAuxiliaryMatrixInQueue.DeQue<float>();
|
||||
DataCopyExtParams copyAuxiliaryMatrixParams{1, static_cast<uint32_t>(gmmSwiglu->tokenLen * sizeof(float)), 0, 0, 0};
|
||||
DataCopyPadExtParams<float> padParams{false, 0, 0, 0};
|
||||
DataCopyPad(weightAuxiliaryMatrixLocal, weightAuxiliaryMatrixGM[vecConfig.curGroupIdx * gmmSwiglu->tokenLen],
|
||||
copyAuxiliaryMatrixParams, padParams);
|
||||
weightAuxiliaryMatrixInQueue.EnQue(weightAuxiliaryMatrixLocal);
|
||||
}
|
||||
|
||||
__aicore__ inline void GMMA8W4PostProcess::UpdateAuxiliaryMatrix(uint32_t loopIdx, VecConfig &vecConfig)
|
||||
{
|
||||
// 更新weightAuxiliaryMatrix
|
||||
if (unlikely(vecConfig.nextUpadteInterVal == 0)) {
|
||||
int64_t loop = gmmSwiglu->groupListLen - vecConfig.curGroupIdx;
|
||||
while (loop--) {
|
||||
int64_t curTemp = groupListGM.GetValue(vecConfig.curGroupIdx);
|
||||
vecConfig.curGroupIdx++;
|
||||
int64_t nextTemp = groupListGM.GetValue(vecConfig.curGroupIdx);
|
||||
if (nextTemp != curTemp) {
|
||||
vecConfig.nextUpadteInterVal = nextTemp - curTemp;
|
||||
break;
|
||||
}
|
||||
}
|
||||
LocalTensor<float> weightAuxiliaryMatrixLocal = weightAuxiliaryMatrixInQueue.DeQue<float>();
|
||||
DataCopyExtParams copyParams{1, static_cast<uint32_t>(gmmSwiglu->tokenLen * sizeof(float)), 0, 0, 0};
|
||||
DataCopyPadExtParams<float> padParams{false, 0, 0, 0};
|
||||
DataCopyPad(weightAuxiliaryMatrixLocal, weightAuxiliaryMatrixGM[vecConfig.curGroupIdx * gmmSwiglu->tokenLen],
|
||||
copyParams, padParams);
|
||||
weightAuxiliaryMatrixInQueue.EnQue(weightAuxiliaryMatrixLocal);
|
||||
}
|
||||
}
|
||||
|
||||
__aicore__ inline void GMMA8W4PostProcess::Process(WorkSpaceSplitConfig &workspaceSplitConfig,
|
||||
int64_t workspaceSplitLoopIdx, TPipe *pipe)
|
||||
{
|
||||
if ASCEND_IS_AIV {
|
||||
if (workspaceSplitLoopIdx >= workspaceSplitConfig.loopCount || workspaceSplitLoopIdx < 0) {
|
||||
return;
|
||||
}
|
||||
VecConfig vecConfig;
|
||||
UpdateVecConfig(blockIdx, vecConfig, workspaceSplitConfig, workspaceSplitLoopIdx, pipe);
|
||||
|
||||
if (blockIdx < vecConfig.usedCoreNum) {
|
||||
mmOutGM = (workspaceSplitLoopIdx % 2 == 0 ? mmOutGM1 : mmOutGM2);
|
||||
LocalTensor<float> weightAuxiliaryMatrixLocal = weightAuxiliaryMatrixInQueue.AllocTensor<float>();
|
||||
LocalTensor<half> mmLocal_fp32 = mmOutQueue.AllocTensor<half>();
|
||||
LocalTensor<float> quantScaleLocal = quantScaleOutQueue.AllocTensor<float>();
|
||||
LocalTensor<int8_t> quantLocal = quantOutQueue.AllocTensor<int8_t>();
|
||||
|
||||
mmOutQueue.EnQue(mmLocal_fp32);
|
||||
quantScaleOutQueue.EnQue(quantScaleLocal);
|
||||
quantOutQueue.EnQue(quantLocal);
|
||||
weightAuxiliaryMatrixInQueue.EnQue(weightAuxiliaryMatrixLocal);
|
||||
PreLoadAuxiliaryMatrix(vecConfig);
|
||||
for (uint32_t outLoopIdx = 0; outLoopIdx < vecConfig.outLoopNum; outLoopIdx++) {
|
||||
vecConfig.innerLoopNum =
|
||||
outLoopIdx == (vecConfig.outLoopNum - 1) ? vecConfig.tailLoopNum : gmmSwiglu->maxProcessRowNum;
|
||||
int32_t eventIdMTE3ToMTE2 = static_cast<int32_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE3_MTE2));
|
||||
SetFlag<HardEvent::MTE3_MTE2>(eventIdMTE3ToMTE2);
|
||||
WaitFlag<HardEvent::MTE3_MTE2>(eventIdMTE3ToMTE2);
|
||||
// 1.matmul中间结果搬入 + 高四位与低四位合并
|
||||
customDataCopyIn(outLoopIdx, mmOutGM, vecConfig, workspaceSplitConfig);
|
||||
|
||||
for (uint32_t innerLoopIdx = 0; innerLoopIdx < vecConfig.innerLoopNum; innerLoopIdx++) {
|
||||
// 2.如果涉及group切换，更新辅助矩阵
|
||||
UpdateAuxiliaryMatrix(innerLoopIdx, vecConfig);
|
||||
// 3. 四步vector计算（辅助矩阵加回、perToken反量化、Swiglu、Quant）
|
||||
VectorCompute(innerLoopIdx, vecConfig, workspaceSplitConfig);
|
||||
}
|
||||
int32_t eventIdVToMTE3 = static_cast<int32_t>(GetTPipePtr()->FetchEventID(HardEvent::V_MTE3));
|
||||
SetFlag<HardEvent::V_MTE3>(eventIdVToMTE3);
|
||||
WaitFlag<HardEvent::V_MTE3>(eventIdVToMTE3);
|
||||
customDataCopyOut(vecConfig, workspaceSplitConfig);
|
||||
}
|
||||
weightAuxiliaryMatrixLocal = weightAuxiliaryMatrixInQueue.DeQue<float>();
|
||||
mmLocal_fp32 = mmOutQueue.DeQue<half>();
|
||||
quantScaleLocal = quantScaleOutQueue.DeQue<float>();
|
||||
quantLocal = quantOutQueue.DeQue<int8_t>();
|
||||
|
||||
weightAuxiliaryMatrixInQueue.FreeTensor(weightAuxiliaryMatrixLocal);
|
||||
mmOutQueue.FreeTensor(mmLocal_fp32);
|
||||
quantScaleOutQueue.FreeTensor(quantScaleLocal);
|
||||
quantOutQueue.FreeTensor(quantLocal);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
__aicore__ inline void GMMA8W4PostProcess::customDataCopyOut(VecConfig &vecConfig,
|
||||
WorkSpaceSplitConfig &workspaceSplitConfig)
|
||||
{
|
||||
LocalTensor<float> quantScaleLocal = quantScaleOutQueue.DeQue<float>();
|
||||
DataCopyParams copyParams_0{1, (uint16_t)(vecConfig.innerLoopNum * sizeof(float)), 0, 0};
|
||||
DataCopyPad(quantScaleOutputGM[workspaceSplitConfig.leftMatrixStartIndex + vecConfig.startIdx], quantScaleLocal,
|
||||
copyParams_0);
|
||||
LocalTensor<int8_t> quantLocal = quantOutQueue.DeQue<int8_t>();
|
||||
DataCopyParams copyParams_1{
|
||||
1, (uint16_t)(vecConfig.innerLoopNum * gmmSwiglu->tokenLen / SWIGLU_REDUCE_FACTOR * sizeof(int8_t)), 0, 0};
|
||||
DataCopyPad(quantOutputGM[(workspaceSplitConfig.leftMatrixStartIndex + vecConfig.startIdx) * gmmSwiglu->tokenLen /
|
||||
SWIGLU_REDUCE_FACTOR],
|
||||
quantLocal, copyParams_1);
|
||||
|
||||
vecConfig.startIdx += vecConfig.innerLoopNum;
|
||||
vecConfig.startOffset = vecConfig.startIdx * gmmSwiglu->tokenLen;
|
||||
quantOutQueue.EnQue(quantLocal);
|
||||
quantScaleOutQueue.EnQue(quantScaleLocal);
|
||||
}
|
||||
|
||||
} // namespace GROUPED_MATMUL_SWIGLU_QUANT
|
||||
#endif // GMM_SWIGLU_QUANT_A8W4_MSD
|
||||
#endif // ASCENDC_GROUPED_MATMUL_SWIGLU_QUANT_A8W4_MSD_AFTER_H
|
||||
@@ -0,0 +1,219 @@
|
||||
/**
|
||||
* Copyright (c) 2025 Huawei Technologies Co., Ltd.
|
||||
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
* CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
* See LICENSE in the root of the software repository for the full text of the License.
|
||||
*/
|
||||
|
||||
/*!
|
||||
* \file grouped_matmul_swiglu_quant_a8w4_msd_pre.h
|
||||
* \brief
|
||||
*/
|
||||
|
||||
#ifndef ASCENDC_GROUPED_MATMUL_SWIGLU_QUANT_A8W4_MSD_PRE_H
|
||||
#define ASCENDC_GROUPED_MATMUL_SWIGLU_QUANT_A8W4_MSD_PRE_H
|
||||
#include "grouped_matmul_swiglu_quant_utils.h"
|
||||
#include "kernel_operator.h"
|
||||
#ifdef GMM_SWIGLU_QUANT_A8W4_MSD
|
||||
namespace GROUPED_MATMUL_SWIGLU_QUANT {
|
||||
using namespace AscendC;
|
||||
#define BUFFER_NUM_A8W4_PRE 1
|
||||
constexpr int TWO = 2;
|
||||
constexpr int EIGHT = 8;
|
||||
constexpr size_t LEN_128 = 128; // 16bit operator
|
||||
constexpr int DATA_BLOCK_SIZE_32 = 32;
|
||||
class GMMA8W4PreProcess {
|
||||
public:
|
||||
__aicore__ inline GMMA8W4PreProcess(){};
|
||||
__aicore__ inline void Init(const GMAddrParams gmAddrParams,
|
||||
const GMMSwigluBaseParams *__restrict gmmSwigluBaseParamsIN);
|
||||
__aicore__ inline void CalculateTaskInfoEachCore(uint32_t &curCoreTaskNum_, uint32_t &curCoreStartOffset_);
|
||||
__aicore__ inline void Process(WorkSpaceSplitConfig &workspaceSplitConfig, int64_t workspaceSplitLoopIdx,
|
||||
TPipe *pipe);
|
||||
__aicore__ inline void CustomInitBuffer(TPipe *pipe);
|
||||
|
||||
private:
|
||||
TQue<QuePosition::VECIN, BUFFER_NUM_A8W4_PRE> vecInQueueX, vecInQueueXBak;
|
||||
TQue<QuePosition::VECOUT, BUFFER_NUM_A8W4_PRE> vecOutQueueA1;
|
||||
TQue<QuePosition::VECOUT, BUFFER_NUM_A8W4_PRE> vecOutQueueA2;
|
||||
TQue<QuePosition::VECOUT, BUFFER_NUM_A8W4_PRE> vecOutQueueA3;
|
||||
TQue<QuePosition::VECOUT, BUFFER_NUM_A8W4_PRE> vecOutQueue0F;
|
||||
TQue<QuePosition::VECOUT, BUFFER_NUM_A8W4_PRE> vecOutQueueRowSum;
|
||||
TBuf<TPosition::VECCALC> tempBuff;
|
||||
const GMMSwigluBaseParams *__restrict gmmSwigluBaseParams;
|
||||
LocalTensor<int8_t> xTensor;
|
||||
LocalTensor<half> xHighHalfTensor;
|
||||
LocalTensor<float> xHighFloatTensor;
|
||||
LocalTensor<half> xLowHalfTensor;
|
||||
LocalTensor<half> xLowHalfTensor2;
|
||||
LocalTensor<int4b_t> xHighI4Tensor;
|
||||
LocalTensor<int4b_t> xLowI4Tensor;
|
||||
LocalTensor<int16_t> xLowI16Tensor;
|
||||
LocalTensor<int64_t> groupListTensor;
|
||||
LocalTensor<float> groupListFTensor;
|
||||
LocalTensor<float> xRowSumTensor;
|
||||
|
||||
GlobalTensor<int8_t> xGM;
|
||||
GlobalTensor<int8_t> yGm;
|
||||
GlobalTensor<int8_t> yGm1;
|
||||
GlobalTensor<int8_t> yGm2;
|
||||
GlobalTensor<int64_t> groupListGM;
|
||||
|
||||
uint32_t vK{0};
|
||||
uint32_t vKAlign{0};
|
||||
uint32_t totalM{0};
|
||||
uint32_t blockDim{0};
|
||||
uint32_t curCoreId{0};
|
||||
uint32_t curCoreTaskNum{0};
|
||||
uint32_t curCoreStartOffset{0};
|
||||
uint32_t curCoreOuterLoopNum{0};
|
||||
uint32_t curCoreInnerTailLoopNum{0};
|
||||
uint32_t groupNum{0};
|
||||
};
|
||||
|
||||
__aicore__ inline void GMMA8W4PreProcess::Init(const GMAddrParams gmAddrParams,
|
||||
const GMMSwigluBaseParams *__restrict gmmSwigluBaseParamsIN)
|
||||
{
|
||||
if ASCEND_IS_AIV {
|
||||
xGM.SetGlobalBuffer((__gm__ int8_t *)gmAddrParams.xGM);
|
||||
yGm1.SetGlobalBuffer((__gm__ int8_t *)gmAddrParams.workSpaceGM);
|
||||
yGm2.SetGlobalBuffer((__gm__ int8_t *)gmAddrParams.workSpaceGM + gmAddrParams.workSpaceOffset1);
|
||||
groupListGM.SetGlobalBuffer((__gm__ int64_t *)gmAddrParams.groupListGM);
|
||||
gmmSwigluBaseParams = gmmSwigluBaseParamsIN;
|
||||
vK = gmmSwigluBaseParams->K;
|
||||
groupNum = static_cast<uint32_t>(gmmSwigluBaseParams->groupNum);
|
||||
// M * K * 7B (1B + 0.5B + 0.5B + 2B + 4B) <= UBsize - 256B
|
||||
blockDim = GetBlockNum() * GetTaskRation();
|
||||
}
|
||||
}
|
||||
|
||||
__aicore__ inline void GMMA8W4PreProcess::CustomInitBuffer(TPipe *pipe)
|
||||
{
|
||||
pipe->InitBuffer(vecInQueueX, BUFFER_NUM_A8W4_PRE, vK * sizeof(int8_t)); // K * 1B
|
||||
pipe->InitBuffer(vecOutQueueA1, BUFFER_NUM_A8W4_PRE, vK * sizeof(int4b_t)); // K * 0.5B
|
||||
pipe->InitBuffer(vecOutQueueA2, BUFFER_NUM_A8W4_PRE, vK * sizeof(int4b_t)); // K * 0.5B
|
||||
pipe->InitBuffer(vecOutQueueA3, BUFFER_NUM_A8W4_PRE, vK * sizeof(half)); // K * 2B
|
||||
// xLowHalfTensor, xLowHalfTensor2 and xHighFloatTensor share the same buffer
|
||||
pipe->InitBuffer(tempBuff, vK * sizeof(float)); // K * 4B
|
||||
constexpr int BUFFER_SIZE_256B = 128 * sizeof(int16_t);
|
||||
pipe->InitBuffer(vecOutQueue0F, BUFFER_NUM_A8W4_PRE, BUFFER_SIZE_256B); // 256B
|
||||
}
|
||||
|
||||
|
||||
__aicore__ inline void GMMA8W4PreProcess::CalculateTaskInfoEachCore(uint32_t &curCoreTaskNum_,
|
||||
uint32_t &curCoreStartOffset_)
|
||||
{
|
||||
// 均分任务数
|
||||
int64_t eachCoreTaskNum = (totalM + blockDim - 1) / blockDim; // 每个核处理的数据量
|
||||
// 尾核任务数
|
||||
int64_t taskNumPertailCore = eachCoreTaskNum - 1;
|
||||
// 实际使用核数
|
||||
int64_t usedCoreNum = totalM >= blockDim ? blockDim : totalM;
|
||||
// 尾核起始索引
|
||||
uint32_t tailCoreIdx = totalM - (eachCoreTaskNum - 1) * usedCoreNum;
|
||||
curCoreId = GetBlockIdx();
|
||||
// 每个核处理的任务数量 = 是否为尾核 ？均分任务数 ：(均分任务数 - 1)
|
||||
curCoreTaskNum_ = curCoreId < tailCoreIdx ? eachCoreTaskNum : eachCoreTaskNum - 1;
|
||||
// 每个核处理的起始偏移地址 = 是否为尾核 ？均分任务数 * blockId : (均分任务数 - 1) * blockId + 尾核起始索引
|
||||
curCoreStartOffset_ =
|
||||
curCoreId < tailCoreIdx ? eachCoreTaskNum * curCoreId : ((eachCoreTaskNum - 1) * curCoreId + tailCoreIdx);
|
||||
}
|
||||
|
||||
__aicore__ inline void GMMA8W4PreProcess::Process(WorkSpaceSplitConfig &workspaceSplitConfig,
|
||||
int64_t workspaceSplitLoopIdx, TPipe *pipe)
|
||||
{
|
||||
if ASCEND_IS_AIV {
|
||||
if (workspaceSplitLoopIdx >= workspaceSplitConfig.loopCount) {
|
||||
return;
|
||||
}
|
||||
yGm = (workspaceSplitLoopIdx % 2 == 0 ? yGm1 : yGm2);
|
||||
CustomInitBuffer(pipe);
|
||||
constexpr int32_t MASK = 128;
|
||||
xTensor = vecInQueueX.AllocTensor<int8_t>();
|
||||
xHighI4Tensor = vecOutQueueA1.AllocTensor<int4b_t>();
|
||||
xLowI4Tensor = vecOutQueueA2.AllocTensor<int4b_t>();
|
||||
xHighHalfTensor = vecOutQueueA3.AllocTensor<half>();
|
||||
const uint32_t xLowHalfOffset = vK * sizeof(half);
|
||||
xLowHalfTensor = tempBuff.GetWithOffset<half>(xLowHalfOffset, 0);
|
||||
xLowHalfTensor2 = tempBuff.GetWithOffset<half>(xLowHalfOffset, xLowHalfOffset);
|
||||
xLowI16Tensor = vecOutQueue0F.AllocTensor<int16_t>();
|
||||
|
||||
Duplicate(xLowI16Tensor, static_cast<int16_t>(0x0F0F), MASK); // get rid of high 4 bits in every int8
|
||||
PipeBarrier<PIPE_V>();
|
||||
const size_t LEN_VK = (vK / 2) / 128;
|
||||
const size_t LAST_LEN_VK = (vK % 256) / 2;
|
||||
const half ONE_SIXTEENTH = static_cast<half>(0.0625f);
|
||||
// groupList仅支持count
|
||||
SetFlag<HardEvent::MTE2_S>(EVENT_ID0);
|
||||
WaitFlag<HardEvent::MTE2_S>(EVENT_ID0);
|
||||
totalM = workspaceSplitLoopIdx < workspaceSplitConfig.loopCount - 1 ? workspaceSplitConfig.notLastTaskSize :
|
||||
workspaceSplitConfig.lastLoopTaskSize;
|
||||
SetFlag<HardEvent::S_MTE2>(EVENT_ID0);
|
||||
WaitFlag<HardEvent::S_MTE2>(EVENT_ID0);
|
||||
CalculateTaskInfoEachCore(curCoreTaskNum, curCoreStartOffset);
|
||||
SetFlag<HardEvent::V_MTE2>(EVENT_ID0); // 0
|
||||
SetFlag<HardEvent::MTE3_V>(EVENT_ID0); // 1
|
||||
SetFlag<HardEvent::MTE3_V>(EVENT_ID1); // 2
|
||||
|
||||
for (uint32_t xloop = 0; xloop < curCoreTaskNum; xloop++) {
|
||||
uint64_t relStartAddr = (xloop + curCoreStartOffset) * vK;
|
||||
uint64_t absStartAddr = workspaceSplitLoopIdx * workspaceSplitConfig.notLastTaskSize * vK + relStartAddr;
|
||||
// 高四位处理开始
|
||||
WaitFlag<HardEvent::V_MTE2>(EVENT_ID0); // 0
|
||||
DataCopy(xTensor, xGM[absStartAddr], vK);
|
||||
SetFlag<HardEvent::MTE2_V>(EVENT_ID0); // 3
|
||||
WaitFlag<HardEvent::MTE2_V>(EVENT_ID0); // 3
|
||||
Cast(xHighHalfTensor, xTensor, AscendC::RoundMode::CAST_NONE, vK);
|
||||
PipeBarrier<PIPE_V>();
|
||||
Muls(xHighHalfTensor, xHighHalfTensor, ONE_SIXTEENTH, vK);
|
||||
PipeBarrier<PIPE_V>();
|
||||
WaitFlag<HardEvent::MTE3_V>(EVENT_ID1); // 2
|
||||
Cast(xHighI4Tensor, xHighHalfTensor, AscendC::RoundMode::CAST_FLOOR, vK);
|
||||
SetFlag<HardEvent::V_MTE3>(EVENT_ID0); // 4
|
||||
WaitFlag<HardEvent::V_MTE3>(EVENT_ID0); // 4
|
||||
DataCopy(yGm[relStartAddr], xHighI4Tensor.ReinterpretCast<int8_t>(), vK / 2);
|
||||
// 高四位处理结束
|
||||
|
||||
// 低四位处理开始
|
||||
SetFlag<HardEvent::MTE3_V>(EVENT_ID1); // 2
|
||||
And(xLowHalfTensor.ReinterpretCast<int16_t>(), xTensor.ReinterpretCast<int16_t>(), xLowI16Tensor, LEN_128,
|
||||
LEN_VK, {1, 1, 1, 8, 8, 0});
|
||||
if (LAST_LEN_VK > 0) {
|
||||
And(xLowHalfTensor[LEN_VK * LEN_128].ReinterpretCast<int16_t>(),
|
||||
xTensor[LEN_VK * LEN_128 * TWO].ReinterpretCast<int16_t>(), xLowI16Tensor, LAST_LEN_VK, 1,
|
||||
{1, 1, 1, 8, 8, 0});
|
||||
}
|
||||
PipeBarrier<PIPE_V>();
|
||||
SetFlag<HardEvent::V_MTE2>(EVENT_ID0); // 0
|
||||
Cast(xLowHalfTensor2.ReinterpretCast<half>(), xLowHalfTensor.ReinterpretCast<int8_t>(),
|
||||
AscendC::RoundMode::CAST_NONE, vK);
|
||||
PipeBarrier<PIPE_V>();
|
||||
const half MINUS_EIGHT = static_cast<half>(-8);
|
||||
Adds(xHighHalfTensor, xLowHalfTensor2, MINUS_EIGHT, vK);
|
||||
PipeBarrier<PIPE_V>();
|
||||
WaitFlag<HardEvent::MTE3_V>(EVENT_ID0); // 1
|
||||
Cast(xLowI4Tensor, xHighHalfTensor.ReinterpretCast<half>(), AscendC::RoundMode::CAST_NONE, vK);
|
||||
SetFlag<HardEvent::V_MTE3>(EVENT_ID1); // 5
|
||||
WaitFlag<HardEvent::V_MTE3>(EVENT_ID1); // 5
|
||||
DataCopy(yGm[relStartAddr + vK / TWO], xLowI4Tensor.ReinterpretCast<int8_t>(), vK / TWO);
|
||||
SetFlag<HardEvent::MTE3_V>(EVENT_ID0); // 1
|
||||
// 低四位处理结束
|
||||
}
|
||||
|
||||
WaitFlag<HardEvent::V_MTE2>(EVENT_ID0); // 0
|
||||
WaitFlag<HardEvent::MTE3_V>(EVENT_ID0); // 1
|
||||
WaitFlag<HardEvent::MTE3_V>(EVENT_ID1); // 2
|
||||
vecInQueueX.FreeTensor(xTensor);
|
||||
vecOutQueueA1.FreeTensor(xHighI4Tensor);
|
||||
vecOutQueueA2.FreeTensor(xLowI4Tensor);
|
||||
vecOutQueueA3.FreeTensor(xHighHalfTensor);
|
||||
vecOutQueue0F.FreeTensor(xLowI16Tensor);
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace GROUPED_MATMUL_SWIGLU_QUANT
|
||||
#endif // GMM_SWIGLU_QUANT_A8W4_MSD
|
||||
#endif // ASCENDC_GROUPED_MATMUL_SWIGLU_QUANT_A8W4_MSD_PRE_H
|
||||
@@ -0,0 +1,635 @@
|
||||
/**
|
||||
* Copyright (c) 2025 Huawei Technologies Co., Ltd.
|
||||
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
* CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
* See LICENSE in the root of the software repository for the full text of the License.
|
||||
*/
|
||||
|
||||
/*!
|
||||
* \file grouped_matmul_swiglu_quant_split_ws.h
|
||||
* \brief
|
||||
*/
|
||||
#ifndef ASCENDC_GROUPED_MATMUL_SWIGLU_QUANT_SPLIT_WS_H
|
||||
#define ASCENDC_GROUPED_MATMUL_SWIGLU_QUANT_SPLIT_WS_H
|
||||
|
||||
#include "grouped_matmul_swiglu_quant_utils.h"
|
||||
namespace GROUPED_MATMUL_SWIGLU_QUANT {
|
||||
/** @brief internal computation class
|
||||
*/
|
||||
|
||||
template <class mmType, bool sync = false, typename CHANNELDTYPE = float>
|
||||
class GMMSwigluSplitWorkSpaceCompute {
|
||||
public:
|
||||
using AT = typename mmType::AT::T;
|
||||
using BT = typename mmType::BT::T;
|
||||
using B = typename mmType::BT;
|
||||
using CT = typename mmType::CT::T;
|
||||
using BiasT = typename mmType::BiasT::T;
|
||||
using WT = int8_t;
|
||||
constexpr static bool transposeX = mmType::AT::isTrans;
|
||||
constexpr static bool transposeW = mmType::BT::isTrans;
|
||||
static constexpr float FLOAT_INF = 3e+99;
|
||||
/** @brief constructor */
|
||||
__aicore__ inline GMMSwigluSplitWorkSpaceCompute(typename mmType::MT &mm_) : mm(mm_)
|
||||
{
|
||||
}
|
||||
|
||||
__aicore__ inline void Init(GM_ADDR x, GM_ADDR weight, GM_ADDR perChannelScale, GM_ADDR perTokenScale,
|
||||
GM_ADDR groupList, GM_ADDR quantOutput, GM_ADDR quantScaleOutput, GM_ADDR workspace,
|
||||
const GMMSwigluBaseParams *__restrict gmmBaseParamsIN,
|
||||
const TCubeTiling *__restrict mmTilingDataIN, const GMMSwiglu *__restrict gmmSwigluIN,
|
||||
TPipe *tPipeIN);
|
||||
__aicore__ inline void Process();
|
||||
|
||||
private:
|
||||
__aicore__ inline void MMCompute(uint32_t groupIdx, MNConfig &mnConfig, uint32_t coreIdx,
|
||||
GlobalTensor<int32_t> &mmOutGM);
|
||||
|
||||
__aicore__ inline void UpdateMnConfig(MNConfig &mnConfig);
|
||||
|
||||
__aicore__ inline void SetMNConfig(const int32_t splitValue, const uint32_t groupIdx, MNConfig &mnConfig);
|
||||
|
||||
__aicore__ inline void SetMKN(const int32_t splitValue, const uint32_t groupIdx, MNConfig &mnConfig);
|
||||
|
||||
__aicore__ inline uint64_t GetWOffset(uint32_t tailN, uint32_t k);
|
||||
|
||||
__aicore__ inline void MNBlockIdxCompute(MNConfig &mnConfig, const uint32_t curBlock, const uint32_t count,
|
||||
const uint32_t thresholdM_dimN);
|
||||
|
||||
template <typename DTYPE_CS>
|
||||
__aicore__ inline void UpdateChannelScale(uint32_t loopidx, VecConfig &vecConfig);
|
||||
|
||||
__aicore__ inline void VectorCompute(uint32_t loopidx, VecConfig &vecConfig);
|
||||
|
||||
template <typename DTYPE_CS>
|
||||
__aicore__ inline void PreLoadTokenAndChannel(LocalTensor<float> &channelScaleLocal, VecConfig &vecConfig);
|
||||
|
||||
__aicore__ inline void UpdateVecConfig(uint32_t blockIdx, VecConfig &vecConfig);
|
||||
|
||||
__aicore__ inline void UpdateWorkSpaceSplitConfig(WorkSpaceSplitConfig &workspaceSplitConfig,
|
||||
int32_t workspaceSplitLoopIdx);
|
||||
|
||||
__aicore__ inline void InitWorkSpaceSplitConfig(WorkSpaceSplitConfig &workspaceSplitConfig);
|
||||
|
||||
__aicore__ inline void customDataCopyIn(uint32_t outLoopIdx, GlobalTensor<int32_t> &mmOutGM, VecConfig &vecConfig);
|
||||
|
||||
__aicore__ inline void customDataCopyOut(VecConfig &vecConfig);
|
||||
|
||||
__aicore__ inline void Dequant(uint32_t loopidx, VecConfig &vecConfig);
|
||||
|
||||
__aicore__ inline void Quant(uint32_t loopidx, VecConfig &vecConfig);
|
||||
|
||||
__aicore__ inline void Swiglu(uint32_t loopidx, VecConfig &vecConfig);
|
||||
|
||||
private:
|
||||
typename mmType::MT &mm;
|
||||
const GMMSwigluBaseParams *__restrict gmmBaseParams;
|
||||
const GMMSwiglu *__restrict gmmSwiglu;
|
||||
const TCubeTiling *__restrict mmTilingData;
|
||||
WorkSpaceSplitConfig workspaceSplitConfig;
|
||||
TPipe *pipe;
|
||||
GlobalTensor<int8_t> xGM;
|
||||
GlobalTensor<int8_t> weightGM;
|
||||
GlobalTensor<CHANNELDTYPE> perChannelScaleGM;
|
||||
GlobalTensor<float> perTokenScaleGM;
|
||||
GlobalTensor<int64_t> groupListGM;
|
||||
GlobalTensor<int8_t> quantOutputGM;
|
||||
GlobalTensor<float> quantScaleOutputGM;
|
||||
GlobalTensor<int32_t> mmOutGM1;
|
||||
GlobalTensor<int32_t> mmOutGM2;
|
||||
// define the que
|
||||
TQue<QuePosition::VECIN, 1> mmOutQueue;
|
||||
TQue<QuePosition::VECIN, 1> perChannelScaleInQueue;
|
||||
TQue<QuePosition::VECOUT, 1> quantOutQueue;
|
||||
TQue<QuePosition::VECOUT, 1> quantScaleOutQueue;
|
||||
TBuf<TPosition::VECCALC> reduceWorkspace;
|
||||
uint32_t blockIdx = 0;
|
||||
int64_t aicCoreNum = 0;
|
||||
int64_t aivCoreNum = 0;
|
||||
GM_ADDR xTensorPtr;
|
||||
GM_ADDR weightTensorPtr;
|
||||
float limited = FLOAT_INF;
|
||||
};
|
||||
|
||||
template <typename mmType, bool sync, typename CHANNELDTYPE>
|
||||
__aicore__ inline void GMMSwigluSplitWorkSpaceCompute<mmType, sync, CHANNELDTYPE>::Init(
|
||||
GM_ADDR x, GM_ADDR weight, GM_ADDR perChannelScale, GM_ADDR perTokenScale, GM_ADDR groupList, GM_ADDR quantOutput,
|
||||
GM_ADDR quantScaleOutput, GM_ADDR workspace, const GMMSwigluBaseParams *__restrict gmmSwigluBaseParamsIn,
|
||||
const TCubeTiling *__restrict mmTilingDataIN, const GMMSwiglu *__restrict gmmSwigluIN, TPipe *tPipeIN)
|
||||
{
|
||||
aicCoreNum = GetBlockNum();
|
||||
aivCoreNum = aicCoreNum * 2;
|
||||
blockIdx = GetBlockIdx();
|
||||
pipe = tPipeIN;
|
||||
xTensorPtr = x;
|
||||
weightTensorPtr = weight;
|
||||
mmTilingData = mmTilingDataIN;
|
||||
gmmBaseParams = gmmSwigluBaseParamsIn;
|
||||
gmmSwiglu = gmmSwigluIN;
|
||||
limited = gmmBaseParams->limited;
|
||||
groupListGM.SetGlobalBuffer((__gm__ int64_t *)groupList, gmmSwiglu->groupListLen);
|
||||
mmOutGM1.SetGlobalBuffer((__gm__ int32_t *)workspace, gmmBaseParams->mLimit * gmmSwiglu->tokenLen);
|
||||
mmOutGM2.SetGlobalBuffer((__gm__ int32_t *)workspace + gmmBaseParams->mLimit * gmmSwiglu->tokenLen,
|
||||
gmmBaseParams->mLimit * gmmSwiglu->tokenLen);
|
||||
if ASCEND_IS_AIV {
|
||||
perChannelScaleGM.SetGlobalBuffer((__gm__ CHANNELDTYPE *)perChannelScale,
|
||||
gmmSwiglu->groupListLen * gmmSwiglu->tokenLen);
|
||||
perTokenScaleGM.SetGlobalBuffer((__gm__ float *)perTokenScale, gmmBaseParams->M);
|
||||
quantOutputGM.SetGlobalBuffer((__gm__ int8_t *)quantOutput,
|
||||
gmmBaseParams->M * gmmSwiglu->tokenLen / SWIGLU_REDUCE_FACTOR);
|
||||
quantScaleOutputGM.SetGlobalBuffer((__gm__ float *)quantScaleOutput, gmmBaseParams->M);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename mmType, bool sync, typename CHANNELDTYPE>
|
||||
__aicore__ inline void GMMSwigluSplitWorkSpaceCompute<mmType, sync, CHANNELDTYPE>::InitWorkSpaceSplitConfig(
|
||||
WorkSpaceSplitConfig &workspaceSplitConfig)
|
||||
{
|
||||
workspaceSplitConfig.M = groupListGM.GetValue(gmmSwiglu->groupListLen - 1);
|
||||
workspaceSplitConfig.loopCount = Ceil(workspaceSplitConfig.M, gmmBaseParams->mLimit);
|
||||
workspaceSplitConfig.notLastTaskSize = gmmBaseParams->mLimit;
|
||||
workspaceSplitConfig.lastLoopTaskSize =
|
||||
workspaceSplitConfig.M - (workspaceSplitConfig.loopCount - 1) * gmmBaseParams->mLimit;
|
||||
workspaceSplitConfig.leftMatrixStartIndex = 0;
|
||||
workspaceSplitConfig.rightMatrixExpertStartIndex = 0;
|
||||
workspaceSplitConfig.rightMatrixExpertNextStartIndex = 0;
|
||||
workspaceSplitConfig.isLastLoop = false;
|
||||
}
|
||||
|
||||
template <typename mmType, bool sync, typename CHANNELDTYPE>
|
||||
__aicore__ inline void GMMSwigluSplitWorkSpaceCompute<mmType, sync, CHANNELDTYPE>::UpdateWorkSpaceSplitConfig(
|
||||
WorkSpaceSplitConfig &workspaceSplitConfig, int32_t workspaceSplitLoopIdx)
|
||||
{
|
||||
workspaceSplitConfig.leftMatrixStartIndex = workspaceSplitLoopIdx * gmmBaseParams->mLimit;
|
||||
workspaceSplitConfig.rightMatrixExpertStartIndex = workspaceSplitConfig.rightMatrixExpertNextStartIndex;
|
||||
workspaceSplitConfig.rightMatrixExpertEndIndex = workspaceSplitConfig.rightMatrixExpertStartIndex;
|
||||
// 计算右专家矩阵的终止索引(rightMatrixExpertEndIndex) 和下一次的起始索引(rightMatrixExpertNextStartIndex)
|
||||
int32_t curTaskNum = 0;
|
||||
int32_t nextTaskNum = 0;
|
||||
while (workspaceSplitConfig.rightMatrixExpertEndIndex < gmmSwiglu->groupListLen) {
|
||||
curTaskNum = groupListGM.GetValue(workspaceSplitConfig.rightMatrixExpertEndIndex) -
|
||||
workspaceSplitConfig.leftMatrixStartIndex;
|
||||
int32_t nextTaskIdx = workspaceSplitConfig.rightMatrixExpertEndIndex >= gmmSwiglu->groupListLen - 1 ?
|
||||
gmmSwiglu->groupListLen - 1 :
|
||||
workspaceSplitConfig.rightMatrixExpertEndIndex + 1;
|
||||
nextTaskNum = groupListGM.GetValue(nextTaskIdx) - workspaceSplitConfig.leftMatrixStartIndex;
|
||||
if (curTaskNum > gmmBaseParams->mLimit) {
|
||||
workspaceSplitConfig.rightMatrixExpertNextStartIndex = workspaceSplitConfig.rightMatrixExpertEndIndex;
|
||||
break;
|
||||
} else if (curTaskNum == gmmBaseParams->mLimit && nextTaskNum > gmmBaseParams->mLimit) {
|
||||
workspaceSplitConfig.rightMatrixExpertNextStartIndex = workspaceSplitConfig.rightMatrixExpertEndIndex + 1;
|
||||
break;
|
||||
} else if (nextTaskNum > gmmBaseParams->mLimit) {
|
||||
workspaceSplitConfig.rightMatrixExpertEndIndex++;
|
||||
workspaceSplitConfig.rightMatrixExpertNextStartIndex = workspaceSplitConfig.rightMatrixExpertEndIndex;
|
||||
break;
|
||||
}
|
||||
workspaceSplitConfig.rightMatrixExpertEndIndex++;
|
||||
}
|
||||
workspaceSplitConfig.isLastLoop = workspaceSplitLoopIdx == workspaceSplitConfig.loopCount - 1 ? true : false;
|
||||
|
||||
if (workspaceSplitConfig.isLastLoop) {
|
||||
workspaceSplitConfig.rightMatrixExpertEndIndex =
|
||||
workspaceSplitConfig.rightMatrixExpertEndIndex >= gmmSwiglu->groupListLen ?
|
||||
gmmSwiglu->groupListLen - 1 :
|
||||
workspaceSplitConfig.rightMatrixExpertEndIndex;
|
||||
}
|
||||
}
|
||||
|
||||
template <typename mmType, bool sync, typename CHANNELDTYPE>
|
||||
__aicore__ inline void GMMSwigluSplitWorkSpaceCompute<mmType, sync, CHANNELDTYPE>::Process()
|
||||
{
|
||||
InitWorkSpaceSplitConfig(workspaceSplitConfig);
|
||||
int32_t parallelNum = 2; // 2: double workspace buffer
|
||||
for (int32_t workspaceSplitLoopIdx = 0; workspaceSplitLoopIdx < workspaceSplitConfig.loopCount;
|
||||
workspaceSplitLoopIdx++) {
|
||||
UpdateWorkSpaceSplitConfig(workspaceSplitConfig, workspaceSplitLoopIdx);
|
||||
GlobalTensor<int32_t> mmOutGM = (workspaceSplitLoopIdx % parallelNum == 0) ? mmOutGM1 : mmOutGM2;
|
||||
|
||||
if ASCEND_IS_AIC {
|
||||
if (workspaceSplitLoopIdx >= parallelNum) { // first parallelNum core no need to wait
|
||||
SyncAll<false>();
|
||||
}
|
||||
MNConfig mnConfig;
|
||||
int32_t prevSplitValue = workspaceSplitConfig.leftMatrixStartIndex;
|
||||
for (uint32_t groupIdx = workspaceSplitConfig.rightMatrixExpertStartIndex, count = 0;
|
||||
groupIdx <= workspaceSplitConfig.rightMatrixExpertEndIndex; ++groupIdx) {
|
||||
UpdateMnConfig(mnConfig);
|
||||
int32_t currSplitValue = static_cast<int32_t>(groupListGM.GetValue(groupIdx));
|
||||
currSplitValue = currSplitValue > (workspaceSplitLoopIdx + 1) * gmmBaseParams->mLimit ?
|
||||
(workspaceSplitLoopIdx + 1) * gmmBaseParams->mLimit :
|
||||
currSplitValue;
|
||||
int32_t splitValue = currSplitValue - prevSplitValue;
|
||||
prevSplitValue = currSplitValue;
|
||||
SetMNConfig(splitValue, groupIdx, mnConfig);
|
||||
if (mnConfig.m <= 0 || mnConfig.k <= 0 || mnConfig.n <= 0) {
|
||||
continue;
|
||||
}
|
||||
mnConfig.blockDimM = Ceil(mnConfig.m, mnConfig.singleM);
|
||||
mnConfig.blockDimN = Ceil(mnConfig.n, mnConfig.singleN);
|
||||
|
||||
uint32_t curCount = count + mnConfig.blockDimM * mnConfig.blockDimN;
|
||||
uint32_t curBlock = blockIdx >= count ? blockIdx : blockIdx + gmmBaseParams->coreNum;
|
||||
uint32_t thresholdM_dimN = THRESHOLD_BLOCK_NUM * mnConfig.blockDimN;
|
||||
|
||||
while (curBlock < curCount) {
|
||||
MNBlockIdxCompute(mnConfig, curBlock, count, thresholdM_dimN);
|
||||
MMCompute(groupIdx, mnConfig, blockIdx, mmOutGM);
|
||||
curBlock += aicCoreNum;
|
||||
}
|
||||
count = curCount % gmmBaseParams->coreNum;
|
||||
}
|
||||
SyncAll<false>();
|
||||
}
|
||||
|
||||
if ASCEND_IS_AIV {
|
||||
VecConfig vecConfig;
|
||||
UpdateVecConfig(blockIdx, vecConfig);
|
||||
if (blockIdx < vecConfig.usedCoreNum) {
|
||||
LocalTensor<float> channelScaleLocal = perChannelScaleInQueue.AllocTensor<float>();
|
||||
LocalTensor<int32_t> mmLocal = mmOutQueue.AllocTensor<int32_t>();
|
||||
LocalTensor<int8_t> quantLocal = quantOutQueue.AllocTensor<int8_t>();
|
||||
LocalTensor<float> quantScaleLocal = quantScaleOutQueue.AllocTensor<float>();
|
||||
mmOutQueue.EnQue(mmLocal);
|
||||
quantScaleOutQueue.EnQue(quantScaleLocal);
|
||||
quantOutQueue.EnQue(quantLocal);
|
||||
PreLoadTokenAndChannel<CHANNELDTYPE>(channelScaleLocal, vecConfig);
|
||||
}
|
||||
SyncAll<false>();
|
||||
if (blockIdx < vecConfig.usedCoreNum) {
|
||||
for (uint32_t outLoopIdx = 0; outLoopIdx < vecConfig.outLoopNum; outLoopIdx++) {
|
||||
vecConfig.innerLoopNum =
|
||||
outLoopIdx == (vecConfig.outLoopNum - 1) ? vecConfig.tailLoopNum : gmmSwiglu->maxProcessRowNum;
|
||||
int32_t eventIdMTE3ToMTE2 = static_cast<int32_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE3_MTE2));
|
||||
SetFlag<HardEvent::MTE3_MTE2>(eventIdMTE3ToMTE2);
|
||||
WaitFlag<HardEvent::MTE3_MTE2>(eventIdMTE3ToMTE2);
|
||||
customDataCopyIn(outLoopIdx, mmOutGM, vecConfig);
|
||||
for (uint32_t innerLoopIdx = 0; innerLoopIdx < vecConfig.innerLoopNum; innerLoopIdx++) {
|
||||
UpdateChannelScale<CHANNELDTYPE>(innerLoopIdx, vecConfig);
|
||||
VectorCompute(innerLoopIdx, vecConfig);
|
||||
}
|
||||
int32_t eventIdVToMTE3 = static_cast<int32_t>(GetTPipePtr()->FetchEventID(HardEvent::V_MTE3));
|
||||
SetFlag<HardEvent::V_MTE3>(eventIdVToMTE3);
|
||||
WaitFlag<HardEvent::V_MTE3>(eventIdVToMTE3);
|
||||
customDataCopyOut(vecConfig);
|
||||
}
|
||||
|
||||
LocalTensor<float> channelScaleLocal = perChannelScaleInQueue.DeQue<float>();
|
||||
LocalTensor<int32_t> mmLocal = mmOutQueue.DeQue<int32_t>();
|
||||
LocalTensor<int8_t> quantLocal = quantOutQueue.DeQue<int8_t>();
|
||||
LocalTensor<float> quantScaleLocal = quantScaleOutQueue.DeQue<float>();
|
||||
perChannelScaleInQueue.FreeTensor(channelScaleLocal);
|
||||
mmOutQueue.FreeTensor(mmLocal);
|
||||
quantScaleOutQueue.FreeTensor(quantScaleLocal);
|
||||
quantOutQueue.FreeTensor(quantLocal);
|
||||
}
|
||||
if (workspaceSplitLoopIdx < workspaceSplitConfig.loopCount - parallelNum) {
|
||||
SyncAll<false>();
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename mmType, bool sync, typename CHANNELDTYPE>
|
||||
template <typename DTYPE_CS>
|
||||
__aicore__ inline void GMMSwigluSplitWorkSpaceCompute<mmType, sync, CHANNELDTYPE>::PreLoadTokenAndChannel(
|
||||
LocalTensor<float> &channelScaleLocal, VecConfig &vecConfig)
|
||||
{
|
||||
DataCopyExtParams copyChannelParams{1, static_cast<uint32_t>(gmmSwiglu->tokenLen * sizeof(DTYPE_CS)), 0, 0, 0};
|
||||
DataCopyPadExtParams<DTYPE_CS> padParams{false, 0, 0, 0};
|
||||
if constexpr (!IsSameType<DTYPE_CS, float>::value) {
|
||||
LocalTensor<DTYPE_CS> dstLocalT = channelScaleLocal.template ReinterpretCast<DTYPE_CS>();
|
||||
DataCopyPad(dstLocalT[gmmSwiglu->tokenLen], perChannelScaleGM[vecConfig.curGroupIdx * gmmSwiglu->tokenLen],
|
||||
copyChannelParams, padParams);
|
||||
int32_t eventIdMTE2ToV = static_cast<int32_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V));
|
||||
SetFlag<HardEvent::MTE2_V>(eventIdMTE2ToV);
|
||||
WaitFlag<HardEvent::MTE2_V>(eventIdMTE2ToV);
|
||||
Cast(channelScaleLocal, dstLocalT[gmmSwiglu->tokenLen], RoundMode::CAST_NONE, gmmSwiglu->tokenLen);
|
||||
} else {
|
||||
DataCopyPad(channelScaleLocal, perChannelScaleGM[vecConfig.curGroupIdx * gmmSwiglu->tokenLen],
|
||||
copyChannelParams, padParams);
|
||||
}
|
||||
perChannelScaleInQueue.EnQue(channelScaleLocal);
|
||||
}
|
||||
|
||||
template <typename mmType, bool sync, typename CHANNELDTYPE>
|
||||
__aicore__ inline void
|
||||
GMMSwigluSplitWorkSpaceCompute<mmType, sync, CHANNELDTYPE>::MMCompute(uint32_t groupIdx, MNConfig &mnConfig,
|
||||
uint32_t coreIdx, GlobalTensor<int32_t> &mmOutGM)
|
||||
{
|
||||
uint32_t tailN = mnConfig.nIdx * mnConfig.singleN;
|
||||
uint32_t curSingleN = mnConfig.nIdx < mnConfig.blockDimN - 1 ? mnConfig.singleN : mnConfig.n - tailN;
|
||||
uint32_t curSingleM =
|
||||
mnConfig.mIdx < mnConfig.blockDimM - 1 ? mnConfig.singleM : mnConfig.m - mnConfig.mIdx * mnConfig.singleM;
|
||||
uint64_t xOffset = mnConfig.mIdx * mnConfig.singleM * mnConfig.k;
|
||||
if constexpr (transposeX) {
|
||||
xOffset = mnConfig.mIdx * mnConfig.singleM;
|
||||
}
|
||||
uint64_t outOffset = mnConfig.mIdx * mnConfig.singleM * mnConfig.n + tailN;
|
||||
xGM.SetGlobalBuffer((__gm__ int8_t *)xTensorPtr + mnConfig.xBaseOffset +
|
||||
workspaceSplitConfig.leftMatrixStartIndex * mnConfig.k);
|
||||
weightGM.SetGlobalBuffer((__gm__ int8_t *)weightTensorPtr + groupIdx * mnConfig.k * mnConfig.n +
|
||||
GetWOffset(tailN, mnConfig.k));
|
||||
if (mnConfig.blockDimM == 1) {
|
||||
weightGM.SetL2CacheHint(CacheMode::CACHE_MODE_DISABLE);
|
||||
} else {
|
||||
weightGM.SetL2CacheHint(CacheMode::CACHE_MODE_NORMAL);
|
||||
}
|
||||
mnConfig.workSpaceOffset = outOffset + mnConfig.yBaseOffset;
|
||||
mm.SetOrgShape(mnConfig.m, mnConfig.n, mnConfig.k);
|
||||
mm.SetSingleShape(curSingleM, curSingleN, mnConfig.k);
|
||||
mm.SetTensorA(xGM[xOffset], transposeX);
|
||||
mm.SetTensorB(weightGM, transposeW);
|
||||
mm.template IterateAll<sync>(mmOutGM[mnConfig.workSpaceOffset], 0);
|
||||
}
|
||||
|
||||
template <typename mmType, bool sync, typename CHANNELDTYPE>
|
||||
__aicore__ inline void GMMSwigluSplitWorkSpaceCompute<mmType, sync, CHANNELDTYPE>::UpdateMnConfig(MNConfig &mnConfig)
|
||||
{
|
||||
if constexpr (B::format == CubeFormat::NZ) {
|
||||
mnConfig.wBaseOffset += AlignUp<16>(mnConfig.k) * AlignUp<32>(mnConfig.n); // 16: nz format last two dim size
|
||||
} else {
|
||||
mnConfig.wBaseOffset += mnConfig.k * mnConfig.n;
|
||||
}
|
||||
mnConfig.nAxisBaseOffset += mnConfig.n;
|
||||
mnConfig.mAxisBaseOffset += mnConfig.m;
|
||||
mnConfig.xBaseOffset += mnConfig.m * mnConfig.k;
|
||||
mnConfig.yBaseOffset += mnConfig.m * mnConfig.n;
|
||||
}
|
||||
|
||||
template <typename mmType, bool sync, typename CHANNELDTYPE>
|
||||
__aicore__ inline void GMMSwigluSplitWorkSpaceCompute<mmType, sync, CHANNELDTYPE>::SetMNConfig(const int32_t splitValue,
|
||||
const uint32_t groupIdx,
|
||||
MNConfig &mnConfig)
|
||||
{
|
||||
SetMKN(splitValue, groupIdx, mnConfig);
|
||||
mnConfig.baseM = BASIC_M;
|
||||
mnConfig.baseN = BASIC_N;
|
||||
mnConfig.singleM = SINGLE_CORE_M;
|
||||
mnConfig.singleN = SINGLE_CORE_N;
|
||||
}
|
||||
|
||||
template <typename mmType, bool sync, typename CHANNELDTYPE>
|
||||
__aicore__ inline void GMMSwigluSplitWorkSpaceCompute<mmType, sync, CHANNELDTYPE>::SetMKN(const int32_t splitValue,
|
||||
const uint32_t groupIdx,
|
||||
MNConfig &mnConfig)
|
||||
{
|
||||
mnConfig.m = static_cast<int64_t>(splitValue);
|
||||
mnConfig.k = gmmBaseParams->K; // tilingData
|
||||
mnConfig.n = gmmBaseParams->N; // tilingData
|
||||
}
|
||||
|
||||
template <typename mmType, bool sync, typename CHANNELDTYPE>
|
||||
__aicore__ inline uint64_t GMMSwigluSplitWorkSpaceCompute<mmType, sync, CHANNELDTYPE>::GetWOffset(uint32_t tailN,
|
||||
uint32_t k)
|
||||
{
|
||||
uint64_t wOffset = 0;
|
||||
if constexpr (mmType::BT::format == CubeFormat::NZ) {
|
||||
wOffset = tailN * AlignUp<16>(k); // 16: nz format last two dim size
|
||||
} else {
|
||||
wOffset = tailN;
|
||||
}
|
||||
return wOffset;
|
||||
}
|
||||
|
||||
template <typename mmType, bool sync, typename CHANNELDTYPE>
|
||||
__aicore__ inline void GMMSwigluSplitWorkSpaceCompute<mmType, sync, CHANNELDTYPE>::MNBlockIdxCompute(
|
||||
MNConfig &mnConfig, const uint32_t curBlock, const uint32_t count, const uint32_t thresholdM_dimN)
|
||||
{
|
||||
mnConfig.mIdx = (curBlock - count) / mnConfig.blockDimN;
|
||||
mnConfig.nIdx = (curBlock - count) % mnConfig.blockDimN;
|
||||
}
|
||||
|
||||
template <typename mmType, bool sync, typename CHANNELDTYPE>
|
||||
__aicore__ inline void GMMSwigluSplitWorkSpaceCompute<mmType, sync, CHANNELDTYPE>::UpdateVecConfig(uint32_t blockIdx,
|
||||
VecConfig &vecConfig)
|
||||
{
|
||||
// 第一步 读取grouplist reduceSum 计算总数据个数
|
||||
vecConfig.M =
|
||||
workspaceSplitConfig.isLastLoop ? workspaceSplitConfig.lastLoopTaskSize : workspaceSplitConfig.notLastTaskSize;
|
||||
// 第二步 计算分核
|
||||
uint32_t eachCoreTaskNum = (vecConfig.M + aivCoreNum - 1) / aivCoreNum;
|
||||
vecConfig.usedCoreNum = vecConfig.M >= aivCoreNum ? aivCoreNum : vecConfig.M;
|
||||
uint32_t tailCoreIdx = vecConfig.M - (eachCoreTaskNum - 1) * vecConfig.usedCoreNum;
|
||||
vecConfig.taskNum = blockIdx < tailCoreIdx ? eachCoreTaskNum : eachCoreTaskNum - 1;
|
||||
vecConfig.startIdx =
|
||||
blockIdx < tailCoreIdx ? eachCoreTaskNum * blockIdx : ((eachCoreTaskNum - 1) * blockIdx + tailCoreIdx);
|
||||
vecConfig.curIdx = vecConfig.startIdx;
|
||||
vecConfig.startOffset = vecConfig.startIdx * gmmSwiglu->tokenLen;
|
||||
vecConfig.curOffset = vecConfig.startOffset;
|
||||
int64_t curStartIdx = vecConfig.startIdx;
|
||||
int64_t prevM = workspaceSplitConfig.leftMatrixStartIndex;
|
||||
for (uint32_t groupIdx = workspaceSplitConfig.rightMatrixExpertStartIndex;
|
||||
groupIdx <= workspaceSplitConfig.rightMatrixExpertEndIndex; groupIdx++) {
|
||||
int64_t currM = groupListGM.GetValue(groupIdx);
|
||||
int64_t tempM = currM - prevM;
|
||||
prevM = currM;
|
||||
if (curStartIdx >= 0 && curStartIdx - tempM < 0) {
|
||||
vecConfig.curGroupIdx = groupIdx;
|
||||
vecConfig.nextUpadteInterVal = tempM - curStartIdx;
|
||||
}
|
||||
curStartIdx -= tempM;
|
||||
}
|
||||
// 第三步 计算总数据量
|
||||
vecConfig.outLoopNum = (vecConfig.taskNum + gmmSwiglu->maxProcessRowNum - 1) / gmmSwiglu->maxProcessRowNum;
|
||||
vecConfig.tailLoopNum = vecConfig.taskNum % gmmSwiglu->maxProcessRowNum ?
|
||||
vecConfig.taskNum % gmmSwiglu->maxProcessRowNum :
|
||||
gmmSwiglu->maxProcessRowNum;
|
||||
pipe->Reset();
|
||||
// 第四步 申请空间
|
||||
pipe->InitBuffer(mmOutQueue, DOUBLE_BUFFER, gmmSwiglu->maxProcessRowNum * gmmSwiglu->tokenLen * sizeof(int32_t));
|
||||
pipe->InitBuffer(perChannelScaleInQueue, DOUBLE_BUFFER, gmmSwiglu->tokenLen * sizeof(float));
|
||||
pipe->InitBuffer(quantOutQueue, DOUBLE_BUFFER,
|
||||
gmmSwiglu->maxProcessRowNum * gmmSwiglu->tokenLen / SWIGLU_REDUCE_FACTOR * sizeof(int8_t));
|
||||
pipe->InitBuffer(quantScaleOutQueue, DOUBLE_BUFFER,
|
||||
AlignUp<int32_t>(gmmSwiglu->maxProcessRowNum, ALIGN_8_ELE) * sizeof(float));
|
||||
// two 32 byte buffer for reduceMax calculation in Quant.
|
||||
pipe->InitBuffer(reduceWorkspace, gmmSwiglu->tokenLen / SWIGLU_REDUCE_FACTOR * sizeof(float) + UB_BLOCK_UNIT_SIZE +
|
||||
UB_BLOCK_UNIT_SIZE);
|
||||
}
|
||||
|
||||
template <typename mmType, bool sync, typename CHANNELDTYPE>
|
||||
__aicore__ inline void GMMSwigluSplitWorkSpaceCompute<mmType, sync, CHANNELDTYPE>::customDataCopyIn(
|
||||
uint32_t outLoopIdx, GlobalTensor<int32_t> &mmOutGM, VecConfig &vecConfig)
|
||||
{
|
||||
LocalTensor<int32_t> _inMMLocal_0 = mmOutQueue.DeQue<int32_t>();
|
||||
DataCopyExtParams copyParams_0{
|
||||
1, static_cast<uint32_t>(vecConfig.innerLoopNum * gmmSwiglu->tokenLen * sizeof(int32_t)), 0, 0, 0};
|
||||
DataCopyPadExtParams<int32_t> padParams_0{false, 0, 0, 0};
|
||||
DataCopyPad(_inMMLocal_0, mmOutGM[vecConfig.curOffset], copyParams_0, padParams_0);
|
||||
mmOutQueue.EnQue(_inMMLocal_0);
|
||||
|
||||
LocalTensor<int32_t> _inMMLocal_1 = mmOutQueue.DeQue<int32_t>();
|
||||
|
||||
Cast(_inMMLocal_1.ReinterpretCast<float>(), _inMMLocal_1, RoundMode::CAST_NONE,
|
||||
vecConfig.innerLoopNum * gmmSwiglu->tokenLen);
|
||||
|
||||
mmOutQueue.EnQue(_inMMLocal_1);
|
||||
LocalTensor<float> _inMMLocal_2 = mmOutQueue.DeQue<float>();
|
||||
int32_t eventIdSToV = static_cast<int32_t>(GetTPipePtr()->FetchEventID(HardEvent::S_V));
|
||||
SetFlag<HardEvent::S_V>(eventIdSToV);
|
||||
for (uint32_t i = 0; i < vecConfig.innerLoopNum; i++) {
|
||||
WaitFlag<HardEvent::S_V>(eventIdSToV);
|
||||
float scale = perTokenScaleGM.GetValue(vecConfig.curIdx + workspaceSplitConfig.leftMatrixStartIndex);
|
||||
SetFlag<HardEvent::S_V>(eventIdSToV);
|
||||
WaitFlag<HardEvent::S_V>(eventIdSToV);
|
||||
Muls(_inMMLocal_2[i * gmmSwiglu->tokenLen], _inMMLocal_2[i * gmmSwiglu->tokenLen], scale, gmmSwiglu->tokenLen);
|
||||
SetFlag<HardEvent::S_V>(eventIdSToV);
|
||||
vecConfig.curIdx++;
|
||||
}
|
||||
WaitFlag<HardEvent::S_V>(eventIdSToV);
|
||||
vecConfig.curOffset = vecConfig.curIdx * gmmSwiglu->tokenLen;
|
||||
mmOutQueue.EnQue(_inMMLocal_2);
|
||||
}
|
||||
|
||||
template <typename mmType, bool sync, typename CHANNELDTYPE>
|
||||
template <typename DTYPE_CS>
|
||||
__aicore__ inline void
|
||||
GMMSwigluSplitWorkSpaceCompute<mmType, sync, CHANNELDTYPE>::UpdateChannelScale(uint32_t loopIdx, VecConfig &vecConfig)
|
||||
{
|
||||
// 更新perChannel
|
||||
if (unlikely(vecConfig.nextUpadteInterVal == 0)) {
|
||||
int64_t loop = gmmSwiglu->groupListLen - vecConfig.curGroupIdx;
|
||||
while (loop--) {
|
||||
int64_t curTemp = groupListGM.GetValue(vecConfig.curGroupIdx);
|
||||
vecConfig.curGroupIdx++;
|
||||
int64_t nextTemp = groupListGM.GetValue(vecConfig.curGroupIdx);
|
||||
if (nextTemp != curTemp) {
|
||||
vecConfig.nextUpadteInterVal = nextTemp - curTemp;
|
||||
break;
|
||||
}
|
||||
}
|
||||
LocalTensor<float> _inChannel = perChannelScaleInQueue.DeQue<float>();
|
||||
DataCopyExtParams copyParams{1, static_cast<uint32_t>(gmmSwiglu->tokenLen * sizeof(DTYPE_CS)), 0, 0, 0};
|
||||
DataCopyPadExtParams<DTYPE_CS> padParams{false, 0, 0, 0};
|
||||
if constexpr (!IsSameType<DTYPE_CS, float>::value) {
|
||||
LocalTensor<DTYPE_CS> dstLocalT = _inChannel.template ReinterpretCast<DTYPE_CS>();
|
||||
DataCopyPad(dstLocalT[gmmSwiglu->tokenLen], perChannelScaleGM[vecConfig.curGroupIdx * gmmSwiglu->tokenLen],
|
||||
copyParams, padParams);
|
||||
int32_t eventIdMTE2ToV = static_cast<int32_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V));
|
||||
SetFlag<HardEvent::MTE2_V>(eventIdMTE2ToV);
|
||||
WaitFlag<HardEvent::MTE2_V>(eventIdMTE2ToV);
|
||||
Cast(_inChannel, dstLocalT[gmmSwiglu->tokenLen], RoundMode::CAST_NONE, gmmSwiglu->tokenLen);
|
||||
} else {
|
||||
DataCopyPad(_inChannel, perChannelScaleGM[vecConfig.curGroupIdx * gmmSwiglu->tokenLen], copyParams,
|
||||
padParams);
|
||||
}
|
||||
perChannelScaleInQueue.EnQue(_inChannel);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename mmType, bool sync, typename CHANNELDTYPE>
|
||||
__aicore__ inline void GMMSwigluSplitWorkSpaceCompute<mmType, sync, CHANNELDTYPE>::VectorCompute(uint32_t loopIdx,
|
||||
VecConfig &vecConfig)
|
||||
{
|
||||
Dequant(loopIdx, vecConfig);
|
||||
Swiglu(loopIdx, vecConfig);
|
||||
Quant(loopIdx, vecConfig);
|
||||
}
|
||||
|
||||
template <typename mmType, bool sync, typename CHANNELDTYPE>
|
||||
__aicore__ inline void GMMSwigluSplitWorkSpaceCompute<mmType, sync, CHANNELDTYPE>::Dequant(uint32_t loopIdx,
|
||||
VecConfig &vecConfig)
|
||||
{
|
||||
// perChanelScale * perTokenScale
|
||||
LocalTensor<float> mmLocal = mmOutQueue.DeQue<float>();
|
||||
LocalTensor<float> perChannelLocal = perChannelScaleInQueue.DeQue<float>();
|
||||
Mul(mmLocal[loopIdx * gmmSwiglu->tokenLen], mmLocal[loopIdx * gmmSwiglu->tokenLen], perChannelLocal,
|
||||
gmmSwiglu->tokenLen);
|
||||
vecConfig.nextUpadteInterVal--;
|
||||
mmOutQueue.EnQue(mmLocal);
|
||||
perChannelScaleInQueue.EnQue(perChannelLocal);
|
||||
}
|
||||
|
||||
template <typename mmType, bool sync, typename CHANNELDTYPE>
|
||||
__aicore__ inline void GMMSwigluSplitWorkSpaceCompute<mmType, sync, CHANNELDTYPE>::Swiglu(uint32_t loopIdx,
|
||||
VecConfig &vecConfig)
|
||||
{
|
||||
// 高阶API swiglu
|
||||
LocalTensor<float> _inMMLocal = mmOutQueue.DeQue<float>();
|
||||
float beta = 1.0f;
|
||||
LocalTensor<float> workspaceLocal = reduceWorkspace.Get<float>();
|
||||
LocalTensor<float> src0Local =
|
||||
_inMMLocal[loopIdx * gmmSwiglu->tokenLen + gmmSwiglu->tokenLen / SWIGLU_REDUCE_FACTOR];
|
||||
LocalTensor<float> src1Local = _inMMLocal[loopIdx * gmmSwiglu->tokenLen];
|
||||
if (limited > 0.0f) {
|
||||
Mins(src0Local, src0Local, limited, gmmSwiglu->tokenLen / 2);
|
||||
PipeBarrier<PIPE_V>();
|
||||
Maxs(src0Local, src0Local, (-1.0f * limited), gmmSwiglu->tokenLen / 2);
|
||||
PipeBarrier<PIPE_V>();
|
||||
Mins(src1Local, src1Local, limited, gmmSwiglu->tokenLen / 2);
|
||||
PipeBarrier<PIPE_V>();
|
||||
}
|
||||
SwiGLU<float, false>(workspaceLocal, src0Local, src1Local, beta, gmmSwiglu->tokenLen / SWIGLU_REDUCE_FACTOR);
|
||||
PipeBarrier<PIPE_V>();
|
||||
DataCopyParams repeatParams{1, static_cast<uint16_t>((gmmSwiglu->tokenLen / SWIGLU_REDUCE_FACTOR) / ALIGN_8_ELE), 0,
|
||||
0};
|
||||
DataCopy(_inMMLocal[loopIdx * gmmSwiglu->tokenLen], workspaceLocal, repeatParams);
|
||||
mmOutQueue.EnQue(_inMMLocal);
|
||||
}
|
||||
|
||||
template <typename mmType, bool sync, typename CHANNELDTYPE>
|
||||
__aicore__ inline void GMMSwigluSplitWorkSpaceCompute<mmType, sync, CHANNELDTYPE>::Quant(uint32_t loopIdx,
|
||||
VecConfig &vecConfig)
|
||||
{
|
||||
LocalTensor<float> _inMMLocal = mmOutQueue.DeQue<float>();
|
||||
uint64_t preOffset = loopIdx * gmmSwiglu->tokenLen;
|
||||
uint64_t halfTokenLen = gmmSwiglu->tokenLen / BISECT;
|
||||
Abs(_inMMLocal[preOffset + gmmSwiglu->tokenLen / BISECT], _inMMLocal[preOffset], halfTokenLen);
|
||||
PipeBarrier<PIPE_V>();
|
||||
// reduceMax
|
||||
LocalTensor<float> workLocal = reduceWorkspace.Get<float>(halfTokenLen);
|
||||
LocalTensor<float> reduceResLocal =
|
||||
reduceWorkspace.GetWithOffset<float>(FLOAT_UB_BLOCK_UNIT_SIZE, halfTokenLen * sizeof(float));
|
||||
LocalTensor<float> reduceTmpLocal = reduceWorkspace.GetWithOffset<float>(
|
||||
FLOAT_UB_BLOCK_UNIT_SIZE, halfTokenLen * sizeof(float) + UB_BLOCK_UNIT_SIZE);
|
||||
ReduceMaxTemplate(reduceResLocal, workLocal, _inMMLocal[preOffset + gmmSwiglu->tokenLen / BISECT], reduceTmpLocal,
|
||||
static_cast<uint32_t>(halfTokenLen));
|
||||
|
||||
int32_t eventIdVToS = static_cast<int32_t>(GetTPipePtr()->FetchEventID(HardEvent::V_S));
|
||||
SetFlag<HardEvent::V_S>(eventIdVToS);
|
||||
WaitFlag<HardEvent::V_S>(eventIdVToS);
|
||||
float quantScale = reduceResLocal.GetValue(0) / QUANT_SCALE_INT8;
|
||||
LocalTensor<float> quantScaleLocal = quantScaleOutQueue.DeQue<float>();
|
||||
quantScaleLocal.SetValue(loopIdx, quantScale);
|
||||
quantScale = 1 / quantScale;
|
||||
int32_t eventIdSToV = static_cast<int32_t>(GetTPipePtr()->FetchEventID(HardEvent::S_V));
|
||||
SetFlag<HardEvent::S_V>(eventIdSToV);
|
||||
WaitFlag<HardEvent::S_V>(eventIdSToV);
|
||||
Muls(_inMMLocal[preOffset], _inMMLocal[preOffset], quantScale, halfTokenLen);
|
||||
PipeBarrier<PIPE_V>();
|
||||
LocalTensor<int8_t> quantLocal = quantOutQueue.DeQue<int8_t>();
|
||||
int32_t dstTempOffset = static_cast<int32_t>(preOffset / BISECT);
|
||||
int32_t srcTempOffset = static_cast<int32_t>(preOffset);
|
||||
int32_t tempCount = static_cast<int32_t>(halfTokenLen);
|
||||
LocalTensor<int8_t> castSpace = reduceWorkspace.Get<int8_t>(UB_BLOCK_UNIT_SIZE);
|
||||
CastFp32ToInt8Template(quantLocal, _inMMLocal, castSpace, dstTempOffset, srcTempOffset, tempCount);
|
||||
mmOutQueue.EnQue(_inMMLocal);
|
||||
quantOutQueue.EnQue(quantLocal);
|
||||
}
|
||||
|
||||
template <typename mmType, bool sync, typename CHANNELDTYPE>
|
||||
__aicore__ inline void
|
||||
GMMSwigluSplitWorkSpaceCompute<mmType, sync, CHANNELDTYPE>::customDataCopyOut(VecConfig &vecConfig)
|
||||
{
|
||||
LocalTensor<float> quantScaleLocal = quantScaleOutQueue.DeQue<float>();
|
||||
DataCopyParams copyParams_0{1, (uint16_t)(vecConfig.innerLoopNum * sizeof(float)), 0, 0};
|
||||
DataCopyPad(quantScaleOutputGM[workspaceSplitConfig.leftMatrixStartIndex + vecConfig.startIdx], quantScaleLocal,
|
||||
copyParams_0);
|
||||
LocalTensor<int8_t> quantLocal = quantOutQueue.DeQue<int8_t>();
|
||||
DataCopyParams copyParams_1{
|
||||
1, (uint16_t)(vecConfig.innerLoopNum * gmmSwiglu->tokenLen / SWIGLU_REDUCE_FACTOR * sizeof(int8_t)), 0, 0};
|
||||
DataCopyPad(quantOutputGM[(workspaceSplitConfig.leftMatrixStartIndex + vecConfig.startIdx) * gmmSwiglu->tokenLen /
|
||||
SWIGLU_REDUCE_FACTOR],
|
||||
quantLocal, copyParams_1);
|
||||
vecConfig.startIdx += vecConfig.innerLoopNum;
|
||||
vecConfig.startOffset = vecConfig.startIdx * gmmSwiglu->tokenLen;
|
||||
quantOutQueue.EnQue(quantLocal);
|
||||
quantScaleOutQueue.EnQue(quantScaleLocal);
|
||||
}
|
||||
|
||||
} // namespace GROUPED_MATMUL_SWIGLU_QUANT
|
||||
#endif // ASCENDC_GROUPED_MATMUL_SWIGLU_QUANT_SPLIT_WS_H
|
||||
@@ -0,0 +1,332 @@
|
||||
/**
|
||||
* Copyright (c) 2025 Huawei Technologies Co., Ltd.
|
||||
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
* CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
* See LICENSE in the root of the software repository for the full text of the License.
|
||||
*/
|
||||
|
||||
/*!
|
||||
* \file grouped_matmul_swiglu_quant_utils.h
|
||||
* \brief
|
||||
*/
|
||||
#ifndef ASCENDC_GROUPED_MATMUL_SWIGLU_QUANT_UTILS_H
|
||||
#define ASCENDC_GROUPED_MATMUL_SWIGLU_QUANT_UTILS_H
|
||||
|
||||
#include "kernel_tiling/kernel_tiling.h"
|
||||
#include "kernel_operator.h"
|
||||
#include "lib/matmul_intf.h"
|
||||
|
||||
#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 220
|
||||
// A8W4 MSD场景
|
||||
#if defined(ORIG_DTYPE_X) && defined(DT_INT8) && ORIG_DTYPE_X == DT_INT8 && defined(ORIG_DTYPE_WEIGHT) && \
|
||||
defined(DT_INT4) && ORIG_DTYPE_WEIGHT == DT_INT4
|
||||
#define GMM_SWIGLU_QUANT_A8W4_MSD
|
||||
using DTYPE_X_A8W4_MSD = AscendC::int4b_t;
|
||||
// A8W8 场景
|
||||
#elif defined(ORIG_DTYPE_X) && defined(DT_INT8) && ORIG_DTYPE_X == DT_INT8 && defined(ORIG_DTYPE_WEIGHT) && \
|
||||
defined(DT_INT8) && ORIG_DTYPE_WEIGHT == DT_INT8
|
||||
#define GMM_SWIGLU_QUANT_A8W8
|
||||
#endif // 场景分类
|
||||
|
||||
#if defined(FORMAT_WEIGHT) && FORMAT_WEIGHT == FORMAT_FRACTAL_NZ
|
||||
constexpr CubeFormat wFormat = CubeFormat::NZ;
|
||||
#elif defined(FORMAT_WEIGHT) && FORMAT_WEIGHT == FORMAT_ND
|
||||
constexpr CubeFormat wFormat = CubeFormat::ND;
|
||||
#endif // weight格式分类
|
||||
|
||||
#endif // 芯片型号分类
|
||||
|
||||
namespace GROUPED_MATMUL_SWIGLU_QUANT {
|
||||
using namespace AscendC;
|
||||
constexpr uint32_t INT8_BITS = 8; // a int8 number has 8 bits
|
||||
constexpr uint32_t UB_BLOCK_UNIT_SIZE = 32; // 32: a block has 32 bytes data
|
||||
constexpr uint32_t THRESHOLD_BLOCK_NUM = 8;
|
||||
constexpr uint32_t UB_BLOCK_DOUBLE_UNIT_SIZE = 64; // 64: a block has 64 bytes data
|
||||
constexpr uint32_t HALF_UB_BLOCK_UNIT_SIZE = UB_BLOCK_UNIT_SIZE / 2; // 2: a float16 data has two bytes
|
||||
constexpr uint32_t FLOAT_UB_BLOCK_UNIT_SIZE = 8; // 2: a float16 data has two bytes
|
||||
constexpr uint32_t SINGLE_CORE_M = 128;
|
||||
constexpr uint32_t SINGLE_CORE_N = 256;
|
||||
constexpr uint32_t SINGLE_CORE_K = 7168;
|
||||
constexpr uint32_t BASIC_M = 128;
|
||||
constexpr uint32_t BASIC_N = 256;
|
||||
constexpr uint32_t BASIC_K = 128;
|
||||
constexpr uint32_t STEP_M = 1;
|
||||
constexpr uint32_t STEP_N = 1;
|
||||
constexpr uint32_t STEP_Ka = 4;
|
||||
constexpr uint32_t STEP_Kb = 4;
|
||||
constexpr uint32_t DEPTH_A1 = 8;
|
||||
constexpr uint32_t DEPTH_B1 = 8;
|
||||
constexpr uint32_t VEC_LEN_ONCE_REPEAT_ELE = 64;
|
||||
constexpr uint32_t VEC_LEN_ONCE_REPEAT_BLOCK = 8;
|
||||
constexpr uint32_t FP32_LEN_64_REPEAT = 4096;
|
||||
constexpr uint32_t REPEAT_64 = 64;
|
||||
constexpr uint32_t REPEAT_8 = 8;
|
||||
constexpr uint32_t BISECT = 2;
|
||||
constexpr uint32_t MOD_32_MASK = 0x1F;
|
||||
constexpr uint32_t MOD_16_MASK = 0x0F;
|
||||
constexpr uint32_t ALIGN_8_ELE = 8;
|
||||
constexpr uint32_t ALIGN_16_ELE = 16;
|
||||
constexpr float QUANT_SCALE_INT8 = 127.0f;
|
||||
constexpr int64_t SWIGLU_REDUCE_FACTOR = 2;
|
||||
constexpr int64_t DOUBLE_BUFFER = 2;
|
||||
constexpr uint8_t NUM_8 = 8;
|
||||
constexpr bool NO_BIAS = false;
|
||||
constexpr int64_t DOUBLE_ROW = 2;
|
||||
constexpr MatmulConfig CUSTOM_CFG_MDL = GetMDLConfig(false, false, 0, true, false, false, true);
|
||||
constexpr MatmulConfig GetMMStaticCFG()
|
||||
{
|
||||
MatmulConfig MM_CFG = CUSTOM_CFG_MDL;
|
||||
MM_CFG.singleCoreM = SINGLE_CORE_M;
|
||||
MM_CFG.singleCoreN = SINGLE_CORE_N;
|
||||
MM_CFG.singleCoreK = SINGLE_CORE_K;
|
||||
MM_CFG.basicM = BASIC_M;
|
||||
MM_CFG.basicN = BASIC_N;
|
||||
MM_CFG.basicK = BASIC_K;
|
||||
return MM_CFG;
|
||||
}
|
||||
|
||||
constexpr static MatmulApiStaticTiling GetMMTiling(const MatmulApiStaticTiling &mmTiling)
|
||||
{
|
||||
MatmulApiStaticTiling tiling = mmTiling;
|
||||
tiling.stepM = STEP_M;
|
||||
tiling.stepN = STEP_N;
|
||||
tiling.stepKa = STEP_Ka;
|
||||
tiling.stepKb = STEP_Kb;
|
||||
tiling.depthA1 = DEPTH_A1;
|
||||
tiling.depthB1 = DEPTH_B1;
|
||||
tiling.isBias = NO_BIAS;
|
||||
return tiling;
|
||||
}
|
||||
|
||||
template <class AT_, class BT_, class CT_>
|
||||
struct MMImplTypeStatic {
|
||||
using AT = AT_;
|
||||
using BT = BT_;
|
||||
using CT = CT_;
|
||||
// bias未被使用但高阶模板参数需要传入
|
||||
using BiasT = MatmulType<AscendC::TPosition::GM, CubeFormat::ND, int32_t>;
|
||||
static constexpr MatmulConfig cfg = GetMMStaticCFG();
|
||||
static constexpr MatmulApiStaticTiling mdl = GetMMTiling(GetMatmulApiTiling<AT, BT, CT, BiasT>(cfg));
|
||||
using MT = matmul::MatmulImpl<AT, BT, CT, BiasT, mdl>;
|
||||
};
|
||||
|
||||
template <class AT_, class BT_, class CT_>
|
||||
struct MMImplType {
|
||||
using AT = AT_;
|
||||
using BT = BT_;
|
||||
using CT = CT_;
|
||||
// bias未被使用但高阶模板参数需要传入
|
||||
using BiasT = MatmulType<AscendC::TPosition::GM, CubeFormat::ND, int32_t>;
|
||||
using MT = matmul::MatmulImpl<AT, BT, CT, BiasT, CUSTOM_CFG_MDL>;
|
||||
};
|
||||
|
||||
struct MNConfig {
|
||||
int64_t m = 0;
|
||||
int64_t k = 0;
|
||||
int64_t n = 0;
|
||||
int64_t baseM = 0;
|
||||
int64_t baseN = 0;
|
||||
int64_t mIdx = 0;
|
||||
int64_t nIdx = 0;
|
||||
int64_t blockDimM = 0;
|
||||
int64_t blockDimN = 0;
|
||||
int64_t singleM = 0;
|
||||
int64_t singleN = 0;
|
||||
int64_t wBaseOffset = 0;
|
||||
int64_t nAxisBaseOffset = 0;
|
||||
int64_t mAxisBaseOffset = 0;
|
||||
int64_t xBaseOffset = 0;
|
||||
int64_t yBaseOffset = 0;
|
||||
int64_t wOutOffset = 0;
|
||||
int64_t workSpaceOffset = 0;
|
||||
};
|
||||
|
||||
struct VecConfig {
|
||||
int64_t M = 0;
|
||||
int64_t usedCoreNum = 0;
|
||||
int64_t startOffset = 0;
|
||||
int64_t curOffset = 0;
|
||||
int64_t startIdx = 0;
|
||||
int64_t curIdx = 0;
|
||||
int64_t taskNum = 0;
|
||||
int64_t curGroupIdx = 0;
|
||||
int64_t outLoopNum = 0;
|
||||
int64_t innerLoopNum = 0;
|
||||
int64_t tailLoopNum = 0;
|
||||
int64_t nextUpadteInterVal = 0;
|
||||
};
|
||||
|
||||
struct WorkSpaceSplitConfig {
|
||||
int64_t M = 0;
|
||||
int64_t loopCount = 0;
|
||||
int64_t leftMatrixStartIndex = 0;
|
||||
int64_t rightMatrixExpertStartIndex = 0;
|
||||
int64_t rightMatrixExpertNextStartIndex = 0;
|
||||
int64_t rightMatrixExpertEndIndex = 0;
|
||||
int64_t notLastTaskSize = 0;
|
||||
int64_t lastLoopTaskSize = 0;
|
||||
bool isLastLoop = false;
|
||||
};
|
||||
|
||||
struct GMAddrParams {
|
||||
// 输入 GM Tensor
|
||||
GM_ADDR xGM; // 左矩阵
|
||||
GM_ADDR weightGM; // 右矩阵
|
||||
GM_ADDR weightScaleGM; // 权重scale
|
||||
GM_ADDR xScaleGM; // 激活scale
|
||||
GM_ADDR weightAuxiliaryMatrixGM; // 权重辅助矩阵
|
||||
GM_ADDR groupListGM; // 分组矩阵
|
||||
// 输出 GM Tensor
|
||||
GM_ADDR yGM; // 输出量化矩阵
|
||||
GM_ADDR yScaleGM; // 输出scale矩阵
|
||||
// workspace GM Tensor
|
||||
GM_ADDR workSpaceGM; // 左矩阵前处理结果矩阵 (double workspace) + 中间处理结果矩阵 (double workspace)
|
||||
int64_t workSpaceOffset1;
|
||||
int64_t workSpaceOffset2;
|
||||
int64_t workSpaceOffset3;
|
||||
};
|
||||
|
||||
template <uint32_t base, typename T = uint32_t>
|
||||
__aicore__ inline auto AlignUp(T a) -> T
|
||||
{
|
||||
if (unlikely(base == 0)) {
|
||||
return a;
|
||||
}
|
||||
return (a + base - 1) / base * base;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__aicore__ inline auto AlignUp(T a, T base) -> T
|
||||
{
|
||||
if (unlikely(base == 0)) {
|
||||
return a;
|
||||
}
|
||||
return (a + base - 1) / base * base;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__aicore__ inline auto AlignDown(T a, T base) -> T
|
||||
{
|
||||
if (unlikely(base == 0)) {
|
||||
return a;
|
||||
}
|
||||
return a / base * base;
|
||||
}
|
||||
|
||||
template <>
|
||||
__aicore__ inline uint32_t AlignUp<4, uint32_t>(uint32_t a)
|
||||
{
|
||||
// to be Multiple of 4, result should be in a format of b(xxxx,x100).
|
||||
// This means last two bits should be zero, requiring that
|
||||
// result = num & b(1111,1100) = num & (~3).
|
||||
// &(~3) operator may reduces num into the range [num, num - 3].
|
||||
// As the result should be no less than a (result >= a), it means num - 3 >= a in the worst case.
|
||||
// In this case, num >= a+3. On the other hand, num should also be less then a+4, otherwise,
|
||||
// the result will not be least multiple of 4 for 3. In other cases like [num, num - 2],
|
||||
// num = a + 3 also satisfies the goal condition.
|
||||
return (a + 3) & ~3; // & ~3: set last two bits of (a+3) to be zero
|
||||
}
|
||||
|
||||
template <>
|
||||
__aicore__ inline uint32_t AlignUp<8, uint32_t>(uint32_t a)
|
||||
{
|
||||
// In general, if we want to get the least multiple of b (b is the power of 2) for a,
|
||||
// it comes to a conclusion from the above comment: result = (a + (b - 1)) & (~b)
|
||||
return (a + 7) & ~7; // & ~7: set last four bits of (a+7) to be zero
|
||||
}
|
||||
|
||||
template <>
|
||||
__aicore__ inline uint32_t AlignUp<16, uint32_t>(uint32_t a)
|
||||
{
|
||||
// In general, if we want to get the least multiple of b (b is the power of 2) for a,
|
||||
// it comes to a conclusion from the above comment: result = (a + (b - 1)) & (~b)
|
||||
return (a + 15) & ~15; // & ~15: set last four bits of (a+15) to be zero
|
||||
}
|
||||
|
||||
template <>
|
||||
__aicore__ inline uint32_t AlignUp<32, uint32_t>(uint32_t a)
|
||||
{
|
||||
// refer to the above comments.
|
||||
return (a + 31) & ~31; // & ~31: set last five bits of (a+31) to be zero}
|
||||
}
|
||||
|
||||
__aicore__ inline void ReduceMaxSmall(const LocalTensor<float> &dstLocal, const LocalTensor<float> &workLocal,
|
||||
const LocalTensor<float> &srcLocal, uint32_t count)
|
||||
{
|
||||
/**
|
||||
* @brief ReduceMaxSmall 此函数仅支持入参count小于4096。
|
||||
*/
|
||||
uint32_t repeat = count / VEC_LEN_ONCE_REPEAT_ELE;
|
||||
uint32_t tailNum = count % VEC_LEN_ONCE_REPEAT_ELE;
|
||||
if (likely(repeat > 0)) {
|
||||
WholeReduceMax(workLocal, srcLocal, VEC_LEN_ONCE_REPEAT_ELE, repeat, 1, 1, VEC_LEN_ONCE_REPEAT_BLOCK,
|
||||
ReduceOrder::ORDER_ONLY_VALUE);
|
||||
PipeBarrier<PIPE_V>();
|
||||
}
|
||||
if (unlikely(tailNum != 0)) {
|
||||
WholeReduceMax(workLocal[repeat], srcLocal[count - tailNum], tailNum, 1, 1, 1, VEC_LEN_ONCE_REPEAT_BLOCK,
|
||||
ReduceOrder::ORDER_ONLY_VALUE);
|
||||
PipeBarrier<PIPE_V>();
|
||||
repeat += 1;
|
||||
}
|
||||
WholeReduceMax(dstLocal, workLocal, repeat, 1, 1, 1, VEC_LEN_ONCE_REPEAT_BLOCK, ReduceOrder::ORDER_ONLY_VALUE);
|
||||
}
|
||||
|
||||
__aicore__ inline void ReduceMaxTemplate(const LocalTensor<float> &dstLocal, const LocalTensor<float> &workLocal,
|
||||
const LocalTensor<float> &srcLocal, const LocalTensor<float> &resTmpLocal,
|
||||
uint32_t count)
|
||||
{
|
||||
/**
|
||||
* @brief 当前算子仅支持[32, 10240]长度的词向量维度N，对应此函数count入参范围在[16, 5120]。
|
||||
* @param [in] count: 本函数支持count范围为[1,8192]。
|
||||
*/
|
||||
if (count <= FP32_LEN_64_REPEAT) {
|
||||
ReduceMaxSmall(dstLocal, workLocal, srcLocal, count);
|
||||
PipeBarrier<PIPE_V>();
|
||||
} else {
|
||||
BlockReduceMax(workLocal, srcLocal, REPEAT_64, VEC_LEN_ONCE_REPEAT_ELE, 1, 1, VEC_LEN_ONCE_REPEAT_BLOCK);
|
||||
PipeBarrier<PIPE_V>();
|
||||
|
||||
BlockReduceMax(workLocal, workLocal, REPEAT_8, VEC_LEN_ONCE_REPEAT_ELE, 1, 1, VEC_LEN_ONCE_REPEAT_BLOCK);
|
||||
PipeBarrier<PIPE_V>();
|
||||
|
||||
WholeReduceMax(resTmpLocal, workLocal, VEC_LEN_ONCE_REPEAT_ELE, 1, 1, 1, VEC_LEN_ONCE_REPEAT_BLOCK,
|
||||
ReduceOrder::ORDER_ONLY_VALUE);
|
||||
PipeBarrier<PIPE_V>();
|
||||
|
||||
ReduceMaxSmall(dstLocal, workLocal, srcLocal[FP32_LEN_64_REPEAT], count - FP32_LEN_64_REPEAT);
|
||||
PipeBarrier<PIPE_V>();
|
||||
|
||||
const BinaryRepeatParams repeatParams = {1, 1, 1, NUM_8, NUM_8, NUM_8};
|
||||
Max(dstLocal, dstLocal, resTmpLocal, 1, 1, repeatParams);
|
||||
}
|
||||
}
|
||||
|
||||
__aicore__ inline void CastFp32ToInt8Template(LocalTensor<int8_t> &dstLocal, LocalTensor<float> &srcLocal,
|
||||
LocalTensor<int8_t> &oneBlockWorkspace, int32_t dstOffset,
|
||||
int32_t srcOffset, int32_t count)
|
||||
{
|
||||
Cast(srcLocal[srcOffset].ReinterpretCast<half>(), srcLocal[srcOffset], RoundMode::CAST_RINT, count);
|
||||
PipeBarrier<PIPE_V>();
|
||||
if ((dstOffset & MOD_32_MASK) == 0) {
|
||||
Cast(dstLocal[dstOffset], srcLocal[srcOffset].ReinterpretCast<half>(), RoundMode::CAST_RINT, count);
|
||||
} else if ((dstOffset & MOD_16_MASK) == 0) {
|
||||
Cast(dstLocal[dstOffset + ALIGN_16_ELE], srcLocal[srcOffset + ALIGN_8_ELE].ReinterpretCast<half>(),
|
||||
RoundMode::CAST_RINT, count - ALIGN_16_ELE);
|
||||
PipeBarrier<PIPE_V>();
|
||||
Cast(oneBlockWorkspace, srcLocal[srcOffset].ReinterpretCast<half>(), RoundMode::CAST_RINT, ALIGN_16_ELE);
|
||||
PipeBarrier<PIPE_ALL>();
|
||||
for (int32_t i = 0; i < ALIGN_16_ELE; i++) {
|
||||
int8_t temp = oneBlockWorkspace.GetValue(i);
|
||||
dstLocal.SetValue(dstOffset + i, temp);
|
||||
}
|
||||
PipeBarrier<PIPE_ALL>();
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace GROUPED_MATMUL_SWIGLU_QUANT
|
||||
|
||||
#endif // ASCENDC_GROUPED_MATMUL_UTILS_H
|
||||
Reference in New Issue
Block a user