nano vllm代码详解

nano vllm是简化版的vllm的实现, 里面实现了大模型推理的常见优化. 包括page attention, kv cache, continuous batching.

文件结构:

E:\NANO-VLLM
│   .gitignore
│   bench.py
│   example.py
│   LICENSE
│   pyproject.toml
│   README.md
│
├───assets
│       logo.png
│
└───nanovllm
    │   config.py
    │   llm.py
    │   sampling_params.py
    │   __init__.py
    │
    ├───engine
    │       block_manager.py
    │       llm_engine.py
    │       model_runner.py
    │       scheduler.py
    │       sequence.py
    │
    ├───layers
    │       activation.py
    │       attention.py
    │       embed_head.py
    │       layernorm.py
    │       linear.py
    │       rotary_embedding.py
    │       sampler.py
    │
    ├───models
    │       qwen3.py
    │
    └───utils
            context.py
            loader.py

1. 配置文件

核心配置文件都在config.py中. llm.py是对外暴露接口的. sampling_params.py是采样的参数配置. engine文件夹定义了调度, runner, block相关的文件. layers是定义模型相关层的文件. models定义模型. utils定义了一些全局contex和加载模型权重的函数.

@dataclass
class Config:
    model: str
    max_num_batched_tokens: int = 16384
    max_num_seqs: int = 512
    max_model_len: int = 4096
    gpu_memory_utilization: float = 0.9
    tensor_parallel_size: int = 1
    enforce_eager: bool = False
    hf_config: AutoConfig | None = None
    eos: int = -1
    kvcache_block_size: int = 256
    num_kvcache_blocks: int = -1

    def __post_init__(self):
        assert os.path.isdir(self.model)
        assert self.kvcache_block_size % 256 == 0
        assert 1 <= self.tensor_parallel_size <= 8
        self.hf_config = AutoConfig.from_pretrained(self.model)
        self.max_model_len = min(self.max_model_len, self.hf_config.max_position_embeddings)
        assert self.max_num_batched_tokens >= self.max_model_len

• max_num_batched_tokens: 一次请求的最大token数
• max_num_seqs: 一次请求的最大seq数目
• max_model_len: 一次请求的seq的最大长度
• tensor_parallel_size: GPU数目
• kvcache_block_size: kv cache的block size.
• num_kvcache_blocks:

2. Sequence

sequence.py: sequence abstraction，用来跟踪一个请求（prompt + 生成内容）的状态、token、以及 KV cache block 的映射关系。Sequence 表示 一个推理请求（一个 prompt + 它生成的 token）。它负责管理：

• 当前生成到哪里了（token_ids）
• prompt 和生成部分的划分
• KV cache 对应的 block（block_table）
• 当前状态（WAITING / RUNNING / FINISHED）
• sampling 参数（temperature 等）

class SequenceStatus(Enum):
    WAITING = auto()
    RUNNING = auto()
    FINISHED = auto()

表示这个序列在调度器里的状态：

• WAITING：还没开始跑（排队）
• RUNNING：正在 decode
• FINISHED：已经结束（遇到 EOS 或 max_tokens）

self.seq_id = next(Sequence.counter): 全局唯一 ID（自增）,self.token_ids = copy(token_ids): 当前所有 token（prompt + generation）, self.last_token = token_ids[-1]: 上一个 token（decode 用）

self.num_tokens = len(self.token_ids)
self.num_prompt_tokens = len(token_ids)

• 初始化时：全部都是 prompt
• 后面 append_token 才增加 completion

self.num_cached_tokens = 0
self.block_table = []

• block_table： 映射 逻辑 block → 物理 KV cache block
• num_cached_tokens： 已经算过 KV cache 的 token 数

self.temperature = sampling_params.temperature
self.max_tokens = sampling_params.max_tokens
self.ignore_eos = sampling_params.ignore_eos

• 每个 sequence 可以有不同 decoding 策略

def num_completion_tokens(self):
    return self.num_tokens - self.num_prompt_tokens

completion token 数

def prompt_token_ids(self):
    return self.token_ids[:self.num_prompt_tokens]
def completion_token_ids(self):
    return self.token_ids[self.num_prompt_tokens:]

prompt / completion 切分

block_size = 256

KV cache 按 block 管理（类似 vLLM）

已缓存 block 数, 前多少 block 已经有 KV cache

def num_cached_blocks(self):
    return self.num_cached_tokens // self.block_size

总共需要多少块才能覆盖已有的 token: ceil(num_tokens / block_size)

def num_blocks(self):
    return (self.num_tokens + self.block_size - 1) // self.block_size

def last_block_num_tokens(self):
    return self.num_tokens - (self.num_blocks - 1) * self.block_size

最后一个 block 的 token 数, 用来判断：

• 当前 block 是否满了

• 是否需要分配新 block

  @property
  def num_blocks(self): # 总共需要多少块才能覆盖已有的 token
  return (self.num_tokens + self.block_size - 1) // self.block_size

  @property
  def last_block_num_tokens(self):
  return self.num_tokens - (self.num_blocks - 1) * self.block_size

  def block(self, i): # 第 i 个 block 的 token
  assert 0 <= i < self.num_blocks
  return self.token_ids[i*self.block_size: (i+1)*self.block_size]

def append_token(self, token_id: int):
    self.token_ids.append(token_id)
    self.last_token = token_id
    self.num_tokens += 1

每生成一个 token：

• 加入 sequence
• 更新 last_token
• token 数 +1
• 这里只是逻辑 token
• KV cache 分配不在这里做（通常在 scheduler / block manager）

Sequence 只是记录：

• 用了哪些 block（block_table）
• 哪些 token 已 cache（num_cached_tokens）

counter = count()

全局变量, 记录当前已经完成了多少sequence.

2. context

@dataclass
class Context:
    is_prefill: bool = False
    cu_seqlens_q: torch.Tensor | None = None
    cu_seqlens_k: torch.Tensor | None = None
    max_seqlen_q: int = 0
    max_seqlen_k: int = 0
    slot_mapping: torch.Tensor | None = None
    context_lens: torch.Tensor | None = None
    block_tables: torch.Tensor | None = None

_CONTEXT = Context()

def get_context():
    return _CONTEXT

def set_context(is_prefill, cu_seqlens_q=None, cu_seqlens_k=None, max_seqlen_q=0, max_seqlen_k=0, slot_mapping=None, context_lens=None, block_tables=None):
    global _CONTEXT
    _CONTEXT = Context(is_prefill, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, slot_mapping, context_lens, block_tables)

def reset_context():
    global _CONTEXT
    _CONTEXT = Context()

单例模式, 创建一个全局的context, 发生的事情是：

• Python 读取文件
• 执行整个文件（包括 _CONTEXT = Context()）
• 把模块对象缓存到sys.modules["my_context_module"]

之后再在别的地方：import my_context_module, Python 会直接从 sys.modules 里拿, 不会重新执行文件代码, 不会再次执行 _CONTEXT = Context()

这里的context是为了设置一些全局变量, 防止某些变量需要通过函数传参层层传递的麻烦.

3. Block和BlockManager

• KV cache 分块（block）: 每block_size=256 tokens 一个 block, 避免连续内存分配（类似分页）

• prefix cache: 如果两个 sequence 有相同前缀：

  seq1: A B C D | E F
  seq2: A B C D | X Y

前 4 个 token 对应的 block 可以共享

• 引用计数（ref_count）
  • 多个 sequence 可以共享一个 block
  • 用 ref_count 管理释放

3.1 Block 类（一个 KV cache 页）

class Block:
    def __init__(self, block_id):
        self.block_id = block_id  # 物理 block ID（显存位置）
        self.ref_count = 0  # 被多少 sequence 使用
        self.hash = -1  # block 对应 token 的 hash（用于 prefix cache）
        self.token_ids = []   # 实际 token（用于校验 hash 冲突）

    def update(self, hash: int, token_ids: list[int]):  # 当 block 写满后才更新 hash
        self.hash = hash
        self.token_ids = token_ids

    def reset(self):  # 分配新 block 时调用
        self.ref_count = 1
        self.hash = -1
        self.token_ids = []

3.2. BlockManager

self.blocks: 所有 block（固定大小 pool）

self.free_block_ids: 空闲 block（deque）

self.used_block_ids: 正在使用的 block

self.hash_to_block_id: prefix cache, hash → block_id

hash 设计:

compute_hash(token_ids, prefix):  # 链式 hash（rolling hash）
	h_i = hash(h_{i-1} + 当前block)

3.2.1 allocate: 给一个 sequence 分配 block

def allocate(self, seq: Sequence):

一个seq会被分成多个block, 比如1个seq_len=10, block_size=4, 那么这个seq会被分成4+4+2, 3个block.

for 每个 block:
    1. 计算 hash
    2. 查 cache（hash_to_block_id）
    3. 命中 → 共享 block
    4. miss → 分配新 block

关键变量：

h = -1
cache_miss = False

一旦 miss，后面全部 miss

逐 block 解析

step1：取 token
token_ids = seq.block(i)
step2：算 hash
h = compute_hash(token_ids, h) if 满block else -1

只有满 block 才参与 cache！

step3：查 cache
block_id = hash_to_block_id.get(h, -1)
step4：判断是否命中
if block_id == -1 or token_ids 不一致:
    cache_miss = True

关键：一旦 miss, 因为 prefix 已经不同了！

if cache_miss:
	后面全部 block 都重新分配！

初始时: self.free_block_ids: deque[int] = deque(range(num_blocks)), 从0~N.

如果miss了:

block_id = self.free_block_ids[0]

self.free_block_ids.remove(block_id)
self.used_block_ids.add(block_id)

如果不满block_size, 直接返回-1

命中情况: 直接复用 KV cache！

seq.num_cached_tokens += block_size
block.ref_count += 1

特殊情况

if block_id in used_block_ids:
	 block 正在用 → ref_count++

否则：从 free list 重新分配（但逻辑上是 reuse）

最后更新 cache

block.update(h, token_ids)
hash_to_block_id[h] = block_id

注意, allocate只在prefill阶段执行一次. 后面的decode都是在做append

3.2.2 deallocate（释放）

for block_id in reversed(seq.block_table):
	从后往前释放（类似栈）
	block.ref_count -= 1

如果没人用了：

if ref_count == 0:
    放回 free_block_ids

注意：不会删除 hash_to_block_id（这是个优化点）

为什么从后往前释放?

从后往前释放是为了符合"后分配先释放（LIFO）"的访问模式，更安全、也更高效

for block_id in reversed(seq.block_table):
    block = self.blocks[block_id]
    block.ref_count -= 1
    if block.ref_count == 0:
        self._deallocate_block(block_id)

也就是：最后一个 block → 倒数第二个 → ... → 第一个 block。核心原因1：符合"生成顺序":

Sequence 的 block 是这样增长的：

block0 → block1 → block2 → ... → blockN

decode 时, 先用 block0（prompt）, 再用 block1, 最后用 blockN（最新生成）. 释放时：最后的 block 最"新"、最可能不被共享. 如果反过来（从前往后）：先处理 block0（ref_count=2）, 再 block1, 逻辑上没错，但不符合访问局部性.

• 后面的block对前面的有依赖性, 因此应该先删除依赖最少的block, 这样如果后面的block可以被复用, 则前面的block肯定可以被复用. 相反, 如果前面的block可以被复用, 后面的Block不一定能复用
• cache locality 更好
• fragmentation 更少
• free list 更稳定

3.2.3. append（decode 核心）

block_table = seq.block_table
last_block = self.blocks[block_table[-1]]

1. 拿到这个 sequence 的 block 列表
2. 找到"当前正在操作的最后一个 block"

# 这个 sequence 用到的所有 block（按顺序）
block_table = seq.block_table

例如：

seq.block_table = [3, 7, 2]

对应：

block 3 → 前缀
block 7 → 中间
block 2 → 当前正在写

block_table[-1]: 取最后一个元素, 当前 sequence 的"最后一个 block"

[3, 7, 2] → 2

self.blocks[block_id]

从全局 block pool 里取出这个 block 对象, 拿到真正的 KV cache block

last_block = self.blocks[block_table[-1]] 意思是：找到当前 sequence 正在写/刚写完的那个 block

完整例子, 假设：

seq.block_table = [3, 7]
block 3: [A B C D]   hash = h1
block 7: [E F _ _]   hash = -1

执行这两行：

block_table = seq.block_table
last_block = self.blocks[block_table[-1]]

得到：last_block = block 7, 接下来会发生什么？

• 如果 append G：

  E F G _

  走：

else:
    assert last_block.hash == -1

什么都不做

• 如果 append H：

  E F G H

走：

elif len(seq) % block_size == 0:

用的就是：

last_block.update(...)

• 如果 append I：
  新 block

用的还是：last_block.hash != -1

def can_append(self, seq):
    return len(self.free_block_ids) >= (len(seq) % self.block_size == 1)

只有在：新开一个 block 时才需要分配。因为len(seq) % self.block_size == 1成立的时候, 代表新的block开始. (len(seq) % self.block_size == 1)的值为1.

• 如果需要新 block → free block 是否 ≥ 1
• 如果不需要 → 永远返回 True

3.2.4. may_append（重点）

每生成 一个新 token（decode step） 就会调用一次：旧 seq + 1 token → 调用 may_append

• block_table: 当前 sequence 用了哪些 block（按顺序）. 例如：block_table = [3, 7, 2]
• last_block: 当前正在写 or 刚写完的 block,
• block 的状态

状态	hash
写入中	-1
已完成	（可复用） != -1

can_append为True才会执行may_append

核心逻辑：分三种情况

情况1：刚进入新 block, 触发时机

if len(seq) % block_size == 1:

刚 append 一个 token，并且：新 token 是这个 block 的第一个 token

例子（block_size=4）
[A B C D]  (满)
append E → len=5 -> %4 == 1 成立

assert last_block.hash != -1

上一个 block 必须已经是"完成态"

block_id = self.free_block_ids[0]

拿一个空闲 block

self._allocate_block(block_id)

将当前block_id标记为使用中（通常会设置 ref_count=1 等）

block_table.append(block_id)

加到 sequence 的 block 列表里

结果:

新 block 开始写：
[E _ _ _]

情况2：刚好写满 block

elif len(seq) % self.block_size == 0:  # 该写入block了
   assert last_block.hash == -1  # 上一个block还没写呢
   token_ids = seq.block(seq.num_blocks-1)  # 当前block的所有ids
   prefix = self.blocks[block_table[-2]].hash if len(block_table) > 1 else -1 # 上一个block的hash, 前缀hash
   h = self.compute_hash(token_ids, prefix)  # 根据当前的tokens和上一个的prefix hash计算当前block的hash
   last_block.update(h, token_ids)  # 更新hash
   self.hash_to_block_id[h] = last_block.block_id

触发时机: block 刚好被填满

例子
[E F G H]  (刚写完)

assert last_block.hash == -1

说明这个 block 之前是"写入中", 也就是还没有被写入, 但是此时又是满的。因此后面我们需要写入该block, 并计算hash值。

token_ids = seq.block(seq.num_blocks - 1)

拿到这个 block 的完整 token

prefix = self.blocks[block_table[-2]].hash if len(block_table) > 1 else -1

前一个 block 的 hash（prefix）

h = self.compute_hash(token_ids, prefix)

计算 prefix-aware hash, 本质：

hash = hash(prefix + 当前 block)
last_block.update(h, token_ids)

更新 block, 保存 hash, 保存 token_ids

self.hash_to_block_id[h] = last_block.block_id

放进 prefix cache

结果:

这个 block 从：写入中 → 可复用 cache

h = compute_hash(...)
last_block.update(h, token_ids)
hash_to_block_id[h] = block_id

这一步才让 block 可被复用

情况3：block 中间, else, 什么都不做

核心设计:

1. 只有"满 block"才参与 cache, 避免 partial block 混乱
2. rolling hash（prefix-aware）, 保证 prefix 一致才能共享
3. 一旦 miss，后面全 miss, 因为 prefix 不一样了
4. ref_count 实现共享, 多 sequence 共用 KV cache
5. append 时才 finalize block, decode 阶段逐步构建 cache

一个seqence来了之后开始不断建立kv cache, 每256个token分配一个kv cache block, 不足就不分配。当前sequence生成结束后, 对应的block的ref减1.注意: 只有当ref为0的时候, 对应的block才会被释放。

kv cache block的服用发生在并发请求的时候。

并发请求时间线：

t1: seq1: A B C D E F   (正在生成)
t2: seq2: A B C D X Y   (新请求进来)

allocate seq2 时：

block0: [A B C D] → 命中 seq1 的 block, 复用发生了！

block1: miss

关键点：seq1 还没结束 → block 还在 → 可以复用

现实中：并发请求非常多（最关键）

比如：

1. 1000 个用户同时问：
   "Write a Python function to...", 前缀一样： 所有人共享 KV cache！

2. streaming / batching, 调度器会这样做：

   step1: 收集多个请求
   step2: 一起 prefill
   step3: 一起 decode

同时存在 → 可复用

3. beam search / sampling: 一个 prompt → 多个分支, 前缀完全一样：
   prompt → 多条生成路径, KV cache 直接共享