数据中心/云端

AI 模型协同设计:硬件友好型 LLM 设计

AI 性能可归结为三个维度: 

  1. 准确性: 模型推理和生成输出结果的效果
  2. 吞吐量: 数据中心每秒可生成的 token 数量
  3. 交互性: 模型对用户的响应速度如何 (主要受延迟影响)

部署必须平衡这三者:如果响应速度缓慢,就会浪费高准确度;如果每个用户的体验滞后,则原始吞吐量就意味着微不足道。因此,实用的系统可同时优化准确性、吞吐量和交互性。

本文重点介绍吞吐量和交互性,以及模型设计选择如何在不牺牲准确性的情况下影响这两者 (我们会在出现准确性权衡的地方标记准确性权衡) 。 

保持精度不变,问题就变成了二维帕累托边界:改进一个通常会使另一个付出代价。目标是将整个边界向外推进,使曲线下方的区域最大化 (请参见下图 1) 。

我们将从 LLM 开始,这是当今最重要的 AI 工作负载。这种权衡有两个方面:一是系统部署人员优先考虑集群吞吐量 (token/ 秒) ,二是用户重视较低的初始 token 和 token 间延迟。token 间延迟与之相反,token/ 秒/ 用户延迟越高,响应速度越快。

单个响应如下所示:prompt =[第一个 token 延迟]+ 第一个 token +[token 间延迟]+ 下一个 token,以此类推。 

对于希望其模型在现代硬件上良好运行的模型开发者来说,这是一个实用的入门指南。思路很简单:与硬件匹配的模型不仅运行速度更快,而且可扩展得更好、成本更低,并得到更广泛的采用。 

首先,考虑部署场景,该场景沿两个轴变化 (请参见下图 2) 。

工作负载从短上下文到长上下文不等,服务目标从以吞吐量为导向 (最大限度地提高 token/ 秒) 到以延迟为导向 (最大限度地减少响应时间) 不等;许多服务目标介于两者之间。 

每个象限需要不同的优化:长上下文、面向吞吐量的服务将大部分时间用于注意力,而面向延迟的服务则增加了模型并行性,以缩短注意力和 FFN,但代价是通信和固定的开销。短上下文、面向吞吐量的服务可在注意力和 FFN 之间更均匀地分配时间,并可大规模受益于并行性 (例如专家并行性) 。

阿姆达尔定律适用:优化一个部件只对其占用的时间有帮助。如果注意力是运行时的 77%,则调整前馈层只会产生微小的收益,而注意力路径是值得付出努力的地方。了解自己的现状可以告诉你应该关注哪些方面。

本文提供了一些简单的经验法则来帮助您尽早做出这些决策,适合那些不想成为系统工程师但最希望从硬件中获得最大益处的人。本系列的后续章节将针对硬件感知设计的不同维度,帮助您避免计算瓶颈,在数据中心规模上顺畅部署,根据您的用例进行设计,并在实践中进行扩展。

更智能的设计。部署速度更快。扩展更广泛。我们开始吧。

对 LLM 中的线性层进行硬件友好的维度计算

Transformer 设计的关键选择是其纵横比:模型宽度与层数 (\(L\)) 之间的平衡。在解码器式模型中,这些因素决定了计算和内存在整个堆栈中的分布方式。宽度本身由两个维度设置,即 MLP 层中的隐藏维度 (\(H\)) 和中间投影维度 (\(H’\)) 。 

\(H\)、\(H’\) 和 \(L\) 共同塑造了模型映射到并行策略的清晰度,以及模型在 GPU 上的扩展能力。本文将重点关注线性层,探讨这三种选择对吞吐量、交互性和可扩展性的影响。 

算术强度的作用

在任何硬件上,可实现的性能都受“超限”模型的限制。工作负载到达的位置取决于其算术强度,定义为每字节内存移动所执行的计算操作数。 

低算术强度的工作负载受显存带宽 (显存受限) 限制;高算术强度的工作负载受设备峰值计算吞吐量 (以 FLOPS (计算受限) 为单位) 限制。当目标是最大吞吐量 (每秒 token 数) 时,您希望将工作负载驱动到计算受限的区域,以便使用硬件的全部数学运算能力。

延迟敏感型解码则恰恰相反,它以低并发率运行且受内存限制,因此缩短内存访问时间可以降低响应延迟。

提高每字节运算量的一种直接方法是增加批量大小,但模型形状也很重要。我们来看看 \(H\) 和 \(H’\) 如何决定 GEMM 是否受计算限制或内存限制。在单个设备上运行时,\(H\) 和 \(H’\) 将 GEMM 的形状设置为: 

\(C = AB\)

按照典型线性代数库中使用的约定,\(A\) 是大小为 \(M \times K\) ( \(M\) rows,\(K\) 列) 的矩阵,\(B\) 是 \(K \times N\)。产品 \(C\) 具有 \(M \times N\) 输出,每个输出都是长度为 \(K\) 的点积,需要 \(M \times N \times K\) 融合乘加 (FMA) 。由于每个 FMA 为 2 FLOP (乘以 1 加 1) ,因此成本为: 

\(\text{FLOPs} = 2 \times M \times N \times K\)

\(\text{Read Bytes} = M \times K \times \text{bytes}_A + N \times K \times \text{bytes}_B\)

\(\text{Write Bytes} = M \times N \times \text{bytes}_C\)

其中 \(\text{bytes}_{A,B,C}\) 是按所用精度设置的每元素字节计数。以下表 1 以 \(A\) 作为输入,\(B\) 作为权重,将 token 数量 (\(\text{Tokens} = \text{concurrency} \times \text{sequence length}\)) 、\(H\) 和 \(H’\) 映射到 GEMM 的每个线性层的 \(M\)、\(N\) 和 \(K\)。 

层名称 投影 ( in = out) GEMM M GEMM N GEMM K
Q/ K/ V 输入线性* H 3H token 3H H
注意力输出线性 H = H token H H
FFN-1 (向上投影) H H token H H
FFN-2 (向下投影) H H token H H

表 1. Transformer 块中线性层的 GEMM 维度 (M,N,K) ,以 Token、H 和 H′ 表示

假设每个 GEMM 都是“正方形”,即 \(\text{Tokens} = H’ = H\)。然后,单个 GEMM 在移动内存的 \(3H^{2}\) 个元素时执行 \(2H^{3}\) FLOP,因此算术强度会随着 \(H\) 的增加而增长。实际结果是:当 \(H\) 或 \(H’\) 较小时,GPU 在移动数据方面所花费的时间比例上要高于数学计算,即使 token 维度较大。 

具体示例以 FFN-2 为例,它特意使用较小的 \(H’=512\) 和 \(H=8192\),在 GB300 上使用 4 位输入和 8 位输出。如下表 2 所示,即使在较大的 token 计数 (高 GEMM-M) 下,此层仍受内存限制,由数据移动主导;写入成本占主导地位,因为输出是 FP8,而输入是 FP4,而小归约维度使 GEMM 受内存限制。

 

M (token)  \(N\)  \(K\)  数学运算 ( s)  FP4 读取 (微秒)  FP8 写入 (微秒) 
256  8192  512  0.14  0.30  0.26 
2048  8192  512  1.15  0.37  2.10 
16384  8192  512  9.16  0.89  16.8 
表 2. 在 GB300 上使用 FFN-2 (NV_ltx_43 = 8192,NV_ltx_44 = 512) 时,假设 FP4 计算峰值为 15 PFLOPS,内存带宽峰值为 8 TB/s,则每个 GEMM 的理论持续时间。每个 token 计数时的内存时间都超过了数学时间,因此层在整个过程中始终受内存限制

下面图 4 中的图形在 GB300 芯片上显示了这一点。当 NVFP4 GEMM 和 token 维度 (GEMM-M) 固定为较大的 8192 时,我们将扫描归约维度 (GEMM-K) (左图) 和投影维度 (GEMM-N) (右图) 。在这两种情况下,当扫描维度较小时,吞吐量会急剧下降,这确认了小型 \(N\) 或 \(K\) 会使硬件未得到充分利用。( \(N\) 和 \(K\) 对应的是 \(H\) 或 \(H’\),具体取决于层;请参阅上面的表 1. ) 

 

对于模型设计师来说,这一点很重要:对于饱和计算,模型尺寸与批量大小一样重要。

准则 1: 对于固定参数模型,首选近平方权重矩阵,并避免使投影或归约维度过小。

但仅凭规模是不够的。为了达到较高的 Tensor Core 利用率,GEMM 的维度还必须清晰地映射到底层平铺几何图形上;对齐不当会导致平铺量化,即使运算强度较高,也会降低吞吐量。

GPU 如何执行 GEMM 

GPU 通过将输出矩阵拆分成由流多处理器 (SM) 计算的图块来执行 GEMM。在最新的 GPU 上,SM 不必单独工作;它们可以在单个更大的图块上进行协作:借助集群 MMA,两个相邻的 SM 在一个图块上联合起来,而协同网格阵列 (CGA) 可让整个 SM 组作为一个集群一起工作。 

协作可改善数据重用,但也会放大有效图块,因此维度必须是较大值的倍数才能保持一致。当某个维度不是该有效图块的倍数时,边缘图块仅部分填充,但仍然启动并运行一个完整图块的计算值;未使用的 (填充) 部分没有任何有用的工作,浪费了周期并降低了吞吐量。图 5 说明了这一点:使用 256 × 128 的基础图块、集群 MMA 和 4 × 2 CGA,当 N 是 256 (集群 MMA 为 2 × 128) 或 512 ( CGA 为 2 × 256) 的倍数时,以精细步骤扫描 GEMM-N 会产生本地最大吞吐量。

为避免这种浪费,请选择大小为图块大小倍数的模型尺寸,并与 GPU 缓存行宽度保持一致。一个 128 的倍数表示安全的便携式地板; 256 (clusterMMA) 或 512 (CGA) 的倍数表示更大的协作图块,并可捕获最多的吞吐量。

准则 2:将模型维度的倍数至少设为 128,以与 GPU 图块大小和缓存行宽度保持一致,并首选 256 或 512,以匹配由集群 MMA 和 CGA 组成的更大图块。

与深度模型相比,更广泛的模型对硬件更友好

对于固定参数预算,与更深的模型相比,更宽的模型通过更高的权重重复使用和更短的顺序关键路径提供更高的算术强度和更低的延迟。这使得更广泛的模型有利于实现吞吐量和延迟导向的服务目标。 

也就是说,长宽比也会影响模型质量:深度有助于表现能力,因此有一个有用的广深带,而不是“越宽越好”。只在准确性保持不变的情况下选择宽度,而不是仅为了使模型更宽而分离层。

准则 3: 当有选择时,与许多较小的运算相比,首选更少、更大的运算;这可更大限度地提高算术强度并提高硬件利用率,同时有利于吞吐量和交互性。换句话说,更广泛的 Transformer 模型比更深层次的模型更适合硬件。

量化作为性能杠杆

量化可提高数学吞吐量并立即减少内存流量,从而帮助计算受限和内存受限的工作。Blackwell 系统支持 NVFP4 以及 FP8 和 FP16/ BF16 等其他位宽格式 (请参见下图 6) 。

NVFP4 专为平衡模型准确性和速度而设计;它对每个 16 值微块应用细粒度 FP8 (E4M3) 刻度,并对每个张量应用二级 FP32 刻度。这种分层扩展可在保持 4 位计算速度的同时大幅减少量化错误,使 NVFP4 能够在各种 LLM 工作负载中实现更高精度的准确性 (请参见下图 7,DeepSeek-R1) 。

为此,NVIDIA 提供了端到端工具:TensorRT Model OptimizerLLM Compressor 支持后训练量化 (PTQ) 、量化感知训练 (QAT) 和高级校准,可将模型量化为 NVFP4,同时将准确性损失降至最低。

因此,大多数矩阵乘法运算(例如线性层)都可以充分利用NVFP4获得更高的吞吐量,同时保持模型保真度。如需了解有关使用较低位宽进行训练的更多信息,请查看使用NVFP4预训练大语言模型(NVIDIA Research,2025年)。

准则 4:在网络中引入计算成本高昂的运算时,请考虑是否可以在部署时对其进行量化。设计受益于低精度执行的层是充分利用现代 GPU 的关键。

大型专家并行可提高吞吐量

在以吞吐量为导向的服务中,目标很简单:在最少的 GPU 上以尽可能快的速度为尽可能多的用户提供服务。由于大多数先进的 LLM 现在都是混合专家模型 (MoE) ,因此要做到这一点,一种非常有效的方法是专家并行 (EP) :运行注意力数据并行,并在 GPU 上分配 FFN 专家。 

使用数据并行 (DP) 处理注意力可以避免张量并行 (TP) 的关键限制,而组合部分结果需要使用成本高昂的 AllReduce。这种开销会随着并发性的增加而增加,并拖累吞吐量,这使得 TP 难以满足需求。DP 注意力扩展更自然:添加 GPU 直接增加全局并发率,而更高的并发率意味着 MoE FFN 中的 GEMM-M 更大。 

假设采用统一 token 路由:

\(\text{GEMM-}M = \frac{\text{global concurrency} \times \text{top-}k}{\#\text{experts}}\)

其中全局并发是指每个 GPU 的并发 token × GPU 数量,top-k 是指每个 token 激活的专家数量 (DeepSeek-R1 为 8) ,而 #experts 是指专家总数 (DeepSeek-R1 为 256) 。由于 GEMM-M 随并发增长,而随着模型稀疏而缩小,因此提高并发性是在稀疏 MoE 模型中提高 GEMM 利用率的关键,而且由于每个 GPU 的并发性受 KV 缓存占用的限制,因此在更多 GPU 上扩大 EP 是提高 EP 的主要手段。

即使每个专家的有效批量 (其 GEMM-M) 很小,无论是受限并发还是不平衡的 token 路由,拓宽 EP 仍然在两个方面有所帮助:

  1. 通过聚合带宽加快执行速度:将专家分布在更多 GPU 上可增加用于加载专家权重的有效内存带宽,从而降低端到端 FFN 延迟。 
  2. 减少每个 GPU 的显存占用:每个 GPU 仅存储专家的子集,从而释放内存,提高每个 GPU 的有效并发率,并提高计算的饱和度。

EP 确实带来了两个挑战,即 all-to-all 通信开销和专家负载不平衡,但 NVIDIA 在 TensorRT-LLM 中的 Wide-EP 功能同时解决了这两个问题:

  • 跨 Blackwell 多节点 NVLink 系统扩展 EP 的高性能多对多核函数  
  • 自适应负载均衡器,可根据实时负载模式重新分配专家流量,从而在 token 和专家分布不均的情况下保持稳定。 

有关详情,请参阅 TensorRT-LLM 文档,了解专家并行扩展。对于面向吞吐量的 MoE 部署,关键的手段是专家并行。

准则 5:拥有大量专家的复杂稀疏 MoE 模型在使用正确的并行策略部署时仍然可以提供高吞吐量,最明显的是通过扩展专家并行范围。

工作流并行设计

当您分解预填充和解码时,工作流并行性就会变得相关,这可能是针对某些模型大小、流量模式和延迟目标的有效服务策略,如我们在《Beyond the Buzz:A Pragmatic Take on Inference Disaggregation》(NVIDIA Research,2025 年)中所分析的那样:使用不同的模型分区运行预填充和解码可以让每个模型都单独进行优化和扩展,提高整体吞吐量。

对于预填充,在保持高吞吐量的同时降低初始标记延迟 (FTL) (尤其是在长上下文中) 需要积极的并行化。分块工作流并行 (CPP) 非常适合这种情况:它将跨 GPU 的模型层和输入上下文拆分成流经工作流的数据块 (请参见下图 8) 。

这让预填充工作节点能够在严格的 FTL 预算范围内处理长序列,而无需采用宽张量并行 (请参见下面的图 9) 。

只有在工作流阶段保持平衡的情况下,CPP 才会产生回报;不平衡的阶段会产生工作流泡沫,导致一些 GPU 处于空闲状态。保持场景平衡的一种有效方法是使用均匀分割的可重复图层模式构建模型。  

准则 6: 设计具有常规、可重复层模式的模型,这些模式易于拆分为平衡的工作流阶段。这使得该模型非常适合分块工作流并行,从而实现更高的吞吐量,同时满足严格的首个 token 延迟目标。

混合并行策略,满足面向延迟的服务目标

到目前为止,我们一直专注于吞吐量;面向延迟的服务遵循不同的约束条件。延迟越低,每个 GPU 所需的工作量就越少 (并发性越低) 或需要更多 GPU。但是,缩小批量最终只会降低注意力延迟:FFN 延迟会受到内存限制,因为 GEMM-M 缩小时 GEMM-N 和 GEMM-K 维度仍然很大,使 \(N \times K\) 矩阵的权重读取成为瓶颈。

在这种低并发机制中,注意力和 FFN 应独立并行化。FFN 通过在 GPU 上使用 TP、EP 或 TP+ EP 分配权重来加速。TP 倾向于使用具有大型专家的粗 MoE 模型,而 EP 更适合细粒度 MoE 模型,否则较小的 \(H’\) 会在 TP 下生成低效的小型 GEMM。

相比之下,注意力则受益于 TP 和 KV 并行。但是,TP 只能扩展到 KV 头的数量;除此之外,必须复制 KV 缓存,从而引入冗余工作。由于现代 MQA/ GQA/ MLA 模型的 KV 头很少,这很快就会成为瓶颈。Helix Parallelism ( NVIDIA 研究,2025 年) 通过在注意力过程中对序列维度中的 KV 缓存进行分片,然后在 FFN 中针对 TPxEP 重复使用相同的 GPU,从而克服了这一问题。Helix 是 TensorRT-LLM 支持的。尽管这些混合方案引入了额外的通信,但 Blackwell NVL72 的高带宽 NVLink 结构吸收了大部分成本,而其余部分可以与计算重叠。

指南 7:设计将注意力和 FFN 并行化解的模型和部署,以便通过战略性地针对主要瓶颈来扩展 GPU,从而提高交互性。

后续步骤

在设计下一个模型时,请将以下指南视为设计检查清单: 

  • 保持尺寸接近方形,并对齐至 128 (理想情况下为 256)
  • 有利宽度高于深度
  • 低精度 (NVFP4) 执行设计模型
  • 使用常规、可重复的图层模式,清晰地映射到工作流并行性上

像这样的小选择可以有意义地提高 GPU 利用率,通过相同的硬件提供更快的推理速度、更高的吞吐量和更好的交互性。为了将它们付诸实践,探索用于 NVFP4 量化的 NVIDIA TensorRT Model Optimizer 和用于专家、工作流以及 Helix 并行的 TensorRT-LLM

标签