Importance of Inference Acceleration
Inference latency of large models directly impacts user experience. Studies show that for every 100ms increase in response time, user satisfaction drops by 3-5%. For real-time conversational applications, a time-to-first-token (TTFT) within 200ms is excellent, within 500ms is acceptable, and over 1s users will noticeably feel the wait.
KV Cache: The Most Basic Acceleration Technique
KV Cache is the core optimization in Transformer inference. Without KV Cache, each new token generation requires recomputing the Key and Value for all previous tokens. KV Cache caches the computed K and V, so each step only needs to compute for the new token:
# KV Cache working principle
class KVCacheAttention:
def __init__(self):
self.k_cache = None
self.v_cache = None
def forward(self, q, k, v, use_cache=True):
if use_cache and self.k_cache is not None:
# Append new K, V to cache
k = torch.cat([self.k_cache, k], dim=-2)
v = torch.cat([self.v_cache, v], dim=-2)
# Update cache
self.k_cache = k
self.v_cache = v
# Compute attention
attn = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(d_k)
attn = F.softmax(attn, dim=-1)
return torch.matmul(attn, v)
# Effect: Reduces O(n²) to O(n), inference speedup 10-100xPagedAttention: Efficient KV Cache Management
PagedAttention proposed by vLLM manages KV Cache in pages, solving the memory fragmentation problem:
# PagedAttention memory management
class PagedKVCache:
def __init__(self, block_size=16, num_blocks=1024):
self.block_size = block_size
self.free_blocks = list(range(num_blocks))
self.block_table = {} # seq_id -> [block_ids]
def allocate(self, seq_id, num_tokens):
num_blocks_needed = (num_tokens + self.block_size - 1) // self.block_size
blocks = self.free_blocks[:num_blocks_needed]
self.free_blocks = self.free_blocks[num_blocks_needed:]
self.block_table[seq_id] = blocks
return blocks
# Effect: Memory utilization increases from 20-40% to nearly 100%Speculative Decoding
Using a small model to "guess" multiple tokens, then verifying them all at once with the large model, achieving 2-3x decoding speedup:
class SpeculativeDecoder:
def __init__(self, target_model, draft_model):
self.target_model = target_model # large model
self.draft_model = draft_model # small model
def generate(self, prompt, max_tokens=100):
tokens = tokenize(prompt)
while len(tokens) < max_tokens:
# 1. Small model generates K candidate tokens
draft_tokens = self.draft_model.generate(tokens, k=5)
# 2. Large model verifies all candidates at once
logits = self.target_model.forward(tokens + draft_tokens)
# 3. Accept matching tokens, reject non-matching
accepted = self._verify_and_accept(logits, draft_tokens)
tokens.extend(accepted)
if len(accepted) < len(draft_tokens):
# Resample from the non-matching position
correction = self.target_model.sample(logits[len(accepted)])
tokens.append(correction)
return detokenize(tokens)Operator Fusion
Fusing multiple small operators into one large operator to reduce memory reads/writes:
# Before fusion: 5 memory reads/writes
# LayerNorm → Linear → Dropout → ReLU → Linear
# After fusion: 1 memory read/write
# FusedMLP (LayerNorm + Linear + GELU + Linear)
# Use torch.compile for automatic fusion
import torch
@torch.compile
def fused_mlp(x, w1, w2, b1, b2):
return torch.nn.functional.linear(
torch.nn.functional.gelu(
torch.nn.functional.linear(x, w1, b1)
),
w2, b2
)
# Effect: Latency reduced by 20-30%, memory bandwidth saved by 40%Technology Comparison
| Technology | Speedup | Implementation Difficulty | Use Cases |
|---|---|---|---|
| KV Cache | 10-100x | Low | All scenarios |
| FlashAttention | 2-4x | Low | Long sequences |
| INT8 Quantization | 2-3x | Medium | GPU inference |
| Speculative Decoding | 2-3x | High | Low-latency scenarios |
| Operator Fusion | 1.2-1.5x | Medium | All scenarios |
| TensorRT-LLM | 3-5x | High | NVIDIA GPU |
Practical Recommendations
- Measure before optimizing: Use a Profiler to find the real bottleneck
- Start simple: KV Cache and FlashAttention are the most cost-effective optimizations
- Prioritize quantization: INT8 quantization has almost no precision loss and significant speedup
- Combine techniques: Combining multiple techniques yields better results
- Continuous monitoring: Monitor latency and throughput after optimization
Summary
Inference acceleration is a systematic engineering effort. It is recommended to start with KV Cache and quantization, then gradually introduce more advanced techniques. The core principles are: reduce computation (quantization, speculative decoding), reduce memory access (FlashAttention, operator fusion), and increase parallelism (continuous batching).