llm-rs is an LLM inference engine written entirely in Rust, from the tokenizer to the CUDA kernels to an OpenAI-compatible server with continuous batching and a paged KV cache. These posts follow its design, its kernels, and how it measures up against vLLM.
-
I benchmarked llm-rs, my Rust inference engine, against vLLM on the same GPUs and models. With one request at a time llm-rs keeps up, and with BF16 weights it can lead. With 256 requests at once vL...
I benchmarked llm-rs, my Rust inference engine, against vLLM on the same GPUs and models. With one request at a time llm-rs keeps up, and with BF16 weights it can lead. With 256 requests at once vLLM is 1.4–3.4x faster. Both results come from a handful of ideas about how a GPU spends its time, and this post explains those ideas, using the benchmark numbers as worked examples.
The setup: Qwen2.5 models (7B in BF16 and 4-bit AWQ, 3B in BF16) on an RTX A5000 (768 GB/s) and an RTX 2000 Ada (224 GB/s). Both servers were capped at 64 concurrent sequences and faced the same client. Prompts were 256 tokens.
One request at a time 256 requests at once 7B BF16, A5000: llm-rs / vLLM 35.5 / 27.0 tok/s 709 / 1,208 tok/s 7B AWQ, A5000: llm-rs / vLLM 41.2 / 94.6 tok/s 550 / 1,878 tok/s 3B BF16, RTX 2000 Ada: llm-rs / vLLM 32.4 / 33.1 tok/s 1,143 / 1,651 tok/s 1. Every kernel is limited by memory or by math
A GPU kernel takes at least the larger of two times:
time ≈ max( bytes moved / memory bandwidth , FLOPs / math throughput )Which one wins depends on arithmetic intensity: FLOPs done per byte read. An RTX A5000 reads about 768 GB/s and does roughly 110 TFLOPS of dense BF16 tensor-core math, a ratio of about 145 FLOPs per byte. A kernel below that ratio waits on memory; above it, on math. This is the roofline model, and nearly everything below follows from it.
2. Decode with one request is a memory problem
To produce one token, the model multiplies every weight matrix by one vector. Each BF16 weight (2 bytes) gets used for one multiply and one add, so 2 FLOPs per 2 bytes: an intensity of 1, more than 100x below the ridge. The math units sit mostly idle; the job is streaming weights.
That gives a hard floor. Qwen2.5-7B in BF16 is about 15 GB, and 15 GB / 768 GB/s = 19.6 ms per token, so no engine can beat about 51 tokens/s for a single stream on that card. The question for single-stream decode is how close a kernel gets to streaming at full bandwidth:
- llm-rs: 27.8 ms per token, about 70% of the floor. Its matrix-vector kernels give each output row to one warp, whose 32 lanes read consecutive memory, the access pattern DRAM serves fastest.
- vLLM: 36.9 ms per token.
That’s why the one-at-a-time column is close: both engines are fighting the same bandwidth limit, and hand-tuned matrix-vector kernels are enough to win it.
3. Batching moves decode from memory toward math
Now serve 64 requests at once. The weights are read once per step and used for all 64 tokens, so the intensity rises to about 64 FLOPs per byte. The step takes barely longer than a single-token step, but produces 64 tokens. This is why throughput at 256-at-once is 20–45x the one-at-a-time number, and why continuous batching is the core idea of modern serving: keep as many requests in each step as possible, admitting new ones as old ones finish.
Two things limit how far batching goes:
- Past the ridge (~145 tokens per step here), a step is bound by math, and the speed of the matrix-multiply kernel decides everything.
- The KV cache. Each request’s past keys and values are re-read every step. For Qwen2.5-7B that is 2 (K and V) x 28 layers x 4 KV heads x 128 dims x 2 bytes = 57 KB per token of context. At 64 requests with ~400 tokens each, that’s ~1.5 GB per step, a tenth of the weights, and it grows with every generated token. Two ideas keep it in check:
- Grouped-query attention (GQA): 28 query heads share 4 KV heads, so the cache is 7x smaller than it would be. A good kernel loads each KV head once for all 7 query heads that use it.
- Paged KV memory (vLLM’s PagedAttention): the cache lives in fixed-size blocks found through a table, so memory isn’t reserved for each request’s maximum length. More requests fit, so batches get bigger.
4. Under load, everything comes down to the GEMM
With 256 requests arriving together, the server also has to process 256 prompts of 256 tokens. A prompt is the extreme case of batching: hundreds of tokens against the same weights, far past the ridge. 65,536 prompt tokens x 2 FLOPs x 7B weights is about 0.9 PFLOP of matrix multiplication.
The profile of llm-rs under that load shows where the time goes:
What the GPU was doing (7B, A5000) BF16 AWQ Matrix multiply (prompts and batched decode) 81% 88% Decode attention 13% 8% Everything else 6% 4% GPU busy 99.7% 99.8% The GPU is never idle, so scheduling isn’t the gap. The gap is how many FLOPs per second the matrix multiply (GEMM) achieves. llm-rs’s hand-written tensor-core GEMM reaches roughly a quarter of the card’s tensor peak on these shapes; vLLM uses cuBLAS, which by my estimate gets about twice that. That one kernel accounts for most of the 1.7x difference at 256-at-once.
5. What makes a GEMM fast
Reaching the tensor cores’ peak is mostly about keeping them fed. Four ideas do most of the work.
Tiling for reuse. A block of threads computes a 128 x 128 tile of the output. It loads a 128-row slice of the input and of the weights into on-chip shared memory, and every loaded value is used 128 times. Bigger tiles mean fewer bytes per FLOP, which pushes the kernel past the ridge point.
Tensor-core instructions. Each warp issues
mma.sync, which multiplies a 16x16 by 16x8 block of BF16 values into FP32 sums in one instruction. Operands come from shared memory vialdmatrix, which loads them straight into the register layout the instruction expects.Pipelining. Loading the next slice of K must overlap with computing on the current one. Asynchronous copies (
cp.async) let a block start fetching step s + 2 while computing step s. Each extra stage gives slow memory more time to arrive, at the cost of more shared memory per block.Bank-conflict-free layouts. Shared memory is split into 32 banks of 4 bytes. If several threads of a warp read different addresses in the same bank, the reads are serialized.
ldmatrixreads one 16-byte chunk from each of 8 consecutive rows. With rows of 64 bytes, those 8 reads pile into the same two groups of banks: a 4-way conflict. Padding each row fixes it but wastes 25% of shared memory. An XOR swizzle fixes it for free, by storing chunk c of row r at slot c XOR ((r / 2) mod 4):bank groups (4 banks each) G0 G1 G2 G3 | G4 G5 G6 G7 row 0 [c0] c1 c2 c3 | row 1 | [c0] c1 c2 c3 row 2 c1 [c0] c3 c2 | row 3 | c1 [c0] c3 c2 row 4 c2 c3 [c0] c1 | row 5 | c2 c3 [c0] c1 row 6 c3 c2 c1 [c0]| row 7 | c3 c2 c1 [c0] chunk 0 of rows 0-7 lands in all 8 groups: one pass instead of fourA fifth idea is easy to miss: everything hot must stay in registers. A warp’s 64 x 32 output tile is 64 FP32 sums per thread. If the compiler can’t keep them in registers, for example because a loop indexing them wasn’t unrolled, they spill to local memory, per-thread storage backed by DRAM. I hit exactly this: a GEMM whose accumulators had spilled was 4.6x slower (9.3 ms vs 2.0 ms per call) and passed every correctness test. The only reliable check is to read the compiled PTX and count
ld.local/st.localinstructions.6. Quantization helps memory-bound work, and can hurt compute-bound work
4-bit weights cut bytes by 4x, so they should help wherever bandwidth is the limit. Two details decide how much.
The format’s overhead. Weights are stored in groups, each with a scale (and often a zero point):
Format Bits per weight BF16 16 AWQ: 4-bit, FP16 scale + zero per 128 4 + 32/128 = 4.25 llm-rs today: 4-bit, f32 scale + offset per 32 4 + 64/32 = 6 The cost of unpacking. Each 4-bit code has to become BF16 before the tensor cores can use it. In the memory-bound regime that work hides behind the memory traffic. In the compute-bound regime it adds instructions to an already-full pipeline, and the FLOP count doesn’t drop at all.
The benchmark shows both effects:
- One at a time (memory-bound), vLLM’s AWQ runs at 9.9 ms per token, within 1.6x of the ~6 ms floor (4.25-bit weights plus a BF16 output layer). llm-rs gets 23.7 ms: its format is 40% fatter, its 4-bit kernel streams at about half bandwidth, and its output layer stays in BF16 (1 GB by itself).
- Under load (compute-bound), llm-rs’s AWQ is slower than its BF16 (550 vs 709 tok/s): it unpacks weights through registers before each multiply. vLLM’s AWQ is faster than its BF16 (1,878 vs 1,208) thanks to Marlin, a GEMM built for 4-bit weights. Marlin reorders the weights at load time so that one 16-byte load hands each thread exactly the fragment it needs. It converts 4-bit codes to BF16 with a few bit operations instead of float math, and it keeps several asynchronous copy stages in flight. Much of a decode batch stays memory-bound, so fewer bytes still win.
7. Reading the scoreboard
With those ideas, each cell of the opening table has a reason:
Result Why llm-rs leads 7B BF16 one at a time Both engines are bandwidth-bound; llm-rs’s matrix-vector kernels stream closer to the 19.6 ms floor vLLM leads 2.3x on AWQ one at a time Compact 4.25-bit format and Marlin kernels near the bandwidth floor vLLM leads 1.4–1.7x on BF16 under load Compute-bound GEMM: cuBLAS gets about twice the tensor throughput vLLM leads 3.4x on AWQ under load Marlin keeps 4-bit fast under load; llm-rs’s unpacking makes 4-bit slower than BF16 So the roadmap for llm-rs is the GEMM, in order:
- Deeper pipelines and swizzled shared memory for the BF16 GEMM.
- BF16 activations end to end, so no conversion pass sits in front of each GEMM.
- Marlin-style 4-bit: reordered weights, bit-trick unpacking, and a compact AWQ-sized format.
8. Takeaways
- Ask which roof you’re under. With one request, decode is about bandwidth; with many requests and prompts, it’s about the GEMM. The same engine can win one and lose the other.
- Batching is the multiplier. Continuous batching and paged KV memory exist to make each weight read serve as many tokens as possible.
- Quantization is a bandwidth tool. It pays off when bytes are the limit, and only if the format is compact and unpacking is nearly free.
- A kernel can be correct and 4x slow. Registers, shared-memory banks, and pipeline depth don’t show up in correctness tests; read the PTX and profile against a ceiling.
-
llm-rs is an LLM inference engine written entirely in Rust: weight loading, tokenizer, sampling, the forward pass, the GPU kernels, and an OpenAI-compatible server. It began as a GPT-2 port with a ...
llm-rs is an LLM inference engine written entirely in Rust: weight loading, tokenizer, sampling, the forward pass, the GPU kernels, and an OpenAI-compatible server. It began as a GPT-2 port with a plain CPU implementation and now runs Llama- and Qwen-family models on NVIDIA GPUs, with continuous batching and a paged KV cache. This post gives the overall design, then the main GPU story: starting with cuTile tile kernels, and then replacing the hot paths with hand-written SIMT kernels. That roughly doubled decode throughput on the same GPU and made long-prompt prefill several times faster.
1. The high-level idea
The project follows one rule: one model, many backends, same math.
- A CPU reference that is simple and readable, and serves as the source of truth.
- A SIMD CPU backend (NEON / AVX2) for real speed on laptops and servers.
- An Apple GPU backend on Metal.
- An NVIDIA GPU backend that keeps weights and the KV cache on the GPU and runs the whole forward pass there.
Each backend is checked against the reference, both kernel by kernel and end to end: greedy GPU output must match the CPU token for token. That check is what made aggressive kernel work safe. A faster kernel that changes the output gets rejected.
On top of the model sits a serving layer inspired by vLLM: a scheduler that batches requests continuously, a paged KV cache with prefix sharing, steps that mix prefill and decode, and an HTTP API compatible with OpenAI clients.
OpenAI-style HTTP API │ scheduler (continuous batching) │ paged KV cache + prefix cache │ executor ─► CPU | Metal | CUDA2. Why decode is a memory problem
Generating one token means reading every weight once to produce one vector. A 0.5B-parameter model in BF16 is about 1 GB of weights. On a GPU with ~450 GB/s of memory bandwidth, that puts the best possible single-stream decode at around 450 tokens per second.
A matrix-vector product does very little math per byte, so decode is really about one thing: reading weights at full bandwidth. That ceiling became the yardstick for every kernel change. “Faster” wasn’t enough; the question was always “what fraction of the ceiling are we at?”
3. First pass: cuTile
cuTile lets you write GPU kernels in Rust as operations on tiles. You describe how tensors are split into blocks and what to do with each block. The compiler decides how threads, shared memory, and tensor-core instructions are used.
That made it fast to get a correct, fully GPU-resident forward pass. Tuning then went a long way:
- Caching compiled kernels and warming them up so compile time didn’t land on the first request.
- Queuing launches without waiting on the host after each one.
- Fusing neighbouring operations (e.g. the gate and up projections with their activation) so the activations aren’t written out and read back.
- Sweeping tile sizes for the decode matvecs.
Together these took decode from about a dozen tokens per second to roughly 160 tok/s, a better than 10x improvement.
Then it stopped improving, at about 35% of the bandwidth ceiling. The GPU was busy almost all the time, so launch overhead wasn’t the problem; the kernels were. Large projections ran reasonably well, but the small ones (attention projections, the output projection) were several times slower than bandwidth allowed. I tried more block shapes, pipelining hints, and split-K reductions. None of them helped.
The underlying problem is that for a small matvec, performance depends almost entirely on how a row maps onto the lanes of a warp and onto memory transactions. The tile abstraction is designed to hide exactly that mapping.
4. Going down a level: SIMT kernels
So the hot paths moved to SIMT kernels: Rust compiled to PTX (via cuda-oxide), with explicit threads, warps, and warp shuffles. This is the classic CUDA style.
The decode matvec becomes the textbook form: one warp per output row. Each of the 32 lanes reads consecutive weights, so the warp as a whole reads one contiguous chunk of memory per step. That access pattern is what lets the hardware reach full bandwidth. At the end, the lanes’ partial sums are combined with warp shuffles, with no shared memory needed.
Two design choices kept this practical:
- Both kernel styles share one stream. SIMT and cuTile kernels can be mixed freely in the same forward pass, so the switch happened one kernel at a time and was A/B tested at every step.
- The SIMT kernels are prebuilt. The special toolchain is only needed when changing a kernel, not to build or use the project.
Decode: from ~35% to ~70% of the ceiling
On the same GPU, the steps that mattered:
Change Approx. decode Share of ceiling cuTile baseline ~160 tok/s ~35% One-warp-per-row matvecs, split-sequence attention, fused small ops ~250 tok/s ~55% No memory allocation inside the decode loop ~300 tok/s ~66% Norms folded into the projections, better work split per warp ~315 tok/s ~70% That’s about 2x the cuTile throughput on identical hardware, and the output still matches the CPU exactly. The number of kernel launches per layer also dropped by more than half. The main lessons:
- Once kernels are fast, the host becomes the bottleneck. With near-bandwidth kernels, the profiler showed the GPU sitting idle a third of every step while the host allocated and freed temporary buffers. Allocating one workspace per request removed that idle time, and run-to-run noise almost disappeared with it.
- Decode attention is about latency, not bandwidth. It works best split into many small chunks of positions, each handled by one warp with an online softmax, followed by a cheap merge step. More chunks mean more warps to hide memory latency, up to the point where merging costs more than it saves. The best chunk size came from measurement, not a guess.
- Fusion can be slower at first. Moving the normalization into the projection kernels removed launches but initially slowed decode. The fused prologue was a chain of dependent memory round trips. Prefetching didn’t fix it, and neither did fixing shared-memory bank conflicts. What worked was issuing several independent loads at once, and then giving each warp more rows so that fewer blocks repeated the prologue.
At the end, the largest projections run at roughly 80–90% of memory bandwidth. What’s left is mostly the smallest projection and latency-bound attention.
Prefill: tensor cores, written by hand
Prefill processes the whole prompt at once, so it is limited by compute, the opposite of decode. It got its own SIMT kernels:
- Short prompts use a batched matvec that reads each weight once for a group of tokens. Its first version was 10x slower than expected because its accumulators spilled out of registers. Unrolling the loops fixed that, and short-prompt prefill became ~2.5x faster.
- Long prompts use a hand-written tensor-core GEMM (BF16 inputs, FP32 accumulation). It is 2.5–5x faster than the cuTile version, and the gap grows with prompt length.
- Prefill attention became flash-style and reads straight from the paged KV cache. It took three attempts. The simple version, where each warp handles its own query row, was slower than cuTile. Sharing key/value tiles across query rows brought it level. A tensor-core version was several times faster on long prompts. Short prompts don’t fill enough tensor-core blocks, so a measured crossover point picks between the two kernels.
End to end, time to first token on a ~1,200-token prompt dropped by more than half.
5. Serving on top of the kernels
The same kernels gained batched versions that know each request’s position and where its KV blocks live. That is what continuous batching and paging need:
- Paged KV cache. Looking up blocks through a table costs essentially nothing once attention fetches one block per chunk.
- Prefix caching. A request that shares a long system prompt with an earlier one starts generating about 4x sooner, because only its new suffix has to be processed.
- Batched sampling. Sampling every request in a batch with one kernel launch, instead of many small launches per request, roughly tripled sampled throughput and cut queueing delay sharply.
Against vLLM on the same GPU and model, llm-rs keeps up at light load. Under heavy load vLLM is still well ahead. The profile points to two causes:
- A fixed cost per step: the scheduler waits on the host for each step’s tokens before planning the next step.
- Batch scaling: the batched decode matvec still re-reads the weights for every group of tokens. At large batch sizes, decode should run on the tensor-core GEMM that prefill already uses.
Those two are the next items on the list.
6. Takeaways
- Start with tiles, finish with warps. A tile-level DSL is the fastest route to a correct GPU implementation. For kernels limited by memory bandwidth or latency, the lane-to-byte mapping is the whole problem, and you have to write it yourself.
- Measure against a ceiling. Tracking the share of peak bandwidth showed when a kernel was done and where the remaining gap was.
- Keep a parity check you trust. Exact agreement with a simple reference made it cheap to try risky ideas and throw away the ones that failed.
- Rust works all the way down. Host code, tile kernels, and SIMT kernels are all written in Rust, and users never need the special toolchain.