rvLLM: The Inference Engine That Lets TPUs Compete by Writing Less Kernel Code

A deep dive into the project that pairs Rust orchestration, handcrafted CUDA, and a flat JAX scan to turn Gemma 4 inference into a compiler problem instead of a kernel-burden problem.

8 to 10 min read • View on GitHub • More from m0at

By Repo Explainer • AI-assisted • Published by Ron Kurti

Published Apr 19, 2026 • Local LLM Inference

A split workshop scene shows two very different approaches to inference. On the left, a dense CUDA machine shop is crowded with gears, launch points, and a technician tightening bolts. On the right, a clean TPU control room carries a single compiled loop through a straight path with a few branching gates. It explains that rvLLM uses manual kernel work on GPU but compiler-shaped structure on TPU. — rvLLM’s surprise is not that it uses fast kernels. It is that it treats TPU speed as a compiler design problem, while the GPU path stays firmly in the world of handcrafted CUDA.

rvLLM: High-performance LLM inference in Rust. Drop-in vLLM replacement.

m0at, Top Contributor · m0at/rvllm

Key Takeaways

rvLLM’s most unusual move is to make TPU inference smaller and more compiler-shaped, not more kernel-heavy.
The project splits its personality by backend: Rust orchestration and CUDA specialization on GPU, JAX and XLA structure on TPU.
Its TPU story centers on a flat scan loop that keeps the whole Gemma 4 pass inside one compiled region.
rvLLM’s comparison with vLLM is really about philosophy, because it optimizes for backend-specific performance rather than universal abstraction.

The surprising part is the TPU

Most inference engines get faster by adding more custom code. rvLLM does the opposite on TPU. It shrinks the implementation down, leans on jax.lax.scan and XLA fusion, and lets the compiler do the work that would otherwise become a pile of handwritten kernels.

That is the thesis that makes the repo interesting. The GPU side is a dense CUDA system with graph capture and fused kernels. The TPU side is almost austere. Same project, same goal, two opposite execution philosophies.

Why this engine exists at all

rvLLM is built for the boring but decisive question behind every serving stack: how many tokens do you get for the money? The repository frames itself as a drop-in vLLM replacement, but its real pitch is narrower and sharper. It wants to win on throughput, latency, and cost when the hardware and model shape line up in its favor.

A two-speed architecture: Rust on the outside, kernels on the inside

The repo is split the way the thesis suggests. Rust crates handle orchestration, runtime control, and the outer engine. CUDA handles the GPU hot path. TPU gets its own JAX-based path, not a thin wrapper over the GPU implementation.

rvLLM’s codebase is not one stack with one optimization strategy. It is a control plane in Rust and two different execution philosophies underneath it.

v3/
  crates/
    rvllm-core/
    rvllm-runtime/
    rvllm-graph/
    rvllm-kernels/
    rvllm-cutlass/

tpu/
  harness/
  ports/

kernels/
  *.cu

The flat-scan trick that makes TPU inference work

This is the centerpiece. In tpu/harness/gemma4_tpu_infer.py, rvLLM uses a single scan loop to step through Gemma 4’s 60 transformer layers. Fifty layers use sliding-window attention. Ten use global attention. A conditional inside the loop picks the right path without ever bouncing back to the host.

A close-up view of a single looped assembly line shows 60 stations arranged in sequence. Most stations are identical, but some are marked as a different attention mode and branch briefly to a side track before rejoining the loop. A muted CPU clock sits outside the factory floor, showing that the host is waiting while the compiled loop keeps moving. It explains how rvLLM keeps the whole forward pass inside one compiled TPU region. — The flat-scan trick keeps the whole forward pass inside one compiled region. That is how rvLLM turns attention switching into a compiler problem instead of a host scheduling problem.

The important detail is not that the loop exists. It is that the loop is the unit of compilation. Once the entire pass sits inside one XLA region, the TPU no longer needs Python to choreograph layer by layer execution. The host gets out of the way.

Megakernels on GPU, compiler fusion on TPU

rvLLM’s decode path makes the backend split even clearer. On GPU, it composes a line of operations into something like a persistent megakernel. On TPU, that same straight-line shape becomes fuel for fusion. The code is different, but the design principle is the same: reduce launch overhead by making the execution path flatter.

Axis	GPU path	TPU path
Execution style	Hand-tuned CUDA and graph capture	Compact JAX and XLA compilation
Optimization lever	Persistent megakernel and fused kernels	Compiler fusion over a single scan loop
Control flow	Specialized kernels and manual runtime structure	Structured loop with conditional attention routing
Philosophy	Win with explicit low-level control	Win by making the compiler see the whole problem

def decode_layer(x, layer_idx, cache):
    x = fused_add_norm_qkv_gemv(x)
    x = rotary_embedding(x)
    x = paged_attention(x, cache, layer_idx)
    x = fused_oproj_add_norm_gateup_gemv(x)
    return x

That snippet is the key bridge. It is simple enough to read as a pipeline, but structured enough for the backend to squeeze into one efficient execution unit. The point is less the individual ops than the absence of chatter between them.

Paged attention and quantization are the real battle

The repo’s performance story does not rest on one clever trick. It rests on memory discipline. Paged attention reconstructs context from non-contiguous KV blocks. INT4 weight-only quantization cuts the cost of matrix math. Together, they reduce the amount of bandwidth the engine has to burn just to keep tokens moving.

Technique	What it fixes	Why it matters
Paged attention	KV cache fragmentation	Keeps long contexts from turning memory into a mess
INT4 weight-only quantization	Expensive linear layers	Shrinks bandwidth and improves fit
Fused decode path	Kernel launch overhead	Turns many small operations into one coherent step

Why rvLLM is different from vLLM

This is not a fair fight in the usual sense, because the projects are aiming at different kinds of generality. vLLM is the broader reference point. rvLLM is more opinionated. It is willing to specialize the backend to the accelerator and the model shape instead of trying to be the most universal abstraction layer in the room.

Question	rvLLM	vLLM
What is it optimizing for?	Backend-specific throughput and tok/s per dollar	Broadly useful inference performance and compatibility
How does it treat TPUs?	As a first-class compiler-shaped target	Not the center of the design story here
Where does speed come from?	Rust orchestration plus hand-tuned GPU kernels or compact JAX on TPU	A mature general-purpose serving stack
What kind of team is it for?	Teams willing to align software with hardware	Teams that want a widely adopted default

That difference matters. rvLLM suggests that inference engines may split into specialists and generalists. The specialists will be harder to port, but easier to optimize. The generalists will be easier to adopt, but harder to push to the edge of a single accelerator’s envelope.

What this says about the next generation of inference engines

rvLLM’s real contribution is not just speed. It is an argument about how speed should be earned. On GPU, you can still win with detailed kernel work and graph replay. On TPU, the better path may be to write less code, shape it carefully, and let the compiler see the whole loop.

That is why the project feels important even before you care about Gemma 4 specifically. It shows that the future of inference infrastructure may not belong to the most abstract engine. It may belong to the engine that is most willing to become hardware-shaped.