CANN pyasc Matmul.iterate_batch 详解一次调用批量计算多块 C 矩阵的高阶矩阵乘接口【免费下载链接】pyasc本项目为Python用户提供算子编程接口支持在昇腾AI处理器上加速计算接口与Ascend C一一对应并遵守Python原生语法。项目地址: https://gitcode.com/cann/pyasc导读asc.language.adv.Matmul.iterate_batch是 CANN pyasc 为 Python 开发者提供的 Matmul 高阶批量迭代接口调用一次即可计算出多个singleCoreM × singleCoreN大小的 C 矩阵分片是编写 BatchMatmulBMM批量矩阵乘算子的核心入口。本文围绕该接口的 Python 签名、Ascend C 函数原型、全量参数语义、Layout 与内存约束、同步/异步用法展开并结合仓库中的 Python 前端实现 与 泛化测试用例 说明其底层行为与正确使用方式。读完本文你将能够基于 pyasc 独立编写出正确的多 Batch Matmul 算子内核。接口定位从单次单块到单次多块在 pyasc 的 Matmul 高阶 API 家族中迭代类接口按计算粒度分为四档iterate_batch位于批量分片这一档接口一次调用产出典型场景Matmul.iterate一块baseM × baseN的 C 矩阵循环内逐片取数Matmul.iterate_all一块singleCoreM × singleCoreN的 C 矩阵单核单块直出Matmul.iterate_batch多块singleCoreM × singleCoreN的 C 矩阵批量 MatmulBMMMatmul.iterate_n_batchN 次 iterate_batch 的计算结果多 BMM 循环批量计算根据文档定义iterate_batch提供批量处理 Matmul 的功能调用一次 iterate_batch可以计算出多个 singleCoreM * singleCoreN 大小的 C 矩阵且该接口只支持 Norm 模板即 BatchMatmul 只支持 Norm 模板接口文档。函数签名与 Ascend C 原型对照pyasc 中该接口存在两个重载overload分别对应不同的参数组织方式# 重载一带 batch_a / batch_b 的标准形式推荐 Matmul.iterate_batch(tensor: BaseTensor, batch_a: int, batch_b: int, en_sequential_write: bool, matrix_stride_a: int 0, matrix_stride_b: int 0, matrix_stride_c: int 0, en_partial_sum: bool False, en_atomic: int 0, sync: bool True, wait_iterate_batch: bool | None None) → None # 重载二纯 Cubecube-only场景的简写形式 Matmul.iterate_batch(tensor: BaseTensor, en_partial_sum, en_atomic, en_sequential_write: bool, matrix_stride_a: int 0, matrix_stride_b: int 0, matrix_stride_c: int 0, sync: bool True) → None两种重载在 Python 侧通过OverloadDispatcher按参数形态分派见 matmul.py最终分别生成asc_MatmulIterateBatchOp与asc_MatmulIterateBatchCubeOnlyOp两个 IR 算子节点。对应的 Ascend C 函数原型如下// 重载一对应同步/异步模板化版本支持输出到 GM 或 UB template bool sync true, bool waitIterateBatch false __aicore__ inline void IterateBatch(const GlobalTensorDstT gm, uint32_t batchA, uint32_t batchB, bool enSequentialWrite, const uint32_t matrixStrideA 0, const uint32_t matrixStrideB 0, const uint32_t matrixStrideC 0, const bool enPartialSum false, const uint8_t enAtomic 0) template bool sync true __aicore__ inline void IterateBatch(const LocalTensorDstT ubCmatrix, uint32_t batchA, uint32_t batchB, bool enSequentialWrite, const uint32_t matrixStrideA 0, const uint32_t matrixStrideB 0, const uint32_t matrixStrideC 0, const bool enPartialSum false, const uint8_t enAtomic 0) // 重载二对应纯 Cube 场景版本 __aicore__ inline void IterateBatch(const GlobalTensorDstT gm, bool enPartialSum, uint8_t enAtomic, bool enSequentialWrite, const uint32_t matrixStrideA 0, const uint32_t matrixStrideB 0, const uint32_t matrixStrideC 0) __aicore__ inline void IterateBatch(const LocalTensorDstT ubCmatrix, bool enPartialSum, uint8_t enAtomic, bool enSequentialWrite, const uint32_t matrixStrideA 0, const uint32_t matrixStrideB 0, const uint32_t matrixStrideC 0)参数语义详解参数类型默认值含义tensorBaseTensor必填C 矩阵类型为 GlobalTensor输出到 GM或 LocalTensor输出到 UBbatch_aint必填左矩阵 A 的 batch 数batch_bint必填右矩阵 B 的 batch 数en_sequential_writebool必填是否开启连续写模式matrix_stride_aint0A 矩阵源操作数相邻 nd 矩阵起始地址间的偏移单位是元素matrix_stride_bint0B 矩阵源操作数相邻 nd 矩阵起始地址间的偏移单位是元素matrix_stride_cint0预留参数开发者无需关注en_partial_sumboolFalse是否将矩阵乘的结果累加于现有的 CO1 数据en_atomicint0是否开启 Atomic 操作取值范围见下文syncboolTrue设置同步或者异步模式wait_iterate_batchbool / NoneNone是否需要通过 wait_iterate_batch 接口等待 iterate_batch 执行结束仅在异步场景下使用en_sequential_write 的输出位置强约束从 matmul.py 的校验逻辑可以确认一个容易踩坑的硬性约束输出位置与连续写模式必须配对。if isinstance(tensor, GlobalTensor): check_type(en_sequential_write, [False], When output to GM, en_sequential_write should be False in iterate_batch) else: check_type(en_sequential_write, [True], When output to GM, en_sequential_write should be True in iterate_batch)即C 矩阵输出到GM 时必须传en_sequential_writeFalse输出到UB 时必须传en_sequential_writeTrue传反会直接抛出ValueError。同时接口文档的约束部分还明确指出BSNGD、SBNGD Layout 格式不支持连续写模式且异步模式不支持 iterate_batch 搬运到 UB 上即异步场景下输出只能落在 GM。en_atomic 的取值Python 侧对en_atomic做了白名单校验check_type(en_atomic, [None, 0, 1, 2, 3], ...)即仅允许None/0/1/2/3matmul.py。其中None表示不显式指定内部按未使能处理0为默认值即不开启 Atomic 操作1/2/3分别对应不同的 Atomic 累加模式。该参数在 IR 层以uint8类型KT.uint8落盘。wait_iterate_batch 与同步/异步syncTrue默认为同步模式iterate_batch调用返回即代表本次批量计算在同步点上完成syncFalse为异步模式此时可通过wait_iterate_batchTrue请求在计算完成后插入等待源码中对异步 立即等待的组合做了显式拦截raise ValueError(Param wait_iterate_batch can be True only when sync is False)matmul.pywait_iterate_batch仅在输出为 GlobalTensor 时生效若为None则内部按False处理异步模式下如需显式同步可调用无参接口Matmul.wait_iterate_batch()其配套约束为仅支持连续输出至 Global Memory且 MixDualMaster 场景下不可用。约束说明Layout、L1 Buffer 与数据类型接口文档给出了非常具体的约束清单这里逐条展开Layout 与 batch 关系BSNGD、SBNGD、BNGS1S2 Layout输入 A、B 矩阵按分形对齐后的多 Batch 数据总和应小于 L1 Buffer 的大小NORMAL Layout 没有该限制但需通过 MatmulConfig 配置 A、B 矩阵多 Batch 数据大小与 L1 Buffer 大小的关系对应get_normal_config的bmm_mode/batch_mode等参数。G 轴对齐约束对 BSNGD、SBNGD、BNGS1S2 Layout左矩阵、右矩阵的 G 轴分别为a_layout_info_g、b_layout_info_g必须满足a_layout_info_g / batch_a b_layout_info_g / batch_bNORMAL Layout 下则要求batch_a、batch_b满足倍数关系。BSNGD 场景限制不支持一次计算多行 SD需要算子程序中循环计算即(a_layout_info_n * a_layout_info_g) / batch_a、(b_layout_info_n * b_layout_info_g) / batch_b均为整数。数据格式支持BSNGD、SBNGD Layout 输入输出只支持 ND 格式数据BNGS1S2、NORMAL Layout 输入支持 ND/NZ 格式数据。Unified Buffer 输出约束输出到 Unified Buffer 时C 矩阵大小Base_m * Base_n应小于分配的 Unified Buffer 内存大小输出到 UB 且单核计算的 N 方向大小single_core_n非 32 字节对齐时C 矩阵的 CubeFormat 仅支持ND_ALIGN格式输出 C 矩阵片时硬件会自动将single_core_n方向上的数据补齐至 32 字节。能力边界不支持量化模式即不支持set_quant_scalar、set_quant_vector接口不支持 MixDualMaster当模板参数enableMixDualMaster设置为 true 时双主模式不能使用该接口不支持 int4b_t 输入A 矩阵、B 矩阵不支持int4b_t类型的输入即 BatchMatmul 不支持 int4b_t 类型的矩阵输入。此外Python 侧对输出 C 矩阵的数据类型同样有白名单[KT.int32, KT.int_, KT.float_, KT.float32, KT.half, KT.float16]matmul.py不在其中的类型会直接报 Tensor type is not supported in iterate_batch。典型调用示例BSNGD Layout 下的多 Batch Matmul接口文档给出了一个完整的 BSNGD 场景内核代码其整体流程为定义 MatmulType → 实例化 Matmul → register_matmul → init(tiling) → 循环内 set_tensor_a/set_tensor_b/set_bias → iterate_batch# 定义matmul type a_type asc.adv.MatmulType(asc.TPosition.GM, asc.CubeFormat.ND, asc.half, False, asc.LayoutMode.BSNGD) b_type asc.adv.MatmulType(asc.TPosition.GM, asc.CubeFormat.ND, asc.half, True, asc.LayoutMode.BSNGD) c_type asc.adv.MatmulType(asc.TPosition.GM, asc.CubeFormat.ND, asc.float, False, asc.LayoutMode.BNGS1S2) bias_type asc.adv.MatmulType(asc.TPosition.GM, asc.CubeFormat.ND, asc.float) mm asc.adv.Matmul(a_type, b_type, c_type, bias_type) asc.adv.register_matmul(pipe, mm) mm.init(tiling) batch_c batch_a if batch_b batch_c: batch_c batch_b g_lay tiling.a_layout_info_g if tiling.b_layout_info g_lay: g_lay tiling.b_layout_info_g for_extent tiling.a_layout_info_b * tiling.a_layout_info_n * g_lay / tiling.batch_num for i in range(for_extent): batch_offset_a i * tiling.a_layout_info_d * batch_a batch_offset_b i * tiling.b_layout_info_d * batch_b mm.set_tensor_a(gm_a[batch_offset_a], is_transpose_a_in) mm.set_tensor_b(gm_b[batch_offset_b], is_transpose_b_in) idx_c i * batch_c if tiling.c_layout_info_g 1 and (tiling.b_layout_info_g ! 1 or tiling.a_layout_info_g ! 1): d tiling.b_layout_info_g if tiling.a_layout_info_g d: d tiling.a_layout_info_g idx_c idx_c // d if tiling.is_bias: batch_offset_bias idx_c * tiling.c_layout_info_s2 mm.ste_bias(gm_bias[batch_offset_bias]) batch_offset_c idx_c * tiling.c_layout_info_s2 if c_type.layout asc.LayoutMode.BNGS1S2: batch_offset_c idx_c * tiling.c_layout_infos2 * tiling.c_layout_info_s1 mm.iterate_batch(tensorgm_c[offsetc], batch_abatch_a, batch_bbatch_b, en_sequential_writeFalse)该示例展示了三个要点Layout 混合A、B 矩阵使用 BSNGDC 矩阵使用 BNGS1S2batch 索引idx_c需按 G 轴比例折算偏移计算A/B 的 batch 偏移分别乘以a_layout_info_d/b_layout_info_dC 的 batch 偏移乘以c_layout_info_s2BNGS1S2 下还需再乘c_layout_info_s1输出到 GM 时en_sequential_writeFalse与源码校验规则一致。纯 Cube 模式与重载二的使用重载二用于纯 Cube只有矩阵计算场景参数顺序为(tensor, en_partial_sum, en_atomic, en_sequential_write, ...)。从源码matmul.py可见其强约束该形式下输出只能是 GlobalTensor否则抛出TypeError(iterate_batch interface under cube-only sence only support output to GM.)且en_sequential_write同样必须为False。这与仓库示例 examples/04_matmul_cube_only/matmul_cube_only.py 展示的纯 Cube 模式输出直落 GM的写法一致——在该示例中 Matmul 实例化未指定 Layout默认 NONE计算完成后通过iterate_all(c_global)直接输出。配合重载二使用的典型流程可参考Matmul.set_batch_num的调用示例在纯 Cube 模式下先用set_batch_num(batch_a, batch_b)在不改变 Tiling 的情况下重新设置 Batch 数再调用mm.iterate_batch(tensorgm_c, en_partial_sumFalse, en_atomic0, en_sequential_writeFalse)完成多 batch Matmul 计算。源码与测试验证NORMAL Layout 的端到端用法仓库中的泛化测试 python/test/generalization/adv/test_matmul_iterate_batch.py 提供了 NORMAL Layout 下iterate_batch的完整端到端样例float32 与 float16 两种 dtype、32/64 尺寸参数化验证。其内核核心代码如下asc.jit def matmul_kernel(a, b, c, tiling, workspace): ... matmul asc.adv.Matmul( aasc.adv.MatmulType(asc.TPosition.GM, asc.CubeFormat.ND, a_global.dtype, False, asc.LayoutMode.NORMAL), basc.adv.MatmulType(asc.TPosition.GM, asc.CubeFormat.ND, b_global.dtype, False, asc.LayoutMode.NORMAL), casc.adv.MatmulType(asc.TPosition.GM, asc.CubeFormat.ND, c_global.dtype, False, asc.LayoutMode.NORMAL), ) asc.adv.register_matmul(pipe, workspace, matmul, tiling) matmul.set_tensor_a(a_global) matmul.set_tensor_b(b_global) matmul.iterate_batch(c_global, BATCH_A, BATCH_B, False) asc.pipe_barrier(asc.PipeID.PIPE_ALL)对应 Host 侧 Tiling 生成test_matmul_iterate_batch.py通过MultiCoreMatmulTiling的set_a_layout/set_b_layout/set_c_layout/set_batch_num配置各矩阵的 Layout 轴信息与最大 Batch 数matmul_tiling host.MultiCoreMatmulTiling(host.get_ascendc_platform()) matmul_tiling.set_a_type(host.TPosition.GM, host.CubeFormat.ND, host_dtype, False) matmul_tiling.set_b_type(host.TPosition.GM, host.CubeFormat.ND, host_dtype, False) matmul_tiling.set_c_type(host.TPosition.VECCALC, host.CubeFormat.ND, host.DataType.DT_FLOAT) matmul_tiling.set_dim(16) matmul_tiling.set_org_shape(m, n, k) matmul_tiling.set_shape(m, n, k) matmul_tiling.set_traverse(host.MatrixTraverse.FIRSTM) matmul_tiling.set_buffer_space(-1, -1, -1) matmul_tiling.set_a_layout(BATCH_A, m, 1, 1, k) matmul_tiling.set_b_layout(BATCH_B, k, 1, 1, n) matmul_tiling.set_c_layout(BATCH, m, 1, 1, n) matmul_tiling.set_batch_num(BATCH) tiling asc.adv.TCubeTiling() matmul_tiling.get_tiling(tiling)测试中的 batch 组合还覆盖了 A、B batch 数相同的常规情形并验证了batch_a batch_b 2时结果与 PyTorch 参考实现a b的一致性torch.allclose(c, matmul, atol1e-3)。这从侧面印证了文档中batch_a、batch_b 需满足倍数关系的约束在 NORMAL Layout 下的具体形态。与 wait_iterate_batch 配合的异步用法异步场景的标准配对写法wait_iterate_batch 接口文档mm asc.adv.Matmul(a_type, b_type, c_type, bias_type) mm.set_tensor_a(gm_a[offset_a]) mm.set_tensor_b(gm_b[offset_b]) if tiling.is_bias: mm.set_bias(gm_bias[offset_bias]) mm.iterate_batch(tensorgm_c[offsetc], batch_abatch_a, batch_bbatch_b, en_sequential_writeFalse) mm.wait_iterate_batch()即iterate_batch发起异步批量计算后调用无参的mm.wait_iterate_batch()等待其返回从而支持连续输出到 Global Memory避免后续搬移与计算的竞争。注意wait_iterate_batch接口的约束同样包含仅支持连续输出至 Global MemoryMixDualMaster 场景不支持。小结iterate_batch 使用自查清单编写使用iterate_batch的 pyasc 算子时建议按以下清单自查模板选择BatchMatmul 仅支持 Norm 模板实例化 Matmul 时通过MatmulType的layout参数指定 LayoutBSNGD / SBNGD / BNGS1S2 / NORMAL输出位置与 en_sequential_write输出 GM 传False输出 UB 传True且 BSNGD/SBNGD 不支持连续写、异步模式不支持输出到 UBbatch 关系BSNGD/SBNGD/BNGS1S2 下满足a_layout_info_g / batch_a b_layout_info_g / batch_bNORMAL 下满足倍数关系内存BSNGD/SBNGD/BNGS1S2 多 Batch 数据总和小于 L1 Buffer输出 UB 时Base_m * Base_n小于分配的 UB 内存非 32 字节对齐时使用ND_ALIGN能力边界不使用量化接口、不使能 MixDualMaster、A/B 不使用 int4b_t异步syncFalse时按需设置wait_iterate_batch或调用wait_iterate_batch()同步。掌握上述要点后即可在 examples 与 python/test/generalization/adv 中已有示例的基础上快速编写属于自己的多 Batch Matmul 昇腾算子。【免费下载链接】pyasc本项目为Python用户提供算子编程接口支持在昇腾AI处理器上加速计算接口与Ascend C一一对应并遵守Python原生语法。项目地址: https://gitcode.com/cann/pyasc创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 CANN pyasc Matmul.iterate_batch 详解：一次调用批量计算多块 C 矩阵的高阶矩阵乘接口
CANN pyasc Matmul.iterate_batch 详解一次调用批量计算多块 C 矩阵的高阶矩阵乘接口【免费下载链接】pyasc本项目为Python用户提供算子编程接口支持在昇腾AI处理器上加速计算接口与Ascend C一一对应并遵守Python原生语法。项目地址: https://gitcode.com/cann/pyasc导读asc.language.adv.Matmul.iterate_batch是 CANN pyasc 为 Python 开发者提供的 Matmul 高阶批量迭代接口调用一次即可计算出多个singleCoreM × singleCoreN大小的 C 矩阵分片是编写 BatchMatmulBMM批量矩阵乘算子的核心入口。本文围绕该接口的 Python 签名、Ascend C 函数原型、全量参数语义、Layout 与内存约束、同步/异步用法展开并结合仓库中的 Python 前端实现 与 泛化测试用例 说明其底层行为与正确使用方式。读完本文你将能够基于 pyasc 独立编写出正确的多 Batch Matmul 算子内核。接口定位从单次单块到单次多块在 pyasc 的 Matmul 高阶 API 家族中迭代类接口按计算粒度分为四档iterate_batch位于批量分片这一档接口一次调用产出典型场景Matmul.iterate一块baseM × baseN的 C 矩阵循环内逐片取数Matmul.iterate_all一块singleCoreM × singleCoreN的 C 矩阵单核单块直出Matmul.iterate_batch多块singleCoreM × singleCoreN的 C 矩阵批量 MatmulBMMMatmul.iterate_n_batchN 次 iterate_batch 的计算结果多 BMM 循环批量计算根据文档定义iterate_batch提供批量处理 Matmul 的功能调用一次 iterate_batch可以计算出多个 singleCoreM * singleCoreN 大小的 C 矩阵且该接口只支持 Norm 模板即 BatchMatmul 只支持 Norm 模板接口文档。函数签名与 Ascend C 原型对照pyasc 中该接口存在两个重载overload分别对应不同的参数组织方式# 重载一带 batch_a / batch_b 的标准形式推荐 Matmul.iterate_batch(tensor: BaseTensor, batch_a: int, batch_b: int, en_sequential_write: bool, matrix_stride_a: int 0, matrix_stride_b: int 0, matrix_stride_c: int 0, en_partial_sum: bool False, en_atomic: int 0, sync: bool True, wait_iterate_batch: bool | None None) → None # 重载二纯 Cubecube-only场景的简写形式 Matmul.iterate_batch(tensor: BaseTensor, en_partial_sum, en_atomic, en_sequential_write: bool, matrix_stride_a: int 0, matrix_stride_b: int 0, matrix_stride_c: int 0, sync: bool True) → None两种重载在 Python 侧通过OverloadDispatcher按参数形态分派见 matmul.py最终分别生成asc_MatmulIterateBatchOp与asc_MatmulIterateBatchCubeOnlyOp两个 IR 算子节点。对应的 Ascend C 函数原型如下// 重载一对应同步/异步模板化版本支持输出到 GM 或 UB template bool sync true, bool waitIterateBatch false __aicore__ inline void IterateBatch(const GlobalTensorDstT gm, uint32_t batchA, uint32_t batchB, bool enSequentialWrite, const uint32_t matrixStrideA 0, const uint32_t matrixStrideB 0, const uint32_t matrixStrideC 0, const bool enPartialSum false, const uint8_t enAtomic 0) template bool sync true __aicore__ inline void IterateBatch(const LocalTensorDstT ubCmatrix, uint32_t batchA, uint32_t batchB, bool enSequentialWrite, const uint32_t matrixStrideA 0, const uint32_t matrixStrideB 0, const uint32_t matrixStrideC 0, const bool enPartialSum false, const uint8_t enAtomic 0) // 重载二对应纯 Cube 场景版本 __aicore__ inline void IterateBatch(const GlobalTensorDstT gm, bool enPartialSum, uint8_t enAtomic, bool enSequentialWrite, const uint32_t matrixStrideA 0, const uint32_t matrixStrideB 0, const uint32_t matrixStrideC 0) __aicore__ inline void IterateBatch(const LocalTensorDstT ubCmatrix, bool enPartialSum, uint8_t enAtomic, bool enSequentialWrite, const uint32_t matrixStrideA 0, const uint32_t matrixStrideB 0, const uint32_t matrixStrideC 0)参数语义详解参数类型默认值含义tensorBaseTensor必填C 矩阵类型为 GlobalTensor输出到 GM或 LocalTensor输出到 UBbatch_aint必填左矩阵 A 的 batch 数batch_bint必填右矩阵 B 的 batch 数en_sequential_writebool必填是否开启连续写模式matrix_stride_aint0A 矩阵源操作数相邻 nd 矩阵起始地址间的偏移单位是元素matrix_stride_bint0B 矩阵源操作数相邻 nd 矩阵起始地址间的偏移单位是元素matrix_stride_cint0预留参数开发者无需关注en_partial_sumboolFalse是否将矩阵乘的结果累加于现有的 CO1 数据en_atomicint0是否开启 Atomic 操作取值范围见下文syncboolTrue设置同步或者异步模式wait_iterate_batchbool / NoneNone是否需要通过 wait_iterate_batch 接口等待 iterate_batch 执行结束仅在异步场景下使用en_sequential_write 的输出位置强约束从 matmul.py 的校验逻辑可以确认一个容易踩坑的硬性约束输出位置与连续写模式必须配对。if isinstance(tensor, GlobalTensor): check_type(en_sequential_write, [False], When output to GM, en_sequential_write should be False in iterate_batch) else: check_type(en_sequential_write, [True], When output to GM, en_sequential_write should be True in iterate_batch)即C 矩阵输出到GM 时必须传en_sequential_writeFalse输出到UB 时必须传en_sequential_writeTrue传反会直接抛出ValueError。同时接口文档的约束部分还明确指出BSNGD、SBNGD Layout 格式不支持连续写模式且异步模式不支持 iterate_batch 搬运到 UB 上即异步场景下输出只能落在 GM。en_atomic 的取值Python 侧对en_atomic做了白名单校验check_type(en_atomic, [None, 0, 1, 2, 3], ...)即仅允许None/0/1/2/3matmul.py。其中None表示不显式指定内部按未使能处理0为默认值即不开启 Atomic 操作1/2/3分别对应不同的 Atomic 累加模式。该参数在 IR 层以uint8类型KT.uint8落盘。wait_iterate_batch 与同步/异步syncTrue默认为同步模式iterate_batch调用返回即代表本次批量计算在同步点上完成syncFalse为异步模式此时可通过wait_iterate_batchTrue请求在计算完成后插入等待源码中对异步 立即等待的组合做了显式拦截raise ValueError(Param wait_iterate_batch can be True only when sync is False)matmul.pywait_iterate_batch仅在输出为 GlobalTensor 时生效若为None则内部按False处理异步模式下如需显式同步可调用无参接口Matmul.wait_iterate_batch()其配套约束为仅支持连续输出至 Global Memory且 MixDualMaster 场景下不可用。约束说明Layout、L1 Buffer 与数据类型接口文档给出了非常具体的约束清单这里逐条展开Layout 与 batch 关系BSNGD、SBNGD、BNGS1S2 Layout输入 A、B 矩阵按分形对齐后的多 Batch 数据总和应小于 L1 Buffer 的大小NORMAL Layout 没有该限制但需通过 MatmulConfig 配置 A、B 矩阵多 Batch 数据大小与 L1 Buffer 大小的关系对应get_normal_config的bmm_mode/batch_mode等参数。G 轴对齐约束对 BSNGD、SBNGD、BNGS1S2 Layout左矩阵、右矩阵的 G 轴分别为a_layout_info_g、b_layout_info_g必须满足a_layout_info_g / batch_a b_layout_info_g / batch_bNORMAL Layout 下则要求batch_a、batch_b满足倍数关系。BSNGD 场景限制不支持一次计算多行 SD需要算子程序中循环计算即(a_layout_info_n * a_layout_info_g) / batch_a、(b_layout_info_n * b_layout_info_g) / batch_b均为整数。数据格式支持BSNGD、SBNGD Layout 输入输出只支持 ND 格式数据BNGS1S2、NORMAL Layout 输入支持 ND/NZ 格式数据。Unified Buffer 输出约束输出到 Unified Buffer 时C 矩阵大小Base_m * Base_n应小于分配的 Unified Buffer 内存大小输出到 UB 且单核计算的 N 方向大小single_core_n非 32 字节对齐时C 矩阵的 CubeFormat 仅支持ND_ALIGN格式输出 C 矩阵片时硬件会自动将single_core_n方向上的数据补齐至 32 字节。能力边界不支持量化模式即不支持set_quant_scalar、set_quant_vector接口不支持 MixDualMaster当模板参数enableMixDualMaster设置为 true 时双主模式不能使用该接口不支持 int4b_t 输入A 矩阵、B 矩阵不支持int4b_t类型的输入即 BatchMatmul 不支持 int4b_t 类型的矩阵输入。此外Python 侧对输出 C 矩阵的数据类型同样有白名单[KT.int32, KT.int_, KT.float_, KT.float32, KT.half, KT.float16]matmul.py不在其中的类型会直接报 Tensor type is not supported in iterate_batch。典型调用示例BSNGD Layout 下的多 Batch Matmul接口文档给出了一个完整的 BSNGD 场景内核代码其整体流程为定义 MatmulType → 实例化 Matmul → register_matmul → init(tiling) → 循环内 set_tensor_a/set_tensor_b/set_bias → iterate_batch# 定义matmul type a_type asc.adv.MatmulType(asc.TPosition.GM, asc.CubeFormat.ND, asc.half, False, asc.LayoutMode.BSNGD) b_type asc.adv.MatmulType(asc.TPosition.GM, asc.CubeFormat.ND, asc.half, True, asc.LayoutMode.BSNGD) c_type asc.adv.MatmulType(asc.TPosition.GM, asc.CubeFormat.ND, asc.float, False, asc.LayoutMode.BNGS1S2) bias_type asc.adv.MatmulType(asc.TPosition.GM, asc.CubeFormat.ND, asc.float) mm asc.adv.Matmul(a_type, b_type, c_type, bias_type) asc.adv.register_matmul(pipe, mm) mm.init(tiling) batch_c batch_a if batch_b batch_c: batch_c batch_b g_lay tiling.a_layout_info_g if tiling.b_layout_info g_lay: g_lay tiling.b_layout_info_g for_extent tiling.a_layout_info_b * tiling.a_layout_info_n * g_lay / tiling.batch_num for i in range(for_extent): batch_offset_a i * tiling.a_layout_info_d * batch_a batch_offset_b i * tiling.b_layout_info_d * batch_b mm.set_tensor_a(gm_a[batch_offset_a], is_transpose_a_in) mm.set_tensor_b(gm_b[batch_offset_b], is_transpose_b_in) idx_c i * batch_c if tiling.c_layout_info_g 1 and (tiling.b_layout_info_g ! 1 or tiling.a_layout_info_g ! 1): d tiling.b_layout_info_g if tiling.a_layout_info_g d: d tiling.a_layout_info_g idx_c idx_c // d if tiling.is_bias: batch_offset_bias idx_c * tiling.c_layout_info_s2 mm.ste_bias(gm_bias[batch_offset_bias]) batch_offset_c idx_c * tiling.c_layout_info_s2 if c_type.layout asc.LayoutMode.BNGS1S2: batch_offset_c idx_c * tiling.c_layout_infos2 * tiling.c_layout_info_s1 mm.iterate_batch(tensorgm_c[offsetc], batch_abatch_a, batch_bbatch_b, en_sequential_writeFalse)该示例展示了三个要点Layout 混合A、B 矩阵使用 BSNGDC 矩阵使用 BNGS1S2batch 索引idx_c需按 G 轴比例折算偏移计算A/B 的 batch 偏移分别乘以a_layout_info_d/b_layout_info_dC 的 batch 偏移乘以c_layout_info_s2BNGS1S2 下还需再乘c_layout_info_s1输出到 GM 时en_sequential_writeFalse与源码校验规则一致。纯 Cube 模式与重载二的使用重载二用于纯 Cube只有矩阵计算场景参数顺序为(tensor, en_partial_sum, en_atomic, en_sequential_write, ...)。从源码matmul.py可见其强约束该形式下输出只能是 GlobalTensor否则抛出TypeError(iterate_batch interface under cube-only sence only support output to GM.)且en_sequential_write同样必须为False。这与仓库示例 examples/04_matmul_cube_only/matmul_cube_only.py 展示的纯 Cube 模式输出直落 GM的写法一致——在该示例中 Matmul 实例化未指定 Layout默认 NONE计算完成后通过iterate_all(c_global)直接输出。配合重载二使用的典型流程可参考Matmul.set_batch_num的调用示例在纯 Cube 模式下先用set_batch_num(batch_a, batch_b)在不改变 Tiling 的情况下重新设置 Batch 数再调用mm.iterate_batch(tensorgm_c, en_partial_sumFalse, en_atomic0, en_sequential_writeFalse)完成多 batch Matmul 计算。源码与测试验证NORMAL Layout 的端到端用法仓库中的泛化测试 python/test/generalization/adv/test_matmul_iterate_batch.py 提供了 NORMAL Layout 下iterate_batch的完整端到端样例float32 与 float16 两种 dtype、32/64 尺寸参数化验证。其内核核心代码如下asc.jit def matmul_kernel(a, b, c, tiling, workspace): ... matmul asc.adv.Matmul( aasc.adv.MatmulType(asc.TPosition.GM, asc.CubeFormat.ND, a_global.dtype, False, asc.LayoutMode.NORMAL), basc.adv.MatmulType(asc.TPosition.GM, asc.CubeFormat.ND, b_global.dtype, False, asc.LayoutMode.NORMAL), casc.adv.MatmulType(asc.TPosition.GM, asc.CubeFormat.ND, c_global.dtype, False, asc.LayoutMode.NORMAL), ) asc.adv.register_matmul(pipe, workspace, matmul, tiling) matmul.set_tensor_a(a_global) matmul.set_tensor_b(b_global) matmul.iterate_batch(c_global, BATCH_A, BATCH_B, False) asc.pipe_barrier(asc.PipeID.PIPE_ALL)对应 Host 侧 Tiling 生成test_matmul_iterate_batch.py通过MultiCoreMatmulTiling的set_a_layout/set_b_layout/set_c_layout/set_batch_num配置各矩阵的 Layout 轴信息与最大 Batch 数matmul_tiling host.MultiCoreMatmulTiling(host.get_ascendc_platform()) matmul_tiling.set_a_type(host.TPosition.GM, host.CubeFormat.ND, host_dtype, False) matmul_tiling.set_b_type(host.TPosition.GM, host.CubeFormat.ND, host_dtype, False) matmul_tiling.set_c_type(host.TPosition.VECCALC, host.CubeFormat.ND, host.DataType.DT_FLOAT) matmul_tiling.set_dim(16) matmul_tiling.set_org_shape(m, n, k) matmul_tiling.set_shape(m, n, k) matmul_tiling.set_traverse(host.MatrixTraverse.FIRSTM) matmul_tiling.set_buffer_space(-1, -1, -1) matmul_tiling.set_a_layout(BATCH_A, m, 1, 1, k) matmul_tiling.set_b_layout(BATCH_B, k, 1, 1, n) matmul_tiling.set_c_layout(BATCH, m, 1, 1, n) matmul_tiling.set_batch_num(BATCH) tiling asc.adv.TCubeTiling() matmul_tiling.get_tiling(tiling)测试中的 batch 组合还覆盖了 A、B batch 数相同的常规情形并验证了batch_a batch_b 2时结果与 PyTorch 参考实现a b的一致性torch.allclose(c, matmul, atol1e-3)。这从侧面印证了文档中batch_a、batch_b 需满足倍数关系的约束在 NORMAL Layout 下的具体形态。与 wait_iterate_batch 配合的异步用法异步场景的标准配对写法wait_iterate_batch 接口文档mm asc.adv.Matmul(a_type, b_type, c_type, bias_type) mm.set_tensor_a(gm_a[offset_a]) mm.set_tensor_b(gm_b[offset_b]) if tiling.is_bias: mm.set_bias(gm_bias[offset_bias]) mm.iterate_batch(tensorgm_c[offsetc], batch_abatch_a, batch_bbatch_b, en_sequential_writeFalse) mm.wait_iterate_batch()即iterate_batch发起异步批量计算后调用无参的mm.wait_iterate_batch()等待其返回从而支持连续输出到 Global Memory避免后续搬移与计算的竞争。注意wait_iterate_batch接口的约束同样包含仅支持连续输出至 Global MemoryMixDualMaster 场景不支持。小结iterate_batch 使用自查清单编写使用iterate_batch的 pyasc 算子时建议按以下清单自查模板选择BatchMatmul 仅支持 Norm 模板实例化 Matmul 时通过MatmulType的layout参数指定 LayoutBSNGD / SBNGD / BNGS1S2 / NORMAL输出位置与 en_sequential_write输出 GM 传False输出 UB 传True且 BSNGD/SBNGD 不支持连续写、异步模式不支持输出到 UBbatch 关系BSNGD/SBNGD/BNGS1S2 下满足a_layout_info_g / batch_a b_layout_info_g / batch_bNORMAL 下满足倍数关系内存BSNGD/SBNGD/BNGS1S2 多 Batch 数据总和小于 L1 Buffer输出 UB 时Base_m * Base_n小于分配的 UB 内存非 32 字节对齐时使用ND_ALIGN能力边界不使用量化接口、不使能 MixDualMaster、A/B 不使用 int4b_t异步syncFalse时按需设置wait_iterate_batch或调用wait_iterate_batch()同步。掌握上述要点后即可在 examples 与 python/test/generalization/adv 中已有示例的基础上快速编写属于自己的多 Batch Matmul 昇腾算子。【免费下载链接】pyasc本项目为Python用户提供算子编程接口支持在昇腾AI处理器上加速计算接口与Ascend C一一对应并遵守Python原生语法。项目地址: https://gitcode.com/cann/pyasc创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考