I read about NVIDIA’s Tile programming model in the release announcement of cuTile.jl v1.0. During a previous project on Sparse Linear Algebra, I wrote all my kernels by reasoning about the data layout, and then translating it to the work each thread had to do. So programming directly on the blocks of data instead sounded like a nice idea that I wanted to try out.

Around the same time, I had been reading papers on LLM serving and inference optimization, like FlashAttention or PagedAttention, which optimize inference by tiling (see where this is going?) the attention kernel instead of computing it all at once. To get a real understanding of these algorithms, I thought building a toy inference engine would be fun.

And so, I started to build TileInfer.jl. It is an LLM inference engine fully implemented in Julia, where the core kernels are written with cuTile.jl. I tried it on small and medium-sized Qwen and Apertus checkpoints, and got them all running at speeds close to the hard limit of what is possible on my system. This post is a way for me to share what I learnt, and give a short introduction to LLM inference optimization.

Just so we are on the same page: TileInfer.jl is meant to be read, measured and poked at, not run in production. If you want to serve models, I’d strongly recommend a well-established engine like vLLM, SGLang or llama.cpp. If for some reason you have a real use-case for a Julia-native inference engine, let me know. Otherwise, I’ll just keep working on this for fun.

High-level design

TileInfer.jl follows the structure of larger engines, but drops two key features for the sake of simplicity:

  • Multi-GPU and multi-node orchestration and routing: I only have one GPU in my PC (and one PC), so I couldn’t begin to test or optimize these even if I wanted to. I’ll get back to this after I invest in a domestic data-center.
  • Multi-user features: batching, KV-cache offloading and prefix caching are all left aside for now. I might revisit this at some point. Currently, the engine only serves a single conversation, and does not cache anything across two sessions.

The main layers left are the following:

┌──────────────────────────────┐
│             CLI              │
├──────────────────────────────┤
│          Generation          │
├──────────────────────────────┤
│         ModelRunner          │
├───────────────┬──────────────┤
│     Model     │   KV cache   │
├───────────────┴──────────────┤
│             Ops              │
└──────────────────────────────┘
  • CLI: pull a model, start a chat session.
  • Generation: turn a prompt into a reply, one token or one verified draft at a time.
  • ModelRunner: run each step from fixed buffers, replayed as a CUDA graph, and sample on the GPU.
  • Model: hold the weights and run the layers.
  • KV cache: keep the keys and values of past tokens, in pages.
  • Ops: the kernels for all operations needed in a forward pass.

Two other design choices matter internally:

  • No neural-network framework. Nothing here needs autodiff, so a model is a plain struct of CuArrays. Every allocation and every kernel launch is visible in the code.
  • Two backends for every op. Each operation has a reference implementation built from CUDA.jl broadcasts and cuBLAS, which is the ground truth in tests, and a cuTile implementation. All custom kernels are tested against the reference, then benchmarked and profiled to see where the gains come from.

Before we dive in, a very quick recap of what “serving an LLM” means

Every reply goes through two phases, and they put very different loads on the GPU:

  • Prefill, or “reading”, runs the whole prompt through the model in one forward pass: all the tokens are known up front, so they are processed in parallel as large matrix-matrix products, and their keys and values are written to the KV cache. The weights are read once for the entire prompt, so prefill is compute-bound, and it sets how long you wait for the first token.

  • Decode then generates the reply one token at a time, because each token depends on the previous one. Each step is a matrix-vector product per layer: it has to read every weight of the model (8 GB for Qwen3-4B in BF16) to do very little arithmetic on a single token, so decode is memory-bound.

Incremental acceleration

The project was built in milestones, each one speeding up the engine a bit more (thanks to survivorship bias). Every number in this post is measured at one reference point: Qwen3-4B in BF16 on an RTX 4080 SUPER, a 1,024-token prompt prefilled in one step, then 128 tokens decoded after it. The table below traces the incremental improvements by switching features on one at a time, so each row adds one feature to the row above it:

Stage Decode Prefill (1,024 tokens) What changed
Reference ops 8.6 tok/s 2.2k tok/s (458 ms) cuBLAS and broadcasts
cuTile kernels 11.2 tok/s (1.30×) 2.9k tok/s (1.30×) fast kernels hidden behind slow attention
Paged flash-attention 72.3 tok/s (6.5×) 10.6k tok/s (3.7×) attention reads pages in place
CUDA graphs 75.9 tok/s (1.05×) 10.6k tok/s (1.00×) one launch per step instead of ~365; only really matters on small models
Fused kernels 78.7 tok/s (1.04×) 12.0k tok/s (1.13×) norms, RoPE and the KV write fused; SwiGLU in the GEMV epilogue

I was surprised to see that before any of the bigger optimizations, swapping the plain CUDA.jl broadcasts and cuBLAS calls for one cuTile kernel per op already gives a 1.3× gain. The largest gain by far comes from the PagedAttention kernel, which reads each page once, in place, through a block table (to be fair, my reference attention kernel is particularly bad). I strongly recommend reading the full paper if you have the time.

CUDA graphs come almost for free, as long as the entire step can be captured: no allocations, no host syncs, etc. On Qwen3-4B, a decode step takes 12.7 ms, long enough to hide most of the kernel launch overheads, so graphs only gain 1.05×. But on a small model like Qwen3-0.6B, graph replay brings a gain of 1.22×.

The fused kernels entry summarizes the smaller optimizations that step-by-step profiling made possible. The goal is simple: make the compute-bound prefill fully use all 80 Streaming Multiprocessors (SM, compute units on the GPU), and push the memory-bound decoding to the bandwidth limit of 736 GB/s. The details are in the history and roadmap page.

Where it stands

Next is a quick comparison with four other engines, in the same configuration on the same machine: plain greedy decoding in BF16, one sequence at a time, prefix caching off, CUDA graphs on everywhere. We benchmark at 1k, the reference point above, and on a longer prompt: an 8,192-token prompt with 1,024 tokens generated.

Qwen3-4B, BF16 Prefill, 1k Decode, 1k Prefill, 8k Decode, 8k
TileInfer 85.6 ms 78.7 tok/s 0.88 s 69.6 tok/s
Nano-vLLM 81.1 ms 78.1 tok/s 0.81 s 69.1 tok/s
mini-SGLang 81.4 ms 79.2 tok/s 0.82 s 70.2 tok/s
vLLM 0.31 88.0 ms 78.9 tok/s 0.84 s 69.8 tok/s
llama.cpp v0.9.0 101 ms 77.0 tok/s 0.94 s 68.0 tok/s

On decode, all five engines are within about 3 % of each other, at 85–88 % of the bandwidth roofline. At batch size 1, everyone hits the same wall: we need to read all the model weights for each token generated.

On prefill, the engines differ more: Nano-vLLM and mini-SGLang, which run cuBLAS GEMMs, prefill 5 % faster than TileInfer at 1k and 7–8 % faster at 8k, where TileInfer’s prefill attention starts to show (faster prefill attention for long prompts is on the to-do list).

Pushing further without cheating: speculative decoding

Speculative decoding uses a cheap proposer to generate a sequence of k tokens, either sequentially (e.g. doing k forward passes with Qwen3-0.6B as a cheap stand-in for Qwen3-4B) or in parallel, generating a block of k tokens at once. The real model then checks them all in a single forward pass, and keeps the longest correct prefix.

The output is unchanged (exactly for greedy decoding, in distribution for sampling), but the computation to get there is very different: where we used to read all model weights for each output token, we can now verify multiple tokens in a single batched forward-pass. If our proposer gets g tokens accepted on average, we only read the model weights once per g output tokens, cutting memory traffic by a factor of up to g. This gives a significant boost to decoding, which is memory-bound as explained earlier.

TileInfer has three proposers: n-gram lookup in the context, a small draft model, and DFlash, which drafts a whole block of 15 tokens in one pass of a small model over the target model’s hidden states.

On a task that mostly copies its input (renaming a type in a Julia file), DFlash takes Qwen3-4B from 79 to 419 tok/s (5.3×). On open-ended text (explaining that same file) it still gets 174 tok/s (2.2×).

Decode, tok/s edit (copies the input) explain (open text)
plain 78.7 78.7
n-gram lookup, k = 7 396 (5.0×) 83 (1.05×)
Qwen3-0.6B draft model 235 (3.0×, k = 7) 113 (1.44×, k = 2)
DFlash, k = 15 419 (5.3×) 174 (2.2×)

vLLM and llama.cpp both support DFlash with the same drafter, so we can compare them side-by-side. At the 1k point, the three engines land close together, and which one wins mostly follows how many drafted tokens each step accepts, which varies between engines through small rounding differences.

At 8k-token prompts, TileInfer is 1.37× faster than vLLM on the edit (95.5 against 69.5 tok/s): the details of why vLLM’s kernel is slower here are explained in the vLLM comparison page, but the TL;DR is: we split the computation of attention on the 8k tokens in chunks, which helps keep all SMs in the GPU busy, while vLLM’s kernel only launches one block of computation per attention head, only using 32 SMs out of the 80 available. This only becomes relevant once reading the model weights is no longer the primary bottleneck, which is why it appears after we add speculative decoding. And it gets worse after INT8 quantization, which comes next, to further reduce the time spent reading model weights.

Pushing further with cheating: INT8 quantization

Since decode is bound by the bytes it reads, the other lever is to read fewer of them. Storing weights in INT8 instead of BF16 halves the footprint: Qwen3-4B’s step reads ~4 GB of weights instead of 8 GB, and the model takes half the memory to store. Obviously, the “cheating” part is that the output is no longer exactly the same, and the loss in precision can degrade model performance.

Symmetric INT8 uses only 255 values, from -127 to 127, so each vector x is stored as INT8 values q and one Float32 scale, with x ≈ scale × q. The scale is picked so that the largest entry of the vector lands on the edge of the range, scale = max |x| / 127, and every other entry is rounded to the nearest multiple of scale. The rounding error of an entry is then at most scale / 2, which is small relative to that vector’s own values. One scale per vector costs 4 bytes next to its hundreds or thousands of INT8 values, which is a few percent of storage at most. The quantization is symmetric (zero stays exactly zero, no offset to store), and TileInfer applies it in two places:

  • Weights: one scale per output channel, for every linear layer and the lm_head (the embedding stays in BF16). Each output channel is a dot product of the input with one column of weights, so scaling a column only rescales its output: the kernel multiplies the INT8 values in, and applies the one scale at the end. With a scale per channel, a channel of small weights does not lose its precision to a channel of large ones.
  • KV cache: one scale per cached vector (one KV head at one position), written once when the position is cached. The attention kernels never dequantize the cache: a key’s scale multiplies its row of scores, and a value’s scale its probability before the product with V.

This table shows how much the INT8 quantization changes the output (ranges are over our two prompts, edit and explain, with ~1,000 positions each), compared to the output under the plain BF16 weights and cache:

Qwen3-4B Mean KL to BF16 Top-1 agreement
INT8 weights 0.0027–0.0038 97.5–98.0 %
INT8 KV cache 0.0029–0.0060 97.9–98.5 %
both 0.0041–0.0156 97.5–98.2 %

The next token is the same at about 98 % of positions. The two errors add up when both are on, mostly on the open-ended prompt. In practice, greedy replies stay the same for a few tokens and then drift apart. Please note that this is a very shallow test, and a proper measure of the “intelligence” lost by quantization needs to run a complete benchmark and evaluate a task success rate of some kind, which I didn’t bother implementing yet.

In the meantime, we can compare the speed:

Qwen3-4B Prefill, 1,024 tokens Decode
BF16 85.6 ms 78.7 tok/s
INT8 weights 90.9 ms (0.94×) 138.2 tok/s (1.76×)
INT8 KV cache 88.1 ms (0.97×) 78.8 tok/s (1.00×)
both 93.0 ms (0.92×) 138.7 tok/s (1.76×)

I didn’t expect the prefill kernel to gain anything as it’s not limited by memory, and it turned out to even slow down a bit: the kernel now needs to convert the INT8 weight tiles to BF16 before going to the tensor cores, which adds additional computation to an already compute-bound kernel.

INT8 weights are close to doubling decode speed under short context, as expected.

The INT8 KV cache does nothing at 1k tokens of context, where the cache is 2 % of the bytes a step reads. It pays off as the context grows and the cache starts to rival the weights in size (4.8 GB at 32k positions, against 8 GB of BF16 weights or 4 GB of INT8 ones):

Decode after N positions BF16 cache INT8 cache
8k, BF16 weights 70.1 tok/s 73.3 tok/s (1.05×)
8k, INT8 weights 113.9 tok/s 122.4 tok/s (1.07×)
32k, BF16 weights 51.4 tok/s 60.9 tok/s (1.18×)
32k, INT8 weights 71.6 tok/s 91.2 tok/s (1.27×)

Quantization and speculative decoding combine: with INT8 weights, DFlash takes the edit task to 906 tok/s and the explanation to 282 tok/s.

Finally, it’s worth noting that INT8 weights also make it possible to fit models that BF16 cannot: Apertus 1.5 8B (16.1 GB in BF16) runs in 8.6 GB at about 80 tok/s. Similarly, the INT8 KV cache doubles the maximum context window we can fit on the same memory footprint.

Returns on CUDA Tiles and cuTile.jl

Using cuTile lets you leave threads, shared memory and a lot of other complicated logic to the compiler. As I mentioned in the introduction, this makes it much easier to reason about what each kernel is doing by thinking everything through the data instead of the compute. This turned out to be even more accurate than I expected, as the PagedAttention kernel is designed to work on blocks / tiles instead of large tensors.

However, like any additional abstraction layer, it sometimes led to unexpected behavior. I was able to track down these issues by profiling with Nsight Compute and poking at the compiler. Here are 3 examples that were caught by benchmarking regressions, but could only (realistically) be diagnosed by digging deeper:

  • Small matrix products don’t run on tensor cores. Below the tensor cores’ native shape (16 × 8), the compiler switches to a slower code path that spills registers. A product with 8 rows took 115 µs where the same product padded to 16 rows took 14 µs: twice the arithmetic, but still 8× faster.

  • Going over 32 kernel parameters is bad: CUDA.jl then falls back to a generic argument-packing path that is much slower (and allocating). The danger is that it happens sooner than you’d think: every array argument expands into a pointer, plus a size and a stride for each dimension. Passing 2-D views instead of 4-D arrays made the KV-cache write kernel 1.7× faster, just by staying under 32 parameters.

  • argmax checks for NaN at every step, even with propagate_nan=false. Logits are never NaN, so a hand-written argmax was 1.2× faster and shaved 2 µs off every sampled token.

The “AI” part

It had been a while since I worked on a Julia project, and I was not sure how to set things up to work efficiently with my Claude subscription. While my regular “dev” setup is pretty barebones (stock Claude Code + a few instructions in CLAUDE.md), I think the performance optimization workflow is worth explaining quickly. Once a first implementation is running correctly, the usual performance optimization loop looks like:

  1. Run a benchmark to get a baseline timing measurement on the kernel we want to optimize.
  2. Run the profiler.
  3. Analyze profile results: where is the bottleneck, is there an unexpected stall, does the kernel spill registers*?
  4. Make a change based on the previous results.
  5. Profile and run the summary script again to see if the issue is solved.
  6. If yes, benchmark again to confirm; if not, go back to step 4.

* Register spilling: If the compiled kernel needs more than the 255 registers available per thread, it has to save some intermediate values to (slow, off-chip) local memory and retrieve them later, which is known as spilling. It is a common problem I encountered while working with cuTile.jl, as it can be hard to estimate how the compiler will unfold the tile operations.

To optimize the performance, I used Claude to automatically iterate through the loop. By default, it is able to read and change the code, run stuff and reason about what to do. So the only change needed from the “manual” setup was to adapt step 3, because the AI agent can’t really open the Nsight GUI and explore the time-graph by itself. Instead, I gave it a script that parses the raw profile output file and prints a short text-based report with the key values for my use case: registers used, occupancy, SM utilization, memory throughput, stall reasons, etc., each linked to the corresponding line in the Julia source of the kernel.

Naturally, Claude can also edit the script to add new values it needs to the report when needed.

While this is not really documented as it is not part of the project’s source code, the different scripts I use are all in the bench/ folder of the repo, so feel free to check it out (or point your agent there to replicate the setup)!

What’s next

The engine still has gaps: prefill always processes the entire prompt in one step, which runs out of memory on prompts longer than ~13k tokens (with DFlash in BF16). Chunked prefill splits a long prompt into chunks, typically large enough to make the overhead negligible. This way, it can process longer prompts with a memory footprint bounded by the chosen chunk size.

Other improvement areas include faster prefill attention for long prompts, an INT8 weight layout that avoids converting tiles in registers, and adaptive draft lengths with tree verification for speculative decoding. Finally, all the things that make an engine serve more than one person: continuous batching and prefix caching across sequences would be very useful and pretty much needed before any serious deployment.

The code is on GitHub, and the documentation goes much deeper into the KV cache, the attention kernels and speculative decoding. If you’re curious to learn more on GPUs going brrrr today, don’t hesitate to check it out!