# Inference Optimization

vLLM, SGLang, TensorRT-LLM, speculative decoding, quantization, and serving at scale.

---

## Inference Engines Comparison

| Engine | Backend | Strengths | Best For |
|---|---|---|---|
| **vLLM** | CUDA, PagedAttention | High throughput, easy API, OpenAI-compatible | Production LLM serving |
| **SGLang** | CUDA, RadixAttention | Structured output, complex control flow, fast | Agent workflows, constrained decoding |
| **TensorRT-LLM** | TensorRT, FP8, plugins | Maximum NVIDIA GPU perf, FP8 | Max throughput on H100/H200 |
| **llama.cpp** | GGUF, CPU/Apple Silicon | Edge, CPU, quantized | Local, edge, mobile |
| **TGI (Text Generation Inference)** | PyTorch, FlashAttention | OpenAI-compatible, streaming, watermarking | HuggingFace ecosystem |
| **MNN/FastLLM** | Mobile/embedded | Lightweight, no deps | Mobile, embedded |

---

## vLLM Deep Dive

### PagedAttention (Key Innovation)
```
Traditional: Contiguous KV cache → fragmentation, OOM
PagedAttention: KV cache in fixed-size blocks (pages)
              Logical blocks → Physical blocks (like OS virtual memory)
```

### Core Features
```python
from vllm import LLM, SamplingParams

llm = LLM(
    model="meta-llama/Meta-Llama-3.1-8B-Instruct",
    tensor_parallel_size=2,        # Multi-GPU
    gpu_memory_utilization=0.9,    # KV cache budget
    max_model_len=8192,            # Context window
    dtype="bfloat16",              # BF16 for H100
    quantization="awq",            # AWQ quantization
    enforce_eager=False,           # CUDA graphs
    max_num_seqs=256,              # Max concurrent
    max_num_batched_tokens=8192,   # Batch size
    enable_prefix_caching=True,    # Share prefix KV
    disable_log_stats=False
)

# OpenAI-compatible API
from vllm.entrypoints.openai import api_server
# Run: python -m vllm.entrypoints.openai.api_server --model ...
```

### Sampling Parameters
```python
params = SamplingParams(
    temperature=0.7,
    top_p=0.9,
    top_k=50,
    max_tokens=2048,
    stop=["</s>", "<|eot_id|>"],
    presence_penalty=0.0,
    frequency_penalty=0.0,
    repetition_penalty=1.0,
    # Structured output
    guided_decoding=GuidedDecodingParams(
        json=MySchema.model_json_schema()
    ),
    # Beam search
    use_beam_search=False,
    best_of=1,
    # Logprobs
    logprobs=5,
    prompt_logprobs=0
)
```

### Advanced Features
```python
# Prefix Caching (share KV for common prefixes)
llm = LLM(..., enable_prefix_caching=True)

# Chunked Prefill (long context)
llm = LLM(..., enable_chunked_prefill=True, max_num_batched_tokens=8192)

# LoRA Support
llm = LLM(..., enable_lora=True, max_lora_rank=64, max_loras=16)
# Runtime: llm.generate(prompt, lora_request=LoRARequest("adapter", 1, "path"))

# Speculative Decoding
llm = LLM(..., speculative_model="draft-model", num_speculative_tokens=5)

# Multi-LoRA
outputs = llm.generate(
    prompts,
    lora_request=LoRARequest("sql_adapter", 1, "./sql_lora")
)
```

---

## SGLang (Structured Generation Language)

### RadixAttention (Prefix Sharing)
```
Radix Tree of KV Cache:
                    Root
                   / | \
                  A  B  C
                 /|   |
                D E   F
Shared prefixes = shared KV cache = memory + speed
```

### Structured Output
```python
import sglang as sgl

@sgl.function
def structured_agent(s, query):
    s += "Analyze: " + query
    
    # Structured output with Pydantic
    class Analysis(sgl.StructuredOutput):
        issue_type: Literal["bug", "feature", "question"]
        severity: Literal["low", "medium", "high", "critical"]
        component: str
        suggested_fix: str
    
    s += sgl.gen("analysis", Analysis, temperature=0.3)
    return s

# Run
runtime = sgl.Runtime(model_path="meta-llama/Meta-Llama-3.1-8B-Instruct")
result = structured_agent.run("User reports login failure")
print(result["analysis"])  # Parsed Pydantic object
```

### Complex Control Flow
```python
@sgl.function
def multi_step_reasoning(s, problem):
    # Step 1: Plan
    s += "Plan: " + sgl.gen("plan", max_tokens=512)
    
    # Step 2: Execute each step
    for i in range(5):
        s += f"Step {i+1}: " + sgl.gen(f"step_{i}", max_tokens=256, stop="\n\n")
        
        # Conditional continuation
        if "DONE" in s[f"step_{i}"]:
            break
    
    # Step 3: Synthesize
    s += "Final Answer: " + sgl.gen("answer", max_tokens=512)
    return s
```

### Parallel & Batch
```python
# Automatic batching
@sgl.function
def batch_agent(s, queries: list[str]):
    fork = s.fork(len(queries))
    for i, q in enumerate(queries):
        fork[i] += "Answer: " + q
        fork[i] += sgl.gen("answer", max_tokens=256)
    return fork.join("answers")

# Automatic continuous batching
runtime = sgl.Runtime(..., chunked_prefill_size=2048)
```

---

## TensorRT-LLM

### FP8 Quantization (H100/H200)
```bash
# Convert to FP8
python examples/quantization/quantize.py \
    --model_dir /path/to/model \
    --output_dir /path/to/fp8_model \
    --dtype fp8 \
    --calib_size 512
```

### Plugin System
```cpp
// Custom kernels for:
// - FlashAttention
// - Fused MLP (SwiGLU)
// - RMSNorm + Residual
// - RoPE
// - Speculative decoding
// - Medusa heads
```

### Inflight Batching
```python
# Continuous batching with iteration-level scheduling
# Max throughput at high concurrency
```

### Deployment
```bash
# Triton Inference Server
tritonserver --model-repository=/models \
    --model-control-mode=explicit \
    --load-model=llama3_8b_fp8
```

---

## Speculative Decoding

### Draft + Target
```python
# Small draft model → Large target model
# Draft generates k tokens → Target verifies in parallel

async def speculative_generate(draft_model, target_model, prompt, k=5):
    tokens = prompt
    while len(tokens) < max_tokens:
        # Draft generates k tokens
        draft_tokens = await draft_model.generate(tokens, max_tokens=k)
        
        # Target scores all k+1 positions in parallel
        target_logits = await target_model.score(tokens + draft_tokens)
        
        # Verify each token
        accepted = []
        for i, (draft_logit, target_logit) in enumerate(zip(draft_logits, target_logits)):
            if accept_token(draft_logit, target_logit):
                accepted.append(draft_tokens[i])
            else:
                break
        
        tokens.extend(accepted)
        
        if len(accepted) < k:
            # Resample from target at rejection point
            next_token = sample(target_logits[len(accepted)])
            tokens.append(next_token)
    
    return tokens
```

### Draft Model Options
| Draft Model | Speedup | Quality Loss |
|---|---|---|
| **Smaller same architecture** | 2-3x | Minimal |
| **Medusa heads** | 2-3x | Minimal |
| **EAGLE (draft from target)** | 2-4x | Minimal |
| **NGram/Simple** | 1.5-2x | Higher |
| **Self-speculative** | 1.5x | None |

### Medusa / EAGLE
```python
# Medusa: Additional heads on target model for drafting
class MedusaModel(nn.Module):
    def __init__(self, base_model, num_heads=5):
        self.base = base_model
        self.medusa_heads = nn.ModuleList([
            nn.Linear(d_model, vocab_size) for _ in range(num_heads)
        ])
    
    def forward(self, x):
        base_logits = self.base(x).logits
        medusa_logits = [h(x) for h in self.medusa_heads]
        return base_logits, medusa_logits

# EAGLE: Draft from target's own hidden states
# More accurate drafting, less quality loss
```

---

## Quantization for Inference

### Post-Training Quantization (PTQ)
```python
# GPTQ (Layer-wise)
from auto_gptq import AutoGPTQForCausalLM

model = AutoGPTQForCausalLM.from_quantized(
    "TheBloke/Llama-3-8B-GPTQ",
    device="cuda:0",
    use_triton=False
)

# AWQ (Activation-aware)
from awq import AutoAWQForCausalLM

model = AutoAWQForCausalLM.from_quantized(
    "TheBloke/Llama-3-8B-AWQ",
    device="cuda:0"
)
```

### Quantization Configs
```python
# vLLM quantization
llm = LLM(
    model="TheBloke/Llama-3-8B-AWQ",
    quantization="awq",
    dtype="half"
)

# TensorRT-LLM FP8
# Best on H100/H200

# llama.cpp GGUF (CPU/Apple Silicon)
# q4_k_m: Best quality/size
# q8_0: Near FP16 quality
```

---

## KV Cache Optimization

### Prefix Caching
```python
# vLLM: Automatic prefix caching
llm = LLM(..., enable_prefix_caching=True)

# Manual: Share KV for system prompt + few-shot
system_prompt = "You are a helpful assistant..."
few_shot = "Example 1...\nExample 2..."

# Both requests share KV for prefix
request1 = system_prompt + few_shot + "Query 1"
request2 = system_prompt + few_shot + "Query 2"
```

### KV Cache Compression
```python
# Quantization
# INT8 KV cache: 2x memory savings
llm = LLM(..., kv_cache_dtype="int8")

# Sliding Window
llm = LLM(..., sliding_window=4096)  # Only keep last 4k tokens

# Offloading (CPU offload)
llm = LLM(..., cpu_offload_gb=10)  # Offload 10GB to CPU RAM
```

### Attention Optimizations
```python
# FlashAttention-2 (H100)
# FlashAttention-3 (H100/H200, FP8)

# PagedAttention (vLLM) - already covered
# FlashDecoding (long context)
# RingAttention (multi-GPU context parallel)
```

---

## Continuous Batching / Inflight Batching

```python
# vLLM: Automatic
# SGLang: Automatic
# TensorRT-LLM: Inflight batching

# Key: Iteration-level scheduling
# New requests join queue, completed exit
# No padding waste, max GPU utilization

# Config
llm = LLM(
    max_num_seqs=256,           # Max concurrent
    max_num_batched_tokens=8192, # Token budget per iteration
    scheduling_policy="fcfs"     # fcfs, priority
)
```

---

## Serving Patterns

### OpenAI-Compatible API
```python
# vLLM, SGLang, TGI all support
from openai import AsyncOpenAI

client = AsyncOpenAI(base_url="http://localhost:8000/v1", api_key="EMPTY")

# Chat
response = await client.chat.completions.create(
    model="meta-llama/Meta-Llama-3.1-8B-Instruct",
    messages=[{"role": "user", "content": "Hello"}],
    stream=True
)

# Completions
response = await client.completions.create(
    model="meta-llama/Meta-Llama-3.1-8B-Instruct",
    prompt="Once upon a time",
    max_tokens=100
)
```

### Load Balancing
```python
# Multiple replicas + load balancer
# vLLM: Data parallel (tensor_parallel_size) + multiple replicas

# Kubernetes
# Deployment: replicas=4, each with tensor_parallel_size=2
# Service: LoadBalancer
# Metrics: vLLM exposes /metrics for Prometheus
```

---

## Monitoring & Metrics

### Key Metrics
| Metric | Target |
|---|---|
| **TTFT (Time to First Token)** | < 100ms (p99) |
| **TPOT (Time per Output Token)** | < 50ms |
| **Throughput** | Tokens/sec/GPU |
| **GPU Utilization** | > 80% |
| **KV Cache Hit Rate** | > 50% (with prefix caching) |
| **Queue Time** | < 1s (p99) |
| **Error Rate** | < 0.1% |

### Prometheus Metrics (vLLM)
```python
# Exposed at /metrics
vllm:num_requests_running
vllm:num_requests_waiting
vllm:gpu_cache_usage_perc
vllm:iteration_tokens_total
vllm:request_latency_seconds
```

---

## Related Notes

- [[ModelArchitectures/Mixture-of-Experts.md]] — MoE inference
- [[ModelArchitectures/Small-Language-Models.md]] — Quantized SLMs
- [[ModelArchitectures/Reasoning-Models.md]] — Reasoning inference
- [[ModelArchitectures/Mixture-of-Experts.md]] — MoE serving
- [[Production/Cost-Optimization.md]] — Inference cost
- [[Production/Scaling-Strategies.md]] — Scaling inference
- [[Production/Deployment-Architecture.md]] — Serving architecture
- [[Production/Reliability.md]] — Circuit breakers, health
- [[Production/Observability.md]] — Metrics, tracing
- [[Frontiers/Research-Directions.md]] — Future inference

---

*Sources: vLLM docs (vllm.ai), SGLang docs (sglang.ai), TensorRT-LLM docs, "PagedAttention" (Kwon et al.), "FlashAttention" (Dao et al.), "GPTQ" (Frantar et al.), "AWQ" (Lin et al.), "Medusa" (Cai et al.), "EAGLE" (Li et al.), "Speculative Decoding" (Leviathan et al.), TGI docs, llama.cpp, TensorRT-LLM benchmarks, NVIDIA H100 FP8 guides.*