从零开始构建 Mini vLLM：深入理解大模型推理优化

最近看到一些mini版本的vLLM实现，比如：

• https://github.com/ovshake/nano-vllm/
• https://github.com/Wenyueh/MinivLLM/tree/main
• https://github.com/ubermenchh/mini-vllm
• https://github.com/skyzh/tiny-llm
• https://github.com/GeeeekExplorer/nano-vllm

下面本文结合：

• https://github.com/ovshake/nano-vllm/blob/main/BLOG.md

来分享下如何从零开始mini版本的vLLM,

当运行 vLLM 这类大模型推理引擎时，背后究竟发生了什么？这个问题驱使开发者构建了nano-vllm——一个极简的、面向教学的高性能 LLM 推理引擎实现。可以把它理解为"傻瓜版 vLLM"，专门用来揭开推理优化的神秘面纱。

这篇文章将带你深入了解大模型推理的各项优化技术，看看它们究竟是如何工作的。

为什么 LLM 推理这么难搞

运行大模型推理，远不只是做几次矩阵乘法那么简单。传统的逐个请求处理方式会造成惊人的 GPU 显存和算力浪费。

大模型生成文本时分为两个阶段：

Prefill 阶段：一次性处理整个 prompt（计算密集型）

Decode 阶段：逐个生成 token（显存密集型）

Decode 阶段尤其值得关注。生成每个新 token 时，模型都需要通过注意力机制回看所有之前的 token。如果不做缓存，就得反复计算相同的内容。这就是 KV cache 存在的意义。

但问题来了：如果按照最大可能序列长度预分配 KV cache 显存，会造成极大的浪费。假设最大长度是 2048 tokens，但实际序列只有 100 tokens，那就浪费了 95% 的显存！

这正是 vLLM 通过 PagedAttention 解决的核心问题，也是 nano-vllm 要实现的关键技术。

整体架构设计

nano-vllm 的代码组织结构如下：

nano_vllm/ ├── engine.py # 主推理引擎 ├── config.py # 模型配置 ├── cache.py # KV cache 实现 ├── sampler.py # Token 采样 ├── core/ │ ├── sequence.py # 请求跟踪 │ ├── scheduler.py # 带优先级的批调度 │ ├── block.py # PagedAttention 的内存块 │ └── block_manager.py # 内存块分配器（类似 OS 内存管理） ├── attention/ │ ├── paged_attention.py # PagedAttention 核心实现 │ └── flash_attention.py # FlashAttention 集成 ├── speculative/ │ └── speculative_decoding.py # 推测解码 ├── educational/ # 可视化学习模式 │ ├── narrator.py # 白话解释 │ ├── xray.py # 张量可视化 │ └── dashboard.py # 实时终端界面 └── model/ ├── loader.py # HuggingFace 模型加载 └── llama.py # Llama 实现（RMSNorm、RoPE、GQA、SwiGLU）

接下来逐个剖析核心优化技术。

PagedAttention：vLLM 的灵魂

传统方案的问题

传统的 KV cache 分配方式，就像为一个人预订整个电影院，"以防"他带来 1999 个朋友。这种做法显然很浪费。

传统方案会根据最大可能长度，为每个序列预分配一大块连续显存，导致：

• 内存碎片化：不同序列在不同时间结束，留下空洞
• 显存浪费：大多数序列永远达不到最大长度
• 批处理受限：GPU 显存装不下太多请求

PagedAttention 的解决方案

PagedAttention 借鉴了操作系统虚拟内存的思想。它不再连续分配，而是把 KV cache 切分成固定大小的内存块（类似内存页）：

# 来自 core/block.py@dataclassclassBlock:"""固定大小的 KV cache 内存块 每个块存储 block_size 个 token 的 KV 状态 """block_id:intblock_size:int=16# 每块 16 个 tokenref_count:int=1# 用于共享（前缀缓存）prefix_hash:Optional[int]=None

每个序列获得一个BlockTable——从逻辑位置到物理块的映射表：

# 来自 core/block.py@dataclassclassBlockTable:"""将逻辑位置映射到物理块 类似虚拟内存中的页表： - 位置 p 的 token 在逻辑块：p // block_size - 块内槽位：p % block_size - 物理块：block_ids[p // block_size] 示例（block_size=16，序列有 35 个 token）： block_table.block_ids = [5, 12, 3] # 3 个物理块 Token 0-15 -> 块 5 Token 16-31 -> 块 12 Token 32-34 -> 块 3（槽位 0-2） """block_ids:List[int]block_size:int=16

BlockManager负责分配管理，就像 OS 管理内存一样：

# 来自 core/block_manager.pyclassBlockManager:"""管理 KV cache 块的分配 使用简单的空闲列表（栈）实现 O(1) 分配/释放 """defallocate_block(self)->int:ifnotself.free_blocks:raiseRuntimeError("KV cache 块用尽！")returnself.free_blocks.pop()deffree_block(self,block_id:int)->None:block=self.blocks[block_id]ifblock.decrement_ref()<=0:self.free_blocks.append(block_id)

Paged Attention 的计算过程

计算注意力时，需要从不连续的块中收集 K 和 V：

# 来自 attention/paged_attention.pydefpaged_attention(query:torch.Tensor,key_cache:torch.Tensor,# [num_blocks, block_size, num_kv_heads, head_dim]value_cache:torch.Tensor,block_tables:List[BlockTable],context_lens:List[int],block_size:int,num_kv_heads:int,)->torch.Tensor:# 为每个序列从块中收集数据forbatch_idxinrange(batch_size):block_table=block_tables[batch_idx]forposinrange(context_len):logical_block=pos//block_size slot_in_block=pos%block_size physical_block=block_table.block_ids[logical_block]# 从缓存中复制gathered_keys[batch_idx,:,pos,:]=key_cache[physical_block,slot_in_block]gathered_values[batch_idx,:,pos,:]=value_cache[physical_block,slot_in_block]# 标准注意力计算attn_weights=torch.matmul(query,gathered_keys.transpose(-2,-1))*scale# ... 应用 mask、softmax 并计算输出

为什么 PagedAttention 如此重要

PagedAttention 带来的好处：

• 接近零显存浪费：只分配实际需要的空间
• 内存共享：相同前缀可以共享块（前缀缓存）
• 更高吞吐量：显存能容纳更多请求，实现更高并行度

连续批处理：告别空闲等待

传统批处理的问题

传统批处理要等待批次中所有序列都完成才能开始新请求。假设有：

• 请求 A：需要生成 50 个 token
• 请求 B：只需生成 5 个 token

请求 B 很快完成，但必须等待请求 A。GPU 只能闲置！

连续批处理的解决方案

nano-vllm 以迭代粒度进行调度：

• 新请求可以中途加入批次
• 完成的请求立即离开
• GPU 保持忙碌状态

调度器的实际运行过程：

# 来自 core/scheduler.pyclassScheduler:"""管理序列的生命周期： - WAITING：排队中 - RUNNING：处理中 - SWAPPED：被抢占 - FINISHED：已完成 """defschedule(self)->SchedulerOutputs:outputs=SchedulerOutputs()# 1. 如果有高优先级请求等待，处理抢占ifself.enable_preemptionandself.block_manager:self._handle_preemption(outputs)# 2. 继续运行中的序列（decode）forseqinself.running:ifseq.is_chunked_prefill():outputs.chunked_prefill_sequences.append(seq)else:outputs.decode_sequences.append(seq)# 3. 从等待队列接纳新序列whilecan_admit_more():seq=self._pop_waiting()seq.status=SequenceStatus.RUNNING outputs.prefill_sequences.append(seq)returnoutputs

引擎在一次迭代中处理这些序列：

# 来自 engine.pydefstep(self)->List[GenerationOutput]:"""连续批处理的一次迭代"""scheduler_outputs=self.scheduler.schedule()# 处理分块 prefillforseq,num_tokensinzip(chunked_prefill_seqs,chunked_prefill_tokens):self._run_chunked_prefill(seq,num_tokens)# 处理完整 prefill（新序列）forseqinprefill_sequences:self._run_prefill(seq)# 处理 decode（批量一起处理！）ifdecode_sequences:self._run_decode(decode_sequences)# 返回完成的序列returnnewly_finished

优先级调度与抢占机制

有时某些请求需要 VIP 待遇。nano-vllm 支持以下特性：

基于优先级的调度

请求带有优先级属性，优先级高的先处理：

# 来自 core/scheduler.pydef_get_priority_key(self,seq:Sequence)->Tuple[int,float,int]:"""堆排序的优先级键。元组值越小，优先级越高"""# 取反优先级，让高值排在前面return(-seq.priority,seq.arrival_time,seq.seq_id)# 使用堆实现 O(log n) 调度heapq.heappush(self._waiting_heap,(priority_key,sequence))

抢占：踢出低优先级请求

当高优先级请求到来但显存不足时，可以抢占低优先级的运行中请求：

# 来自 core/scheduler.pydef_handle_preemption(self,outputs):"""为高优先级等待序列抢占低优先级序列"""highest_waiting=self._peek_waiting()whilenotself.block_manager.can_allocate(blocks_needed)andself.running:# 找到优先级最低的运行序列lowest_running=min(self.running,key=lambdas:s.priority)ifhighest_waiting.priority>lowest_running.priority:# 抢占！释放块并重置以便重新计算self.running.remove(lowest_running)self.block_manager.free_sequence_blocks(lowest_running.block_table)lowest_running.reset_for_recompute()self._push_waiting(lowest_running)

被抢占的序列回到等待队列，稍后会重新进行 prefill。这是基于重计算的抢占（相比交换到 CPU 内存），实现更简单，实践中效果也不错。

前缀缓存：共享通用前缀

许多请求的开头都是相同的系统提示词。为什么要重复计算相同的 KV cache？

工作原理

块根据 token 内容和在序列中的位置进行哈希：

# 来自 core/block.pydefhash_token_block(token_ids:Tuple[int,...],parent_hash:Optional[int]=None)->int:"""包含整个前缀链的累积哈希 这确保只有在整个前缀匹配时才共享块 """ifparent_hashisNone:returnhash(token_ids)returnhash((parent_hash,token_ids))

新序列到来时，检查其前缀块是否已存在：

# 来自 core/block_manager.pydefallocate_blocks_with_prefix_caching(self,token_ids:List[int]):"""分配块，尽可能复用缓存的前缀块"""parent_hash=Noneforblock_idxinrange(num_full_blocks):block_tokens=tuple(token_ids[start:end])cache_key=(parent_hash,block_tokens)ifcache_keyinself.prefix_cache:# 缓存命中！复用现有块cached_block_id=self.prefix_cache[cache_key]self.blocks[cached_block_id].increment_ref()# 引用计数block_table.append_block(cached_block_id)else:# 缓存未命中 - 分配新块block_id=self.allocate_block()self.prefix_cache[cache_key]=block_id block_table.append_block(block_id)parent_hash=self.blocks[block_id].prefix_hashreturnblock_table,shared_prefix_len

引用计数确保块在仍被其他序列使用时不会被释放。

分块 Prefill：避免长 Prompt 阻塞

一个很长的 prompt（比如 4000 tokens）在 prefill 时会阻塞整个批次。分块 prefill 将其拆分成更小的片段：

# 来自 engine.pydef_run_chunked_prefill_paged(self,seq:Sequence,num_tokens:int):"""处理一块 prompt tokens"""start_pos=seq.num_prefilled_tokens end_pos=start_pos+num_tokens chunk_tokens=seq.prompt_token_ids[start_pos:end_pos]# 为这一块分配块# ...# 只对这一块进行前向传播logits=self.model(input_ids,block_kv_cache=...,start_positions=[start_pos])# 更新进度seq.num_prefilled_tokens=end_pos# 只有在所有 prompt tokens 都处理完后才采样ifseq.num_prefilled_tokens>=len(seq.prompt_token_ids):next_token=self.sampler.sample(logits)seq.append_token(next_token.item())

调度器控制每次迭代 prefill 多少 token：

# max_prefill_tokens 限制每次迭代的计算量ifprompt_len<=prefill_budget:outputs.prefill_sequences.append(seq)# 完整 prefillelse:outputs.chunked_prefill_sequences.append(seq)# 部分 prefilloutputs.chunked_prefill_tokens.append(prefill_budget)

FlashAttention：高效的注意力计算

标准注意力会具化完整的 N×N 注意力矩阵。对于 2048 token 的序列，那就是 400 万个元素！FlashAttention 使用分块技术避免这个问题。

在 nano-vllm 中的集成

# 来自 attention/flash_attention.pydefflash_attention(query,key,value,causal=True):"""使用 FlashAttention，显存复杂度 O(N) 而非 O(N^2)"""# FlashAttention 期望：[batch, seq_len, num_heads, head_dim]query=query.transpose(1,2)key=key.transpose(1,2)value=value.transpose(1,2)output=flash_attn_func(query,key,value,causal=causal)returnoutput.transpose(1,2)# 统一接口，带降级方案defattention(query,key,value,use_flash_attn=True,causal=True):ifuse_flash_attnandFLASH_ATTN_AVAILABLE:returnflash_attention(query,key,value,causal)# 降级到 PyTorch SDPA（也是优化过的！）returnF.scaled_dot_product_attention(query,key,value,is_causal=causal)

FlashAttention 在模型的注意力层中使用：

# 来自 model/llama.pyclassLlamaAttention(nn.Module):def__init__(self,config,layer_idx,use_flash_attn=True):self.use_flash_attn=use_flash_attnandis_flash_attn_available()defforward(self,hidden_states,...):# ... 计算 Q、K、V 并应用 RoPE ...# 使用统一注意力接口（如可用则用 FlashAttention）attn_output=unified_attention(query=query_states,key=key_states,value=value_states,use_flash_attn=self.use_flash_attn,causal=True,)

推测解码：草稿与验证

Decode 很慢，因为一次只生成一个 token。如果能在大模型的一次前向传播中生成多个 token 呢?

核心思路

1. 使用小而快的草稿模型生成 K 个候选 token
2. 大模型验证所有 K+1 个位置，只需一次前向传播
3. 接受匹配的 token，拒绝不匹配的并重新采样

# 来自 speculative/speculative_decoding.pydef_speculative_step(self,current_ids,target_kv_cache,draft_kv_cache,remaining_tokens):"""一次推测解码步骤"""K=self.config.num_speculative_tokens# 步骤 1：生成 K 个草稿 token（快速！）draft_tokens,draft_probs=self._generate_draft_tokens(current_ids,draft_kv_cache,K)# 步骤 2：用目标模型验证（一次前向传播处理 K+1 个 token！）verify_ids=[[current_ids[-1]]+draft_tokens]target_logits=self.target_model(verify_ids,kv_cache=target_kv_cache)target_probs=F.softmax(target_logits,dim=-1)# 步骤 3：使用拒绝采样接受/拒绝accepted_tokens=[]fori,draft_tokeninenumerate(draft_tokens):target_prob=target_probs[0,i,draft_token].item()draft_prob=draft_probs[i]# 如果目标概率 >= 草稿概率则接受（保持目标分布！）acceptance_prob=min(1.0,target_prob/draft_prob)ifrandom()<acceptance_prob:accepted_tokens.append(draft_token)else:# 从调整后的分布重新采样resampled=sample_from_adjusted(target_probs[0,i],draft_prob,draft_token)accepted_tokens.append(resampled)break# 第一次拒绝后停止# 如果全部接受，再采样一个额外 token！iflen(accepted_tokens)==len(draft_tokens):bonus_token=sample(target_probs[0,-1])accepted_tokens.append(bonus_token)returnaccepted_tokens

神奇之处：无质量损失

这是拒绝采样——数学上保证输出分布与目标模型完全相同。没有任何近似！

加速效果取决于：

• 草稿模型速度（应该比目标模型快约 10 倍）
• 接受率（越高表示每次目标前向传播获得更多 token）
• K 值（更多推测 = 更大潜在收益）

Llama 模型的完整实现

nano-vllm 包含从头编写的 Llama 实现，具备所有现代特性：

RMSNorm（替代 LayerNorm）

# 来自 model/llama.pyclassRMSNorm(nn.Module):"""均方根归一化 - 比 LayerNorm 更简单"""defforward(self,x):rms=torch.sqrt(x.pow(2).mean(dim=-1,keepdim=True)+self.eps)returnx/rms*self.weight

旋转位置编码（RoPE）

# 来自 model/llama.pydefapply_rotary_pos_emb(q,k,cos,sin):"""通过旋转 Q 和 K 向量来编码位置 旋转公式：q_rotated = q * cos + rotate_half(q) * sin 这让模型能通过点积学习相对位置 """q_embed=(q*cos)+(rotate_half(q)*sin)k_embed=(k*cos)+(rotate_half(k)*sin)returnq_embed,k_embed

分组查询注意力（GQA）

# 来自 model/llama.pyclassLlamaAttention(nn.Module):"""GQA：KV head 数量少于 Q head，节省显存"""def__init__(self,config):self.num_heads=config.num_attention_heads# 例如 32self.num_kv_heads=config.num_key_value_heads# 例如 8self.num_kv_groups=self.num_heads//self.num_kv_heads# = 4# Q 投影比 K、V 投影更大self.q_proj=nn.Linear(hidden,num_heads*head_dim)self.k_proj=nn.Linear(hidden,num_kv_heads*head_dim)# 更小！self.v_proj=nn.Linear(hidden,num_kv_heads*head_dim)

SwiGLU MLP

# 来自 model/llama.pyclassLlamaMLP(nn.Module):"""SwiGLU：output = down(silu(gate(x)) * up(x))"""defforward(self,x):gate=F.silu(self.gate_proj(x))# Swish 激活up=self.up_proj(x)returnself.down_proj(gate*up)# 门控线性单元

教学模式：边看边学

这是最有意思的特性之一！nano-vllm 包含多种教学模式，解释推理过程中发生的事情：

解说模式

提供白话解说，就像专家带着观看手术：

python -m nano_vllm.cli --model TinyLlama/TinyLlama-1.1B-Chat-v1.0\--prompt"The capital of France is"--narrate

输出示例：

═══════════════════════════════════════════════════════════════════ 推理剖析 - 教学模式 ═══════════════════════════════════════════════════════════════════ Prompt: "The capital of France is" Model: TinyLlama/TinyLlama-1.1B-Chat-v1.0 ═════ 第一幕：分词 ═════ 将 prompt 转换为模型能理解的数字... "The capital of France is" ↓ 分词器（BPE 算法） [The] [capital] [of] [France] [is] → [450, 7483, 310, 3444, 338] ═════ 第二幕：PREFILL 阶段 ═════ 模型一次性读取整个 prompt... 通过 22 层处理 5 个 token ✓ 并行计算（所有 token 一起处理） ✓ 构建 KV cache ═════ 第三幕：DECODE 阶段 ═════ 现在逐个生成 token... 步骤 1：预测第 6 个 token │ 前 5 个预测： │ Paris ████████████████████ 82.3% │ the ███ 7.1% │ located ██ 4.2% └── 采样：「Paris」（82.3%）

X-Ray 模式

显示张量形状和数学运算：

python -m nano_vllm.cli --model TinyLlama/TinyLlama-1.1B-Chat-v1.0\--prompt"Hello"--xray

仪表盘模式

实时终端界面显示进度（需要rich库）：

python -m nano_vllm.cli --model TinyLlama/TinyLlama-1.1B-Chat-v1.0\--prompt"Hello"--dashboard

交互式教程

逐步学习体验：

python -m nano_vllm.cli --tutorial

快速上手

安装

pipinstall-e.# 可选：FlashAttention（加速推理）pipinstallflash-attn --no-build-isolation

基本使用

# 单条 promptpython -m nano_vllm.cli --model TinyLlama/TinyLlama-1.1B-Chat-v1.0\--prompt"Hello, world"# 多条 prompt（连续批处理）python -m nano_vllm.cli --model TinyLlama/TinyLlama-1.1B-Chat-v1.0\--prompt"The capital of France is"\--prompt"The largest planet is"\--prompt"Python is a"# 优先级调度python -m nano_vllm.cli --model TinyLlama/TinyLlama-1.1B-Chat-v1.0\--prompt"Low priority task"--priority1\--prompt"High priority task"--priority10# 推测解码python -m nano_vllm.speculative.cli\--target-model TinyLlama/TinyLlama-1.1B-Chat-v1.0\--draft-model TinyLlama/TinyLlama-1.1B-Chat-v1.0\--prompt"The future of AI is"\--num-speculative-tokens5

Python API

fromnano_vllm.engineimportLLMEngine engine=LLMEngine(model_path="TinyLlama/TinyLlama-1.1B-Chat-v1.0",use_paged_attention=True,enable_prefix_caching=True,use_flash_attn=True,)# 单条生成output=engine.generate("What is machine learning?",max_tokens=100)# 带优先级的批量生成engine.add_request("Prompt 1",max_tokens=50,priority=1)engine.add_request("Prompt 2",max_tokens=50,priority=10)# 更高优先级outputs=engine.run_to_completion()

核心收获

构建 nano-vllm 带来的几点领悟：

显存是瓶颈：大多数 LLM 推理优化都在解决显存问题，而非计算问题。

OS 概念的应用：PagedAttention 本质上就是 KV cache 的虚拟内存。

批处理的复杂性：连续批处理远比简单地"把东西放进批次"复杂得多。

推测的威力：在昂贵的前向传播中获得多个 token 带来巨大收益。

细节决定成败：因果掩码、位置 ID、引用计数……无穷无尽的边界情况。

参考资料

• vLLM 论文
• vLLM 代码仓库
• PagedAttention 博客
• FlashAttention 论文
• 推测解码论文