•2025年10月
ai-systems

超越 MatMul 霸权：AI 计算的“众篮之策”与张量结构优化

深入探讨 AI 领域对矩阵乘法（MatMul）的过度依赖问题。文章从“不要把所有鸡蛋放在一个篮子里”的视角，揭示将所有张量运算“扁平化”为 MatMul 的性能陷阱，并提出面向未来的结构化张量优化方案。

在人工智能（AI）的计算世界中，矩阵乘法（MatMul, or GEMM）无疑是王者。从卷积神经网络（CNN）到庞大的 Transformer 模型，几乎所有核心计算都可以最终归结为大规模的 MatMul 操作。得益于 GPU、TPU 等专用硬件的飞速发展以及 cuBLAS、oneMKL 等高度优化的计算库，我们已经能以惊人的速度执行矩阵乘法。这种强大的能力使得一种设计范式深入人心：将所有形式的张量（Tensor）运算都“扁平化”（Flattening）或“展开”（Unfolding）为二维矩阵，然后调用最优的 MatMul 内核解决问题。

这种策略看似一劳永逸，宛如一个无所不包的“超级篮子”。然而，正如谚语所警示的，“不要把所有的鸡蛋都放在一个篮子里”。当我们盲目地将所有张量计算都塞进 MatMul 这一个篮子时，我们可能正在为看似极致的峰值性能付出沉重的、隐藏的代价。

“扁平化”的代价：被破坏的数据结构与局部性

张量是多维数据结构，其维度天然蕴含着丰富的结构信息和数据局部性。例如，在图像处理中，一个像素及其邻域在空间上是紧密相关的；在时间序列分析中，相邻的数据点在时间上具有连续性。然而，为了将这些高维张量适配于二维的 MatMul 运算，我们不得不执行“扁平化”操作。

以 CNN 中的卷积操作为例，一个 3x3 的卷积核在特征图上滑动，其计算本质上是高度局部的，每次只涉及一小片输入区域。但为了套用 MatMul，我们通常使用 im2col (image-to-column) 这样的技术，将输入特征图的局部区域块（patches）提取出来，重新排列成一个巨大的矩阵。这个过程不仅引入了显著的数据冗余（一个像素可能被复制到多个列中），更严重的是，它彻底破坏了数据原始的空间局部性。

当数据结构被破坏后，硬件的缓存系统便难以发挥其最大效用。原本在内存中连续存放、可以被高效预取和缓存的邻近数据，在 im2col 之后变得分散。这导致计算单元在执行 MatMul 时，不得不频繁地从速度较慢的主存（DRAM）中抓取数据，造成所谓的“内存墙”瓶颈。尽管计算单元本身可能非常强大，但其大部分时间都消耗在了等待数据的空转上。我们为了追求计算的统一性，牺牲了数据访问的效率，这正是将所有“鸡蛋”放入 MatMul 单一篮子的首要风险。

另辟蹊径：拥抱张量原生结构

“众篮之策”的核心思想是，我们应该拥有一个更丰富的计算工具箱，针对不同特性的张量操作采用最适合的优化策略，而不是一刀切地诉诸 MatMul。这意味着我们需要重新拥抱并利用张量的原生结构，在不破坏数据局部性的前提下进行计算。

这种方法被称为面向结构的张量优化，其关键在于“分块张量”（Tiled Tensors）或“分块计算”。其理念是将大规模的张量运算分解为一系列在小型、连续的张量块（Tensor Tiles）上进行的操作。这些张量块足够小，可以完全载入到 CPU 或 GPU 的高速缓存（如 L1/L2 Cache）中。

一旦数据进入缓存，计算单元就可以在极高带宽、极低延迟的环境下对其进行反复利用，从而实现计算强度与内存访问的平衡。这不仅适用于卷积，对于自注意力机制、池化、归一化等众多操作同样有效。

实践参数与考量清单

在评估是否应超越 MatMul 范式时，可以考虑以下几个具体的可落地要点：

  1. 缓存利用率与数据复用：监控硬件性能计数器中的缓存命中率（Cache Hit Rate）和内存带宽占用。如果一个计算任务的算术强度（Arithmetic Intensity）不高，但内存带宽已达瓶颈，这便是一个明确的信号：当前的计算方式未能有效复用数据，可能是“扁平化”导致的恶果。此时，转向分块张量计算，通过循环分块（Loop Tiling）等技术提高数据复用，是首选的优化方向。

  2. 操作的内在稀疏性与局部性：并非所有运算都像稠密矩阵乘法那样全局相关。例如，图神经网络（GNN）中的消息传递、部分 CNN 中的深度可分离卷积（Depthwise Separable Convolution）等，其计算模式具有高度的局部性或稀疏性。强行将其转换为稠密的 MatMul 会引入大量无效计算（例如与零相乘）。为这些操作设计专门的、感知稀疏/局部模式的计算内核，远比使用通用 MatMul 更高效。

  3. 动态形状与填充开销：在处理自然语言等可变长度输入时，模型经常需要处理动态变化的张量形状。基于 MatMul 的系统为了批量处理，通常需要将短序列填充（Padding）到统一长度，这会造成相当大的计算资源浪费。而一个更加灵活、支持异构或动态分块的张量计算引擎，则可以更优雅地处理此类问题，减少不必要的开销。

迈向统一与多元的未来

反思对 MatMul 的过度依赖，并不意味着要全盘否定它。MatMul 依然是 AI 计算中最高效的工具之一。关键在于，我们的 AI 软件栈和编译器需要具备识别并应用多种计算策略的能力。

这正是像 Chris Lattner 创办的 Modular 等前沿公司正在努力的方向。他们致力于构建一个统一的 AI 引擎，它不再将开发者锁定在特定硬件厂商的、碎片化的库（如 cuBLAS, MKL）之上，而是提供一个更具可组合性、可移植性的平台。在这个平台上，系统可以根据操作特性、硬件架构和数据布局，智能地选择最优的执行策略——有时是调用高度优化的 MatMul，有时则是动态生成一个感知数据局部性的分块张量内核。

总而言之，AI 系统性能优化的下一站，将是从“一招鲜吃遍天”的 MatMul 霸权，走向一个更加多元、更加智能、更加尊重数据内在结构的“众篮时代”。只有为不同类型的“鸡蛋”精心准备好各自合适的“篮子”，我们才能真正突破内存墙的束缚，释放硬件的全部潜力。