PyPTO-Pro 批量矩阵乘降维 2-D 收缩模式内核内批量塌缩与 wrapper 边界实战指南【免费下载链接】pypto-gymPyPTO-Gym 是基于 PyPTO 编程框架构建的算子与模型样例仓库项目地址: https://gitcode.com/cann/pypto-gym本文是 PyPTO-Gym 仓库中 PyPTO-Pro 算子知识库KBbatched-2d-contraction.md 的深度展开面向在昇腾 NPU 上用 PyPTO-Pro 语言实现带 batch 维度的矩阵乘类算子如多 Batch 线性层、共享权重 GQA/MLA 投影等的开发者。读完本文你将掌握批量维度塌缩为 2-D 收缩的适用判定、在内核内部而非 host 侧完成塌缩的决策规则、与其强绑定的 wrapper 边界约束以及如何正确地参考保留的 2-D 矩阵乘样例代码。模式概述什么是批量 2-D 收缩PyPTO-Pro 的 Cube 单元矩阵乘加速单元处理的是二维分形fractal数据布局下的矩阵乘。当算子面临带批量维度的收缩计算时最直接的想法往往是在 host 端把张量 reshape 成二维再送进内核。然而在 PyPTO-Pro 的单内核交付single-kernel delivery模型下这个做法会付出真实的设备开销。本模式要解决的问题是如下一类计算out[..., M, N] left[..., M, K] right[K, N]其结构特征是左操作数带前置批量维度右操作数被所有 batch 共享形状中没有 batch 维并且批量维度可以塌缩而不改变收缩的语义——因为收缩只发生在 K 维上batch 维与 M、N 维是互相独立的。该模式在 KB 的模式选择器 pattern-index.md 中被归类为当 host 侧 reshape 可以把批量收缩降维为 2-D 时使用验证状态为conceptual only概念性指导下文详述。适用条件Applies when在决定采用本模式之前必须逐条核对以下结构条件全部成立才适用左操作数有前置批量维度left的形状为[..., M, K]其中...是一个或多个 batch 维度右操作数被所有 batch 共享right的形状为[K, N]不包含 batch 维即同一份权重/矩阵被每个 batch 复用批量维度可以塌缩而不改变收缩语义由于矩阵乘只在 K 维上做收缩batch 维与 M、N 维相互独立out[b, :, :] left[b, :, :] right对每个 b 独立成立因此把[..., M, K]视作[batch_product * M, K]不会改变任何输出元素的值。反之若右操作数也随 batch 变化即right[b, K, N]这种形状每个 batch 使用的是不同的权重批量维度无法通过简单塌缩共享同一份右操作数此时不得套用本模式——这正是原文档明确列出的禁用条件之一。决策规则在内核内塌缩而不是在 host 上模式的决策规则非常明确将前置 batch 轴塌缩使收缩以 2-D 的[batch_product * M, K] x [K, N]形式运行并在输出时恢复[..., M, N]的形状。关键约束是塌缩发生的位置Collapse it inside the kernel, not on the host.原因在于batch product 本质上是对形状做的一次算术运算batch_product prod(batch_dims)内核完全可以根据真实形状计算扁平偏移flat offset直接按偏移索引操作数根本不需要物理上重新排布数据。如果把塌缩放到 host 端就会触发一次reshape——而当输入不是连续contiguous张量时还会连带触发contiguous()把整个张量物化materialize成一个可被测量到的设备内核。这直接违反 KB 中强制生效的 wrapper 边界约束 wrapper-boundary.md。该边界约束对所有算子类强制生效优先级高于本页面上的任何建议。即reshape、contiguous()、to()等一切数据形状与 dtype 处理都必须发生在pl.jit内核内部wrapper 只允许做参数校验、读取shape/ndim/dim()/size()/stride()/dtype/device/layout/numel()/storage_offset()等元数据、推导 Python 整数、用torch.empty分配当前契约声明的输出以及启动唯一一次内核。为什么 host 端 reshape 是禁区wrapper 边界的代价分析原文档与 wrapper 边界约束共同回答了为什么不能偷懒在 host 上做的问题。这不是风格偏好而是单内核交付边界同时决定了调用方的端到端设备成本代价之一是设备时间。从该规则蒸馏出的事实测量显示在 host 侧做大量数据塑形的 wrapper其设备时间占比可达总设备时间的10% 到 62%。一个真实的反模式示例wrapper-boundary.md 中收录的改名后的生成 wrapper展示了四次测量到的设备算子包围一次内核启动def op_wrapper(input_tensor, dim-1, ...): x_fp32 input_tensor.to(torch.float32) # cast - measured x_transposed x_fp32.movedim(dim, -1).contiguous() # transposecopy - measured x_2d x_transposed.reshape(M, D_full) y_2d torch.empty(M, D_out, ...) op_kernel(x_2d, y_2d, ...) # the actual work y_transposed y_2d.reshape(*non_dim_shape, D_out) y y_transposed.movedim(-1, dim).contiguous() # transposecopy - measured return y.to(out_dtype) # cast - measured内核被写成期望规范化的 FP32 连续二维输入host 就去制造这个输入——而这每一次便利都在测量窗口内被计费为一个设备内核。该约束还指出两个持久结论内核时间与可调用时间可能反向变动某算子内核提速约三分之一端到端反而变慢因为 wrapper 的膨胀快于内核的缩小以及host 侧塑形并不罕见在某次运行的生成内核中.to()、.contiguous()、.reshape()是最主要的调用。代价之二是兼容性风险。评估容器的 CANN 环境与开发机不同实测中 eval CANN 9.1.0 对 dev 9.2.0其算子清单并非开发机的超集一个在本地运行正常的 wrapper.to(torch.float32)在真实交付中于所有 fp16/bf16 用例上抛出aclnnInplaceCopy failed, error code is 561103EZ1013: aclnnInplaceCopy_1_CastAiCore cannot be found只有 fp32 用例通过。而一个不派发任何设备算子的 wrapper 则没有这种依赖。代价之三是视图 reshape的迷惑性。对已连续张量的纯视图 reshape 确实零成本、不会出现在 profile 中——但约束明确指出如果你无法从 profile 中证明 host 端 reshape 没有产生任何设备操作那它就产生了。成本与合规是两回事任何 profile 结果都不能豁免这条边界。内核体回归纯 2-D 收缩模式塌缩之后内核体就退化为普通的 2-D 收缩直接遵循 KB 中的 cube-only.md 纯 Cube 收缩模式无向量归约、无激活、无 cast、无逐元素 epilogue。其数据流为GM → Mat → Left/Right → Acc → GM即从 GM 搬运操作数到 L1Mat 内存再搬到 L0A/L0B在 L0C 中累加最后写回 GM。参考的完整单 block 骨架来自该模式引用的保留实现为with pl.section_cube(): a_mat a_l1.current() b_mat b_l1.current() a_left a_l0a.current() b_right b_l0b.current() out_acc acc.current() pl.load(a_mat, a, [0, 0]) pl.load(b_mat, b, [0, 0]) pl.move(a_left, a_mat) pl.move(b_right, b_mat) pl.matmul(out_acc, a_left, b_right) pl.store(out, out_acc, [0, 0])当 K 维需要切分为多个 K block 时cube-only 模式给出的累加语义为只有最后一个 K block 使用Finalphase之前的每个 block 都必须用Partial当目标 SDK 使用AccPhase时for block in K_blocks: is_first block 0 is_last block K_blocks - 1 phase Final if is_last else Partial if is_first: matmul(acc, left, right, phasephase) else: matmul_acc(acc, acc, left, right, phasephase)这覆盖了单 block 情形matmul(..., Final)、多 block 的首块Partial、所有中间块Partial以及仅最后一块Final。需要说明的是cube-only 模式中单 K block 已有保留的可运行参考验证而通用 K-loop 仅属概念性指导实现前应针对检测到的 SDK 版本对照官方示例验证确切的 API 签名与 buffer 轮转。保留的 2-D 参考实现及其正确读法原文档明确提醒KB 目前没有保留任何端到端实现批量塌缩的内核。唯一保留的矩阵乘参考实现 matmul_float_mmad_impl.py 只是2-D 收缩参考——它在 host 端准备布局而这正是本模式现在告诉你不要做的。该实现是单 tile 的 FP32 矩阵乘冒烟形状M, K, N 32, 16, 48均为 16 的 cube 分形倍数内核签名如下pl.jit(auto_mutexTrue) def matmul_float_mmad_kernel( a: pl.Tensor[[M, K], pl.DT_FP32], # [M, K] ( x) b: pl.Tensor[[K, N], pl.DT_FP32], # [K, N] ( y.T; wrapper transposes y[N,K]) out: pl.Tensor[[M, N], pl.DT_FP32], # [M, N] a b x y.T, FP32 accum in L0C ):其内核体通过pl.make_tile_group声明 L1Mat、L0ALeft/NZ、L0BRight/ZN、L0CAccfractal1024四个 tile group每个 tile group 绑定独立的 mutex_id然后在pl.section_cube()内完成load → move → matmul → store。该实现的验证记录为Ascend a5 (950) NPU 上 FP32 与 golden 的max_abs_diff 0.000e00逐位一致PASS配套 golden 见 matmul_float_mmad_golden.py。阅读该样例的正确姿势是读它的内核体kernel body不要读它的边界boundary。原因在于其 drivermatmul_float_mmad_wrapper做了y.t().contiguous()、torch.zeros分配输出、torch.npu.synchronize()等 host 侧操作——这些是让文件能独立运行的 harness 行为不是被许可的交付形态。样例头注释与 examples/README.md 都明确Copy the kernel. Do not copy the driver.复制内核不要复制驱动。交付 wrapper 只允许调用torch.empty分配输出其余布局工作全部放入内核。输出形状恢复与边界条件在原文档的决策规则中塌缩之外还有三条补充要求缺一不可恢复输出形状时保持原始 batch 轴顺序塌缩是[..., M, K] → [batch_product * M, K]输出必须按原来的...轴顺序还原为[..., M, N]。轴顺序是语义的一部分任何重排都会改变结果的元素位置右操作数也随 batch 变化时不得套用本模式right含 batch 维意味着每个 batch 的收缩对象不同塌缩会错误共享右操作数在目标 SDK 上验证动态形状与启动几何launch geometry塌缩后的扁平索引基于真实形状的算术动态 batch 形状下的偏移计算、tile 划分与block_dim/核数配置都需要在目标 SDK 上实测确认不能照搬静态形状的假设。相关约束可参考 tiling.mdtile 形状、布局与内存空间合法性和 tail-validshape.md动态维度与尾窗。验证状态与使用注意事项原文档对本模式的验证状态给出的是Conceptual概念性这一点必须在使用前充分理解没有保留任何端到端实现批量塌缩的内核保留的matmul_float_mmad只是 2-D 收缩参考且其在 host 上准备布局的做法已被本页面否定因此在实现批量塌缩之前应先参考目标 SDK 官方的 reshape 与动态形状示例将其作为实现的起点。这与 KB 模式选择器 pattern-index.md 对conceptual only状态的定义一致使用其中的公式与决策规则把任何研究性代码当作未经验证的输入来审视不要把它当作经过验证的起点。同时应遵循 KB 保留门retention gate见 README.md只有当某个模式被路由或选择器链接、其用途可跨实验/分支/环境复用、且技术声明引用了官方文档路径或保留的可审阅验证工件时才可进入生产选择器。KB 的完整性校验脚本 check_kb_integrity.py 会检查引用路径是否可解析、验证状态声明是否与保留工件匹配采用任何样例前都应拉取声明的分支并重新运行该校验与目标硬件上的正确性验证——历史验证记录并不能确立对另一 PyPTO/CANN 版本或平台的兼容性kernel-index.md。何时选择本模式路由与决策路径在 KB 的路由体系ROUTER.md中路由依据是计算的形状而非算子名。设计阶段遇到批量收缩时推荐决策路径为判定是否为out[..., M, N] left[..., M, K] right[K, N]结构且右操作数被所有 batch 共享若是选择本模式并在内核内完成批量塌缩与扁平偏移计算内核体按 cube-only.md 的纯 Cube 数据流编写全程守住 wrapper-boundary.md 的边界wrapper 只做元数据读取、参数校验、torch.empty分配输出与单次内核启动输出侧恢复原始 batch 轴顺序并在目标 SDK 上验证动态形状与启动几何。这套路径的核心收益在于把形状算术留在内核内以真实形状与扁平偏移表达避免了 host 端 reshape/contiguous 带来的设备算子开销与 CANN 版本兼容风险同时让内核体复用已验证的 2-D Cube 收缩骨架是 PyPTO-Pro 单内核交付模型下处理批量矩阵乘的首选形态。【免费下载链接】pypto-gymPyPTO-Gym 是基于 PyPTO 编程框架构建的算子与模型样例仓库项目地址: https://gitcode.com/cann/pypto-gym创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考