Elk: Exploring the Efficiency of Inter-core Connected AI Chips with Deep Learning Compiler Techniques

• 作者/机构: Yiqi Liu (UIUC), Yuqi Xue (UIUC), Noelle Crawford (UIUC), Jilong Xue (Microsoft Research), Jian Huang (UIUC)

A1 主要贡献

为了满足深度学习（DL）模型（如大型语言模型）日益增长的计算需求，各种AI芯片被开发出来。典型的AI芯片拥有许多并行核心，每个核心都有本地SRAM。为利用这种并行性，DL编译器将张量算子划分为片（tile），并将每片映射到一个核心。然而，由于片上SRAM大小有限，AI芯片通常采用HBM等片外存储器来支持更大的模型。

核心问题: 片外存储器带宽的增长速度远慢于计算性能，这成为了一个瓶颈。为了缓解这一问题，核间互连AI芯片（ICCA）应运而生。ICCA芯片通过核间链路使核心可以直接访问其他核心的SRAM（如图1），从而聚合了大量的片上存储空间和高带宽。然而，在ICCA芯片上高效运行DL模型面临着三大性能因素之间的根本性冲突：计算（每核执行）、通信（核间数据交换）和I/O（片外数据访问）。如图2所示，这三者存在资源竞争：
1. 片上存储容量竞争：为计算分配更多空间可以提高计算效率，但减少了为预加载数据准备的缓冲空间，从而影响I/O效率；反之亦然。
2. 互连带宽竞争：片上互连带宽被核间数据交换和从HBM到核心的数据加载共享。
3. 存储器访问竞争：每个核心的本地SRAM需要同时服务于本地计算和来自其他核心的远程访问请求。

研究目标: 本文提出Elk，一个深度学习编译器框架，旨在通过联合优化上述三个性能因素来最大化ICCA芯片的效率。Elk将这些性能因素形式化为可配置的编译器参数，并构建了一个全局权衡空间。

创新与贡献:
* 首次识别性能挑战：本文首次系统性地识别并分析了带有片外HBM的核间互连AI芯片在高效利用其硬件特性时面临的性能挑战。
* 开发Elk编译器框架：Elk将性能因素（计算、通信、I/O）结构化为编译器中的可配置参数，使得可以通过编译器技术来探索和优化硬件性能。
* 新的调度与分配算法：
* 提出了一种新的归纳式算子调度策略，用于优化HBM数据加载和片上执行的重叠。
* 设计了一种新的成本感知片上存储器分配算法，以在执行空间和预加载空间之间进行权衡。

• 通用硬件映射接口：为推广Elk的设计，构建了一个通用接口，可将优化的端到端执行计划映射到主流的ICCA芯片架构上。

• 全面的评估框架：

  • 基于真实的IPU-POD4硬件构建了一个全功能的仿真框架，用于评估带有HBM的ICCA芯片。
  • 构建了首个支持流行网络拓扑和多种带宽行为的ICCA芯片硬件模拟器。

• 设计空间探索：利用Elk和ICCA芯片模拟器，对ICCA芯片的设计空间进行了探索，并提出了相关见解。

A3 背景知识与关键观察

2.1 ICCA芯片架构

以Graphcore IPU MK2为例。为了便于介绍图1所示的ICCA芯片，我们以Graphcore IPU MK2【索引29，Graphcore Colossus Mk2 IPU，2021，HCS’21】为例。一个IPU芯片拥有1472个核心，这些核心并行独立执行。每个核心配备624KB的本地SRAM，总计片上存储容量达到896MB。所有核心通过高带宽、低延迟的链路互连。每个核心能够以5.5GB/s的速度访问任何其他核心的本地存储器，从而提供了约8TB/s的聚合全连接（all-to-all）带宽【索引27，Dissecting the Graphcore IPU Architecture via Microbenchmarking，2020，arXiv】。巨大的片上存储空间通过存储更多算子甚至整个模型来提高片上数据复用。全连接的互连结构使得每个核心都能以高带宽独立访问片上数据。如果多个核心同时从同一个核心接收/发送不同数据，互连会以全带宽顺序服务每一次数据传输。除了IPU，其他ICCA芯片如SambaNova SN40L【索引46，SambaNova SN40L: Scaling the AI Memory Wall with Dataflow and Composition of Experts，2024，arXiv】和Tenstorrent【索引52，Meet Grayskull，2023，Tenstorrent】则采用了基于Mesh的片上互连。总体而言，ICCA芯片架构能够实现可扩展的性能，并缓解像LLM这类内存密集型DL工作负载的内存带宽瓶颈。

通过HBM扩展ICCA芯片。为了支持尺寸超过片上容量的大型模型，我们可以使用高带宽存储器（HBM）【索引39，HBM3E，2024，Micron】等片外存储模块来扩展存储容量。许多ICCA芯片已经集成了片外存储器【索引37, 46】。如图1所示，它们将HBM控制器连接到片上互连，因此每个控制器可以直接向每个核心发送数据，类似于核心之间的数据交换。核心通过互连与HBM控制器通信来访问HBM。HBM控制器合并来自核心的内存请求，从HBM加载数据，并将其发送给核心。

2.2 带有HBM的ICCA芯片的执行模型

数据加载与执行流程。在执行一个DL模型之前，所有必需的数据（如模型权重）都被加载到HBM中。ICCA芯片会顺序执行模型中的每个算子，首先将其所需数据从HBM预加载到片上存储器，然后执行片上计算。

通过双缓冲重叠计算与I/O。为了最大化计算吞吐量，编译器将片上SRAM作为双缓冲来管理，以重叠片上执行和片外HBM访问。编译器将每个核心的SRAM划分为一个执行空间（用于存储当前正在执行的算子）和一个预加载空间（用于存储从HBM预加载的算子）。当一个算子正在执行时，ICCA芯片可以同时将其他算子从HBM预加载到片上存储器。预加载时，HBM控制器利用互连将数据分发到各核心。每个核心需要为此预留足够的本地存储空间。当预加载空间满时，预加载将停止，导致HBM带宽未被充分利用。

片上执行模型。在ICCA芯片上有多种并行执行模型可以执行张量算子【索引34，Scaling Deep Learning Computation over the Inter-Core Connected Intelligence Processor with T10，2024，SOSP’24；索引40，Efficient Large-Scale Language Model Training on GPU Clusters Using MegatronLM，2021，SC’21；索引74，ROLLER: Fast and Efficient Tensor Compilation for Deep Learning，2022，OSDI’22】，所有这些模型都需要大量的计算、片上存储和通信资源。在这些执行模型中，编译器会将一个张量算子的计算划分为小块（tile）【索引34, 74】，并将每个tile映射到一个核心。为了执行一个tile，每个核心必须通过互连从HBM或其他核心获取所需数据到其本地存储器。例如，图3(a)中一个MatMul算子被划分为四个tile，所有核心都需要“Input 2”来进行各自的计算。在一些执行模型中【索引40】，这个共享张量会由HBM控制器通过互连直接广播到每个核心，如图3(b)所示。这需要更多的本地存储空间，但减少了核间访问。另一些执行模型【索引34, 74】允许一个核心在执行期间从其他核心访问共享张量，如图3(c)所示。这需要较少的每核本地空间，但增加了核间访问。在每核执行之后，一个算子可能需要将多个核心的部分结果规约（reduce）为最终结果，这些核心将通过互连交换部分结果。

2.3 使用带有HBM的ICCA芯片的挑战

资源竞争导致性能权衡。为了最大化ICCA芯片的性能，我们必须（1）实现更快的每核执行，（2）利用更多的HBM带宽及时预加载所需数据，以及（3）减少核间数据共享开销。然而，由于这些性能指标存在冲突的资源需求，很难同时将它们最大化。

片上存储空间竞争。由于片上存储空间的竞争，我们无法同时最大化每核执行性能和HBM带宽利用率。如图2中①所示，每个核心为当前执行的算子预留一个执行空间，为预加载的算子预留一个预加载空间。为了加速每核执行，需要一个更大的执行空间（见§3.1和图5）。为了防止HBM利用率不足，需要一个更大的预加载空间（见§3.2和图6）。在有限的片上存储器中，我们无法同时扩展这两个空间。

互连带宽竞争。由于互连带宽的竞争，我们无法同时最大化HBM带宽利用率和最小化核间数据共享开销。如图2中②所示，片上互连同时承载用于核间数据共享的核心到核心流量和用于预加载的HBM控制器到核心的流量。当两种流量都很重时，互连会发生拥塞（见§3.3和图8）。

存储器访问竞争。由于存储器访问的竞争，我们无法同时最大化每核执行性能和最小化核间数据共享开销。如图2中③所示，每个核心的本地存储器被其自身用于计算一个tile，同时也被其他核心用于核间数据共享。例如在IPU上，每个核心在执行像MatMul这样的DL算子时会以全速读取其本地存储器（128位/周期【索引19，Tile Vertex ISA，2022，Graphcore】），任何其他访问都会暂停执行。发生竞争时，该核心上的tile执行会以较慢的速度从本地存储器读取数据，甚至完全暂停。远程核心也可能因SRAM带宽下降而受到影响。

3. Elk中的性能权衡

将性能因素映射到编译器决策。如图4所示，Elk将各个性能因素映射到编译器决策上。首先，增加每核执行空间可以通过更大的tile尺寸实现更快的每核执行（§3.1）。其次，增加预加载的算子数量可以更好地重叠片上执行和片外HBM加载，从而提高HBM带宽利用率（§3.2）。第三，增加每个预加载算子的预加载空间可以减少核间数据交换开销和存储器访问竞争，因为共享数据可以提前在核心上复制，以减少对其他核心的按需访问开销（§3.3）。

3.1 更大的执行空间可实现更快的每核执行

执行空间与性能的关系。划分一个算子的方式有很多种【索引70, 71, 74】，会产生具有不同每核tile大小、执行时间和核间数据交换流量的划分方案。通常，每核更大的执行空间可以支持更大的tile大小，从而提高每核执行性能，因为更大的tile意味着更高的每核数据复用和更大的计算粒度，以及更少的核间数据访问。图5展示了执行时间与执行空间之间的相关性。我们选择了来自不同规模的流行LLM（Llama-2-13B【索引53】，Gemma-2-27B【索引51】，OPT-30B【索引66】）的代表性算子。对于每个算子，我们绘制了由一个先进的ICCA芯片编译器【索引34】在不同SRAM大小约束下生成的划分方案的执行时间。结果显示，执行速度更快的方案需要更多的每核执行空间。

现有编译器的局限性。现有编译器专注于在给定的执行空间大小下实现更高的片上执行性能。然而，它们无法通过协调§2.3中提到的内存空间竞争来找到一个合适的执行空间大小。此外，同一模型中的算子具有不同的内存与时间相关性。因此，我们需要根据每个算子的性能特性来调整执行空间大小，而不是在整个模型执行过程中分配一个固定大小的执行空间。

3.2 预加载更多算子可提高HBM带宽利用率

算子计算强度的多样性。DL模型中的算子具有不同的计算强度（即每字节执行的浮点运算次数，或FLOPs）。一些算子由于更多的片上数据复用而计算密集（例如，使用模型参数的算子），而另一些则是内存密集型（例如，KV缓存【索引45】，在一个批次内的请求之间没有数据复用）。

多样性导致利用率不佳。不同算子之间多样的HBM访问和执行时间导致了计算和HBM带宽利用率的次优。如果当前执行的算子执行时间短，而下一个算子的HBM加载时间长，那么当前算子会在下一个算子完成预加载之前结束，导致计算停顿。类似地，如果下一个算子在当前算子完成之前就预加载完毕，HBM带宽就会被闲置。

预加载更多算子以平滑带宽需求。为了提高HBM带宽利用率，我们可以预加载更多的算子。这也提高了计算利用率，因为片上会有更多数据准备就绪，未来的执行不太可能停顿。然而，预加载更多算子需要更大的预加载空间。图6显示了在LLM推理中，不同每核预加载空间大小下HBM带宽需求随时间的变化。带宽需求量化为防止片上执行停顿所需的最小HBM带宽。当预加载空间较小时，由于预加载机会不足，带宽需求剧烈波动。当预加载空间较大时，可以预加载更多算子，这平滑了带宽需求，减少了计算/内存的空闲时间，并提升了整体性能。

3.3 更大的单算子预加载空间可减少核间数据访问量

预加载广播与执行时访问的权衡。如§2.2所述，核心间共享的数据既可以在预加载期间由HBM控制器广播，也可以在执行期间从对等核心访问。更大的预加载空间允许在预加载时进行更多的广播，从而在执行时减少按需访问。此外，随着核间访问的减少，每个核心上的内存访问竞争也会减少。

减少核间带宽需求。图7显示，扩大预加载空间可以减少核间带宽需求。对于每个算子，我们选择在给定执行空间大小内容纳的最快执行方案。MinPreload让每个核心在执行时从其他核心访问所有共享数据，这需要最小的预加载空间。MaxPreload让HBM控制器在预加载时广播尽可能多的共享数据，这需要最大的预加载空间。我们分析了每个核心的核间带宽需求（核间传输量 / 每核执行时间）。

平滑总互连流量。尽管在预加载时进行更多广播会增加HBM控制器到核心的流量，但预加载流量可以机会性地与正在进行的核间流量交错，以减少竞争。图8显示了每个核心的总互连带宽需求（定义为 (核间传输量 + HBM到核心传输量) / 每核执行时间 + HBM加载时间）随时间的变化。纯粹依赖核间传输会导致流量压力剧烈波动，造成互连利用不足或拥塞。在预加载时进行更多广播，通过将流量分散到预加载和执行时间内，减少了波动。

A2 方法细节

我们设计了Elk，一个用于探索ICCA芯片效率的编译器框架。Elk通过配置预加载的算子数量、每核执行空间大小、每个算子的预加载空间大小以及算子的预加载顺序来自动权衡性能因素。Elk的设计概览如图9所示。我们使用表1来展示Elk的哪个设计组件处理§3中的每个性能权衡。

4.1 设计概览

两级搜索空间。对于一个DL模型，Elk通过探索一个两级搜索空间来调度算子的预加载和执行。首先，对于每个算子，Elk探索在其执行之前或期间预加载未来算子的所有可能数量（§4.2）。其次，对于每个预加载算子的数量，Elk通过在执行空间和预加载空间之间进行权衡来优化片上内存分配（§4.3）。

预加载顺序优化。为了减少核间数据交换开销并支持更大的执行空间，Elk允许以不同的顺序预加载算子。Elk通过搜索所有有希望的顺序来找到最优的预加载顺序。对于每个顺序，Elk应用算子调度策略并进行性能评估。为了减少搜索开销，Elk会剪枝那些会导致片上内存溢出的顺序（§4.4）。最后，Elk为整个模型生成一个优化的端到端计划。该计划指定了每个算子的预加载和执行计划。然后，代码生成器将此计划转换为硬件的可执行程序（§4.5）。

4.2 两级归纳式算子调度

调度算法目标。该调度算法通过决定在每个算子的片上执行之前或期间预加载未来算子的数量（即预加载数量），来最小化DL模型的端到端执行时间。每个预加载数量代表了片上执行速度和HBM带宽利用率之间的一个权衡点（图4）。

搜索空间与算法。优化空间对于算子数量是指数级的。假设一个DL模型中有$N$个算子，片上内存最多能容纳$K$个算子。每个算子的执行可以与1到$K$个算子的预加载重叠。因此，所有算子的预加载数量有$K^N$种组合。我们开发了一个时间复杂度为$O(NK)$的算法来解决这个问题。其核心思想是，由于DL模型中的算子通常因数据依赖而顺序执行，我们可以利用执行顺序，为每个算子归纳地推导出最优的预加载数量，而不是探索所有组合。

归纳调度过程。我们可以从第一个算子开始，为每个后续算子找到最优的预加载数量，也可以从最后一个算子开始，为每个前面的算子进行调度。对于每个算子，我们基于已经调度好的算子，探索所有可能的预加载数量，并选择能最小化“从开始到当前”或“从当前到结束”执行时间的预加载数量。由于两种归纳方向是等价的，我们关注第二种。归纳的基例是平凡的，因为最后一个算子没有后续算子可以预加载（即预加载数量总是0）。对于归纳步骤，如图10所示，在调度Op5时，Elk会枚举所有可能的预加载数量。对于每个数量，Elk调用成本感知的片上内存分配算法（§4.3）来确定相关算子的执行/预加载空间大小和Op5的估计执行时间。例如，预加载数量0意味着不重叠Op5的执行与任何预加载，Op5的执行时间最短，但从Op5到模型结束的总执行时间可能不是最优的。而预加载数量1和2虽然延长了Op5的执行时间，但总执行时间可能更优。Elk选择能产生最短“从当前到结束”时间的预加载数量。

调度与时间估计。调度完Op5的执行后，Elk会安排其预加载发生在其执行之前或Op6的预加载之前（取较早者），以保证数据依赖。前一个算子（Op4）的调度将依赖于Op5的预加载时间，该时间被估计为（1）来自roofline模型【索引60】的HBM访问时间和（2) 来自§4.3中成本模型的互连传输时间的最大值。我们的算法复杂度为$O(NK)$，因为它遍历$N$个算子，每个算子最多有$K$个预加载数量。

4.3 成本感知的片上内存分配

两级权衡空间。在§4.2中，调度一个算子时，Elk会为每个预加载数量优化性能。给定当前正在执行的算子和一组待预加载的算子，Elk定义了一个执行/通信时间与内存消耗之间的两级权衡空间。首先，存在两种算子内权衡：（1）对于当前执行的算子，Elk用内存空间换取执行时间（§3.1）。（2）对于每个预加载的算子，Elk用内存空间换取该算子的核间数据交换开销（§3.3）。其次，存在一个算子间权衡：由于不同算子有不同的内存-时间权衡，我们将更多内存分配给那些从更大执行/预加载空间中获益更多的算子。

探索权衡空间。Elk分两个阶段探索这个两级空间。首先，对于每个算子，Elk找到所有在时间和内存之间的帕累托最优权衡方案。其次，Elk根据帕累托最优方案和总片上内存容量，联合决定所有算子的执行/预加载空间大小。

算子内权衡（执行）。对于当前执行的算子，Elk集成了现有的编译器技术【索引7, 16, 34, 46, 70, 74】来枚举其所有划分方案。对于每个方案，Elk使用成本模型估计其执行时间，并使用tile大小估计其执行空间。Elk检查所有方案以找到帕累托最优曲线上的方案。

成本模型。Elk使用一个精确的成本模型来快速估计每核执行和核间传输的性能。对于每种算子类型（如MatMul），我们随机生成不同形状的tile，并在目标设备上用一个核心运行每个tile。然后，我们使用tile形状作为输入，分析得到的执行时间作为输出，来拟合一个线性树模型【索引10】。对于核间传输，我们为每个网络链接拟合一个模型，使用传输量作为输入，传输时间作为输出。图12显示Elk可以准确预测IPU芯片的执行和传输时间。

算子内权衡（预加载）。对于每个预加载的算子，其执行态（execute-state）方案已在归纳调度的前一步骤中确定。由于该算子当前未执行，Elk为其分配一个节省内存的预加载态（preload-state）方案。当开始执行时，一个数据分发阶段通过互连分发所需数据，将算子从预加载态转换为执行态。这节省了该算子的预加载空间，但代价是额外的核间数据交换开销。Elk通过估计它们的预加载空间大小和数据分发时间，找到每个预加载算子的帕累托最优预加载态方案。

算子间权衡。在有限的片上内存下，Elk在执行和预加载的算子之间联合权衡内存分配。它最小化总时间，该时间由（1）执行时间、（2）数据分发时间、（3）预加载和执行重叠导致的互连竞争开销，以及（4）本地SRAM访问和核间访问之间的内存访问竞争开销决定。由于枚举所有可能的方案组合不切实际，Elk使用一种基于每个算子内存-成本效率的启发式算法。Elk从每个算子的最快方案开始。如果总空间需求超过内存容量，Elk会迭代地搜索能容纳在片上内存中的最佳方案组合。在每个搜索步骤中，Elk为每个算子检查帕累托最优曲线上的下一个内存占用更小的方案，并选择“成本效益”最高的算子（即其下一个方案具有最大的 Δ = 减少的空间大小 / 增加的时间开销 比率）进行更新，如图11所示。当总内存需求不超过可用容量时，Elk停止搜索。

4.4 预加载顺序置换

重排预加载顺序的好处。Elk允许算子以不同于执行顺序的顺序进行预加载。这有两个好处。首先，重排有助于缓解互连拥塞。通过机会性地重新安排重度预加载流量，以避开互连的“高峰时段”。其次，通过将一些大型算子的预加载推迟到更晚的时间，我们可以通过减少其在片上SRAM中大内存占用的生命周期，为执行节省更多空间。例如，在图13中，如果按顺序预加载，执行空间为总内存的1/2；如果重排预加载，执行空间可达总内存的5/6。

生成有效的预加载顺序。由于大型模型包含数千个算子，测试所有预加载顺序（$N!$种）是不现实的。然而，大多数顺序都是无效的，因为它们会导致片上内存溢出。Elk通过遵循归纳式算子调度顺序扫描所有算子，并逐步选择下一个要预加载的算子，来枚举所有有效的预加载顺序。图14展示了一个例子，该过程通过增量构建，生成了一个所有有效预加载顺序的后缀树。

剪枝有效顺序搜索空间。鉴于LLM的独特性，Elk可以进一步剪枝候选顺序。首先，许多算子（如softmax）从HBM预加载的数据很少或没有。Elk只关注重排那些HBM加载量大的算子。其次，一个LLM由相同的Transformer层组成。Elk只在一个层内重排预加载，并将相同的顺序应用于相同的层。通过这些规则，Elk将搜索空间从$O(K^C)$减少到$O(k^c)$，其中$k$和$c$远小于$K$和$C$。对于每个生成的预加载顺序，Elk调用§4.2中的算子调度过程，形成一个$O(N K k^{c} P^2)$的搜索空间。Elk在所有预加载顺序中选择最佳的端到端计划。

4.5 映射到硬件

抽象编程模型。Elk生成的执行计划指定了所有算子的预加载顺序和每个算子的划分方案。Elk将该计划映射到一个抽象的编程模型，该模型可应用于通用的带有片外内存的ICCA芯片。如图15所示，Elk抽象出两个在编译期间生成的关键设备函数：（1）preload_async(op=i)命令所有核心根据预加载态划分方案从HBM请求Op i的数据。（2）execute(op=i)根据执行态方案在所有核心上运行Op i。

执行流程与同步。如图15所示，preload_async(op=2)请求HBM控制器将Op2的数据分发到每个核心的SRAM。当数据分发完成时，控制器会在每个核心SRAM中已分发数据的末尾附加一个done_preload_op_2标签。然后，当execute(op=2)被调用时，它会分3步运行。首先，它通过验证done_preload_op_2标签的值来等待preload_async(op=2)完成。其次，每个核心调用distribute_data从对等核心复制共享数据，从预加载态转换为执行态。第三，每个核心调用local_execute来计算一个tile。硬件通过单向同步强制执行三条规则：（1）execute的调用会阻塞所有未来的preload_async和execute，直到其完成。（2）为了强制执行预加载顺序，所有preload_async顺序执行。（3）preload_async(op=i)不阻塞除execute(op=i)之外的任何execute。

5. 实现细节

Elk编译器框架。我们实现Elk为一个通用的编译器框架，可以支持不同的ICCA芯片实现。
1. Elk前端：将PyTorch等ML框架的模型作为输入，转换为ONNX图，从中获取层信息、算子定义和张量形状。
2. 执行计划生成：接收由不同并行执行模型【索引34, 40, 74】编译器生成的单算子划分方案作为输入。对于每个方案，Elk决定tile到核心的映射策略（例如，对于全连接网络是顺序映射，对于mesh网络是维度顺序路由），并通过成本模型估计其计算、内存和互连成本。然后，Elk运行调度、分配和重排程序来生成优化的端到端执行计划。
3. 代码生成：根据目标硬件和选定的划分方案，生成用于计算每个tile的内核代码和核间数据传输操作。计算代码使用供应商提供的库【索引21】中的代码模板。

Elk的可扩展性。Elk将大型模型的搜索空间剪枝到$O(N K k^{c} P^2)$的复杂度。如表2所示，随着模型规模的增长，N线性增长，但其他参数（C, H, P, K）增长较慢或独立变化，因此Elk的搜索空间大小随DL模型大小次线性增长。Elk可以在一台32核CPU上在5分钟内为大型LLM生成端到端计划（图16）。

仿真框架。由于我们能访问的ICCA芯片（IPU-POD4）没有HBM，我们使用真实的IPU-POD4构建了一个仿真框架。该框架使用一个公认的内存模拟器【索引32，DRAMsim3: A Cycle-Accurate, Thermal-Capable DRAM Simulator，2020，IEEE Computer Architecture Letters】来获取HBM访问延迟。框架在IPU-POD4上执行Elk生成的端到端计划，其中一个核心充当HBM控制器，向其他核心广播“HBM数据”，并通过延迟每次广播来模拟HBM延迟。

模拟框架。为了进行敏感性分析和设计空间探索，我们构建了一个ICCA芯片的事件驱动模拟器，模拟所有核心、网络链接和片外HBM访问。它支持全连接和2D mesh网络拓扑。我们使用真实的IPU仿真器来验证我们的模拟器。

A4 实验环境与结果

6.1 实验环境

• 工作负载：评估了不同规模LLM的推理解码阶段（见表2），使用了不同的批处理大小和序列长度。还测试了一个稳定扩散模型（图23）和LLM训练（图24）。
• 硬件配置（仿真器）：仿真了每个ICCA芯片配备4个HBM3E模块【索引39】，共4个ICCA芯片，总HBM带宽为16TB/s。使用了4个IPU MK2芯片，总计5888个核心和3.5GB片上内存。
• 硬件配置（模拟器）：默认模拟4个芯片和16TB/s的HBM带宽。核心配置（计算和本地SRAM）以及网络链接的延迟/带宽与仿真器设置相同。
• 软件配置：Elk编译器框架，事件驱动模拟器。
• 基线：
  • Basic：遵循现有DL编译器，最大化执行空间，用剩余空间预加载下一个算子。
  • Static：扩展了最先进的ICCA芯片编译器T10【索引34】，为整个DL模型找到最佳的静态预加载和执行空间大小。
  • Elk-Dyn：Elk的部分设计，禁用了预加载顺序置换（§4.4）。
  • Elk-Full：完整的Elk设计。
  • Ideal：理论上的roofline性能，预加载和执行有各自独立的互连和全尺寸片上内存。

6.2 端到端性能

总体性能：如图17所示，在我们的仿真器上，Elk-Full的性能平均比Basic高1.87倍，比Static高1.37倍，达到了Ideal性能的94.84%。Elk的性能随着批处理大小和序列长度的增加而良好扩展。

推理延迟分解：在图18(a)中，我们将总时间分解为四类。Basic的预加载和执行重叠很差。Static通过预加载更多算子增加了11.26倍的重叠时间，但受限于固定的空间大小。Elk-Dyn通过动态调整内存分配实现了更好的重叠，但受到互连拥塞的影响。通过平均2.9步的编辑距离重排预加载，Elk-Full消除了Elk-Dyn 87.65%的互连拥塞开销，并将非重叠的预加载时间减少到总时间的0.037%。

6.3 硬件资源利用率

• HBM带宽：图18(b)显示，Basic只使用了34.7%的带宽。Static利用了46.42%。Elk-Dyn达到51.97%。Elk-Full通过预加载重排进一步达到了62.40%的利用率，接近Ideal的64.38%。
• 互连带宽：图18(c)显示，Elk-Full达到了89.52%的利用率，因为它通过重排预加载来匹配算子执行，从而缓解了互连拥塞。
• FLOPS：在图18(d)中，Elk-Full达到了81.06 TFLOPS。尽管我们的仿真器理论上提供更高的TFLOPS，但LLM推理是带宽受限的，实际TFLOPS受限于片上数据传输。Elk-Full的TFLOPS已接近Ideal。

6.4 ICCA芯片的设计空间探索

我们使用ICCA芯片模拟器来探索不同网络拓扑、互连带宽、HBM带宽和计算能力（FLOPS）的性能影响。
1. HBM带宽的影响：如图19和图20所示，更高的HBM带宽可以提高每token延迟，但由于更高的互连拥塞，收益会递减。基于Mesh的网络比全连接网络更容易出现互连拥塞。图21显示，在各种HBM带宽下，Elk-Full都能最好地利用互连。
2. 互连与HBM带宽的协同扩展：如图22所示，互连和HBM带宽应协同扩展以避免性能瓶颈。当HBM带宽较低时，增加互连带宽的收益有限。当HBM带宽较高时，性能随互连带宽扩展，而Elk-Full能最好地利用这两种带宽。
3. 核心数量扩展性：如图23所示，Elk在扩展ICCA芯片核心数量时，能够为ML推理工作负载带来可扩展的性能。对于计算密集的DiT-XL模型，Elk-Full的优势不如LLM明显，但仍优于其他设计并接近理想性能。
4. 对ML训练的启示：如图24所示，对于计算密集的训练工作负载，扩展FLOPS比扩展带宽更重要。因此，ICCA芯片应专注于扩展FLOPS，并可以搭配更便宜的内存（如GDDR/DDR）来降低制造成本。

A7 补充细节

7. 讨论与未来工作

• 将Elk应用于GPU：最新的NVIDIA GPU也使用核间链路连接其流多处理器（SM）。由于其聚合的SM间带宽接近HBM带宽，它将面临严重的互连拥塞。未来工作希望将Elk扩展到GPU，并研究优化GPU互连架构的设计空间。
• 将Elk应用于MoE：Elk可以支持动态专家混合（MoE）模型。在编译时，Elk会基于一个通用专家优化执行计划。在执行时，芯片使用Elk给出的划分方案和运行时选择的专家索引来预加载专家张量。
• 将Elk应用于其他优化目标：通过替换基于性能的成本模型，Elk可以被调整以支持优化其他目标（例如，功耗）。
• 将Elk应用于其他执行模型：对于SambaNova芯片支持的空间流水线执行模型，可以修改Elk的搜索算法来探索其新的调度空间（例如，决定每个芯片的流水线阶段数和每个阶段的核心数）。

8. 相关工作

• 深度学习编译器：许多DL编译器【索引8, 12, 25, 65, 69, 74】是为没有核间链路的架构设计的。一些编译器只在片上服务模型，未考虑片外内存【索引33, 34, 43】。Elk可以利用它们来生成每个算子的划分方案。
• 使用算子融合的ML优化：DL框架通过融合多个算子来提高片上数据复用。由于ICCA芯片具有大的分布式SRAM，可以缓冲整个中间张量，因此融合需求不那么迫切。Elk专注于解决分布式SRAM分配、核间互连拥塞及其对HBM预加载的影响等独特挑战。
• ICCA芯片编译器：数据流编译器优化了片上执行和通信，但它们对片外HBM的使用方式不同。例如，T10【索引34】在优化片上执行时没有考虑片外内存。Elk通过综合考虑片外HBM、片上SRAM、核间互连和每核执行，为通用ICCA芯片最大化性能。
• 分布式模型执行：先前工作侧重于设备间通信的优化。Elk则针对芯片内优化，并额外考虑HBM数据访问及其对每核SRAM和核间链路使用的影响。
• 基于ML和SAT的编译器优化：Elk的代码生成可以利用基于LLM的内核代码优化工作。对于调度问题，Elk的启发式算法比SAT求解器更高效，因为SAT求解器的运行时会随布尔变量数量呈指数增长。

A5 结论

本文研究了支持片外内存的通用ICCA芯片的性能权衡，并开发了一个DL编译器框架Elk来探索ICCA芯片的效率。Elk还支持ICCA芯片架构的设计空间探索。我们使用ICCA仿真器和模拟器展示了Elk的能力。

A6 附录

本附录提供了Elk的编译、模拟和评估框架的源代码，并指导读者探索Elk如何提高各种ICCA芯片上的模型服务性能。
* 算法：归纳式张量算子调度，成本感知的片上内存分配，以及ICCA芯片设计空间探索。
* 神经网络模型：Llama2-13B, Gemma2-27B, OPT-30B, Llama2-70B, and DiT-XL。
* 运行环境：Ubuntu 20.04或更新版本，Python 3.10。
* 指标：执行时间，硬件利用率。
* 访问方式：源代码可从Zenodo (https://doi.org/10.5281/zenodo.16541972) 或 GitHub (https://github.com/platformxlab/elk.git) 下载。
* 依赖：需要x86机器，至少200GB主内存和20GB磁盘空间。
* 实验流程：提供了benchmark_scripts/generate_data_from_sim.py脚本来运行所有测试用例，大约需要30小时。运行evaluate.py脚本来收集数据并绘制所有图表。