OCRia 推理栈：单卡 B200 提速近 10 倍

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 Repla