LlamaIndex 披露 OCRia-1 推理栈：单张 B200 上 token 生成提速近 10 倍

LlamaIndex：产品、工程与评测·2026-11-04 00:00· 32 分钟前AI 评分44

OCRia 背后的推理栈

AI 导读

LlamaIndex 披露 OCRia-1 文档解析模型的推理优化方案，在单张 B200 GPU 上把 token 生成速度从约 285 tokens/sec 提升到约 2800 tokens/sec，接近 10 倍。方案结合了混合精度量化（BF16、FP8、NVFP4，视觉编码器不量化）与 DFlash2 块扩散投机解码，后者接受率超过 12 个 token。

正文 · 原文

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

于是，我们开始动手折腾，最终在 token 生成上实现了接近 10 倍的提升，在单张 B200 GPU 上从 约 285 tokens/sec 提升到约 2800 tokens/sec。

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

为什么解码主导延迟

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

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

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

在我们的批大小下，GPU 每一步花费过多时间搬运模型权重和状态，而非实际数学计算。这些约束帮助塑造了有效的优化方向：减少每个解码步骤搬运的数据、每次目标模型前向传播生成更多 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/ 的文章。

每个目标前向传递生成更多 token：投机解码

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

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

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

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

在我们的优化进程中，此时在 batch 8 到 12 下，我们每用户每秒可生成超过 1400 个 token，而基础模型约为每秒 200 个 token。接下来的问题是，我们还能不能进一步提升？要在不诉诸更激进的量化或模型剪枝、以免损失质量的前提下提升，需要做些什么？

移除剩余开销

对于下一组改进，我们必须深入挖掘一些复杂的实现细节。所以请耐心听我们逐步叠加 token 生成过程中一些扎实的收益。

让投机验证更便宜

一种流行的推理优化技术是应用自定义 kernel，例如 FlashAttention，它在推理期间优化 GPU 上的特定操作。在我们的场景中，我们注意到来自 dao-lab 的 ReplaySSM 技术非常契合我们优化技术栈推理的需求。但要理解这项技术，我们需要先介绍一些关键背景。

OCRia 的架构中使用了 Gated Delta networks（GDN）。与在 KV-cache 中保留先前键和值的完整注意力层不同，GDN 层以固定大小的循环状态概括历史。在正常解码过程中，每个 token 都需要从 HBM 读取该状态、更新它并写回。计算量相对较小，但移动状态的开销很大，速度受内存带宽限制。当投机 token 被拒绝时，这意味着需要移动大量数据来还原状态。

ReplaySSM 在保留循环状态检查点的同时，还保留一个近期输入的小缓冲区。它从检查点和缓冲输入计算输出，仅在周期性 flush 时更新完整检查点。当投机 token 被拒绝时，回滚会调整缓冲区指针以保留已接受的前缀，从而避免全状态恢复。

我们将 ReplaySSM（迄今为止仍是一系列提交给 vLLM 的 draft PR）适配到投机解码下的 GDN 层，然后对得到的推理步骤进行了性能分析。在32 个并发请求下，一次解码步骤耗时 6.07 ms，而 ReplaySSM 的验证和 flush kernel 在模型的 24 个 GDN 层中共占 2.10 ms。减少状态流量是正确的方向，但该操作的实现已成为剩余延迟的主要部分。

随后我们优化了 kernel 本身。从上游 Triton 实现的每层 72–80 µs 起步。然后自定义 CUDA 实现将其降至 23–24 µs，在同一基准测试中比原始 kernel 快约 3.3 倍。

避免不必要的内存传输

另一个机会是kernel 融合。当连续操作作为独立 kernel 运行时，它们的中间结果通常会被写入 HBM，并立即被下一个操作读回。通过将它们合并，我们可以将这些中间结果保留在 GPU 寄存器或共享内存中，避免昂贵的 HBM 往返。

我们将此应用于 MLP 的 SiLU-and-multiply 激活及其 FP8 量化，也应用于门控 RMSNorm，随后进行 FP8 量化。我们还将 NVFP4 激活量化集成到产生这些激活的 kernel 中，并为 GDN 层实现了更快的因果卷积。仅 SiLU-and-quantization 融合一项，在我们的测量中就将吞吐量提升了 2.3–2.6%。

修复 CPU 侧准备工作

随着我们更深入地分析，我们意识到很大一部分延迟来自每个解码步骤的 CPU 侧准备工作。在启动 attention 和 GDN kernel 之前，推理引擎会准备描述批次的元数据：序列长度、token 偏移、缓存位置以及用于推测的状态索引。这些信息每一步都会变化，因此需要随着请求的推进、完成或到达而更新。

问题在于我们计算相同信息的频率过高。vLLM 将层组织为 KV-cache 组，并为每个组构建元数据。对于我们的目标模型和草稿模型，存在四种不同的缓存布局：24 个 GDN 层、8 个目标 attention 层、5 个草稿滑动窗口层和 1 个草稿全注意力层。分组规则使用最小的桶作为组大小，在本例中为一。这导致我们得到 38 个组，每个组都会触发一次元数据构建，尽管其中大量信息是共享的。

我们将其改为每个兼容的构建器类和缓存规范只构建一次共享元数据，然后为其他组提供带有各自块表和槽映射的浅拷贝。元数据仍然每一步都会更新，但公共工作不再为每一层重复执行。我们还将这种共享应用到了 FlashInfer 解码包装器上。

GDN 元数据构建本身也需要关注。在 24 个组中重复执行它，每一步会发出大约 290 次小型 GPU 启动，此外还有同步和复制开销。我们将与组无关的计算移至 NumPy，使用 CPU 上已有的数组。由于草稿请求占据批次的前几行，它们的掩码可以变成简单的前缀切片，并且一些元数据张量可以直接引用现有缓冲区。这些改动共同将 GDN 元数据工作减少到每一步一次 GPU 启动。

所有小事项

总而言之，我们最终又实现了 17 项优化，从针对我们模型架构的新 Fused Triton CUDA kernel，到选择特定的现有 kernel 和 cuDNN 组件。每一项微优化单独只能带来 2% 到 5% 的吞吐量提升，但组合在一起，它们贡献了超过 60% 的额外吞吐量提升和延迟降低。

我们的方法是观察并进行详细分析，以识别小的瓶颈和机会，寻找未融合的操作。然后我们寻找任何可以调整的引擎级旋钮，如果不存在，我们就实现它或修改引擎源代码。最后，我们仔细测量和基准测试，以确保每个操作的数值等价性，然后循环回去再次重复整个过程。

这些优化与量化和推测协同工作：我们减少了传输的字节数，每次目标前向传播生成了更多 token，并降低了每一步的 GPU 和 CPU 开销。最终，这些组合起来让我们从约 285 tokens/秒提升到约 2800 tokens/秒，提升近10 倍。

今天和明天可用

模型推理优化并非孤立进行的过程。我们需要考虑模型架构的具体细节以及生产推理中预期的流量模式，以便在吞吐量和交互性之间做出正确的权衡，并通过大量观察和测量来确保我们不会损失模型能力。

对于普遍适用的优化，我们计划回馈 vLLM 以回馈社区，并支持可跨模型普遍适用的优化。

与此同时，你今天就可以在 LlamaParse 上使用 OCRia-1，它为我们所有层级提供支持。

来源：LlamaIndex：产品、工程与评测 · llamaindex.ai