来源论文: https://arxiv.org/abs/2303.08989 生成时间: Mar 05, 2026 14:37

量子电路模拟加速：基于 Tensor Core 的 SGEMM 仿真与自动精度选择技术深度解析

0. 执行摘要

随着量子计算研究进入“量子优越性”时代，经典计算机对量子电路的模拟已成为验证量子算法和评估量子处理器性能的关键工具。在众多的模拟方法中，张量网络收缩（Tensor Network Contraction）因其在处理大规模低深度量子电路方面的卓越能力而备受青睐。然而，张量收缩的核心瓶颈在于高密度的矩阵乘法（GEMM）运算，这类运算往往占据了总计算时间的 90% 以上。

本文深入探讨了一项前沿研究，该研究提出了一种名为 cuMpSGEMM 的库，利用 NVIDIA A100 等现代 GPU 中的 Tensor Cores（张量核心）来加速单精度矩阵乘法（SGEMM）。研究的核心创新点在于：1）通过误差修正技术在混合精度硬件上模拟全 FP32 精度；2）引入基于输入张量指数分布统计的自动精度选择方案，在 FP16 和 TF32 模式间动态切换，并支持自动缩放以防止欠流。实验结果显示，该方法在模拟 Google Sycamore 53 标位量子电路和 9×9 随机量子电路（RCS）时，吞吐量提升了 1.44 至 1.86 倍，同时保持了极高的模拟保真度。这一成果为量子化学计算、量子机器学习等高性能计算（HPC）领域提供了重要的硬件加速范式。

1. 核心科学问题，理论基础，技术难点与方法细节

1.1 核心科学问题：精度与性能的权衡

量子电路模拟本质上是对高维希尔伯特空间中态矢量演化的追踪。对于张量网络模拟器而言，计算过程转化为一系列张量收缩操作。这些张量通常需要保持单精度（FP32）以确保模拟的保真度（Fidelity）。

然而，现代 GPU 的峰值性能主要由混合精度单位（如 Tensor Cores）提供。例如，NVIDIA A100 在 FP16 Tensor Core 上的理论峰值比标准 FP32 SIMT 核心高出 10 倍以上。直接使用低精度（FP16 或 TF32）会导致严重的精度损失，尤其是在深层量子电路中，累积误差会迅速降低结果的保真度。因此，如何在利用 Tensor Cores 的超高吞吐量的同时，回收并维持 FP32 的计算精度，是本作解决的核心科学问题。

1.2 理论基础：张量网络与 TTGT 算法

量子电路可以完美地表示为张量网络。每个量子门是一个张量，量子比特线则是连接这些张量的“边”。计算某个比特串出现的概率，等同于对该网络进行缩并（Contraction）。

为了利用优化的线性代数库（如 cuBLAS），研究采用了 TTGT (Transpose-Transpose-GEMM-Transpose) 算法：

  1. Transpose (Reshape): 将输入张量的索引重新排列，把需要收缩的索引合并为矩阵的一维，保留索引合并为另一维。
  2. GEMM: 执行矩阵乘法。
  3. Transpose: 将结果矩阵重新映射回张量结构。

由于张量收缩涉及大量的中间产物，GEMM 的效率直接决定了模拟器的生死。

1.3 技术难点：Tensor Core 的局限性

Tensor Core 虽然强大，但在执行 SGEMM 时面临两个主要挑战：

  1. 尾数精度不足：FP16 只有 10 位尾数，TF32 也只有 10 位，远低于 FP32 的 23 位。这会导致截断误差随矩阵规模线性增加。
  2. 指数范围限制：FP16 的指数位仅为 5 位，表示范围极窄（约 $6 \times 10^{-5}$ 到 $6.5 \times 10^4$）。在量子模拟中，振幅（Amplitude）会随着比特数增加而呈指数级减小（$2^{-(n-1)/2}$），极易发生欠流（Underflow）。

1.4 方法细节：SGEMM 仿真与误差修正

作者采用了基于补偿求和（Compensated Summation）的误差修正技术。具体公式如下：

设 $A_{F32}$ 和 $B_{F32}$ 是输入矩阵，目标是计算 $C_{F32} \approx A_{F32} \cdot B_{F32}$。

  1. 拆分输入：将 $A_{F32}$ 拆分为高位部分 $A_{low}$（转换为低精度格式）和误差部分 $\Delta A_{low} = (A_{F32} - toF32(A_{low})) \times 2^{11}$。对 $B$ 同理。
  2. 多重收缩：利用 Tensor Core 计算四项乘积： $$C_{F32} \approx A_{low} \cdot B_{low} + (\Delta A_{low} \cdot B_{low} + A_{low} \cdot \Delta B_{low}) / 2^{11}$$
  3. 舍入模式控制：作者在之前的研究（Ootomo et al. 2022）中发现，Tensor Core 内部的舍入模式（通常是 Round-to-Zero）是精度损失的主因。通过使用 FP32 SIMT 核心辅助进行特定模式的加法，可以实现对 FP32 尾数精度的完美回收。

1.5 动态精度选择逻辑

为了在性能最高的 FP16TCEC（FP16 误差修正）和指数范围更广的 TF32TCEC 之间平衡，作者设计了自动选择机制：

  • 阶段 1：全局统计。扫描输入张量，计算最大指数 $e_{max}$ 和处于 FP16 表示范围内的元素数量 $N_1$。如果绝大多数元素未欠流，直接使用 FP16TCEC。
  • 阶段 2：缩放检测。如果检测到大量欠流，但整体分布紧凑，则启动“带缩放的 FP16TCEC”。通过将所有元素整体放大（例如乘以 $2^{k}$），使其落入 FP16 的有效表示区，计算完成后再缩小回来。
  • 兜底方案：如果指数分布太散，则退回到 TF32TCEC 甚至标准 FP32。

2. 关键 Benchmark 体系与性能数据解析

2.1 测试环境

  • GPU: NVIDIA A100-SXM4-40GB (Ampere 架构)。
  • 理论峰值: FP32 SIMT 为 19.5 TFlop/s；TF32 Tensor Core 为 156 TFlop/s；FP16 Tensor Core 为 312 TFlop/s。

2.2 随机量子电路（RQC）实验数据

作者在 $7 \times 7$、$8 \times 8$、$9 \times 9$ 和 $10 \times 10$ 的矩形点阵量子电路上进行了测试。

电路规模精度模式吞吐量提升 (vs cuBLAS)保真度 (Fidelity)
$9 \times 9$ (depth 33)AUTO-0.51.86x$\sim 1.0$ (极高)
$10 \times 10$ (depth 33)AUTO-0.51.68x$\sim 1.0$
Google Sycamore (53 qubits)AUTO-0.51.44x$\sim 1.0$

数据深度分析： 在 $9 \times 9$ 电路中，由于矩阵规模巨大且指数分布相对集中，自动选择机制成功激活了 FP16TCEC。相比于传统的 cuBLAS SGEMM，其计算耗时显著下降。值得注意的是，即使加入了指数统计的额外开销（约占总时间的 5-10%），由于计算核心吞吐量的巨大红利，总时间依然大幅缩短。

2.3 精度验证（Fidelity Evaluation）

实验通过计算经典模拟得到的比特串概率与 FP64（双精度）基准结果的相对误差来评估。结果显示：

  • TF32TC (无修正)：在电路深度超过 20 层后，保真度迅速跌落至 0.5 以下，完全无法使用。
  • FP16TCEC / TF32TCEC (有修正)：保真度曲线与标准 FP32 重合，维持在 0.9999 以上。这证明了作者的误差修正算法在长链条张量收缩中能够有效抑制舍入误差的累积。

2.4 指数统计开销分析

作者展示了统计核函数的带宽效率。对于大于 $2^{11}$ 规模的矩阵，统计核函数的内存带宽利用率达到了 A100 峰值的 90% 以上（> 1500 GB/s）。这证明了通过精心优化的 CUDA Kernel，自动精度选择带来的延迟可以被有效地掩盖在计算密集的 GEMM 之后。

3. 代码实现细节与复现指南

3.1 cuMpSGEMM 库的设计理念

该研究最实用的贡献之一是开源库 cuMpSGEMM。它通过 LD_PRELOAD 机制拦截 cuBLAS 调用。这意味着用户不需要修改一行源代码，只需在执行原有的量子模拟程序前设置环境变量，即可享受加速。

3.2 关键实现技术

  • WMMA API: 使用 NVIDIA 的 Warp-level Matrix Multiply-Accumulate API 直接操纵 Tensor Core。
  • Kernel Preemption (内核抢占/终止)：为了最小化动态选择的延迟，作者采用了一种“多路发射，择优录取”的策略。程序会并发或连续启动统计 Kernel 和计算 Kernel，计算 Kernel 会读取显存中的 mode_flag_buffer，如果不匹配则立即退出（exit with nothing）。
  • 不规则形状优化: 量子模拟中常出现 $(2, 2^N, 2)$ 这种极其瘦长的矩阵乘法。cuMpSGEMM 针对此类不规则形状实现了专门的 Kernel，利用 Shared Memory 优化带宽。相比 cuBLAS 提升了近 10 倍的带宽效率（见附录 Figure 12）。

3.3 复现指南

  1. 获取代码： GitHub Repo: https://github.com/enp1s0/cuMpSGEMM
  2. 编译环境：
    • CUDA Toolkit 11.0+
    • C++14 兼容编译器
    • NVIDIA A100 或更高版本的 GPU (需要计算能力 8.0+ 以支持 TF32)
  3. 集成步骤：
    # 编译库
make
# 运行现有的量子模拟器（如基于 CuPy 或 PyTorch 的张量收缩程序）
export LD_PRELOAD=/path/to/libcumpsgemm.so
export CUMPSGEMM_MODE=AUTO # 开启自动精度选择
python your_quantum_sim.py

4. 关键引用文献与局限性评论

4.1 关键引用文献

  1. [1] Arute et al. (2019): Google Sycamore 量子优越性论文，定义了模拟的目标标准。
  2. [14] Markidis et al. (2018): 首次提出在 Tensor Core 上利用补偿求和实现精度回收的先驱工作。
  3. [19] Ootomo & Yokota (2022): 本文作者的前序工作，深入分析了 Tensor Core 内部舍入模式对精度的影响，并提出了改进的 TCEC 算法。
  4. [13] Liu et al. (2021): 神威·太湖之光上的量子模拟工作，使用了类似的混合精度思想，但采用了预先分析而非动态选择。

4.2 局限性评论

尽管该工作表现出色，但仍存在以下局限：

  1. 硬件依赖性强：该方法高度依赖 NVIDIA Tensor Core 的架构细节。虽然思想可推广，但在 AMD (CDNA/ROCm) 或 Intel (XMX) 平台上需要重新实现底层 Kernel。
  2. 小规模矩阵无效：对于 $N < 512$ 的小型矩阵，统计和缩放的开销会抵消 Tensor Core 的收益。论文中通过阈值判断跳过了这类情况，但在某些极碎碎的张量网络中，加速效果有限。
  3. 内存压力：误差修正（TCEC）由于需要存储中间的 $\Delta A$ 和 $\Delta B$，实际上增加了算术强度和寄存器压力。在显存带宽受限（Memory-bound）的场景下，性能提升幅度会收缩。
  4. 自动缩放的边界：虽然引入了 Scaling 机制，但如果张量元素的动态范围（Dynamic Range）跨度本身超过了 FP16 的表示极限（即 $e_{max} - e_{min} > 31$），Scaling 也无法救回那些发生下溢的极小值。

5. 补充内容：从量子化学视角看此技术的应用前景

作为量子化学科研人员，我们不应仅将此技术视为一种“加速器”，而应看到其对经典电子结构计算的潜移默化影响。

5.1 在耦合簇理论（CCSD/CCSD(T)）中的应用

耦合簇理论是计算化学的“金标准”，其计算量主要集中在四阶张量的收缩。传统的 CCSD 实现大量依赖于矩阵乘法。如果能将 cuMpSGEMM 的自动精度选择引入到诸如 NWChem 或 ORCA 的 GPU 加速模块中，研究人员可以在不牺牲能量精度的前提下，将计算规模从几十个原子扩展到上百个原子。

5.2 密度矩阵重整化群（DMRG）

DMRG 是处理强关联体系的有力工具，其核心在于对矩阵乘积态（MPS）的更新和收缩。这与量子电路模拟的张量网络结构高度相似。由于 DMRG 涉及大量的奇异值分解（SVD）前后的矩阵运算，精度的维持至关重要。本文的自动精度选择技术可以完美适配 DMRG 中不同迭代阶段对精度的动态需求。

5.3 展望：超越单精度的仿真

目前的论文重点在于仿真 FP32。在量子化学中，有时需要双精度（FP64）以获得化学精度（1 kcal/mol）。未来的研究方向可以探索是否能利用相同的误差补偿思想，通过 4 次或更多次的低精度 Tensor Core 收缩，来模拟 FP64。如果能实现 FP64 级别的自动精度选择，那将彻底改变高性能计算化学的版图。

5.4 总结

Ootomo 等人的这项工作向我们展示了“软硬件协同设计”的力量。它没有盲目等待硬件厂商提供更高精度的核心，而是通过巧妙的数学算法从现有的、为人工智能设计的硬件中“挤出”了科学计算所需的精度。这对于目前算力受限、算法复杂的量子模拟和量子化学领域具有极强的启发意义。