[Paper] FalconGEMM：通过低复杂度矩阵乘法超越硬件峰值

发布: 4天前 (2026年5月7日 GMT+8 19:41)

7 分钟阅读

原文: arXiv

Source: arXiv - 2605.06057v1

概述

本文介绍了 FalconGEMM，一个跨平台框架，使“低复杂度”矩阵乘法算法（LCMAs）在实际深度学习工作负载中变得可行。通过自动生成、优化并为特定硬件目标选择最佳算法，FalconGEMM 在 GPU 和 CPU 上始终超越传统 GEMM 库乃至其他 LCMA 方案的性能。

可移植部署模块 – 生成硬件特定代码，使相同的 LCMA 能在 GPU（H100、A100）、ARM CPU 和 x86 CPU 上运行，无需手动重新调优。
组并行优化 – 新颖的片上数据复用和并行资源调度，降低内存带宽压力，最大化计算利用率。
轻量决策模块 – 一个分析性能模型，预测任意矩阵形状和硬件配置下最快的算法，实现运行时选择。
全面评估 – 在最先进的 GEMM 库（cuBLAS、CUTLASS、Intel MKL）上实现 7.6 %–17.9 % 的加速，在大型语言模型（LLM）训练和推理任务中相较于竞争的 LCMA 框架（如 AlphaTensor）实现 12.4 %–55.6 % 的提升。

Algorithm Catalog – 作者收集了一套 LCMA（例如 Strassen‑like、基于 Winograd 的以及更新的张量分解方法），其算术复杂度低于经典的 O(n³) GEMM。
Code Generation – 部署引擎解析目标硬件的指令集、内存层次结构和并行执行模型，然后生成优化的内核（CUDA、HIP、AVX‑512、NEON 等）。
Group‑Parallel Optimizations – 将内核组织成“组”，在片上缓冲区（共享内存、L1 缓存或寄存器）中共享中间结果。这降低了相同子乘积从 DRAM 中读取的次数。
Analytical Performance Model – 决策模块通过考虑计算吞吐量、内存带宽以及 LCMA 递归深度带来的额外开销来估算运行时间。随后选择能够最小化预测时间的算法（及其分块参数）。
Run‑time Dispatch – 在推理或训练启动时，FalconGEMM 使用实际矩阵尺寸和硬件统计信息查询模型，选择最佳内核并在无需开发者干预的情况下启动。

平台	数据类型	相较于 cuBLAS / MKL 的加速	相较于 AlphaTensor 的加速
NVIDIA H100	FP16	+15.2 %	+32.8 %
NVIDIA A100	BF16	+12.7 %	+28.4 %
ARM Neoverse	FP32	+9.3 %	+18.5 %
Intel Xeon (AVX‑512)	FP64	+7.6 %	+12.4 %

突破峰值性能：在多个 LLM 层（例如 transformer 注意力和前馈块）中，FalconGEMM 通过算法上减少算术工作量，超越了底层硬件的理论峰值 FLOPs。
跨形状的鲁棒性：决策模块能够在经典 GEMM（针对小的方阵）和 LCMA（针对 token‑wise 操作中常见的高瘦或宽短矩阵）之间正确切换。
低开销：分析模型仅增加 < 0.5 % 的运行时开销，使该框架适用于批量训练和对延迟敏感的推理。

LLM 训练流水线 – 更快的矩阵乘法直接转化为 GPU 使用时长的减少，从而降低大规模模型预训练的云成本。
边缘推理 – 在基于 ARM 的服务器甚至移动 SoC 上，FalconGEMM 能在不牺牲电池寿命的前提下提供更高的设备端语言模型吞吐量。
框架集成 – 由于部署模块生成标准的 CUDA/HIP/AVX 内核，现有深度学习库（PyTorch、TensorFlow、JAX）可以通过薄包装器直接用 FalconGEMM 内核替换其 GEMM 调用，实现即插即用。
硬件无关优化 – 拥有异构设备（GPU + CPU）舰队的公司不再需要分别手动调优内核；FalconGEMM 能自动适配，简化模型部署的 CI/CD 流程。

数值稳定性 – 某些 LCMAs（例如 Strassen‑type）会引入额外的舍入误差；论文指出在 FP64 下精度略有下降，这在某些科学工作负载中可能不可接受。
内存占用 – 递归算法需要额外的临时缓冲区；在内存受限的设备上，这可能限制可处理的矩阵规模。
模型泛化 – 分析性能模型在一组代表性的 GPU/CPU 上进行校准；将其扩展到新兴加速器（TPU、定制 ASIC）需要额外的性能分析。
未来方向 – 作者计划加入混合精度自动调优，基于运行时误差度量探索自适应递归深度，并开源该框架以促进社区驱动的内核扩展。