团队优化OCRia-1推理栈，单卡B200生成速度提升近10倍

在训练和部署我们最新的文档解析模型 OCRia-1 时，我们注意到了一个反复出现的瓶颈：推理速度不够快。这与客户的体验一致——在任何文档解析工作流中，延迟始终是首要诉求之一。

于是我们着手进行调优，最终在单块 B200 GPU 上实现了 token 生成速度近 10 倍的提升，从约 285 tokens/秒提升到约 2800 tokens/秒。

这意味着我们可以更快地迭代基准测试，并更快地向用户交付结果。我们通过组合大量较小的改进实现了这一目标：减少推理过程中移动的字节数、在每次目标前向传播中生成更多 token，以及降低每个步骤的整体 GPU 和 CPU 开销。其中一些改进借助成熟实践（量化、投机解码）轻松实现，而另一些则要求我们深入挖掘实现细节。

OCRia 主要由图像输入加文本提示驱动，输出文档的结构化文本表示。内容密集的页面最终需要更长的输出才能正确表示。因此，我们的优化重点转向提升单块 GPU（此处为 B200）上并发请求的生成速度。

LLM 的推理主要包含两个阶段：prefill 和 decode。在 prefill 阶段，输入被处理并构建内部模型状态以生成响应。在 decode 阶段，这些输入通过模型传递，输出则是一系列“步骤”，每一步采样下一个 token（自回归解码）。这里的区别在于，prefill 是计算受限操作，而 decode 是内存受限操作。

Williams 等人的 roofline 模型捕捉了这两种操作之间的平衡。根据内存带宽和计算能力，性能要么受限于内存，要么受限于计算。对于我们的 B200 GPU 和我们的工作负载，我们主要处于内存带宽受限区域。

在我们的批次大小下，GPU 在每个步骤中花费过多时间移动模型权重和状态，而非实际数学计算。这些约束帮助我们明确了有效的优化方向：减少每个 decode 步骤中移动的数据、每次目标模型前向传播生成更多 token，以及尽可能从推理引擎中移除开销。

量化是 AI 社区中的常见实践。新模型发布后，几小时内就会出现量化版本。特别是，量化可以对内存使用产生重大影响。从 BF16 → FP8 将每个模型权重的大小从 2 字节减少到 1 字节。这减少了 GPU 需要移动的数据量，并且在使用支持低精度计算的硬件和内核时，还可以提高整体计算吞吐量。

简单地对模型中每个参数应用相同的量化可能会损害性能。模型的不同层对量化的敏感度不同，因此我们使用了混合精度量化，将敏感层保持在高精度，同时更积极地压缩其他层。

为了保持准确性，我们使用激活感知量化，它考虑权重与真实输入产生的激活之间的交互，而非孤立地简单舍入权重。为此，我们使用包含分布内数据的校准集。这有助于根据模型的预期用途调整缩放因子，并减少量化对其输出的影响。我们混合使用了 BF16、FP8 和 NVFP4 量化权重。

最后，我们通过评估来比较不同的量化层和精度组合，在留出样本上衡量推理速度和准确性。针对我们的具体情况，我们选择不对视觉编码器进行量化以保持视觉准确性，并通过严格的评估确保量化在我们的目标任务上实际上是无损的，同时提高了推理吞吐量。

如需深入了解训练后量化，我们推荐阅读 https://developer.nvidia.com/blog/optimizing-llms-for-performance-and-accuracy-with-post-training-quantization/ 上的文章。

像 OCRia 这样的自回归模型每步生成一个 token，每个新 token 都需要对模型进行一次完整的前向传播才能生成。即使 KV-cache 避免了对先前 token 的重新计算，我们仍然需要在每个步骤读取模型权重。在小批次大小下，这意味着移动大量数据只为了生成一个 token，而大量可用计算资源未被使用，从而将生成速度降低到数据在 GPU 上来回移动的速度。

投机解码是一种旨在改变这种动态的技术，它添加一个较小的第二模型（drafter），提前生成（推测）若干未来 token。然后我们使用原始模型在单次前向传播中并行验证这些 token。使用正确的验证程序可确保接受的 token 与目标模型的输出分布匹配，丢弃不匹配目标分布的 token。这一点很重要，因为这从数学上保证了推测是 100% 无损的。

我们使用了 DFlash，这是一个轻量级块扩散模型，使用小型基于 transformer 的扩散模型在单次前向传播中起草多个 token。它以从目标模型提取的隐藏状态特征为条件进行预测，利用目标模型对上下文的表示来提出一块未来 token 以供验证，即对下一组 token 进行“推测”。

我们最初训练的 DFlash 模型实现了 9 个 token 的接受率。当我们转向 DFlash2 时，进一步推高到超过 12 个 token 的接受率。DFlash2 增加了局部卷积和候选选择器以改进起草预测，同时保留了并行的起草模型前向传播。我们复用了现有的 DFlash1 权重来初始化 DFlash2 模型，并在此基础上进行训练。

（原文较长，此处译出前 6000 字，完整内容请点「打开原文」。）