AI Engineering: Inference Optimization

Inference optimization focuses on making machine learning models in production faster, cheaper, and computationally efficient without degrading response quality. While model development prioritizes training metrics and predictive accuracy, production environments demand strict latency guarantees, high throughput, and cost-effective compute allocation. This chapter systematically covers the architectural, algorithmic, hardware, and system-level techniques required to optimize large language model (LLM) inference workloads.

1. Fundamentals of Inference Optimization

In a production environment, predictive capability alone does not guarantee utility. A model that delivers an accurate prediction outside an acceptable operational window fails to provide business value.

Production Latency Failure Case:
Consider a model trained to predict financial stock prices for the next trading day. If the end-to-end inference latency of the serving pipeline causes the prediction to arrive two days late, the output becomes useless regardless of how accurate the underlying prediction was.

Inference optimization targets three interrelated dimensions:

• Latency: Minimizing the time taken to generate a complete response for a user request.
• Throughput: Maximizing the aggregate volume of tokens or requests served across the infrastructure per unit of time.
• Cost: Reducing the operational expense (compute hardware, memory, power) incurred per transaction or token generated.

The Economic Shift: Closed-Source APIs vs. Self-Hosted Open-Source

For proprietary foundational model providers, infrastructure economics dictates that aggregate inference revenue must exceed aggregate training amortized costs:

Where:

• $T$ is the total cost of training the base model.
• $P$ is the pricing charged per inference call.
• $N$ is the total volume of inference calls executed over the model lifecycle.

For organizations leveraging open-source foundational models (such as LLaMA, Mistral, or Falcon), the capital expenditure of initial pre-training is zero ($T \approx 0$). The engineering objective shifts toward designing custom inference servers that beat commercial API costs by maximizing utilization and minimizing serving latency.

2. Hardware Bottlenecks: Compute-Bound vs. Memory Bandwidth-Bound

Inference performance is constrained by physical hardware limitations. Every computational workload executed on modern accelerators encounters one of two primary bottlenecks: compute bounds or memory bandwidth bounds.

Computational Definitions

• Compute-Bound Tasks: Operations where the total execution duration is limited by the raw mathematical processing speed of the compute units (floating-point operations per second, or FLOPs). The processor operates near peak capacity, waiting very little for memory transfers.
• Memory Bandwidth-Bound Tasks: Operations where execution duration is governed by the speed at which data can be transported between storage memory (e.g., CPU DRAM or GPU High Bandwidth Memory) and processing registers (on-chip SRAM). The compute cores sit idle while waiting for weights and activations to arrive.

Arithmetic Intensity and the Roofline Model

The dividing line between these two performance constraints is quantified by Arithmetic Intensity:

The Roofline Model visualizes system performance by plotting attainable computational performance against arithmetic intensity:

• When arithmetic intensity is below the inflection threshold, throughput is limited by the slope of the line, which represents memory bandwidth. Hardware speedups require chips with higher memory bandwidth.
• When arithmetic intensity is above the threshold, throughput plateaus at the hardware's theoretical peak compute ceiling. Hardware speedups require chips with higher raw FLOPs capabilities.

3. Dissecting Transformer Inference: Prefill vs. Decode Stages

Autoregressive transformer inference consists of two structurally distinct phases: the Prefill stage and the Decode stage.

The Prefill Phase (Prompt Processing)

During prefill, the model consumes the entire sequence of input prompt tokens ($X_1, \dots, X_T$) at once:

• Parallel Execution: Because the full prompt context is available, self-attention and feed-forward layers process all prompt tokens concurrently via dense matrix-matrix multiplications ($\text{GEMM}$).
• Compute-Bound Characteristics: Processing hundreds or thousands of tokens simultaneously yields high arithmetic intensity, fully occupying modern tensor cores. Execution speed is dictated by peak tensor FLOP capacity.

The Decode Phase (Token Generation)

Once the prompt is prefilled, generation shifts to an autoregressive loop generating one token at a time:

• Sequential Execution: Generating output token $X_{t+1}$ strictly depends on all prior tokens ($X_1, \dots, X_t$).
• Memory Bandwidth-Bound Characteristics: For every single token generated, the model reads billions of parameter weights from high-bandwidth device memory (HBM) into fast on-chip memory (SRAM), executes a matrix-vector product ($\text{GEMV}$), and outputs a single token logit.
• Because the parameters are loaded to produce only one token per request, arithmetic intensity drops drastically. The memory bus saturates while compute cores remain largely underutilized.

Production Architectural Decoupling

Because Prefill and Decode exhibit opposite performance profiles, hosting both workloads on identical physical nodes leads to resource contention and suboptimal hardware utilization.

Production inference systems increasingly decouple these phases across physically distinct clusters:

• Prefill Clusters: Configured with high-FLOP nodes optimized for parallel dense computation.
• Decode Clusters: Configured with high-bandwidth memory nodes optimized for high transfer speeds and memory capacity.

4. Serving Paradigms and Key Performance Metrics

Online vs. Batch Inference APIs

Inference traffic generally falls into two serving paradigms:

• Online Serving (Latency-Prioritized): Designed for real-time human interaction (e.g., interactive chat assistants). Requests are processed immediately upon arrival to minimize responsiveness metrics.
• Batch Serving (Cost- and Throughput-Prioritized): Designed for asynchronous workloads (e.g., offline summarization, bulk data labeling). Requests are accumulated and grouped into large processing batches, trading higher latency for optimal GPU core saturation and reduced per-token processing costs.

Streaming Mode and the Validation Trade-Off

Online APIs often implement token streaming, emitting tokens over server-sent events (SSE) the moment they are generated:

• Advantage: Substantially reduces perceived latency by giving the end user immediate visual feedback.
• Trade-Off: Prevents post-generation validation or moderation checks. In non-streaming configurations, a system can generate five tokens, evaluate them for coherence or safety, and rerun generation if thresholds fail. Once an initial token streams to the client, the generation path cannot be undone.

Latency Decomposition: TTFT and TPOT

End-to-end latency cannot be treated as an atomic scalar value; it is broken into two distinct phases:

The mathematical formula for total generation latency is:

Where:

• $\text{TTFT}$ (Time to First Token) corresponds directly to the duration of the prefill stage plus network queuing. It scales with input sequence length.
• $\text{TPOT}$ (Time Per Output Token) represents the time delta between consecutive generated tokens during the decode phase.
• $N_{\text{output}}$ is the total number of output tokens produced.

The Misleading Nature of Averages: Percentile Profiling

Reporting latency via arithmetic means ($\mu$) obscures tail-end performance degradation caused by memory stalls, dynamic batch reconfigurations, or network variations.

Illustrative Failure of Averages:
Assume an inference node processes 10 requests. Nine requests complete in $100\text{ ms}$ each. The tenth request encounters a network partition or large context burst and takes $3,000\text{ ms}$.

• Arithmetic Mean: $\frac{(9 \times 100) + 3000}{10} = 390\text{ ms}$

The average ($390\text{ ms}$) fails to represent the standard user experience ($100\text{ ms}$) while minimizing the severity of the tail latency ($3,000\text{ ms}$).

Inference engineering mandates using non-parametric percentile metrics:

• $P_{50}$ (Median): Represents the latency threshold under which $50\%$ of requests complete.
• $P_{90} / P_{95}$: Identifies operational latency for standard degraded executions.
• $P_{99}$: Directly exposes tail latency bottlenecks and hardware resource contention under edge conditions.

5. System Efficiency Metrics: Throughput, Goodput, MFU, and MBU

Throughput vs. Latency Dynamics

• Throughput measures the volume of work an inference service completes per unit of time, quantified as Tokens Per Second (TPS) or Requests Per Minute (RPM).
• System throughput and single-user latency operate in an inverse trade-off curve. Processing requests concurrently across batched tensor inputs increases total system throughput (improving cost efficiency) but increases individual user latency due to hardware queuing and memory multiplexing.

Goodput and Service-Level Objectives (SLOs)

Raw throughput counts all tokens produced by the cluster. Goodput counts only the subset of requests that successfully meet predefined Service Level Objectives (SLOs):

If an inference system processes $100\text{ RPM}$, but individual operational constraints dictate that $\text{TTFT} \le 200\text{ ms}$ and $\text{TPOT} \le 50\text{ ms}$, any request breaching these bounds is excluded from the goodput metric. A system processing $100\text{ RPM}$ where only 30 requests meet the latency SLOs has a goodput of $30\text{ RPM}$.

Hardware Utilization: MFU vs. MBU

Standard monitoring utilities (such as nvidia-smi) report time-based GPU utilization, indicating what percentage of a sampling window the GPU kernel was active. This is misleading: a GPU running an unoptimized kernel that executes one operation per cycle reports $100\%$ temporal utilization despite wasting hardware compute resources.

Rigorous efficiency demands tracking Model FLOPs Utilization (MFU) and Model Bandwidth Utilization (MBU):

1. Model FLOPs Utilization (MFU)

Measures the ratio of observed computational throughput against theoretical peak floating-point capabilities:

2. Model Bandwidth Utilization (MBU)

Measures the ratio of memory moved against the physical memory bandwidth ceiling:

Where achieved memory bandwidth during decoding is approximated by:

Where:

• $P_{\text{count}}$ is the total number of active model parameters.
• $B_{\text{param}}$ is the bytes allocated per parameter (e.g., $2\text{ bytes}$ for FP16, $1\text{ byte}$ for INT8).

Compute Workload (Prefill):       High MFU  | Low MBU
Memory-Bound Workload (Decode):   Low MFU   | High MBU

6. AI Accelerators and the Memory Hierarchy

Modern machine learning workloads run predominantly on dedicated accelerators (e.g., GPUs, TPUs) designed for high-density matrix multiplication. Unlike general-purpose CPUs with large caches and low core counts, AI accelerators prioritize massive parallelism via thousands of compact arithmetic units.

Accelerator Memory Tiers

1. Host CPU DRAM: Holds system memory. Transferring model activations or weights from CPU DRAM to the GPU over the PCIe bus introduces major latency, with bandwidth limited to tens of GB/s.
2. GPU Device Memory (HBM/GDDR): High-density memory located directly on the accelerator substrate. Modern HBM3 interfaces achieve bandwidths exceeding $3\text{ TB/s}$, yet this remains a primary bottleneck during decode phases.
3. GPU On-Chip SRAM (L1, L2, Register Files): Memory integrated directly within compute streaming multiprocessors. SRAM boasts bandwidths exceeding $10\text{--}20\text{ TB/s}$, but capacity is physically constrained to small allocations (e.g., $50\text{--}100\text{ MB}$).

Standard frameworks (PyTorch, TensorFlow) often introduce intermediate register writes to HBM between successive neural network operations. Advanced GPU programming languages (such as CUDA or OpenAI Triton) give engineers direct control over this memory hierarchy, enabling fused computational execution entirely within on-chip SRAM.

Power Considerations: Peak Draw vs. TDP

• Maximum Power Draw: The absolute instantaneous electrical power consumed under maximum computational stress.
• Thermal Design Power (TDP): The maximum continuous heat output that the physical server cooling apparatus is engineered to dissipate under prolonged workloads. Exceeding sustained cooling thresholds triggers thermal throttling, degrading GPU frequencies and causing unexpected latency spikes.

7. Model-Level Optimizations: Compression and Speculative Decoding

Model Compression: Quantization, Distillation, and Pruning

To decrease memory bandwidth pressure during autoregressive decoding, model sizes can be compressed:

1. Quantization: Converts parameter weights and activations from high-precision floating point formats to lower-bit representations:
   • FP16 ($16\text{ bits} = 2\text{ bytes}$) $\rightarrow$ INT8 ($8\text{ bits} = 1\text{ byte}$) $\rightarrow$ INT4 ($4\text{ bits} = 0.5\text{ bytes}$).
   • Performance Impact: Reducing bit precision decreases $B_{\text{param}}$, directly reducing memory bandwidth demand during the decode phase. This allows twice as many parameters to fit within the same memory throughput envelope without changing arithmetic operational counts.
2. Knowledge Distillation: Trains a compact "student" model to mimic the output distribution and hidden states of a large "teacher" model, compressing capacity into fewer parameters.
3. Weight Pruning: Removes parameter weights that contribute minimally to the loss surface. While historically effective for sparse neural networks, unstructured pruning is rarely used on modern dense tensor cores because non-contiguous memory access patterns cause memory alignment penalties.

Speculative Decoding

Speculative decoding bypasses the sequential bottleneck of autoregressive generation by pairing a small, fast draft model with a large, expressive target model.

Verification and Acceptance Mechanics

1. Let the target sequence be indexed up to step $T$. The lightweight draft model predicts $K$ speculative tokens autoregressively:

2. The large target model processes all $K$ proposed tokens simultaneously in parallel via a single prefill-like forward pass, computing exact joint likelihoods across the draft sequence.
3. The engine accepts the longest contiguous subsequence $J \le K$ (from left to right) where the draft predictions match or satisfy the target model's acceptance criteria.
4. The target model generates an additional replacement token $X_{T+J+1}$ conditioned on the validated prefix.
5. Efficiency Bounds:
   • Best Case: All $K$ tokens match; the system advances by $K + 1$ tokens in the time it takes the target model to run a single decode step.
   • Worst Case: The first draft token fails validation ($J = 0$); the target model rejects the draft sequence and emits its own verified token. The system advances by exactly 1 token, maintaining baseline autoregressive performance.

def speculative_decoding_step(draft_model, target_model, context, K):
    # 1. Draft model generates K speculative tokens sequentially
    draft_tokens = []
    curr_context = context.copy()
    for _ in range(K):
        next_token = draft_model.predict_next(curr_context)
        draft_tokens.append(next_token)
        curr_context.append(next_token)

    # 2. Target model evaluates the proposed sequence in parallel
    target_logits = target_model.evaluate_parallel(context, draft_tokens)
    
    # 3. Determine longest accepted contiguous prefix
    accepted_tokens = []
    for i, token in enumerate(draft_tokens):
        if target_model.verify(target_logits[i], token):
            accepted_tokens.append(token)
        else:
            # Rejection encountered: sample replacement token from target distribution
            correct_token = target_model.sample_from_logits(target_logits[i])
            accepted_tokens.append(correct_token)
            return accepted_tokens  # Early return on first divergence

    # If all K tokens accepted, sample K+1 bonus token from final target logit
    bonus_token = target_model.sample_from_logits(target_logits[-1])
    accepted_tokens.append(bonus_token)
    return accepted_tokens

Operational Failure Modes

Speculative decoding relies on exploiting idle GPU compute cores during memory-bandwidth-bound decode steps.

If an inference system runs at maximum concurrency under saturated batch sizes where MFU is already maxed out, spare computational headroom disappears. In that setting, running parallel verification on the target model introduces resource contention that can degrade rather than improve serving latency.

Parallel Decoding: Medusa

Rather than relying on an external secondary draft model, architectures like Medusa introduce multiple specialized prediction heads on top of the original target model base:

• Each specialized head is trained to predict future token offsets (e.g., Head 1 predicts $X_{t+1}$, Head 2 predicts $X_{t+2}$, Head $K$ predicts $X_{t+K}$).
• Candidates are evaluated using tree-based attention structures, and consistency is verified in parallel using the Jacobi iteration method, allowing speedups up to $1.9\times$ without maintaining a separate draft model pipeline.

8. Attention Mechanism Optimization and KV Caching

The KV Cache: Memory-for-Compute Trade-Off

In standard causal multi-head self-attention, generating token $X_{t}$ requires computing dot-product attention between the current token's query vector and the key/value vectors of all preceding tokens:

Without caching, computing token $X_t$ requires reprojecting $X_1, \dots, X_{t-1}$ to regenerate their corresponding key ($K$) and value ($V$) projections at every step. This repeats quadratic work over time, resulting in an algorithmic complexity of $O(N^2)$ across a sequence of length $N$.

The Key-Value (KV) Cache stores the computed key and value activation tensors for all historical tokens in device memory. Each subsequent decode step only needs to project the single newest token into $K$ and $V$, append those vectors to the cache, and evaluate attention over the cached history—reducing inference complexity to $O(N)$ operations.

Without KV Cache:
Step 1: Compute KV for [X_1]
Step 2: Recompute KV for [X_1, X_2]
Step 3: Recompute KV for [X_1, X_2, X_3]
Total Complexity: O(N^2)

With KV Cache:
Step 1: Compute and store KV for [X_1]
Step 2: Retrieve cached KV[X_1]   + Compute & append KV[X_2]
Step 3: Retrieve cached KV[X_1:2] + Compute & append KV[X_3]
Total Complexity: O(N)

Computing the Memory Footprint of the KV Cache

While the KV cache removes redundant projections, its memory usage grows linearly with both context length and batch size.

The memory footprint (in bytes) of the KV cache is calculated as:

Where:

• $2$ accounts for the two separate tensors stored: Key and Value.
• $B$ is the runtime serving batch size.
• $L$ is the cumulative sequence context length (input tokens $+$ generated tokens).
• $N_{\text{layers}}$ is the number of transformer layers in the model architecture.
• $d_{\text{model}}$ is the hidden dimensionality of the model.
• $P_{\text{bytes}}$ is the precision format storage cost (e.g., $2\text{ bytes}$ for FP16, $1\text{ byte}$ for FP8).

For large contexts and high batch counts, the KV cache can quickly exceed the memory footprint of the static model parameter weights themselves.

Architectural Variants: MHA, MQA, and GQA

To control the memory footprint of the KV cache, modern architectures modify the structural ratio between Query heads and Key-Value heads:

Multi-Head Attention (MHA)       Multi-Query Attention (MQA)      Grouped-Query Attention (GQA)
Each Q has a dedicated KV        All Qs share 1 single KV         Qs grouped; each group shares KV

Q1 Q2 Q3 Q4 Q5 Q6 Q7 Q8          Q1 Q2 Q3 Q4 Q5 Q6 Q7 Q8          Q1 Q2 Q3 Q4 Q5 Q6 Q7 Q8
 |  |  |  |  |  |  |  |           \  \  \  \  /  /  /  /           \  \  /  /   \  \  /  /
K1 K2 K3 K4 K5 K6 K7 K8                    KV1                        KV1          KV2
V1 V2 V3 V4 V5 V6 V7 V8
[Largest Cache Size]             [Smallest Cache Size]            [Balanced Footprint & Quality]

1. Multi-Head Attention (MHA): Every individual query head maintains its own independent key and value head. For $H$ attention heads, the KV cache maintains $H$ distinct sets of tensors. This maximizes representational capacity, but produces the largest memory footprint.
2. Multi-Query Attention (MQA): Collapses all $H$ query heads into a single shared key and value head. This reduces the KV cache size by a factor of $H\times$, greatly lowering memory bandwidth requirements. However, it can cause quality degradation on nuanced or information-dense prompts.
3. Grouped-Query Attention (GQA): Divides query heads into $G$ distinct functional groups, where each group shares a single set of key and value heads. For example, an architecture with 8 query heads might be split into 2 groups of 4 heads, each sharing one KV head. This balances the cache footprint with overall model quality, making it the standard choice in architectures such as LLaMA-2/3.

Advanced Attention Strategies

• Cross-Layer Attention: Adjacent transformer layers frequently learn correlated, near-identical attention patterns. Cross-layer attention shares KV caches across successive layers (e.g., Layer 1 and Layer 2 share one cache), cutting total allocation in half.
• Local Window Attention: Constrains token attention strictly to a local neighborhood window of length $W$, bounding the KV cache size to $W$ instead of the full sequence length $L$.
• Interleaved Attention: Alternates between local window attention layers (to capture fine-grained token patterns) and global attention layers (to maintain wide contextual dependencies), keeping context broad while reducing overall cache volume.

Production Case Study: Character.ai
Serving long multi-turn conversations creates long sequence contexts ($L$), which can quickly lead to memory exhaustion. By combining Multi-Query Attention (MQA), cross-layer cache sharing, and interleaved local-global attention windows, Character.ai achieved a $20\times$ aggregate reduction in runtime KV cache memory allocations, drastically lowering serving costs for long contexts.

9. Hardware-Aware Execution: Kernel Fusion and Compilation

Memory-Aware Cache Management: PagedAttention

Traditional inference engines allocate memory for the KV cache as contiguous physical blocks sized to the maximum theoretical sequence length of the request. Because actual sequence lengths vary, this practice leads to two forms of memory waste:

• Internal Fragmentation: Pre-allocated memory slots reserved for maximum output lengths sit empty when generation finishes early.
• External Fragmentation: Scattered, non-contiguous allocations cannot be combined to serve incoming requests.

PagedAttention (introduced in vLLM) solves this by adapting the concept of virtual memory and paging from operating systems:

• The KV cache is divided into fixed-size physical memory blocks called pages.
• Non-contiguous memory pages are dynamically allocated and mapped via a page table as output tokens are generated.
• This dynamic allocation eliminates fragmentation, keeping unused reserved memory close to zero and enabling higher batch sizes within the same physical memory footprint.

FlashAttention: Hardware-Aware Kernel Design

Standard self-attention computes intermediate matrices ($S = Q K^T$, $P = \text{softmax}(S)$) and writes them back to GPU HBM before reading them again to multiply by $V$. For a sequence of length $N$, these intermediate matrices require $O(N^2)$ memory writes and reads.

FlashAttention reorganizes the attention computation to be hardware-aware:

1. Tiling: Splits the large $Q, K, V$ matrices into smaller blocks that fit directly inside fast on-chip SRAM cache memory.
2. Kernel Fusion: Computes the entire attention sequence (matrix multiply, scale, causal mask, online softmax normalization, value aggregation) within SRAM in a single pass before writing the final result back to HBM.
3. Recomputation: Avoids storing large $N \times N$ attention weight matrices in HBM altogether, removing memory bandwidth bottlenecks and speeding up execution without changing the mathematical output.

Compiler Optimization and Lowering

Executing high-level deep learning code involves multiple abstraction layers:

Key compiler optimizations include:

• Operator Fusion: Merges adjacent elementwise operations into a single execution pass (e.g., combining BiasAdd, GELU, and LayerNorm), avoiding round-trip memory access.
• Loop Tiling: Reorders nested loops to match the physical layout of memory (row-major vs. column-major), maximizing data reuse in cache.
• Vectorization: Rewrites operations to execute across wide vector registers (e.g., 128-bit or 256-bit instructions), processing multiple contiguous data elements per cycle.
• Lowering: The progressive translation of high-level computational graphs into target-specific hardware instructions (such as NVIDIA PTX/SASS machine instructions).

10. Service-Level Optimization: Advanced Batching and Decoupling

Batching Strategies: Static, Dynamic, and Continuous

Serving requests individually (Batch Size = 1) leaves GPU hardware significantly underutilized. Batching combines multiple concurrent inputs to increase operational efficiency.

1. Static Batching

The server waits until a fixed number of requests ($N$) arrive before launching inference. Early requests sit idle waiting for later arrivals, causing unpredictable queuing delays.

2. Dynamic Batching

Introduces a maximum time window ($T_{\text{max}}$) alongside the target batch size ($N$). The batch runs as soon as $N$ requests arrive or $T_{\text{max}}$ elapses—whichever happens first. This bounds maximum queuing delay while still batching adjacent requests.

3. Continuous Batching (Iteration-Level / Cellular Batching)

Because token generation is autoregressive, different requests within the same batch finish at different times based on their output lengths. Under static and dynamic batching, early-completing requests remain padded in memory until the longest request in the batch finishes, wasting compute cycles on padding tokens.

Continuous batching operates at the iteration level rather than the request level:

The Public Transit Analogy:
Think of continuous batching like a bus route. Instead of waiting at the terminal until every passenger finishes their round trip, the bus drops off individual passengers at their stops and immediately fills empty seats with new passengers waiting at the curb.

Whenever any request emits an end-of-sequence (<EOS>) token, it is returned to the client and evicted from the batch. A newly arrived request is inserted into that vacated slot on the very next token generation step, keeping active batch slots saturated.

Decoupling Prefill and Decode Instances

Because prefill is compute-bound and decode is memory-bandwidth-bound, running both phases on the same machine creates hardware interference. Large prefill passes interrupt decode iterations, causing latency spikes for active streams.

Modern serving architectures decouple these phases into dedicated physical clusters:

The ratio of prefill-to-decode worker nodes is dynamically scaled based on the target workload:

• Long Input Prompts / Short Outputs (e.g., Document Analysis): Deploy a higher ratio of prefill workers (e.g., $2:1$) to handle prompt processing loads and maintain low Time-to-First-Token ($\text{TTFT}$).
• Short Input Prompts / Long Outputs (e.g., Code Generation): Deploy a higher ratio of decode workers (e.g., $1:3$) to prioritize Time-Per-Output-Token ($\text{TPOT}$) and sustain aggregate generation throughput.

Prompt Caching

Many LLM applications repeatedly send shared text across requests, such as system prompts, few-shot examples, or retrieved context documents.

Request 1: [ System Instructions (2,000 tokens) ] + [ User Query A (50 tokens) ]
Request 2: [ System Instructions (2,000 tokens) ] + [ User Query B (30 tokens) ]
Request 3: [ System Instructions (2,000 tokens) ] + [ User Query C (45 tokens) ]

Without caching, the system recomputes keys and values for those identical $2,000$ tokens on every single query.

Prompt Caching stores the pre-computed KV cache states for deterministic prefix tokens in memory:

• When a new request arrives, the engine checks for cached prefixes, loads the existing KV representations directly, and runs the prefill pass only on the new query tokens.
• For a service handling 1,000,000 daily queries with a 2,000-token system prompt, this avoids recomputing $2\text{ billion}$ tokens every day, drastically lowering both $\text{TTFT}$ and serving costs.

11. Distributed Inference: Tensor, Pipeline, and Sequence Parallelism

When model parameter counts exceed the memory capacity of a single accelerator, the workload must be distributed across multiple devices using parallelization strategies.

Tensor Parallelism (Intra-Node Matrix Splitting)

Tensor Parallelism (TP) splits individual weight matrices across multiple processing units within a single layer, enabling cooperative matrix multiplication.

Consider a linear layer computation $Y = X W$, where weight matrix $W$ is split across two devices using Column-Parallel Linear Splitting:

1. The parameter matrix $W$ is partitioned column-wise into two submatrices:

2. Input tensor $X$ is broadcast identically to both Device 1 and Device 2.
3. Each device performs an independent, smaller matrix multiplication:

4. The partial activations are concatenated along the feature dimension using an All-Gather collective communication call to produce the full output $Y = \begin{bmatrix} Y_1 & Y_2 \end{bmatrix}$.

Because tensor parallelism requires frequent cross-device communication at every single transformer layer, it requires high-speed interconnects (such as NVLink at $\ge 900\text{ GB/s}$). Distributing TP across standard network switches or slow PCIe buses introduces communication latency that quickly outweighs the parallel compute gains.

Pipeline Parallelism (Inter-Node Layer Splitting)

Pipeline Parallelism (PP) partitions the model sequentially by grouping entire layers across different devices:

If implemented naively on full batches, pipeline parallelism causes device underutilization: while Device 4 processes the final layer, Devices 1 through 3 sit completely idle.

To maintain pipeline efficiency, the input batch is broken into smaller micro-batches, allowing upstream devices to begin processing subsequent inputs as soon as their current stage completes:

Because communication only occurs at stage boundaries between groups of layers, pipeline parallelism has much lower bandwidth requirements than tensor parallelism, making it practical across inter-node network links.

Context and Sequence Parallelism

When sequence contexts grow to tens or hundreds of thousands of tokens, intermediate activation tensors during the prefill phase can exceed single-device memory limits even if the model weights fit comfortably:

• Context Parallelism: Partitions the input prompt sequence along the context length dimension across multiple nodes. The self-attention operation is coordinated using ring-style communication primitives (Ring-Attention), allowing systems to process massive context windows without memory exhaustion.
• Sequence Parallelism: Extends tensor parallelism by splitting operations that are traditionally replicated across devices (such as LayerNorm and Dropout) along the sequence dimension, removing redundant activation storage and lowering peak memory usage.

12. Summary and Key Takeaways

• The Objective of Inference Optimization: Production serving requires balancing three competing constraints: minimizing latency ($P_{50}/P_{99}$), maximizing aggregate system throughput (Tokens/sec), and minimizing operating cost per request.
• Divergent Hardware Bottlenecks: Workloads are bound by either compute capability (FLOPs ceiling) or memory bandwidth (bus transfer speeds). The governing metric is Arithmetic Intensity (FLOPs executed per byte of memory read).
• The Transformer Phase Dichotomy: Large language model inference consists of two fundamentally different phases:
  • Prefill: Parallel, high arithmetic intensity, dense matrix-matrix operations ($\text{GEMM}$), strictly compute-bound.
  • Decode: Sequential token generation, low arithmetic intensity, matrix-vector operations ($\text{GEMV}$), strictly memory-bandwidth-bound.
• Latency Decomposition: Total serving latency is formulated as $\text{Latency} = \text{TTFT} + (\text{TPOT} \times N_{\text{output}})$. Average latency metrics obscure tail-end degradation; production SLAs must be engineered against non-parametric percentiles ($P_{90}, P_{95}, P_{99}$).
• Hardware Utilization Metrics: Simple GPU operational metrics are often misleading. True performance must be measured using Model FLOPs Utilization (MFU) during compute-heavy prefill phases and Model Bandwidth Utilization (MBU) during memory-heavy decode phases.
• Speculative Decoding Mechanics: Exploits underutilized compute cores during the memory-bound decode phase by having a fast draft model generate candidate tokens that a large target model verifies in parallel, improving generation speed without altering model output quality.
• KV Cache Optimization: Caching past Key and Value activation vectors drops autoregressive decode complexity from $O(N^2)$ to $O(N)$. Modern architectures use Grouped-Query Attention (GQA) to achieve an optimal balance between the low cache footprint of MQA and the high expressive capacity of MHA.
• Memory-Aware Execution: Techniques like PagedAttention eliminate internal and external memory fragmentation by managing the KV cache using non-contiguous virtual memory pages. FlashAttention fuses attention operations within on-chip SRAM to eliminate intermediate reads and writes to high-bandwidth memory.
• Continuous Batching: Unlike static or dynamic batching, continuous batching works at the individual iteration level, immediately evicting completed requests and inserting new ones to eliminate idle padding overhead.
• Distributed Scaling Paradigms: Massive models rely on hybrid parallelization schemes: Tensor Parallelism splits individual layers across intra-node GPUs via high-bandwidth NVLink, while Pipeline Parallelism divides layers sequentially across networked nodes using micro-batched execution pipelines.