Optimizing Deep Learning Models on GPUs A Cross-Stack Survey from Sampling Kernels to Auto-Parallelized Clusters

Introduction

The dominant cost of contemporary machine learning is no longer algorithmic novelty but execution efficiency on accelerators. Training a frontier LLM consumes tens of thousands of GPU-years, and serving one at scale is a continuous, latency-sensitive workload measured in millions of tokens per second. In both regimes the GPU is the unit of account, and the central engineering question is the same: given a fixed model and fixed hardware, how close can we get to the hardware's roofline?

The answer, for a naive implementation, is "not very close." A straightforward PyTorch model expressed as a sequence of framework operators leaves large amounts of performance on the table. Each operator launches a separate GPU kernel, round-trips its inputs and outputs through high-bandwidth memory (HBM), and is tuned — if at all — for an average case rather than the specific shapes at hand. Meanwhile the arithmetic units sit idle waiting on memory, occupancy is low, and inter-device communication is scheduled pessimistically. The discipline of GPU optimization for deep learning exists to recover this lost performance, and over the past decade it has matured from a collection of hand-written CUDA kernels into a layered stack of automated compilers, schedulers, and runtime systems.

1.1Three exemplars, three layers

This survey is organized around three recent systems that, taken together, illustrate the breadth of the field. They were chosen because each attacks a different layer of the stack and each is state-of-the-art in its niche.

• Qrita1 operates at the operator/kernel layer. It replaces the sort-and-filter implementation of Top-k and Top-p truncation — the sampling step executed on every decoded token of every LLM request — with a pivot-based selection algorithm. Two ideas do the work: a Gaussian σ-truncation that shrinks the candidate set using the statistical distribution of the logits, and a quaternary pivot search with duplicate bookkeeping that halves the number of search iterations while guaranteeing deterministic output. Implemented in Triton, Qrita reaches up to 2× the throughput of the samplers in vLLM, SGLang, and FlashInfer at half the memory, and is now the default GPU Top-p sampler in vLLM.
• Mirage2 and MPK3 operate at the graph/compiler and runtime layers. Mirage is a multi-level superoptimizer that searches a space of μGraphs — representations that jointly describe computation at the kernel, thread-block, and thread levels of the GPU hierarchy — and verifies candidate rewrites for equivalence probabilistically. MPK pushes this to its logical conclusion by fusing an entire LLM's forward pass, including cross-GPU communication, into a single persistent megakernel, cutting inference latency by 1.2× to 6.7×.
• Alpa4 operates at the parallelization layer. It automates the notoriously manual task of sharding a large model across a cluster by decomposing the search into two levels — an integer-linear-program (ILP) formulation for intra-operator (data and tensor) parallelism nested inside a dynamic program for inter-operator (pipeline) parallelism — so that a single decorator turns single-device code into a distributed training job. Alpa's core auto-sharding algorithm has since been upstreamed into XLA.

These three systems share a design philosophy — replace expert hand-tuning with automated search over a well-defined space — but they never touch the same code. Qrita optimizes one operator; Mirage rewrites subgraphs and fuses kernels; Alpa decides how the whole graph is distributed. A production LLM deployment uses all three kinds of optimization simultaneously, which is precisely why a stack-wide taxonomy is useful: it lets us reason about which techniques compose, which conflict, and where the remaining opportunities lie.

1.2Contributions

1. A six-layer taxonomy (Table 2, Figure 1) that classifies GPU optimization techniques for deep learning by the layer of the stack they operate on and the performance principle they exploit.
2. A unified treatment of more than ninety techniques and systems beyond the three exemplars, from vendor libraries (cuDNN, CUTLASS) through auto-schedulers (TVM, Ansor), IO-aware kernels (FlashAttention, FlashInfer), automatic parallelizers (Mesh-TensorFlow, GSPMD, Alpa, Unity, Galvatron), memory systems (ZeRO, FSDP, PagedAttention), quantization methods (GPTQ, AWQ, SmoothQuant, FP8, QuaRot/SpinQuant), and modern serving techniques (chunked-prefill and disaggregated serving, multi-adapter LoRA serving, and advanced speculative decoding).
3. A cross-cutting analysis of search and cost modeling (Section 4), showing that the ILP/DP decomposition of Alpa, the probabilistic-equivalence search of Mirage, and the learned cost models of Ansor are three points in one design space.
4. A discussion of composition, phase ordering, determinism, and portability (Section 5) that frames the open problems of the field.
5. A set of proposed directions for automatic cross-stack optimization (Section 6) — global error-budget allocation, differentiable phase-order-free autotuning, online reoptimization as control, and transferable / co-searched cost models — each positioned against its nearest prior art.

1.3Scope and method

We focus on techniques that improve wall-clock throughput, latency, or memory footprint of deep learning workloads on NVIDIA-class GPUs, with an emphasis on the transformer and LLM workloads that dominate current practice. We deliberately treat algorithmic model compression (pruning, distillation) only where it interacts with GPU execution, and we do not cover accelerator microarchitecture or training-data efficiency. The three exemplar systems are described from their primary sources1234; the surveyed systems are cited to their originating publications throughout.

Background: The GPU Execution Model

To understand why the techniques in this survey work, one must first understand where GPU performance is lost. This section is a compact refresher.

2.1Hardware hierarchy

A modern data-center GPU is a throughput machine built from dozens to hundreds of streaming multiprocessors (SMs). Work is issued as a grid of thread blocks; each block is resident on one SM and executes as warps of 32 threads in lockstep (SIMT). Performance is governed by a steep memory hierarchy: per-thread registers (fastest, tiny), per-block shared memory / L1 (fast, kilobytes), a chip-wide L2 cache, and off-chip HBM (large, but with latency and bandwidth orders of magnitude worse than registers). Dedicated matrix units — Tensor Cores — provide the bulk of arithmetic throughput for the low-precision matrix multiplications at the heart of deep learning, but only if data can be fed to them fast enough. Across GPUs, high-speed interconnects (NVLink within a node, InfiniBand across nodes) carry the collective communication required by distributed execution.

2.2The roofline and the two regimes

The roofline model5 captures the first-order constraint: a kernel's achievable performance is bounded by min(π, β·I), where π is peak compute, β is memory bandwidth, and I is arithmetic intensity (FLOPs per byte of memory traffic). Kernels split into two regimes:

• Compute-bound (high I): large matrix multiplications, convolutions. Here the win comes from better tiling, higher Tensor-Core utilization, and lower precision. Vendor libraries and template kernels (Section Layer 1) dominate.
• Memory-bound (low I): element-wise operations, normalizations, activations, attention with short sequences, and — critically — the token-by-token decode phase of LLM inference. Here the win comes from avoiding memory traffic: fusing operators so intermediates never reach HBM, and redesigning algorithms (as Qrita does for sampling) so they touch less data.

2.3Sources of inefficiency

Five recurring inefficiencies motivate essentially every technique in this survey, and it is worth naming them explicitly (Table 1). Each layer of the taxonomy in Section 3 can be read as a family of responses to one or more of them.

2.4The modern accelerator: Hopper and Blackwell

The optimizations in this survey aim at a moving target: each GPU generation adds hardware features that reshape what is cheap. NVIDIA's Hopper (H100)88 introduced the Tensor Memory Accelerator (TMA) for bulk asynchronous HBM↔SRAM transfers, warpgroup-level wgmma matrix instructions, and thread-block clusters with distributed shared memory (DSMEM) that let neighboring blocks share on-chip memory — features that FlashAttention-311 exploits directly for its asynchronous, low-precision pipeline. Blackwell (B200)89 adds a second-generation Transformer Engine with native FP4 and microscaling formats, higher HBM bandwidth, and faster NVLink for larger coherent GPU domains. These features do not change the roofline model of Section 2.2, but they move its corners: more on-chip memory and asynchrony raise achievable arithmetic intensity, and narrower numeric formats raise peak FLOPs — which is why analytical partitioning models such as that of Pope et al.87 must be re-derived for each generation. The stack optimizes against hardware that is itself a fast-moving part of the design.

A Taxonomy of GPU Optimization Techniques

We organize the field into six layers, ordered roughly from the narrowest scope (a single GPU kernel) to the broadest (an entire cluster and serving system). Figure 1 places the three exemplar systems and the major surveyed systems within this stack, and Table 2 summarizes each layer's goal, principle, and representatives. The remaining sections treat the layers in turn.

Operator- and Kernel-Level Optimization

The lowest layer optimizes individual GPU kernels: the compiled routines that execute a single operation such as a matrix multiply, a normalization, an attention block, or a sampling step. Because a kernel is where FLOPs and bytes are actually spent, this layer sets the performance ceiling for everything above it.

1.1Vendor libraries and template kernels

The historical baseline is a hand-tuned vendor library. cuDNN6 and cuBLAS provide expertly optimized primitives for convolution and matrix multiplication, selected at runtime from a set of pre-tuned implementations. CUTLASS7 exposes the underlying building blocks as C++ templates, letting developers assemble custom high-performance GEMM and convolution kernels with control over tiling, data movement, and Tensor-Core instruction selection. These libraries are extremely fast on the shapes they were tuned for, but they are closed to fusion (a library call is an opaque kernel boundary) and cannot anticipate every shape or operator a new model introduces.

1.2Kernel programming models: Triton

Triton8 is a Python-embedded language and compiler that lets developers write GPU kernels at the granularity of thread blocks operating on tiles, while the compiler handles intra-block scheduling, memory coalescing, and shared-memory allocation. Triton hit a practical sweet spot: kernels are far easier to write than raw CUDA yet reach a large fraction of hand-tuned performance, which is why it has become the implementation substrate for much recent work — including Qrita, whose entire Top-k / Top-p pipeline is written in Triton.

1.3IO-aware kernels: FlashAttention

Attention is the canonical example of a memory-bound operator made fast by algorithmic redesign. FlashAttention9 observes that the standard implementation materializes the full N × N attention matrix in HBM, dominating both memory and time for long sequences. By tiling the computation and fusing softmax with the two matrix multiplies — keeping partial results in shared memory and recomputing rather than storing the attention matrix — FlashAttention makes attention IO-aware and sub-quadratic in memory. FlashAttention-210 improves work partitioning and parallelism, and FlashAttention-311 exploits asynchrony and low-precision (FP8) on Hopper-class GPUs. The lesson — fuse to avoid HBM round-trips, recompute cheap quantities rather than store them — recurs at every layer of this survey.

1.4Auto-scheduling: TVM and Ansor

Rather than hand-tune each kernel, auto-schedulers search the space of loop transformations (tiling, reordering, unrolling, vectorization) automatically. TVM12 pioneered a separation between the mathematical definition of an operator and its schedule, then searched schedules with a learned cost model (AutoTVM). Ansor13 removed the need for hand-written schedule templates by sampling from a hierarchical space of complete programs and prioritizing promising candidates with a learned model, matching or beating vendor libraries on many operators. Auto-scheduling established the template that later superoptimizers refine: define a search space, define a cost model, search.

1.5Superoptimization: TASO and Mirage

Superoptimization goes beyond scheduling a fixed computation — it rewrites the computation itself into an equivalent but faster form, verifying equivalence rather than trusting a fixed rule set. TASO14 automatically generated and verified graph substitutions (e.g., fusing a convolution and a following operation, or reassociating matrix products) for tensor computation graphs, replacing years of manual rewrite-rule engineering. PET15 extended this to partially equivalent transformations with automatic correction.

Mirage2, our second exemplar, is the current state of the art in this line. Its key idea is the μGraph: a representation that describes a tensor program simultaneously at three levels of the GPU hierarchy — the kernel level (grids of blocks), the thread-block level (tiles in shared memory), and the thread level (register computation). By searching over μGraphs, Mirage discovers optimizations that cut across these levels — for example, a novel fusion of attention and subsequent operations, or a reorganization of a normalization — that neither a graph-level rewriter nor a single-kernel scheduler could find in isolation. Because the search generates candidate programs aggressively, Mirage verifies each candidate's equivalence to the original probabilistically, evaluating both on random inputs in a finite field to make the check tractable while keeping the false-acceptance probability negligible. The result is a superoptimizer that regularly outperforms expert-written and vendor kernels on LLM building blocks.

1.6Specialized operator algorithms: Qrita

Some operators are best served not by generic search but by a better algorithm. Qrita, our first exemplar, is such a case. The operator is LLM sampling truncation: on every decoded token, the model's output logits over a vocabulary of 10⁵ or more tokens must be truncated to the Top-k highest-probability tokens (keep the k largest) and/or to the Top-p nucleus (the smallest set whose cumulative probability exceeds p)16 before a token is drawn. Existing LLM engines implement this by sorting the entire vocabulary, which is exact but wasteful: sorting is O(V log V), memory-heavy, and a poor fit for the GPU, and it runs in the innermost loop of decoding.

Qrita replaces the sort with pivot-based selection, extending an idea (RTop-k) originally developed for node selection in graph neural networks. Two techniques make it both fast and correct:

• Gaussian σ-truncation. Logit values across a large vocabulary are approximately normally distributed. Qrita exploits this to compute, from the mean and standard deviation of the logits, a threshold that discards the vast majority of tokens that cannot possibly be in the target set — before any search begins — shrinking the candidate space dramatically at negligible cost.
• Quaternary pivot search with duplication handling. Rather than a binary search for the cutoff value, Qrita evaluates four pivots per iteration, halving the number of iterations needed to converge on the exact threshold. Careful bookkeeping of duplicate logit values (ties at the boundary) guarantees that the output is deterministic — the same input always yields the same truncated set — which distinguishes Qrita from prior stochastic approximations that trade correctness for speed.

Algorithm 1 sketches the resulting selection procedure. Evaluated against the Top-k / Top-p kernels of vLLM, SGLang, and FlashInfer, Qrita delivers up to 2× the throughput and half the memory of sorting-based samplers while producing identical output, and improves end-to-end serving throughput by up to 1.4×. It is now the default GPU Top-p sampler in vLLM. Qrita is a reminder that the operator layer is not only about scheduling known computations faster — sometimes the biggest win is a better algorithm for a step everyone else treated as fixed.

A minimal Triton kernel skeleton for such a fused reduction is shown in Listing 1; the point is that the entire candidate-reduction runs inside one kernel, so the logits are read from HBM once and the truncated set is produced without materializing a sorted copy.

import triton
import triton.language as tl

@triton.jit
def topk_reduce(logits_ptr, out_ptr, V, k, BLOCK: tl.constexpr):
    pid = tl.program_id(0)
    offs = pid * BLOCK + tl.arange(0, BLOCK)
    mask = offs < V
    x = tl.load(logits_ptr + offs, mask=mask, other=-float("inf"))

    # Gaussian sigma-truncation: cheap statistics over the tile
    mu = tl.sum(x, axis=0) / V
    var = tl.sum(tl.where(mask, (x - mu) * (x - mu), 0.0), axis=0) / V
    tau = mu + 2.0 * tl.sqrt(var)          # coarse candidate threshold

    keep = x >= tau                        # prune the tile in registers
    # ... pivot-based count/select loop over survivors (deterministic) ...
    tl.store(out_ptr + offs, tl.where(keep, x, -float("inf")), mask=mask)

1.7Tile-based kernel DSLs

Between the low-level control of raw CUDA and the tile abstractions of Triton, a family of embedded domain-specific languages now targets the exact primitives that modern accelerators reward. ThunderKittens46 exposes a small, opinionated set of tile data structures (16×16 register tiles), an asynchronous thread-block template, and grid-level helpers that map directly onto the warp, block, and grid levels of the GPU hierarchy; with these it matches or beats hand-written kernels — including a competitive FlashAttention-3 implementation — while remaining short enough to read. FlexAttention47 takes the complementary compiler route inside PyTorch: the user writes an attention score modification as a small function and the framework generates a fused, FlashAttention-style kernel, recovering bespoke-kernel performance for the long tail of attention variants (ALiBi, sliding-window, soft-capping, document masking) without writing one. Both illustrate a maturing middle layer — enough abstraction to be productive, enough control to hit the roofline — that increasingly sits beneath the systems of Layers 2 and 6.

1.8Linear-attention and state-space kernels

Not every sequence model is built from softmax attention. Linear-attention and state-space models replace the quadratic attention matrix with a recurrence that is linear in sequence length, but their efficiency on GPUs depends entirely on the kernel. Mamba72 makes a selective state-space model (S6) practical with a hardware-aware parallel-scan kernel that keeps the expanded recurrent state in fast SRAM and fuses the scan, discretization, and output projection so the state never round-trips through HBM. Mamba-273 goes further with structured state-space duality (SSD), which recasts the recurrence as operations on semiseparable matrices that map onto matrix multiplication — letting the model exploit Tensor Cores directly and train several times faster than the original selective scan. These kernels show that the Layer-1 principles — fuse, stay on-chip, and shape the computation as a matmul — apply just as forcefully to sub-quadratic architectures as to attention.

1.9Trainable sparse attention

Sliding-window and attention-sink schemes (Section 4.4) impose sparsity on a model trained with full attention; a newer line makes the sparsity native — learned during training and paired with a hardware-aligned kernel. Native Sparse Attention (NSA)96 gives each query three parallel branches — coarse compressed blocks, a few finely selected blocks, and a local sliding window — combined by a learned gate, with a kernel that loads KV by GQA group and streams only the selected blocks through SRAM, so the sparsity yields real speedups in both training and decoding rather than merely saving memory. MoBA97 routes each query to a small top-k set of key/value blocks through a gating network — an attention analogue of mixture-of-experts — preserving causality and switching a few final layers back to full attention for quality. Unlike the fixed patterns of earlier sparse transformers, these methods co-design the sparse pattern, the training objective, and the GPU kernel together, which is what lets them stay near-lossless while cutting long-context attention cost.

1.10Decode-phase attention

FlashAttention (Section 1.3) parallelizes across the query dimension, which is ideal for training and prefill where many query tokens are processed at once. In autoregressive decoding, however, only one query token is generated per step against a long KV cache, so there is little query-level parallelism and the GPU is left underused. Flash-Decoding addresses this by splitting attention over the KV-length dimension: the KV cache is partitioned into chunks processed in parallel across many SMs, whose partial softmax results are then combined, restoring occupancy for long-context decode. FlashDecoding++102 extends the idea with asynchronous softmax using a unified maximum, optimized flat-GEMM kernels for the tall-and-thin matrices typical of decoding, and heuristic dataflow selection, further accelerating both phases. Decode-phase attention is a reminder that the right parallelization dimension depends on the workload: what is optimal for prefill is not optimal for the single-token decode step.

Graph- and Compiler-Level Optimization

Where Layer 1 makes a single kernel fast, Layer 2 optimizes the computation graph as a whole: the directed acyclic graph of operators that constitutes a model. Its central weapon is operator fusion — merging adjacent operators into a single kernel so that intermediate tensors never leave on-chip memory — together with classical compiler techniques: layout assignment, constant folding, common-subexpression elimination, and algebraic simplification.

2.1Production graph compilers

XLA17 (the compiler behind JAX and, historically, TensorFlow) performs aggressive fusion of element-wise and reduction operators and is the backend on which Alpa builds. TensorRT18 is NVIDIA's inference compiler, fusing layers, selecting precision, and picking tuned kernels for a target GPU; TensorRT-LLM specializes this for transformer inference. In the PyTorch ecosystem, TorchDynamo captures Python execution into a graph and TorchInductor19 lowers it — emitting Triton kernels for GPUs — bringing compilation to eager-mode users with minimal code change. ONNX Runtime20 provides a portable graph optimizer across hardware backends. These systems rely largely on curated rewrite rules: fast, predictable, but limited to transformations their authors anticipated.

2.2Search-based graph rewriting

The superoptimizers of the operator layer — TASO14, PET15, and Mirage2 — also live partly at this layer, because the substitutions they discover are graph rewrites. The trend they embody is the replacement of hand-curated rule sets with automatically generated and verified rewrites, and the erosion of the boundary between "graph optimization" and "kernel optimization": Mirage's μGraphs deliberately span both, because the best fusion often depends on how the fused kernel will be scheduled internally.

2.3The limit case: whole-model megakernels (MPK)

Fusion has a natural endpoint. If fusing two kernels avoids one HBM round-trip, why not fuse all of them? MPK (Mirage Persistent Kernel)3, the successor to Mirage and our runtime-layer exemplar, does exactly this: it compiles an entire LLM's forward pass into a single persistent megakernel. Rather than launching thousands of small kernels per decoded token — each with launch overhead and each synchronizing through HBM — MPK expresses the whole model as a task graph executed by a fixed pool of persistent workers and schedulers resident on the SMs, so that computation and inter-GPU communication (all-reduce, all-gather) happen inside one kernel launch. The developer defines the model's computation graph in a few dozen lines of Python, chaining fused layers such as a fused RMSNorm-plus-Linear block, and MPK generates the CUDA megakernel. Because the decode phase is dominated by launch overhead and memory-bound element-wise work, this end-to-end fusion cuts inference latency by 1.2×–6.7×. MPK is the clearest demonstration of a theme this survey returns to: the layers are not truly separate, and the most aggressive systems collapse operator, graph, and runtime optimization into one.

Listing 2 illustrates the megakernel programming model at a high level.

import mirage as mi

mpk = mi.PersistentKernel(world_size=W, mpi_rank=r,
                          num_workers=96, num_local_schedulers=48,
                          num_remote_schedulers=0,
                          meta_tensors=[step, tokens])

x = mpk.attach_input(torch_tensor=hidden, name="hidden")
mpk.rmsnorm_linear_layer(input=x, weight_norm=w_norm,
                         weight_linear=w_qkv, output=attn_in,
                         grid_dim=(96,1,1), block_dim=(128,1,1))
# ... chain the remaining fused layers of the model ...
mpk.compile()   # generate a single fused CUDA megakernel
mpk()           # run: all compute + communication in one launch

2.4Task-mapping tensor compilers

Between the schedule search of Ansor13 and the hand-written tiles of a kernel DSL sits a class of compilers that make the mapping of computation onto the GPU hierarchy a first-class, programmable object. Hidet90 introduces a task-mapping programming paradigm in which the assignment of tensor-program tasks to threads, warps, and blocks is expressed and optimized explicitly, giving the compiler fine control over tiling and memory movement and yielding faster kernels than schedule-only auto-schedulers on many operators. This is the compiler-side counterpart to the tile abstractions of ThunderKittens46: both recognize that on modern GPUs where data and work sit in the memory and execution hierarchy matters as much as the arithmetic, and both make that placement an explicit, searchable dimension rather than a fixed heuristic.

Parallelization and Distributed Execution

When a model or its activations exceed a single GPU's memory, or when training must finish in reasonable time, work must be split across many GPUs. This layer decides how, and its central difficulty is that partitioning introduces communication that can easily erase the benefit of added hardware. Table 3 summarizes the major strategies.

3.1Manual and semi-automatic parallelism

The workhorses of large-model training are hand-designed. Megatron-LM21 introduced the tensor-parallel partitioning of transformer layers that most large models still use. GPipe22 established synchronous pipeline parallelism with micro-batching, and PipeDream23 added asynchronous and 1F1B scheduling to keep pipeline bubbles small. GShard24 scaled mixture-of-experts models via all-to-all expert routing. Combining these — so-called 3D parallelism — is powerful but requires expert effort to tune for each model and cluster.

3.2Toward automatic parallelization

A line of work seeks to automate the choice of parallel strategy. Mesh-TensorFlow25 let users annotate tensor dimensions for partitioning; GSPMD26 generalized this into an XLA-level sharding system driven by lightweight annotations and propagation. These reduce effort but still lean on human guidance for the overall plan.

3.3Alpa: two-level automatic parallelization

Alpa4, our third exemplar, removes the human from the loop by observing that the space of parallel strategies has a natural hierarchy that mirrors the hardware. Alpa distinguishes:

• Intra-operator parallelism — data and tensor parallelism that partition the axes of individual operators. This maps to tightly-coupled, high-bandwidth device groups (e.g. GPUs within a node connected by NVLink), where the frequent collective communication it requires is cheap.
• Inter-operator parallelism — pipeline parallelism that assigns disjoint stages of the graph to different device groups. This maps to the looser, lower-bandwidth connections between nodes, where only stage-boundary activations cross the wire.

The two levels are solved with different tools (Algorithm 2). For a given graph partition into stages and a device assignment, Alpa formulates the best intra-op plan for each stage as an integer linear program (ILP) that minimizes compute plus communication cost subject to consistency constraints on tensor shardings. It then wraps this in a dynamic program over the space of graph-into-stage partitions and device-group assignments (inter-op), choosing the slicing that minimizes the pipelined end-to-end cost. This decomposition tames a combinatorial space that is intractable to search jointly, and it lets a single @parallelize decorator (Listing 3) turn single-device JAX code into a distributed job. Alpa achieves strong scaling on models with billions of parameters, and its auto-sharding core has been merged into XLA.

import alpa

@alpa.parallelize          # auto data + operator + pipeline parallelism
def train_step(state, batch):
    def loss_fn(params):
        out = state.forward(params, batch["x"])
        return jnp.mean((out - batch["y"]) ** 2)
    grads = grad(loss_fn)(state.params)
    return state.apply_gradient(grads)

state = create_train_state()
for batch in data_loader:      # now runs across the whole cluster
    state = train_step(state, batch)

3.4Beyond Alpa

Contemporary and follow-on systems broaden the automatic-parallelization search. Unity27 jointly optimizes algebraic transformations and parallelization in one substitution-based search; Galvatron28 automatically selects among multiple parallelism dimensions for transformer training with a decision-tree / dynamic-programming search. The common thread is the same as at Layers 1–2: a well-defined space, a cost model, and an efficient search — applied here to the placement of computation across a cluster.

3.5Sequence and context parallelism

The parallelism dimensions above shard the model and the batch, but not the sequence — and for long-context training the activations and attention of a single long sequence exceed one device. Sequence parallelism in Megatron56 partitions the otherwise-replicated LayerNorm and dropout activations along the sequence axis, cutting activation memory alongside tensor parallelism. Ring Attention54 distributes a sequence across devices and streams key/value blocks around a ring, overlapping the block-wise attention computation with the communication of the next block so that context length scales with the device count at near-zero attention memory per device. DeepSpeed-Ulysses55 instead partitions along the attention-head dimension and uses all-to-all collectives to transpose between sequence- and head-parallel layouts, achieving high communication efficiency for extreme sequence lengths. These techniques add a fourth axis — sequence — to the data / tensor / pipeline decomposition that Alpa and its successors search.

3.6Pipeline schedules and communication overlap

Two persistent inefficiencies of distributed training are the pipeline bubble — idle time at the start and end of each stage — and the exposed cost of collective communication. Zero-bubble pipeline parallelism57 splits the backward pass into its input-gradient and weight-gradient components and schedules them to fill the bubble, reaching under the right conditions a schedule with no pipeline stall at all. DeepSeek-V3's DualPipe58 pushes further with a bidirectional schedule that overlaps the forward and backward computation-communication phases across a pair of micro-batch streams, hiding most of the cross-node all-to-all traffic of expert parallelism. At the collective level, Flux59 over-decomposes a communication operation and its dependent GEMM into fine-grained tiles and fuses them into a single kernel, so that tensor-parallel communication is hidden inside computation rather than serialized after it. The unifying idea is to treat communication not as a barrier between compute phases but as work to be interleaved with them.

3.7Mixture-of-experts parallelism

Sparsely-gated mixture-of-experts (MoE) layers scale parameter count at fixed per-token compute by routing each token to a few of many experts (Section 3.1 introduced GShard's24 expert sharding). The systems challenge is the resulting all-to-all communication that dispatches tokens to their experts and gathers the results, whose cost and imbalance can dominate. DeepSpeed-MoE74 provides an end-to-end training-and-inference stack combining expert, tensor, and data parallelism with a distilled, compressed inference path. Tutel75 makes the parallelism adaptive: a single invariant data / parameter layout lets it switch parallelism and pipelining strategy at zero cost as the runtime token-to-expert workload shifts, delivering large speedups at thousand-GPU scale. At the frontier, DeepSeek-V358 pairs its DualPipe schedule with dedicated cross-node all-to-all kernels (released as DeepEP) so that fine-grained expert parallelism achieves near-zero exposed communication. MoE turns model sparsity into a communication-scheduling problem.

3.8Collective communication and interconnect

Every parallelism strategy above ultimately runs on a substrate of collective operations — all-reduce, all-gather, reduce-scatter, all-to-all — executed over the GPU interconnect (NVLink and NVSwitch within a node, InfiniBand across nodes). The NVIDIA Collective Communication Library (NCCL)103 is the de-facto substrate: it detects the interconnect topology and implements ring- and tree-based algorithms whose bandwidth and latency determine what the overlap techniques of Section 3.6 have to hide. Because a single ring is limited by its weakest link, a line of work synthesizes topology-aware collectives instead of using fixed algorithms. Blink104 packs spanning trees over the links actually available to a job (via integer programming) to reach near-optimal link utilization on heterogeneous topologies, and TACCL105 synthesizes collective algorithms from a high-level communication sketch and a profiled topology, outperforming NCCL on all-gather and all-to-all and speeding up end-to-end MoE training. The interconnect and its collectives are the floor beneath every distributed technique: no amount of overlap helps if the collective itself is suboptimal for the topology.

Memory Optimization

GPU memory is scarce and often the binding constraint: a model may be fast enough but simply not fit. This layer trades compute or bandwidth for capacity.

4.1Recomputation and state sharding

Gradient checkpointing (activation recomputation)29 stores only a subset of activations during the forward pass and recomputes the rest during the backward pass, cutting activation memory at the cost of extra FLOPs — a direct compute-for-memory trade. ZeRO30 attacks the memory of the training state itself: its three stages progressively shard optimizer state (ZeRO-1), gradients (ZeRO-2), and parameters (ZeRO-3) across data-parallel workers, so that no single GPU holds the full state, while gathering parameters just-in-time for each layer. PyTorch's FSDP31 brings the same fully-sharded data-parallel approach into the core framework. Because these techniques shard state rather than computation, they compose cleanly with the parallelism of Layer 3.

4.2Offloading

When even sharding is insufficient, state can be pushed to slower, larger memory. ZeRO-Offload and ZeRO-Infinity32 move optimizer state and parameters to CPU RAM and NVMe, enabling models far larger than aggregate GPU memory. For inference, FlexGen33 schedules weights and KV cache across GPU, CPU, and disk to run large models on a single commodity GPU at high throughput for latency-tolerant, batch-oriented workloads. Offloading trades bandwidth (and latency) for capacity, and is most effective when computation can be overlapped with the slower transfers.

4.3KV-cache management

In autoregressive inference, the dominant memory consumer is the KV cache: the keys and values of all past tokens, kept to avoid recomputation. PagedAttention, the core of vLLM34, applies the operating-system idea of paging to this cache: it stores KV in fixed-size, non-contiguous blocks referenced through a block table, eliminating the fragmentation and over-reservation of contiguous allocation and allowing near-zero-waste memory use and prefix sharing across requests. FlashInfer35 generalizes KV storage into composable block-sparse formats with JIT-compiled attention templates and load-balanced scheduling, and has become a shared attention backend across vLLM, SGLang, and MLC. Since KV-cache management sits at the boundary of memory and serving, it recurs in Layer 6.

4.4Attention-level KV-cache reduction

Paging (Section 4.3) manages the KV cache efficiently; a complementary line of work makes the cache smaller at the source by changing the attention architecture or its precision. Multi-query attention (MQA)48 shares a single key/value head across all query heads, shrinking the cache by the head count at some quality cost; grouped-query attention (GQA)49 interpolates by sharing key/value heads within groups, recovering most of the quality while keeping the cache small, and is now standard in open LLMs. Multi-head latent attention (MLA)50, introduced in DeepSeek-V2, instead compresses keys and values into a low-rank latent vector that is cached and re-expanded on the fly, cutting KV memory further while preserving multi-head expressivity. Orthogonally, KV-cache quantization stores the cache in low precision: KIVI52 uses tuning-free asymmetric 2-bit quantization — per-channel for keys, per-token for values — while KVQuant53 combines per-channel keys, pre-RoPE quantization, and non-uniform levels to reach near-lossless sub-4-bit caches enabling million-token contexts. Finally, attention sinks (StreamingLLM)51 exploit the observation that keeping a few initial “sink” tokens plus a sliding window lets a model stream over unbounded input with a bounded cache. These methods reduce the very quantity that Layer 6 must serve.

4.5Low-bit optimizer state and gradient projection

The techniques above shard or offload the training state; a complementary family shrinks it numerically or structurally. 8-bit optimizers91 store Adam's first- and second-moment statistics in 8 bits via block-wise dynamic quantization, roughly halving optimizer memory while matching 32-bit convergence. ZeRO++92 attacks the communication cost of sharded training (Section 3) with quantized all-gather of weights, hierarchical data remapping, and quantized gradient reduction, cutting ZeRO's30 communication volume up to fourfold on bandwidth-limited clusters. GaLore93 takes a structural route: it projects gradients into a low-rank subspace, runs the optimizer there, and projects the update back, so optimizer state scales with the rank rather than the full parameter count — enabling full-parameter pre-training of a 7B model on a single consumer GPU, complementary to the adapter-based QLoRA60. Memory optimization thus reaches into the numerics and rank of the training state itself.

4.6KV-cache eviction and token dropping

Quantization (Section 4.4) keeps every token but in fewer bits; an orthogonal strategy keeps fewer tokens at full precision. H2O98 observes that a small set of “heavy-hitter” tokens accounts for most attention mass and evicts the rest with a low-cost scoring policy, bounding the cache while largely preserving quality. SnapKV99 compresses the prompt's cache before generation by selecting, per head, the KV positions that the end of the prompt actually attends to, cutting cache size with minimal accuracy loss on long inputs. PyramidKV100 allocates the cache budget non-uniformly across layers — more in lower layers where attention is diffuse, less in higher layers where it concentrates — following the pyramidal way information funnels through the network. These eviction methods trade a controlled, usually small quality loss for a hard cap on KV memory, and compose with the paging and quantization of the surrounding subsections.

Numerical Precision and Sparsity

The cheapest FLOP is the one performed in fewer bits — or skipped entirely. This layer reduces the cost of arithmetic by lowering precision or exploiting structural zeros.

5.1Reduced-precision training and inference

Mixed-precision training36 performs most computation in FP16 (or BF16) while keeping a master copy of weights and select reductions in FP32, with loss scaling to preserve small gradients — roughly doubling throughput and halving memory on Tensor-Core hardware at no accuracy loss. Newer hardware adds FP8 formats, exploited for both training and inference (e.g. in FlashAttention-311), pushing the trade further.

5.2Quantization

Post-training quantization compresses trained weights (and sometimes activations) to 8-bit or 4-bit integers. GPTQ37 uses second-order information to quantize LLM weights to 3–4 bits with minimal accuracy loss. SmoothQuant38 makes activation quantization tractable by migrating quantization difficulty from activations to weights, enabling efficient INT8 weight-and-activation inference. AWQ39 protects the salient weight channels that matter most for accuracy, achieving strong 4-bit weight-only quantization. These methods are largely orthogonal to the other layers: a quantized model still benefits from fused kernels, paging, and parallelism.

5.3Sparsity and structural compression

Pruning removes weights; when the remaining pattern matches hardware support — notably the 2:4 structured sparsity accelerated by Ampere-and-later Tensor Cores40 — it yields real speedups rather than merely smaller storage. Knowledge distillation and low-rank adaptation reduce effective model size, and interact with GPU execution by changing the shapes and intensities of the underlying kernels. Precision and sparsity are the layer where model and system optimization most visibly meet.

5.4The low-precision frontier: FP8 training and rotation-based quantization

The precision trade is being pushed below eight bits on both the training and inference sides. On the training side, DeepSeek-V358 demonstrated a production-scale FP8 recipe: fine-grained, tile- and block-wise scaling keeps the dynamic range of activations and weights within FP8 while accumulating in higher precision, training a frontier model at markedly lower cost — and hardware vendors have since standardized still-narrower microscaling formats (MXFP8 / MXFP6 / MXFP4 and NVFP4) that attach shared block scales to 4–8-bit elements. On the inference side, the difficulty of 4-bit quantization is dominated by activation outliers. QuaRot61 multiplies weights and activations by randomized Hadamard rotations that provably preserve the network's output while spreading outliers across channels, enabling outlier-free 4-bit weights, activations, and KV cache; SpinQuant62 replaces the random rotations with matrices learned by Cayley optimization, closing most of the remaining W4A4 gap to full precision. For memory-bound fine-tuning, QLoRA60 back-propagates through a frozen 4-bit (NF4) base model into small LoRA adapters, making single-GPU fine-tuning of very large models routine. Rotation and low-rank structure are becoming the standard tools for reaching the sub-8-bit regime without retraining.

5.5Extreme quantization: sub-4-bit and ternary weights

Below four bits, uniform integer quantization breaks down and two richer families take over. Vector (codebook) quantization maps groups of weights to entries of a learned codebook: QuIP#76 combines incoherence processing via randomized Hadamard transforms with an E8-lattice codebook to make 2-bit models usable, and AQLM77 pushes additive multi-codebook quantization to a strong 2-bit accuracy / size trade-off — both well beyond what GPTQ37 reaches at that width. A different route trains extreme precision in from the start: BitNet b1.5878 constrains every weight to the ternary set {−1, 0, +1}, replacing most multiplications with additions and, at scale, matching full-precision perplexity while slashing memory and energy. The sub-4-bit regime rewards structure — codebooks or ternary constraints — over ever-finer uniform grids.

5.6One-shot pruning and contextual sparsity

Where quantization shrinks each weight, pruning removes weights entirely. SparseGPT79 shows that a GPT-scale model can be pruned to 50–60% unstructured sparsity (or 2:4 semi-structured, for Tensor-Core acceleration40) in one shot, without retraining, by solving a layer-wise reconstruction with efficient Hessian updates; Wanda80 reaches comparable quality with a far simpler criterion — the product of weight magnitude and input-activation norm — and no weight update at all. A dynamic alternative exploits contextual sparsity: Deja Vu81 observes that for any given input only a small, predictable subset of attention heads and MLP neurons matters, and predicts that subset on the fly to skip the rest, yielding wall-clock speedups without quality loss. PowerInfer82 turns this skew into a GPU–CPU inference engine that keeps hot neurons resident on the GPU and cold ones on the CPU. Sparsity, like precision, pays off only when the remaining pattern is one the hardware can actually skip.

Inference-Serving and Runtime Optimization

The top layer optimizes the serving system: how many requests share the GPU, how their state is reused, and how the per-token runtime overhead is hidden. This is where several lower-layer techniques — Qrita's sampler, FlashInfer's attention, PagedAttention's KV cache, MPK's megakernel — are actually deployed.

6.1Batching and scheduling

Static batching wastes capacity when requests have different lengths. Orca41 introduced continuous (iteration-level) batching, in which finished sequences leave the batch and new ones join at each decoding step, keeping the GPU saturated. vLLM34 combines continuous batching with PagedAttention to raise serving throughput by an order of magnitude over earlier systems.

6.2Prefix and prompt reuse

Many requests share prefixes — system prompts, few-shot examples, chat history. SGLang42 introduces RadixAttention, which stores the KV cache in a radix tree keyed by token prefix so that shared prefixes are computed once and reused across requests, cutting redundant prefill work substantially for structured and multi-turn workloads.

6.3Reducing the length of the decode loop

Speculative decoding4344 runs a small "draft" model to propose several tokens that the large model then verifies in parallel, accepting a run of tokens per expensive forward pass and thereby amortizing the memory-bound decode step. Medusa45 attaches multiple decoding heads to the model itself to generate and verify candidate continuations without a separate draft model. These methods reduce the number of sequential large-model steps.

6.4Hiding launch and dispatch overhead

Because the decode phase issues many tiny kernels, per-launch overhead is significant. CUDA Graphs capture a fixed sequence of kernel launches and replay them as a single unit, removing per-launch CPU cost; persistent kernels keep a worker pool resident on the SMs to process a stream of tasks without relaunching. MPK's megakernel3 (Section 2.3) is the extreme form of this idea, and Qrita's fused sampler1 removes one of the small, memory-bound kernels that would otherwise punctuate every decode step. The serving layer is thus the integration point where the whole stack's optimizations are composed into a single deployed system.

6.5Disaggregated and chunked-prefill serving

LLM inference has two phases with opposite hardware profiles: a compute-bound prefill over the prompt and a memory-bound decode over generated tokens. Interleaving them on one GPU means each interferes with the other's latency target. Two strategies resolve this. Chunked prefill, introduced by Sarathi-Serve64, splits a long prefill into small chunks and piggybacks decode iterations onto each, keeping the GPU saturated while avoiding the generation stalls a monolithic prefill would cause. Disaggregated serving instead assigns the phases to separate GPU pools: DistServe63 and Splitwise65 place prefill and decode on different devices — each tuned and scaled for its own latency target — and transfer the KV cache between them, improving goodput under tight time-to-first-token and time-per-output-token constraints. Mooncake66 builds a KVCache-centric architecture around this split, adding a disaggregated, tiered KV store that maximizes prefix-cache reuse across a cluster. Disaggregation turns the phase mismatch from a scheduling headache into an architectural boundary.

6.6Advances in speculative decoding and multi-adapter serving

Speculative decoding (Section 6.3) has been sharpened considerably. EAGLE67 drafts at the feature level and verifies a tree of candidate tokens, raising the accepted-token rate well beyond a naive draft model; its successors extend the draft tree and its training. Lookahead decoding68 removes the draft model entirely, using a Jacobi-style parallel iteration to generate and verify n-grams from the target model itself. Hydragen69 attacks the shared-prefix case directly, decomposing attention into a batched prefix part and per-sequence suffixes so that a common prompt is read from memory once for the whole batch. A parallel concern is serving many fine-tuned variants: S-LoRA70 and Punica71 keep a single copy of the base model resident and batch requests across thousands of LoRA adapters using custom gather-scatter kernels, so multi-tenant adapter serving costs little more than serving the base model alone. Together these techniques attack the two remaining costs of the decode loop: its sequential length and its per-tenant duplication.

6.7Request scheduling and live migration

Continuous batching (Section 6.1) decides what runs each step; a scheduling layer above it decides which requests occupy scarce GPU memory and for how long. Because generation length is unknown in advance, head-of-line blocking and fragmentation are endemic. FastServe83 applies token-level preemptive, multi-level scheduling to cut the latency tail of long generations. Llumnix84 adds efficient live migration of a running request and its KV cache across GPU instances, unifying load-balancing, de-fragmentation, prioritization, and auto-scaling under one dynamic policy. Together with chunked prefill64, these schedulers treat the cluster's aggregate GPU memory as a managed, migratable resource rather than a static partition.

6.8Constrained and structured decoding

Many deployments need output that conforms to a schema — JSON, a regular expression, or a grammar — and enforcing it by sample-validate-retry is slow. Constrained decoding instead masks the next-token distribution to only grammar-valid tokens at each step. Outlines85 compiles a regular expression or context-free grammar into a finite-state machine and builds an index over the vocabulary so the per-token mask is a cheap lookup, adding negligible overhead while guaranteeing well-formed output. XGrammar86 extends this to full context-free grammars with a pushdown-automaton backend and precomputed token-mask caches, reaching near-zero-overhead structured generation that integrates with batched serving. Structured decoding is where the serving layer meets the application contract: correctness of form, enforced at the level of the sampler.

6.9Hierarchical and cross-request KV reuse

RadixAttention42 reuses shared prefixes within the GPU's KV cache; two lines of work extend reuse in space and time. Prompt Cache94 generalizes prefix reuse to modular segments: a schema marks reusable prompt modules (system messages, templates, documents) whose attention states are precomputed once and spliced into any prompt that contains them, with positional bookkeeping that keeps the reuse exact. CachedAttention / AttentionStore95 extends reuse across turns and time: it persists a conversation's KV cache into a tiered store (HBM, host memory, disk) when a session goes idle and reloads it on resumption, so a multi-turn chat prefills only the new tokens; scheduler-aware placement and positional-encoding decoupling keep the cached state valid and fast to fetch. Together with PagedAttention's34 intra-GPU management, these turn the KV cache into a managed hierarchy spanning the whole memory system.

6.10Tree-structured speculative inference

The speculative methods of Section 6.6 mostly verify a single drafted sequence per step; verifying a tree of candidates raises the expected number of tokens accepted per expensive forward pass. SpecInfer101 generalizes speculation to a token tree: several small draft models propose many continuations at once, and a single tree-based verification pass over the large model accepts the longest valid path through the tree, with a custom attention kernel that computes the tree's shared prefixes only once. This tree-verification mechanism is what the later EAGLE variants67 build on with dynamically shaped draft trees, and it composes with continuous batching so speculation and multi-request serving coexist. Tree speculation reframes the decode bottleneck as a search over candidate continuations verified in parallel.

6.11Elastic and long-context serving

Long-context requests strain a serving system unevenly: a million-token prompt's KV cache can exceed any single GPU's memory, and the resource demand swings sharply between the prefill and decode phases of the same request. Static parallelism cannot track that variance. LoongServe106 introduces elastic sequence parallelism: it adjusts the degree of sequence parallelism per request and per phase in real time, distributes the decoding of a single batch across multiple instances (multi-master distributed decoding) to avoid whole-request KV migration, and overlaps partial-decode communication with computation — improving GPU compute, communication, and memory efficiency together for variable-length long-context workloads. This closes the loop with the sequence parallelism of Section 3.5 and the disaggregation of Section 6.5: elasticity treats the parallelism degree itself as a runtime scheduling decision rather than a fixed deployment choice.

Search, Cost Models, and Autotuning

A striking commonality runs through the exemplars and much of the surveyed work: modern GPU optimization is increasingly search over a defined space guided by a cost model, rather than fixed heuristics. The differences between systems are largely differences in (i) what the search space is, (ii) how correctness of a candidate is guaranteed, and (iii) how cost is estimated. Table 4 places the three exemplars and several surveyed systems in this design space.

Three observations follow. First, correctness guarantees trade against search-space size: Alpa's plans are semantics-preserving by construction so it can search freely; Mirage generates such aggressive rewrites that it must fall back to probabilistic verification; Qrita, operating on a single well-understood operator, can afford an exact and deterministic guarantee, which is precisely its selling point over stochastic samplers. Second, cost models range from closed-form to learned: Qrita's Gaussian model of the logit distribution is analytic; Alpa's cost model sums analytical communication and compute terms; Ansor learns its model from measurements. Third, decomposition is what makes search tractable: Alpa's nesting of ILP inside DP, and Qrita's pre-truncation before pivot search, are both moves that shrink an intractable space into a searchable one. Recognizing GPU optimization as a search problem clarifies why advances often come from better spaces and cheaper verification, not just faster kernels.

Discussion: Composition, Trade-offs, and Open Challenges

5.1The layers compose — but not freely

The practical power of the taxonomy is that a real deployment stacks techniques from every layer: a model is quantized (Layer 5), fused and compiled (Layers 1–2), parallelized across a cluster (Layer 3) with sharded and paged memory (Layer 4), and served with continuous batching and an optimized sampler (Layer 6). Yet composition is not free. A choice at one layer constrains another: an aggressive Layer-3 sharding changes the operator shapes that Layer-1 auto-scheduling must tune; a Layer-2 megakernel that fuses everything (MPK) leaves little for a separate Layer-6 batching scheduler to do and must therefore absorb that responsibility. This is the classic phase-ordering problem of compilers, now spread across an entire systems stack, and it is largely unsolved: today's systems optimize each layer in sequence with hand-designed interfaces between them, and joint optimization across layers — as Unity27 attempts for algebra and parallelism, or as Mirage attempts across kernel and block levels — remains an active frontier.

5.2Determinism, reproducibility, and correctness

Performance optimizations frequently perturb numerics or ordering, and the field is increasingly aware that correctness is a feature. Qrita's explicit guarantee of deterministic, sort-identical output1 is a direct response to earlier samplers that bought speed with stochastic approximation; Mirage's probabilistic equivalence check2 is an engineered bound on the risk that aggressive rewriting introduces. As optimizations grow more aggressive, principled correctness guarantees — formal, probabilistic, or exact-by-construction — become part of the value proposition rather than an afterthought.

5.3Portability and the cost of automation

Hand-tuned kernels are fast but brittle: a new GPU generation, a new attention variant, or a new sparsity pattern can invalidate months of work. The automated systems surveyed here — auto-schedulers, superoptimizers, auto-parallelizers — exist largely to make performance portable across hardware and models by re-searching rather than re-engineering. The cost is search time and the difficulty of building accurate cost models; the benefit is that a single system such as Alpa or Mirage adapts to new targets with far less human effort. The trajectory of the field is clearly toward more automation, with expert kernels reserved for the few hottest operators.

5.4Open challenges

We highlight four. (1) Cross-layer joint optimization: principled co-optimization of parallelism, fusion, memory, and precision, rather than sequential per-layer passes. (2) Search cost: making superoptimization and auto-parallelization fast enough to run routinely, perhaps with transferable or learned cost models. (3) Correctness at scale: extending exact or probabilistic guarantees to the whole optimized stack. (4) Hardware co-evolution: FP8, structured sparsity, and megakernel-friendly scheduling change the target underneath every technique, so the taxonomy's contents will keep shifting even as its structure endures.

5.5Where the frontier is moving

The systems of the 2024–2026 wave sharpen three tensions already latent in the taxonomy. First, communication is now co-equal with computation: MoE all-to-all (Section 3.7), sequence-parallel ring exchange (3.5), disaggregated prefill/decode transfer (6.5), and quantized collectives (ZeRO++) all treat the network as something to be scheduled and overlapped rather than merely tolerated. Second, the numeric floor keeps dropping: FP8 training, FP4 and microscaling hardware, rotation-based 4-bit inference, and 2-bit or ternary weights push the precision–quality trade further, raising the correctness and reproducibility stakes of Section 5.2. Third, the KV cache has become a first-class, hierarchical resource — paged, quantized, architecturally reduced (MLA/GQA), disaggregated, and cached across requests and turns — because for long-context and multi-turn serving it, not raw FLOPs, is the binding constraint. None of these dissolve the six-layer structure; they populate it with denser, more automated, and more communication-aware machinery.

Proposed Directions: Automatic Cross-Stack Optimization

The taxonomy of Section 3 exposes a structural limitation that no single surveyed system escapes: each technique optimizes within one layer and chooses its approximations independently of the others. Ansor13 searches kernel schedules; Alpa4 searches parallelization; GPTQ37 selects a quantization; H2O98 sets an eviction aggressiveness — each without knowing what the others decided. Section 4 argued that these are all instances of search over a defined space guided by a cost model, and Section 5 flagged phase ordering — that applying the same optimizations in different orders yields different, non-monotonically related results — as an open problem. This section sketches four directions that follow directly from those observations by treating the whole stack, rather than one layer, as the unit of automatic optimization. These are proposals, not evaluated systems; for each we state the nearest prior art so the novelty claim can be checked, and we are explicit about which are genuinely new and which are extensions.

6.1Global error-budget allocation

Most of Layers 4–6 trade a little accuracy for speed, but each trade is made in isolation: the quantizer, the KV-cache evictor, and the speculative sampler each pick their own operating point, and their errors interact in ways none of them models. We propose making correctness a shared resource that is allocated, not spent locally. The user declares a single global error budget ε — for example a bound on the KL divergence of the output distribution, or a maximum tolerated accuracy drop δ — and an allocator distributes ε across all lossy layers at once, solving maximize throughput subject to composed-error ≤ ε as a constrained optimization (a Lagrangian / knapsack over the per-layer (speed, error) Pareto curves that each technique already exposes). The core novelty is that correctness becomes a globally optimized budget rather than a set of independent local decisions. The hard part — and the real research contribution — is that error does not compose linearly: the interaction of quantization, sparsity, and speculation is non-additive, so the method needs a learned error-composition model that predicts end-to-end degradation from a vector of per-layer operating points. Nearest prior art: HAWQ107 allocates bit-width across layers by Hessian sensitivity — but only within quantization, and with an analytic per-layer sensitivity rather than a learned cross-layer composition model. Extending one shared budget across heterogeneous lossy transforms appears to be new.

import torch

# Each lossy stage exposes a Pareto front of operating points
# (speedup, local_error) measured offline:  stages[name] -> Tensor[P, 2]

class ErrorComposer(torch.nn.Module):
    """Learned: per-stage local errors -> end-to-end degradation.
    Captures the NON-ADDITIVE interaction of quant x sparsity x speculation.
    Trained on (chosen op-points, measured true KL) pairs."""
    def __init__(self, n_stages, h=64):
        super().__init__()
        self.net = torch.nn.Sequential(
            torch.nn.Linear(n_stages, h), torch.nn.GELU(),
            torch.nn.Linear(h, h),        torch.nn.GELU(),
            torch.nn.Linear(h, 1),        torch.nn.Softplus())   # error >= 0
    def forward(self, local_errs):                # [B, n_stages]
        return self.net(local_errs).squeeze(-1)   # composed end-to-end error

def allocate(stages, composer, eps, lmbda=1.0, steps=300):
    """max total speedup  s.t.  composed_error <= eps  (one shared budget)."""
    names  = list(stages)
    logits = {k: torch.zeros(len(v), requires_grad=True) for k, v in stages.items()}
    opt = torch.optim.Adam(list(logits.values()), lr=0.05)
    for _ in range(steps):
        opt.zero_grad()
        w        = {k: torch.softmax(logits[k], 0) for k in names}   # soft op-point
        speedup  = sum((w[k] * stages[k][:, 0]).sum() for k in names)
        local    = torch.stack([(w[k] * stages[k][:, 1]).sum() for k in names])
        composed = composer(local.unsqueeze(0))                      # learned, not sum
        loss = -speedup + lmbda * torch.relu(composed - eps)         # Lagrangian
        loss.backward(); opt.step()
    return {k: int(logits[k].argmax()) for k in names}   # op-point chosen per stage

6.2Differentiable, phase-order-free cross-stack autotuning

Phase ordering exists because the stack is applied as a pipeline of independent passes — quantize, then fuse, then shard — whose order is fixed by convention. We propose dissolving the ordering by construction: collect the entire stack configuration into one vector — tile sizes, fusion boundaries, per-tensor precision, sharding dimensions, KV-cache policy, speculation depth — relax the discrete choices with a Gumbel-softmax reparameterization into a continuous space, and train a graph-neural-network surrogate over the computation graph (annotated with the configuration) to predict latency and memory from profiling data. Optimization then proceeds by gradient descent or Bayesian optimization over the unified configuration vector, so there is no sequence of passes and hence no phase-ordering problem. Nearest prior art: DiffTune tunes cost-model parameters with a differentiable surrogate; DARTS108 makes architecture search differentiable; Unity27 and Alpa4 jointly optimize algebraic transformation and parallelization but over a discrete, layer-limited space. The proposed novelty is a single continuous relaxation spanning kernel, precision, parallelism, and serving knobs together, which is what removes phase ordering rather than merely re-ordering the passes better.

import torch, torch.nn.functional as F

# ONE configuration vector spans every layer at once -> no pass ordering.
KNOBS = {
    "tile_m":    [64, 128, 256],
    "fuse_attn": [0, 1],
    "precision": ["fp16", "fp8", "int4"],
    "shard_dim": ["data", "tensor", "pipe"],
    "kv_policy": ["full", "paged", "quant", "evict"],
    "spec_depth":[0, 2, 4, 8],
}

class LatencySurrogate(torch.nn.Module):
    """GNN over the config-annotated compute graph -> (latency, memory).
    Differentiable in the config embedding; trained on profiled runs."""
    def forward(self, graph, cfg_embed): ...      # -> latency, memory

def autotune(graph, surrogate, embed, mem_cap, steps=500):
    # a learnable logit vector for EVERY knob = the whole stack in one tensor
    logits = {k: torch.zeros(len(v), requires_grad=True) for k, v in KNOBS.items()}
    opt = torch.optim.Adam(list(logits.values()), lr=0.1)
    for _ in range(steps):
        opt.zero_grad()
        soft = {k: F.gumbel_softmax(logits[k], tau=0.7, hard=True) for k in KNOBS}
        cfg  = torch.cat([soft[k] @ embed[k] for k in KNOBS])   # continuous relaxation
        latency, memory = surrogate(graph, cfg)
        loss = latency + 1e3 * F.relu(memory - mem_cap)         # min latency, fit memory
        loss.backward(); opt.step()             # gradients flow to ALL knobs jointly
    return {k: KNOBS[k][int(logits[k].argmax())] for k in KNOBS}

6.3Online reoptimization as a control problem

Every autotuner in this survey optimizes offline for an assumed workload, but serving traffic drifts: the distribution of prompt lengths, the composition of each batch, and the prefix-cache hit rate all change over time, so a configuration that was optimal at deployment slowly ceases to be. We propose closing the loop and framing reoptimization as online control. Each configuration — a kernel variant, a chunking policy, a KV-cache budget, a speculation depth — is an arm of a contextual bandit whose context is the observed workload; candidate configurations are trialed through shadow execution (run on mirrored traffic in the background, never on the latency-critical path) so exploration cannot cause a user-visible regression. The novelty is treating optimization as a telemetry-driven control loop rather than a one-shot compilation. Nearest prior art: Tutel75 adapts MoE parallelism to runtime load and adaptive-chunking schedulers react to batch composition, but neither unifies whole-stack configuration selection with safe online exploration under a bandit objective.

import torch

# Reoptimization as online control: pick a stack config per workload context,
# explore ONLY via shadow execution (mirrored traffic, off the hot path).
class LinUCBController:
    def __init__(self, configs, d, alpha=1.0):
        self.configs = configs                          # candidate stack configs
        self.A = [torch.eye(d) for _ in configs]
        self.b = [torch.zeros(d) for _ in configs]
        self.alpha = alpha
    def select(self, ctx):                              # ctx = workload features
        scores = []
        for A, b in zip(self.A, self.b):
            theta = torch.linalg.solve(A, b)
            bonus = self.alpha * (ctx @ torch.linalg.solve(A, ctx)).sqrt()
            scores.append(theta @ ctx + bonus)          # optimism under uncertainty
        return int(torch.stack(scores).argmax())
    def update(self, arm, ctx, reward):                 # reward = goodput (or -latency)
        self.A[arm] += torch.outer(ctx, ctx)
        self.b[arm] += reward * ctx

def control_loop(server, ctl):
    while True:
        ctx  = server.workload_features()   # prompt-len dist, batch mix, cache hit-rate
        arm  = ctl.select(ctx)
        r    = server.shadow_run(ctl.configs[arm])   # SAFE: never on live latency path
        ctl.update(arm, ctx, r)
        if ctl.confident(arm, ctx):
            server.promote(ctl.configs[arm])         # swap into production

6.4Transferable cost models and architecture–system co-search

Two further directions are, we note honestly, more extensions than departures, but they address problems this survey raised directly. First, Section 2.4 observed that a cost model must be re-derived for every GPU generation; a hardware-conditioned meta cost model — a surrogate conditioned on a compact hardware descriptor (SM count, HBM bandwidth, Tensor-Core throughput, interconnect topology, cache sizes) and adapted to a new device with a few calibration measurements via meta-learning109 or a hypernetwork — would let autotuning transfer across Hopper→Blackwell rather than restart. Second, NSA96 and MoBA97 co-design one attention architecture with its kernel by hand; an architecture–system co-search would jointly search architectural knobs (GQA group size, MLA rank, sparse block size, MoE granularity) and system knobs (tiling, sharding, precision) under a single wall-clock objective measured on the target hardware — differing from conventional NAS, which optimizes accuracy or FLOPs, by optimizing measured end-to-end system time including the lower layers.

Common to all four is that the genuine contribution — and the main technical risk — is a single learned model: the error-composition predictor of 6.1, the GNN latency surrogate of 6.2, the bandit value function of 6.3, the transferable cost model of 6.4. Each turns a problem the current stack solves by fixed convention or per-layer heuristic into one solved by search over a model of the whole system. If the taxonomy's lesson is that GPU optimization has become search guided by a cost model, its natural next step is to make that search cross-layer, differentiable, and online.

Conclusion

Getting deep learning to run efficiently on GPUs is a stack-wide effort, and this survey has organized that effort into six layers — operator, graph, parallelization, memory, precision, and serving — unified by a common methodology of search over a defined space under a cost model with an explicit correctness guarantee. Our three exemplars make the breadth concrete: Qrita1 rethinks a single sampling operator and wins 2× throughput with an exact, deterministic algorithm; Mirage and MPK23 search and fuse tensor programs until an entire model becomes one megakernel, cutting latency several-fold; and Alpa4 automates the partitioning of a model across a cluster with a two-level ILP-in-DP search. Around them, dozens of systems — FlashAttention, TVM/Ansor, Megatron-LM, ZeRO/FSDP, PagedAttention, SGLang, FlashInfer, and modern quantization — fill out the stack. The layers compose in every production system yet interact in ways that current tools handle only sequentially, leaving cross-layer joint optimization, cheaper search, end-to-end correctness guarantees, and hardware co-design as the defining challenges of the next several years of GPU optimization for deep learning.

References

1. J. Park, S. Kim, A. Cheung, and I. Stoica. Qrita: High-performance Top-k and Top-p algorithm for GPUs using pivot-based truncation and selection. arXiv preprint arXiv:2602.01518, 2026.
2. M. Wu, X. Cheng, S. Liu, C. Shi, J. Ji, K. Ao, P. Velliengiri, X. Miao, O. Padon, and Z. Jia. Mirage: A multi-level superoptimizer for tensor programs. 19th USENIX Symposium on Operating Systems Design and Implementation (OSDI), 2025. arXiv:2405.05751.
3. X. Cheng, Z. Zhang, Y. Zhou, J. Ji, et al. MPK: A compiler and runtime for mega-kernelizing tensor programs. 20th USENIX Symposium on Operating Systems Design and Implementation (OSDI), 2026. arXiv:2512.22219.
4. L. Zheng, Z. Li, H. Zhang, Y. Zhuang, Z. Chen, Y. Huang, Y. Wang, Y. Xu, D. Zhuo, E. P. Xing, J. E. Gonzalez, and I. Stoica. Alpa: Automating inter- and intra-operator parallelism for distributed deep learning. 16th USENIX Symposium on Operating Systems Design and Implementation (OSDI), 2022.
5. S. Williams, A. Waterman, and D. Patterson. Roofline: An insightful visual performance model for multicore architectures. Communications of the ACM, 52(4):65–76, 2009.
6. S. Chetlur, C. Woolley, P. Vandermersch, J. Cohen, J. Tran, B. Catanzaro, and E. Shelhamer. cuDNN: Efficient primitives for deep learning. arXiv preprint arXiv:1410.0759, 2014.
7. NVIDIA. CUTLASS: CUDA templates for linear algebra subroutines. github.com/NVIDIA/cutlass, 2017–.
8. P. Tillet, H. T. Kung, and D. Cox. Triton: An intermediate language and compiler for tiled neural network computations. ACM SIGPLAN Workshop on Machine Learning and Programming Languages (MAPL), 2019.
9. T. Dao, D. Y. Fu, S. Ermon, A. Rudra, and C. Ré. FlashAttention: Fast and memory-efficient exact attention with IO-awareness. Advances in Neural Information Processing Systems (NeurIPS), 2022. arXiv:2205.14135.
10. T. Dao. FlashAttention-2: Faster attention with better parallelism and work partitioning. arXiv preprint arXiv:2307.08691, 2023.
11. J. Shah, G. Bikshandi, Y. Zhang, V. Thakkar, P. Ramani, and T. Dao. FlashAttention-3: Fast and accurate attention with asynchrony and low-precision. Advances in Neural Information Processing Systems (NeurIPS), 2024.
12. T. Chen, T. Moreau, Z. Jiang, L. Zheng, E. Yan, M. Cowan, H. Shen, L. Wang, Y. Hu, L. Ceze, C. Guestrin, and A. Krishnamurthy. TVM: An automated end-to-end optimizing compiler for deep learning. 13th USENIX Symposium on Operating Systems Design and Implementation (OSDI), 2018.
13. L. Zheng, C. Jia, M. Sun, Z. Wu, C. H. Yu, A. Haj-Ali, Y. Wang, J. Yang, D. Zhuo, K. Sen, J. E. Gonzalez, and I. Stoica. Ansor: Generating high-performance tensor programs for deep learning. 14th USENIX Symposium on Operating Systems Design and Implementation (OSDI), 2020.
14. Z. Jia, O. Padon, J. Thomas, T. Warszawski, M. Zaharia, and A. Aiken. TASO: Optimizing deep learning computation with automatic generation of graph substitutions. 27th ACM Symposium on Operating Systems Principles (SOSP), 2019.
15. H. Wang, J. Zhai, M. Gao, Z. Ma, S. Tang, L. Zheng, Y. Li, K. Rong, Y. Chen, and Z. Jia. PET: Optimizing tensor programs with partially equivalent transformations and automated corrections. 15th USENIX Symposium on Operating Systems Design and Implementation (OSDI), 2021.
16. A. Holtzman, J. Buys, L. Du, M. Forbes, and Y. Choi. The curious case of neural text degeneration. International Conference on Learning Representations (ICLR), 2020.
17. Google. XLA: Optimizing compiler for machine learning. openxla.org/xla, 2017–.
18. NVIDIA. TensorRT: A high-performance deep learning inference optimizer and runtime. developer.nvidia.com/tensorrt, 2016–.
19. J. Ansel, E. Yang, H. He, N. Gimelshein, et al. PyTorch 2: Faster machine learning through dynamic Python bytecode transformation and graph compilation. 29th ACM International Conference on Architectural Support for Programming Languages and Operating Systems (ASPLOS), 2024.
20. ONNX Runtime developers. ONNX Runtime: Cross-platform, high-performance ML inferencing and training accelerator. onnxruntime.ai, 2018–.
21. M. Shoeybi, M. Patwary, R. Puri, P. LeGresley, J. Casper, and B. Catanzaro. Megatron-LM: Training multi-billion parameter language models using model parallelism. arXiv preprint arXiv:1909.08053, 2019.
22. Y. Huang, Y. Cheng, A. Bapna, O. Firat, M. X. Chen, D. Chen, H. Lee, J. Ngiam, Q. V. Le, Y. Wu, and Z. Chen. GPipe: Efficient training of giant neural networks using pipeline parallelism. Advances in Neural Information Processing Systems (NeurIPS), 2019.
23. D. Narayanan, A. Harlap, A. Phanishayee, V. Seshadri, N. R. Devanur, G. R. Ganger, P. B. Gibbons, and M. Zaharia. PipeDream: Generalized pipeline parallelism for DNN training. 27th ACM Symposium on Operating Systems Principles (SOSP), 2019.
24. D. Lepikhin, H. Lee, Y. Xu, D. Chen, O. Firat, Y. Huang, M. Krikun, N. Shazeer, and Z. Chen. GShard: Scaling giant models with conditional computation and automatic sharding. arXiv preprint arXiv:2006.16668, 2020.
25. N. Shazeer, Y. Cheng, N. Parmar, D. Tran, A. Vaswani, P. Koanantakool, P. Hawkins, H. Lee, M. Hong, C. Young, R. Sepassi, and B. Hechtman. Mesh-TensorFlow: Deep learning for supercomputers. Advances in Neural Information Processing Systems (NeurIPS), 2018.
26. Y. Xu, H. Lee, D. Chen, B. Hechtman, Y. Huang, R. Joshi, M. Krikun, D. Lepikhin, A. Ly, M. Maggioni, et al. GSPMD: General and scalable parallelization for ML computation graphs. arXiv preprint arXiv:2105.04663, 2021.
27. C. Unger, Z. Jia, W. Wu, S. Lin, M. Baines, C. E. Q. Narvaez, V. Ramakrishnaiah, et al. Unity: Accelerating DNN training through joint optimization of algebraic transformations and parallelization. 16th USENIX Symposium on Operating Systems Design and Implementation (OSDI), 2022.
28. X. Miao, Y. Wang, Y. Jiang, C. Shi, X. Nie, H. Zhang, and B. Cui. Galvatron: Efficient transformer training over multiple GPUs using automatic parallelism. Proceedings of the VLDB Endowment, 16(3):470–479, 2022.
29. T. Chen, B. Xu, C. Zhang, and C. Guestrin. Training deep nets with sublinear memory cost. arXiv preprint arXiv:1604.06174, 2016.
30. S. Rajbhandari, J. Rasley, O. Ruwase, and Y. He. ZeRO: Memory optimizations toward training trillion parameter models. International Conference for High Performance Computing, Networking, Storage and Analysis (SC), 2020.
31. Y. Zhao, A. Gu, R. Varma, L. Luo, C.-C. Huang, M. Xu, L. Wright, H. Shojanazeri, M. Ott, S. Shleifer, et al. PyTorch FSDP: Experiences on scaling fully sharded data parallel. Proceedings of the VLDB Endowment, 16(12):3848–3860, 2023.
32. S. Rajbhandari, O. Ruwase, J. Rasley, S. Smith, and Y. He. ZeRO-Infinity: Breaking the GPU memory wall for extreme scale deep learning. International Conference for High Performance Computing, Networking, Storage and Analysis (SC), 2021.
33. Y. Sheng, L. Zheng, B. Yuan, Z. Li, M. Ryabinin, D. Y. Fu, Z. Xie, B. Chen, C. Barrett, J. E. Gonzalez, P. Liang, C. Ré, I. Stoica, and C. Zhang. FlexGen: High-throughput generative inference of large language models with a single GPU. International Conference on Machine Learning (ICML), 2023.
34. W. Kwon, Z. Li, S. Zhuang, Y. Sheng, L. Zheng, C. H. Yu, J. E. Gonzalez, H. Zhang, and I. Stoica. Efficient memory management for large language model serving with PagedAttention. 29th ACM Symposium on Operating Systems Principles (SOSP), 2023.
35. Z. Ye, L. Chen, R. Lai, W. Lin, Y. Zhang, S. Wang, T. Chen, B. Kasikci, V. Grover, A. Krishnamurthy, and L. Ceze. FlashInfer: Efficient and customizable attention engine for LLM inference serving. Proceedings of Machine Learning and Systems (MLSys), 2025. arXiv:2501.01005.
36. P. Micikevicius, S. Narang, J. Alben, G. Diamos, E. Elsen, D. Garcia, B. Ginsburg, M. Houston, O. Kuchaiev, G. Venkatesh, and H. Wu. Mixed precision training. International Conference on Learning Representations (ICLR), 2018.
37. E. Frantar, S. Ashkboos, T. Hoefler, and D. Alistarh. GPTQ: Accurate post-training quantization for generative pre-trained transformers. International Conference on Learning Representations (ICLR), 2023.
38. G. Xiao, J. Lin, M. Seznec, H. Wu, J. Demouth, and S. Han. SmoothQuant: Accurate and efficient post-training quantization for large language models. International Conference on Machine Learning (ICML), 2023.
39. J. Lin, J. Tang, H. Tang, S. Yang, W.-M. Chen, W.-C. Wang, G. Xiao, X. Dang, C. Gan, and S. Han. AWQ: Activation-aware weight quantization for LLM compression and acceleration. Proceedings of Machine Learning and Systems (MLSys), 2024.
40. A. Mishra, J. A. Latorre, J. Pool, D. Stosic, D. Stosic, G. Venkatesh, C. Yu, and P. Micikevicius. Accelerating sparse deep neural networks. arXiv preprint arXiv:2104.08378, 2021.
41. G.-I. Yu, J. S. Jeong, G.-W. Kim, S. Kim, and B.-G. Chun. Orca: A distributed serving system for transformer-based generative models. 16th USENIX Symposium on Operating Systems Design and Implementation (OSDI), 2022.
42. L. Zheng, L. Yin, Z. Xie, C. Sun, J. Huang, C. H. Yu, S. Cao, C. Kozyrakis, I. Stoica, J. E. Gonzalez, C. Barrett, and Y. Sheng. SGLang: Efficient execution of structured language model programs. Advances in Neural Information Processing Systems (NeurIPS), 2024. arXiv:2312.07104.
43. Y. Leviathan, M. Kalman, and Y. Matias. Fast inference from transformers via speculative decoding. International Conference on Machine Learning (ICML), 2023.
44. C. Chen, S. Borgeaud, G. Irving, J.-B. Lespiau, L. Sifre, and J. Jumper. Accelerating large language model decoding with speculative sampling. arXiv preprint arXiv:2302.01318, 2023.
45. T. Cai, Y. Li, Z. Geng, H. Peng, J. D. Lee, D. Chen, and T. Dao. Medusa: Simple LLM inference acceleration framework with multiple decoding heads. International Conference on Machine Learning (ICML), 2024.
46. B. Spector, S. Arora, A. Singhal, D. Y. Fu, and C. Ré. ThunderKittens: Simple, fast, and adorable AI kernels. International Conference on Learning Representations (ICLR), 2025. arXiv:2410.20399.
47. H. He, D. Guessous, Y. Liang, and J. Dong. FlexAttention: The flexibility of PyTorch with the performance of FlashAttention. PyTorch Blog, 2024.
48. N. Shazeer. Fast transformer decoding: One write-head is all you need. arXiv preprint arXiv:1911.02150, 2019.
49. J. Ainslie, J. Lee-Thorp, M. de Jong, Y. Zemlyanskiy, F. Lebrón, and S. Sanghai. GQA: Training generalized multi-query transformer models from multi-head checkpoints. Conference on Empirical Methods in Natural Language Processing (EMNLP), 2023. arXiv:2305.13245.
50. DeepSeek-AI. DeepSeek-V2: A strong, economical, and efficient mixture-of-experts language model. arXiv preprint arXiv:2405.04434, 2024.
51. G. Xiao, Y. Tian, B. Chen, S. Han, and M. Lewis. Efficient streaming language models with attention sinks. International Conference on Learning Representations (ICLR), 2024. arXiv:2309.17453.
52. Z. Liu, J. Yuan, H. Jin, S. Zhong, Z. Xu, V. Braverman, B. Chen, and X. Hu. KIVI: A tuning-free asymmetric 2-bit quantization for KV cache. International Conference on Machine Learning (ICML), 2024. arXiv:2402.02750.
53. C. Hooper, S. Kim, H. Mohammadzadeh, M. W. Mahoney, Y. S. Shao, K. Keutzer, and A. Gholami. KVQuant: Towards 10 million context length LLM inference with KV cache quantization. Advances in Neural Information Processing Systems (NeurIPS), 2024. arXiv:2401.18079.
54. H. Liu, M. Zaharia, and P. Abbeel. Ring attention with blockwise transformers for near-infinite context. arXiv preprint arXiv:2310.01889, 2023.
55. S. A. Jacobs, M. Tanaka, C. Zhang, M. Zhang, S. L. Song, S. Rajbhandari, and Y. He. DeepSpeed Ulysses: System optimizations for enabling training of extreme long sequence transformer models. arXiv preprint arXiv:2309.14509, 2023.
56. V. Korthikanti, J. Casper, S. Lym, L. McAfee, M. Andersch, M. Shoeybi, and B. Catanzaro. Reducing activation recomputation in large transformer models. Proceedings of Machine Learning and Systems (MLSys), 2023. arXiv:2205.05198.
57. P. Qi, X. Wan, G. Huang, and M. Lin. Zero bubble pipeline parallelism. International Conference on Learning Representations (ICLR), 2024. arXiv:2401.10241.
58. DeepSeek-AI. DeepSeek-V3 technical report. arXiv preprint arXiv:2412.19437, 2025.
59. L.-W. Chang, W. Bao, Q. Hou, C. Jiang, N. Zheng, Y. Zhong, X. Zhang, Z. Song, Z. Jiang, H. Lin, X. Jin, and X. Liu. Flux: Fast software-based communication overlap on GPUs through kernel fusion. arXiv preprint arXiv:2406.06858, 2024.
60. T. Dettmers, A. Pagnoni, A. Holtzman, and L. Zettlemoyer. QLoRA: Efficient finetuning of quantized LLMs. Advances in Neural Information Processing Systems (NeurIPS), 2023. arXiv:2305.14314.
61. S. Ashkboos, A. Mohtashami, M. L. Croci, B. Li, M. Jaggi, D. Alistarh, T. Hoefler, and J. Hensman. QuaRot: Outlier-free 4-bit inference in rotated LLMs. Advances in Neural Information Processing Systems (NeurIPS), 2024. arXiv:2404.00456.
62. Z. Liu, C. Zhao, I. Fedorov, B. Soran, D. Choudhary, R. Krishnamoorthi, V. Chandra, Y. Tian, and T. Blankevoort. SpinQuant: LLM quantization with learned rotations. International Conference on Learning Representations (ICLR), 2025. arXiv:2405.16406.
63. Y. Zhong, S. Liu, J. Chen, J. Hu, Y. Zhu, X. Liu, X. Jin, and H. Zhang. DistServe: Disaggregating prefill and decoding for goodput-optimized LLM serving. 18th USENIX Symposium on Operating Systems Design and Implementation (OSDI), 2024. arXiv:2401.09670.
64. A. Agrawal, A. Panwar, J. Mohan, N. Kwatra, B. S. Gulavani, and R. Ramjee. SARATHI: Efficient LLM inference by piggybacking decodes with chunked prefills. arXiv preprint arXiv:2308.16369, 2023.
65. P. Patel, E. Choukse, C. Zhang, A. Shah, Í. Goiri, S. Maleki, and R. Bianchini. Splitwise: Efficient generative LLM inference using phase splitting. ACM/IEEE International Symposium on Computer Architecture (ISCA), 2024. arXiv:2311.18677.
66. R. Qin, Z. Li, W. He, M. Zhang, Y. Wu, W. Zheng, and X. Xu. Mooncake: A KVCache-centric disaggregated architecture for LLM serving. arXiv preprint arXiv:2407.00079, 2024.
67. Y. Li, F. Wei, C. Zhang, and H. Zhang. EAGLE: Speculative sampling requires rethinking feature uncertainty. International Conference on Machine Learning (ICML), 2024. arXiv:2401.15077.
68. Y. Fu, P. Bailis, I. Stoica, and H. Zhang. Break the sequential dependency of LLM inference using lookahead decoding. International Conference on Machine Learning (ICML), 2024. arXiv:2402.02057.
69. J. Juravsky, B. Brown, R. Ehrlich, D. Y. Fu, C. Ré, and A. Mirhoseini. Hydragen: High-throughput LLM inference with shared prefixes. arXiv preprint arXiv:2402.05099, 2024.
70. Y. Sheng, S. Cao, D. Li, C. Hooper, N. Lee, S. Yang, C. Chou, B. Zhu, L. Zheng, K. Keutzer, J. E. Gonzalez, and I. Stoica. S-LoRA: Serving thousands of concurrent LoRA adapters. Proceedings of Machine Learning and Systems (MLSys), 2024. arXiv:2311.03285.
71. L. Chen, Z. Ye, Y. Wu, D. Zhuo, L. Ceze, and A. Krishnamurthy. Punica: Multi-tenant LoRA serving. Proceedings of Machine Learning and Systems (MLSys), 2024. arXiv:2310.18547.
72. A. Gu and T. Dao. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint arXiv:2312.00752, 2023.
73. T. Dao and A. Gu. Transformers are SSMs: Generalized models and efficient algorithms through structured state space duality. International Conference on Machine Learning (ICML), 2024. arXiv:2405.21060.
74. S. Rajbhandari, C. Li, Z. Yao, M. Zhang, R. Y. Aminabadi, A. A. Awan, J. Rasley, and Y. He. DeepSpeed-MoE: Advancing mixture-of-experts inference and training to power next-generation AI scale. International Conference on Machine Learning (ICML), 2022. arXiv:2201.05596.
75. C. Hwang, W. Cui, Y. Xiong, Z. Yang, Z. Liu, H. Hu, Z. Wang, R. Salas, J. Jose, P. Ram, J. Chau, P. Cheng, F. Yang, M. Yang, and Y. Xiong. Tutel: Adaptive mixture-of-experts at scale. Proceedings of Machine Learning and Systems (MLSys), 2023. arXiv:2206.03382.
76. A. Tseng, J. Chee, Q. Sun, V. Kuleshov, and C. De Sa. QuIP#: Even better LLM quantization with Hadamard incoherence and lattice codebooks. International Conference on Machine Learning (ICML), 2024. arXiv:2402.04396.
77. V. Egiazarian, A. Panferov, D. Kuznedelev, E. Frantar, A. Babenko, and D. Alistarh. Extreme compression of large language models via additive quantization. International Conference on Machine Learning (ICML), 2024. arXiv:2401.06118.
78. S. Ma, H. Wang, L. Ma, L. Wang, W. Wang, S. Huang, L. Dong, R. Wang, J. Xue, and F. Wei. The era of 1-bit LLMs: All large language models are in 1.58 bits. arXiv preprint arXiv:2402.17764, 2024.
79. E. Frantar and D. Alistarh. SparseGPT: Massive language models can be accurately pruned in one-shot. International Conference on Machine Learning (ICML), 2023. arXiv:2301.00774.
80. M. Sun, Z. Liu, A. Bair, and J. Z. Kolter. A simple and effective pruning approach for large language models. International Conference on Learning Representations (ICLR), 2024. arXiv:2306.11695.
81. Z. Liu, J. Wang, T. Dao, T. Zhou, B. Yuan, Z. Song, A. Shrivastava, C. Zhang, Y. Tian, C. Ré, and B. Chen. Deja Vu: Contextual sparsity for efficient LLMs at inference time. International Conference on Machine Learning (ICML), 2023. arXiv:2310.17157.
82. Y. Song, Z. Mi, H. Xie, and H. Chen. PowerInfer: Fast large language model serving with a consumer-grade GPU. arXiv preprint arXiv:2312.12456, 2023.
83. B. Wu, Y. Zhong, Z. Zhang, S. Liu, F. Liu, Y. Sun, G. Huang, X. Liu, and X. Jin. Fast distributed inference serving for large language models. arXiv preprint arXiv:2305.05920, 2023.
84. B. Sun, Z. Huang, H. Zhao, W. Xiao, X. Zhang, Y. Li, and W. Lin. Llumnix: Dynamic scheduling for large language model serving. 18th USENIX Symposium on Operating Systems Design and Implementation (OSDI), 2024. arXiv:2406.03243.
85. B. T. Willard and R. Louf. Efficient guided generation for large language models. arXiv preprint arXiv:2307.09702, 2023.
86. Y. Dong, C. F. Ruan, Y. Cai, R. Lai, Z. Xu, Y. Zhao, and T. Chen. XGrammar: Flexible and efficient structured generation engine for large language models. Proceedings of Machine Learning and Systems (MLSys), 2025. arXiv:2411.15100.
87. R. Pope, S. Douglas, A. Chowdhery, J. Devlin, J. Bradbury, A. Levskaya, J. Heek, K. Xiao, S. Agrawal, and J. Dean. Efficiently scaling transformer inference. Proceedings of Machine Learning and Systems (MLSys), 2023. arXiv:2211.05102.
88. NVIDIA. NVIDIA H100 Tensor Core GPU architecture (Hopper). Whitepaper, 2022.
89. NVIDIA. NVIDIA Blackwell architecture. Technical brief, 2024.
90. Y. Ding, C. H. Yu, B. Zheng, Y. Liu, Y. Wang, and G. Pekhimenko. Hidet: Task-mapping programming paradigm for deep learning tensor programs. 28th ACM International Conference on Architectural Support for Programming Languages and Operating Systems (ASPLOS), 2023.
91. T. Dettmers, M. Lewis, S. Shleifer, and L. Zettlemoyer. 8-bit optimizers via block-wise quantization. International Conference on Learning Representations (ICLR), 2022. arXiv:2110.02861.
92. G. Wang, H. Qin, S. A. Jacobs, C. Holmes, S. Rajbhandari, O. Ruwase, F. Yan, L. Yang, and Y. He. ZeRO++: Extremely efficient collective communication for giant model training. International Conference on Learning Representations (ICLR), 2024. arXiv:2306.10209.
93. J. Zhao, Z. Zhang, B. Chen, Z. Wang, A. Anandkumar, and Y. Tian. GaLore: Memory-efficient LLM training by gradient low-rank projection. International Conference on Machine Learning (ICML), 2024. arXiv:2403.03507.
94. I. Gim, G. Chen, S.-s. Lee, N. Sarda, A. Khandelwal, and L. Zhong. Prompt Cache: Modular attention reuse for low-latency inference. Proceedings of Machine Learning and Systems (MLSys), 2024. arXiv:2311.04934.
95. B. Gao, Z. He, P. Sharma, Q. Kang, D. Jevdjic, J. Deng, X. Yang, Z. Yu, and P. Zuo. Cost-efficient large language model serving for multi-turn conversations with CachedAttention. USENIX Annual Technical Conference (ATC), 2024. arXiv:2403.19708.
96. J. Yuan, H. Gao, D. Dai, J. Luo, L. Zhao, Z. Zhang, Z. Xie, Y. X. Wei, L. Wang, Z. Xiao, et al. Native sparse attention: Hardware-aligned and natively trainable sparse attention. Annual Meeting of the Association for Computational Linguistics (ACL), 2025. arXiv:2502.11089.
97. E. Lu, Z. Jiang, J. Liu, Y. Du, T. Jiang, C. Hong, S. Liu, W. He, E. Yuan, Y. Wang, et al. MoBA: Mixture of block attention for long-context LLMs. arXiv preprint arXiv:2502.13189, 2025.
98. Z. Zhang, Y. Sheng, T. Zhou, T. Chen, L. Zheng, R. Cai, Z. Song, Y. Tian, C. Ré, C. Barrett, Z. Wang, and B. Chen. H2O: Heavy-hitter oracle for efficient generative inference of large language models. Advances in Neural Information Processing Systems (NeurIPS), 2023. arXiv:2306.14048.
99. Y. Li, Y. Huang, B. Yang, B. Venkitesh, A. Locatelli, H. Ye, T. Cai, P. Lewis, and D. Chen. SnapKV: LLM knows what you are looking for before generation. Advances in Neural Information Processing Systems (NeurIPS), 2024. arXiv:2404.14469.
100. Z. Cai, Y. Zhang, B. Gao, Y. Liu, T. Liu, K. Lu, W. Xiong, Y. Dong, B. Chang, J. Hu, et al. PyramidKV: Dynamic KV cache compression based on pyramidal information funneling. arXiv preprint arXiv:2406.02069, 2024.
101. X. Miao, G. Oliaro, Z. Zhang, X. Cheng, Z. Wang, R. Y. Y. Wong, Z. Chen, D. Arfeen, R. Abhyankar, and Z. Jia. SpecInfer: Accelerating generative LLM serving with speculative inference and token tree verification. 29th ACM International Conference on Architectural Support for Programming Languages and Operating Systems (ASPLOS), 2024. arXiv:2305.09781.
102. K. Hong, G. Dai, J. Xu, Q. Mao, X. Li, J. Liu, K. Chen, Y. Dong, and Y. Wang. FlashDecoding++: Faster large language model inference on GPUs. Proceedings of Machine Learning and Systems (MLSys), 2024. arXiv:2311.01282.
103. S. Jeaugey and NVIDIA. NCCL: Optimized primitives for collective multi-GPU communication. developer.nvidia.com/nccl, 2017–.
104. G. Wang, S. Venkataraman, A. Phanishayee, N. Devanur, J. Thelin, and I. Stoica. Blink: Fast and generic collectives for distributed ML. Proceedings of Machine Learning and Systems (MLSys), 2020. arXiv:1910.04940.
105. A. Shah, V. Chidambaram, M. Cowan, S. Maleki, M. Musuvathi, T. Mytkowicz, J. Nelson, O. Saarikivi, and R. Singh. TACCL: Guiding collective algorithm synthesis using communication sketches. 20th USENIX Symposium on Networked Systems Design and Implementation (NSDI), 2023. arXiv:2111.04867.
106. B. Wu, S. Liu, Y. Zhong, P. Sun, X. Liu, and X. Jin. LoongServe: Efficiently serving long-context large language models with elastic sequence parallelism. ACM SIGOPS 30th Symposium on Operating Systems Principles (SOSP), 2024. arXiv:2404.09526.
107. Z. Dong, Z. Yao, A. Gholami, M. W. Mahoney, and K. Keutzer. HAWQ: Hessian AWare Quantization of neural networks with mixed-precision. IEEE/CVF International Conference on Computer Vision (ICCV), 2019. arXiv:1905.03696.
108. H. Liu, K. Simonyan, and Y. Yang. DARTS: Differentiable architecture search. International Conference on Learning Representations (ICLR), 2019. arXiv:1806.09055.
109. C. Finn, P. Abbeel, and S. Levine. Model-agnostic meta-learning for fast adaptation of deep networks. International Conference on Machine Learning (ICML), 2017. arXiv:1703.03400.

@techreport{doan2026gpusurvey,
  title   = {Optimizing Deep Learning Models on GPUs: A Cross-Stack
             Survey from Sampling Kernels to Auto-Parallelized Clusters},
  author  = {Doan, Chi},
  year    = {2026},
  institution = {Independent Researcher, Hanoi, Vietnam},
  type    = {Preprint / Technical Report}
}