import pickle import torch import torch.distributed as dist from multiprocessing.synchronize import Event from multiprocessing.shared_memory import SharedMemory from nanovllm.config import Config from nanovllm.engine.sequence import Sequence from nanovllm.models.qwen3 import Qwen3ForCausalLM from nanovllm.layers.sampler import Sampler from nanovllm.utils.context import set_context, get_context, reset_context from nanovllm.utils.loader import load_model class ModelRunner: """模型运行器：负责模型推理、KV cache 管理、CUDA Graph 捕获和张量并行通信。 在张量并行（TP）模式下: - Rank 0 是主进程，负责采样和与引擎通信。 - Rank > 0 是工作进程，通过共享内存（SharedMemory）接收指令。 - 所有进程共享同一个模型和 KV cache 的分片。 生命周期: 1. 初始化: 加载模型 → warmup → 分配 KV cache → (可选)捕获 CUDA Graph 2. 推理: 接收序列 → 准备输入 → 运行模型 → 采样 token 3. 退出: 释放资源 """ def __init__(self, config: Config, rank: int, event: Event | list[Event]): self.config = config hf_config = config.hf_config self.block_size = config.kvcache_block_size self.enforce_eager = config.enforce_eager self.world_size = config.tensor_parallel_size self.rank = rank self.event = event # 初始化分布式进程组（NCCL 后端），所有 GPU 通过 TCP 通信 dist.init_process_group("nccl", "tcp://localhost:2333", world_size=self.world_size, rank=rank) torch.cuda.set_device(rank) # 加载模型权重 default_dtype = torch.get_default_dtype() torch.set_default_dtype(hf_config.dtype) torch.set_default_device("cuda") self.model = Qwen3ForCausalLM(hf_config) load_model(self.model, config.model) self.sampler = Sampler() # Warmup: 运行一次前向传播以确定模型本身的显存占用 self.warmup_model() # 根据剩余显存分配 KV cache self.allocate_kv_cache() # 捕获 CUDA Graph 以加速 decode 阶段的小批量推理 if not self.enforce_eager: self.capture_cudagraph() torch.set_default_device("cpu") torch.set_default_dtype(default_dtype) # 张量并行时，rank > 0 的工作进程进入消息循环 if self.world_size > 1: if rank == 0: # 主进程创建共享内存，工作进程打开它 self.shm = SharedMemory(name="nanovllm", create=True, size=2**20) dist.barrier() else: dist.barrier() self.shm = SharedMemory(name="nanovllm") self.loop() # 工作进程在此循环，直到收到 exit 指令 def exit(self): """释放所有资源并退出。""" if self.world_size > 1: self.shm.close() dist.barrier() if self.rank == 0: self.shm.unlink() # 只有创建者需要 unlink if not self.enforce_eager: del self.graphs, self.graph_pool torch.cuda.synchronize() dist.destroy_process_group() def loop(self): """工作进程的主循环：等待主进程指令，执行对应方法。""" while True: method_name, args = self.read_shm() self.call(method_name, *args) if method_name == "exit": break def read_shm(self): """从共享内存读取主进程发送的方法调用指令。""" assert self.world_size > 1 and self.rank > 0 self.event.wait() # 等待主进程通知 n = int.from_bytes(self.shm.buf[0:4], "little") method_name, *args = pickle.loads(self.shm.buf[4:n+4]) self.event.clear() return method_name, args def write_shm(self, method_name, *args): """将方法调用指令写入共享内存，通知工作进程。""" assert self.world_size > 1 and self.rank == 0 data = pickle.dumps([method_name, *args]) n = len(data) self.shm.buf[0:4] = n.to_bytes(4, "little") self.shm.buf[4:n+4] = data for event in self.event: event.set() def call(self, method_name, *args): """调用指定方法。TP 模式下主进程先通知工作进程，再本地执行。""" if self.world_size > 1 and self.rank == 0: self.write_shm(method_name, *args) method = getattr(self, method_name, None) return method(*args) def warmup_model(self): """预热模型：运行一次最大批量的前向传播。 目的是让 PyTorch 分配所有内部缓存（cuBLAS workspace 等）， 然后通过 empty_cache 释放临时显存，这样后续的 peak memory 统计 就只包含模型权重，从而准确计算 KV cache 可用空间。 """ torch.cuda.empty_cache() torch.cuda.reset_peak_memory_stats() max_num_batched_tokens, max_model_len = self.config.max_num_batched_tokens, self.config.max_model_len seq_len = min(max_num_batched_tokens, max_model_len) num_seqs = min(max_num_batched_tokens // seq_len, self.config.max_num_seqs) seqs = [Sequence([0] * seq_len) for _ in range(num_seqs)] for seq in seqs: seq.num_scheduled_tokens = seq_len self.run(seqs, True) torch.cuda.empty_cache() def allocate_kv_cache(self): """根据剩余 GPU 显存分配 KV cache。 计算公式: 可用显存 = 总显存 × gpu_memory_utilization - 非模型占用 其中非模型占用 = 已用显存 - peak（模型权重）+ current（当前模型张量） KV cache 形状: (2, num_layers, num_blocks, block_size, num_kv_heads/head_dim) 其中第一维 2 分别对应 K 和 V cache。 """ 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) # 每个块占用的字节数：2(K+V) × 层数 × block_size × KV头数 × head_dim × dtype字节数 block_bytes = 2 * hf_config.num_hidden_layers * self.block_size * num_kv_heads * head_dim * hf_config.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 张量，形状为 (2, num_layers, num_blocks, block_size, num_kv_heads, head_dim) self.kv_cache = torch.empty(2, hf_config.num_hidden_layers, config.num_kvcache_blocks, self.block_size, num_kv_heads, head_dim) # 将 KV cache 的视图绑定到模型中每个 Attention 层 layer_id = 0 for module in self.model.modules(): 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 def prepare_block_tables(self, seqs: list[Sequence]): """将序列的 block_table 列表填充为等长的二维张量，用于 GPU 计算。""" max_len = max(len(seq.block_table) for seq in seqs) # 用 -1 填充短序列的 block_table 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 def prepare_prefill(self, seqs: list[Sequence]): """准备 prefill 阶段的模型输入张量。 Prefill 阶段需要处理多个序列的 prompt tokens，所有 token 被拼接成一个连续的输入。 使用 cu_seqlens（累积序列长度）来标记每个序列的边界，供 flash_attn_varlen 使用。 关键数据: - input_ids: 所有序列的 token ID 拼接。 - positions: 每个 token 的位置 ID（考虑前缀缓存偏移）。 - cu_seqlens_q/k: 查询和键值的累积序列长度。 - slot_mapping: 将每个 token 映射到 KV cache 中的物理存储位置。 - block_tables: 前缀缓存命中时需要 block_table 来从 KV cache 读取已缓存的 K/V。 """ input_ids = [] positions = [] cu_seqlens_q = [0] cu_seqlens_k = [0] max_seqlen_q = 0 max_seqlen_k = 0 slot_mapping = [] block_tables = None for seq in seqs: start = seq.num_cached_tokens # 跳过已缓存的 token seqlen_q = seq.num_scheduled_tokens end = start + seqlen_q seqlen_k = end # KV 的长度是从 0 到 end（包括缓存前缀） input_ids.extend(seq[start:end]) positions.extend(range(start, end)) 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 阶段没有 block_table continue # 计算 slot_mapping：每个 token 对应 KV cache 中的哪个 slot start_block = start // self.block_size end_block = (end + self.block_size - 1) // self.block_size for i in range(start_block, end_block): slot_start = seq.block_table[i] * self.block_size if i == start_block: slot_start += start % self.block_size if i != end_block - 1: slot_end = seq.block_table[i] * self.block_size + self.block_size else: slot_end = seq.block_table[i] * self.block_size + end - i * self.block_size slot_mapping.extend(range(slot_start, slot_end)) if cu_seqlens_k[-1] > cu_seqlens_q[-1]: # 前缀缓存命中时，KV 长度 > Q 长度 block_tables = self.prepare_block_tables(seqs) 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 def prepare_decode(self, seqs: list[Sequence]): """准备 decode 阶段的模型输入张量。 Decode 阶段每个序列只处理 1 个 token（最新生成的 token）。 模型从 KV cache 中读取之前所有的 K/V 向量来做注意力计算。 关键数据: - input_ids: 每个序列的最新 token ID。 - positions: 每个 token 的位置 ID（序列长度 - 1）。 - slot_mapping: 新 token 的 KV 写入位置。 - context_lens: 每个序列的上下文总长度。 - block_tables: KV cache 块映射表。 """ input_ids = [] positions = [] slot_mapping = [] context_lens = [] for seq in seqs: input_ids.append(seq.last_token) positions.append(len(seq) - 1) context_lens.append(len(seq)) # slot = 最后一个块的起始位置 + 该块内已有 token 数 - 1 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 def prepare_sample(self, seqs: list[Sequence]): """准备采样所需的温度参数张量。""" temperatures = [seq.temperature for seq in seqs] temperatures = torch.tensor(temperatures, dtype=torch.float32, pin_memory=True).cuda(non_blocking=True) return temperatures @torch.inference_mode() def run_model(self, input_ids: torch.Tensor, positions: torch.Tensor, is_prefill: bool): """运行模型前向传播。 对于 decode 阶段的小批量（<=512），使用 CUDA Graph 加速： CUDA Graph 将整个计算图"录制"下来，后续只需回放即可，避免了 CPU 端的 kernel launch 开销，对 decode（每个 step 计算量很小）尤为有效。 """ if is_prefill or self.enforce_eager or input_ids.size(0) > 512: # 直接运行：prefill（批量动态）、eager 模式、或大批量 decode return self.model.compute_logits(self.model(input_ids, positions)) else: # 使用 CUDA Graph 回放加速小批量 decode bs = input_ids.size(0) context = get_context() # 选择 >= bs 的最小预捕获图大小 graph = self.graphs[next(x for x in self.graph_bs if x >= bs)] 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 # 回放图（比重新执行快，跳过了 Python/PyTorch 调度开销） graph.replay() return self.model.compute_logits(graph_vars["outputs"][:bs]) def run(self, seqs: list[Sequence], is_prefill: bool) -> list[int]: """执行一次完整的推理步骤：准备输入 → 模型前向 → 采样。 Returns: 采样得到的 token ID 列表（仅 rank 0 返回有效值）。 """ 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 @torch.inference_mode() def capture_cudagraph(self): """预捕获不同批量大小的 CUDA Graph。 CUDA Graph 要求输入张量的地址不变（同一个内存池），所以需要预分配 固定大小的输入缓冲区，并为每个 batch size 录制一个图。 预捕获的 batch size: [1, 2, 4, 8, 16, 32, ..., max_bs] 运行时选择 >= 实际 batch size 的最小预捕获图。 """ 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) # 要捕获的 batch size 列表 self.graph_bs = [1, 2, 4, 8] + list(range(16, max_bs + 1, 16)) self.graphs = {} self.graph_pool = None # 逆序捕获：先捕获大的 batch size，共享同一个 graph pool 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]) # 捕获计算图 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, )