Inference Engineering: What It Is and How to Learn It

What is inference engineering?

Inference engineering is the work of making trained models, especially large language models, serve predictions quickly, cheaply and reliably. It balances latency, throughput, memory, cost and output quality using techniques such as KV caching, batching, quantization and optimized GPU kernels. The best way to learn it is to build an LLM's inference path in order, starting with the forward pass and generation, then adding the KV cache, measurement, kernels, quantization and finally batching.

The Inference Engineering study plan covers attention variants, KV-cache operations, quantization and inference systems in PyTorch, and the CUDA Kernels and Triton Kernels plans cover the kernels underneath. This guide points to them where they fit the build order.

What does an inference engineer do?

An inference engineer makes models run efficiently in production. The job is to serve as many requests as possible, as fast as users need, on the least hardware, without degrading output quality. In practice that means measuring where time and memory go, then changing the model, the kernels or the serving system to improve it.

Inference engineering overlaps with ML engineering, MLOps and GPU programming, but its focus is narrower. It is about the runtime behavior of a model that has already been trained.

Why is LLM inference hard to make fast?

LLM inference is hard to make fast because text is generated one token at a time, and each step must read every model weight from GPU memory. For small batches, generation is limited by memory bandwidth rather than compute. Meanwhile the KV cache grows with every token and every concurrent request, so memory runs out before compute does.

Prefill and decode are different workloads

Every LLM request has two phases, and most inference techniques target one of them or the trade-off between them.

• Prefill processes the whole prompt at once. Many tokens go through the model in parallel, so the GPU runs large matrix multiplications and is usually compute-bound.
• Decode generates the response one token at a time. Each step pushes a single new token per request through every layer, so the main operations become matrix-vector products, and with small batches decode is usually bound by memory bandwidth.

Why is decode memory-bound?

Take a 7-billion-parameter model stored in 16-bit precision. Its weights occupy about 14 GB, and at batch size 1 every decode step reads all 14 GB to produce one token. On a GPU with about 3.35 TB/s of memory bandwidth, which is the published figure for an H100 SXM, reading 14 GB takes at least 4.2 ms. That sets a rough ceiling of around 240 tokens per second for a single request, before counting KV-cache reads or any overhead. The GPU's arithmetic units sit mostly idle, because each weight is used for only a couple of operations per step.

This is why batching and quantization help so much. Serving several requests in one batch reuses each weight read for several tokens, and storing weights in 4 bits instead of 16 cuts the bytes read per step by roughly four.

The useful concept here is arithmetic intensity, which is the number of floating-point operations per byte moved. At batch size 1, decode performs about one multiply-add per weight read, which is a very low intensity. A GPU like the H100 needs a few hundred operations per byte before compute becomes the limit, so small-batch decode stays memory-bound.

How fast does the KV cache grow?

During generation, each layer stores the keys and values of every previous token so they are not recomputed. That store is the KV cache. Its size is the product below, where the leading 2 counts keys and values separately. The script prints 2.0 and 0.5, the size in GiB of one sequence with and without grouped-query attention.

def kv_cache_bytes(layers, kv_heads, head_dim, seq_len, batch, bytes_per_value=2):
    return 2 * layers * kv_heads * head_dim * seq_len * batch * bytes_per_value

full = kv_cache_bytes(layers=32, kv_heads=32, head_dim=128, seq_len=4096, batch=1)
gqa = kv_cache_bytes(layers=32, kv_heads=8, head_dim=128, seq_len=4096, batch=1)
print(full / 2**30, gqa / 2**30)

For a 7B-class configuration with 32 layers and 32 heads of dimension 128, each token costs 0.5 MiB of cache in 16-bit precision, so one 4,096-token sequence needs 2 GiB. Thirty such requests would need about 60 GiB for the cache alone. Grouped-query attention with 8 key-value heads cuts that by four, which is why many recent models use it.

Which metrics matter in LLM inference?

The core LLM inference metrics are time to first token, time per output token, throughput, memory use, cost per token and output quality. They trade off against each other, so larger batches raise throughput but can slow each request. Always report which metric you optimized and what happened to the others.

Report latency as a distribution, such as the median and the 99th percentile, rather than as a single average. Tail latency is often what breaks a service-level target.

What should you know before learning inference engineering?

Before learning inference engineering, you should be comfortable with PyTorch, understand the transformer well enough to implement attention, and know basic GPU concepts such as memory bandwidth and parallel threads. Linear algebra and basic profiling help. You do not need to be a CUDA expert to start, but you will need some kernel knowledge later.

• You need PyTorch and a firm grip on tensor shapes, because most inference work starts in PyTorch before it moves to kernels.
• You need transformer internals, meaning attention, positional encodings, normalization layers and the feed-forward block. Cracking Deep Learning covers the fundamentals.
• You need GPU basics, such as threads, the memory hierarchy and why moving memory is expensive. The early sections of the CUDA Kernels plan cover them, and our guide on how to learn CUDA gives an order.
• You need measurement habits, which means timing code correctly with warm-up runs and repeated measurements.

What is the best order to learn inference engineering?

The best order is to build an LLM's inference path from the inside out. Start with a correct forward pass, add token generation, then the KV cache, then measure where time goes. Only then optimize, moving through kernels, quantization, batching and scheduling, and finally speculative decoding and multi-GPU serving. Each stage depends on the previous one.

Learning in this order means every optimization has a baseline to beat and a reference to check against. It also mirrors how inference engines are built, with correctness first, then memory, then speed.

Stage 1: Build a correct forward pass

Build a decoder-only transformer that turns token IDs into next-token logits, with embeddings, positional information, attention, normalization and the feed-forward block. Every later optimization must preserve this output, and without a verified forward pass you cannot tell a speedup from a bug.

To verify it, load released weights for a small open model and match a reference implementation's logits on the same tokens. Then measure perplexity on held-out text, which relies on the cross-entropy loss. Build LLM from Scratch covers language modeling end to end, and the Research section rebuilds Attention Is All You Need stage by stage, starting with positional encoding.

Stage 2: Add autoregressive generation and sampling

Build a generation loop that feeds each new token back in, with greedy decoding, temperature, top-k and top-p (nucleus) sampling, and stop conditions. Generation is where inference cost comes from. Sampling also affects quality, and some optimizations must preserve the sampling distribution exactly.

To verify it, check that greedy decoding matches a reference implementation token for token. For sampling, fix the random seed and check that the probabilities sum to one after filtering. Greedy decoding is an argmax over the vocabulary and sampling starts with a stable softmax, so you can write both as kernels in the Argmax and Softmax problems.

Stage 3: Add a KV cache

Build a per-layer cache of keys and values. Prefill fills it for the whole prompt in one pass, and each decode step computes keys and values only for the new token, appends them and attends over the full cache. Without a cache, each step recomputes keys and values for the entire prefix, so generation cost grows quadratically. The cache trades memory for compute, and that memory becomes the main limit on how many requests fit on a GPU.

To verify it, confirm that cached and uncached generation give the same outputs. In this single-head NumPy example, no_cache recomputes keys and values for the whole prefix at every step, while the loop below it computes each token's key and value once and appends them. The script prints True.

import numpy as np

def attend(q, K, V):
    s = K @ q / np.sqrt(q.shape[-1])
    w = np.exp(s - s.max())
    return (w / w.sum()) @ V

d = 64
rng = np.random.default_rng(0)
Wq, Wk, Wv = (rng.normal(size=(d, d)) / np.sqrt(d) for _ in range(3))
x = rng.normal(size=(10, d))

no_cache = [attend(x[t] @ Wq, x[:t + 1] @ Wk, x[:t + 1] @ Wv) for t in range(len(x))]

K_cache, V_cache, with_cache = [], [], []
for t in range(len(x)):
    K_cache.append(x[t] @ Wk)
    V_cache.append(x[t] @ Wv)
    with_cache.append(attend(x[t] @ Wq, np.array(K_cache), np.array(V_cache)))

print(np.allclose(no_cache, with_cache))

One detail catches many learners. With rotary position embeddings, keys are rotated by their position before they are cached, so each cached key keeps the position it was created at. The Inference Engineering plan covers KV-cache operations in PyTorch.

Stage 4: Shrink the cache with attention variants

Build multi-query attention (MQA), where all query heads share one key-value head, and grouped-query attention (GQA), where groups of query heads share key-value heads. Then add sliding-window attention, which keeps only the most recent tokens in the cache.

The KV cache scales with the number of key-value heads, and the calculation above shows GQA with 8 key-value heads cutting cache size by four compared with 32. Many recent open models, including Llama 3 and Mistral 7B, use GQA. DeepSeek-V2 introduced multi-head latent attention, which compresses keys and values into a smaller latent vector. To verify your versions, check cache size against the formula and confirm that a GQA layer with as many key-value heads as query heads reproduces standard multi-head attention. The Inference Engineering plan covers these attention variants.

Stage 5: Measure before you optimize

Build a benchmark harness that reports TTFT, TPOT, throughput and peak memory across a fixed set of prompt lengths, output lengths and batch sizes. Optimization without measurement is guesswork. The harness also tells you which phase dominates, since long prompts make prefill matter and long outputs make decode matter.

To verify the harness, run each configuration several times after warm-up runs and check that the results are stable. Synchronize the GPU before stopping a timer, because GPU work runs asynchronously. Every CUDA and Triton problem on TensorTonic reports kernel timings across several input sizes after a correct submission, and you can compare the same kernel on several GPUs, from a T4 up to an H100 and H200, which shows how differences in bandwidth and compute change the result.

Stage 6: Optimize the kernels that dominate

Build faster versions of the operations that take the most time, which are matrix-vector products for decode, matrix multiplications for prefill, attention, normalization and activations. Fuse small operations to avoid extra round trips to memory.

At small batch sizes, decode is dominated by matrix-vector products that stream weights from memory, while prefill is dominated by matrix multiplications. Attention over long contexts adds its own memory traffic, which FlashAttention reduces by computing attention in on-chip tiles with an online softmax. Verify each kernel by matching a reference implementation's outputs within a tolerance, then measure speedups across sizes. The CUDA Kernels plan covers the operations in an LLM's decode path.

• Matrix-Vector Multiplication is the core of batch-1 decode.
• Matrix Multiplication and Tiled Matrix Multiplication are the core of prefill.
• RMS Normalization and Layer Normalization run in every block.
• Swish is the activation inside SwiGLU feed-forward layers.
• Scaled Dot-Product Attention is the starting point for fused attention.

The Triton Kernels plan adds fused kernels such as online softmax and FlashAttention forward.

Stage 7: Quantize weights and the KV cache

Build weight-only quantization to 8 and 4 bits, then KV-cache quantization. Start with simple per-channel or per-group rounding before trying methods such as GPTQ or AWQ, which choose quantized values to reduce error. Because small-batch decode is memory-bound, fewer bytes per weight means faster tokens as well as more free memory, and quantizing the KV cache lets more requests fit at once.

Verify quality as well as speed. Compare perplexity and task accuracy with the 16-bit model, and check memory use against the figure you expected. A faster model that answers worse is not an optimization. The Inference Engineering plan covers quantization in PyTorch.

Stage 8: Batch and schedule many requests

Build static batching first, then continuous batching, where the scheduler adds and removes requests at every decode step instead of waiting for a whole batch to finish. Add a paged KV cache that allocates memory in fixed-size blocks, and prefix caching for shared prompt prefixes.

Batching raises arithmetic intensity, which turns memory-bound decode into better GPU utilization. Continuous batching stops short requests from waiting for long ones, and paged KV caches reduce memory lost to fragmentation so that more requests fit at once. To verify a scheduler, plot throughput against latency as load increases. A good one raises throughput while keeping TTFT and TPOT within a target. You can practice designing these serving systems on ML System Design, where you lay out queues, schedulers, caches and capacity on an interactive canvas.

Stage 9: Add speculative decoding

In speculative decoding, a small draft model proposes several tokens and the large target model checks them in one forward pass, accepting the longest correct run. The target model then processes several tokens per step instead of one, which uses the GPU's spare compute during memory-bound decode. With the acceptance rule from the original papers, the output distribution matches the target model exactly.

To verify it, confirm that greedy outputs are identical with and without speculation, and measure the acceptance rate and speedup at different draft lengths.

Stage 10: Serve across multiple GPUs

Learn tensor parallelism, which splits each layer's matrices across GPUs, and pipeline parallelism, which assigns different layers to different GPUs. Large models do not fit on one GPU, and splitting them changes latency, throughput and communication costs. To verify a setup, measure how throughput and latency change as you add GPUs and find the point where communication starts to dominate. ML System Design lets you work through multi-GPU serving setups, including how many GPUs a target load needs.

How do you know if an inference optimization worked?

An inference optimization worked if it improves the target metric on a realistic workload, keeps the other metrics within limits, and leaves output quality unchanged or acceptably close. Measure before and after with the same prompts, lengths, batch sizes and hardware, using warm-up runs and repeated measurements. A reliable benchmarking routine follows these steps.

1. Fix the workload by using the same prompts, input and output lengths, and batch sizes for every comparison.
2. Warm up first, because the first runs include compilation and memory allocation.
3. Synchronize before timing, because GPU calls return before the work finishes.
4. Repeat each run and report medians and tail percentiles instead of a single number.
5. Vary the inputs, since an optimization that helps short prompts may hurt long ones.
6. Check quality every time by comparing outputs or perplexity with the unoptimized baseline.
7. Test on more than one GPU, because bandwidth and compute differ and a speedup on one GPU may not transfer.

Which papers explain inference engineering, in build order?

The papers below introduced the core techniques, listed in the same order as the build stages. Read each one when you reach its stage rather than before.

If you want to practice implementing papers in general, our guide on how to read an ML paper and implement it explains the workflow.

What does an 8-week plan to learn inference engineering look like?

This suggested plan follows the build order at about 10 hours a week. Treat it as a guide rather than a guarantee, because stages such as kernels and batching can take longer, especially if you are new to CUDA.

Which inference engineering projects belong in a portfolio?

The strongest inference engineering portfolio projects measure trade-offs instead of just running a model. A KV cache and quantization study, a kernel benchmarked across GPUs, or a small serving engine with continuous batching each show that you can make a model faster and prove it without breaking its outputs.

• A latency, memory and quality study measures how a KV cache, GQA and quantization change TTFT, TPOT, memory and perplexity for one small model, and presents the results in a single table.
• A decode-path kernel suite has naive and optimized versions of GEMV, RMS normalization and attention kernels, timed across several GPUs.
• A mini serving engine adds continuous batching and a paged KV cache to your own model and shows a throughput-versus-latency plot.
• A speculative decoding study reports acceptance rates and speedups across draft lengths, with a check that greedy outputs are unchanged.

Our guide on how to build an ML portfolio covers how to present projects like these.

How do interviews test inference engineering skills?

Interviews for AI engineer and GPU performance roles often test inference through explanation, calculation and design. Expect to explain prefill versus decode, estimate KV-cache memory, reason about memory-bound decode, implement attention or a KV cache, and design a serving system for a target load.

Typical questions ask how much memory the KV cache needs for a given model and context length, why batching improves throughput but not always latency, and how you would serve a model to a given number of users within a latency target. TensorTonic's timed mock assessments include an LLM Internals round for Research Engineer candidates, and ML System Design covers the serving design round.

What are the common mistakes when learning inference engineering?

• Many learners optimize before measuring, so build the benchmark harness first.
• Timing without synchronizing measures the wrong thing, because GPU work is asynchronous.
• Reporting only averages hides tail latency, which often matters more.
• Ignoring quality turns a speedup into a regression, so every quantization or approximation needs a quality check.
• Treating prefill and decode as one workload leads to the wrong fix, because they have different bottlenecks.
• Benchmarking a single configuration misleads you, since results change with prompt length, output length, batch size and GPU.
• Skipping the fundamentals makes kernels and schedulers hard to reason about without a verified forward pass and KV cache underneath.

Practice this

Sources

• Vaswani et al., “Attention Is All You Need”, arXiv:1706.03762 (2017).
• Su et al., “RoFormer: Enhanced Transformer with Rotary Position Embedding”, arXiv:2104.09864 (2021).
• Holtzman et al., “The Curious Case of Neural Text Degeneration”, arXiv:1904.09751 (2019).
• Shazeer, “Fast Transformer Decoding: One Write-Head is All You Need”, arXiv:1911.02150 (2019).
• Ainslie et al., “GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints”, arXiv:2305.13245 (2023).
• DeepSeek-AI, “DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model”, arXiv:2405.04434 (2024).
• Dao et al., “FlashAttention”, arXiv:2205.14135 (2022), and Dao, “FlashAttention-2”, arXiv:2307.08691 (2023).
• Dettmers et al., “LLM.int8()”, arXiv:2208.07339 (2022).
• Xiao et al., “SmoothQuant”, arXiv:2211.10438 (2022).
• Lin et al., “AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration”, arXiv:2306.00978 (2023).
• Frantar et al., “GPTQ: Accurate Post-Training Quantization for Generative Pre-trained Transformers”, arXiv:2210.17323 (2022).
• Yu et al., “Orca: A Distributed Serving System for Transformer-Based Generative Models”, OSDI (2022).
• Kwon et al., “Efficient Memory Management for Large Language Model Serving with PagedAttention”, SOSP (2023), arXiv:2309.06180.
• Leviathan, Kalman and Matias, “Fast Inference from Transformers via Speculative Decoding”, arXiv:2211.17192 (2023).
• Chen et al., “Accelerating Large Language Model Decoding with Speculative Sampling”, arXiv:2302.01318 (2023).
• Shoeybi et al., “Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism”, arXiv:1909.08053 (2019).
• NVIDIA, H100 Tensor Core GPU specifications, for the memory bandwidth of the SXM variant.

The code outputs above were produced by running the snippets as shown. Hardware figures are approximate and ignore overheads.

Related guides

• How to learn CUDA for machine learning
• CUDA practice problems: from vector add to attention
• CUDA interview questions for ML engineers
• ML system design interview practice