devops

[Paper] 层次精度与递归用于加速 MXU 上的对称线性求解

发布: (2026年1月13日 GMT+8 07:46)
7 min read
原文: arXiv

Source: arXiv - 2601.08082v1

Overview

本文提出了一种 portable mixed‑precision Cholesky solver，针对现代 Matrix Processing Units (MXUs) 如 NVIDIA Tensor Cores (H200) 和 AMD Matrix Cores (MI300X)。通过将经典的 Cholesky 算法重新组织为 hierarchical recursive structure，作者能够将大部分工作转移到高速低精度 GEMM 单元上，同时将数值敏感的对角块保留在高精度中，从而在不牺牲精度的前提下实现显著的加速。

关键贡献

  • 层次递归 将 Cholesky、TRSM 和 SYRK 分解为嵌套子问题，暴露出用于 MXU 加速 GEMM 的最大并行性。
  • 细粒度混合精度数据布局：大块非对角子块使用 FP16 计算，对角子块保持 FP64 以保证数值稳定性。
  • 基于 Julia 的实现，利用多重分派和动态类型推断编写硬件无关代码，同时在底层调用供应商优化的内核。
  • 性能突破：相较于全精度基线，SYRK 提速最高 27×，TRSM 提速 5×；整体上 Cholesky 相比 cuSOLVER FP64 提速 5×。
  • 精度保证：混合精度求解器的精度约为纯 FP16 方法的 ≈100×，同时保留约 ≈88 % 的 FP16 峰值吞吐量。
  • 跨供应商验证：在 NVIDIA（H200）和 AMD（MI300X）GPU 上均观察到相似的提升，展示了真正的可移植性。

方法论

  1. 递归分解 – Cholesky 分解被表示为一种递归，将矩阵划分为四个象限。每一步递归会生成更小的 TRSM 和 SYRK 任务，这些任务会再次递归细分，直至符合 MXU 瓦片大小。
  2. 精度分层 – 自定义数据结构为每个瓦片标记目标精度。非对角瓦片（对舍入误差不敏感）在送入 MXU GEMM 引擎前会被转换为 FP16。对角瓦片保持 FP64，并使用标准的高精度内核进行处理。
  3. Julia 前端 – 该算法使用纯 Julia 编写，利用多重分派在运行时自动选择合适的精度内核。代码调用 cuBLAS / rocBLAS 完成底层的 GEMM、TRSM 和 SYRK 操作，使高层逻辑保持硬件无关。
  4. 误差控制 – 在每个递归层级之后，算法执行残差校正步骤（本质上是高精度的细化），以限制 FP16 瓦片引入的误差，确保最终解满足 FP64 精度要求。

结果与发现

内核基准 (FP64)混合精度 (FP16‑FP64)加速比精度（相对于 FP64）
SYRKcuBLAS FP64Recursive FP16‑FP6427×≈100× 优于纯 FP16
TRSMcuBLAS FP64Recursive FP16‑FP645×在 FP64 容差范围内
Cholesky (overall)cuSOLVER FP64Recursive mixed‑precision5×88 % 的 FP16 峰值速度，FP64 级别精度

同样的模式在 AMD MI300X 上也成立：混合精度求解器实现了类似的加速比和误差界限，证明递归 + 精度分层策略并非依赖于单一厂商的 MXU 实现。

实际意义

  • 加速的科学工作负载 – 依赖对称求解（例如有限元分析、气候模拟、高斯过程回归）的应用程序，在不重写整个代码库的情况下，可实现数量级的壁钟时间缩短。
  • 机器学习流水线 – 许多机器学习算法（例如核方法、二阶优化器）需要 Cholesky 分解；混合精度求解器可以使这些步骤与纯 FP32 训练循环竞争。
  • 以 GPU 为中心的库 – Julia 实现展示了一条清晰的路径，使库作者能够公开混合精度线性代数原语，并自动调度到最佳可用的 MXU 内核。
  • 能源效率 – 在 MXU 上以 FP16 进行更多计算可降低每 FLOP 的功耗，这对大规模 HPC 集群和配备 GPU 加速器的边缘 AI 设备非常有价值。
  • 可移植性 – 由于高级算法与硬件无关，开发者可以以最小的代码更改面向未来的 MXU 架构（例如即将推出的 NVIDIA Hopper 或 AMD CDNA 后继产品）。

限制与未来工作

  • 内存带宽压力 – 递归分块可能会在极大矩阵上增加数据移动，在内存受限的系统上可能限制可扩展性。
  • 精度调优开销 – 目前确定最佳块大小和精度划分需要经验性调优；自适应运行时调度器可以实现自动化。
  • 稀疏扩展 – 当前工作聚焦于稠密对称矩阵；将递归方法推广到稀疏 Cholesky 仍是一个未解决的挑战。
  • 更广泛的语言支持 – 虽然 Julia 提供了简洁的实现，但将相同方法集成到 C++/Python 生态系统中将有助于推广。

总体而言，本文展示了一种令人信服的方案，能够在保持工程师所依赖的数值保证的前提下，最大限度地挖掘现代 MXU 的性能潜力。

作者

  • Vicki Carrica
  • Rabab Alomairy
  • Evelyne Ringoot
  • Alan Edelman

论文信息

  • arXiv ID: 2601.08082v1
  • 分类: cs.DC, cs.ET, cs.MS, cs.PF
  • 出版时间: 2026年1月12日
  • PDF: 下载 PDF
Back to Blog

相关文章

阅读更多 »