利用 NVIDIA Blackwell Ultra 提升 Softmax 效率

LLM 的上下文长度正在迅速增长，架构也在转向复杂的注意力方案，例如 Multi-Head Latent Attention (MLA) 和 Grouped Query Attention (GQA)。因此，AI 的“思考速度”越来越多地不是由矩阵乘法的巨大吞吐量所决定，而是由 softmax 函数的超越数学运算所支配。

超越函数指的是无法表示为有理系数多项式方程之根的函数。因此，它们“超越”了加法和乘法等基本代数运算——而这些正是 Tensor Cores 擅长的运算。在 softmax 的具体语境中，这些超越函数里计算开销最大的是在 Special Function Units (SFUs) 上执行的自然指数函数。在 NVIDIA 汇编指令（SASS）中，该函数通过 MUFU.EX2 指令调用。这种架构上的分离在注意力模块内造成了 softmax 瓶颈：强大的矩阵引擎被迫空闲，等待 SFU 数据通路对注意力分数进行归一化。

NVIDIA Blackwell Ultra 通过将 SFU 吞吐量提升至标准 NVIDIA Blackwell 架构的两倍，缓解了这一瓶颈。

本文深入探讨注意力循环中 softmax 的机制，介绍 Blackwell Ultra 的硬件优化如何消除流水线停顿，并提供一个基准测试，让你亲自衡量原始 MUFU.EX2 的加速效果。

注意力如何工作

现代大型语言模型的一个基础组件是注意力机制，它使模型能够将静态 token 向量动态转换为具有上下文感知能力的动态表示。其核心是一个通过让 token 相互调整重要性来重新加权信息的过程。为了促进这种交互，序列中的每个 token 都会被投影为三种功能角色：

• Query：表示当前 token 为理解自身上下文而正在寻找的内容。
• Key：表示一个 token 的特征画像，供其他 token 用于匹配。序列中位于前面的 token 具有一些 key，用来表明它们与 query 的具体相关性。
• Value：它保存实际的信息内容。一旦 query 与 key 之间的匹配得到确认，Value 就是传递给原始 token 的具体数据。

下方图 1 展示了 attention 的实际运作。这里有两个句子，它们以两种不同的定义使用了 “dog” 一词。起初，我们可以看到两个 “dog” 提及的 embeddings（在多维空间中捕捉含义和细微差别的数值向量）是相同的。

注意力机制的运作方式是：模型计算“dog”查询向量与序列中所有其他标记的键向量之间的点积。

如果“dog”的查询向量与“lazy”的键向量高度匹配，就表明二者具有高度相关性。正是这种交互使得“dog”这个词能够引入其相邻词的特定值。到这一轮结束时，“dog”的原始向量已经在物理上被其相邻词的内容更新，从一个通用的词典定义演变为一个上下文化嵌入，能够“理解”它指的是一只嗜睡的动物，还是一年中酷热的高峰时段。

softmax 与注意力的关系

Softmax 是关键的决策阶段，它将原始兼容性分数转换为可操作的权重。一旦查询向量与键向量之间的初始点积计算完成，所得分数就会通过 softmax 函数，被归一化为总和恰好为 1 的概率。这一步决定了模型的“注意力范围”，实际上是在决定应优先关注哪些标记、忽略哪些标记。没有 softmax，模型就无法客观地权衡其收集到的信息，从而导致数据混合变得难以管理且充满噪声。

然而，softmax 操作是长上下文 AI 中出现“性能悬崖”的主要来源。由于序列中的每个 token 都必须与其他每个 token 进行比较，一个包含 8,192 个 token 的序列会产生一个巨大的 [8,192 x 8,192] 注意力矩阵。对该矩阵进行归一化需要数十亿次超越函数计算，并且计算量会随序列长度呈二次方增长。这会形成一个瓶颈，其中大量的超越数学运算可能使整个推理流水线停滞。

Blackwell Ultra 专注于加速这些指数计算，专门用于缓解这一数学瓶颈，并确保系统能够处理大上下文窗口所需的大规模归一化，同时不牺牲吞吐量。

缓解 Blackwell Ultra 中的 softmax 瓶颈

通过在 Blackwell Ultra 架构中将 SFU 的指数运算吞吐量提高一倍，NVIDIA 正在缓解这一瓶颈，并使处理流水线更加均衡高效。这会带来更快的整体性能，尤其适用于高度依赖注意力机制的任务。

下图 2 展示了标准注意力机制中固有的顺序依赖性，该机制通常称为注意力循环，并在上一代 NVIDIA Blackwell（GB200）上运行。请注意，流式多处理器（SM）会加载两个并发运行注意力循环的线程块。这些独立的注意力循环用两种不同深浅的绿色表示。

该流水线由三个必须按顺序执行的不同阶段组成：

• BMM1（分数计算）：Tensor Cores 执行矩阵乘法，以计算原始注意力分数，即 logits。
• Softmax（归一化）：流水线切换到 SFU，使用指数函数将这些分数归一化为概率。
• BMM2（上下文聚合）：流水线返回 Tensor Cores，将概率与值向量相乘。

该时间线展示了 Blackwell GPU 在执行注意力内核期间固有的延迟约束。由于第二次矩阵乘法（BMM2）作用于 softmax 的输出，因此必须等到归一化完成后才能开始。

Blackwell GPU 的 SFU 吞吐量较低，迫使 Tensor Cores 在分数计算（BMM1）和上下文聚合（BMM2）之间处于空闲状态。这种依赖关系阻止流水线充分利用计算资源，并延长了 softmax 操作的持续时间

下一个时间线如图 3 所示，展示了 Blackwell Ultra GPU 在 NVIDIA GB300 NVL72 和 NVIDIA HGX B300 系统中对同一指令序列将 SFU 吞吐量翻倍所产生的直接影响。

从视觉上看，softmax 块的宽度减少了近 50%，反映出硬件能够以两倍速率处理 MUFU 指令。

softmax 延迟的降低收紧了整个流水线。BMM1 与 BMM2 之间的间隔被大幅缩小，使 Tensor Cores 能够在查询-键乘法与概率-值乘法之间切换时仅发生最小程度的停顿。其结果是主循环更加紧凑，高性能矩阵引擎在总执行时间中处于活跃状态的占比更高，直接转化为更高的整体推理吞吐量。

MUFU.EX2 性能基准测试

为了通过实证方式验证 MUFU 流水线的理论吞吐量，我们可以构建一个合成微基准测试。以下内核代码隔离了指数指令，以便在不受全局内存延迟或其他算术运算干扰的情况下测量原始周期数。

该测试框架会启动一个线程网格，其中每个线程都会执行一个密集的 MUFU.EX2 指令循环。通过对执行过程计时并将其与时钟频率进行比较，你可以直接计算有效指令吞吐量，并验证前面提到的带宽饱和点。

步骤 1：克隆以下仓库以获取 exp2-bg300.cu 基准测试。

git clone https://github.com/jamieliNVIDIA/mufu_ex2_bench.git
cd mufu_ex2_bench

步骤 2：使用以下方式编译（GB300 使用 sm100f，GB200 使用 sm103a）。

nvcc -O3 -gencode=arch=compute_103a,code=sm_103a --extended-lambda -o /tmp/exp2-gb300.out exp2-gb300.cu

样本结果

我们看到，在所有测试的数据类型中，GB300 的 FLOPs 性能约为 GB200 的 2 倍，这与 SFU 吞吐量翻倍一致。

正文：Blackwell（GB200）

exp2 BF16x2 2454 Gop/s (4908 GFLOPS)
exp2 BF16 4938 Gop/s
exp2 FP32 4943 Gop/s

正文：Blackwell Ultra（GB300）

exp2 BF16x2 4996 Gop/s (9992 GFLOPS)
exp2 BF16 9738 Gop/s
exp2 FP32 Time:  10024 Gop/s

Blackwell 与 Blackwell Ultra 中的注意力前向传播性能对比

从 Blackwell 过渡到 Blackwell Ultra 带来了有针对性的计算吞吐量提升，这一提升由 SFU 性能提升 2 倍所驱动。此次硬件升级直接加速了 DeepSeek-V3 等模型的前向传播（FPROP）流水线。

FPROP 是输入数据在神经网络中“向前”传递的过程——从输入层经过隐藏层到达输出层——以生成预测。每当模型生成一个新词时，都必须运行一次完整的 FPROP。

下图 4 显示，通过将 SFU 的吞吐量提升一倍，GB300 大幅缩短了注意力块中 softmax 层的执行时间。这种更快的归一化意味着 GPU 花在处理注意力分数上的时间更少，而将更多时间用于利用高速矩阵引擎进行下一层计算，从而直接提高前向传播的整体速度。

基准测试结果显示，FP8 运算的 FPROP 吞吐量提升约 35%。这一增益在 FP8 中尤为明显，因为矩阵数学运算已经极其快速。在这种低精度场景下，softmax 所花费的时间在总步骤时间中所占比例更大。

入门

DeepSeek-V3 在 Blackwell Ultra 上的性能动态凸显了推理中的一个关键但常被忽视的瓶颈：非线性运算的计算成本。

通过优化和压缩注意力机制，最先进的模型有效提高了 softmax 运算相对于标准线性计算的密度，从而暴露出 SFU 成为总吞吐量的制约因素。

Blackwell Ultra 直接解决了这一瓶颈。通过将这些专用单元的吞吐量提高一倍，Blackwell Ultra 打通了此前迫使强大的 Tensor Cores 闲置的超越函数拥堵。基准测试结果证实了这一影响，显示 FP8 前向传播提升了 35%。

对于现代高度优化的架构而言，实现更快推理的路径并不只是更快的 Tensor Cores，还在于确保非线性数学单元足够快，能够跟上整体节奏。

访问 NVIDIA 的 trtllm-gen 仓库，了解更多基准测试以及如何在工作负载中利用这一 SFU 加速的信息。将 MUFU.EX2 的 SFU 吞吐量提升一倍，只是 Blackwell Ultra 实现快速注意力速度的众多特性之一。NVIDIA 极致的硬件-软件协同设计通过以下技术加速了完整的注意力循环：

• 通过 LDTM.STAT 将关键的“find-max”归约卸载到 Tensor Memory 控制器。
• 使用 CUDNN 优化性能。
• 使用 NVFP4 优化 KVCache 数据移动。

敬请关注 NVIDIA 技术博客，获取未来发布的文章。

致谢

特别感谢 cuDNN 工程团队创建基准测试并构建软件优化，使这一前沿性能成为可能。