LlamaIndex 团队在部署其最新文档解析模型 OCRia-1 时发现，推理延迟是客户最关心的问题之一。经过一系列优化，他们成功将单张 B200 GPU 上的 token 生成速度从约 285 tokens/s 提升到约 2800 tokens/s，实现了近 10 倍的改进。这一成果不仅加快了内部基准测试的迭代速度，也让最终用户能够更快获得解析结果。
OCRia 的推理过程主要分为两个阶段：预填充（prefill）和解码（decode）。预填充是计算密集型操作，而解码是内存带宽密集型操作。对于 B200 和该团队的工作负载，解码阶段严重受限于内存带宽——GPU 花费大量时间在移动模型权重和状态上，而非实际计算。因此，优化方向聚焦于减少每步解码移动的数据量、在单次目标模型前向传播中生成更多 token，以及消除推理引擎中的各种开销。
首先，团队采用了工作负载感知的量化策略。他们并未简单地对所有参数进行统一量化，而是根据各层对量化的敏感度使用混合精度：敏感层保留较高精度，其余层则更激进地压缩。通过使用真实输入分布进行激活感知量化校准，并评估不同量化组合的精度与速度，最终在视觉编码器上保持不量化以保护视觉准确度，同时证明量化在目标任务上几乎无损。他们混合使用了 BF16、FP8 和 NVFP4 量化权重。
其次，投机解码是提升吞吐的关键。传统自回归模型每步只生成一个 token，且每一步都需要读取全部模型权重，导致在较小批量下计算资源大量闲置。团队引入了一个小型草稿模型 DFlash，它基于目标模型的隐藏状态特征，通过一次前向传播并行预测多个未来 token。目标模型随后在一次前向传播中并行验证这些 token，并保证接受结果与目标模型分布一致，从而实现数学上无损的加速。最初的 DFlash 接受率为 9 个 token，升级到 DFlash2 后（增加了局部卷积和候选选择器），接受率超过 12 个 token。此时，在批量大小为 8 到 12 时，每用户吞吐已从约 200 tokens/s 提升至 1400 tokens/s 以上。
为进一步挖掘性能，团队深入优化了投机验证环节。OCRia 架构使用了门控 Delta 网络（GDN），它不像全注意力层那样保留 KV 缓存，而是用固定大小的循环状态总结历史。在投机解码中，当草稿 token 被拒绝时，传统方法需要恢复完整状态，造成大量内存搬运。团队适配了 ReplaySSM 技术：保存循环状态的检查点以及近期输入缓冲区，仅在周期刷新时更新完整检查点；当拒绝发生时，只需调整缓冲区指针来保留已接受前缀，避免全状态恢复。随后他们优化了该 kernel，将每层耗时从上游 Triton 实现的 72–80 微秒降低到自定义 CUDA 实现的 23–24 微秒，提速约 3.3 倍。
此外，团队还进行了算子融合以减少不必要的内存传输。例如，将 MLP 中的 SiLU 激活与乘法及其 FP8 量化融合，将门控 RMSNorm 与 FP8 量化融合，并将 NVFP4 激活量化集成到产生激活的 kernel 中。仅 SiLU 与量化融合一项就带来了 2.3%–2.6% 的吞吐提升。同时，他们还为 GDN 层实现了更快的因果卷积。
CPU 端的准备工作也被证明是延迟的重要来源。推理引擎在每个解码步骤前都需要构建描述批次的元数据（序列长度、token 偏移、缓存位置等）。由于 vLLM 按 KV 缓存组组织层，而目标模型和草稿模型共有四种不同的缓存布局（24 个 GDN 层、8 个目标注意力层、5 个草稿滑动窗口层和 1 个草稿全注意力层），分组规则导致生成了 38 个组，每组都触发一次元数据构建，尽管大量信息是共享的。团队改为对兼容的构建器类和缓存规格只构建一次共享元数据，然后为其他组提供带有自身块表和槽映射的浅拷贝。对于 GDN 元数据，他们将与组无关的计算移到 NumPy 上，并利用草稿请求位于批次前几行的特点简化掩码，最终将每步约 290 次小型 GPU 启动减少到仅 1 次。
最后，团队还实施了另外 17 项优化，包括针对模型架构定制的 Fused Triton CUDA kernel，以及选择特定的现有 kernel 和 cuDNN 组件。每项微优化单独只能带来 2%–5% 的吞吐提升，但叠加起来效果显著。这些优化共同构成了 OCRia 高效推理栈的基础，并且部分改进（如 ReplaySSM 适配）仍以草稿 PR 形式贡献给 vLLM 社区。