Technology · Discipline

GPU inference engine engineering

I built a batched continuous-decode engine with paged key-value memory, admission control, canary-gated binding and typed refusals, and I tried the vendor's industrial stack first and closed it on measurements.

Rust and CUDA · about 16,400 linesCapabilities modelOne NVIDIA Blackwell GPU
Grammar-bounded output · read only the rows that can win
the tied embedding matrix: the full head reads all 593.5 MiBlit: the trie node's edges, at most 28
Dense head kept in the build as a cross-check
112 KiBread by the sparse head to rank the same candidates
  • 518 µsremoved per token by the on-GPU decode loop, up to 696 under load
  • 13.4×faster prefill in 256-token chunks, identical answers
  • 503 MiBCUDA context, reported apart so admission stays honest
  • 2.53×less device memory with NVFP4 weights, about 5 percent slower

Measuring the industrial path first

Capabilities model is a 0.6-billion-parameter tool-discovery model that had to serve two thousand concurrent agents. It had only the memory a resident 30-billion-parameter model left free on one NVIDIA Blackwell GPU. I had the vendor's industrial stack tried first. Its weight conversion corrupted some tensors by up to a thousand times. That is Trying the industrial path first.

TensorRT-LLM · requests per second
Card to itself492Beside the 30B model56
The memory left beside the 30B model is the only memory this model ever has.

I can take a model from weights to a custom inference engine whose speed and numerics I can defend.

I build custom only after measuring the standard option.

The whole decode loop on the GPU

I then directed the build of a purpose-built engine in Rust and CUDA, with its lab. It uses hand-written FP32-accumulating kernels with no cuBLAS and no tensor cores. Batch-one decode is a matrix-vector product, so tensor cores would sit idle, and staying off cuBLAS removes its silent TF32 default.

I put the entire generate-until-stop loop on the GPU as one CUDA graph: a seed sampling step, then a conditional WHILE node.

Host-driven · 451 launches per forward pass

the host sees every token and issues every step

One CUDA graph · the loop lives on the GPU
tokenpositiongrammar nodeemitted tokensno-repeat state
Loop state stays in device memory and the sampling kernel writes the loop condition, so the host sees nothing until the answer is done. Output is identical to the host-driven loop.

The output layer is grammar-bounded. The grammar knows the next token is one of the current trie node's outgoing edges, so the engine lists those candidates first and computes logits for those rows only. Prefill runs all prompt tokens through the weights in one pass per chunk.

Batching without running out of cache

The batched engine replaces three things that do not survive batching. KV memory is a paged pool with a per-sequence block table. Every step is a set of rows, each carrying its own token, sequence and position, so decode rows from many sequences and a slice of one long prompt share a launch and long prompts never block short ones. Grammar state is 16 bytes per sequence: no-repeat becomes a 64-bit set of used tools tested against a static per-edge set of the tools below it, so sequences share the trie read-only and never write to each other's state.

Paged key-value pool · 32-token pages, a block table per sequence
Short discovery call · 5 pages
2,662-token request · 86 pages
Admission reserves the pages for the whole prompt and answer up front, so an admitted request never runs out of cache mid-answer.

Admission refuses with a typed reason and the numbers when the pages are not there. The engine computes its memory need without allocating, tested against a real load. A failed load now frees what it took; before, four retries drained the card to 2 MiB.

The same bits in any batch

Determinism is structural. Reductions use a fixed 32-lane shuffle tree, or a fixed 256-thread tree for the head, and attention sums positions inside one thread. The batched engine reuses the serial engine's row kernels. That story is Bit-identical across batch compositions.

Alone or in a batch of 512
Bit-identity gates12 of 12Rank-one results stable across shuffled batches2,345 of 2,345Ranked lists matching an independent implementation484 of 484

Every speed change had to be bit-identical to its reference: register-tiled projections, hardware FP8 decoding in place of a divergent table lookup, row-tiled attention that stages a shared prompt prefix once for eight rows, and fused launches. The one value-changing switch, BF16 key-value storage with hand-written round-to-nearest-even, went in only after a full determinism re-run.

Speed that counts end to end

A change stays only if it moves the 469-query end-to-end time. My first row-tiled attention added a launch per row in decode and tiled prompts with no prefix worth sharing. I fixed both and kept the rule. NVFP4 is the footprint option and FP8 is the default build.

First row-tiled attention · time against the change before it
Before1×On a microbenchmark2× fasterOn the 469-query evaluation6× slower
Launches per step
Before fusion479Fused339
Bars in the first set are time, so shorter is faster. CUDA-graph replay for the batched step was priced at 1.26× at best, and fusion was chosen instead.

Beside the 30B model the engine used 2,740 MiB and served a burst of 2,048 clients at 85.9 requests per second. The engine serves one architecture. On mixed long requests it sustained 6.9 requests per second against a target of 200; the limit is key-value memory beside the 30B model, and A router beside a production model explains it.