Block Scheduler GMM ASWT With Tail Split【免费下载链接】ops-tensorops-tensor 是 CANN Compute Architecture for Neural Networks算子库中提供张量类计算的基础算子库采用模块化设计支持灵活的算子开发和管理。项目地址: https://gitcode.com/cann/ops-tensor代码位置功能说明Grouped Matmul 的 BlockScheduler 组件用于 QGMM MX Tensor API kernel。调度器按 group 逐次更新问题规模在 group 间延续物理核分配位置以均衡负载并在末组计算量较小时利用空闲核拆分 M/N tail tile。框架参考Block Scheduler 公共框架特殊约束适用路径适用于 grouped matmul 场景当前由GemmUniversal的 QGMM MX 特化使用不处理 split-KK 维仅作为 shape 字段保留。扫描策略使用 SWAT 扫描策略沿 M 轴以GMM_WINDOW_LEN为窗口组织 tile并在奇数行反向扫描 N 轴以提升局部性。Tail SplitTail split 由调用方在末组按需触发。SetTailAlign配置 M/N tail 的最小对齐粒度UpdateTailTile根据剩余空闲核数拆分末尾 M/N tail tile。特殊静态常量常量说明GMM_WINDOW_LENSWAT 扫描的 M 轴窗口长度当前为 4INNER_AXIS_MIN_SPLIT_VAL内轴最小切分值当前为 128特殊类型别名别名含义TupleShape问题规模Shapeint64_t, int64_t, int64_t, int64_tBlockCoordtile 坐标Coordint64_t, int64_t, int64_t, int64_t特殊数据结构Paramsstruct Params { int32_t baseM; int32_t baseN; };参数说明baseMM 轴基础 tile 大小baseNN 轴基础 tile 大小特殊成员方法构造函数__aicore__ inline BlockSchedulerGmmSwatWithTailSplit(const Params params) __aicore__ inline BlockSchedulerGmmSwatWithTailSplit(int32_t baseM, int32_t baseN, int32_t baseK)功能初始化 M/N 轴基础 tile 大小。三参数构造函数保留baseK以兼容 kernel 侧统一构造形式调度器内部不使用 K 轴切分。UpdateNextProblem 函数__aicore__ inline void UpdateNextProblem(const TupleShape problemShape)功能刷新当前 group 的m/n/k问题规模。重新计算 M/N tile 数、tail tile 大小、轮次和起止物理核。从上一 group 的结束核后继续分配实现 group 间负载均衡。UpdateBaseM 函数__aicore__ inline void UpdateBaseM(uint32_t baseM)功能更新 M 轴基础 tile 大小配合 kernel 侧 M 轴均衡逻辑使用。SetTailAlign 函数__aicore__ inline void SetTailAlign(uint32_t mTailAlign, uint32_t nTailAlign)功能配置 M/N tail split 的对齐粒度。UpdateTailTile会结合该粒度与剩余空闲核数计算 tail 拆分数。UpdateTailTile 函数__aicore__ inline void UpdateTailTile() __aicore__ inline void UpdateTailTile(uint32_t mTailCnt, uint32_t nTailCnt)功能无参版本根据当前 group 的 tail 大小和空闲核数自动计算 M/N tail split。有参版本按调用方指定的mTailCnt/nTailCnt更新 tail split 状态。GetTileIdx 函数__aicore__ inline bool GetTileIdx(BlockCoord blockCoord)功能返回当前核本轮需要处理的 tile 坐标。当当前核无更多 tile 时返回false。GetBlockShape 函数__aicore__ inline TupleShape GetBlockShape(const BlockCoord blockCoord)功能根据 tile 坐标返回当前 tile 的实际 M/N 形状。在 tail split 场景下返回值的第 3/4 维携带 M/N split 偏移当拆分后当前核没有有效工作量时返回{0, 0, 0, 0}。GetEndBlockIdx 函数__aicore__ inline int64_t GetEndBlockIdx() const功能返回当前 group 分配后的结束物理核索引。kernel 可据此判断末组是否有空闲核可用于 tail split。调度流程构造 scheduler - SetTailAlign - 每个 group 调用 UpdateNextProblem - 末组按需调用 UpdateTailTile - GetTileIdx 获取 tile 坐标 - GetBlockShape 获取 tile 形状和 tail split 偏移适用场景QGMM MX Tensor API kernel。group 间 M/N 规模动态变化的 grouped matmul。末组 tile 数少于可用核数需要利用空闲核拆分 M/N tail 的场景。【免费下载链接】ops-tensorops-tensor 是 CANN Compute Architecture for Neural Networks算子库中提供张量类计算的基础算子库采用模块化设计支持灵活的算子开发和管理。项目地址: https://gitcode.com/cann/ops-tensor创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考