立志写出让纯夏都能看懂的教程讲解

终于来到了 model runner 的学习，在学习了 PageAttention 的基础管理和调度逻辑后，我们终于准备窥见模型推理计算的“真面目”了，ModelRunner 层负责了模型的加载、推理执行和多 GPU 调度，是连接 Scheduler 和底层 CUDA 算子的中间层。

简单来说，

1
2
3
4
Scheduler  →  ModelRunner.run(seqs, is_prefill)  →  Model  →  Sampler  →  token_ids
              ├── prepare_prefill / prepare_decode   （构造输入）
              ├── run_model                          （执行推理）
              └── sampler                            （采样 token）

多 GPU 通信机制

虽然示例代码的 qwen3-8b 用单机单卡也能跑，nano-vllm 还是保留了 vLLM 的多卡通信的基本机制 —— 利用 共享内存 实现多卡通信和管理。

具体来说，主进程 rank0 负责调度和运行，并把其他进程运行需要的方法名和参数写入到提前创建好的共享内存 shared memory 中，其他进程 rankX 在事件循环中持续等待，直到收到 rank0 写入 共享内存 的事件，从里面读取内容，然后执行模型计算的算子调用，模型内部算子通过 tensor parallel NCCL 来通信保证中间结果一致，最终只有 rank0 采样出 token_ids 结果并返回给 scheduler。

1
2
3
4
5
6
7
8
9
rank0                        rank1, rank2, ...
  │                               │
  ├─ call("run", seqs, True) ─→  loop() 中阻塞等待 event
  │   ├─ write_shm(方法名, 参数)    │
  │   │   序列化 → 共享内存          │
  │   │   event.set() ──────────→ event.wait() 唤醒
  │   │                            read_shm() → 反序列化
  │   │                            调用对应方法
  │   └─ 自己也执行同一方法

KV Cache 分配

我们在上一章中知道：KV Cache 在 vLLM 中是以 Block 形式存储的，从而避免外部碎片和减少内部碎片。那么在具体代码实现中，我们究竟是怎么存储的这些 KV Cache 呢？model_runner 就是实际上分配 KV Cache 内存的人：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23

def allocate_kv_cache(self):
    config = self.config
    hf_config = config.hf_config
    free, total = torch.cuda.mem_get_info()
    used = total - free
    peak = torch.cuda.memory_stats()["allocated_bytes.all.peak"]
    current = torch.cuda.memory_stats()["allocated_bytes.all.current"]
    num_kv_heads = hf_config.num_key_value_heads // self.world_size
    head_dim = getattr(hf_config, "head_dim", hf_config.hidden_size // hf_config.num_attention_heads)
    block_bytes = 2 * hf_config.num_hidden_layers * self.block_size * num_kv_heads * head_dim * hf_config.torch_dtype.itemsize
    config.num_kvcache_blocks = int(total * config.gpu_memory_utilization - used - peak + current) // block_bytes
    assert config.num_kvcache_blocks > 0
    
    # KV Cache 是一个 6 维的张量
    self.kv_cache = torch.empty(2, hf_config.num_hidden_layers, config.num_kvcache_blocks, self.block_size, num_kv_heads, head_dim)
    layer_id = 0
    for module in self.model.modules():
      # 把每一层分配好的 KV Cache 替换到模型中
        if hasattr(module, "k_cache") and hasattr(module, "v_cache"):
            module.k_cache = self.kv_cache[0, layer_id]
            module.v_cache = self.kv_cache[1, layer_id]
            layer_id += 1

这里对模型 KV Cache 维度的计算需要注意下：

  • 2 是因为 K, V 需要分别存储；
  • num_hidden_layers 代表了有多少个 attention 层，每层都有自己的 KV Cache；
  • num_kvcache_blocks 是通过 int(total * config.gpu_memory_utilization - used - peak + current) // block_bytes 计算的，代表了在推理过程中，我们允许 每一层 分配的最多的 KV Cache Blocks 数量。config.gpu_memory_utilization 是我们运行 nano-vllm 时可以指定的比例，譬如显存 12G，运行时指定 0.5 那么最多能够使用 6G 显存。used 是已经使用的显存，比如存放的权重等。(peak-current) 是在计算运行时占用的临时峰值显存，比如拿来存 activations。这里整体的逻辑就是：可分配量 = 显存上限 − ( 静态占用 + 运行时额外 Peak )。block_bytes 是是 一个 block_id 在所有层加起来 的字节数；
  • block_size 是每个 block 存多少个 token 的 KV；
  • num_kv_heads 是每个 attention 的在单卡存储的 KV 头数量。这里是通过 hf_config.num_key_value_heads // self.world_size 计算的，因为 Tensor Parallelism 中，KV heads 会被均匀分配到多个 GPU 上，所以每张卡只负责 总head数 / GPU数；
  • head_dim 是每个头的 K/V 的维度

推理流程

在梳理了基础知识后，我们终于可以开始正式思考整个执行流程了，为了方便，我们假设整个推理是单卡执行的。

在 llm_engine 中，我们之前有执行 token_ids = self.model_runner.call("run", seqs, is_prefill) 来让 rank0 显卡调用推理逻辑，run() 的具体逻辑是这样的：

  1. 根据入参判断是 prefill 还是 decode，并分别调用 prepare_prefill 和 prepare_decode 给模型构造需要的数据；
  2. rank0 显卡收集准备各个 sequence 的采样参数 temperatures；
  3. 调用 run_model 拿到 next token logits 结果；
  4. rank0 显卡根据 logits 结果采样 next token；
  5. 返回采样后的 token_ids。
1
2
3
4
5
6
7
def run(self, seqs: list[Sequence], is_prefill: bool) -> list[int]:
    input_ids, positions = self.prepare_prefill(seqs) if is_prefill else self.prepare_decode(seqs)
    temperatures = self.prepare_sample(seqs) if self.rank == 0 else None
    logits = self.run_model(input_ids, positions, is_prefill)
    token_ids = self.sampler(logits, temperatures).tolist() if self.rank == 0 else None
    reset_context()
    return token_ids

整个逻辑还是非常清晰简洁的，我们下面将对这几个步骤做更细致的代码解析。

prepare_prefill & prepare_decode

在看具体代码之前，明确我们的目标：prepare 阶段实际上就是为后面 FlashAttention 的算子计算准备需要的数据内容。

我们先看 prefill 阶段的代码，整体思路还是很简单的，整体都是围绕：将各个 sequence 需要计算的 token_id 拼接起来分配到显存中 这一任务进行的。

稍微需要注意的是 cu_seqlens_q 和 cu_seqlens_k ，这两是一个前缀和数组，因为我们传入底层算子接口的 input_ids 是需要将各个 sequence 的 token 全部拼接起来的序列，那么要区分各个 sequence 序列的 token 边界自然是一个前缀和数组了。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
def prepare_prefill(self, seqs: list[Sequence]):
			...
       for seq in seqs:
           seqlen = len(seq)
           # 跳过已经缓存过的 token (我们上一章讲过的 Prefix Cache)
           input_ids.extend(seq[seq.num_cached_tokens:])
           # 同理，只记录未缓存过的位置
           positions.extend(list(range(seq.num_cached_tokens, seqlen)))
           # 实际需要计算的 Q 长度
           seqlen_q = seqlen - seq.num_cached_tokens
           # 实际需要的 K 长度
           seqlen_k = seqlen
           # cu_xx 是一个前缀和序列，
           # 因为我们在 input_ids 等输入里把多个 sequence 拼成一个序列，
           # 我们需要这个 cu_xx 来区分不同 sequence 任务
           cu_seqlens_q.append(cu_seqlens_q[-1] + seqlen_q)
           cu_seqlens_k.append(cu_seqlens_k[-1] + seqlen_k)
           max_seqlen_q = max(seqlen_q, max_seqlen_q)
           max_seqlen_k = max(seqlen_k, max_seqlen_k)
           if not seq.block_table:    # warmup
               continue
           for i in range(seq.num_cached_blocks, seq.num_blocks):
               start = seq.block_table[i] * self.block_size
               if i != seq.num_blocks - 1:
                   end = start + self.block_size
               else:
                   end = start + seq.last_block_num_tokens 
               # KV 写入的物理显存位置
               slot_mapping.extend(list(range(start, end)))
       if cu_seqlens_k[-1] > cu_seqlens_q[-1]:    # prefix cache
           block_tables = self.prepare_block_tables(seqs)
       
       # 将所有数据放到 GPU 中
       input_ids = torch.tensor(input_ids, dtype=torch.int64, pin_memory=True).cuda(non_blocking=True)
       positions = torch.tensor(positions, dtype=torch.int64, pin_memory=True).cuda(non_blocking=True)
       cu_seqlens_q = torch.tensor(cu_seqlens_q, dtype=torch.int32, pin_memory=True).cuda(non_blocking=True)
       cu_seqlens_k = torch.tensor(cu_seqlens_k, dtype=torch.int32, pin_memory=True).cuda(non_blocking=True)
       slot_mapping = torch.tensor(slot_mapping, dtype=torch.int32, pin_memory=True).cuda(non_blocking=True)
       set_context(True, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, slot_mapping, None, block_tables)
       return input_ids, positions

接着是 decode 阶段，这个就想对简单多了，因为每个 sequence 都只有 1 个新 token 需要计算。这里需要注意下 context_lens ，其记录了 sequence 任务目前的 token 长度是多少。之所以要记录这个，主要因为底层的 attention 算子需要根据这个值判断 seq.block_tables 中哪些是块是真正 Cache 有效的。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
def prepare_decode(self, seqs: list[Sequence]):
  ...
    for seq in seqs:
      # 单个 token
        input_ids.append(seq.last_token)
        positions.append(len(seq) - 1)
        # 保存 sequence 的有效长度，
        # KV Cache 根据这个来决定界限
        context_lens.append(len(seq))
        slot_mapping.append(seq.block_table[-1] * self.block_size + seq.last_block_num_tokens  - 1)
    input_ids = torch.tensor(input_ids, dtype=torch.int64, pin_memory=True).cuda(non_blocking=True)
    positions = torch.tensor(positions, dtype=torch.int64, pin_memory=True).cuda(non_blocking=True)
    slot_mapping = torch.tensor(slot_mapping, dtype=torch.int32, pin_memory=True).cuda(non_blocking=True)
    context_lens = torch.tensor(context_lens, dtype=torch.int32, pin_memory=True).cuda(non_blocking=True)
    block_tables = self.prepare_block_tables(seqs)
    set_context(False, slot_mapping=slot_mapping, context_lens=context_lens, block_tables=block_tables)
    return input_ids, positions

在 decode 阶段，我们知道 sequence.block_table 记录了每个 seq 的 KV Cache 存储的 table 物理索引。为了让我们能够构造一个二维 tensor，我们把所有需要 decode 的 seqs 内存长度对齐，不足的补充 -1。

1
2
3
4
5
def prepare_block_tables(self, seqs: list[Sequence]):
    max_len = max(len(seq.block_table) for seq in seqs)
    block_tables = [seq.block_table + [-1] * (max_len - len(seq.block_table)) for seq in seqs]
    block_tables = torch.tensor(block_tables, dtype=torch.int32, pin_memory=True).cuda(non_blocking=True)
    return block_tables

prepare_sample

因为每个 sequence 任务我们都可以指定一个 temperature 采样值，所以我们在正式采样前需要整理下，代码很简单，依然和前面 prepare 阶段的逻辑一脉相承：将各个 sequence 任务需要的参数整合进一个 list 中：

1
2
3
4
5
6
def prepare_sample(self, seqs: list[Sequence]):
    temperatures = []
    for seq in seqs:
        temperatures.append(seq.temperature)
    temperatures = torch.tensor(temperatures, dtype=torch.float32, pin_memory=True).cuda(non_blocking=True)
    return temperatures

run_model

这里我们进入了模型前向推理的部分，这里区分 prefill 和 decode 的作用终于开始体现了出来：

  • prefill 阶段我们直接常规 forward
  • decode 阶段我们使用 CUDA Graph 做加速。
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
@torch.inference_mode()
def run_model(self, input_ids: torch.Tensor, positions: torch.Tensor, is_prefill: bool):
    if is_prefill or self.enforce_eager or input_ids.size(0) > 512:
        return self.model.compute_logits(self.model(input_ids, positions))
    else:
        bs = input_ids.size(0)
        context = get_context()
        graph = self.graphs[next(x for x in self.graph_bs if x >= bs)]
        # 将输入复制到 record 的变量中
        graph_vars = self.graph_vars
        graph_vars["input_ids"][:bs] = input_ids
        graph_vars["positions"][:bs] = positions
        graph_vars["slot_mapping"].fill_(-1)
        graph_vars["slot_mapping"][:bs] = context.slot_mapping
        graph_vars["context_lens"].zero_()
        graph_vars["context_lens"][:bs] = context.context_lens
        graph_vars["block_tables"][:bs, :context.block_tables.size(1)] = context.block_tables
       # replay 执行 forward 计算
        graph.replay()
        return self.model.compute_logits(graph_vars["outputs"][:bs])

什么是 CUDA Graph？

把 GPU 操作 capture 下来，形成一张静态的执行图，之后只需要 replay() 就可以把整个图提交给 GPU 进行计算，避免了反复执行 CPU → GPU 调用导致的 kernel launch overhead。

1
2
3
4
5
6
7
录制阶段（capture）：
  with torch.cuda.graph(graph):
      outputs = model(input_ids, positions)    # 不真正执行，只记录操作序列

回放阶段（replay）：
  graph_vars["input_ids"][:bs] = real_input    # 往录制时的 tensor 里填真实数据
  graph.replay()                                # 一次性提交所有 kernel

这里推荐下 NV 自己的一篇入门 Blog 来更进一步了解：https://developer.nvidia.com/blog/cuda-graphs/

为什么只有 decode 使用 CUDA Graph 加速？

主要是因为 CUDA Graph 是一张静态图，它要求输入 tensor 的 shape 是不能变化的。 prefill 阶段每次 sequence 的长度都是没法确定的，而 decode 时 seq 的输入长度始终固定是 1,只有 batch size 变化因此可以通过填充向上取整的方式来匹配预先生成好的 graph。

capture_cudagraph

让我们回到具体的代码来看看到底是怎么准备 CUDA Graph 这个阶段的。capture 的代码会在 ModelRunner 初始化的时候就被调用执行。

由于静态图的缘故，输入和输出的内存地址是需要固定的，因此需要 self.graph_vars 来保证我们在 capture 阶段使用的变量生命周期能够覆盖整个后续的推理过程。在后续使用的时候，我们也是需要将输入拷贝进 self.graph_vars 来进行推理的（参考上面的 run_model()）。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
@torch.inference_mode()
def capture_cudagraph(self):
    config = self.config
    hf_config = config.hf_config
    max_bs = min(self.config.max_num_seqs, 512)
    max_num_blocks = (config.max_model_len + self.block_size - 1) // self.block_size
    input_ids = torch.zeros(max_bs, dtype=torch.int64)
    positions = torch.zeros(max_bs, dtype=torch.int64)
    slot_mapping = torch.zeros(max_bs, dtype=torch.int32)
    context_lens = torch.zeros(max_bs, dtype=torch.int32)
    block_tables = torch.zeros(max_bs, max_num_blocks, dtype=torch.int32)
    outputs = torch.zeros(max_bs, hf_config.hidden_size)
    self.graph_bs = [1, 2, 4, 8] + list(range(16, max_bs + 1, 16))
    self.graphs = {}
    self.graph_pool = None

# 选择不同 batch_size
    for bs in reversed(self.graph_bs):
        graph = torch.cuda.CUDAGraph()
        set_context(False, slot_mapping=slot_mapping[:bs], context_lens=context_lens[:bs], block_tables=block_tables[:bs])
        outputs[:bs] = self.model(input_ids[:bs], positions[:bs])    # warmup
        with torch.cuda.graph(graph, self.graph_pool):
            outputs[:bs] = self.model(input_ids[:bs], positions[:bs])    # capture
        if self.graph_pool is None:
            self.graph_pool = graph.pool()
        self.graphs[bs] = graph
        torch.cuda.synchronize()
        reset_context()

    self.graph_vars = dict(
        input_ids=input_ids,
        positions=positions,
        slot_mapping=slot_mapping,
        context_lens=context_lens,
        block_tables=block_tables,
        outputs=outputs,
    )

context 是什么

我们在上面的各种 prepare 函数中，经常能看见 set_context(), get_context() 的身影。Context 类主要存储了底层 FlashAttention 算子计算 attention 时需要的一些变量（实际上就是把除了 input_ids, position_ids 以外的变量都放进这个全局变量里）：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
@dataclass
class Context:
    is_prefill: bool = False
    # varlen attention 的 Q 序列边界
    cu_seqlens_q: torch.Tensor | None = None
    # varlen attention 的 K 序列边界
    cu_seqlens_k: torch.Tensor | None = None
    max_seqlen_q: int = 0
    max_seqlen_k: int = 0
    # KV 写入 Cache 的物理位置
    slot_mapping: torch.Tensor | None = None
    # 每个 sequence 的有效 KV 长度
    context_lens: torch.Tensor | None = None
    # 每个 sequence 物理块页表
    block_tables: torch.Tensor | None = None

当然，你要说我就爱一路透传参数难道不行吗？也可以，只是每一层函数封装都要重复写，麻烦一些，nano-vllm 用的全局变量更方便一些。

sampler()

当完成模型的 forward 推理后，我们拿到了 token logits，还需要把这些 logtis 做 temperature-softmax 然后采样得到预测的 next token id。

关于 temperature 的作用，我们在第二节的时候就已经很详细的分析过了。

简单来说，当我们在 Softmax 引入 Temperature T 后，我们便有了缩放模型生成的原始未归一化的 logits 的能力，直接改变最终生成 token 的概率分布形态，更加主动的影响 token 截取数量。此时，新的公式如下：

$$P(X_i)=\frac{exp(z_i / T)}{\sum_jexp(z_j/T)} $$

可以看到，就是除以了一个 T，

  • 当 T=1 时：
    公式等同于标准 softmax 函数。此时不产生缩放干预；
  • 当 T<1 时（降低随机性）：
    将 logits 除以小于 1 的数值，会扩大 logits 之间的绝对差值。经过指数运算后，高分与低分之间的差异被呈指数级放大。概率质量会高度集中于排名靠前的 token。当 T→0 时，基本等效于 argmax 运算，系统将确定性地选择最高分 token；
  • 当 T>1 时（增加随机性）：
    将 logits 除以大于 1 的数值，会衰减 logits 之间的绝对差值。指数运算后的结果使整体概率分布被压缩并变得平缓，逐渐趋近于均匀分布。原本处于边缘位置、低 logit 的候选 token 会被分配到相对更高的采样概率，从而提高了统计上的不确定性。

Sampler 整体的代码基本就是在计算上面的公式，唯一值的说的就是最后的采样代码：

  1. torch.empty_like(probs).exponential_(1) 创建一个与 probs 形状相同的张量，并从指数分布 Exp(1) 中采样填充，得到一组随机噪声；
  2. .clamp_min_(1e-10) 将噪声张量中的值截断到最小 1e-10，防止后续出现除以零的情况；
  3. probs.div_(...) 将原始概率逐元素除以噪声，计算 probs[i] / noise[i] ；
  4. .argmax(dim=-1) 取上一步结果中最大值的下标，这就是采样出的 token。

在数学上，这个做法等价于按概率取样（具体数学证明看后续补充吧）。

1
2
3
4
5
6
7
8
9
10
11
12
13
class Sampler(nn.Module):

    def __init__(self):
        super().__init__()

    @torch.compile
    def forward(self, logits: torch.Tensor, temperatures: torch.Tensor):
		    # sequence 除以各自对应的 temperature
        logits = logits.float().div_(temperatures.unsqueeze(dim=1))
        # sequence 做 softmax
        probs = torch.softmax(logits, dim=-1)
        sample_tokens = probs.div_(torch.empty_like(probs).exponential_(1).clamp_min_(1e-10)).argmax(dim=-1)
        return sample_tokens

说些什么吧！

beaudar