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

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

The Inference Stack Behind OCRia

AI 导读

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

正文 · 原文

While training and deploying our latest document parsing model, OCRia-1, we noticed a recurring bottleneck: inference was not fast enough. And this lined up with our experience customers had as well, latency is always one of the top asks for any document parsing workflow.

So, we set out to work tinkering, and came back with a nearly 10x improvement in token generation, going from ~285 tokens/sec to ~2800 tokens/sec on a single B200 GPU.

This means we can iterate on benchmarks faster and deliver results faster to users. We accomplished this by combining a ton of smaller improvements: reducing the bytes moved during inference, producing more tokens per target forward pass, and reducing the overall GPU and CPU overhead of each step. Some of these were easy wins with established practices (quantization, speculative decoding), while others required us to dig very deep into implementation details.

Why decoding dominates latency

OCRia is largely driven by image inputs with a text prompt, and outputs a structured textual representation of the document. Pages that are dense end up requiring longer outputs to represent them properly. Our primary focus for optimization then became increasing our generation speed across concurrent requests on a single GPU (B200 in this case).

Inference in an LLM largely consists of two phases: prefill and decode. During prefill, inputs are processed and internal model state is built to generate a response. During decode, those inputs are passed through the model, and the output is then a series of “steps” where the next token is sampled (autoregressive decoding). The difference here is that prefill is a compute bound operation, whereas decode is a memory bound operation.

The roofline model from Williams et al. captures the balance between these two operations. Performance is either going to be memory bound or compute bound, depending on memory bandwidth and compute capacity. For our B200 GPUs and our workload, we are largely in the memory-bandwidth-bound region.

At our batch sizes, the GPU spends too much time each step moving model weights and state around vs. actual math. These constraints helped shaped the optimizations that worked: reducing data moved at each decode step, generating more tokens per target-model pass, and removing as much overhead as possible from the inference engine.

Reducing the data moved: workload-aware quantization

Quantization is a common practice in the AI community. New models come out, and within hours, quantized versions appear. In particular, quantization can have a large impact on memory usage. Going from BF16 → FP8 reduces the size of each model weight from 2 bytes to 1 byte. This reduces the amount of data the GPU needs to move and, when using hardware and kernels that support lower-precision computation, can also increase overall compute throughput.

Naively applying the same quantization to every parameter in the model can hurt performance. Different layers of the model have different sensitivities to quantization, so we used mixed precision quantization, keeping sensitive layers at higher precision while compressing the others more aggressively.

To preserve accuracy, we use activation-aware quantization, which considers how the weights interact with the activations produced by real inputs, rather than simply rounding weights in isolation. For this we use a calibration set with in-distribution data. This helps tune the scaling factors to the model’s intended use and reduce the impact of quantization on its outputs. We used a mix of BF16, FP8, and NVFP4 quantization weights.

Finally, we use evaluation to compare different combinations of quantized layers and precisions, measuring both inference speed and accuracy on held-out examples. For our specific case we choose not to quantize the vision encoder to preserve visual accuracy, and through rigorous evals, ensured that our quantization was effectively lossless on our target taskswhile increasing inference throughput.

For a deeper read on post-training quantization, we recommend reading the post from https://developer.nvidia.com/blog/optimizing-llms-for-performance-and-accuracy-with-post-training-quantization/

Generate more tokens per target pass: speculative decoding

Autoregressive models like OCRia produce one token per step, where each new token requires a full forward pass over the model to generate it. Even with the KV-cache avoiding re-computation of previous tokens, we still need to read the model weights for every step. At small batch sizes this means moving a large amount of data to produce just one token, while much of the available compute remains unused, reducing the speed of the generation to how fast we can move data back and forth across the GPU.

Speculative decoding is a technique that aims to change that dynamic by adding a second smaller model (drafter) that produces (speculates) several future tokens in advance. Then we use the original model to verify those tokens in parallel during a single forward pass. Using the correct validation procedure ensures that the accepted tokens match the target models output distribution, discarding the ones that does not match the target distribution. This is important because this creates a mathematical guarantee that speculation is 100% lossless.

We used DFlash, a lightweight block-diffusion model that drafts multiple tokens in a single forward pass using a small transformer based diffusion model. It conditions predictions on hidden-state features extracted from the target model, using the target’s representation of the context to propose a block of future tokens for verification, “speculating” on the next set of tokens.

We originally trained a DFlash model that achieved an acceptance rate of 9 tokens. We pushed this even further when we moved to DFlash2, with acceptance rates of over 12 tokens. DFlash2 adds local convolutions and a candidate selector to improve draft predictions while retaining a parallel draft-model forward pass. We re-used the existing DFlash1 weights to initialize our DFlash2 model and trained from there.

At this point in our optimization journey we were producing more than 1400 tokens per second per user at batch 8 to 12, whereas the base model was around 200 tokens per second. The question then became, can we improve it even more? What would be needed to improve it without resorting to more aggressive quantization or model pruning to avoid losing quality?

Removing remaining overhead

For the next set of improvements, we had to dig pretty deep into some complex implementation details. So bear with as we start stacking on some solid wins during token generation.

Make speculative verification cheaper

A popular inference optimization technique is applying custom kernels, like FlashAttention, that optimize specific operations on your GPU during inference. In our case, we noticed that the ReplaySSM technique from dao-lab was well suited to our needs around optimizing inference in our stack. But to understand the technique, we need to cover some key pieces of background.

OCRia uses Gated Delta networks (GDN) in its architecture. Unlike full attention layers, which retain previous keys and values in a KV-cache, GDN layers summarize the history in a fixed-size recurrent state. During normal decode, each token requires reading that state from HBM, updating it and writing it back. The computation is relatively small, but moving the state is expensive and the speed is limited by the memory bandwidth. When speculated tokens are rejected, this means a lot of data has to move around to revert states.

ReplaySSM keeps a checkpoint of the recurrent state alongside a small buffer of recent inputs. It computes outputs from the checkpoint and buffered inputs, updating the full checkpoint only at periodic flushes. When speculative tokens are rejected, rollback adjusts the buffer pointer to retain the accepted prefix, avoiding a full-state restore.

We adapted ReplaySSM (to date still as a series of draft PRs to vLLM) to the GDN layers under speculative decoding, then profiled the resulting inference step. At 32 concurrent requests, a decode step took 6.07 ms, and the ReplaySSM verification and flush kernels accounted for 2.10 ms across the model’s 24 GDN layers. Reducing state traffic was the right direction, but the implementation of that operation had become a major part of the remaining latency.

We then optimized the kernel itself. Starting from the upstream Triton implementation at 72–80 µs per layer. Then a custom CUDA implementation brought it down to 23–24 µs, roughly 3.3× faster than the original kernel in the same benchmark.

Avoiding unnecessary memory transfers

Another opportunity was kernel fusion. When consecutive operations run as separate kernels, their intermediate results are normally written to HBM and immediately read back by the next operation. By combining them, we can keep those intermediate results in GPU registers or shared memory avoiding expensive trips to the HBM.

We applied this to the MLP’s SiLU-and-multiply activation and its FP8 quantization, and also to the gated RMSNorm, followed by FP8 quantization. We also integrated NVFP4 activation quantization into the kernels producing those activations and implemented a faster causal convolution for the GDN layers. The SiLU-and-quantization fusion alone improved throughput by 2.3–2.6% in our measurements.

Fixing CPU-side preparation

As we profiled deeper, we realized a large part of the latency was coming from the CPU-side preparation of each decode step. Before launching the attention and GDN kernels, the inference engine prepare metadata describing the batch: sequence lengths, token offsets, cache locations and the state indices used for speculation. This information changes every step, so it needs to be updated as requests advance, finish or arrive.

The problem was how often we were computing the same information. vLLM organizes layers into KV-cache groups and builds metadata for each group. With our target and draft models, there were four different cache layouts: 24 GDN layers, 8 target attention layers, 5 draft sliding-window layers and 1 draft full-attention layer. The grouping rule used the smallest bucket as the group size, which in this case was one. That left us with 38 groups, each triggering a metadata build, even though much of the information was shared.

We changed this to build the shared metadata once per compatible builder class and cache specification, then give the other groups a shallow copy with their own block table and slot mapping. The metadata still gets updated every step, but the common work is no longer repeated for every layer. We also applied this sharing to the FlashInfer decode wrappers.

The GDN metadata build itself needed attention. Repeating it across 24 groups was issuing around 290 small GPU launches per step, alongside synchronization and copying overhead. We moved the group-independent calculations to NumPy using arrays already available on the CPU. Since drafted requests occupy the first rows of the batch, their masks could become simple prefix slices, and some metadata tensors could reference existing buffers directly. Together, these changes reduced the GDN metadata work to one GPU launch per step.

All the small stuff

Altogether, we ended up implementing 17 more optimizations, ranging from new Fused Triton CUDA kernels specific to our model architecture, to selecting specific existing kernels and cuDNN components. Each micro optimization alone could only contribute a small 2% to 5% increase in throughput but composed together they contributed more than 60% in additional throughput and latency reduction.

Our methodology was to observe and do detailed profiling to identify small bottlenecks and opportunities looking for un-fused operations. Then we looked for any engine level knob we could tune, and if it does not exist, we implement it or modify the engine source code. Lastly, we measure and benchmark carefully to guarantee numerical equivalence in every operation before looping back and doing it all over again.

These optimizations work together with quantization and speculation: we reduced the bytes moved, produced more tokens per target forward pass, and reduced the GPU and CPU overhead of each step. By the end, this composed together to bring us from ~285 tokens/sec to ~2800 tokens/sec, an improvement of almost 10x.

Available today and tomorrow

Model inference optimization is not a process done in isolation. We need to take into account the specifics of the model architecture and the traffic patterns we expect for production inference so we can make the correct tradeoff between throughput and interactivity, and a lot of observations and measurement to guarantee we are not losing mode capabilities.

For the optimizations that apply generally, we plan to contribute back to vLLM to give back to the community, and enable optimizations that work generally across models.

In the meantime, you can use OCRia-1 today on LlamaParse, powering all our tiers.

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