AI Inference Engineering: Detailed Explanation with Examples

1. Overview

When you call a chatbot’s API and wait for its response to appear gradually on your screen, you’re witnessing the inference process of a Large Language Model (LLM). This article explains how inference differs from training, why it consumes enormous resources, and the techniques engineers use to make it faster and cheaper: KV cache, continuous batching, quantization, speculative decoding. By the end, you’ll understand the numbers TTFT/TPOT in benchmarks, and know how to write Java code to measure performance of a real inference endpoint.

Note: This article doesn’t dive deep into transformer internals (attention, feed-forward…). The focus is on infrastructure perspective — how a request “flows” through an inference system.

2. What is Inference?

Training is the process of learning weights from data — costing thousands of GPUs over many weeks, but done only once. Inference is after the model has learned; you feed a prompt and get output — this happens every time there’s a request, with millions of requests daily in production. That’s why, even though training costs more in total spend once, inference determines long-term operational costs and user experience (latency).

Key difference: during training, the model sees the entire sentence to calculate loss at once. During text generation inference, the model must generate one token at a time, and each token depends on the previous one — called autoregressive decoding.

2.1. Two Phases: Prefill and Decode

Every inference request goes through 2 phases:

• Prefill: the model reads the entire input prompt at once, computing attention for all input tokens. This phase leverages GPU parallelism well because it’s compute-bound.
• Decode: the model generates output tokens sequentially, one token per step, based on all previous context. This phase runs sequentially; each step only computes for 1 new token (memory-bound — bottlenecked by memory bandwidth, not computation).

Here’s a simple Java illustration of a basic decode loop (simulated, not a real model) to see the sequential nature:

public class AutoregressiveDecoder {

    public List<String> generate(List<String> prompt, int maxNewTokens) {
        List<String> context = new ArrayList<>(prompt);
        List<String> generated = new ArrayList<>();

        for (int step = 0; step < maxNewTokens; step++) {
            String nextToken = predictNextToken(context); // 1 forward pass
            if (nextToken.equals("<EOS>")) break;

            context.add(nextToken);   // new token becomes input for next step
            generated.add(nextToken);
        }
        return generated;
    }

    private String predictNextToken(List<String> context) {
        // simulation: real model runs forward pass through entire transformer
        return "token_" + context.size();
    }
}

Each call to predictNextToken, the model must re-read all existing context to compute attention — without optimization, cost grows with sentence length. This is the motivation behind KV cache in section 4.

3. Measuring Inference Performance: TTFT and TPOT

The two most important metrics when benchmarking an LLM:

• TTFT (Time To First Token): time from sending request to receiving the first token — reflects prefill phase latency.
• TPOT (Time Per Output Token): average time to generate each subsequent token during the decode phase.

public record InferenceMetrics(Duration ttft, Duration avgTpot, int totalTokens) {

    public Duration totalLatency() {
        return ttft.plus(avgTpot.multipliedBy(totalTokens - 1));
    }
}

The totalLatency() method shows the real formula: total time ≈ TTFT (1 pass) + TPOT × (remaining tokens). With a 200-token response, TPOT slowing by just 10ms/token adds an extra 2 seconds total — that’s why TPOT gets more scrutiny than TTFT in serving benchmarks.

Example real metrics (7B model, batch size 1, A100 GPU):
TTFT  ≈ 80ms
TPOT  ≈ 15ms/token
200 token → 80ms + 15ms × 199 ≈ 3.07s

4. KV Cache: Avoid Recomputing What You Already Know

Without caching, each decode step must recompute attention for the entire token sequence from the beginning — complexity grows quadratically with length. KV cache stores the key/value vectors of already-processed tokens, so the next step only computes for the new token, then reuses the old cache.

public class KvCache {

    private final Map<Integer, float[]> keyCache = new HashMap<>();
    private final Map<Integer, float[]> valueCache = new HashMap<>();

    public void store(int position, float[] key, float[] value) {
        keyCache.put(position, key);
        valueCache.put(position, value);
    }

    public List<float[]> allKeysUpTo(int position) {
        return IntStream.rangeClosed(0, position)
            .mapToObj(keyCache::get)
            .toList();
    }
}

store is called once per new token generated; allKeysUpTo lets the next decode step retrieve all previously-computed keys without re-running forward pass from scratch. This is why a long request (long context) drains VRAM quickly — the cache scales with batch_size × sequence_length × num_layers × hidden_size.

Note: KV cache is the main reason why long context windows (32K, 128K tokens) consume enormous RAM on GPU, not computation. This is also why techniques like paged attention (used in vLLM) exist — managing cache like virtual memory (pages) to avoid fragmentation.

5. Continuous Batching: Group Requests to Increase Throughput

GPU processing by batches is more efficient than handling single requests one at a time. But each request has different output length — request A finishes at token 50, request B continues until token 300. Static batching (waiting for entire batch to finish before accepting the next) wastes GPU because it waits for the slowest request. Continuous batching solves this by: the moment a request in the batch completes, its “slot” is immediately replaced with a new request from the queue, no need to wait for the whole batch.

public class ContinuousBatchScheduler {

    private final int batchSlots;
    private final Queue<InferenceRequest> waiting = new LinkedList<>();
    private final List<InferenceRequest> active;

    public ContinuousBatchScheduler(int batchSlots) {
        this.batchSlots = batchSlots;
        this.active = new ArrayList<>(Collections.nCopies(batchSlots, null));
    }

    public void onStep() {
        for (int slot = 0; slot < batchSlots; slot++) {
            InferenceRequest req = active.get(slot);
            if (req == null || req.isFinished()) {
                active.set(slot, waiting.poll()); // fill slot immediately with next request
            } else {
                req.decodeNextToken();
            }
        }
    }
}

onStep() runs each decode step: if any slot is empty or request just finished, it immediately takes the next waiting request from waiting. This way, GPU always runs at full batch capacity, no “idle” slots waiting for slow requests — this is the core mechanism in vLLM, TGI (Text Generation Inference), and TensorRT-LLM.

6. Quantization: Trade Accuracy for Speed

Original models typically store weights in FP16 or BF16 (16-bit). Quantization reduces storage bits (INT8, INT4) to lower VRAM usage and increase throughput, in exchange for slightly reduced accuracy.

7B parameter model, storage comparison:
- FP16: ~14 GB VRAM
- INT8: ~7 GB VRAM
- INT4: ~3.5 GB VRAM

When serving with engines like vLLM, you just declare the configuration:

model: meta-llama/Llama-3-8B
quantization: awq
dtype: float16
gpu_memory_utilization: 0.9

quantization: awq specifies using the AWQ method (Activation-aware Weight Quantization) — an INT4 technique that preserves accuracy better than naive quantization by prioritizing keeping the most important weights intact.

Note: Very aggressive quantization (INT4 or lower) can increase model hallucination or reduce reasoning ability. Always benchmark output quality (not just speed) after quantizing, especially for quality-sensitive tasks like code generation or calculation.

7. Speculative Decoding: Predict Ahead to Save Steps

Sequential decode is slow because each step only generates 1 token. Speculative decoding uses a smaller, faster model (“draft model”) to predict several next tokens, then the large model (“target model”) only needs to verify those tokens in one forward pass instead of many.

Draft model predicts: "Ha" "Noi" "la" "thu" "do"
Target model verifies all 5 tokens at once in 1 forward pass:
  - If 4/5 tokens match the large model's distribution → keep 4 tokens, drop token 5, decode from there.
  - If wrong from start → costs exactly 1 forward pass like normal decode, no waste.

This technique doesn’t reduce output quality (because always verified by target model), just reduces how many sequential forward passes you need to run on the large model — most effective when draft model predicts correctly often.

8. Practice: Measuring TTFT/TPOT from an LLM Serving Endpoint

Suppose you have a vLLM server running OpenAI-compatible API at http://localhost:8000. Add WebClient dependency:

<dependency>
    <groupId>org.springframework.boot</groupId>
    <artifactId>spring-boot-starter-webflux</artifactId>
</dependency>

Write a client to measure TTFT using streaming response:

@Service
public class LlmBenchmarkClient {

    private final WebClient webClient;

    public LlmBenchmarkClient(WebClient.Builder builder) {
        this.webClient = builder.baseUrl("http://localhost:8000").build();
    }

    public InferenceMetrics measure(String prompt) {
        Instant start = Instant.now();
        AtomicReference<Instant> firstTokenTime = new AtomicReference<>();
        AtomicInteger tokenCount = new AtomicInteger();

        webClient.post()
            .uri("/v1/completions")
            .bodyValue(Map.of("prompt", prompt, "stream", true, "max_tokens", 100))
            .retrieve()
            .bodyToFlux(String.class)
            .doOnNext(chunk -> {
                if (firstTokenTime.get() == null) {
                    firstTokenTime.set(Instant.now());
                }
                tokenCount.incrementAndGet();
            })
            .blockLast();

        Duration ttft = Duration.between(start, firstTokenTime.get());
        Duration total = Duration.between(start, Instant.now());
        Duration avgTpot = tokenCount.get() > 1
            ? total.minus(ttft).dividedBy(tokenCount.get() - 1)
            : Duration.ZERO;

        return new InferenceMetrics(ttft, avgTpot, tokenCount.get());
    }
}

doOnNext is called each time a new chunk (token) arrives — the first call records firstTokenTime, which is exactly TTFT. After stream completes (blockLast), we calculate average TPOT as total time minus TTFT, divided by remaining tokens.

8.1. Testing with JUnit 5 and AssertJ

@SpringBootTest
class LlmBenchmarkClientTest {

    @Autowired
    LlmBenchmarkClient client;

    @Test
    void shouldMeetLatencySla() {
        InferenceMetrics metrics = client.measure("Explain KV cache in 2 sentences.");

        assertThat(metrics.ttft()).isLessThan(Duration.ofMillis(500));
        assertThat(metrics.avgTpot()).isLessThan(Duration.ofMillis(50));
        assertThat(metrics.totalTokens()).isGreaterThan(0);
    }
}

This test doesn’t check response content, only latency SLA — exactly on-target for “inference engineering”: ensuring the system responds fast enough, not evaluating model quality (that’s a separate evaluation process outside this article’s scope).

Note: This test needs a real LLM server running (or mock WebClient with MockWebServer). In CI, separate such tests into integration tests, not running with regular unit tests, since they depend on external infrastructure.

9. Common Mistakes and When NOT to Build Your Own Inference Engine

• Rewriting batching/KV cache from scratch: unless you’re doing specialized ML infrastructure research, don’t rebuild what vLLM, TGI, TensorRT-LLM have already optimized — you’ll spend months just to reach their existing performance.
• Benchmarking only throughput, ignoring tail latency: high throughput but bad p99 latency still ruins user experience. Always measure p50/p95/p99, not just average.
• Quantizing without testing quality again: as mentioned in section 6, reducing bits always has tradeoffs — measuring speed without measuring quality is a common mistake.
• Using static batching for variable real-world traffic: if requests vary greatly in output length, static batching wastes GPU seriously — choose an engine with continuous batching from the start.
• Not separating prefill-heavy and decode-heavy workloads: some production systems run prefill and decode nodes separately (disaggregated serving) to optimize resources — if your traffic has clear patterns (e.g., RAG with very long prompt, short output), consider this architecture.

10. Conclusion

Inference engineering is an optimization problem between latency, throughput, and GPU cost — not about training models. Understanding prefill/decode, KV cache, continuous batching, quantization, and speculative decoding helps you read benchmarks, choose the right serving engine (vLLM, TGI…), and measure your system with correct metrics (TTFT, TPOT) instead of just eyeballing “is it fast or slow”. The next step naturally is to try changing quantization or gpu_memory_utilization in actual config and observe how these metrics change on your own workload.