CATLASS GMM_sliceM_perToken_Dequant

原型设计

名称/Name	类型/Class	数据类型/Dtype	维度/Dims	格式/Format	描述/Description
matA	inTensor	int8	[m, k]	ND	左矩阵
matB	inTensor	int8	[groupCount, n, k]	ND	右矩阵，支持转置
groupList	inTensor	int	[groupCount]	ND	m轴方向分组大小，累加和列表
scale	inTensor	bf16	[groupCount, n]	ND	perChannel量化系数
perTokenScale	inTensor	bf16	[m]	ND	perToken量化系数
matD	outTensor	bf16	[m, n]	ND	输出矩阵

样例实现

CATLASS GMM_sliceM_perToken_Dequant样例算子是基于CATLASS Gemm Api实现的亲和昇腾AtlasA2硬件的GMM算子，算子的结构可以分为以下几部分

Example组装，grouped_matmul_slice_m_per_token_dequant.cpp；
Kernel实现，grouped_matmul_slice_m_per_token_dequant_multistage_workspace.hpp；
Block组件，包含：
- 通用的mmad组件block_mmad_preload_async_with_callback.hpp
- 定制的后处理组件block_epilogue_per_token_dequant.hpp；
Tile组件，除一些基本Tile_copy和tile_mmad组件外，重点关注组件：
- blockMmad中CopyGmToL1A使用的CopyGmToL1GMMPTD
- 后处理中TileRowBroadcastMul
- 后处理中TileBroadcastOneBlk
- 后处理中TileOneBlkColumnBroadcastMul

Example组装

构造输入

计算各输入输入数据len
计算各输入输入数据size
host侧生成各输入初始值，这里groupList（M轴前缀和）使用GenerateGroupList()随机生成
构造device侧输入

组装blockMmad

定义各输入的layout，layout详见
设置DispatchPolicy为MmadAtlasA2PreloadAsyncWithCallback，即选取blockMmad组件
设置L1TileShape和L0TileShape，用于切基本块分核和L1/L0上切块。需要注意TileShape和DispatchPolicy的设置有约束关系，例如当前L1::M=128、L1::K=128时，l0AStages不能超过4，以免超出L0A大小、在样例编译时被blockMmad侧静态校验拦截。TileShape计算可参考TileShape约束
设置blockMmad输入输出Type
为了专门使用CopyGmToL1GMMPTD来做CopyGmToL1A的动作，重新定义并组装了新的TileCopyMmad，若不使能CopyGmToL1GMMPTD可以跳过此步骤，BlockMmad会使用默认TileCopyMmad。
使用上述模板入参组装BlockMmad

组装blockEpilogue

设置EpilogueDispatchPolicy为EpilogueAtlasA2PerTokenDequant，即选取block_epilogue组件
定义后处理相关入参的Type：ScaleType、PerTokenScaleType、DType
定义后处理各计算步骤的Type：RowBroadcastMulType、BroadcastOneBlkType、OneBlkColumnBroadcastMulType
定义后处理单次计算的Tile大小：<m0,n0>
定义后处理计算使用的Tile-TileRowBroadcastMul：单个Tile块<m0,n0>和scale片段<1,n0>，先将scale片段<1,n0>Broadcast到<m0,n0>，再做elementwise乘法
定义后处理计算使用的Tile-TileBroadcastOneBlk：perTokenScale片段<m0,1>做Broadcast到<m0, blk>
定义后处理计算使用的Tile-TileOneBlkColumnBroadcastMul：单个Tile块<m0,n0>和perTokenScale的Broadcast片段<m0,blk>，先将perTokenScale片段<m0,blk>Broadcast到<m0,n0>，再做elementwise乘法。和上面TileBroadcastOneBlk绑定使用
定义后处理搬运使用的Tile-TileCopy，偏特化位置
定义后处理计算基本块时、按照<m0,n0>切块后的swizzle策略-EpilogueHorizontalTileSwizzle，源码位置，目前提供EpilogueIdentityTileSwizzle和EpilogueHorizontalTileSwizzle
使用上述模板入参组装BlockEpilogue

组装和执行kernel

定义kernel中基本块的swizzle策略，详见swizzle_explanation
组装kernel-GroupedMatmulSliceMPerTokenDequantMultiStageWorkspace，需要引用对应Kernel文件
将kernel组装入适配器并实例化
构造入参arguments
校验入参arguments，当前kernel内无实现、直接返回true
计算算子需要的workspace大小，kernel内实现
申请workspace
适配器初始化算子，特别注意，对于涉及核间同步的算子，需要初始化fftsAddr并在MatmulAdapter执行时调用传入

// Prepare FFTS address
    uint64_t fftsAddr{0};
    uint32_t fftsLen{0};
    RT_CHECK(rtGetC2cCtrlAddr(&fftsAddr, &fftsLen));
    ...
    matmulOp(stream, aicCoreNum, fftsAddr);

精度校验和空间释放

Kernel实现

Kernel主要结构体和函数

struct Params：执行时需要的参数
struct Arguments：host侧入参
static Params ToUnderlyingArguments：将host侧入参Arguments解析为Params，由适配器在初始化算子时调用
static size_t GetWorkspaceSize：根据Arguments计算算子执行时需要的workspace
void operator()AscendC::AIC：AIC执行代码
void operator()AscendC::AIV：AIV执行代码
struct AicWaitFunc：封装AIC等待AIV的MTE3搬运完成的核间同步。AIC上由于使用MmadAtlasA2PreloadAsyncWithCallback方案，需要将callback传入blockMmad侧，由blockMmad决定调用核间同步的时机。
struct AicSetFunc：封装AIV等待AIC的FIXPIPE搬运完成的核间同步。

AIC流程

初始化BlockScheduler和BlockMmad，代码
初始化GlobalTensor：gmA/groupList/gmC，代码
获取当前AIC序号coreIdx、AIC总数coreNum，代码
初始化计算参数：A矩阵在GM上地址偏移gmGroupOffsetA、B矩阵在GM上地址偏移gmGroupOffsetB、workspace上的输出数据排布layoutC、WORKSPACE_STAGES操作对应stageId和stageUsed、当前group起始AIC序号startCoreIdx，代码
group循环，代码
- 计算当前group的M，代码
- 判断是否使能A矩阵L2Cache绕过，代码
- 计算当前核在当前group基本块中的起始idx，代码
- 当前group中、当前核的基本块循环，代码
  - blockMmad所需入参计算，代码
  - 根据是否为异步方案调用不同blockMmad，代码，当前样例为异步方案，callback需要传入blockMmad内安排调用
  - 更新stageId，代码
- 更新gmGroupOffsetA、gmGroupOffsetB、startCoreIdx，代码
异步方案blockMmad完成遗留的计算，代码

AIV计算流程与AIC一致，在调用blockMmad()处改为调用blockEpilogue()，并在kernel代码内做核间同步，不需要将callBack传入处改为调用blockEpilogue内

基本块分核方案

group内对[currentM, N]按照[L1TileShape::M, L1TileShape::N]切基本块分核
group间连续分核，达成不同AIC的负载均衡
相关变量：
- groupIdx：当前group序号
- coreLoops：当前group的block数
- startCoreIdx：起始block的aic序号（当前group内）
- startLoopIdx：当前核的起始block序号（当前group内）
- loopIdx：当前核需要处理的block序号（当前group内）
示例：

Workspace方案

当前workspaceStages=2的配置下，使用基本块双缓冲，每个AIC分配两个基本块大小的Workspace。对于任意shape的输入，申请Workspace大小固定。

static size_t GetWorkspaceSize(const Arguments &args)
{
    size_t lenWorkspace = static_cast<size_t>(L1TileShape::M) * L1TileShape::N *
        args.aicCoreNum * WORKSPACE_STAGES;
    size_t sizeWorkspace = lenWorkspace * sizeof(uint32_t);
    return sizeWorkspace;
}

（库上还有一种方案为申请完整的 promblemShape.m * promblemShape.n 的Workspace，在小shape场景可以节约Workspace，参考）

AIC/AIV核间同步方案

AIC上对于单个block的Mmad可以拆分为L1Tile搬入、L1Tile计算、L0C搬出三个动作，由于blockMmad有preload预载和async异步的特性，operator()<AscendC::AIC>内调用blockMmad()时，仅完成当前block的L1Tile搬入和大部分L1Tile计算，而剩余的L1Tile计算和L0C搬出会在下次调用blockMmad()时执行。所以需要将AIC/AIV核间同步相关的callBack入参传入blockMmad()，由blockMmad()内决定调用时机。

单个AIC和分配的两个AIV在处理GM上的一个个block的流程图如下：

BlockMmad实现

待完善

BlockEpilogue实现

待完善