CANN pypto-gym 实战transpose_quant_batch_matmul 算子的 MXFP8 量化批量矩阵乘法带转置实现解析【免费下载链接】pypto-gymPyPTO-Gym 是基于 PyPTO 编程框架构建的算子与模型样例仓库项目地址: https://gitcode.com/cann/pypto-gymtranspose_quant_batch_matmul 是 CANN pypto-gym 仓库中基于 PyPTO 编程框架实现的实验性量化矩阵乘算子它在一次 kernel 内完成「批量矩阵乘法 MXFP8 块量化 多种 perm 转置组合 FP16/BF16 输出」并支持 M 轴动态化是学习 PyPTO scaled_mm、E8M0 缩放因子布局与 tile 配置的典型样例。阅读本文后你将掌握该算子的数学语义、ShapeConfig 各字段含义、permX2 与 b_trans 的映射关系、M 轴切分实现策略以及如何通过仓库内测试用例与 Golden 实现完成精度验证。算子定位与产品支持情况该算子位于仓库 src/pypto_gym/ops/pypto_tensor/experimental/matmul/transpose_quant_batch_matmul/与 gmm_mxfp8、quant_matmul_reduce_sum、quant_batch_matmul 等同属实验性 matmul 算子族见 matmul 目录总览。产品支持情况如下产品支持情况Ascend 950PR支持Atlas A3 训练系列 / Atlas A3 推理系列不支持Atlas A2 训练系列 / Atlas A2 推理系列不支持注意该算子仅面向 Ascend 950PR测试用例也通过pytest.mark.soc(950)做了平台限定见 test_transpose_quant_batch_matmul.py在 A2/A3 平台上不可直接运行。算子语义与数学公式transpose_quant_batch_matmul实现基于 MXFP8 量化的批量矩阵乘法带转置。对每个 batch 索引 bb ∈ [0, B-1]计算公式为$$ out[:, b, :] permY\left(permX1(x1)[:, b, :] \times permX2(x2)[b, :, :]\right) $$其中各张量含义为x1形状[M, B, K]FP8 格式M 轴动态x2形状[B, K, N]permX2[0,1,2]或[B, N, K]permX2[0,2,1]permX1[1,0,2]x1 从[M,B,K]变为[B,M,K]permX2[0,1,2]或[0,2,1]x2 的布局决定 scaled_mm 的 b_trans 参数permY[1,0,2]输出从[B,M,N]变为[M,B,N]。核心参数共 4 个参数名说明示例值M左矩阵的行维度动态8 ~ 32768K矩阵乘的公共维度128N右矩阵的列维度512Bbatch 维度编译期固定128输入 x1形状[M, B, K]FP8M 轴动态输入 x2形状[B, K, N]或[B, N, K]取决于 permX2输出 out形状[M, B, N]FP16 或 BF16。输入输出规格输入参数参数名类型形状描述x1Tensor[M, B, K]左矩阵FP8E4M3 或 FP8E5M2x2Tensor[B, K, N]或[B, N, K]右矩阵FP8E4M3 或 FP8E5M2x1ScaleTensor[M, B, K//64, 2]左矩阵缩放因子FP8E8M0x2ScaleTensor[B, K//64, N, 2]或[B, N, K//64, 2]右矩阵缩放因子FP8E8M0输出参数参数名类型形状描述outTensor[M, B, N]输出矩阵FP16 或 BF16MX 量化约束⚠️ K 必须是 64 的倍数这是 MXFP8 块量化格式的硬性要求。调用前必须验证的约束检查清单K 轴对齐K % 64 0perm 组合permX2 决定 x2 和 x2Scale 的形状Scale 形状符合 MXFP8 格式要求输出 dtypedtype1→FP16dtype27→BF16B 轴固定batch_size 在编译期固定。MXFP8 量化说明MX 量化Microscaling Quantization是一种基于块缩放的量化格式量化块大小每 64 个元素共享一个缩放因子缩放因子格式FP8E8M08 位纯指数隐含 mantissa1.0数据格式FP8E4M3FN 或 FP8E5M2x1Scale形状[M, B, K//64, 2]x2Scale形状取决于 permX2。在仓库测试中缩放因子的生成使用torch.float8_e8m0fnu类型数据范围约束在[0.9, 1.1]附近FP8 数据则通过torch.float8_e4m3fn/torch.float8_e5m2生成见 test_transpose_quant_batch_matmul.py。perm 组合说明permX2[0, 1, 2]K,N 顺序x2 形状[B, K, N]x2Scale 形状[B, K//64, N, 2]scaled_mm 参数无 b_trans计算路径[M,K] × [K,N] → [M,N]。permX2[0, 2, 1]N,K 反序x2 形状[B, N, K]x2Scale 形状[B, N, K//64, 2]scaled_mm 参数b_transTrue, scale_b_transTrue计算路径[M,K] × [N,K]^T → [M,N]。这一映射在 kernel 实现中直接体现实现文件 transpose_quant_batch_matmul_impl.py 依据permX2 [0,1,2]判断是否给pypto.scaled_mm传入b_transTrue, scale_b_transTrue硬件内部完成转置避免额外数据搬运。核心实现M 轴切分与非并行循环kernel 入口为transpose_quant_batch_mat_mul_kernel通过pypto.frontend.jit装饰其整体结构如下pypto.frontend.jit( pass_options{ auto_mix_partition: 1, cube_l1_reuse_setting: {-1: 2}, cube_nbuffer_setting: {-1: 4}, vec_nbuffer_setting: {-2: 1, -1: 2}, }, runtime_options{stitch_function_max_num: 1024, device_sched_mode: 1}, ) def transpose_quant_batch_mat_mul_kernel(x1, x2, x1Scale, x2Scale, out, tile_config): ...ShapeConfig 数据结构ShapeConfig是 dataclass字段含义如下见 transpose_quant_batch_matmul_impl.py字段说明默认值ori_shape原始形状[M, K, N]M 在 kernel 中动态必填batch_sizeB 轴大小编译期固定必填m_tile_shapecube 运算 M 维 tile必填k_tile_shapecube 运算 K 维 tile必填n_tile_shapecube 运算 N 维 tile必填vector_tile_shape向量运算 tile 形状必填num_k_groupsK 轴分组数1num_n_groupsN 轴分组数1in_dtype输入数据类型DT_FP8E4M3out_dtype输出数据类型DT_BF16permX1x1 的 perm[1,0,2]permX2x2 的 perm[0,1,2]permY输出 perm[1,0,2]description测试用例描述M 轴切分策略避免 IR 爆炸实现采用「M 轴切分 非并行循环」策略kernel docstring 中明确说明其动机外层 LOOP_M非并行将 M 切分为 m_chunk_size 的 tile 依次迭代内层 LOOP_B非并行遍历 B 个 batchreshapex1 重排为[M, B*K]x1Scale 重排为[M, B*K//64, 2]每次迭代从重排后的二维张量上按行区间[m_begin:m_end]与列区间[begin:end]切片将 mm_result 通过pypto.assemble写入 local_out 的[m_begin, out_pos]最终将 local_out[M, B*N]reshape 回[M, B, N]。选择非并行循环而非 parallelTrue 的原因在于并行循环会触发 LoopUnroll 与 ExpandFunction 展开导致 IR 图爆炸、编译缓慢而把循环迭代放到运行时处理可以保持 IR 图小巧加快编译速度。核心循环体batch 切分 scaled_mm 调用如下x1_reshape pypto.reshape(x1, [M, B * K], inplaceTrue) x1_scale_reshape pypto.reshape(x1Scale, [M, B * K // 64, 2], inplaceTrue) local_out pypto.Tensor(shape(M, B * N), dtypeout_dtype) for b_idx in range(B): begin b_idx * K end (b_idx 1) * K x1_slice x1_reshape[:, begin:end] x1_scale_slice x1_scale_reshape[:, begin // 64:end // 64, :] x2_slice x2[b_idx, :, :] x2_scale_slice x2Scale[b_idx, :, :, :] if permX2 [0, 1, 2]: mm_result pypto.scaled_mm(x1_slice, x2_slice, out_dtype, x1_scale_slice, x2_scale_slice) else: mm_result pypto.scaled_mm(x1_slice, x2_slice, out_dtype, x1_scale_slice, x2_scale_slice, b_transTrue, scale_b_transTrue) out_pos b_idx * N pypto.assemble(mm_result, [0, out_pos], local_out) out[:, :, :] pypto.reshape(local_out, [M, B, N])实现特点小结MXFP8 量化数据支持 FP8E4M3FN 和 FP8E5M2缩放因子使用 FP8E8M0M 轴动态化同一编译 kernel 支持不同 M 大小运行时通过ori_shape[0]获取B 轴切分非并行 LOOP_B 逐 batch 独立计算避免 parallel 展开导致的 IR 膨胀perm 组合支持通过 permX2 决定 scaled_mm 的 b_trans 参数硬件内完成转置输出 dtype 灵活支持 FP16 和 BF16 输出Cache 策略输入 tensor 使用 NONE_CACHEABLE减少 L1 缓存占用MXFP8 数据一次性读取无需缓存复用。调用示例test1小 M b_trans 配置test_transpose_quant_batch_matmul( ShapeConfig( ori_shape[8, 128, 512], # [M, K, N] batch_size128, # B轴大小 m_tile_shape[256, 256], k_tile_shape[128, 128], n_tile_shape[256, 256], vector_tile_shape[1, 128, 256, 32], in_dtypepypto.DT_FP8E5M2, out_dtypepypto.DT_BF16, permX1[1, 0, 2], permX2[0, 2, 1], # N,K反序 → b_transTrue permY[1, 0, 2], descriptiontest1 ) )test3大 M 标准配置test_transpose_quant_batch_matmul( ShapeConfig( ori_shape[8192, 128, 512], # [M, K, N] batch_size128, m_tile_shape[256, 256], k_tile_shape[128, 128], n_tile_shape[256, 256], vector_tile_shape[1, 128, 256, 32], in_dtypepypto.DT_FP8E4M3, out_dtypepypto.DT_FP16, permX1[1, 0, 2], permX2[0, 1, 2], # K,N顺序 → 无b_trans permY[1, 0, 2], descriptiontest3 ) )注意仓库测试文件 test_transpose_quant_batch_matmul.py 中实际注册的用例 tile 配置与 README 示例略有差异如 test1 使用m_tile_shape[128,128]、n_tile_shape[512,512]并且通过FILTERED_CONFIGS过滤掉了 test3大 M 用例默认不在 pytest 参数化中执行仅在__main__手动运行从源码结构看这是为了控制常规 CI 测试时长。优化配置参考NBuffer 配置cube_nbuffer_setting{-1:2}vec_nbuffer_setting{-2:1,-1:16}Cache 策略NONE_CACHEABLE用于所有输入 tensorCube L1 复用cube_l1_reuse_setting{-1:2}。kernel 内实际 pass_options 为cube_nbuffer_setting{-1:4}、vec_nbuffer_setting{-2:1,-1:2}与 README 记录的历史优化配置略有出入实际以 实现文件 为准。性能数据配置MKNB预估时间test18128512128~108 ustest2128128512128~122 ustest38192128512128~9583 ustest432768128512128~40635 us以上为文档记录的预估时间Ascend 950PR 平台实际性能随环境、tile 配置与编译选项变化应以实测为准。精度验证与 Golden 实现容差设置相对容差RTOL1e-3绝对容差ATOL1e-3对比工具numpy.testing.assert_allclose。验证方法Golden 实现纯 PyTorch 实现作为精度基准三态标记[PRECISION_PASS]或[PRECISION_FAIL]对比工具numpy.testing.assert_allclose。Golden 实现在 transpose_quant_batch_matmul_golden.py其compute_golden_result完整复现了量化的逆过程可作为理解 MXFP8 语义的参考实现FP8 → FP32x.float()E8M0 Scale → FP32对 x1 应用permute(permX1)[M,B,K] → [B,M,K]对 x2 应用permute(permX2)缩放因子 reshape 后按块用torch.repeat_interleave(..., repeats32)广播为逐元素缩放每 64 元素一个缩放因子展开为 2×32 结构反量化x * scaletorch.matmul批量矩阵乘输出permute(permY)并按 dtype 标志转 FP16dtype1或 BF16dtype27。测试用例test_transpose_quant_batch_matmul.py的整体流程为生成 MXFP8 输入 → 计算 Golden → 包装器把张量搬到 NPU 并分配输出 → 调用 kernel →assert_allclose(golden, result, rtol1e-3, atol1e-3)比对。包装器transpose_quant_batch_matmul中 N 的取值逻辑也印证了 permX2 与形状的关系N x2.shape[-1] if permX2 [0,1,2] else x2.shape[1]。测试用例一览测试名称MKNBpermX2说明test18128512128[0,2,1]小 Mb_trans 验证test2128128512128[0,1,2]中 M标准验证test38192128512128[0,1,2]大 M 性能验证test432768128512128[0,1,2]超大规模验证常见问题Q1为什么 permX2[0,2,1] 需要 b_transTruepermX2[0,2,1] 表示 x2 的布局是[B,N,K]b_transTrue让 scaled_mm 在硬件内部完成转置scale_b_transTrue同时处理 scale tensor 的转置这样避免了额外的数据搬运操作。Q2M 轴动态化如何实现编译期不固定 M 值同一 kernel 可处理不同 M 的输入B 轴在编译期固定batch_size 参数通过循环切分M 轴动态化减少重编译开销适配不同输入规模从源码实现看M 通过tile_config.ori_shape[0]在运行时读取配合非并行 LOOP 切片完成。Q3为什么输入 tensor 使用 NONE_CACHEABLEMXFP8 数据一次性读取无需缓存复用NONE_CACHEABLE 减少 L1 缓存占用为输出 tensor 腾出空间对单次读取场景无性能损失。参考文档SPEC.md- 详细需求规格API_REPORT.md- API 映射分析DESIGN.md- 详细设计文档matmul 算子族总览kernel 实现精度验证测试Golden 参考实现版本历史版本日期说明v1.02026-06-03初始版本支持 MXFP8 批量矩阵乘法带转置【免费下载链接】pypto-gymPyPTO-Gym 是基于 PyPTO 编程框架构建的算子与模型样例仓库项目地址: https://gitcode.com/cann/pypto-gym创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考