Inference Optimization

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

Inference Engines Comparison

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

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

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

# 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

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

@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

# 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)

# 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

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

Inflight Batching

# Continuous batching with iteration-level scheduling
# Max throughput at high concurrency

Deployment

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

Speculative Decoding

Draft + Target

# 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

Medusa / EAGLE

# 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)

# 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

# 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

# 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

# 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

# FlashAttention-2 (H100)
# FlashAttention-3 (H100/H200, FP8)

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

Continuous Batching / Inflight Batching

# 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

# 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

# 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

Prometheus Metrics (vLLM)

# 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

• Mixture of Experts — MoE inference
• Small Language Models — Quantized SLMs
• Reasoning Models — Reasoning inference
• Mixture of Experts — MoE serving
• Cost Optimization — Inference cost
• Scaling Strategies — Scaling inference
• Deployment Architecture — Serving architecture
• Reliability — Circuit breakers, health
• Observability — Metrics, tracing
• Research Directions — 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.