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-100x

PagedAttention: 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

TechnologySpeedupImplementation DifficultyUse Cases
KV Cache10-100xLowAll scenarios
FlashAttention2-4xLowLong sequences
INT8 Quantization2-3xMediumGPU inference
Speculative Decoding2-3xHighLow-latency scenarios
Operator Fusion1.2-1.5xMediumAll scenarios
TensorRT-LLM3-5xHighNVIDIA GPU

Practical Recommendations

  1. Measure before optimizing: Use a Profiler to find the real bottleneck
  2. Start simple: KV Cache and FlashAttention are the most cost-effective optimizations
  3. Prioritize quantization: INT8 quantization has almost no precision loss and significant speedup
  4. Combine techniques: Combining multiple techniques yields better results
  5. 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).