量化存储墙（三）：GEMM EMA 下限解析解以及硬件静态资源分配设计

news/2025/11/16 18:42:53/文章来源:https://www.cnblogs.com/devil-sx/p/19228698

Roofline 缺失的一角：EMA

计算-存储工艺各自单独演化发展步调不一， Memory Wall 在 AI 计算时代越来越显著^[1]。Roofline Model 将计算/存储，软件/硬件用一个简洁优雅的统一模型概括，然而带入具体数值求解时，硬件的极限算力、带宽物理资源是个静态常量，冯诺依曼体系往往不涉及 recomputation^[2]，模型 FLOPs 也是固定的，而片外访存量（External Memory Access, EMA）并不那么直观，往往需要通过仿真方法获取。这是由于片外访存求解涉及硬件片内存储设计、软件 intra/inter-kernel 的复杂数据复用关系和调度策略，是一个极其复杂的 NP 问题。

Orojenesis ：静态资源 Matters

不过换句话说，如果能够简化硬件/软件建模存在求解的可能。比如 decode 阶段由于自回归 token-per-token 丧失了 token 并行度，降低数据复用率，在单用户情形下 EMA 等同数据量。

而在 Prefill 阶段，2024 ISCA 来自 NVIDIA 的 Orojenesis ^[3] DSE 工作通过简化的 "Snowcat" 模型，即对片上仅用一个片上容量参数建模，并将 workload 限制在单个 GEMM kernel（以及简单 fusion，但不是本文重点不展开了），通过遍历式搜索调度 “SRAM 容量-调度策略” 空间，得到一个函数曲线。当片内容量逼近无限大的时候，所有数据都可以放在 SRAM 之中，下限即是“输入-权重-输出” 变量大小之和；而当片内容量逼近 0 的时候，等同于计算单元直接连接片外访存的访存量，其数值近似为 FLOPs 除以一个表示计算单元空间复用度的常数。

这种简化仅仅考虑静态资源片上 SRAM 容量，代表一种 EMA 下限，而实际和理论之差则是靠架构改进。Orojenesis 团队也在多张 NV 实体卡上 CUTCLASS 环境测试对比，可以看到在 SRAM 容量较大时测试结果和仿真值数量级和趋势接近，目测误差在 50-100% 之间。

Orojenesis 的函数曲线揭示了化简硬件和 workload 对 EMA 建模理论化的可能性，本文尝试给出该关系的解析解。

单 GEMM-kernel 优化方程和求解

定义 GEMM 的三个维度为 T (Token)、O (Output Channel)、 I (Input Channel)。定义片内映射维度为 [Ts,Os,Is]，对于三种不同的片外循环次顺序 [Tt, Ot, It] （Output Stationary）、 [Ot, It, Tt] （Weight Stationary）、[Tt, It, Ot] （Input Stationary）。此部分建模可参考 timeloop。

比如片外 Output Stationary 的 nested-loop 表示：

---- off-chip loop ----
for tt in Ttfor ot in Otfor it in It---- inner-chip loop ----for ts in Tsfor os in Osfor is in Isy[tt*Ts+ts][ot*Os+os] += w[ot][it][os][is] * x[tt*Ts+ts][it*Is+is]

定义数据变量精度 bit 数为 \(p\) 以及 SRAM 容量为 \(M\)，各个变量的 EMA 开销以及 buffer 占用量如下：

	Size	EMA (OS)	EMA (WS)	EMA (IS)	Buffer
Input Feature	Si=TxIxp	OtxSi	OtxSi	Si	TsxIs
Weight	Sw=OxIxp	TtxSw	Sw	TtxSw	OsxIs
Output Feature	So=TxOxp	So	ItxSo	ItxSo	TsxOs

以 Output Stationary 为例，其优化问题形式为：

\[\begin{align*} min& \quad EMA =O_{t}\times S_{i}+ T_{t}\times S_{w}+ S_{o}\\ s.t.&\quad (T_{s}\times I_{s} + O_{s}\times I_{s} + T_{s}\times O_{s})\times p\le M \\ s.t.&\quad T_{s}\times T_{t}=T\\ s.t.&\quad O_{s}\times O_{t}=O\\ s.t. &\quad I_{s}\times I_{t}=I \end{align*} \]

其解为:

\[\begin{align*}\\ O_t &= \sqrt{\frac{O T p S_w}{M S_i}} \\ O_s &= \sqrt{\frac{M O S_i}{T p S_w}} \\\\ T_t &= \sqrt{\frac{O T p S_i}{M S_w}} \\ T_s &= \sqrt{\frac{M T S_w}{O p S_i}} \\ I_s &\to 0^+ \\ I_t &\to \infty \\ \text{EMA}_{\text{min}} &= S_o + 2 \sqrt{\frac{O T p S_i S_w}{M}}=S_{o} + 2 \sqrt{ \dfrac{ S_{i} S_{w} S_{o} }{ M } } \end{align*} \]

因为 Output Stationary，此时 Input Channel 分配在片外不会影响整体 EMA，因此最优解将所有的 Input Channel 维度分配在片外，塞下更多 Token 和 Output Channel 在片内，以便减少 Input Feature 和 Weight 的 EMA 访存。这个解和实际的误差来源三方面：

仅仅考虑 EMA，片内没有任何 Input Channel 维度，会增大 Output Feature 片内 buffer 访问次数，实际矩阵计算单元设计一般回考虑 Input Channel 维度；
内积形式实践中为了精度-访存量的 trade-off，累加和往往在高精度内实现计算，而转为低精度访存交互，Input Channel 片内维度分配越少，对精度损坏越大；
变量数值都是整数。

但不妨先按这一假设讨论，定义 \(C = 2 \sqrt{ \dfrac{ S_{i} S_{w} S_{o} }{ M } }\) 。GEMM 维度高度对称，以此类推三种 Stationary 下的 EMA 理论最优结果：

\[\begin{align*} &\min EMA_{OS} = S_{o} + C\\ &\min EMA_{WS} = S_{w} + C \\ &\min EMA_{IS} = S_{i} + C \\ \end{align*} \]

不同 stationary 的理论最优解仅与相应变量数据大小相关，该结论对于片内片外访存都适用。具体来说，哪个变量大用哪种 stationary ，这验证了之前设计加速器的 insight，比如 dense vison prediction task 的浅层网络是典型的 activation-heavy 性质，应当用 OS 或 IS，考虑前文提到的精度 trade-off ，仅有 OS 唯一可行选项。

依此总结为 stationary 无关的通用解析式，仅当片上存储约束取等号时有效：

\[EMA = \min{(S_{i},S_{w},S_{o})}+ 2 \sqrt{ \dfrac{ S_{i} S_{w} S_{o} }{ M } }=\min(S_{i},S_{w},S_{o})+ FLOPs \sqrt{\frac{p^3}{M}} \]

至此得出 EMA 的解析解，我们带入具体数值和 Orojenesis 结果对比。当片上容量大于 1 KB 时，二者相对误差大致在 10% 一下，基本符合，验证结论正确性。

近似推导结论

EMA 解析解由两项组成，一项 \(\min(S_{i},S_{w},S_{o})\) 是 Stationary 下对应的数据大小，另一项 \(FLOPs \sqrt{\frac{p^3}{M}}\) 则是其余变量贡献正比于 FLOPs。T/I/O 三个维度，数据变量是两个维度乘积，而 FLOPs 是三个维度乘积，主要 EMA 实际由第二项贡献，当 \(FLOPs \sqrt{\frac{p^3}{M}} >> \min(S_{i},S_{w},S_{o})\) 时，有：

\[EMA \approx FLOPs \sqrt{\frac{p^3}{M}} \]

片内容量-EMA 经验公式

该公式揭示 EMA 和片上容量 M 满足 \(EMA\propto FLOPs\times M^{-0.5}\) 关系。\(FLOPs\) 含义代表 memory 直连计算单元时的访存上限，改写公式为 \(EMA\propto EMA_{max}\times M^{-0.5}\)。

本文单 GEMM-kernel workload 代表 intra-kernel 复用度分析，至于 inter-kernel 本篇系列文章第二篇^[4]曾推导过简化 Layer-Fusion 容量和 EMA 的关系满足 \(EMA \propto EMA_{max} \times M^{-1}\)，而 Flash Attention ^[5]的理论上限也满足 \(EMA \propto EMA_{max} \times M^{-1}\) 关系。大致总结一种经验公式：

\[\begin{align*} EMA &\propto EMA_{\text{max}} \times M^{-\alpha} \\ \text{where, } & \\ \alpha &= \begin{cases} 0.5 & \text{if Intra-Kernel} \\ 1 & \text{if Inter-Kernel} \end{cases} \end{align*} \]

芯片静态资源分配和精度 Scaling

\[\begin{align*} EMA &\approx FLOPs \sqrt{\frac{p^3}{M}}\\ \text{Software Intensity} &= \frac{FLOPs}{EMA} \approx \sqrt{\frac{M}{p^3}} \\ \underbrace{\frac{FLOPs}{\sqrt{M/p}}}_{\text{on chip}}&\approx\underbrace{\frac{EMA}{p}}_{\text{off chip}} \end{align*} \]

此公式左侧代表片内计算和 SRAM 存储资源，右侧代表片外带宽资源。当静态资源分配符合此公式时，片内和片外取得平衡，workload 在 roofline 上取转折点位置，计算和带宽资源利用率最优。假设某个设计符合此平衡，且固定片外带宽供应商，精度 bit 每变为原来的 1/2，\(\frac{FLOPs}{\sqrt{M}}\) 应变为原来的 \(\sqrt{2}\approx 1.4\) 倍。

https://www.cnblogs.com/devil-sx/p/18563196 ↩︎
https://www.zhihu.com/question/1954115162412942829/answer/1954961870525030613 ↩︎
https://ieeexplore.ieee.org/stamp/stamp.jsp?tp=&arnumber=10609642 ↩︎
https://www.cnblogs.com/devil-sx/p/18759528#fn7 ↩︎
https://arxiv.org/pdf/2307.08691 ↩︎