读 Nano-vLLM 源码：一个极简 LLM 推理引擎是如何工作的

1
nanovllm/
2
├── llm.py                    # 用户入口（仅5行）
3
├── engine/                   # 引擎核心
4
│   ├── llm_engine.py        # LLM引擎主类
5
│   ├── scheduler.py         # 请求调度器
6
│   ├── block_manager.py     # KV Cache块管理
7
│   ├── model_runner.py      # 模型执行器
8
│   └── sequence.py          # 序列数据结构
9
├── layers/                   # 神经网络层
10
│   ├── attention.py         # 注意力层
11
│   ├── linear.py            # 线性层
12
│   ├── rotary_embedding.py  # 旋转位置编码
13
│   ├── layernorm.py         # 层归一化
14
│   ├── activation.py        # 激活函数
15
│   ├── sampler.py           # 采样器
16
│   └── embed_head.py        # 嵌入头
17
├── models/                   # 模型定义
18
│   └── qwen3.py             # Qwen3模型
19
└── utils/                    # 工具函数
20
    ├── context.py           # 上下文管理
21
    └── loader.py            # 模型加载器

1
class BlockManager:
2
    def __init__(self, num_blocks, block_size):
3
        self.num_blocks = num_blocks      # 物理块总数
4
        self.block_size = block_size      # 每块存储的token数
5
        self.free_blocks = set(range(num_blocks))
6

7
    def allocate(self, num_blocks):
8
        """分配物理块"""
9
        allocated = []
10
        for _ in range(num_blocks):
11
            if not self.free_blocks:
12
                raise MemoryError("Out of blocks")
13
            block_id = self.free_blocks.pop()
14
            allocated.append(block_id)
15
        return allocated
16

17
    def free(self, block_ids):
18
        """释放物理块"""
19
        for block_id in block_ids:
20
            self.free_blocks.add(block_id)

1
class ModelRunner:
2
    def __init__(self, model, block_size):
3
        self.model = model
4
        self.block_size = block_size
5

6
    def run(self, sequences, blocks):
7
        """执行模型推理"""
8
        # 1. 准备输入tensor
9
        input_ids = self.prepare_input_ids(sequences)
10
        positions = self.prepare_positions(sequences)
11

12
        # 2. 构建KV Cache索引
13
        block_tables = self.prepare_block_tables(blocks)
14

15
        # 3. 执行模型forward
16
        hidden_states = self.model.forward(
17
            input_ids, positions, block_tables
18
        )
19

20
        # 4. 采样生成下一个token
21
        next_tokens = self.sample(hidden_states)
22
        return next_tokens

1
class Sequence:
2
    def __init__(self, prompt, sampling_params):
3
        self.prompt = prompt
4
        self.sampling_params = sampling_params
5
        self.token_ids = []          # 生成的token序列
6
        self.block_table = []        # 逻辑块->物理块映射
7
        self.status = "waiting"      # 状态：waiting/running/swapped/finished
8

9
    def append_token(self, token_id):
10
        """追加新生成的token"""
11
        self.token_ids.append(token_id)
12
        # 检查是否需要新块
13
        if len(self.token_ids) % self.block_size == 0:
14
            self._allocate_new_block()
15

16
    def is_finished(self):
17
        """检查是否完成"""
18
        return (
19
            len(self.token_ids) >= self.max_tokens or
20
            token_id in self.stop_tokens
21
        )

1
用户请求 → LLMEngine.generate()
2
              ↓
3
         创建Sequence对象
4
              ↓
5
         加入Scheduler.waiting队列
6
              ↓
7
    ┌──── Scheduler.schedule() ←────┐
8
    │         ↓                     │
9
    │   分配BlockManager物理块       │
10
    │         ↓                     │
11
    │   ModelRunner.run() 执行推理   │
12
    │         ↓                     │
13
    │   生成新token，更新Sequence    │
14
    │         ↓                     │
15
    │   检查是否完成                 │
16
    │    ├── 完成 → 移至完成队列      │
17
    │    └未完成 → 继续循环           │
18
    └─────────────────────────────────┘

1
class LLMEngine:
2
    def generate(self, prompts, sampling_params):
3
        # 1. 创建Sequence
4
        sequences = [Sequence(p, sampling_params) for p in prompts]
5

6
        # 2. 加入等待队列
7
        for seq in sequences:
8
            self.scheduler.add(seq)
9

10
        # 3. 推理循环
11
        while not self.scheduler.is_finished():
12
            # 调度
13
            scheduler_output = self.scheduler.schedule()
14

15
            # 执行模型
16
            output = self.model_runner.run(
17
                scheduler_output.running,
18
                scheduler_output.blocks
19
            )
20

21
            # 更新序列
22
            self._update_sequences(output)
23

24
        return [seq.token_ids for seq in sequences]

1
class Attention:
2
    def forward(self, hidden_states, positions, block_tables):
3
        # 1. 计算Q, K, V
4
        q = self.q_proj(hidden_states)
5
        k = self.k_proj(hidden_states)
6
        v = self.v_proj(hidden_states)
7

8
        # 2. 应用旋转位置编码
9
        q, k = self.rotary_emb(q, k, positions)
10

11
        # 3. 写入KV Cache到物理块
12
        self._write_to_blocks(k, v, block_tables)
13

14
        # 4. 读取历史KV Cache
15
        k_cache, v_cache = self._read_from_blocks(block_tables)
16

17
        # 5. 计算注意力分数
18
        attn_output = self._compute_attention(q, k_cache, v_cache)
19

20
        return self.o_proj(attn_output)