CANN/ops-tensor量化矩阵乘StreamK内核

Kernel Qbmm StreamK【免费下载链接】ops-tensorops-tensor 是 CANN Compute Architecture for Neural Networks算子库中提供张量类计算的基础算子库采用模块化设计支持灵活的算子开发和管理。项目地址: https://gitcode.com/cann/ops-tensor代码位置功能说明QBMM MX StreamK Kernel 是面向 MxFP4/MxFP8 量化矩阵乘的GemmUniversal特化实现。该 Kernel 复用 MX 量化场景的BlockMmad复用非量化 Matmul StreamK 的调度器和后处理组件通过 AIC AIV 协同完成 K 轴切分后的中间结果归约。继承自Kernel Matmul 基础框架特殊约束量化格式支持支持 MX 量化输入A/B 矩阵可使用 MxFP4 或 MxFP8 数据类型Scale 张量使用fp8_e8m0_t类型。Scale 因子要求必须提供两个 Scale 因子scaleAGmAddrA 矩阵的 per-token scale对应算子输入perTokenScalescaleBGmAddrB 矩阵的 per-group scale对应算子输入scale计算模式支持 AIC AIV 双核协同AIC 核执行 MX 量化矩阵乘计算DP tile 直接写 CSK tile 写 workspaceAIV 核执行 StreamK 后处理从 workspace 读取多个 K 切分结果累加后写回 CBatch 限制当前 StreamK 特化仅支持batch 1。多 Batch、Batch 广播和 Batch 偏移不在该 Kernel 中处理。Workspace 必需必须提供 workspace 用于存储 StreamK 的 K 轴切分中间结果AIC 在 SK 模式下把 partial sum 写入 workspaceAIV 从 workspace 读取 partial sum 并执行 Add 归约AIV 将最终结果写回输出 GMBlockScheduler 限制仅支持BlockSchedulerMatmulStreamK调度器调度策略与非量化 StreamK 保持一致支持 DPSK 混合模式。BlockEpilogue 限制复用非量化BlockEpilogueMatmulStreamK。workspace 写入和读取必须使用同一套 stride 约定尤其在ND_FIXPIPE_1_2模式下workspace 的 N 维 stride 需要按 epilogue 的 Fixpipe 策略对齐。L0C2Out 模式支持两种 Fixpipe 输出模式ON_THE_FLYworkspace 按紧凑 N stride 写入和读取ND_FIXPIPE_1_2workspace 的 N stride 按 32B 对齐保证 AIC 写入和 AIV 读取约定一致特殊模板组件组件要求ProblemShapeShapem, n, k, batch其中batch 1BlockMmadBlockMmadMatmulWithScaleMx..., KernelQbmmMultiBlockStreamK, ...BlockEpilogueBlockEpilogueMatmulStreamKWorkspaceType, OutType, MatmulMultiBlockWithStreamK...BlockSchedulerBlockSchedulerMatmulStreamKProblemShape特殊类型别名类型说明ProblemShape问题形状类型包含 m、n、k、batchBlockMmadOpMX 量化 BlockMmad 组件BlockEpilogueStreamK 后处理组件BlockSchedulerStreamK 调度器BlockMmadParamsBlockMmadOp::Params包含 A/B/C/Bias/Scale 地址BlockEpilogueParamsBlockEpilogue::Params包含 C 和 workspace 地址BlockSchedulerParamsBlockScheduler::Params包含 usedCoreNum、baseM/baseN/baseK、singleCoreK、kL1 等调度参数特殊数据结构Paramsstruct Params { ProblemShape problemShape; // 问题 shape (m, n, k, batch) BlockMmadParams blockMmadParams; // MX BlockMmad 参数 BlockEpilogueParams epilogueParams; // StreamK epilogue 参数 BlockSchedulerParams schParams; // StreamK scheduler 参数 QBMMStreamKParams qbmmParams; // QBMM StreamK 特有参数 };QBMMStreamKParamsstruct QBMMStreamKParams { uint32_t scaleKL1; // Scale 在 L1 上的 K 方向长度 uint32_t dbL0C; // L0C 双缓冲标志 };Scale L1 约束scaleKL1 schParams.kL1scaleKL1 % schParams.kL1 0scaleKL1与schParams.kL1均以原始 K 轴元素数计量不使用压缩后的 scale 元素数计量约束原因BlockMmad使用scaleKL1 / kL1计算一个 Scale L1 buffer 可覆盖的 K window 数并据此控制 scale buffer 复用与释放。如果scaleKL1 kL1或scaleKL1不是kL1的整数倍scale 搬运和同步节奏可能与 A/B 的 K window 节奏不一致。K 轴切分与 Scale group 对齐约束schParams.singleCoreK必须按 MX scale group 对齐即schParams.singleCoreK % 64 0非尾部 K split 不允许从一个 MX scale group 中间开始每个非尾 split 的真实 K 起点应落在 64 的整数倍上尾部 K split 可以小于singleCoreK但其起点仍必须满足 64 对齐scale 读取长度按尾部真实 K 长度ceil(curK / 64) * 2计算该约束由 host tiling 保证如果外部直接实例化该 kernel也必须提供满足上述条件的singleCoreK约束原因MxFP4/MxFP8 的 scale 以每 64 个 K 元素为一个 group并在 scale tensor 上以ceil(K / 64) * 2存储。StreamK 沿 K 轴切分时scale offset 必须由真实 K 起点对应的 scale group 推导若singleCoreK不是 64 对齐后续 split 可能切在 scale group 中间导致 A/B 数据与 scale group 错位。BlockMmadParamsQBMM MX StreamK 特有struct Params { GM_ADDR aGmAddr; // A 矩阵 GM 地址 GM_ADDR bGmAddr; // B 矩阵 GM 地址 GM_ADDR cGmAddr; // C 矩阵 GM 地址 GM_ADDR biasGmAddr; // Bias GM 地址可选 GM_ADDR scaleAGmAddr; // A 矩阵 Scale GM 地址 GM_ADDR scaleBGmAddr; // B 矩阵 Scale GM 地址 };说明是否启用 Bias 由biasGmAddr ! nullptr判断不额外传递 bias 标志。特殊成员方法构造函数__aicore__ inline GemmUniversal()功能构造 QBMM MX StreamK Kernel 对象。析构函数__aicore__ inline ~GemmUniversal()功能析构 QBMM MX StreamK Kernel 对象。Init函数__aicore__ inline void Init(Params const params)功能初始化 Kernel提取问题规模、GM 地址、workspace 地址和 QBMM 参数。执行流程保存problemShape和usedCoreNum设置 A/B/C/Bias/ScaleA/ScaleB 的 GM 地址设置 workspace GM 地址根据biasGmAddr判断是否启用 Biasoperator函数__aicore__ inline void operator()(Params const params)功能执行 QBMM MX StreamK Kernel。公共流程调用Init(params)初始化参数检查usedCoreNum和 batchbatch 不为 1 时直接返回创建BlockSchedulerMatmulStreamK获取 L1/L0 tile 形状、MN tile 数量和 K 切分数量AIC 核执行流程根据 blockIdx 判断是否参与实际计算初始化 MXBlockMmad构建 A/B/C/Bias/ScaleA/ScaleB 的 Tensor API Layout根据当前 tile 坐标 slice GM TensorDP 模式下直接写 CSK 模式下写 workspace并在需要时预取下一轮 SK tile计算完成后通过跨核同步标志通知 AIVAIV 核执行流程等待 AIC 同步标志初始化BlockEpilogueMatmulStreamK从 workspace 读取 K 轴切分的 partial sum执行 Add 归约和类型转换将最终结果写回 C GMWorkspace 布局写入约定AIC 在 SK tile 上把 partial sum 写入 workspace。workspace 的基地址按 core、tile 和 K 切分编号组织保证每个 K 切分结果互不覆盖。stride 约定workspace 的 N 维 stride 需要和 epilogue 保持一致auto workspaceStrideColumn0 BlockEpilogue::DispatchPolicy::fixpOpti MatMulL0C2Out::ND_FIXPIPE_1_2 ? CeilAlign(GetMNK_N(singleCoreShape), BLOCK_BYTE_SIZE) : GetMNK_N(singleCoreShape);说明ON_THE_FLY模式使用紧凑 N strideND_FIXPIPE_1_2模式使用 32B 对齐后的 N strideAIC 写 workspace 和 AIV 读 workspace 必须使用同一个 stride否则非 32 对齐 N 场景会出现读写错位DPSK 混合策略调度策略复用非量化 StreamKtailMNTileNum (mTileNum * nTileNum) % usedCoreNum; tileNum (mTileNum * nTileNum - tailMNTileNum) tailMNTileNum * skKTileNum;说明DPData Parallel模式前mTileNum * nTileNum - tailMNTileNum个 MN tile 完整计算结果直接写 CSKStreamK模式尾部tailMNTileNum个 MN tile 沿 K 轴拆成多个 tilepartial sum 写 workspaceAIV 归约对 SK 模式的多个 K 切分结果执行 Add输出最终 C调用示例组件组装using AType fp8_e4m3fn_t; using BType fp8_e5m2_t; using OutType half; using BiasType float; using WorkspaceType float; using ProblemShape AscendC::Te::Shapeint64_t, int64_t, int64_t, int64_t; using DispatchPolicy Blaze::Gemm::MatmulWithScaleMx0, false, Blaze::Gemm::KernelQbmmMultiBlockStreamK; using EpilogueDispatchPolicy Blaze::Gemm::MatmulMultiBlockWithStreamK; using BlockMmad Blaze::Gemm::Block::BlockMmad DispatchPolicy, AType, LayoutA, BType, LayoutB, OutType, LayoutC, BiasType, LayoutC; using BlockEpilogue Blaze::Gemm::Block::BlockEpilogueMatmulStreamK WorkspaceType, OutType, EpilogueDispatchPolicy; using BlockScheduler Blaze::Gemm::Block::BlockSchedulerMatmulStreamKProblemShape; using QbmmStreamKKernel Blaze::Gemm::Kernel::GemmUniversal ProblemShape, BlockMmad, BlockEpilogue, BlockScheduler;参数准备using Params typename QbmmStreamKKernel::Params; Params params{ {m, n, k, 1}, // problem shapebatch 固定为 1 {aGM, bGM, cGM, biasGM, scaleAGM, scaleBGM}, {cGM, workspaceGM}, {usedCoreNum, baseM, baseN, baseK, singleCoreK, kL1}, {scaleKL1, dbL0C} // scaleKL1 kL1 and scaleKL1 % kL1 0 };Kernel 执行QbmmStreamKKernel qbmm; qbmm(params);数据流存储层次GM(A/B/ScaleA/ScaleB/Bias) ↓ Tensor API Layout BlockSchedulerMatmulStreamK ↓ BlockMmadMatmulWithScaleMx..., KernelQbmmMultiBlockStreamK ↓ DP tile: L0C → GM(C) SK tile: L0C → Workspace(float partial sum) ↓ BlockEpilogueMatmulStreamK(AIV) ↓ GM(C)AIC/AIV 协同流程AIC: 计算 DP tile → 直接写 C AIC: 计算 SK tile → 写 workspace AIC: CrossCoreSetFlag 通知 AIV ↓ AIV: CrossCoreWaitFlag 等待 AIC 完成 AIV: 从 workspace 读取多个 K 分片 AIV: Add 归约 Cast/ReLU 等后处理 AIV: 写回 C【免费下载链接】ops-tensorops-tensor 是 CANN Compute Architecture for Neural Networks算子库中提供张量类计算的基础算子库采用模块化设计支持灵活的算子开发和管理。项目地址: https://gitcode.com/cann/ops-tensor创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考