企业官网建设流程全解析

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 NPU（910B2 / 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规模，再在较小的矩阵上做精确分解。流水线分为四步矩阵乘法与两阶段分解：

1. MatMul1：Y = A @ G，其中A{B,S,D}、高斯矩阵G{D,R}，得到Y{B,S,R}。这一步把S×D的大矩阵投影到S×R的"高瘦"矩阵；
2. TSQR：对高瘦矩阵Y{B,S,R}做 Tall-Skinny QR 分解，得到正交基Q{B,S,R}与上三角R因子（fp32 精度，内部使用 Householder QR）；
3. MatMul2：A' = Qᵀ @ A，得到压缩后的A'{B,R,D}；
4. Jacobi：对A'{B,R,D}做 Jacobi 旋转 SVD（fp32），得到U'{B,R,R}、奇异值S{B,R}、右奇异向量Vᵀ{B,D,R}；
5. MatMul3：U = 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 规格与参数约束

维度定义

输入输出

输入

输出

从 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 <= 168）；
• M ≥ N*8（高瘦条件）；
• blockSize需满足N*2 ≤ blockSize ≤ M/4且M % blockSize == 0，block 数为 2 的幂。

当不满足上述约束时 tiling 返回GRAPH_FAILED，因此实际使用中建议参考 test_svd.py 中已验证的 shape（如[48, 64*1024, 160]、[48, 128, 16]）。

三个核心算子：Ascend 实现与 PyTorch 接口映射

Jacobi：旋转迭代 SVD

Jacobi 算子的入口内核（jacobi.cpp）根据 tiling 中vecInstructionSize（32 或 64）实例化SVD::JacobiBase<float, 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：分治式高瘦 QR

TSQR 内核（tsqr.cpp）是一个两阶段算法：

• Forward（向量核 AIV 执行）：先把M×N输入按blockSize切成numBlocks个块，逐块调用QRHouseholderSingleVec子内核做 Householder QR；随后按log2(numBlocks)层迭代，把上一层成对的 R 因子再次 QR 合并，最终得到最终 R；
• Reorder + Backward（Cube 核 AIC 执行）：通过RunMatmul（封装MatmulImpl，支持是否转置 B）逐层把 Q 因子组合回原尺寸，期间用CrossCoreSetFlag/WaitFlag做跨核同步。

内核以KERNEL_TYPE_MIX_AIC_1_2混跑 AIC/AIV，batch 维在主循环中逐个处理（tsqr.cpp）。PyTorch 侧封装（tsqr.cpp）支持可选的block_size参数，不传时由 tiling 按M/N自动推导（默认 16，满足N*2与M/4上限约束，见 tsqr_tiling.cpp）。

SvdLowrank：随机化低秩 SVD

svd_lowrank的 C++ 实现（svd_lowrank.cpp）完整复刻了 PyTorch_lowrank.py的算法骨架：

1. 校验q ∈ [0, min(m, n)]；
2. 若未显式传入随机矩阵omega，则生成G = randn({n, q})（Gauss 矩阵）；
3. Y = A @ G（aclnnMatmul，cubeMathType=0即 KEEP_DTYPE，保持 FP32 精度）；
4. Q, R = aclnnTsqr(Y, block_size)，其中block_size在m > 1024时取 1024、否则 32，且不小于n*2；
5. tempA = Qᵀ @ A（aclnnMatmul）；
6. tempU, S, V = aclnnJacobi(tempA, num_iterations=5)；
7. U = Q @ tempUᵀ（aclnnMatmul），返回(U, S, V)。

注意：参数M（PyTorch 低秩 SVD 中的中心化矩阵）暂不支持，传入时会打印TORCH_WARN警告。另外从签名看niter参数已接收但当前实现固定使用 5 次 Jacobi 迭代，q才是决定秩的关键参数。

目录结构详解

experimental/svd/ ├── ascendc/ # AscendC 自定义算子 │ ├── jacobi/ │ │ ├── op_host/ # 算子定义与 tiling（jacobi_def/tiling） │ │ └── op_kernel/ # Jacobi 内核实现（impl_part1/2/3 + base + svd_common） │ ├── tsqr/ │ │ ├── op_host/ # 算子定义与 tiling（tsqr_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 类型按实际设备调整）：

path="ops-transformer" cd $path bash build.sh --make_clean --experimental -j16 --pkg --soc=ascend910b --ops=svd

--experimental表示编译 experimental 目录下的算子，--ops=svd指定仅构建 svd 算子组，--pkg生成安装包，-j16控制并行度。

步骤 2：编译安装 PyTorch 包装层

cd $path/experimental/svd/torch_npu_linalg bash build_and_install.sh

build_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], dtype=torch.float32, device=device) >>> 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], dtype=torch.float32, device=device) >>> 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, q=16) >>> 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-5；
• TSQR：(128,16)、(1,128,16)、(48,128,16)、(2,1,3,128,16)及大 shape(48,64*1024,160)，误差阈值1e-5~1e-2；
• SVD 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（读取随机投影矩阵），用于构造低秩测试样本。

特性说明

1. 当前版本算子仅支持FP32 实数值输入输出；
2. 所有npu_linalg系列算子支持BSN 数据布局（即(batch, M, N)语义）；
3. 精度已通过与 PyTorch 参考实现对比验证，测试脚本见各算子的测试目录（即 test_svd.py）。

此外，__init__.py（npu_linalg/init.py）会把torch.ops.npu_linalg下所有算子动态挂载到torch_npu模块，因此除torch.ops.npu_linalg.svd(...)外，也可尝试torch_npu.svd(...)等价调用（若挂载失败会打印警告并提示回退到torch.ops方式）。

单融合算子性能参考

README 给出了各算子独立测试脚本在 910B2 与 950PR（David）上的 profile 结果（见 README.md 性能表）：

从数据看，在文档记录的测试环境（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/MatMul），torch_npu_linalg层负责与 PyTorch 生态对齐，整个链路支持(batch, S, D)布局的 FP32 输入，并配套了可复现的精度测试与性能 profile。对于需要在 NPU 上执行长序列低秩分解的开发者，可直接按上文编译安装流程接入，并参考测试用例的形状约束规划输入规模。

【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库，实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer