重新设计推理芯片： 从 Nvidia GPU 的缺陷到 OpenAI Jalapeño

TL;DR

假如我重生成为某基模型团队的 Infra 负责人, 那么一定会对 Nvidia 的 GPU 及整个系统有各种不满, 然后自己做芯片...

对于 Nvidia 网络部分就不用谈了, RoCE各种缺陷, 然后 DPU 也有一堆性能和安全问题, 也难怪今年 Nvidia 又开始吹 Scale-in 了, 但是一个从来没做过云的团队这方面想成熟起码还要 5 年. 举个例子, 最近我们做了一件事情, 在我们 MaaS 在线业务中的某个 PD 分离的场景中, 将传输协议用了 CIPU 的 eRDMA 替换直接把 TTFT 降低了 40%, 直接等同于维持 TTFT SLA的情况下可以省大量的 Prefill 计算服务器...

对于 GPU 微架构,去年就对此详细写了一篇 《Inside Nvidia GPU: 谈谈Blackwell的不足并预测一下Rubin的微架构》(英文版: https://github.com/zartbot/blog/issues/3 )批评 Nvidia Blackwell微架构的各种缺陷, 老黄还在 NV 内部转发过...

当时建议 SM 增加用于调度的超标量核, 这不 OpenAI Jalapeño 就有了. 另一个问题是我一直给 Nvidia 在提, 你们 GPU 架构中, 大量的L2 Cache占用了很多芯片面积, 但是似乎对它又很难去做一些控制, 缓存层次结构比较厚, 为了维持 UMA 实质上在 L2 两个 Partition 上会增加200~400 cycle 的延迟, 而且对外 I/O 的交互也无法直接写到 SMEM 内,  这下好了 OpenAI Jalapeño 直接砍掉 L2 Cache. 还有前几个月在卷 Anthropic 的一个面试题, 实际上也是在准备和自研训练和推理芯片相关的任务, 并且在内网也详细分析了各种加速器架构以及该怎么做...

今天这篇分析 Jalapeño  的文章, 我假设我自己是某基模芯片团队的负责人, 从模型的需求以及当前 Nvidia GPU 的缺陷出来, 逐步从第一性原理开始推导该如何设计, 当然也会在公司内部逐渐复现整个设计过程.  通过这种方式来阐述 Jalapeño  的芯片架构以及背后的一些工具链... 本文目录如下:

1. 从 Nvidia GPU 谈起 1.1 推理场景下的理论峰值 1.2 架构决定E2E延迟 1.3 Nvidia GPU 的缺陷 2. 如何设计推理芯片 2.1 推理芯片第一性原理 2.2 推理 Infra 现状 2.2.1 不同阶段的性能需求 2.2.2 KVCache 的代价 2.3 推理芯片体系结构 2.3.1 编程接口的变化 2.3.2 体系结构的取舍 3. OpenAI Jalapeño 芯片架构 3.1 芯片架构 Overview  3.1.1 Core Slice 架构 3.1.2 NOC 架构 3.2 Core Slice架构详解 3.3 内存子系统 3.4 片上网络 3.6 系统架构 4. 软件架构 5. 芯片设计 5.1 关键前提: 先换语言, 再让 AI 上场 5.2 设计如何闭环 6. 未来展望 6.1 对比 Nvidia GPGPU架构 6.2 OpenAI roadmap 6.3 一些分析总结

⚠️ 很多更详细的评估以及复现整个设计流程的工作就不公开了... 由于微信公众号排版不适合电脑前阅读, 因此在 Github 上放了一个重新排版的

中文版: https://zartbot.github.io/blog/arch/jalapeno/index.html

英文版: https://zartbot.github.io/blog/arch/jalapeno/en.html

1. 从 Nvidia GPU 谈起

首先想起的是一个段子, 也是现状.

产品管理的兄弟: 隔壁NV做了, 为啥咱不做?

也是产品管理的兄弟: 隔壁NV都没做, 为啥我们能做?

还是产品管理的兄弟: 隔壁NV做了, 为啥我们还要做?

依旧是产品管理的兄弟: 隔壁NV做了, 我们怎么敢做?

大多数做加速器的厂商, 满脑子的是抄 NV 作业, 然后美其名曰 CUDA 生态兼容, 但再怎么抄整体落后 NV 两到三代... 另外一个极端就是总想搞一些奇奇怪怪的体系结构, 试图弯道超车.... 然而几乎所有的芯片厂商和 NVidia 最大的差距在于全栈的能力, 从模型结构和算法, 再到软件生态, 最后到芯片实现, 其中哪些地方可以有 trade-off, 每次交易赢了什么输了什么, 很少有人能够讲清楚. 恰好我从算法到芯片全栈都懂, 因此可以完整的来阐述整个芯片设计的过程以及 Nvidia GPU 的缺陷.

1.1 推理场景下的理论峰值

GPU一直以来"靠大量并发 warp 隐藏访存延迟"的机制处理大量的数据,  这个机制本质上依赖大 batch. Batch 小了, 可调度的 warp 就少, 流水线填不满, 于是远离 roofline. 这便是推理场景中经常遇到的问题, 延迟隐藏机制失效.

既然推理阶段很大程度上是 Memory bound 的, OpenAI 做了一个很简单的计算, 假设累计 ScaleUP 域有 N 个芯片, 每个芯片 HBM 带宽为 M , 估计一个模型参数规模为 K, 那么可以简单的假设计算没有任何瓶颈时, 理论上输出 tokens 的最大速度为 

也就是说理论上, 不加投机解码, 1T 参数的模型应该能过做到  1000~2000

tokens/s/user.  但现阶段大多数推理框架在B300一类的平台上, 做多只能 100~200

tokens/s/user, 只有 10% 的理论性能. 因此这一页在说: 就算把 HBM 带宽用到理论极限, 得到的延迟仍然有一个下界. 也就是说, 单靠加带宽解决不了低延迟问题. 既然带宽不是唯一变量, 那么一定是体系结构上的问题了...

1.2 架构决定E2E延迟

在一年多前一篇文章里我暗示过, 从处理器的 workload 来看, 在现代计算机体系结构中, CPU的核心设计哲学与GPU截然不同: GPU通过海量并行线程掩盖延迟以最大化吞吐量, 而CPU则是不惜耗费巨大的芯片面积与晶体管功耗预算, 穷尽一切微架构手段将单线程执行延迟 (Single-Thread Latency) 压缩到极致.

CPU通过分支预测提前探路、多发射与乱序执行见缝插针地并行计算, 依托多级缓存拉近数据距离, 并利用硬件缓存一致性实现多核间纳秒级的就地数据共享与免内存同步, 从而掩盖了单核计算、主存访问及多核协作的等待延迟.

与之相对, GPU将绝大部分晶体管预算倾斜给密集的算术逻辑单元, 通过SIMT架构组织起数以万计的并发线程, 依托硬件级超快线程上下文切换在某些线程等待访存时瞬间调度其他就绪线程, 但是它根本不在乎单条指令的快慢, 而是用绝对的并行度将数据流的整体处理延迟彻底淹没在海量的计算任务中.

而对于 LLM 推理来看, 它刚好落在下表中 A* 这个位置. batch=1 时一层就是几个 GEMV, 同时还要在毫秒量级穿过数十层, 还要在每层中间做完大量的并行计算后的结果同步.

同时 OpenAI 在 HotChips 上的 Session 有一页专门阐述了这个问题: 我们实际测量到的延迟, 很大程度上是由底层架构本身决定的.

图中的 长延迟路径 → 操作数迟到 → 计算单元阻塞 为传统的 roofline model 补上了另一个维度, 对于 GPU 芯片而言并不是简单的 Compute bound 和 memory bound 的 workload 区分, 而是还要加上一个 阻塞, 此时没有任何资源被打满, 但所有资源都在等. 利用率掉下去, 而两条天花板都还很远.

1.3 Nvidia GPU 的缺陷

我们在这一节详细展开阐述一下 OAI 那一页的观点.

首先我们来谈 OAI 提到的Unified memory subsystems tend to highly contended paths, 它采用统一 L2 + 全局编址, 任何核等价访问任何地址. 但是这样的统一访问虽然编程上隐藏了硬件的复杂性, 但是也带来了大量的工程约束.  例如 SM 太多带宽需求太大, 从 Ampere 开始 L2 拆成了 2 个 Partition,  在 Hopper 上 cross partition 会增加 200 cycles 的延迟. 而在 Blackwell / Rubin 上还会带来更多的影响, 每个 Die 有两个 L2 partition, Dual-Die 会有 4 个 Partition. 在 Blackwell 上会增加接近 400 Cycles 的延迟, Rubin 还需要考虑到 HBM4 的累计带宽已经快接近 L2 带宽了, 同时还要受到 NV HBI 的影响, 我个人认为 Unified Memory Access 将会导致 HBM 22TB/s的实际带宽利用率只有 60%~70%. 并且当并发请求超出端口服务能力时, 产生的排队延迟 (Queuing Delay) 其开销远超数据传输本身.

另一个问题是在一些数据同步和 warp 调度上, 在现代的GPU架构中, 每个计算核有了独立 PC、自由异步推进, 并且配合 TMA 这些异步内存访问器件增加了吞吐, 但正因为计算核是独立的, 硬件就无法预知谁在什么时候到达, 于是一次全局 fence 必须真的去"问一圈", 甚至在一些算子上需要整个system level cross ScaleUP/ScaleOut network 进行 barrier. 这就是 OAI 谈到的 Independent cores out of sync make global memory fences expensive.

最后Centralized resources mediating network access讲的是当成百上千个核心需要向外部网络发起通信时, 若通过集中式的 DMA 控制器、MMIO 接口或共享网卡端口进行收发, 仲裁逻辑 (Arbiter) 会成为严重的吞吐瓶颈与延迟热点, 从而将硬件本身的极速传输能力彻底拖垮. 这个涉及到一系列复杂的互联系统的问题, NOC / ScaleUP / ScaleOut 都有各种问题,  例如 KVCache 拷贝时采用的 CopyEngine 或者 ScaleOut NIC 上 doorbell的约束, 以及很多通知无法支持直接写到 SMEM, 因此只能写到一些 Remote GMEM, 然后 CUDA Core polling GMEM.

正是因为这些原因, 我过去几年才会对 Nvidia 提出 SM 内需要一个超标量核让编程者能够更好的控制指令发射和调度, 同时在L2 上给用户更多的控制能力, 进一步和 SM 内的 SMEM 融合...  而恰好 OpenAI Jalapeño 也作出了这样的选择, 下一章我们将从推理业务第一性原理的角度逐渐展开, 如何设计一颗适合推理的芯片.

2. 如何设计推理芯片

我们仔细分析发现前一节所讲的缺陷,  它们全都是通用性的代价., 需要强调, 这并没有在说 GPGPU 设计得差, 因为每一条决定当初都是对的, 它是 Nvidia 成功的关键, CUDA SIMT的编程范式在过去十多年极大的简化了大家编写并行计算程序的复杂度.  但是这个前提在 LLM 推理场景带来了大量的问题, 然后伴随着 Coding Agent 的发展, 各种 Auto-Kernel(eg. KDA)百花齐放, 似乎我们也不需要为这些通用性付出高昂的代价了...

但这些仅仅是从芯片体系结构的视角, 更重要的是, 我们需要考虑从整个推理业务视角作为第一性原理.

2.1 推理芯片第一性原理

从第一性原理来看, 以一个推理服务经营的视角主要就是基础设施成本和供应链约束下, 满足用户体验的SLO下尽可能多的输出 tokens. 而供应链在北美最大的约束可能主要是在电力供应上.

因此 OpenAI 的演讲稿故意不把chips used,throughput per chip和TTFT列为主目标. 它选择的是 request-level user experience 和 energy per request.

吞吐能够通过并发掩盖某些等待, 单个任务的依赖链却不能凭空并行. 因而使用TTLT, 即time to last token, 表示整个请求完成时间, 同时用TBT, 即 time between tokens, 近似流式交互速度. 但我们注意到, 降低 batch 通常会改善TBT, 却让权重和网络复用下降, 每token能耗上升. 提高 batch 则通常提高tokens/s/kW, 但请求排队和每用户token间隔变长. 因此存在一个响应延迟(TTLT) vs 能效(tokens/Joule)的 Pareto 前沿. 形式化的描述为:

而 LLM 推理, 实际上是 CPU 界研究了三十年的单线程延迟问题...

那么从体系结构的视角来看得出一个很简单的结论:

尽量提高 Data Locality, 降低数据搬运的成本.

从E2E延迟的视角, 进一步提高指令执行的并行性, 降低空闲等待的时间.

接下来几个小节我们将从推理的技术细节详细展开分析...

2.2 推理 Infra 现状

2.2.1 不同阶段的性能需求

通常对于一个推理系统来看, 主要业务分为Prefill, Decode, SpecDecode 三块, 可能我们原来的定性的结论是 Prefill是 Compute Bound, Decode是 Memory Bound, 但实际上我们从 E2E 延迟的视角还需要考虑 SpecDecode 时的 Draft 模型和 Verify 两个阶段, 它们都会处在 roofline 的不同位置.

OpenAI在 hotchips 的session 上也谈到这个问题, 但是其表述是不完善的. 例如在第三格, 主标题是Spec-verify, 副标题是Decode,但的自回归decode与的批量verify处在roofline上两个不同位置, 相差倍算术强度.

因此我们严格的把它分为四块, 为了简化所分析的问题, 此时我们以单个 Request 的处理进行分析

1. Prefill: Prefill一次处理完整输入上下文. 对 hidden-dim 为, 长度为  的普通Transformer模型, projection 与 FFN 的主计算量近似随  增长, full attention 的 score 部分随  增长. 大  带来足够的 M 维度, 权重可以在许多 token 之间复用, 所以 tensor engine 容易进入高利用率区. 因此它是一个 Compute-bound 的计算, 即 compute high, memory BW low, 另外对于通信而言, 它的 collective 消息随  增长, 属于带宽受限而非延迟受限, 因而容易平滑调度.

2. Decode: 自回归模型在 Decode 阶段每个序列每次只处理 1 个新 token, 即 , 矩阵计算退化为 GEMV, 这意味着一份权重只服务一行, 在短或中等 context 下通常受 Weight HBM 带宽限制, 在长 context 下会逐渐转为 KVCache HBM 带宽限制. 在分布式系统中, 与此同时每层边界仍要做一次collective, 消息只有大小, 小消息 collective 的网络同步延迟可能成为第一瓶颈.

3. Draft Model: 通常和模型相关. 经典自回归草稿模型和EAGLE类方法需要执行多次串行起草步, 因而经常受权重/KVCache带宽和同步延迟限制. DFlash使用一次块并行扩散主干网络前向传播同时生成整个草稿块. DSpark使用一次并行主干网络前向传播, 再叠加一个低成本Markov或RNN串行头. 对DFlash/DSpark而言, 草稿权重只需按主干网络调用读取, 算术强度可随  提升, 瓶颈可能转向矩阵算力, 目标隐藏特征/KV注入带宽,  vocabulary projection 或轻量串行采样延迟.

4. Verify: 验证的物理token数也不一定等于固定的 . 线性草稿通常验证一个前缀, 树形草稿需要验证全部树形节点, DSpark则按置信度和实时硬件容量为每个请求选择不同的验证长度. 稠密权重可以跨本轮物理验证token复用, 但是MoE目标模型可能因为去重后的专家数量随token预算增长而失去大部分权重复用. 高效验证注意力必须让一个前缀KV分块被同一请求的全部候选查询或树形节点复用.

总结成一个表如下:

可以看到在不同阶段对于芯片的计算, 内存, 网络带宽, 网络同步延迟 的需求会存在巨大的不同. 如果我们采用专用的芯片应对每个阶段, 对于资源的配比也会出现问题. OpenAI 也谈了这个问题:

OAI阐述的是当前各种 PD 分离和 AFD/SpecDecode 这样的 LPU 部署带来的缺陷. 原页面沿用它自己的三分法, 所以只画了prefill, draft, verify三个比例. 实际上按照我们前面所分的4个阶段分别放入4种专用设备, 看似能让每种芯片都针对自己的kernel做到最优. 问题在于设备台数在部署后是离散且相对固定的. 设4个池的能力为, 到达工作量为, 下标依次对应prefill, decode, speculate与verify, 则可服务速率受最紧张的那个池限制:

例如我们来分析 SpecDecode 接受率动态变化的影响,  与 此消彼长, 而拨动它们的不是请求. 一个请求的输出token只有两条出路, 走纯decode, 或者走 Draft + Verify. 记输出token中经 SpecDecode 路径提交的比例为, 为每秒已提交的输出token数, 则

其中是一轮 Draft Model 的成本, 自回归为步而块并行为1次宽pass, 则正比于该轮的物理 token 预算. 不由请求混合决定:  塌陷把它压向0, 高并发下调度器按前缀存活概率裁掉低置信后缀也把它压向0, 而接受率回升时它又回到1.

于是在同一批请求, 同一个模型, 同一个context length下, decode池与verify池的负载可以走向两个相反的极端. decode池必须按配齐, draft与verify池必须按配齐, 而这两个峰值不会同时到来: 资源预留的是峰值之和, 用的是其中一组.

四个阶段的限制项落在不同资源轴上: prefill是矩阵算力, decode是HBM权重带宽加逐层同步延迟, 自回归起草是draft权重带宽加次同步, verify是去重expert的权重流量加EP all-to-all突发. 所以这四种专用芯片不是同一款芯片的大小号, 而是算力, 带宽与网络配比互不相同的四种芯片. 接受率上升时, verify的次数下降; context length上升时, prefill与KV的成本提高; 模型换成expert占比高的MoE时, verify的网络与HBM压力又增加. 一个pool成为瓶颈时, 另外两个pool仍要支付package, HBM, I/O, 网络与散热的基线成本. 由此得到一个容易漏掉的后果: 即使运维愿意把设备从闲置的池搬进 瓶颈池, 搬过去也无法使用, 因为prefill芯片的HBM与算力之比本来就不是按decode配的. 固定的设备台数比例等于同时锁死五个资源维度的比例.

2.2.2 KVCache 的代价

在传统的推理部署中, 我们通常会考虑 PD 分离的部署方式, 而此时大量的 KVCache 的搬运也带来了极大的功耗以及网络同步的损失. OpenAI 在 hotchips 中的阐述如下:

这一页把架构选择压成一个二选一: 异构要么发生在系统之间, 代价是 KV 跟着不同的阶段走, 这就是现在 PD 分离的做法, 但是每个阶段的网络和同步所带来的影响以及KVCache 搬运的功耗是巨大的.另一种选择是将不同阶段的 KVCache 移动放在芯片之内, 代价是活跃资源的比例跟着不同的阶段变化.

通常 KVCache 的大小为  和序列长度成正比, 即便是一些 Linear Attention 的算法也需要在多次交互中为每一轮留下 Snapshot. 数据的搬运会导致 ScaleUP/ScaleOut 以及 NOC 对原有的集合通信的性能影响, 同时还要考虑大量的 KVCache 从 Prefill 节点的 HBM 读出, 再从 Decode 节点HBM写入, 依然会影响到原有的推理算子的执行. 具体的分析可以从时间比, 每轮边界, HBM容量与事务语义 四个方向展开进行量化的分析.

由于涉密这部分就不公开多说了...还记得本文开头所讲的我们 MaaS 服务将传输协议换成 CIPU eRDMA 后, 推理引擎未做任何代码修改 TTFT 下降了 40%...

最后对于芯片的设计, 如果用几种独立加速器, KV cache 需要在设备之间移动. Jalapeño 的策略是把异构能力放在一颗平衡芯片内部:

当前阶段需要 tensor compute, 就激活更多计算单元.

当前阶段受 HBM 限制, 就强调本地内存路径.

当前阶段通信密集, 就使用 collective 和 scale-up 网络.

暂时不需要的模块进行 power gating.

KV cache 尽可能留在本地.

2.3 推理芯片体系结构

几个月前 David Patterson 有一篇文章《Challenges and Research Directions for Large Language Model Inference Hardware》[1]在阐述推理硬件的体系结构, 对于内存子系统, 使用 HBM / 3D-DRAM / HBF 有很多讨论, 但是为了简化问题, 例如新的 HBM 可以构建 PNM 和 Customized Based Die 将 Memory Controller 和 HBM PHY 从计算 Die 中移出获得更多的算力资源, 或者像 AMD 那样使用 XCD 和 FCD 的 3D 封装, 可以进一步提高 XCD 的良率以及更大的 L2 Cache 和互联NOC, 我们暂时不讨论这部分的内容, 更多的关注在计算 Die 内部的微架构以及互联架构.

具体的内容在公司内部也有一个详细的文档分析了各种算力芯片的微架构, 包括 Nvidia GPGPU / LPU / Google TPU 和华为 Ascend 950. 并且针对 Agent based compiler 这些工作进行了一系列探索.

2.3.1 编程接口的变化

前几个月我在做一个 Anthropic 公司招聘用的性能工程相关的家庭作业面试题, 任务代码在: https://github.com/anthropics/original_performance_takehome , 大概的内容是在一个自定义的 VLIW + SIMD 处理器上进行性能优化. 通过使用Claude-Opus模型进行优化, 后来刷到第二以后因为差距已经在2%以内, 以及一些计算资源和 token 费用的问题就没有继续做下去了, 仿佛现在还是第三...

通过这个题目让我意识到它就是在设计一个专用 ASIC 芯片用于推理和训练的编译系统, 似乎原有的 CUDA 生态和 SIMT 的编程接口虽然对人友好, 但带来的资源开销是巨大的, 然而有了 Agent as compiler 了以后, 并且伴随着 Ligeng‘s KDA 这样的工作, Kernel 优化似乎可以通过 Agent 很好的完成了, 似乎再也不需要去照顾人类的心智缺陷, 而充分的发挥硬件的能力.

那么原有的 SIMT 的抽象和 GPGPU 在推理场景下的优势也就不存在了, 通常一个模型的架构在预训练阶段就已经定死, 我们会有两个月的时间充分的对推理平台进行基于 Agent 的调优. 这也是 OpenAI Jalapeño 体系结构上解除的一道枷锁.

2.3.2 体系结构的取舍

去年的blog也讨论过一个问题, Nvidia 的 Warp 调度机制和延迟隐藏机制并不适合推理场景. 另一方面在文章《谈谈GPU的内存模型及互联网络设计》中也谈到了 Nvidia GPGPU 复杂的内存语义, 数据路径上存在 Genric Proxy 和 Async Proxy, 如果需要 ScaleOut 还要额外的同步机制, 然后整个 NOC 和内存层次结构也越来越重, 而 Warp 调度的机制来看并不能让我满意, 每一次数据搬运都需要显式编排, 而编排本身有开销, DMA 描述符的构造、启动、以及等待完成的 barrier 等等...

假设一个 tile 的计算只需要 500 个周期, 而 DMA 编排 + barrier 开销是 100 个周期, 那么固定开销占了 20%. 但是在 GPU/TPU 的典型场景 (大 batch, 一个 tile 算几千周期) 里, 同样的 100 周期只占 2%, 完全可以忽略. 很不幸的是, Decode 阶段正好落在前一种情况: batch 小, 每个 tile 的计算量小, 于是同步开销从"可忽略"变成"主导项". 这时候把同步从"显式 barrier"换成"硬件依赖跟踪 (乱序 + cache)" 就有了正收益.

因此, 我希望他们能在 SM 内加一个超标量核进一步来控制延迟, 正好这次 OpenAI Jalapeño 中的 Out-of-Order Core 正好也验证了这个判断.

同时, 前一段时间也对 Blackwell 系列进行了完整的指令 Cycle 级的测量和完整的芯片级分析, 对外公开了《Dissecting SM_120 through Microbenchmarking》[2]. 而对于 B300 片上网络 / L2 Cache 和功耗管理上也发现了大量的问题. 简而言之就是整个内存层次结构太重了, 而片上网络中, 为了维持 L2 的带宽, 很早就被设计成了 2 个 Partition, Nvidia 期望继续维持 UMA 的编程模型来隐藏不同 L2 Partition 访问的延迟差, 但是在 Multi-Die 的 Blackwell / Rubin 上代价太大, MultiDie的结构使得它单颗 GPU 有了 4 个 L2 Partition, 再维持 UMA 在推理阶段引入的访问内存损失就多达几百个周期, 显然是得不偿失的.

另一方面是在 RDMA 这些场景, 很多通知机制我希望能够 bypass L2 直接写到 SMEM, 这样 SM 就可以很快速的 polling SMEM 获得相对确定性的延迟, 而当前 polling GMEM 很有可能会被 warp schedueler yield导致更大的延迟. 并且 L2 Cache 的机制我们是很难控制的, 很有可能我们关键的一些同步用的 counter 被 SM 大量的内存访问 evict, 我一直也在跟 Nvidia 提什么时候能够把 L2 Cache 也划分出一部分作为 SMEM 那样可以自己控制的器件....

还有进一步的一些问题, 例如由于同步机制的影响, 无论是 intra-chip 还是 inter-chip 的集合通信性能都会受到延迟的影响, 这也导致了很多场景下 TP 并行的规模无法进一步扩展成为显著的影响点. 而这些问题似乎都来自 UMA 的影响, 那么很自然的一个想法, 能否做成一个 NUMA 结构, 并将关键的集合通信相关的处理在 NOC 和互联(ScaleUP/Scale Out)进行充分的优化?

当然我能理解 Nvidia 作为一个商业公司要去追求 peak througnput 的原因, 但是对于推理系统而言, 我们从目标函数出发, 有两条独立的推论, 分别落在微架构和系统架构上.

推论 A: 有效 token 必须靠"低延迟下的高利用率"取得, 不能靠堆 batch, 具体来看:

1. 缩短访问内存的路径: 对于内存延迟放弃隐藏的策略, 而通过更少的内存层次消除, 可以接受 NUMA 的代价

2. 更高的单指令流的ILP: 使用乱序核 + 硬件 L1 + 预取的方式, 而不依赖 occupancy

3. 适当的矩阵单元大小: 保证小batch时也不会有性能损失

4. 消除 Kernel Launch 与同步的固定开销

推论 B: 芯片面积不再是稀缺资源, 闲置整机才是浪费.

1. 拒绝异构的各种分离技术, 同构集群避免专用资源调度瓶颈

2. 不以 peak FLOPS 为目标, 充分考虑推理各阶段的资源需求, 做到算力/内存带宽/通信的平衡

3. 允许部分的暗硅, 保证芯片满足各阶段的需求

4. 平衡整个芯片的功耗, 并以它作为芯片设计目标

5. 降低通信的延迟, 特别是长尾延迟

整个在体系结构上的决策可以概括为下面这张图:

而这些体系结构方面的思考和 OpenAI 不谋而合, 具体的分析我们将在下一章展开.

3. OpenAI Jalapeño 芯片架构

Jalapeño 应该是工业界第一款 AI 加速的软硬件协同设计的芯片平台, 硬件上从第一行 RTL 到最终 tapeout 仅用了 9 个月, 当然这和整个芯片架构也紧密相关, 由于没有复杂的同步和硬件调度机制, 这使得验证的周期加快了很多.

← 返回资讯流