描述

1.1 需求背景和描述
Transformer 里最经典的 FFN 结构：

本融合算子需要再FFN结构里面支持一个融合算子（下面蓝色部分）：

即QuantMatmul+Gelu+DynamicMxQuant(或者DynamicHif4Quant)融合成一个独立的算子。
将FNN计算过程中的QBMM、GELU、DynamicmxQuant三个单算子，融合合成一个算子，提升模型的推理性能。

2 约束和周边影响评估
2.1 支持芯片型号

3 算子功能及定义
3.1 算子功能

• 支持- 图模式
• 支持确定性计算
• 不涉及反向测试

3.2 接口定义
aclnn接口定义：

aclnnStatus QuantMatmulActivationQuantWeightNzGetWorkspaceSize(
const aclTensor *x1,
const aclTensor *x2,
const aclTensor *x1ScaleOptional,
const aclTensor *x2Scale,
const aclTensor *biasOptional,
bool transposeX1,
bool transposeX2,
int64_t groupSize
char *activationType,
char *quantMode,
char *roundMode,
int64_t scaleAlg,
double dstTypeMax,
const aclTensor *yOut,
const aclTensor *yScaleOut,
uint64_t *workspaceSize,
aclOpExecutor **executor);

aclnnStatus QuantMatmulActivationQuantWeightNz(
void *workspace,
uint64_t workspaceSize,
aclOpExecutor *executor,
aclrtStream stream);

Aten 接口定义：
func: quant_matmul_activation_quant(
"Tensor x1, Tensor x2, *, Tensor? x1_scale=None, Tensor? x2_scale=None, Tensor? bias=None, "
"int? x1_dtype=None, int? x2_dtype=None, int? x1scale_dtype=None, int? x2scale_dtype=None, "
"bool? transpose_x1=None, bool? transpose_x2=None, int[]? group_sizes=None, str activation_type="gelu_tanh", "
"int? y_dtype=None, str quant_mode="mx", str round_mode="rint", int? scale_alg=0, float dst_type_max=0.0)"
" -> (Tensor, Tensor)"

4 详细方案设计
4.1 计算逻辑
整体计算逻辑
算子核心逻辑包含矩阵乘法（Cube Core计算）与Vector计算。Matmul计算为qbmm，vector计算部分主要是gelu激活并进行mxdynamicquant输出得到量化的y与scale。

4.2 tiling方案
1）ASWT
实现ASWT（ASW Tiling）的切分的通用模板，针对weight较小场景的实现全载模板
将 MN 的输出空间划分为 baseMbaseN 的基本块（BasicBlock），以 4行窗口组（WINDOW_LEN=4）为调度单位，采用 S形走位（偶数行左到右，奇数行右到左）在多核间分配计算任务，最大化 L1/L2 缓存复用，并通过尾窗口重切分（TailTile）和边缘负载均衡（EdgeOptimize）实现近100%的核利用率按窗口（WINDOW_LEN=4）切分M，得到mCoreNum和mTailCoreNum,便于后期尾窗口重切分。

2) Full Load
对于输入可完整缓存在L1中的场景，可使用全载模板，在不同轮询中使输入始终驻留在L1中，从而减少整体搬运耗时。
4.3 kernel方案

//伪代码实现
single batch：
for blockIdx in totalblock:
if ASCEND_IS_AIC{
            if (isSync) {
                WaitForVector();
            }
            mmadOp_(gmBlockA, gmBlockB, gmBlockScaleA, gmBlockScaleB, gmBlockBias, locOutUb, singleShape);
            NotifyVector();
        }
       isSync = true;
        if ASCEND_IS_AIV {
            WaitForCube();
            epilogueOp_({baseM, baseN, 0, 0},
            {mPos * n + nPos, mPos * CeilDiv(n, static_cast<int64_t>(32))  + CeilDiv(nPos, static_cast<int64_t>(32)), 0, 0, 0});
            NotifyCube();
        }

关联的Issue

https://gitcode.com/cann/ops-nn/issues/3984

测试

文档更新

更新了quant_matmul_activation_quant文件
更新了aclnn_quant_matmul_activation_quant.md文件
更新了docs/zh/menu_aclnn_api.md文件
更新了docs/zh/menu_torch_api.md文件
更新了docs/zh/op_api_list.md docs/zh/op_list.md文件
更新了docs/zh/torch_api_list.md文件
更新了matmul/quant_matmul_activation_quant/docs/torchapi_quant_matmul_activation_quant.md文件

类型标签

• Bug修复
• 新特性
• 性能优化
• 文档更新
• 其他，请描述：

AI/Agent生成声明

• AI辅助编写