Announcing Carat: An inference engine designed for Gemma 4

Getting started with inference is easy. Rent a few NVIDIA GPUs, run vLLM or SGLang, and you’re good to go. But if you already know which model you’ll run, how much performance are you leaving on the table by using an engine built for everything else too?

We built Carat, a C++/CUDA inference engine for Gemma 4, to find out. Gemma’s fixed sliding window and shared KV heads let us replace page lookups with ring buffers and group attention queries into larger matrix multiplications. Here’s what we learned.

Replacing page lookups with a ring buffer

One of the core things an inference engine needs to optimize is KV cache reuse for attention layers. The KV cache for a session includes intermediate state from the attention calculations, and we store it to avoid recomputing it on every forward pass of the model to generate a new token.

As context windows grow larger, the bandwidth used to transfer the KV cache begins to dominate relative to the model weights, and so KV cache optimization is one of the biggest targets for increasing model throughput.

Unlike standard attention that attends across all previous tokens, Gemma uses a 1024-token sliding window for attention in 50 of its 60 transformer layers.

Building Carat specifically for Gemma lets us carry model-specific decisions through the cache layout, attention kernels and runtime.

SGLang builds indices for the recent window from the request's token mapping, and then translates those into local-cache locations. The attention kernel then loads those indices and reads KV values at the resulting addresses. The slots can be scattered, and paging allows for efficient use of memory rather than reserving the maximum context length in memory.

Carat instead exploits the token-window architecture specific to the model and gives each request a fixed ring for its local history. A token’s absolute position determines its slot modulo the window size.

When the window advances, a new token reuses an old slot without shifting the remaining data.

This, alongside other improvements to how attention is handled in Carat, reduced the time-per-output-token (TPOT) from 35.92 ms to 24.96 ms¹.

More FLOPs can be faster

The sliding-window ring only helps for Gemma’s local layers. The ten global layers still need the full context, and their attention matrices grow with the conversation length.

Looking at inputs from our dataset, Gemma 4’s global query-key multiplication did not scale smoothly with input length. It was spiky. One input with 6,005 keys took 0.830 ms which was longer than expected and global attention layers became a bottleneck. We experimented with padding the K matrix here to see if more FLOPs could improve performance.

The figure shows the last 64 columns of the key matrix. Switching from Exact to Padded adds eleven columns that are not useful in the attention calculation: the preallocated cache tail is zero-initialized and the causal mask gives those columns zero weight, so the output is unchanged.

Rounding the dimensions of K up to the next multiple of 64 and masking the future positions resulted in a 7.9x performance improvement for QK calculations.

const int key_count = context_length + proposal_count; // 6,005

// BF16 matrix multiplication through cuBLASLt.
qk_matmul(queries, key_cache, scores, key_count);

// Each query can attend only up to its own position.
causal_softmax(scores, query_positions, key_count);

pv_matmul(scores, value_cache, output, key_count);

At 6,005 keys, cuBLASLt uses a conventional kernel with ordinary floating-point multiply-add instructions. At 6,016 keys, it used a Hopper Tensor Core kernel with HGMMA instructions, which is built to be significantly more efficient at matrix multiplications.

Adding a few masked columns and increasing slightly the number of FLOPs in the calculation can significantly improve the speed by pushing it to take a different route through the hardware. This matches our experience optimizing inference, counting FLOPs is a poor predictor of execution time, and the fastest configuration is often counter-intuitive.

Doing it all in one kernel can be slower

Up until now, we’ve focused on making the target model’s forward pass faster. Speculative decoding lets the target model check several proposed tokens together in a single forward pass, amortizing the cost of reading weights and loading the KV cache from memory.

Speculative verification sits between ordinary decoding and prompt processing. Instead of attending with one new token or thousands, the target model checks a handful of proposed positions together.

For Gemma's global-attention layers, the SGLang version we tested uses one fused attention kernel for each query head. Each head processes its four query positions against the cached history in chunks, calculating attention scores, applying softmax, and combining values within the same kernel.

Carat instead groups the eight query heads that share each KV head. With four positions per head, each group has 32 query rows to process against the same cached keys and values. Carat performs that work using batched matrix multiplications: one to calculate attention scores, followed by softmax, then another to combine the values.

SGLang keeps the calculation fused and split by head, while Carat groups heads and positions and separates the calculation into matrix operations. At this verification shape, those matrix operations run much faster, even after accounting for the extra memory traffic from storing intermediate scores.

In a controlled global-attention benchmark, SGLang’s production-style attention core took 1.60 ms, compared with 0.099 ms for Carat. Even with the padding optimization from the previous section disabled, Carat took 0.165 ms, 9.7× faster.

Across the complete serving workload, Carat produced 350.8 output tokens/s versus SGLang’s 209.8, a 67.2% increase, while reducing TPOT from 15.22 ms to 7.46 ms.

Design the inference engine around the model

Carat demonstrates that the best performance means designing the engine around the model, all the way down.

Abstractions exist throughout the inference stack, and they serve a purpose - allowing software to be reusable and easy to maintain.

Agents allow us to build much faster than before, and using them we're able to build more from scratch on a per-model basis, reducing the wasted performance lost from unnecessary abstractions.

Building Carat based on that thesis delivers 67% more output tokens per second and 51% lower TPOT than SGLang in our benchmarks.