Month 5 · 原创项目 · mini-vLLM

从零写一个能 serve 一个请求的 paged-attention engine。这是简历最重的那张牌。

到目前为止，你都在读别人的代码。这个月反过来： 把 M2-M4 学到的东西用自己的代码重新发明一遍。 发明的过程会暴露所有你以为懂、其实没懂的地方—— 这正是这一步的价值，也是这门课最贵的一张牌。

01项目意义 · 为什么"重新发明轮子"是顶投资

有人会说：vLLM 已经写完了，开源、文档全、社区活跃，你写一个慢 3 倍的 toy 有什么用？ 这个问题问得对——但答案不在最终产物，而在制造过程。

"读" 是识别，"写" 是召回

人脑对信息有两种掌握程度：识别 (recognition) 和 召回 (recall)。 读一段代码、点头说"嗯我懂这个 if 在干嘛"，是识别。 让你关掉编辑器、自己写一遍同样功能的代码，是召回。 这两个差一个数量级——你以为懂的 80%，召回时只剩 20%。

M1-M4 你都在做识别：读 vLLM 源码、看论文、画状态图。 M5 强制你切换到召回模式：合上 vLLM 的 tab，在空白的 engine.py 里从 0 写。 写出来你会发现一堆"咦，这里 vLLM 是怎么处理的来着？" 那些就是你的盲点。

你会发现的具体盲点

下面这些是历届做 mini-vLLM 的人真实写到的总结，几乎每个人都中招：

• KV cache 的 shape 你以为懂，写第一行 torch.zeros(...) 时才发现： layer × 2 (K/V) × seq × n_kv_heads × head_dim？还是 layer × seq × 2 × ...？vLLM 的 layout 跟 HF 不一样， paged 之后又不一样。这种细节读 10 遍都印不进脑子，写一遍就再也忘不了。
• logits 切片你以为是 [:, -1, :]，写完一跑全是乱码—— 才发现 HF 的 generate() 默认会丢掉一个 bos token，你手写时漏了； 或 chat template 没套，模型在做 continuation 而不是 instruct。
• "prefill 和 decode 用同一个 forward" 你点过头，自己写时才意识到 decode 时 position_ids = [seq_len - 1] 而不是 [0, ..., seq_len-1]， KV cache 只查不写 vs 写一格。
• block table 你画过图，自己写 BlockManager 时才发现： block 0 是所有序列共享的特殊块还是普通块？free 时要不要 zero-out？vLLM 用 ref count 还是用 owner？
• scheduler 的 token budget——你以为是 max_batch_size × max_seq_len， 自己写时才发现 vLLM 是每步固定 token 数，prefill 和 decode 共享这个预算。这是一个不写一遍体会不到的设计。

每个盲点都是 vLLM 源码里你滑过去没注意的一个细节。 你的笔记本会从这些"咦"积累出一份比任何博客都准的 vLLM 注解。

简历价值：可验证的 hard signal

在 infra/serving 的招聘市场里，最稀缺的不是"懂论文的人"，而是能交付端到端系统的人。 "我读完了 vLLM 全部源码"——招聘方只能信你。 "我从零写了一个 500 行的 mini-vLLM，github 公开仓库，benchmark 表在 README"—— 招聘方可以 5 分钟自己验证。这个差别在 hiring 里叫 hard signal vs soft signal，重要性差一个数量级。

更进一步：当面试官问你"PagedAttention 跟 OS 虚拟内存的区别"， 读过 vLLM 的人能答出 3 个；写过 BlockManager 的人能答出 3 个 + 自己踩过的 1 个反直觉点。 后者是 senior signal，能让你的薪资段跳一档。

时间预算：3 周够吗？

v0.5 的目标定得很谨慎，3 周（约 60-80 小时）对单人完全够。 历届做完的人里，最快 1.5 周 v0.1 跑通、3 周 v0.5 收尾； 最慢的人卡在 v0.1 一周，主要因为 HF 模型 load 不熟。 本章后面的"急救包"专门治这些卡点。

02项目范围 · 明确做什么和不做什么

第一周开工前，把这张表抄一份贴在显示器边上。每次冒出"要不要加 X"的念头，看一眼。

为什么是这些"不做"？三条标准：

1. 不在学习路径上：比如量化是另一个 6 个月的领域，强塞进来稀释主线。
2. 工程量远大于学习量：比如多模型架构支持，写一遍 Mistral 适配跟 Llama 几乎一样，但要花一周做适配层。
3. 依赖你还没掌握的东西：比如多 GPU TP 要 NCCL 通讯调试，先把单卡跑通再说。

03架构 · 4 个核心模块

模块清单 · 每个一段

开工前把每个模块的 spec 写在 design.md 里，不要写完一个改一个。下面是参考规格。

请求生命周期 · 一图看穿

04每周里程碑 · v0.1 → v0.5

下面这张时间线是硬性目标。每周末必须有一个可演示的版本。 "差不多写完了，差一点收尾"不算通过——下一周开始前能 python demo.py 跑出预期输出才算。

05关键设计决策 · 每个都要算一遍

开工前把这 6 个决策都写下来 + 算清楚，存进 design.md。 不算清楚就开写，第 2 周一定要回头改架构。

5.1 block_size 选多少 (= 16)

vLLM 默认 16。你也选 16，但要能讲清楚为什么。trade-off 的两端：

16 是 attention kernel tile size 跟内部碎片的甜点。 内部碎片公式：(block_size / 2) / avg_seq_len，假设 avg_seq_len = 256，block 16 浪费 ≈ 3%。 没什么动力换。

5.2 KV cache 池多大 (Llama 3.1 8B fp16 on A10 24GB)

这是最值得手算一遍的题。算式拆成几步：

(a) 单 token KV 大小

Llama 3.1 8B 架构：32 layers, 8 KV heads (GQA), head_dim = 128, fp16 (2 bytes)。

per_token_kv = 2 (K+V) × 32 × 8 × 128 × 2 = 131072 bytes ≈ 128 KB

(b) 单 block 大小 (block_size=16)

per_block = 128 KB × 16 = 2 MB

(c) 可用 KV 显存

• A10 总显存：24 GB
• 模型权重 fp16：8B × 2 bytes = 16 GB
• activation 与临时缓冲：~2 GB
• 留给 KV pool：24 - 16 - 2 = 6 GB

(d) block 总数

num_blocks = 6 GB / 2 MB = 3072 块

(e) 等价并发能力

• 每块 16 token = 总 49152 token 容量
• 假设 avg 2048 token/seq → 同时容纳 ~24 个并发请求
• vLLM 实测约 30-40，因为 GPU memory 估算更激进、prefix 共享、partial block 复用。

把这一坨算完写进 design.md，再写成函数 compute_kv_pool_size(gpu_mem, model)。 启动时打印这些数字——出问题第一时间就知道是 OOM 边界。

5.3 token budget · 一步算多少 token

scheduler 的核心数。设过小 → 每步 batch 小、吞吐低；设过大 → 单 step 太久、tail latency 高。 vLLM 默认 2048-8192，依硬件而定。

从 prefill throughput 反推合理值：

• A10 上 Llama 8B 朴素 prefill 吞吐 ≈ 5000 tok/s
• 目标 step time ≤ 100ms (TPOT 目标)
• token_budget ≈ 5000 × 0.1 = 500 tok/step 是下限

但 decode 是 memory-bound，500 token 的 decode batch 几乎不增加耗时； 所以实际值定 1024-2048 之间，先用 2048，benchmark 后再调。

5.4 preempt 策略 · recompute vs swap

选 recompute 的理由：你已经有 prefill path，只需要清掉 block 让 request 回到 waiting， 20 行代码以内能搞定。swap 要写 host 内存池、异步拷贝、跨 step 状态机，单独是一周工作量。

5.5 attention backend

三种可选。建议组合：

不要自己写 paged attention CUDA kernel——M4 学的 Triton 知识够你看懂 FlashAttention， 不够你写一个生产级 paged kernel。把这个时间留给 scheduler 和 benchmark。

5.6 sampling · temperature + top_p + top_k

完整 sampling 公式：

1. logits → logits / temperature (temperature=0 → greedy)
2. top_k 截断：保留 top k 个 logits，其余设为 -inf
3. softmax → 概率分布
4. top_p (nucleus) 截断：累积概率超过 p 的低概率位置设为 0，重归一
5. multinomial 采样

PyTorch 全部能搞，~20 行。OpenAI 协议默认 temperature=1, top_p=1。 注意 top_k 在 top_p 之前，符合 HuggingFace transformers 的惯例 (尽管 vLLM 顺序略有差异——cross-check)。

06完整代码骨架 · v0.1 单请求版

下面这段 ~150 行是可以直接 copy-paste 跑起来的 (Linux + CUDA + transformers + flash-attn)。 v0.1 不追求美，追求跑出文本。读完每一节再去敲。

6.1 imports 与 setup

# mini_vllm/v01.py
"""mini-vLLM v0.1 — naive single-request decoder.
No paging, no batching, no HTTP. Just prove the model speaks.
"""
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

MODEL_NAME = "meta-llama/Llama-3.1-8B-Instruct"
DEVICE = "cuda"
DTYPE = torch.float16
MAX_NEW_TOKENS = 200

注意 torch.float16 在 A10 上跑得通；如果你拿到的卡是 A100/H100， 推荐 torch.bfloat16 数值更稳。Llama 官方权重对 bf16 更友好。

6.2 load model · 一次性

def load_model():
    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_NAME,
        torch_dtype=DTYPE,
        attn_implementation="flash_attention_2",  # 关键：用 FA2
        device_map=DEVICE,
    )
    model.eval()
    return tokenizer, model

tokenizer, model = load_model()
config = model.config
print(f"Loaded {MODEL_NAME} · {sum(p.numel() for p in model.parameters())/1e9:.1f}B params")

attn_implementation="flash_attention_2" 在 transformers 4.40+ 是支持的； 如果你的版本不支持，换成 "sdpa"，慢一点但跑得通。

6.3 prepare prompt · chat template

def prepare_prompt(user_msg: str) -> torch.Tensor:
    messages = [
        {"role": "system", "content": "You are a helpful assistant."},
        {"role": "user", "content": user_msg},
    ]
    # 套 chat template 是关键的一步；缺了 instruct 模型输出会乱
    prompt_ids = tokenizer.apply_chat_template(
        messages,
        add_generation_prompt=True,   # 末尾加 assistant 的 prompt header
        return_tensors="pt",
    ).to(DEVICE)
    return prompt_ids  # shape: [1, prompt_len]

prompt_ids = prepare_prompt("用一句话解释 PagedAttention。")
print("Prompt tokens:", prompt_ids.shape[1])

6.4 朴素 KV cache · 最 dumb 版

def make_kv_cache(max_len: int):
    """Naive contiguous KV cache: [num_layers, 2, max_len, n_kv_heads, head_dim]
    Reality check: layout 必须跟 HF Llama forward 接受的 past_key_values 形状匹配。
    HF 用 list[tuple[K, V)] 而不是单 tensor — 我们包一层。
    """
    n_layers = config.num_hidden_layers
    n_kv_heads = config.num_key_value_heads
    head_dim = config.hidden_size // config.num_attention_heads
    # 每层一对 (K, V)，shape [batch=1, n_kv_heads, max_len, head_dim]
    kv = []
    for _ in range(n_layers):
        K = torch.zeros(1, n_kv_heads, max_len, head_dim, dtype=DTYPE, device=DEVICE)
        V = torch.zeros(1, n_kv_heads, max_len, head_dim, dtype=DTYPE, device=DEVICE)
        kv.append((K, V))
    return kv

这个 layout 跟 vLLM 的 paged layout 完全不同。v0.1 用的是 HF 标准 contiguous， v0.2 才换成 paged。一步一步来。

6.5 prefill · 一次性塞 prompt

@torch.no_grad()
def prefill(prompt_ids):
    """跑完整 prompt 一次 forward, 返回 logits[最后 token] 和填好的 KV。"""
    out = model(
        input_ids=prompt_ids,
        use_cache=True,
        return_dict=True,
    )
    logits = out.logits[:, -1, :]            # [1, vocab]
    past_kv = out.past_key_values            # tuple of (K, V)
    return logits, past_kv

logits, past_kv = prefill(prompt_ids)
print("Prefill done, KV layers:", len(past_kv), "K shape:", past_kv[0][0].shape)

注意 past_key_values 是 HF 模型 forward 的第二份返回。 v0.1 直接信任 HF 的 layout，自己不动它；v0.2 才会自己管。

6.6 decode loop · 一次一个 token

@torch.no_grad()
def decode_step(token_id: int, past_kv):
    """单 token forward，返回 next logits 和更新后的 KV。"""
    input_ids = torch.tensor([[token_id]], device=DEVICE)
    out = model(
        input_ids=input_ids,
        past_key_values=past_kv,
        use_cache=True,
        return_dict=True,
    )
    return out.logits[:, -1, :], out.past_key_values

def sample_greedy(logits):
    return int(logits.argmax(dim=-1).item())

def generate(prompt_ids, max_new_tokens=MAX_NEW_TOKENS):
    logits, past_kv = prefill(prompt_ids)
    eos = tokenizer.eos_token_id
    out_tokens = []
    for _ in range(max_new_tokens):
        next_tok = sample_greedy(logits)
        if next_tok == eos:
            break
        out_tokens.append(next_tok)
        # 流式打印
        print(tokenizer.decode([next_tok], skip_special_tokens=True), end="", flush=True)
        logits, past_kv = decode_step(next_tok, past_kv)
    print()
    return out_tokens

generate(prompt_ids)

跑通这段就是 v0.1 验收通过。跑通之后再继续，不要先跳到 v0.2。

07完整代码骨架 · v0.2 加 BlockManager

v0.1 用的是 HF 的 contiguous KV。v0.2 换成 paged。核心改动： (a) 自己管 KV pool tensor；(b) 写 BlockManager；(c) 改 forward 路径不再用 past_key_values，而是自己 gather。

7.1 BlockManager · ~80 行

# mini_vllm/block_manager.py
from collections import deque
from dataclasses import dataclass, field
from typing import Dict, List

BLOCK_SIZE = 16

@dataclass
class BlockManager:
    num_blocks: int
    free_blocks: deque = field(default_factory=deque)
    block_tables: Dict[str, List[int]] = field(default_factory=dict)
    seq_lens: Dict[str, int] = field(default_factory=dict)

    def __post_init__(self):
        self.free_blocks = deque(range(self.num_blocks))

    def can_allocate(self, n_tokens: int) -> bool:
        n_blocks = (n_tokens + BLOCK_SIZE - 1) // BLOCK_SIZE
        return len(self.free_blocks) >= n_blocks

    def allocate(self, req_id: str, n_tokens: int) -> List[int]:
        """初次分配 (prefill 时调一次)"""
        n_blocks = (n_tokens + BLOCK_SIZE - 1) // BLOCK_SIZE
        if len(self.free_blocks) < n_blocks:
            raise RuntimeError(f"OOM: need {n_blocks} blocks, have {len(self.free_blocks)}")
        blocks = [self.free_blocks.popleft() for _ in range(n_blocks)]
        self.block_tables[req_id] = blocks
        self.seq_lens[req_id] = n_tokens
        return blocks

    def append_token(self, req_id: str) -> int:
        """decode 时每生成一个 token 调一次; 必要时新增一个 block."""
        seq_len = self.seq_lens[req_id]
        new_len = seq_len + 1
        # 当前 block 还没满 → 不用新分配
        if new_len % BLOCK_SIZE != 1 or new_len == 1:
            self.seq_lens[req_id] = new_len
            return self.block_tables[req_id][-1]
        # 需要新 block
        if not self.free_blocks:
            raise RuntimeError("OOM during append")
        new_block = self.free_blocks.popleft()
        self.block_tables[req_id].append(new_block)
        self.seq_lens[req_id] = new_len
        return new_block

    def free(self, req_id: str):
        for b in self.block_tables.pop(req_id, []):
            self.free_blocks.append(b)
        self.seq_lens.pop(req_id, None)

    def get_block_table(self, req_id: str) -> List[int]:
        return self.block_tables[req_id]

    def stats(self):
        return {
            "free": len(self.free_blocks),
            "used": self.num_blocks - len(self.free_blocks),
            "util": 1 - len(self.free_blocks) / self.num_blocks,
        }

7.2 paged KV pool · 全局一坨大 tensor

# mini_vllm/kv_pool.py
import torch

class PagedKVCache:
    def __init__(self, num_blocks, num_layers, n_kv_heads, head_dim, dtype, device):
        # Layout: [num_layers, 2 (K/V), num_blocks, BLOCK_SIZE, n_kv_heads, head_dim]
        # 这是 FlashAttention 2 的 page_table API 友好的 layout
        self.cache = torch.zeros(
            num_layers, 2, num_blocks, BLOCK_SIZE, n_kv_heads, head_dim,
            dtype=dtype, device=device,
        )

    def get_k(self, layer):
        return self.cache[layer, 0]  # [num_blocks, BLOCK_SIZE, n_kv_heads, head_dim]

    def get_v(self, layer):
        return self.cache[layer, 1]

7.3 paged attention 调用

# mini_vllm/paged_attention.py
# 调 FlashAttention 2 的 with_kvcache 接口 —— 它原生支持 page_table。
# 注意：这是真实 API；具体签名请对照你装的 flash-attn 版本。
from flash_attn import flash_attn_with_kvcache

def paged_decode_attention(
    q,                  # [1, 1, n_q_heads, head_dim]  当前 token 的 query
    k_cache,            # [num_blocks, BLOCK_SIZE, n_kv_heads, head_dim]
    v_cache,            # 同上
    k_new, v_new,       # [1, 1, n_kv_heads, head_dim]  当前 token 的 K/V (写入用)
    block_table,        # [1, max_blocks]  本 batch 的 block 索引
    cache_seqlens,      # [1]              本 batch 已写入的 token 数
):
    """FA2 自动:
       1. 把 k_new/v_new 写入 k_cache/v_cache 的 cache_seqlens 位置
       2. 用 block_table 做 paged gather
       3. 算 attention 输出
    """
    out = flash_attn_with_kvcache(
        q=q,
        k_cache=k_cache,
        v_cache=v_cache,
        k=k_new,
        v=v_new,
        block_table=block_table,
        cache_seqlens=cache_seqlens,
        causal=True,
    )
    return out  # [1, 1, n_q_heads, head_dim]

7.4 改造 Llama forward · 替换 attention 层

HF Llama 自带 attention 层不支持 page_table。两种选项：

• (简) Monkey-patch LlamaAttention.forward，把它替换成调上面的 paged_decode_attention。
• (更干净) 复制 modeling_llama.py 到本地，改 attention 层。

v0.2 推荐 monkey-patch，等 v0.3 scheduler 稳定后再 refactor。 具体怎么 patch 不在这一节展开——这是 transformers 内部 API， 建议读 vLLM 的 vllm/model_executor/models/llama.py看它如何把 attention 接到自己的 paged kernel。

08完整代码骨架 · v0.3 加 Scheduler

v0.3 是整个项目最难的一步。这周不要碰 HTTP、不要碰 sampling 复杂逻辑， 全力打磨调度器。

8.1 Request dataclass

# mini_vllm/request.py
from dataclasses import dataclass, field
from enum import Enum
from typing import List, Optional
import uuid

class RequestState(Enum):
    WAITING = "waiting"          # 还没 prefill
    RUNNING = "running"          # 已经在 batch 里 decode
    PREEMPTED = "preempted"      # 被抢占，回到 waiting
    FINISHED = "finished"

@dataclass
class SamplingParams:
    temperature: float = 1.0
    top_p: float = 1.0
    top_k: int = -1              # -1 表示不截断
    max_tokens: int = 256
    stop_token_ids: List[int] = field(default_factory=list)

@dataclass
class Request:
    request_id: str
    prompt_token_ids: List[int]
    sampling_params: SamplingParams
    state: RequestState = RequestState.WAITING
    output_token_ids: List[int] = field(default_factory=list)
    # KV cache 长度 = prompt + output 已 commit 到 cache 的 token 数
    num_computed_tokens: int = 0
    arrival_time: float = 0.0

    @property
    def total_len(self) -> int:
        return len(self.prompt_token_ids) + len(self.output_token_ids)

    @property
    def is_prefill(self) -> bool:
        return self.num_computed_tokens < len(self.prompt_token_ids)

    @staticmethod
    def new(prompt_ids: List[int], params: SamplingParams) -> "Request":
        return Request(
            request_id=str(uuid.uuid4()),
            prompt_token_ids=prompt_ids,
            sampling_params=params,
        )

8.2 SchedulerOutput · 一次 step 的计划

# mini_vllm/scheduler.py
from dataclasses import dataclass, field
from collections import deque
from typing import List, Dict
from .request import Request, RequestState

@dataclass
class SchedulerOutput:
    # 本 step 要跑的 requests
    scheduled_reqs: List[Request] = field(default_factory=list)
    # 每个请求本步要算的 token 数（prefill 时可能 chunked）
    num_tokens_per_req: Dict[str, int] = field(default_factory=dict)
    # 本步抢占了谁
    preempted_req_ids: List[str] = field(default_factory=list)
    # 本步完成了谁 (EOS / max_tokens)
    finished_req_ids: List[str] = field(default_factory=list)

8.3 Scheduler 主体 · ~120 行

class Scheduler:
    def __init__(self, block_manager, token_budget: int = 2048, max_running: int = 64):
        self.block_manager = block_manager
        self.token_budget = token_budget
        self.max_running = max_running
        self.waiting: deque[Request] = deque()
        self.running: List[Request] = []

    def add_request(self, req: Request):
        self.waiting.append(req)

    def schedule(self) -> SchedulerOutput:
        out = SchedulerOutput()
        budget_remaining = self.token_budget

        # ----- 阶段 1: running 队列继续 decode -----
        new_running = []
        for req in self.running:
            if budget_remaining < 1:
                break
            # decode 一定要先确保 KV 有空间 (需要 append 一个 token)
            if not self.block_manager.can_allocate(1):
                # 显存压力 → 抢占队尾
                victim = self.running[-1] if self.running[-1] is not req else None
                if victim is not None:
                    self._preempt(victim)
                    out.preempted_req_ids.append(victim.request_id)
                    new_running = [r for r in new_running if r.request_id != victim.request_id]
                else:
                    break
            out.scheduled_reqs.append(req)
            out.num_tokens_per_req[req.request_id] = 1
            budget_remaining -= 1
            new_running.append(req)
        self.running = new_running

        # ----- 阶段 2: waiting 队列入场 (prefill) -----
        while self.waiting and budget_remaining > 0 and len(self.running) < self.max_running:
            req = self.waiting[0]
            need = len(req.prompt_token_ids) - req.num_computed_tokens
            # chunked prefill: 一次最多塞 budget_remaining 个 token
            chunk = min(need, budget_remaining)
            if not self.block_manager.can_allocate(chunk):
                break
            if req.state == RequestState.WAITING:
                # 第一次 prefill，分配 block
                self.block_manager.allocate(req.request_id, chunk)
                req.state = RequestState.RUNNING
                self.running.append(req)
            self.waiting.popleft()
            out.scheduled_reqs.append(req)
            out.num_tokens_per_req[req.request_id] = chunk
            budget_remaining -= chunk

        return out

    def _preempt(self, req: Request):
        """recompute preempt: 丢 KV，回 waiting，重置 num_computed_tokens."""
        self.block_manager.free(req.request_id)
        req.state = RequestState.PREEMPTED
        req.num_computed_tokens = 0
        # 已生成的 output tokens 保留 (假设它们值得保留)
        # 简化：直接当 prompt 部分重 prefill
        self.waiting.appendleft(req)

    def mark_finished(self, req: Request):
        req.state = RequestState.FINISHED
        self.block_manager.free(req.request_id)
        self.running = [r for r in self.running if r.request_id != req.request_id]

8.4 Engine 主循环 · 串起来

# mini_vllm/engine.py
from .scheduler import Scheduler
from .block_manager import BlockManager
from .request import Request, SamplingParams, RequestState
from .runner import ModelRunner

class Engine:
    def __init__(self, model_name: str, num_blocks: int):
        self.runner = ModelRunner(model_name, num_blocks)
        self.block_manager = BlockManager(num_blocks=num_blocks)
        self.scheduler = Scheduler(self.block_manager)
        self.tokenizer = self.runner.tokenizer

    def add_request(self, prompt: str, params: SamplingParams) -> str:
        # 套 chat template
        prompt_ids = self.tokenizer.apply_chat_template(
            [{"role": "user", "content": prompt}],
            add_generation_prompt=True,
        )
        req = Request.new(prompt_ids, params)
        self.scheduler.add_request(req)
        return req.request_id

    def step(self) -> List[tuple]:
        """跑一步, 返回 [(req_id, new_token_id, finished_bool), ...]"""
        sched_out = self.scheduler.schedule()
        if not sched_out.scheduled_reqs:
            return []
        # 执行 forward
        logits_per_req = self.runner.execute(sched_out)
        # 采样
        outputs = []
        for req in sched_out.scheduled_reqs:
            if req.is_prefill:
                # prefill chunk 没结束就不采样
                req.num_computed_tokens += sched_out.num_tokens_per_req[req.request_id]
                if req.num_computed_tokens < len(req.prompt_token_ids):
                    continue
            new_tok = self.runner.sample(logits_per_req[req.request_id], req.sampling_params)
            req.output_token_ids.append(new_tok)
            self.block_manager.append_token(req.request_id)
            finished = (
                new_tok == self.tokenizer.eos_token_id
                or len(req.output_token_ids) >= req.sampling_params.max_tokens
            )
            if finished:
                self.scheduler.mark_finished(req)
            outputs.append((req.request_id, new_tok, finished))
        return outputs

09完整代码骨架 · v0.4 加 FastAPI

v0.4 把 Engine 套上 OpenAI 协议。重点是异步——HTTP request 等 streaming 时， engine 主循环还在跑 step；两边用 asyncio queue 桥接。

9.1 server 主体

# mini_vllm/server.py
import asyncio
import json
import time
import uuid
from typing import AsyncIterator
from fastapi import FastAPI, Request as HTTPRequest
from fastapi.responses import StreamingResponse
from pydantic import BaseModel
from .engine import Engine
from .request import SamplingParams

app = FastAPI()
engine: Engine = None  # 启动时填
# req_id -> asyncio.Queue 用来推 token 给等待中的 HTTP 处理函数
output_queues: dict[str, asyncio.Queue] = {}

class ChatMessage(BaseModel):
    role: str
    content: str

class ChatCompletionRequest(BaseModel):
    model: str
    messages: list[ChatMessage]
    temperature: float = 1.0
    top_p: float = 1.0
    max_tokens: int = 256
    stream: bool = False

@app.on_event("startup")
async def startup():
    global engine
    engine = Engine("meta-llama/Llama-3.1-8B-Instruct", num_blocks=3000)
    asyncio.create_task(engine_loop())

async def engine_loop():
    """背景任务：不停 step, 把 token 推到对应 req 的 queue"""
    while True:
        outputs = engine.step()
        for req_id, tok, finished in outputs:
            q = output_queues.get(req_id)
            if q is not None:
                await q.put((tok, finished))
        if not outputs:
            await asyncio.sleep(0.001)  # 让出 event loop

@app.post("/v1/chat/completions")
async def chat_completions(req: ChatCompletionRequest):
    user_msg = req.messages[-1].content  # 简化
    params = SamplingParams(
        temperature=req.temperature,
        top_p=req.top_p,
        max_tokens=req.max_tokens,
    )
    req_id = engine.add_request(user_msg, params)
    output_queues[req_id] = asyncio.Queue()

    if req.stream:
        return StreamingResponse(
            stream_sse(req_id, req.model),
            media_type="text/event-stream",
        )
    # 非流式：等到 finished
    tokens = []
    while True:
        tok, finished = await output_queues[req_id].get()
        tokens.append(tok)
        if finished:
            break
    output_queues.pop(req_id, None)
    text = engine.tokenizer.decode(tokens, skip_special_tokens=True)
    return {
        "id": f"chatcmpl-{uuid.uuid4().hex[:8]}",
        "object": "chat.completion",
        "created": int(time.time()),
        "model": req.model,
        "choices": [{"index": 0, "message": {"role": "assistant", "content": text}, "finish_reason": "stop"}],
    }

async def stream_sse(req_id: str, model_name: str) -> AsyncIterator[str]:
    cmpl_id = f"chatcmpl-{uuid.uuid4().hex[:8]}"
    try:
        while True:
            tok, finished = await output_queues[req_id].get()
            text = engine.tokenizer.decode([tok], skip_special_tokens=True)
            chunk = {
                "id": cmpl_id,
                "object": "chat.completion.chunk",
                "created": int(time.time()),
                "model": model_name,
                "choices": [{"index": 0, "delta": {"content": text}, "finish_reason": "stop" if finished else None}],
            }
            yield f"data: {json.dumps(chunk)}\n\n"
            if finished:
                yield "data: [DONE]\n\n"
                break
    finally:
        output_queues.pop(req_id, None)

跑起来：

uvicorn mini_vllm.server:app --host 0.0.0.0 --port 8000

验证：

import openai
client = openai.OpenAI(base_url="http://localhost:8000/v1", api_key="not-needed")
stream = client.chat.completions.create(
    model="mini-vllm",
    messages=[{"role": "user", "content": "为什么 KV cache 要分页？"}],
    stream=True,
)
for chunk in stream:
    print(chunk.choices[0].delta.content or "", end="", flush=True)

10"对账"方法论 · 跟 vLLM 对照

这是这个月最有价值的副产物：每写完一个模块，立刻打开 vLLM 同位置的代码对照。 盲目读 vLLM 学到的是知识；写完自己版本再对账，学到的是设计 trade-off。

对账清单 · 文件级别

你会在对账中发现的典型差异

• vLLM 的 BlockManager 有 ref_count，你的没有。 它是为 prefix caching 准备的——多个序列共享同一 prompt 时，前缀 block 引用计数 > 1。 你简化掉是对的，但要在 README 里写"v0.5 不支持 prefix sharing"。
• vLLM 的 schedule() 有 ~600 行，你的 ~100 行。 多在：chunked prefill 边界、多种 sampling param 的分组、speculative decoding hook、metrics。 每一段都是一个 PR 解决一个真实问题——读 PR 记录比读最终代码更长见识。
• vLLM 用多进程 (engine + 多个 worker)。 你单进程跑没问题，但 v0.5 benchmark 时会发现 GIL 影响 tokenization 吞吐。 这是 v0.6 优化方向之一，先记在 NOTES.md。
• vLLM 的 prepare_inputs 用 Triton。 把 SchedulerOutput 翻译成 GPU tensor 这一步，vLLM 写了 Triton kernel 加速。 你用 Python + torch ops，慢 10×——这是单请求 latency 差距的来源之一。

每一条都写进 NOTES.md。这份 NOTES.md 就是你 blog 第二篇的草稿。

11Benchmark · 跟 vLLM 跑同样测试

这是 v0.5 的硬通货。没有 benchmark，整个项目少 80% 价值。

11.1 benchmark protocol

• 硬件：单卡 A10 (24GB) 或 A100 (40GB)。文档里写清楚。
• 模型：Llama 3.1 8B Instruct, fp16。
• 工作负载：ShareGPT 数据集 1000 prompts (vLLM benchmark 标配)。 或者 azure-llm-trace 真实生产 trace。
• 并发设置：1, 8, 16, 32, 64 并发分别跑。
• 指标：
  • Throughput (tokens/sec, requests/sec)
  • TTFT (Time To First Token, p50/p95/p99)
  • TPOT (Time Per Output Token, p50/p95)
  • 显存峰值占用
• warmup：跑 50 个 request warmup 再开始 metric 记录。
• 重复：每个配置至少 3 次取中位数。

11.2 预期数字 · 跟 vLLM 对比

11.3 差距来源 · 4 个真正可分析的

1. CUDA graph 缺失 (~30% 损失)。 vLLM 默认在 decode 阶段用 CUDA graph capture，把 launch overhead 从 ~50us/op 压到 ~5us。 decode 一步要 launch 几十个 kernel，graph capture 直接省一半。 mitigation：v0.6 加 torch.cuda.graph。
2. Triton prepare_inputs (~20% 损失)。 把 SchedulerOutput → flat GPU tensors 这一步，Python 跑 ~5ms/step。 vLLM 写了 Triton kernel，~0.5ms/step。32 并发时这是大头。 mitigation：先 profile 确认是瓶颈，再决定是否值得写。
3. Kernel 选择差异 (~15% 损失)。 vLLM 在 H100 上用 FlashAttention 3, 在 A100/A10 上用 FA2 但有 fused QKV projection； 你只用了 stock FA2。 mitigation：v0.6 适配 FA3 + 自己写 fused projection。
4. 多 worker GIL 影响 (~10% 损失)。 tokenization 和 detokenization 在主线程跑会跟 engine step 抢 GIL。 vLLM 把 tokenizer worker 独立进程。 mitigation：v0.6 加 tokenizer subprocess。

把上面 4 点写进 README 的 "Performance Gap Analysis" 一节。能讲清楚 4 个差距 = 你比 90% 候选人懂 serving。

12README 模板

README 是你的项目对外的脸。下面这份模板照着写，能在 1-2 小时内出第一版。

# mini-vLLM

A 500-line educational LLM serving engine, built from scratch in 3 weeks.
Implements paged KV cache and continuous batching à la vLLM.

**Not a production engine.** Built to demonstrate I can build one.

![demo gif](docs/demo.gif)

## Why this exists

vLLM is ~50K lines and hard to learn from. This is a minimal implementation
of the same three core ideas — PagedAttention, continuous batching, OpenAI API —
in <500 lines, designed to be read in one sitting.

## Architecture

[insert architecture SVG/PNG here, same as §03 in tutorial]

- `block_manager.py` — fixed-size KV blocks, allocate/free/append
- `scheduler.py` — FCFS continuous batching, recompute preempt
- `runner.py` — model forward + sampling
- `engine.py` — main loop wiring above
- `server.py` — OpenAI-compatible FastAPI server

## Quickstart

```bash
pip install -r requirements.txt
python -m mini_vllm.server  # localhost:8000
```

```python
import openai
client = openai.OpenAI(base_url="http://localhost:8000/v1", api_key="x")
stream = client.chat.completions.create(...)
```

## Benchmarks (A10 24GB, Llama 3.1 8B, ShareGPT 1K)

| Metric | mini-vLLM | vLLM 0.6 | Gap |
|---|---|---|---|
| Throughput @ 32 concurrent | 900 tok/s | 3000 tok/s | 3.3× |
| TTFT p50 (1K prompt) | 150 ms | 80 ms | 1.9× |
| TPOT p95 | 35 ms | 20 ms | 1.8× |

## Performance gap analysis

1. **No CUDA graph (~30%)** — see [issue #1]
2. **Python prepare_inputs (~20%)** — see [issue #2]
3. **Stock FA2, no fused QKV (~15%)** — see [issue #3]
4. **Single-process tokenizer (~10%)** — see [issue #4]

## Design decisions

See [`design.md`](./design.md) — block_size choice, KV pool sizing,
token budget, preempt policy, attention backend.

## Limitations (intentional)

- No multi-GPU / TP / PP
- No quantization (fp16 only)
- No LoRA
- No prefix caching
- No speculative decoding
- Llama 3.1 8B only

## What I learned

See blog post: ["I wrote a mini-vLLM in 3 weeks"](https://...)

## Credits

Inspired by [vLLM](https://github.com/vllm-project/vllm).
Built as part of [Infra Learning Path](https://weishuz.com/infra-learning-path/).

填空就是了。重点是不要藏短板——把不做的写清楚，反而让做了的部分显得有重量。

13blog 第 2 篇 · 写作模板

blog 是 README 之外的叙事。READE 答"是什么"，blog 答"为什么 + 怎么做的 + 你能学到什么"。

结构 · 5 个段落

1. 钩子 (200 字)。 举例："我花了 3 周写一个 500 行 Python 的 LLM serving engine，比 vLLM 慢 3 倍，但学到了 10 倍。下面是 3 个发现。" 钩子的目的是让人想读下去，不是让人佩服你。
   技巧：开头给一个具体数字 + 一个反直觉 claim。
2. 对照 (600 字)。 vLLM 在内存 / 调度 / kernel 三层各做了什么，你做了什么。 画一张表（你的 vs vLLM）。这一节让读者知道你不是 vibes，而是真懂。
   技巧：写实际的代码 diff 或对比代码片段——比文字描述强 5×。
3. 差距分析 (800 字)。 把 §11.3 的 4 个差距源展开成 4 段，每段：现象 → profiler 看到的图 → 原因 → 怎么补。 这一节是全文 hardest signal，能写完 = 你真做过。
   技巧：贴 nsight / py-spy / nvprof 截图，比文字可信度高 10×。
4. 意外发现 (500 字)。 写的过程中发现自己 M2-M4 漏掉的至少 3 个细节。 举例："我以为 vLLM 的 schedule() 一直返回 100% 利用率的 batch，写完才发现它经常吐 partial batch 来避免长尾。" 这一节最 humble，但最让面试官印象深。
   技巧：用第一人称 + 具体 commit hash，不用"我们"或"大家都知道"。
5. 结论 (200 字)。 "用到再学"在这个项目里的实操：哪些事是必须用到才学得会的，哪些事是读论文就够的。 这一节给读者带走点东西，决定他/她会不会分享这篇。
   技巧：结尾给读者一个具体下一步（"如果你想自己写一个，从这个 repo fork 开始"）。

14常见卡点 · 急救包

下面是历届做这个项目最常踩的 6 个坑，按出现频率排序。卡住时先查这里。

14.1 "模型 load 不了 / 巨慢"

症状：from_pretrained 跑 10 分钟还没出 progress bar；或报 401。

诊断：

• Llama 3.1 模型需要 HF 申请 access。第一次 load 要 huggingface-cli login。
• HF 镜像问题：用 HF_ENDPOINT=https://hf-mirror.com 加速。
• 磁盘空间：8B fp16 = 16GB 下载。df -h ~/.cache/huggingface 看看。
• 先本地 cache 验证：AutoTokenizer.from_pretrained(MODEL_NAME) 单独跑，确认 tokenize 文件下完了。

14.2 "fp16 输出乱码"

症状：模型能跑，但输出全是 !!!!! 或 ### ### ### 或乱码 token。

诊断顺序：

1. chat template 套了吗？ 90% 的乱码是这个。打印 tokenizer.decode(prompt_ids[0]) 看 prompt 是不是带了 <|begin_of_text|>、<|start_header_id|> 这些标记。
2. 试 bf16。Llama 官方权重对 bf16 数值更稳。如果你在 A10 上 (支持 bf16) 直接换。
3. attention_mask 对吗？ prefill 时要 [[1, 1, 1, ..., 1]]，decode 时要 [[1, ..., 1]] 长度 = 累积已 cache + 当前 token。
4. position_ids 对吗？ decode 时是 [[cur_pos]] 不是 [[0]]。

14.3 "KV cache 越来越慢"

症状：第 10 个 token 100ms, 第 100 个 token 500ms, 第 500 个 token 几秒。

诊断：99% 是你在每 step 做 torch.cat 拼 KV，O(n²) 复杂度。

修法：预分配最大长度的 KV tensor，每 step 原地写一行 (kv[:, :, cur_pos:cur_pos+1, :] = new_kv)。 这是 v0.1 → v0.2 升级的核心动机：paged 天然是预分配。

14.4 "throughput 不增长"

症状：v0.3 加了 scheduler，benchmark 32 并发跟单请求一样慢。

诊断：

• 看 sched_out.scheduled_reqs 长度——是不是真的每 step 都 batch > 1?
• 看 token_budget 是不是太小，把请求都拒了？
• 看 ModelRunner 的 forward——你是不是一次只跑一个 req，而不是 stack 起来 batch forward？

最常见错误：scheduler 选了 16 个 req，但 runner 用了 for req in reqs: forward(req)。 正确做法是把 16 个 req 的 input torch.cat 成一个 batch forward。

14.5 "OOM 在第二个请求"

症状：第一个 req 跑完，第二个 req 来了立刻 OOM。

诊断：KV 没释放。检查：

• scheduler.mark_finished(req) 有没有调 block_manager.free(req.request_id)?
• block_manager.free 有没有把 block 放回 free_blocks?
• 有没有可能两个 req 拿到了同一个 block_id (race condition)?

调试技巧：每 step 末尾打印 block_manager.stats()。util 应该波动而不是单调上升。

14.6 "FastAPI streaming 客户端断了 engine 还在跑"

症状：用户 curl ctrl-C，engine 继续生成到 max_tokens，浪费算力。

修法：在 stream_sse 的 finally 里 abort engine 的 request:

try:
    ...
finally:
    output_queues.pop(req_id, None)
    engine.abort_request(req_id)  # 你需要在 Engine 实现这个

abort_request 内部：找到对应 Request，标记 FINISHED，让 scheduler 下次 step 时 free 掉。

15后续延伸 · v0.6 及以后

v0.5 交付后，下面这些方向哪个对你的下一份工作价值最高？挑 1-2 个继续做。

16仓库交付物 · checklist

v0.5 收尾时对照这张表。少一项扣 README 一颗星。

• README.md：架构图 + Quickstart + benchmark 表 + 差距分析 + 限制列表 + Credits。
• design.md：§05 的 6 个设计决策，每个 200-500 字。
• NOTES.md：跟 vLLM 的对照笔记，至少 10 条具体 diff。
• benchmarks/：benchmark 脚本 + 数据 CSV + 画图代码 + 输出图片。
• tests/：至少覆盖 BlockManager 和 Scheduler 的关键路径 (allocate/free/preempt/edge case)。
• docs/demo.gif：30 秒视频，从 python -m mini_vllm.server 到收到 streaming token。
• Blog 第 2 篇已发布，README 里链接到位。
• GitHub repo 公开，topics 加 llm-inference vllm pagedattention。

17本页自检

勾选状态会保存在你的浏览器 localStorage 里，下次打开继续。