LLM Inference Optimization: техники ускорения генерации

Введение: LLM Inference Optimization

Проблема

LLM inference is expensive:
  - GPU memory: $10K-$100K per GPU
  - Compute: 1000s of GPU hours per inference
  - Latency: 1-10 seconds per response
  - Throughput: 10-100 tokens/sec per GPU

Solution: Multi-technique Optimization
  - Quantization: reduce precision
  - Caching: reuse computations
  - Parallelism: use multiple GPUs
  - Architecture: efficient models

Факт: Proper optimization can reduce inference cost by 10x and latency by 5x.

Quantization

FP16 vs BF16 vs INT8

Precision Levels:
  FP32: 32 bits (baseline)
  FP16: 16 bits (2x speedup)
  BF16: 16 bits (2x speedup, better stability)
  INT8: 8 bits  (4x speedup, some accuracy loss)
  INT4: 4 bits  (8x speedup, more accuracy loss)
  FP8:  8 bits  (new standard for LLMs)

# Quantization example with bitsandbytes
import torch
from transformers import AutoModelForCausalLM

# Load model in 4-bit quantized format
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-2-7b",
    load_in_4bit=True,           # 4-bit quantization
    load_in_8bit=False,
    bnb_4bit_compute_dtype=torch.float16,
    bnb_4bit_quant_type="nf4",   # Normalized FP4
    bnb_4bit_use_double_quant=True,
)

# Memory usage:
# FP16:  7B params × 2 bytes = 14 GB
# INT8:  7B params × 1 byte = 7 GB
# INT4:  7B params × 0.5 byte = 3.5 GB

NF4 (Normalized FP4)

# Normalized FP4: better than regular INT4
from bitsandbytes.nn import Linear4bit

# NF4 distributes quantization levels based on normal distribution
# Better for LLM weights (which are normally distributed)

# Comparison:
# INT4:  uniform quantization
# NF4:  normal-distribution-aware quantization
# Result: NF4 has ~1% better accuracy at same bitwidth

AWQ (Activation-aware Weight Quantization)

# AWQ: quantize weights based on activation importance
from awq import AutoAWQForCausalLM

model = AutoAWQForCausalLM.from_pretrained("meta-llama/Llama-2-7b")
quant_config = {
    "zero_point": True,
    "q_group_size": 128,
    "w_bit": 4,
    "version": "GEMM",
}

# Quantize
model.quantize(tokenizer, quant_config=quant_config)

# Save
model.save_quantized("llama-2-7b-awq")

# AWQ principle:
# - Some activations cause large outputs (important)
# - Protect weights that affect important activations
# - Quantize less important weights more aggressively

KV Cache Optimization

What is KV Cache?

KV Cache stores key and value tensors for each token
This avoids recomputing them for each new token

Memory usage:
  - Per layer: 2 × batch_size × num_heads × head_dim × seq_len × 2 bytes
  - For Llama-2-7B:
    - 32 layers × 2 × 1 × 32 heads × 128 dim × 4096 seq × 2 bytes
    - ≈ 21 GB for batch=1, seq=4096

Problem: KV cache dominates memory for long sequences

Paged Attention (vLLM)

# vLLM: PagedAttention for efficient KV cache management
import vllm

llm = vllm.LLM(model="meta-llama/Llama-2-7b")

prompts = [
    "Hello, my name is",
    "The capital of France is",
    "The largest planet in our solar system is",
]

outputs = llm.generate(prompts, sampling_params)

# PagedAttention:
# - Treats KV cache as memory pages
# - Allocates contiguous blocks for non-contiguous KV tensors
# - Achieves near-zero memory fragmentation
# - 2-4x better memory efficiency than standard caching

KV Cache Compression

# KV Cache eviction strategies:
# 1. Discard oldest tokens (sliding window)
# 2. Keep most important tokens (attention-based)
# 3. Compress multiple tokens into one (summary)

def compress_kv_cache(kv_cache, compression_ratio=0.5):
    """
    Compress KV cache by keeping top-k attention scores
    """
    # Calculate attention scores for each token
    attention_scores = compute_attention_scores(kv_cache)
    
    # Keep top-k tokens
    k = int(len(attention_scores) * compression_ratio)
    top_k_indices = torch.topk(attention_scores, k).indices
    
    # Compress remaining tokens
    compressed_kv = compress_remaining(kv_cache, top_k_indices)
    
    return compressed_kv

# Trade-off:
# - Higher compression → faster inference, lower quality
# - Typical ratio: 0.3-0.7

Speculative Decoding

Basic Idea

Speculative decoding uses a small model to draft tokens,
then verifies them with the large model.

Small model:  fast, generates N draft tokens
Large model:  slow, verifies N tokens in parallel

Speedup: ~N× for N draft tokens (theoretical)

# Speculative decoding example
def speculative_decode(draft_model, target_model, prompt, num_drafts=5):
    """
    Draft tokens with small model, verify with large model
    """
    # Draft tokens
    draft_tokens = draft_model.generate(prompt, max_new_tokens=num_drafts)
    
    # Verify with target model (parallel verification)
    target_output = target_model.generate(
        prompt + draft_tokens,
        verify_only=True,
    )
    
    # Accept or reject draft tokens
    accepted_tokens = []
    for draft, target in zip(draft_tokens, target_output):
        if draft == target:
            accepted_tokens.append(draft)
        else:
            accepted_tokens.append(target)
            break
    
    return accepted_tokens

# Acceptance rate:
# - 80%+ → 2-3x speedup
# - 60%+ → 1.5-2x speedup
# - <40% → no speedup

Medusa

# Medusa: multiple decoding heads for speculative decoding
from medusa import MedusaLM

model = MedusaLM.from_pretrained("meta-llama/Llama-2-7b")

# Add multiple decoding heads
# Each head predicts the next N tokens
# During inference, all heads vote on next token

# Result:
# - 2-3x speedup without fine-tuning
# - Simple to implement
# - Works with any base model

Tensor Parallelism

Splitting Across GPUs

Tensor parallelism splits model weights across GPUs

For Llama-2-7B with 4 GPUs:
  - Each GPU holds 1/4 of each layer's weights
  - All GPUs compute in parallel
  - Results are synchronized after each layer

Memory per GPU:
  - FP16: 14 GB / 4 = 3.5 GB
  - Plus KV cache: ~5 GB
  - Total: ~8.5 GB per GPU (A100-40GB is sufficient)

# Tensor parallelism with vLLM
import vllm

llm = vllm.LLM(
    model="meta-llama/Llama-2-7b",
    tensor_parallel_size=4,  # Use 4 GPUs
)

# Tensor parallelism steps:
# 1. Split weight matrices across GPUs
# 2. Each GPU computes its portion
# 3. All-Reduce to combine results
# 4. Proceed to next layer

Pipeline Parallelism

Pipeline parallelism splits layers across GPUs

For Llama-2-7B with 4 GPUs:
  - GPU 0: layers 0-7
  - GPU 1: layers 8-15
  - GPU 2: layers 16-23
  - GPU 3: layers 24-31

Bubble problem:
  - GPUs wait for previous stage to finish
  - Solution: micro-batching

Throughput:
  - Pipeline parallelism: higher throughput
  - Tensor parallelism: lower latency

Continuous Batching

What is Continuous Batching?

Traditional batching:
  - Wait for all requests to finish
  - Process as a batch
  - High latency for individual requests

Continuous batching (Iterative Batching):
  - Start processing requests immediately
  - New requests injected when GPU is free
  - Lower latency, higher throughput

vLLM implements continuous batching:
  - Throughput: 2-3x higher than HuggingFace
  - Latency: 20-50% lower

# vLLM continuous batching
import vllm

llm = vllm.LLM(model="meta-llama/Llama-2-7b")

# Requests are processed as soon as GPU is available
# No need to wait for batch to complete
outputs = llm.generate(
    prompts,
    sampling_params,
    use_tqdm=True,
)

# Continuous batching benefits:
# - Higher GPU utilization
# - Lower tail latency (p99)
# - Better throughput under variable-length requests

Kernel Optimization

Flash Attention

# Flash Attention: memory-efficient attention
import torch
from flash_attn import flash_attn_func

def efficient_attention(q, k, v):
    """
    Flash Attention reduces memory complexity
    from O(n²) to O(n)
    """
    # q, k, v: [batch, heads, seq_len, head_dim]
    output = flash_attn_func(q, k, v)
    return output

# Memory comparison:
# Standard attention: O(seq_len²) for attention matrix
# Flash Attention: O(seq_len) for intermediate results
# Result: 2-4x faster for long sequences

Triton Kernels

# Triton: custom CUDA kernels in Python
import triton
import triton.language as tl

@triton.jit
def quantize_kernel(weight, output, block_size):
    """
    Custom quantization kernel in Triton
    """
    pid = tl.program_id(0)
    start_idx = pid * block_size
    
    # Load weights
    weights = tl.load(weight + start_idx, block_size)
    
    # Quantize
    scales = tl.max(weights) / 7.0  # INT4 range
    quantized = (weights / scales * 7.0).to(tl.int4)
    
    # Store
    tl.store(output + start_idx, quantized)

# Triton benefits:
# - Write kernels in Python
# - Compile to efficient CUDA
# - No C++/CUDA expertise needed

Serving Frameworks

vLLM

# vLLM: high-throughput LLM serving
import vllm

llm = vllm.LLM(
    model="meta-llama/Llama-2-7b",
    tensor_parallel_size=4,
    gpu_memory_utilization=0.9,  # Use 90% of GPU memory
    max_num_batched_tokens=4096,
)

outputs = llm.generate(prompts, sampling_params)

# vLLM features:
# - PagedAttention
# - Continuous batching
# - Tensor parallelism
# - Speculative decoding
# - 24x higher throughput than HuggingFace

TensorRT-LLM

# TensorRT-LLM: NVIDIA's optimization library
import tensorrt_llm

# Build engine
builder = tensorrt_llm.Builder()
builder.model_name = "llama-2-7b"
builder.num_layers = 32
builder.num_heads = 32
builder.hidden_size = 4096

# Optimize
config = builder.builder_config
config.precision = "float16"
config.tensor_parallel = 4
config.trt_platform = "linux"

# Build and deploy
engine = builder.build_engine()
tensorrt_llm.runtime.infer(engine, inputs)

# TensorRT-LLM features:
# - Layer fusion
# - Quantization (FP8, INT8, INT4)
# - Custom kernels
# - Up to 3x faster than vLLM on NVIDIA GPUs

TGI (Text Generation Inference)

# TGI: HuggingFace's serving framework
from text_generation import Client

client = Client("http://localhost:8080")
response = client.generate(
    prompt="Hello, my name is",
    max_new_tokens=100,
    temperature=0.7,
)

# TGI features:
# - Continuous batching
# - Tensor parallelism
# - Quantization (AWQ, GPT-Q)
# - Docker deployment
# - Easy to use

Заключение

LLM inference optimization requires multiple techniques:

Quantization:

FP16 → INT8 → INT4
AWQ for better quality
Target: 4-8x memory reduction

KV Cache:

PagedAttention (vLLM)
Compression for long sequences
Target: 2-4x memory reduction

Parallelism:

Tensor parallelism: lower latency
Pipeline parallelism: higher throughput
Target: scale to larger models

Serving:

vLLM: best general-purpose
TensorRT-LLM: best NVIDIA performance
TGI: easiest to deploy

Best practices:

Start with vLLM + FP16
Add quantization if memory constrained
Use tensor parallelism for large models
Monitor GPU utilization