本地LLM部署方案：vLLM与Ollama进阶

# PagedAttention key concepts
# 1. Block: fixed-size KV-Cache unit (e.g., 16 tokens)
# 2. Block Table: maps logical blocks to physical GPU memory blocks
# 3. Copy-on-Write: for parallel sampling, share blocks until modification

# Memory savings example:
# Model: Llama-2-7B, max_seq_len=2048, batch_size=32
# Traditional: 32 * 2048 * 2 * 32 * 128 * 2 bytes = ~32 GB
# PagedAttention: only allocate for actual tokens, ~5-8 GB typical

# vLLM Python API for batch inference
from vllm import LLM, SamplingParams

llm = LLM(
    model="Qwen/Qwen2.5-7B-Instruct",
    tensor_parallel_size=2,
    gpu_memory_utilization=0.9,
    max_model_len=8192,
)

sampling_params = SamplingParams(
    temperature=0.7,
    top_p=0.9,
    max_tokens=1024,
    repetition_penalty=1.1,
    stop=["<|endoftext|>"],
)

# Batch inference
prompts = [
    "Explain quantum computing in simple terms.",
    "Write a Python function to sort a list.",
    "What are the benefits of microservices?",
]

outputs = llm.generate(prompts, sampling_params)

for output in outputs:
    print(f"Prompt: {output.prompt[:50]}...")
    print(f"Generated: {output.outputs[0].text}")
    print(f"Tokens/s: {len(output.outputs[0].token_ids) / output.metrics.finished_time:.1f}")
    print("---")

# Client code — fully compatible with OpenAI SDK
from openai import OpenAI

client = OpenAI(
    base_url="http://localhost:8000/v1",
    api_key="dummy",  # vLLM doesn't require a real key
)

# Chat completion
response = client.chat.completions.create(
    model="Qwen/Qwen2.5-7B-Instruct",
    messages=[
        {"role": "system", "content": "你是一个专业的技术顾问。"},
        {"role": "user", "content": "如何设计一个高可用的微服务架构？"},
    ],
    temperature=0.7,
    max_tokens=1024,
    stream=True,
)

for chunk in response:
    if chunk.choices[0].delta.content:
        print(chunk.choices[0].delta.content, end="", flush=True)

import httpx

OLLAMA_BASE = "http://localhost:11434"

# Generate completion
response = httpx.post(
    f"{OLLAMA_BASE}/api/generate",
    json={
        "model": "qwen2.5:7b",
        "prompt": "Explain Docker in 3 sentences.",
        "stream": False,
        "options": {
            "temperature": 0.7,
            "num_ctx": 4096,
        },
    },
    timeout=60.0,
)

print(response.json()["response"])

# Chat API
response = httpx.post(
    f"{OLLAMA_BASE}/api/chat",
    json={
        "model": "qwen2.5:7b",
        "messages": [
            {"role": "system", "content": "You are a helpful assistant."},
            {"role": "user", "content": "What is Kubernetes?"},
        ],
        "stream": False,
    },
    timeout=60.0,
)

print(response.json()["message"]["content"])

# Embeddings
response = httpx.post(
    f"{OLLAMA_BASE}/api/embed",
    json={
        "model": "nomic-embed-text",
        "input": ["Hello world", "Machine learning is great"],
    },
)

embeddings = response.json()["embeddings"]
print(f"Embedding dimension: {len(embeddings[0])}")

from awq import AutoAWQForCausalLM
from transformers import AutoTokenizer

model_path = "Qwen/Qwen2.5-7B-Instruct"
quant_path = "qwen2.5-7b-awq"

# Load model
model = AutoAWQForCausalLM.from_pretrained(model_path)
tokenizer = AutoTokenizer.from_pretrained(model_path)

# Quantize
quant_config = {
    "zero_point": True,
    "q_group_size": 128,
    "w_bit": 4,
    "version": "GEMM",  # or "GEMV" for batch_size=1
}

model.quantize(
    tokenizer,
    quant_config=quant_config,
    calib_data="wikitext",  # Calibration dataset
)

# Save quantized model
model.save_quantized(quant_path)
tokenizer.save_pretrained(quant_path)

# Deploy with vLLM
# vllm serve ./qwen2.5-7b-awq --quantization awq

def estimate_gpu_memory(
    num_params_billion: float,
    precision_bits: int = 16,
    batch_size: int = 1,
    seq_len: int = 2048,
    num_layers: int = 32,
    hidden_size: int = 4096,
    num_kv_heads: int = 8,
    head_dim: int = 128,
):
    """Estimate GPU memory requirements for LLM inference."""

    # Model weights
    weight_bytes = num_params_billion * 1e9 * (precision_bits / 8)

    # KV Cache (per token per layer: 2 * num_kv_heads * head_dim * precision_bytes)
    kv_per_token = 2 * num_kv_heads * head_dim * (precision_bits / 8)
    kv_cache_bytes = batch_size * seq_len * num_layers * kv_per_token

    # Activation memory (rough estimate)
    activation_bytes = batch_size * seq_len * hidden_size * 4  # FP32 activations

    total_gb = (weight_bytes + kv_cache_bytes + activation_bytes) / (1024**3)

    print(f"Model weights: {weight_bytes / (1024**3):.1f} GB")
    print(f"KV Cache: {kv_cache_bytes / (1024**3):.1f} GB")
    print(f"Activations: {activation_bytes / (1024**3):.1f} GB")
    print(f"Total estimate: {total_gb:.1f} GB")

    return total_gb

# Example: Qwen2.5-7B with INT4 quantization
estimate_gpu_memory(
    num_params_billion=7,
    precision_bits=4,
    batch_size=8,
    seq_len=4096,
)
# Model weights: 3.3 GB
# KV Cache: 2.0 GB
# Activations: 0.5 GB
# Total estimate: 5.8 GB

import time
import httpx

def benchmark_throughput(
    base_url: str,
    model: str,
    num_requests: int = 100,
    prompt_tokens: int = 128,
    max_tokens: int = 256,
):
    """Benchmark LLM serving throughput."""
    prompt = "Explain the concept of " * (prompt_tokens // 5)

    start = time.perf_counter()
    total_output_tokens = 0

    for i in range(num_requests):
        response = httpx.post(
            f"{base_url}/v1/completions",
            json={
                "model": model,
                "prompt": prompt,
                "max_tokens": max_tokens,
                "temperature": 0.7,
            },
            timeout=120.0,
        )
        data = response.json()
        total_output_tokens += data["usage"]["completion_tokens"]

    elapsed = time.perf_counter() - start

    print(f"Total time: {elapsed:.1f}s")
    print(f"Requests/s: {num_requests / elapsed:.1f}")
    print(f"Output tokens/s: {total_output_tokens / elapsed:.1f}")
    print(f"Avg latency: {elapsed / num_requests * 1000:.0f}ms")

# Benchmark
benchmark_throughput(
    base_url="http://localhost:8000",
    model="Qwen/Qwen2.5-7B-Instruct",
    num_requests=50,
)