03. 5 个 Mini 实验：把直觉变成数字

# experiment1_prefix_cache.py
import time
from vllm import LLM, SamplingParams

SYSTEM = "你是一个助手。" * 200      # 大约 1k tokens
USER_QUERIES = [f"用户问题 {i}" for i in range(10)]

def run(enable_prefix_caching: bool):
    llm = LLM(
        model="Qwen/Qwen2.5-0.5B-Instruct",
        enforce_eager=True,
        enable_prefix_caching=enable_prefix_caching,
        gpu_memory_utilization=0.5,
    )
    params = SamplingParams(max_tokens=50, temperature=0)
    prompts = [SYSTEM + q for q in USER_QUERIES]

    # 第一次跑：填 cache
    llm.generate(prompts[:1], params)

    # 计时：后续 9 次
    t0 = time.perf_counter()
    llm.generate(prompts[1:], params)
    dt = time.perf_counter() - t0
    print(f"prefix_caching={enable_prefix_caching}: 9 requests in {dt*1000:.0f}ms")

run(enable_prefix_caching=False)
run(enable_prefix_caching=True)

# 把 max-num-batched-tokens 设小（不易混入长 prefill）
vllm serve Qwen/Qwen2.5-0.5B-Instruct \
    --enforce-eager \
    --gpu-memory-utilization 0.5 \
    --max-num-batched-tokens 1024 \
    --port 8001 &

# 设大（一个长 prefill 可能占满整步）
vllm serve Qwen/Qwen2.5-0.5B-Instruct \
    --enforce-eager \
    --gpu-memory-utilization 0.5 \
    --max-num-batched-tokens 16384 \
    --port 8002 &

# experiment2_tpot.py：同时发短请求 + 一个长请求
import asyncio
import time
import httpx

async def short_req(client, port):
    t0 = time.perf_counter()
    await client.post(f"http://localhost:{port}/v1/completions", json={
        "model": "Qwen/Qwen2.5-0.5B-Instruct",
        "prompt": "Hi, how are you?",
        "max_tokens": 30,
        "temperature": 0,
    })
    return time.perf_counter() - t0

async def long_req(client, port):
    await client.post(f"http://localhost:{port}/v1/completions", json={
        "model": "Qwen/Qwen2.5-0.5B-Instruct",
        "prompt": "重复内容 " * 4000,    # ~8k tokens
        "max_tokens": 10,
        "temperature": 0,
    })

async def run(port):
    async with httpx.AsyncClient(timeout=60) as client:
        # 同时发 5 个短请求 + 1 个长请求
        tasks = [short_req(client, port) for _ in range(5)] + [long_req(client, port)]
        results = await asyncio.gather(*tasks)
    print(f"port {port}: short req latencies = {[f'{r*1000:.0f}ms' for r in results[:5]]}")

asyncio.run(run(8001))   # 小 budget
asyncio.run(run(8002))   # 大 budget

# baseline
vllm serve Qwen/Qwen2.5-0.5B-Instruct \
    --enforce-eager \
    --gpu-memory-utilization 0.5 2>&1 | grep -E "GPU blocks|KV cache" &

sleep 60 && pkill -f "vllm serve" && sleep 5

# fp8 KV
vllm serve Qwen/Qwen2.5-0.5B-Instruct \
    --enforce-eager \
    --gpu-memory-utilization 0.5 \
    --kv-cache-dtype fp8 2>&1 | grep -E "GPU blocks|KV cache"

# experiment4_preempt.py
import asyncio
import httpx

# 启动 server 时强制小 KV：
# vllm serve Qwen/Qwen2.5-0.5B-Instruct \
#     --enforce-eager --gpu-memory-utilization 0.2 \
#     --max-num-seqs 256

async def long_req(client, i):
    return await client.post("http://localhost:8000/v1/completions", json={
        "model": "Qwen/Qwen2.5-0.5B-Instruct",
        "prompt": f"请求 {i} " + "上下文 " * 500,
        "max_tokens": 500,
        "temperature": 0.7,
    })

async def run():
    async with httpx.AsyncClient(timeout=120) as client:
        # 同时发 200 个长请求
        tasks = [long_req(client, i) for i in range(200)]
        await asyncio.gather(*tasks)

asyncio.run(run())

watch -n 0.5 'curl -s localhost:8000/metrics | grep -E "vllm:num_preemptions|vllm:num_running|vllm:num_waiting"'

# baseline
python benchmarks/benchmark_throughput.py \
    --model Qwen/Qwen2.5-0.5B-Instruct \
    --num-prompts 100 \
    --input-len 256 \
    --output-len 256 \
    --enforce-eager

# 用 ngram spec
python benchmarks/benchmark_throughput.py \
    --model Qwen/Qwen2.5-0.5B-Instruct \
    --num-prompts 100 \
    --input-len 256 \
    --output-len 256 \
    --enforce-eager \
    --speculative-config '{"method": "ngram", "num_speculative_tokens": 3, "prompt_lookup_max": 4}'

实验 X：[一句话目标]
观察：[最关键的 1-2 个数字]
结论：[这告诉我们关于 vLLM 的什么]
踩坑：[过程中遇到的意外，怎么解决的]

# 1. 性能 benchmark（4 组）
for fmt in fp16 fp8 awq gptq; do
    case $fmt in
        fp16) MODEL=meta-llama/Llama-2-7b-hf;        QUANT="" ;;
        fp8)  MODEL=nm-testing/llama-2-7b-fp8;       QUANT="--quantization fp8" ;;
        awq)  MODEL=TheBloke/Llama-2-7B-AWQ;         QUANT="--quantization awq_marlin" ;;
        gptq) MODEL=TheBloke/Llama-2-7B-GPTQ;        QUANT="--quantization gptq_marlin" ;;
    esac

    echo "=== $fmt ==="
    vllm serve $MODEL $QUANT \
        --gpu-memory-utilization 0.9 \
        --max-num-seqs 256 \
        --port 8000 &
    SERVER_PID=$!
    sleep 90   # 等 compile + warmup

    python benchmarks/benchmark_serving.py \
        --model $MODEL \
        --dataset-name sharegpt \
        --dataset-path ShareGPT_V3_unfiltered_cleaned_split.json \
        --num-prompts 500 \
        --request-rate 10 \
        --result-filename results_$fmt.json

    # 取关键 metric
    curl -s :8000/metrics | grep -E "vllm:(kv_cache_usage|num_requests_running)" \
        > runtime_$fmt.txt

    kill $SERVER_PID && sleep 30
done

# 2. 精度 benchmark
for fmt in fp16 fp8 awq gptq; do
    lm_eval --model vllm \
        --model_args pretrained=$MODEL,quantization=$QUANT \
        --tasks wikitext,arc_easy,hellaswag \
        --batch_size auto \
        --output_path eval_$fmt.json
done

NUM_GPUS=$(nvidia-smi -L | wc -l)
for tp in 1 2 4 8; do
    [ $tp -gt $NUM_GPUS ] && continue

    vllm serve meta-llama/Llama-2-13b-hf \
        --tensor-parallel-size $tp \
        --max-num-seqs 256 \
        --port 8000 &
    SERVER_PID=$!
    sleep 120   # 等 NCCL init + compile

    # 多 QPS 曲线（找到拐点）
    for qps in 5 10 20 50 100; do
        python benchmarks/benchmark_serving.py \
            --model meta-llama/Llama-2-13b-hf \
            --dataset-name sharegpt \
            --num-prompts 300 \
            --request-rate $qps \
            --result-filename results_tp${tp}_qps${qps}.json
    done

    # nsys 看 NCCL 占比（取 QPS=20 这个中等负载）
    nsys profile -t cuda,nvtx,osrt \
        --capture-range cudaProfilerApi --capture-range-end stop \
        -o profile_tp${tp} \
        python -c "
from vllm import LLM, SamplingParams
import torch
llm = LLM('meta-llama/Llama-2-13b-hf', tensor_parallel_size=$tp, enforce_eager=True)
torch.cuda.cudart().cudaProfilerStart()
llm.generate(['Hello'] * 32, SamplingParams(max_tokens=100))
torch.cuda.cudart().cudaProfilerStop()
"

    kill $SERVER_PID && sleep 30
done

# 用 nsys-ui 或 ncu 打开 profile_tp*.nsys-rep，看：
#   - NCCL AllReduce kernel 占 timeline 的 % 比例
#   - 不同 TP 下 forward 总时长

# 1. 启动服务
vllm serve Qwen/Qwen2.5-0.5B-Instruct \
    --max-num-seqs 256 \
    --max-num-batched-tokens 8192 \
    --port 8000 &
SERVER_PID=$!
sleep 60

# 2. 用 py-spy 采样 scheduler 进程的 Python 栈
SCHEDULER_PID=$(pgrep -f "EngineCore" | head -1)
sudo py-spy record \
    -o profile_async.svg \
    --pid $SCHEDULER_PID \
    --duration 30 \
    --rate 1000 \
    --threads &
PYSPY_PID=$!

# 3. 同时跑大 batch workload 制造 scheduler 压力
python benchmarks/benchmark_serving.py \
    --model Qwen/Qwen2.5-0.5B-Instruct \
    --dataset-name sharegpt \
    --num-prompts 500 \
    --request-rate 50

wait $PYSPY_PID

# 4. 看 profile_async.svg：
#    - 找 schedule() 调用栈，看它在 30s 里占多少
#    - 看 update_from_output / preempt 等子函数分布
#    - 看 scheduler thread 状态：busy% 多少

# experiment_async.py
import torch
from torch.profiler import profile, ProfilerActivity
from vllm import LLM, SamplingParams

llm = LLM(
    model="Qwen/Qwen2.5-0.5B-Instruct",
    enforce_eager=True,   # 关掉 CUDA Graph，让 CPU 时间清晰可见
    max_num_seqs=256,
    gpu_memory_utilization=0.9,
)

# 大 batch 制造调度压力
prompts = ["写一段 200 字介绍 vLLM 的文章。"] * 200
params = SamplingParams(max_tokens=200, temperature=0)

with profile(
    activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
    record_shapes=True,
    with_stack=True,
) as prof:
    outs = llm.generate(prompts, params)

print(prof.key_averages().table(sort_by="cpu_time_total", row_limit=30))
prof.export_chrome_trace("trace_async.json")
# 用 chrome://tracing 或 https://ui.perfetto.dev 打开 trace_async.json

线程 EngineCore-main  ▓▓▓▓░░▓▓▓▓░░▓▓▓▓░░▓▓▓▓     ← schedule() 间歇执行
线程 Worker-0 GPU     ░░▓▓▓▓░░▓▓▓▓░░▓▓▓▓░░▓▓▓▓   ← forward 间歇执行
                       ^ schedule 与 forward 错开，几乎完全重叠

# Instance A：prefill 节点（专配大 batch token，小 seq 数）
CUDA_VISIBLE_DEVICES=0 vllm serve meta-llama/Llama-2-7b-hf \
    --port 8000 \
    --max-num-seqs 8 \
    --max-num-batched-tokens 16384 \
    --kv-transfer-config '{
        "kv_connector": "PyNixlConnector",
        "kv_role": "kv_producer",
        "kv_rank": 0,
        "kv_parallel_size": 2
    }' &

# Instance B：decode 节点（专配多并发，小 batch token）
CUDA_VISIBLE_DEVICES=1 vllm serve meta-llama/Llama-2-7b-hf \
    --port 8001 \
    --max-num-seqs 128 \
    --max-num-batched-tokens 4096 \
    --kv-transfer-config '{
        "kv_connector": "PyNixlConnector",
        "kv_role": "kv_consumer",
        "kv_rank": 1,
        "kv_parallel_size": 2
    }' &

# 简化的"前端路由"参考：vllm/examples/online_serving/disaggregated_prefill.sh
# 它把请求第一次发到 prefill 节点，拿到 KV 后路由到 decode 节点继续

# Baseline：TP=2 单实例（同时做 prefill+decode）
CUDA_VISIBLE_DEVICES=0,1 vllm serve meta-llama/Llama-2-7b-hf \
    --tensor-parallel-size 2 \
    --port 8002 &

# 跑 2 组 benchmark
for port in 8000 8002; do
    python benchmarks/benchmark_serving.py \
        --model meta-llama/Llama-2-7b-hf \
        --base-url http://localhost:$port \
        --dataset-name sharegpt \
        --num-prompts 300 \
        --request-rate 30 \
        --result-filename results_port${port}.json
done

total_hbm = 80 GB (H100)
   minus 模型权重         (~16 GB Llama-3-8B BF16)
   minus CUDA buffer       (~2 GB)
   minus activation 预算   (~4-6 GB, 与 max_num_batched_tokens 相关)
   minus CUDA Graph buffer (~1-2 GB)
   minus profiling 留 5%   (~4 GB)
= KV 可用                  (~50 GB)

→ KV blocks (BF16) = 50 / single_block_bf16
→ KV blocks (FP8)  = 50 / single_block_fp8 = 50 / (single_block_bf16 / 2) = 100 / single_block_bf16

ratio = 100 / 50 = 2.0  ← 理论值