DeepStack：3D堆叠AI芯片的"设计导航仪"，10万倍加速寻最优解

当AI模型遇上"内存墙"

想象你正在驾驶一辆超级跑车，它的发动机拥有惊人的马力，可以在几秒内从0加速到100公里。但有一个问题：油箱的输油管太细了，发动机无法获得足够的燃料，导致它无法发挥全部性能。

这正是当前AI芯片面临的困境。随着大语言模型（LLM）的规模从数十亿参数增长到数千亿甚至数万亿参数，计算能力的需求呈指数级增长。但与此同时，内存带宽却成为了严重的瓶颈——这就是业界常说的"内存墙"（Memory Wall）问题。

现代AI模型不仅需要存储海量的参数，还需要在推理过程中频繁地读取和写入数据。传统的2D芯片架构中，计算单元和内存单元位于同一平面上，数据需要通过有限的引脚在两者之间传输。随着模型越来越大，数据传输的开销已经远远超过了实际计算的开销。

3D堆叠：打破内存墙的曙光

为了突破内存墙，业界开始转向3D堆叠技术。这种技术的核心思想是：不再把所有组件平铺在一个平面上，而是像搭积木一样，把计算芯片和内存芯片垂直堆叠起来。

通过先进的混合键合（Hybrid Bonding）和封装技术，3D堆叠可以在垂直方向上建立大量的高速连接。这意味着计算单元和内存单元之间的距离从原来的几毫米缩短到了几十微米，数据传输的带宽和能效都得到了数量级的提升。

3D堆叠带来的好处是显而易见的：

更高的内存带宽：垂直连接的数量远远超过平面引脚的数量，带宽可以提升10倍甚至更多。

更低的访问延迟：数据不需要走"远路"，传输延迟大幅降低。

更高的能效：短距离传输意味着更低的功耗，对于数据中心来说，这意味着更低的电费账单和更小的碳足迹。

更大的内存容量：可以在有限的空间内集成更多的内存，满足超大模型的需求。

然而，3D堆叠也带来了新的挑战。

分布式3D推理：复杂度呈指数级增长

当模型规模达到数百亿甚至数万亿参数时，单个3D芯片已经无法容纳全部模型。这时候就需要分布式推理：把模型分割成多个部分，分布在多个3D芯片上，让它们协同工作。

这就引入了一个极其复杂的优化问题。设计一个分布式3D堆叠AI系统，需要考虑的因素包括：

硬件层面：

• 每个3D堆叠应该有多少层DRAM？
• 层与层之间的垂直连接如何配置？
• 芯片之间的互连网络应该用什么拓扑？
• 计算资源和内存资源如何分配？
• 功耗和散热如何管理？

系统层面：

• 模型应该如何切分？按层切分？按张量切分？
• 使用什么并行策略？数据并行？模型并行？流水线并行？
• 如何调度计算和通信，最大化重叠？
• 批处理大小如何设置？

这些因素相互影响、相互制约。改变一个参数可能会对其他多个参数产生连锁反应。更糟糕的是，设计空间的规模是惊人的：仅仅考虑几个关键参数的组合，设计点的数量就可能达到10的14次方（100万亿）级别。

在这个巨大的设计空间中，如何找到最优的配置？传统的暴力搜索显然不可行——即使每秒评估一个设计点，也需要数万年才能遍历整个空间。

DeepStack：专为3D堆叠AI系统设计空间探索而生

针对这一挑战，研究团队提出了DeepStack，一个准确且高效的性能建模工具和框架，专门用于分布式3D堆叠AI系统的早期设计空间探索（Design Space Exploration, DSE）。

核心设计理念

DeepStack的设计理念可以概括为："在保持准确性的前提下，尽可能快地评估设计点"。

传统的模拟器（如基于周期的模拟器）虽然准确，但运行速度极慢——评估一个设计点可能需要数小时甚至数天。这对于探索巨大的设计空间来说是完全不可接受的。

DeepStack采用了一种新颖的建模方法，通过抽象和近似，在准确性和速度之间取得了绝佳的平衡。它能够在毫秒级别评估一个设计点，同时保持与详细模拟器相当（误差在2-12%范围内）的准确性。

硬件级建模：精细捕捉3D内存语义

DeepStack在硬件层面提供了极其精细的建模能力，专门针对3D堆叠内存的特性：

事务感知带宽建模：传统的带宽模型往往假设数据可以连续不断地传输，但现实中内存访问是以事务为单位的。DeepStack精确建模了事务的开销，包括行激活、预充电、读写切换等操作的延迟，从而准确预测实际的可用带宽。

Bank激活约束：DRAM内部被划分为多个Bank，每个Bank有自己的行缓冲器。DeepStack建模了Bank级别的并行性和冲突，准确反映多线程访问模式下的性能。

缓冲限制：3D堆叠芯片中的片上缓冲是有限的。DeepStack建模了缓冲容量对数据重用和通信开销的影响，帮助设计者找到最优的缓冲配置。

热-功耗建模：3D堆叠虽然带来了性能提升，但也带来了散热挑战。DeepStack集成了热模型和功耗模型，可以预测不同工作负载下的温度和功耗，确保设计在热约束范围内。

系统级建模：全面的并行化和调度策略

在系统层面，DeepStack支持全面的并行化策略和执行调度：

数据并行：多个芯片同时处理不同的输入批次，适合批处理场景。

模型并行：模型的不同层分布在不同芯片上，适合超大模型。

流水线并行：计算和通信重叠，提高资源利用率。

张量并行：单个层的计算分布在多个芯片上，减少每芯片的内存需求。

混合并行：上述策略的组合，适应复杂的模型和硬件配置。

DeepStack还建模了各种调度策略，包括静态调度和动态调度，以及计算-通信重叠优化。

创新技术：双阶段网络抽象和Tile级重叠

DeepStack引入了多项创新技术来实现高效建模：

双阶段网络抽象（Dual-Stage Network Abstraction）：传统的网络建模要么太粗略（只考虑带宽和延迟），要么太详细（逐个数据包模拟）。DeepStack采用双阶段方法：在粗粒度阶段，使用分析模型快速估计通信时间；在细粒度阶段，针对关键路径进行更精确的模拟。这种方法既保证了速度，又保证了关键路径的准确性。

Tile级计算-通信重叠（Tile-Level Compute-Communication Overlap）：在AI推理中，计算和通信往往可以重叠进行——当一个芯片在计算当前批次时，可以同时在后台传输下一批次的数据。DeepStack在Tile（计算块）级别建模这种重叠，准确预测实际的执行时间。

验证：与真实硬件和详细模拟器的高度一致

DeepStack的准确性经过了多重验证：

内部3D设计：研究团队使用自己设计的3D堆叠芯片进行验证，DeepStack的预测与实际测量结果高度一致。

NS-3后端：NS-3是一个广泛使用的网络模拟器。DeepStack与NS-3的对比显示，误差仅为2.12%，证明了网络建模的准确性。

vLLM在8xB200 GPU上的实测：vLLM是一个流行的推理服务框架，B200是NVIDIA的高端GPU。DeepStack预测与vLLM在8卡B200系统上的实际运行结果相比，误差为12.18%。考虑到这是一个完全不同的硬件平台，这个误差水平已经非常出色。

速度提升：最关键的是，DeepStack相比最先进的详细模拟器，运行速度提升了10万倍（100,000x）。这意味着原本需要数小时的模拟，现在只需要几秒钟。

设计空间探索：在100万亿个设计点中寻找最优解

有了高效的评估工具，DeepStack进一步引入了层次化的设计空间搜索策略，以应对巨大的设计空间。

设计空间的规模

DeepStack探索的设计空间包括：

• 3D堆叠DRAM的层数（2层、4层、8层、16层……）
• DRAM垂直连接的配置（TSV数量、带宽、延迟）
• 芯片间互连拓扑（全连接、环形、网格、胖树……）
• 计算-内存分配比例
• 并行化策略（数据并行度、模型并行度、流水线深度）
• 批处理大小
• 调度策略参数

这些参数的组合产生了约2.5 × 10^14（250万亿）个设计点。

层次化搜索策略

面对这个巨大的空间，DeepStack采用了层次化的搜索策略：

第一层：粗粒度剪枝。使用非常快速的分析模型，快速排除明显不合理的设计。比如，如果内存容量不足以容纳模型，或者带宽低到无法满足计算需求，这些设计点会被立即剔除。这一步可以将设计空间缩小几个数量级。

第二层：中等精度评估。对剩余的设计点使用DeepStack的中等精度模式进行评估，进一步筛选出候选集。

第三层：精细评估。对最终的候选集使用最高精度的模式进行详细评估，找到最优设计。

第四层：局部优化。在最优设计附近进行局部搜索，微调参数，寻找可能的改进。

这种层次化策略确保计算资源被用在最有希望的设计点上，而不是浪费在明显不合理的配置上。

实验结果：9.5倍吞吐量提升

DeepStack的设计空间探索带来了显著的收益：

相比基线设计，吞吐量提升高达9.5倍。这意味着在相同的硬件成本下，可以处理9.5倍的推理请求；或者在相同的吞吐量需求下，只需要原来1/9.5的硬件。

这一提升来自于并行化策略和3D架构的协同优化。DeepStack发现，最优配置往往不是直觉上的选择——比如，它发现批处理大小对架构选择的影响比预填充/解码阶段的区分更为根本。

关键发现：

1. 批处理大小驱动架构分化：不同的批处理大小需要完全不同的硬件配置。高批处理场景需要更大的内存带宽和更多的计算单元，而低批处理场景则对延迟更敏感。

2. 并行策略与硬件架构紧密耦合：选择不合适的并行策略，即使硬件再强大也无法发挥性能。DeepStack发现，不完整的调度搜索会导致永久性的次优性能——软件调优无法弥补硬件配置的错误。

3. 3D堆叠层数存在甜蜜点：更多的DRAM层带来更高的带宽和容量，但也带来更高的成本和散热挑战。DeepStack帮助找到特定工作负载下的最优层数。

4. 互连拓扑的重要性被低估：在分布式系统中，芯片间的互连往往是性能瓶颈。DeepStack的优化显示，投资更好的互连网络（如全连接或胖树拓扑）往往比单纯增加计算单元更有效。

实际应用价值

DeepStack对AI芯片设计和部署具有重要价值：

芯片架构师

对于正在设计下一代AI芯片的架构师，DeepStack提供了一个强大的探索工具。它可以在流片（Tape-out）之前快速评估不同的架构选择，避免昂贵的试错。

比如，架构师可以问："如果我把DRAM从4层增加到8层，性能会提升多少？值得吗？"DeepStack可以在几秒钟内给出答案，而不是等待数周的模拟。

系统部署工程师

对于部署AI推理服务的工程师，DeepStack可以帮助优化系统配置。给定特定的模型和工作负载，DeepStack可以推荐最优的并行策略、批处理大小、调度参数等。

研究人员

对于研究新型AI架构的研究人员，DeepStack提供了一个快速验证想法的平台。研究人员可以快速测试新的并行策略、调度算法、内存层次结构等，而无需构建昂贵的硬件原型。

开源与未来展望

研究团队计划将DeepStack开源，以支持未来的研究。开源将包括：

• DeepStack的性能建模框架
• 预训练的模型（用于不同硬件配置的快速评估）
• 设计空间搜索工具
• 基准测试套件

开源DeepStack将降低3D堆叠AI系统设计的门槛，让更多的研究者和工程师能够参与到这个领域的创新中来。

未来的研究方向包括：

支持更多硬件类型：目前的DeepStack主要针对3D堆叠DRAM，未来可以扩展到HBM、CXL内存、存内计算等新兴技术。

训练时间建模：目前的DeepStack主要关注推理，未来可以扩展到分布式训练场景的建模。

自动优化：结合机器学习技术，实现全自动的设计空间探索和架构优化。

多目标优化：除了性能，同时考虑成本、功耗、可靠性等多个目标，找到帕累托最优的设计。

结语：让芯片设计从"艺术"走向"科学"

芯片架构设计长期以来被认为是一门"艺术"——依赖工程师的经验和直觉，通过反复试错来找到好的设计。DeepStack代表了向"科学"的转变：通过准确的建模和系统的搜索，让计算机帮助我们找到最优解。

在AI芯片竞争日益激烈的今天，设计效率至关重要。能够在几小时内探索数百万个设计点，而不是花费数月进行手工评估，这种能力可能是竞争胜负的关键。

DeepStack不仅是一个工具，更是一种方法论——它展示了如何通过建模和优化来应对复杂的系统设计挑战。随着AI模型和硬件架构的持续发展，这种方法论将变得越来越重要。

3D堆叠技术为突破内存墙带来了希望，而DeepStack则为驾驭这种技术提供了导航。在这个250万亿个设计点构成的迷宫中，DeepStack是那盏指引我们找到最优路径的明灯。