Inference Acceleration

Appearance

Computation Graph Optimization

At a glance The computation graph (IR) is the hub of an inference engine — constant folding, dead-code elimination, operator fusion, layout propagation, CUDA graph capture, and prefix caching are all graph-level optimizations. This article covers the ONNX/MLIR/XLA HLO IRs, static vs. dynamic shapes, and the TVM Relax/Inductor/XLA/Triton compiler stack.

Computation Graph Optimization ​

Concept Definition: From Operators to Graphs ​

A deep learning model is, at its core, a directed acyclic graph (DAG) of operators — each node is a matmul/conv/relu, and each edge is a tensor data flow. Computation graph optimization is the process of rewriting, simplifying, and scheduling this DAG so the inference engine can "precompute what can be precomputed, fuse what can be fused, and delete what can be deleted" before execution.

Two key insights for graph optimization:

  1. Graph optimization is a "macro-level" optimization — one level above Kernel Fusion and Custom Kernels, covering operator selection, layout conversion, memory reuse, and control-flow simplification;
  2. Graph optimization for LLM inference differs from traditional CV models — LLM sequence lengths are dynamic, there is autoregressive control flow, and the KV cache must be reused across steps; many optimizations from classical compiler stacks (like ONNX Runtime) don't apply.

LLM inference engines (vLLM, TensorRT-LLM, llama.cpp) each have a set of graph optimization passes; the core ones are: constant folding, dead-code elimination, operator fusion, CUDA graph capture, prefix caching.

1. Computation Graph IRs: The Compiler's Hub ​

Graph optimization needs an intermediate representation (IR) — something that lets users write models and lets compilers manipulate them. Mainstream IRs:

IRCreatorCharacteristicsTypical Use
ONNXMicrosoft + FacebookNeutral format, cross-frameworkModel export/exchange
MLIRLLVMMulti-level, extensibleCompiler research stacks
XLA HLOGoogleXLA compiler IRJAX/TF
TorchScript / FX / ExportPyTorchPyTorch-internal graphsPyTorch ecosystem
TVM RelaxApache TVMIR purpose-built for LLM inferenceResearch-grade inference engines
TensorRT networkNVIDIAClosed-source but high-performanceTensorRT inference
PaddlePaddle IRBaiduPaddle-internalPaddle ecosystem

The Different Roles of ONNX and PyTorch IRs

  • ONNX is an "exchange format" — HuggingFace models are exported to ONNX for cross-engine use (both ONNX Runtime: Cross-Platform and TensorRT support ONNX import);
  • PyTorch IRs are "working formats" — the graph is dynamic during training, and torch.export converts it to a static graph for Inductor compilation at deployment. Common LLM inference paths: HF model → optimum exports ONNX → ONNX Runtime deployment; or HF model → torch.compile → Inductor optimization → deployment.

2. Classic Graph Optimization Passes ​

1. Constant Folding ​

Some nodes in the graph have "fixed inputs + fixed weights → fixed outputs" and can be computed offline:

text
Before: x_fixed → Multiply → Linear → ...    (Multiply+Linear computed on every inference)
After:  replace with a precomputed constant → ConstValue → ...

Example: the cos/sin tables of RoPE are constants and can be precomputed and baked into the graph.

2. Dead Code Elimination (DCE) ​

Some nodes in the graph don't affect the final output (e.g. an unused residual path) — delete them outright:

text
Before: x → LayerA → LayerB → out
              ↓
             LayerC (no downstream consumer)
After:  x → LayerA → LayerB → out  (LayerC removed)

3. Operator Fusion ​

See Kernel Fusion and Custom Kernels. At the graph level, the fusion pass marks operator combinations that "can be merged into one kernel" and hands them to the backend (Inductor, Triton, CUDA) to generate the fused kernel.

4. Layout Conversion + Layout Propagation ​

Different hardware and different operators prefer different memory layouts:

text
NCHW (PyTorch default)  vs.  NHWC (recommended for Tensor Cores)  vs.  channels-last

The graph optimization pass decides where to insert transposes — as few as possible, and ideally fused into the epilogue of adjacent matmuls. Layout propagation is a core pass of XLA/Inductor.

5. Memory Reuse (Memory Planning) ​

The graph optimization pass analyzes the lifetime of every tensor and reuses the same memory:

text
Before: Tensor1 takes 1GB, Tensor2 takes 1GB → 2GB total
After:  Tensor1 and Tensor2 have non-overlapping lifetimes → share 1GB

KV cache management in LLM inference and the paged memory management of vLLM and PagedAttention are the ultimate form of this idea.

3. Graph Optimizations Specific to LLM Inference ​

1. CUDA Graph Capture ​

CUDA graphs (introduced by NVIDIA in 2017) let GPU operations be "recorded" and replayed, eliminating per-kernel launch overhead. LLM inference launches dozens of kernels per generated token, and launch overhead can reach 10-20%.

text
Normal execution: CPU → launch kernel1 → kernel2 → kernel3 → ... → kernelN
                                ↑ launch overhead on every kernel (µs-level)
CUDA graph:       record once → replay the whole graph afterwards; the CPU issues a single launch
  • vLLM: use_cuda_graph=True on by default;
  • TensorRT-LLM: a core feature;
  • Limitation: CUDA graphs require fixed shapes at capture time → the decode phase captures one graph per batch size.

The Tension Between CUDA Graphs and Dynamic Shapes

CUDA graphs require the recorded shapes to stay fixed during replay. But LLM inference batch sizes change (continuous batching):

  • Solution 1: bucket by batch size (capture 1/2/4/8/16/...), pick the nearest bucket at runtime;
  • Solution 2: pad to the next bucket size (wastes compute but keeps graph reuse);
  • Solution 3: dynamic graphs (conditional nodes in CUDA graphs, supported from H100).

vLLM combines 1+2; TensorRT-LLM uses 1.

2. Prefix Caching ​

Multiple requests sharing the same prompt prefix reuse the already-computed KV cache:

text
Before: request 1 [system_prompt + user1_msg] → compute the entire KV cache
        request 2 [system_prompt + user2_msg] → compute the entire KV cache again
After:  cache the KV of system_prompt; requests 1 and 2 both reuse it
  • vLLM: enable_prefix_caching=True on by default;
  • SGLang: built-in RadixAttention implements an efficient prefix tree;
  • Gain: with long system prompts (4K+ tokens), TTFT drops by 50-90%.

3. Graph Optimizations for Speculative Decoding ​

Speculative decoding has a small model generate candidate tokens first, and a large model verify them in batch — see Speculative Decoding and Medusa/EAGLE. From the graph optimization angle:

  • The small and large models execute in parallel, so the batch verification must be graph-captured;
  • Tree decoding (tree-shaped multi-token guesses) requires sparse attention mask patterns → kernels like FlashInfer support this.

4. Weight Prepacking ​

Convert the weight storage layout into a hardware-friendly format once at load time:

  • Unpacking and rearranging INT4 weights;
  • Transposing BMM weights (so matmul takes cuBLAS's fastest path);
  • Pre-allocating the paged layout of the KV cache.

All of this happens "at graph load time" with zero cost "at inference time."

4. Static Shapes vs. Dynamic Shapes ​

Traditional DNNs ​

CV models (ResNet, ViT) have fully static shapes — 224×224 input, batch 32 fixed throughout. Graph optimization can do heavy static reasoning:

  • Constant folding can precompute all shape-related constants;
  • Operator selection takes the fastest Tensor Core path;
  • Memory planning is precise to the byte.

The LLM Dilemma ​

LLM sequence lengths are dynamic:

  • Prompt lengths range from 100 to 8000;
  • Decode grows the length by 1 per generated token;
  • The KV cache must expand dynamically with generation.

Solution paths:

ApproachHowCost
Padded shapePad all sequences to max_lenWastes compute
Mask-basedCompute at the true length; attention mask hides paddingSaves compute but kernels get complex
Bucketed shapeBucket sequences by length; one CUDA graph per bucketBalanced — the industry mainstream
Dynamic shapeThe compiler supports truly dynamic shapesHard to implement; performance may suffer

vLLM's Practical Approach

vLLM combines bucketed + mask-based:

  • Prefill phase: bucket by prompt length (e.g. 128/256/512/1024/2048/...), optimizing each bucket separately;
  • Decode phase: all batches share the max_gen_len bucket; the attention mask hides finished requests;
  • Plus PagedAttention to handle the dynamic length of the KV cache. This is the de facto standard of shape management for industrial LLM inference.

5. The Compiler Stack Landscape ​

The modern compiler stack for LLM inference:

1. PyTorch Inductor (torch.compile) ​

  • Input: PyTorch eager models;
  • Output: Triton kernels (GPU) or C++ kernels (CPU);
  • Strength: seamless integration with PyTorch training;
  • LLM suitability: scenarios where training and inference are unified (e.g. using torch.compile to speed up inference of HuggingFace models);
  • Limitation: weak at dynamic shapes and continuous batching.

2. XLA (Google) ​

  • Input: JAX/TF HLO;
  • Output: GPU/CPU/TPU kernels;
  • Strength: stable, cross-hardware;
  • LLM suitability: Google-internal large model training/inference; the JAX ecosystem (e.g. MaxText, EasyLM).

3. TVM Relax (Apache) ​

  • Input: Relax IR (purpose-built to support LLM dynamic shapes);
  • Output: Triton/CUDA kernels;
  • Strength: research-friendly, supports LLM dynamic shapes;
  • LLM suitability: research-grade inference engines (the MLC-LLM project uses TVM);
  • Limitation: ecosystem weaker than PyTorch.

4. Triton (OpenAI) ​

  • Both a kernel DSL and a compiler;
  • Input: Python-like kernels;
  • Output: efficient CUDA;
  • Strength: the benchmark for hand-written kernels — nearly every LLM inference engine uses it;
  • Details: Kernel Fusion and Custom Kernels.

5. TensorRT / TensorRT-LLM (NVIDIA) ​

  • Input: ONNX / HF models;
  • Output: closed-source high-performance kernels;
  • Strength: optimal on NVIDIA hardware, full INT8/FP8/FP4 support;
  • LLM suitability: TensorRT-LLM is the flagship for LLM inference (see TensorRT and GPU Inference).

6. Graph Optimization in Practice Across LLM Inference Engines ​

The graph optimization stacks of mainstream LLM inference engines:

EngineGraph Optimization StackDetails
vLLMPyTorch eager + Marlin/FlashAttention/FlashInfer + CUDA graphsvLLM and PagedAttention
TensorRT-LLMTensorRT + in-flight batching + INT8/FP8/FP4 kernelsTensorRT and GPU Inference
SGLangRadixAttention + CUDA graphs + FlashInfer—
ONNX RuntimeONNX graphs + classic fusion passes + IO-aware kernelsONNX Runtime: Cross-Platform
OpenVINOIntel graph optimization stack + CPU/GPU/VPU backendsOpenVINO and CPU Inference
llama.cppHand-written C++ graphs + GGUF quantizationllama.cpp and GGUF

Why LLM Inference Graph Optimization Differs from Traditional DNNs

Traditional DNN graph optimization focuses on operator fusion + constant folding + static-shape scheduling — compiling the whole graph once.

LLM inference adds a dimension: autoregressive control flow. Every generated token must:

  1. Schedule the next batch (continuous batching);
  2. Update the KV cache;
  3. Check for stop tokens;
  4. Re-enter graph execution.

This loop + control flow structure invalidates many optimizations of classical compiler stacks (like ONNX Runtime) — a "dynamic graph executor" like vLLM/TensorRT-LLM becomes necessary.

7. Engineering Trade-offs ​

DimensionOptionsRules of Thumb
Graph vs. eagerStatic graph (TensorRT) vs. dynamic graph (vLLM)LLM serving uses vLLM (dynamic); offline batch processing uses TensorRT-LLM (static)
CUDA graphs on or offOn (vLLM default) vs. offOn for 99% of scenarios; off only while debugging
Prefix caching on or offOnNearly free of side effects — always on for long prompts
Bucket sizeMultiple tiers (16/32/64/128)Large batches use large buckets, but don't have too many buckets (each capture takes memory)
Compilation time vs. inference performanceLong (TensorRT takes hours) vs. short (vLLM takes seconds)Online serving uses vLLM; offline extremes use TensorRT-LLM
Stacking quantizationGraph optimization after quantizationWeight-Only Quantization and Mixed Precision + graph optimization is the golden combo

8. Future Directions ​

Research hotspots in LLM inference graph optimization:

  1. MLIR applied to LLM inference — multi-level IR lets the graph, kernel, and machine code be optimized in layers;
  2. Native dynamic-shape support — championed by TVM Relax;
  3. Cross-engine graph reuse — MLC-LLM compiles with TVM so the same graph runs on GPU/CPU/phone/WebGPU;
  4. Auto-scheduling — Ansor and MetaSchedule let compilers automatically search for optimal kernel parameters;
  5. CUDA graph conditional nodes (H100+) — even dynamic branches can be graph-captured.

9. Trade-offs ​

  • Static vs. dynamic: static graphs are fast, but LLM inference is dynamic — vLLM's bucketed CUDA graphs are the practical answer;
  • Compilation time vs. performance: online serving (vLLM) trades away compilation time for flexibility; offline (TensorRT-LLM) demands ultimate performance and accepts long compilation;
  • Graph optimization vs. kernel optimization: graph optimization makes macro decisions (what to fuse, which layout); kernel optimization does micro execution (Triton/CUDA);
  • Open-source vs. closed-source: the open stack (vLLM/Triton/TVM) has a strong ecosystem and is customizable; the closed stack (TensorRT) delivers ultimate performance but is a black box.

Further Reading ​

References ​