CANN ops-transformer 长序列线性代数算子npu_linalg 的 TSQR、Jacobi SVD 与低秩 SVD 实战指南【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer导读本文围绕 experimental/svd 目录下的 ShadowKV 长序列线性代数组件展开讲解该项目如何在 CANN ops-transformer 框架上为 Ascend NPU910B2 / 950PR即 David实现 PyTorch 线性代数算子tsqr、svd与svd_lowrank的高性能自定义算子。文章覆盖低秩 SVD 算法流水线、算子 I/O 规格、AscendC 内核与 PyTorch 包装层源码解析、编译安装流程、精度测试以及单融合算子性能数据读者读完可掌握如何在自己的 NPU 推理/训练链路中接入并验证这套npu_linalg算子。背景为什么长序列需要专门的低秩 SVD 算子ShadowKV 方法的核心思想是为 KV Cache 选取高价值 key/value 子集从而降低长上下文推理时的显存与带宽压力。该方法的算子侧依赖低秩 SVD把形状为(B, S, D)的长序列矩阵分解为U diag(S) Vᵀ其中S是序列长度可达百万量级D是隐藏维度分解得到的低秩近似用于后续的 key 选择。PyTorch 原生线性代数库面向通用 CPU/GPU 场景在 NPU 上要么没有实现要么性能不足本项目Ascend CANN / cann-ops-transformer为此提供了专用实现目标是支持 PyTorch 线性代数操作在 NPU Ascend 910B/C、Ascend 950PR 上高效运行。从 目录结构 看该组件位于仓库experimental/svd分两层ascendc/AscendC 自定义算子源码jacobi/与tsqr/负责在 NPU 上完成数值计算torch_npu_linalg/PyTorch 包装层把自定义算子封装为torch.ops.npu_linalg.*并提供精度测试。低秩 SVD 计算流水线文档给出了完整的低秩 SVD 状态机见 README.md。其核心思想是不直接对巨大的(B, S, D)矩阵做 SVD而是先用高斯随机矩阵把问题压缩到秩R规模再在较小的矩阵上做精确分解。流水线分为四步矩阵乘法与两阶段分解MatMul1Y A G其中A{B,S,D}、高斯矩阵G{D,R}得到Y{B,S,R}。这一步把S×D的大矩阵投影到S×R的高瘦矩阵TSQR对高瘦矩阵Y{B,S,R}做 Tall-Skinny QR 分解得到正交基Q{B,S,R}与上三角R因子fp32 精度内部使用 Householder QRMatMul2A Qᵀ A得到压缩后的A{B,R,D}Jacobi对A{B,R,D}做 Jacobi 旋转 SVDfp32得到U{B,R,R}、奇异值S{B,R}、右奇异向量Vᵀ{B,D,R}MatMul3U Q U把左奇异向量从压缩空间投影回原始空间得到U{B,S,R}。其中 TSQR 内部又是一个HouseholderQR ⟷ MatMul4的迭代过程。整条流水线在图中被标注为三色Low-rank SVD主流程、CUB矩阵乘法MatMul1~4、VEC向量分解Gauss、Jacobi、HouseholderQR。从源码看svd_lowrank的实现正是按aclnnMatmul → aclnnTsqr → aclnnMatmul → aclnnJacobi → aclnnMatmul的顺序串起来的见 svd_lowrank.cpp。算子 I/O 规格与参数约束维度定义参数含义取值范围Bbatch size零个或多个 batch 维≥ 1可多维S序列长度127 S 8M对应源码MD隐藏维度31 D 2K对应源码NRSVD 秩15 R D输入输出输入TensorLayoutdType描述ABSDFP32形状(*, S, D)*为零个或多个 batch 维且 S D输出TensorLayoutdType描述UBSRFP32左正交矩阵SBRFP32奇异值VᵀBDRFP32右正交转置矩阵从 tiling 源码确认的形状约束Jacobi 的 tiling 逻辑jacobi_tiling.cpp要求输入(..., M, N)中M ≤ N即矩阵必须为宽矩阵或方阵——当M ≤ N时奇异值个数sSize M、U形状为(M, M)、V形状为(M, N)否则U为(M, N)、V为(N, N)。因此低秩 SVD 流水线里特意先做Qᵀ A把(S, D)压缩成(R, D)的宽矩阵再交给 Jacobi这与 tiling 约束是自洽的。同时 tiling 中硬性检查输入/输出 dtype 必须为DT_FLOAT否则直接报错。TSQR 的 tiling 逻辑tsqr_tiling.cpp对形状有更细的约束M ≥ 128 N ≥ 16最小形状M ≤ 8*1024*1024 N ≤ 160最大形状N 另要求N % 8 0 N 168M ≥ N*8高瘦条件blockSize需满足N*2 ≤ blockSize ≤ M/4且M % blockSize 0block 数为 2 的幂。当不满足上述约束时 tiling 返回GRAPH_FAILED因此实际使用中建议参考 test_svd.py 中已验证的 shape如[48, 64*1024, 160]、[48, 128, 16]。三个核心算子Ascend 实现与 PyTorch 接口映射AscendPyTorch 参考PyTorch NPU Linalg 接口功能Tsqrtorch.linalg.qr(A)torch.ops.npu_linalg.tsqr(A)高瘦矩阵 QR 分解Jacobitorch.linalg.svd(A)torch.ops.npu_linalg.svd(A)Jacobi 旋转 SVDSvdLowranktorch.svd_lowrank(A, q6, niter2)torch.ops.npu_linalg.svd_lowrank(A, q6, niter2)低秩 SVDJacobi旋转迭代 SVDJacobi 算子的入口内核jacobi.cpp根据 tiling 中vecInstructionSize32 或 64实例化SVD::JacobiBasefloat, float, float, float, N模板主体逻辑拆分为jacobi_impl_part1/2/3.h三个实现文件共享jacobi_base.h与svd_common.h。tiling 阶段jacobi_tiling.cpp会按mSize/2与可用 Vector Core 数决定使用核数并把行数据按 256B mask 对齐、按 UB 容量计算 bucket 大小最终把numAICores设置为所用向量核的约一半C1V2 模式。PyTorch 侧封装svd.cpp直接把输入A和迭代次数num_iterations传给aclnnJacobi输出U/S/V形状分别为(*, minDim, minDim)、(*, minDim)、(*, minDim, n)其中minDim min(m, n)。TSQR分治式高瘦 QRTSQR 内核tsqr.cpp是一个两阶段算法Forward向量核 AIV 执行先把M×N输入按blockSize切成numBlocks个块逐块调用QRHouseholderSingleVec子内核做 Householder QR随后按log2(numBlocks)层迭代把上一层成对的 R 因子再次 QR 合并最终得到最终 RReorder BackwardCube 核 AIC 执行通过RunMatmul封装MatmulImpl支持是否转置 B逐层把 Q 因子组合回原尺寸期间用CrossCoreSetFlag/WaitFlag做跨核同步。内核以KERNEL_TYPE_MIX_AIC_1_2混跑 AIC/AIVbatch 维在主循环中逐个处理tsqr.cpp。PyTorch 侧封装tsqr.cpp支持可选的block_size参数不传时由 tiling 按M/N自动推导默认 16满足N*2与M/4上限约束见 tsqr_tiling.cpp。SvdLowrank随机化低秩 SVDsvd_lowrank的 C 实现svd_lowrank.cpp完整复刻了 PyTorch_lowrank.py的算法骨架校验q ∈ [0, min(m, n)]若未显式传入随机矩阵omega则生成G randn({n, q})Gauss 矩阵Y A GaclnnMatmulcubeMathType0即 KEEP_DTYPE保持 FP32 精度Q, R aclnnTsqr(Y, block_size)其中block_size在m 1024时取 1024、否则 32且不小于n*2tempA Qᵀ AaclnnMatmultempU, S, V aclnnJacobi(tempA, num_iterations5)U Q tempUᵀaclnnMatmul返回(U, S, V)。注意参数MPyTorch 低秩 SVD 中的中心化矩阵暂不支持传入时会打印TORCH_WARN警告。另外从签名看niter参数已接收但当前实现固定使用 5 次 Jacobi 迭代q才是决定秩的关键参数。目录结构详解experimental/svd/ ├── ascendc/ # AscendC 自定义算子 │ ├── jacobi/ │ │ ├── op_host/ # 算子定义与 tilingjacobi_def/tiling │ │ └── op_kernel/ # Jacobi 内核实现impl_part1/2/3 base svd_common │ ├── tsqr/ │ │ ├── op_host/ # 算子定义与 tilingtsqr_def/tiling │ │ └── op_kernel/ # TSQR 内核 QRHouseholderSingleVec │ └── CMakeLists.txt └── torch_npu_linalg/ # PyTorch 包装层 ├── npu_linalg/ │ ├── csrc/ # svd.cpp / tsqr.cpp / svd_lowrank.cpp / ops_common.* / ops_def_registration.cpp │ └── __init__.py # 把 torch.ops.npu_linalg 挂载到 torch_npu ├── tests/ │ ├── utils/utils.py # 测试工具低秩矩阵构造、omega 生成 │ └── test_svd.py # 三个算子的精度测试 ├── build_and_install.sh # 一键编译安装 wheel ├── requirements.txt └── setup.py # wheel 打包NpuExtension编译、安装与使用步骤 1编译安装 CANN 算子在 ops-transformer 仓库根目录执行以 910B2 为例SOC 类型按实际设备调整pathops-transformer cd $path bash build.sh --make_clean --experimental -j16 --pkg --socascend910b --opssvd--experimental表示编译 experimental 目录下的算子--opssvd指定仅构建 svd 算子组--pkg生成安装包-j16控制并行度。步骤 2编译安装 PyTorch 包装层cd $path/experimental/svd/torch_npu_linalg bash build_and_install.shbuild_and_install.sh内部执行python3 setup.py build bdist_wheel后对dist/*.whl做pip3 install --force-reinstall --no-deps。setup.py 通过torch_npu.utils.cpp_extension.NpuExtension把npu_linalg/csrc/*.cpp编译为npu_linalg.npu_linalg_lib扩展头文件路径取自torch_npu/include/third_party/acl/inc因此前提是先安装好 torch、torch_npu 与 CANN 环境。步骤 3精度测试pytest $path/experimental/svd/torch_npu_linalg/tests/test_svd.py示例代码 import torch import torch_npu import npu_linalg device torch.device(npu:6) A torch.randn([32, 16], dtypetorch.float32, devicedevice) U, S, Vt torch.ops.npu_linalg.svd(A, 2) U.shape, S.shape, Vt.shape (torch.Size([32, 16, 0]), torch.Size([32, 0]), torch.Size([32, 0, 0])) torch.dist(A.cpu(), U.cpu() torch.diag_embed(S.cpu()) Vt.cpu().T) tensor(3.0957e-06) A torch.randn([128, 16], dtypetorch.float32, devicedevice) Q, R torch.ops.npu_linalg.tsqr(A) Q.shape, R.shape (torch.Size([128, 16]), torch.Size([16, 16])) torch.dist(A.cpu(), Q.cpu() R.cpu()) tensor(1.0486e-06) U, S, Vt torch.ops.npu_linalg.svd_lowrank(A, q16) torch.dist(A.cpu(), U.cpu() torch.diag_embed(S.cpu()) Vt.cpu().T) tensor(3.0957e-06)说明README 示例中展示的U/S/Vt形状输出为文档当时的运行结果快照实际输出形状以当前算子实现为准svd的minDim min(m, n)决定了U与V的具体维度。torch.dist结果数量级在1e-6左右说明U diag(S) Vᵀ能很好重建原矩阵。精度验证测试用例设计test_svd.py 覆盖了三类用例Jacobi SVD(1,16,16)、(1,160,256)、(7,168,1024)、(2,1,3,168,512)带多维 batch、(48,160,512)等迭代次数 5~10重建误差阈值1e-4~1e-5TSQR(128,16)、(1,128,16)、(48,128,16)、(2,1,3,128,16)及大 shape(48,64*1024,160)误差阈值1e-5~1e-2SVD Lowrank(48,64*1024,512)、(48,128*1024,512)秩 160验证重建误差 ≤1e-2。测试通过np.random.seed(0)固定输入并把 NPU 结果搬回 CPU 后与 PyTorch 重建公式比对A ≈ U diag(S) V或A ≈ Q R。utils.py 提供create_low_rank_matrix用两个随机矩阵相乘构造低秩矩阵与get_omega读取随机投影矩阵用于构造低秩测试样本。特性说明当前版本算子仅支持FP32 实数值输入输出所有npu_linalg系列算子支持BSN 数据布局即(batch, M, N)语义精度已通过与 PyTorch 参考实现对比验证测试脚本见各算子的测试目录即 test_svd.py。此外__init__.pynpu_linalg/init.py会把torch.ops.npu_linalg下所有算子动态挂载到torch_npu模块因此除torch.ops.npu_linalg.svd(...)外也可尝试torch_npu.svd(...)等价调用若挂载失败会打印警告并提示回退到torch.ops方式。单融合算子性能参考README 给出了各算子独立测试脚本在 910B2 与 950PRDavid上的 profile 结果见 README.md 性能表torch 算子Shape [B, S, D]torch 耗时 (s)torch.npu 910B2 耗时 (s)910B2 加速比torch.npu 950PR 耗时 (s)950PR 加速比qr[48, 64K, 168]16013.212.28.3419.2svd[48, 168, 512]45.80.361270.765.4svd.lowrank[48, 64K, 512]77914.4548.7688.9从数据看在文档记录的测试环境910B2 与 950PR下qr加速 12~19 倍、svd加速 65~127 倍、svd_lowrank加速 54~89 倍。需要说明的是这些数据为 README 记录的 profile 结果具体收益依赖设备型号、算子版本、shape 与测试脚本环境部署时应以本仓库测试脚本在目标环境的实测为准。总结experimental/svd为 ShadowKV 方法提供了完整的低秩线性代数算子栈TSQR 负责把S×R高瘦矩阵高效分解Jacobi 在压缩后的R×D宽矩阵上做精确 SVD两者再通过三次 MatMul 组合成完整的svd_lowrank流水线。从架构上看AscendC 层承担数值计算与 tiling 优化Vector Core 跑 Forward、Cube Core 跑 Backward/MatMultorch_npu_linalg层负责与 PyTorch 生态对齐整个链路支持(batch, S, D)布局的 FP32 输入并配套了可复现的精度测试与性能 profile。对于需要在 NPU 上执行长序列低秩分解的开发者可直接按上文编译安装流程接入并参考测试用例的形状约束规划输入规模。【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 
长河君 · 2026/9/18 17:33:30 · 阅读 1.2k CANN ops-transformer 长序列线性代数算子：npu_linalg 的 TSQR、Jacobi SVD 与低秩 SVD 实战指南
CANN ops-transformer 长序列线性代数算子npu_linalg 的 TSQR、Jacobi SVD 与低秩 SVD 实战指南【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer导读本文围绕 experimental/svd 目录下的 ShadowKV 长序列线性代数组件展开讲解该项目如何在 CANN ops-transformer 框架上为 Ascend NPU910B2 / 950PR即 David实现 PyTorch 线性代数算子tsqr、svd与svd_lowrank的高性能自定义算子。文章覆盖低秩 SVD 算法流水线、算子 I/O 规格、AscendC 内核与 PyTorch 包装层源码解析、编译安装流程、精度测试以及单融合算子性能数据读者读完可掌握如何在自己的 NPU 推理/训练链路中接入并验证这套npu_linalg算子。背景为什么长序列需要专门的低秩 SVD 算子ShadowKV 方法的核心思想是为 KV Cache 选取高价值 key/value 子集从而降低长上下文推理时的显存与带宽压力。该方法的算子侧依赖低秩 SVD把形状为(B, S, D)的长序列矩阵分解为U diag(S) Vᵀ其中S是序列长度可达百万量级D是隐藏维度分解得到的低秩近似用于后续的 key 选择。PyTorch 原生线性代数库面向通用 CPU/GPU 场景在 NPU 上要么没有实现要么性能不足本项目Ascend CANN / cann-ops-transformer为此提供了专用实现目标是支持 PyTorch 线性代数操作在 NPU Ascend 910B/C、Ascend 950PR 上高效运行。从 目录结构 看该组件位于仓库experimental/svd分两层ascendc/AscendC 自定义算子源码jacobi/与tsqr/负责在 NPU 上完成数值计算torch_npu_linalg/PyTorch 包装层把自定义算子封装为torch.ops.npu_linalg.*并提供精度测试。低秩 SVD 计算流水线文档给出了完整的低秩 SVD 状态机见 README.md。其核心思想是不直接对巨大的(B, S, D)矩阵做 SVD而是先用高斯随机矩阵把问题压缩到秩R规模再在较小的矩阵上做精确分解。流水线分为四步矩阵乘法与两阶段分解MatMul1Y A G其中A{B,S,D}、高斯矩阵G{D,R}得到Y{B,S,R}。这一步把S×D的大矩阵投影到S×R的高瘦矩阵TSQR对高瘦矩阵Y{B,S,R}做 Tall-Skinny QR 分解得到正交基Q{B,S,R}与上三角R因子fp32 精度内部使用 Householder QRMatMul2A Qᵀ A得到压缩后的A{B,R,D}Jacobi对A{B,R,D}做 Jacobi 旋转 SVDfp32得到U{B,R,R}、奇异值S{B,R}、右奇异向量Vᵀ{B,D,R}MatMul3U Q U把左奇异向量从压缩空间投影回原始空间得到U{B,S,R}。其中 TSQR 内部又是一个HouseholderQR ⟷ MatMul4的迭代过程。整条流水线在图中被标注为三色Low-rank SVD主流程、CUB矩阵乘法MatMul1~4、VEC向量分解Gauss、Jacobi、HouseholderQR。从源码看svd_lowrank的实现正是按aclnnMatmul → aclnnTsqr → aclnnMatmul → aclnnJacobi → aclnnMatmul的顺序串起来的见 svd_lowrank.cpp。算子 I/O 规格与参数约束维度定义参数含义取值范围Bbatch size零个或多个 batch 维≥ 1可多维S序列长度127 S 8M对应源码MD隐藏维度31 D 2K对应源码NRSVD 秩15 R D输入输出输入TensorLayoutdType描述ABSDFP32形状(*, S, D)*为零个或多个 batch 维且 S D输出TensorLayoutdType描述UBSRFP32左正交矩阵SBRFP32奇异值VᵀBDRFP32右正交转置矩阵从 tiling 源码确认的形状约束Jacobi 的 tiling 逻辑jacobi_tiling.cpp要求输入(..., M, N)中M ≤ N即矩阵必须为宽矩阵或方阵——当M ≤ N时奇异值个数sSize M、U形状为(M, M)、V形状为(M, N)否则U为(M, N)、V为(N, N)。因此低秩 SVD 流水线里特意先做Qᵀ A把(S, D)压缩成(R, D)的宽矩阵再交给 Jacobi这与 tiling 约束是自洽的。同时 tiling 中硬性检查输入/输出 dtype 必须为DT_FLOAT否则直接报错。TSQR 的 tiling 逻辑tsqr_tiling.cpp对形状有更细的约束M ≥ 128 N ≥ 16最小形状M ≤ 8*1024*1024 N ≤ 160最大形状N 另要求N % 8 0 N 168M ≥ N*8高瘦条件blockSize需满足N*2 ≤ blockSize ≤ M/4且M % blockSize 0block 数为 2 的幂。当不满足上述约束时 tiling 返回GRAPH_FAILED因此实际使用中建议参考 test_svd.py 中已验证的 shape如[48, 64*1024, 160]、[48, 128, 16]。三个核心算子Ascend 实现与 PyTorch 接口映射AscendPyTorch 参考PyTorch NPU Linalg 接口功能Tsqrtorch.linalg.qr(A)torch.ops.npu_linalg.tsqr(A)高瘦矩阵 QR 分解Jacobitorch.linalg.svd(A)torch.ops.npu_linalg.svd(A)Jacobi 旋转 SVDSvdLowranktorch.svd_lowrank(A, q6, niter2)torch.ops.npu_linalg.svd_lowrank(A, q6, niter2)低秩 SVDJacobi旋转迭代 SVDJacobi 算子的入口内核jacobi.cpp根据 tiling 中vecInstructionSize32 或 64实例化SVD::JacobiBasefloat, float, float, float, N模板主体逻辑拆分为jacobi_impl_part1/2/3.h三个实现文件共享jacobi_base.h与svd_common.h。tiling 阶段jacobi_tiling.cpp会按mSize/2与可用 Vector Core 数决定使用核数并把行数据按 256B mask 对齐、按 UB 容量计算 bucket 大小最终把numAICores设置为所用向量核的约一半C1V2 模式。PyTorch 侧封装svd.cpp直接把输入A和迭代次数num_iterations传给aclnnJacobi输出U/S/V形状分别为(*, minDim, minDim)、(*, minDim)、(*, minDim, n)其中minDim min(m, n)。TSQR分治式高瘦 QRTSQR 内核tsqr.cpp是一个两阶段算法Forward向量核 AIV 执行先把M×N输入按blockSize切成numBlocks个块逐块调用QRHouseholderSingleVec子内核做 Householder QR随后按log2(numBlocks)层迭代把上一层成对的 R 因子再次 QR 合并最终得到最终 RReorder BackwardCube 核 AIC 执行通过RunMatmul封装MatmulImpl支持是否转置 B逐层把 Q 因子组合回原尺寸期间用CrossCoreSetFlag/WaitFlag做跨核同步。内核以KERNEL_TYPE_MIX_AIC_1_2混跑 AIC/AIVbatch 维在主循环中逐个处理tsqr.cpp。PyTorch 侧封装tsqr.cpp支持可选的block_size参数不传时由 tiling 按M/N自动推导默认 16满足N*2与M/4上限约束见 tsqr_tiling.cpp。SvdLowrank随机化低秩 SVDsvd_lowrank的 C 实现svd_lowrank.cpp完整复刻了 PyTorch_lowrank.py的算法骨架校验q ∈ [0, min(m, n)]若未显式传入随机矩阵omega则生成G randn({n, q})Gauss 矩阵Y A GaclnnMatmulcubeMathType0即 KEEP_DTYPE保持 FP32 精度Q, R aclnnTsqr(Y, block_size)其中block_size在m 1024时取 1024、否则 32且不小于n*2tempA Qᵀ AaclnnMatmultempU, S, V aclnnJacobi(tempA, num_iterations5)U Q tempUᵀaclnnMatmul返回(U, S, V)。注意参数MPyTorch 低秩 SVD 中的中心化矩阵暂不支持传入时会打印TORCH_WARN警告。另外从签名看niter参数已接收但当前实现固定使用 5 次 Jacobi 迭代q才是决定秩的关键参数。目录结构详解experimental/svd/ ├── ascendc/ # AscendC 自定义算子 │ ├── jacobi/ │ │ ├── op_host/ # 算子定义与 tilingjacobi_def/tiling │ │ └── op_kernel/ # Jacobi 内核实现impl_part1/2/3 base svd_common │ ├── tsqr/ │ │ ├── op_host/ # 算子定义与 tilingtsqr_def/tiling │ │ └── op_kernel/ # TSQR 内核 QRHouseholderSingleVec │ └── CMakeLists.txt └── torch_npu_linalg/ # PyTorch 包装层 ├── npu_linalg/ │ ├── csrc/ # svd.cpp / tsqr.cpp / svd_lowrank.cpp / ops_common.* / ops_def_registration.cpp │ └── __init__.py # 把 torch.ops.npu_linalg 挂载到 torch_npu ├── tests/ │ ├── utils/utils.py # 测试工具低秩矩阵构造、omega 生成 │ └── test_svd.py # 三个算子的精度测试 ├── build_and_install.sh # 一键编译安装 wheel ├── requirements.txt └── setup.py # wheel 打包NpuExtension编译、安装与使用步骤 1编译安装 CANN 算子在 ops-transformer 仓库根目录执行以 910B2 为例SOC 类型按实际设备调整pathops-transformer cd $path bash build.sh --make_clean --experimental -j16 --pkg --socascend910b --opssvd--experimental表示编译 experimental 目录下的算子--opssvd指定仅构建 svd 算子组--pkg生成安装包-j16控制并行度。步骤 2编译安装 PyTorch 包装层cd $path/experimental/svd/torch_npu_linalg bash build_and_install.shbuild_and_install.sh内部执行python3 setup.py build bdist_wheel后对dist/*.whl做pip3 install --force-reinstall --no-deps。setup.py 通过torch_npu.utils.cpp_extension.NpuExtension把npu_linalg/csrc/*.cpp编译为npu_linalg.npu_linalg_lib扩展头文件路径取自torch_npu/include/third_party/acl/inc因此前提是先安装好 torch、torch_npu 与 CANN 环境。步骤 3精度测试pytest $path/experimental/svd/torch_npu_linalg/tests/test_svd.py示例代码 import torch import torch_npu import npu_linalg device torch.device(npu:6) A torch.randn([32, 16], dtypetorch.float32, devicedevice) U, S, Vt torch.ops.npu_linalg.svd(A, 2) U.shape, S.shape, Vt.shape (torch.Size([32, 16, 0]), torch.Size([32, 0]), torch.Size([32, 0, 0])) torch.dist(A.cpu(), U.cpu() torch.diag_embed(S.cpu()) Vt.cpu().T) tensor(3.0957e-06) A torch.randn([128, 16], dtypetorch.float32, devicedevice) Q, R torch.ops.npu_linalg.tsqr(A) Q.shape, R.shape (torch.Size([128, 16]), torch.Size([16, 16])) torch.dist(A.cpu(), Q.cpu() R.cpu()) tensor(1.0486e-06) U, S, Vt torch.ops.npu_linalg.svd_lowrank(A, q16) torch.dist(A.cpu(), U.cpu() torch.diag_embed(S.cpu()) Vt.cpu().T) tensor(3.0957e-06)说明README 示例中展示的U/S/Vt形状输出为文档当时的运行结果快照实际输出形状以当前算子实现为准svd的minDim min(m, n)决定了U与V的具体维度。torch.dist结果数量级在1e-6左右说明U diag(S) Vᵀ能很好重建原矩阵。精度验证测试用例设计test_svd.py 覆盖了三类用例Jacobi SVD(1,16,16)、(1,160,256)、(7,168,1024)、(2,1,3,168,512)带多维 batch、(48,160,512)等迭代次数 5~10重建误差阈值1e-4~1e-5TSQR(128,16)、(1,128,16)、(48,128,16)、(2,1,3,128,16)及大 shape(48,64*1024,160)误差阈值1e-5~1e-2SVD Lowrank(48,64*1024,512)、(48,128*1024,512)秩 160验证重建误差 ≤1e-2。测试通过np.random.seed(0)固定输入并把 NPU 结果搬回 CPU 后与 PyTorch 重建公式比对A ≈ U diag(S) V或A ≈ Q R。utils.py 提供create_low_rank_matrix用两个随机矩阵相乘构造低秩矩阵与get_omega读取随机投影矩阵用于构造低秩测试样本。特性说明当前版本算子仅支持FP32 实数值输入输出所有npu_linalg系列算子支持BSN 数据布局即(batch, M, N)语义精度已通过与 PyTorch 参考实现对比验证测试脚本见各算子的测试目录即 test_svd.py。此外__init__.pynpu_linalg/init.py会把torch.ops.npu_linalg下所有算子动态挂载到torch_npu模块因此除torch.ops.npu_linalg.svd(...)外也可尝试torch_npu.svd(...)等价调用若挂载失败会打印警告并提示回退到torch.ops方式。单融合算子性能参考README 给出了各算子独立测试脚本在 910B2 与 950PRDavid上的 profile 结果见 README.md 性能表torch 算子Shape [B, S, D]torch 耗时 (s)torch.npu 910B2 耗时 (s)910B2 加速比torch.npu 950PR 耗时 (s)950PR 加速比qr[48, 64K, 168]16013.212.28.3419.2svd[48, 168, 512]45.80.361270.765.4svd.lowrank[48, 64K, 512]77914.4548.7688.9从数据看在文档记录的测试环境910B2 与 950PR下qr加速 12~19 倍、svd加速 65~127 倍、svd_lowrank加速 54~89 倍。需要说明的是这些数据为 README 记录的 profile 结果具体收益依赖设备型号、算子版本、shape 与测试脚本环境部署时应以本仓库测试脚本在目标环境的实测为准。总结experimental/svd为 ShadowKV 方法提供了完整的低秩线性代数算子栈TSQR 负责把S×R高瘦矩阵高效分解Jacobi 在压缩后的R×D宽矩阵上做精确 SVD两者再通过三次 MatMul 组合成完整的svd_lowrank流水线。从架构上看AscendC 层承担数值计算与 tiling 优化Vector Core 跑 Forward、Cube Core 跑 Backward/MatMultorch_npu_linalg层负责与 PyTorch 生态对齐整个链路支持(batch, S, D)布局的 FP32 输入并配套了可复现的精度测试与性能 profile。对于需要在 NPU 上执行长序列低秩分解的开发者可直接按上文编译安装流程接入并参考测试用例的形状约束规划输入规模。【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考