CANN/ops-tensor QBMM立方体内核

Kernel Qbmm Cube【免费下载链接】ops-tensorops-tensor 是 CANN Compute Architecture for Neural Networks算子库中提供张量类计算的基础算子库采用模块化设计支持灵活的算子开发和管理。项目地址: https://gitcode.com/cann/ops-tensor代码位置功能说明Fixpipe 量化 Batch Matmul Kernel仅支持 AIC 计算。该 Kernel 组合BlockMmadA8W8FixpipeQuant与BlockSchedulerQuantBatchMatmulV3完成量化 A/B 矩阵乘、Batch 广播、Bias 处理、scale 反量化与 GM 输出。BlockMmadA8W8FixpipeQuant沿用 A8W8 路径命名实际输入类型由AType/BType决定。继承自Kernel Matmul 基础框架特殊约束量化格式支持支持 Tensor API Cube Mmad/Fixpipe 路径可处理的量化输入类型典型包括int8_tA8W8 输入L0C 通常累加为int32_thifloat8_tHiFloat8 输入L0C 通常累加为floatfp8_e4m3fn_t、fp8_e5m2_tFP8 输入支持 E4M3FN/E5M2 同型或混合组合L0C 通常累加为floatC 输出由CType指定支持half、bfloat16_t、float或int32_t最终以 Tensor API Fixpipe 静态检查为准Scale 因子要求Scale 模式由x1QuantMode和x2QuantMode决定取值来自QuantMode枚举值取值说明DEFAULT0x0不启用对应 scalePERTENSOR_MODE0x1per-tensor scalar scalePERCHANNEL_MODE0x2per-channel scalePERTOKEN_MODE0x4per-token scaleMX_PERGROUP_MODE0x8MX per-group scalePERBLOCK_MODE0x10per-block scalePERGROUP_MODE0x20per-group scalex1QuantMode描述 A 矩阵的 scale 模式x2QuantMode描述 B 矩阵的 scale 模式。当前 Kernel 支持以下组合A 矩阵模式 (x1QuantMode)B 矩阵模式 (x2QuantMode)适用场景与 Scale 处理DEFAULTDEFAULTA、B 均不提供反量化 scale。用于CType int32_t的场景Block 将 L0C 累加结果直接写入 GM。DEFAULTPERCHANNEL_MODEA 不提供 scaleB 的每个 N 通道分别使用一个 scale。scaleBGmAddr指向 B 的 per-channel scale 数组Kernel 根据当前 block 负责的 N 范围截取对应 scale Tensor 并传给 BlockBlock 将 scale 搬入 L1Fixpipe 搬出结果时按通道完成反量化。DEFAULTPERTENSOR_MODEA 不提供 scale整个 B 矩阵共用一个 scalar scale。Kernel 从scaleBGmAddr读取该值将其转换并封装为 Fixpipe 使用的uint64_t scaleScalar_再传给 Block。支持的 scale 存储类型为uint64_t/int64_t、bfloat16_t或uint32_t。Block 在 Fixpipe 搬出结果时使用该 scalar scale 完成反量化。PERTENSOR_MODEPERTENSOR_MODEA、B 各自提供一个 per-tensor scalar scale二者的存储类型均为float。Kernel 分别从scaleAGmAddr和scaleBGmAddr读取 scale计算两者乘积将乘积转换并封装为 Fixpipe 使用的uint64_t scaleScalar_再传给 Block。Block 在 Fixpipe 搬出结果时使用该乘积 scale 完成反量化。其他组合当前模板未提供对应 scale 数据流。当输出CType int32_t时无论传入 Block 的是 scalar scale 还是 per-channel scaleBlock 都不会使用 scale也不会执行 Fixpipe 反量化而是将 L0C 中的int32_t累加结果直接写入 GM。计算模式仅支持 AIC 模式不支持 AIV 计算AIV 核在Init中直接返回。BlockMmad 限制仅支持调度策略为MatmulWithScaleFixpipeQuant的BlockMmad即BlockMmad::DispatchPolicy::ScheduleType必须为KernelMmadWithScaleFixpipeQuant。BlockScheduler 限制使用BlockSchedulerQuantBatchMatmulV3调度器支持 block 切分、尾块切分和多 Batch 维度遍历。Batch 维度限制支持 4 维 BatchbatchA1/A2/A3/A4、batchB1/B2/B3/B4、batchC1/C2/C3/C4。多 Batch 路径按 C 的 Batch 空间遍历并根据 A/B 到 C 的广播倍数计算 GM 地址偏移。Atomic Add 模式可选 Atomic Add 模式IS_ATOMIC_ADD true。Kernel 在计算开始时配置 atomic add结束时恢复 atomic none。模板参数模板定义template class ProblemShape, // 问题形状类型 class BlockMmad, // Fixpipe Quant BlockMmad class BlockEpilogue, // 后处理组件通常为空 class BlockScheduler // BlockSchedulerQuantBatchMatmulV3 调度器模板参数说明参数说明ProblemShape问题形状类型包含 m、n、k、bBlockMmadFixpipe Quant BlockMmad基于MatmulWithScaleFixpipeQuantBlockEpilogue后处理组件当前路径通常使用BlockEpilogueEmptyBlockSchedulerQBMM V3 调度器特殊类型别名类型说明BlockMmadParamsBlockMmad 参数类型包含 GM 地址AType/BType/CType/BiasType继承自 BlockMmad 的数据类型X2ScaleTypeKernel 内部使用的 scalar scale 类型固定为uint64_tScaleGmTypeBlockMmad 中BTypeTuple的第 1 个类型WEIGHT_NZB 矩阵是否为 NZ 格式继承自 BlockMmadTRANS_AA 矩阵是否转置继承自 BlockMmadTRANS_BB 矩阵是否转置继承自 BlockMmadIS_ATOMIC_ADD是否启用 Atomic Add 模式继承自 BlockMmad::DispatchPolicyC0_SIZEA/B C0 对齐大小MakeLayoutA/MakeLayoutB/MakeLayoutCA/B/C GM Tensor Layout 构建器特殊数据结构Paramsstruct Params { ProblemShape problemShape; // 问题 shape (m, n, k, batch) BlockMmadParams mmadParams; // BlockMmad 参数包含 GM 地址 BlockSchedulerParams schParams; // scheduler 参数 QBMMTiling qbmmParams; // QBMM Cube 参数 };QBMMTilingstruct QBMMTiling { uint32_t batchA1, batchA2, batchA3, batchA4; // A 矩阵 Batch 维度 uint32_t batchB1, batchB2, batchB3, batchB4; // B 矩阵 Batch 维度 uint32_t batchC1, batchC2, batchC3, batchC4; // C 矩阵 Batch 维度 uint32_t biasThreeDim; // Bias 是否按 Batch 展开 uint32_t x1QuantMode; // x1 scale 量化模式 uint32_t x2QuantMode; // x2 scale 量化模式 uint32_t kAL1; // A 的 L1 K 轴切分 uint32_t kBL1; // B 的 L1 K 轴切分 uint32_t nBufferNum; // A/B L1 缓冲数量 uint32_t baseM, baseN, baseK; // L0 tile 形状 uint32_t isBias; // 是否启用 bias uint32_t dbL0C; // L0C 双缓冲标志 };BlockMmadParamsstruct Params { GM_ADDR aGmAddr; // A 矩阵 GM 地址 GM_ADDR bGmAddr; // B 矩阵 GM 地址 GM_ADDR cGmAddr; // C 矩阵 GM 地址 GM_ADDR biasGmAddr; // Bias GM 地址可选 GM_ADDR scaleAGmAddr; // x1 scale GM 地址 GM_ADDR scaleBGmAddr; // x2 scale GM 地址 };执行流程Kernel 对外通过operator()执行完整 QBMM Cube 计算内部流程概括如下解析Params设置 A/B/C/Bias/Scale 的 GM 地址和 scale 模式。按baseM/baseN/baseK/kAL1/kBL1/nBufferNum/x2QuantMode/isBias/dbL0C初始化 BlockMmad。由BlockSchedulerQuantBatchMatmulV3负责 block 坐标、尾块和 Batch 遍历。单 Batch 场景直接遍历 block多 Batch 场景按 C 的 4 维 Batch 空间遍历并根据广播关系更新 A/B/C 地址偏移。per-channel scale 场景向 BlockMmad 传入当前 N 分片 scale Tensorper-tensor scale 场景传入折算后的 scalar scale。Atomic Add 模式下Kernel 在计算前后分别配置和恢复 atomic 状态。调用示例组件组装// 以下以 int8_t A/B 为例可按 Tensor API 支持组合替换为 hifloat8_t 或 fp8 类型。 using AType int8_t; using BType int8_t; using CType bfloat16_t; using BiasType int32_t; using X2ScaleType uint64_t; using LayoutA AscendC::Te::NDExtLayoutPtn; using LayoutB AscendC::Te::NDExtLayoutPtn; using LayoutC AscendC::Te::NDExtLayoutPtn; using LayoutBias AscendC::Te::NDExtLayoutPtn; using ProblemShape AscendC::Te::Shapeint64_t, int64_t, int64_t, int64_t; using BTypeTuple AscendC::Std::tupleBType, X2ScaleType; using DispatchPolicy Blaze::Gemm::MatmulWithScaleFixpipeQuant0, false; using BlockMmad Blaze::Gemm::Block::BlockMmad DispatchPolicy, AType, LayoutA, BTypeTuple, LayoutB, CType, LayoutC, BiasType, LayoutBias; using BlockScheduler Blaze::Gemm::Block::BlockSchedulerQuantBatchMatmulV3 ProblemShape, 0, LayoutA, LayoutB, AType; using BlockEpilogue Blaze::Gemm::Block::BlockEpilogueEmpty; using QBMMKernel Blaze::Gemm::Kernel::GemmUniversal ProblemShape, BlockMmad, BlockEpilogue, BlockScheduler;参数准备using Params typename QBMMKernel::Params; Params params; params.problemShape {m, n, k, batch}; params.mmadParams.aGmAddr aGM; params.mmadParams.bGmAddr bGM; params.mmadParams.cGmAddr cGM; params.mmadParams.biasGmAddr biasGM; params.mmadParams.scaleAGmAddr pertokenScaleGM; params.mmadParams.scaleBGmAddr scaleGM; params.schParams { baseM, baseN, mTailTile, nTailTile, mBaseTailSplitCnt, nBaseTailSplitCnt, mTailMain, nTailMain }; params.qbmmParams { batchA1, batchA2, batchA3, batchA4, batchB1, batchB2, batchB3, batchB4, batchC1, batchC2, batchC3, batchC4, biasThreeDim, x1QuantMode, x2QuantMode, kAL1, kBL1, nBufferNum, baseM, baseN, baseK, isBias, dbL0C };Kernel 执行QBMMKernel kernel; kernel(params); // 或 kernel.Run(params)数据流单 Batch 流程Params ↓ 解析 GM 地址和 scale ↓ BlockScheduler 获取 block ↓ Slice A/B/C/Bias/Scale ↓ BlockMmadA8W8FixpipeQuant ↓ GM 输出多 Batch 流程4 维 Batch C 空间遍历 ↓ 按广播倍数计算 A/B Batch 偏移 ↓ 更新 A/B/C GM 地址偏移 ↓ 进入单 Batch block 处理流程Scale 处理流程Scale 的处理分为 Kernel 准备和 Block 搬出两个阶段。Kernel 根据x1QuantMode/x2QuantMode确定 scale 的来源和形态并将 scalar scale 或当前 N 分片的 per-channel scale Tensor 传给 Block。Mmad 计算始终先在 L0C 中完成累加仅当输出类型不是int32_t时Block 才在结果搬出阶段使用 scale 进行 Fixpipe 反量化。CType int32_tMmad 完成后Block 将 L0C 中的int32_t累加结果直接写入 GM。该路径不使用传入 Block 的 scale不执行 Fixpipe 反量化也不把累加结果转换为浮点类型。CType half/bfloat16_t/float且 B 使用 per-channel scaleKernel 根据当前 block 的 N 方向起始位置和长度从scaleBGmAddr中取得对应通道的 scale Tensor。Block 将该 Tensor 从 GM 搬入 L1Fixpipe 搬出 L0C 结果时对每个 N 通道应用各自的 scale并转换为CType后写入 GM。CType half/bfloat16_t/float且仅 B 使用 per-tensor scaleKernel 从scaleBGmAddr读取整个 B 矩阵共用的 scalar scale将其转换并封装为uint64_t scaleScalar_后传给 Block。Fixpipe 搬出 L0C 结果时对所有元素使用同一个 scale并转换为CType后写入 GM。CType half/bfloat16_t/float且 A、B 均使用 per-tensor scaleKernel 从scaleAGmAddr和scaleBGmAddr分别读取 A、B 的floatscalar scale先计算二者的乘积再将乘积转换并封装为uint64_t scaleScalar_后传给 Block。Fixpipe 使用该乘积 scale 完成反量化和输出类型转换然后将结果写入 GM。性能优化建议Block 与 L1 配置baseM/baseN/baseK应与 BlockMmad 的 L0/L1 切分匹配。kAL1/kBL1可根据 A/B 复用关系配置详见 Block Mmad A8W8 Fixpipe Quant。nBufferNum 4可提高搬运流水并行度nBufferNum 2支持 A/B 不同 K-L1 窗口。Batch 广播多 Batch 场景下A/B/C 的 4 维 Batch 参数需满足广播关系。biasThreeDim 1时Bias 随 C Batch 偏移否则所有 Batch 复用同一 Bias。尾块切分尾轮由 scheduler 更新尾块切分用于提升尾块场景下的核利用率。mTailTile/nTailTile应与 scheduler 参数保持一致。【免费下载链接】ops-tensorops-tensor 是 CANN Compute Architecture for Neural Networks算子库中提供张量类计算的基础算子库采用模块化设计支持灵活的算子开发和管理。项目地址: https://gitcode.com/cann/ops-tensor创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考