nano-vllm tech analysis
nano-vllm 技术分析
项目地址:https://github.com/GeeeekExplorer/nano-vllm
基准分支:main
精确版本:812eb1c
文章目录
1. 项目概述
nano-vllm 是一个从零构建的轻量级大语言模型推理框架,以约 1,200 行纯 Python 代码实现了工业级推理引擎 vLLM 的核心功能,并在某些场景下取得了超越 vLLM 的吞吐量。项目的核心价值在于:
- 工程可读性:完整的高性能推理引擎实现精简至千行量级,非常适合学习和研究
- 性能对标:在 RTX 4070 Laptop(8GB)上运行 Qwen3-0.6B 模型,吞吐量达 1,434 token/s,对标 vLLM 的 1,362 token/s
- 生产可用特性:包含前缀缓存、张量并行、CUDA Graph、Flash Attention 等生产级优化
依赖栈(pyproject.toml):
| 依赖 | 版本要求 | 用途 |
|---|---|---|
torch |
≥2.4.0 | 深度学习框架 |
triton |
≥3.0.0 | GPU 内核编写 |
transformers |
≥4.51.0 | 分词器、模型配置加载 |
flash-attn |
最新 | 高效注意力计算 |
xxhash |
最新 | 高速哈希(前缀缓存) |
Python 版本要求: >=3.10, <3.13
2. 整体架构设计
用户接口层
┌─────────────────────────────────────────────┐
│ LLM / LLMEngine │
│ generate(prompts, sampling_params) │
└──────────────┬──────────────────────────────┘
│ 调度循环
┌──────────────▼──────────────────────────────┐
│ Scheduler(调度器) │
│ ├── 等待队列(waiting deque) │
│ ├── 运行队列(running deque) │
│ └── BlockManager(KV Cache 块内存管理) │
└──────────────┬──────────────────────────────┘
│ 序列批次
┌──────────────▼──────────────────────────────┐
│ ModelRunner(模型执行器) │
│ ├── prepare_prefill / prepare_decode │
│ ├── run_model(CUDA Graph / Eager) │
│ └── Sampler(采样器) │
└──────────────┬──────────────────────────────┘
│ NCCL / SharedMemory(多 GPU)
┌──────────────▼──────────────────────────────┐
│ Qwen3ForCausalLM(模型) │
│ ├── VocabParallelEmbedding │
│ ├── Qwen3DecoderLayer × N │
│ │ ├── RMSNorm(+残差合并) │
│ │ ├── Qwen3Attention │
│ │ │ ├── QKVParallelLinear │
│ │ │ ├── RoPE │
│ │ │ └── FlashAttention │
│ │ ├── RMSNorm(+残差合并) │
│ │ └── Qwen3MLP(SiluAndMul) │
│ └── ParallelLMHead │
└─────────────────────────────────────────────┘
代码架构层级:
nanovllm/
├── __init__.py → 导出 LLM, SamplingParams
├── llm.py → LLM(LLMEngine 的空子类)
├── config.py → Config(数据类,系统配置)
├── sampling_params.py → SamplingParams(采样参数)
├── engine/
│ ├── llm_engine.py → LLMEngine(主引擎)
│ ├── scheduler.py → Scheduler(请求调度)
│ ├── block_manager.py → BlockManager(KV Cache 管理)
│ ├── model_runner.py → ModelRunner(GPU 执行)
│ └── sequence.py → Sequence(请求序列抽象)
├── layers/
│ ├── attention.py → Attention + Triton KV 写入核
│ ├── linear.py → 张量并行线性层族
│ ├── activation.py → SiluAndMul
│ ├── layernorm.py → RMSNorm(含残差合并)
│ ├── rotary_embedding.py → RoPE
│ ├── embed_head.py → 词表并行嵌入 & LM Head
│ └── sampler.py → Gumbel-Max 采样器
├── models/
│ └── qwen3.py → Qwen3 模型实现
└── utils/
├── context.py → 推理上下文(全局单例)
└── loader.py → SafeTensors 权重加载器
nanovllm在技术实现方面需要重点关注ModelRunner的实际任务处理(3.4 模型执行器(ModelRunner))和Qwen3模型多卡并行实现部分(4. 神经网络层实现分析)的设计。
3. 核心模块深度分析
3.1 配置系统
文件: nanovllm/config.py
@dataclass
class Config:
model: str # 模型权重路径(必须是目录)
max_num_batched_tokens: int = 16384 # 单批次最大 token 数
max_num_seqs: int = 512 # 单批次最大并发序列数
max_model_len: int = 4096 # 最大上下文长度
gpu_memory_utilization: float = 0.9 # GPU 显存利用率目标
tensor_parallel_size: int = 1 # 张量并行度(1-8)
enforce_eager: bool = False # 是否禁用 CUDA Graph
kvcache_block_size: int = 256 # KV Cache 块大小(token 数)
num_kvcache_blocks: int = -1 # KV Cache 块总数(运行时计算)
hf_config: AutoConfig | None = None # HuggingFace模型配置
eos: int = -1 # 序列结束标记,在加载tokenize后读取赋值
关键约束与初始化逻辑(__post_init__):
kvcache_block_size必须是 256 的倍数,以保证 CUDA 访问对齐- hf_config通过
AutoConfig.from_pretrained自动读取 HuggingFace 模型配置 max_model_len取配置值与模型max_position_embeddings的最小值
num_kvcache_blocks 初始为 -1,由 ModelRunner.allocate_kv_cache() 在 GPU 内存探测后动态计算并回写。
该配置类在 LLMEngine 的初始化中被使用:读取未配置的权重路径,并根据 HuggingFace 模型配置自动补全其他配置项。
3.2 全局上下文(Context)
文件: nanovllm/utils/context.py
@dataclass(slots=True)
class Context:
is_prefill: bool = False # 是否处于预填充阶段,预填充阶段会进行块分配和预录制CUDA图,非预填充阶段则直接执行单步解码
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 # 预填充和解码阶段每个token对应的cache槽位映射,用于将当前step的kv写入正确的位置
context_lens: torch.Tensor | None = None # 解码阶段每个序列的上下文长度,用于注意力计算和块分配
block_tables: torch.Tensor | None = None # 预填充和解码阶段每个序列的显存块分配表,用于块稀疏注意力计算
全局上下文Context是一个在每次forward开始前设置、结束后清空的临时全局状态,它解决的是模型内部的Attention、LMHead等层在forward时需要感知当前这批请求的调度信息,具体的参数依赖可以参见nanovllm/layers/attention.py(67~74行)中对flash_attn的flash_attn_varlen_func和flash_attn_with_kvcache的入参。如果不使用全局上下文变量Context则必须将这些变量从ModelRunner中侵入模型结构的形式传递下去。
3.3 请求调度器(Scheduler)
文件: nanovllm/engine/scheduler.py
调度器实现了 vLLM 论文中提出的连续批处理(Continuous Batching)调度策略,维护两个双端队列:
waiting: deque[Sequence] # 待 prefill 的请求
running: deque[Sequence] # 正在 decode 的请求
调度决策流程(schedule() → tuple[list[Sequence], bool]):
schedule() 调用时:
阶段一:Prefill 优先
for seq in waiting:
if 超出 max_num_batched_tokens 或无法分配 KV Cache块:
break
→ 分配 KV Cache(含前缀缓存命中检测)
→ 移入 running 队列
→ 加入当次批次
若有 prefill 序列 → 返回 (prefill_seqs, is_prefill=True)
阶段二:Decode
for seq in running:
while 无法追加新 KV Cache 块:
抢占(preempt)running 末尾的序列
→ 追加 KV Cache 槽位
→ 加入当次批次
→ 返回 (decode_seqs, is_prefill=False)
设计要点:
- Prefill 优先策略:每个调度周期只做一件事,要么 prefill,要么 decode,不混合。静态场景下,初次schedule调用时会处理用户的prompt将其写入KV cache中,再次调用schedule时会执行decode针对KV cache中的prompt和已生成的前文(包括最新生成的上一个token)生成最新的一个token。实际在线服务计算是一个动态的场景,会不断有sequence进入等待队列,这时schedule方法会在显存足够条件下优先执行等待队列中的prefill。
- 抢占机制(Preemption):当 decode 阶段内存不足时,将
running队列中尚未被调度处理的序列强制抢占(running.pop(),即 deque 尾部的序列):释放其 KV Cache 块,状态重置为WAITING,重新插入waiting队列头部。这保证了系统不会死锁,但被抢占的序列需要重新 prefill(当前实现不保存 KV Cache)。 - 结束后处理(
postprocess):对每条序列 append 新生成的 token,检查是否遇到 EOS 或达到max_tokens上限,如是则释放 KV Cache 并从 running 队列移除。
3.4 KV Cache 块管理器(BlockManager)
文件: nanovllm/engine/block_manager.py
这是整个系统中实现最为精妙的模块,实现了带前缀缓存的分页 KV Cache 管理。
数据结构:
class Block:
block_id: int # 物理块 ID(GPU 显存中的索引)
ref_count: int # 引用计数(共享块的多序列共享)
hash: int # 内容哈希(-1 表示未完成/脏块)
token_ids: list[int] # 块内 token ID(用于哈希验证)
class BlockManager:
blocks: list[Block] # 所有物理块
hash_to_block_id: dict[int, int] # 内容哈希 → 物理块 ID(前缀缓存索引)
free_block_ids: deque[int] # 空闲块队列
used_block_ids: set[int] # 已使用块集合
哈希计算(链式前缀哈希):
@classmethod
def compute_hash(cls, token_ids: list[int], prefix: int = -1):
h = xxhash.xxh64()
if prefix != -1:
h.update(prefix.to_bytes(8, "little")) # 将前一块哈希混入
h.update(np.array(token_ids).tobytes())
return h.intdigest()
前缀哈希形成哈希链:块 i 的哈希 = hash(块 i 的 tokens + 块 i-1 的哈希)。这确保了相同前缀内容的序列能命中同一物理块,即使序列本身不同,也能共享前缀的 KV Cache。默认配置的block大小是256,意味着命中同一物理块至少需要两次请求前256个token完全一致才行,这种情况主要针对的是固定的超过256长度的system prompt或者用户侧由于网络原因等的重复请求。
分配逻辑(allocate):
for 每个块 i in [0, num_blocks):
计算该块哈希(仅对完整块,即 len == block_size 时)
查 hash_to_block_id:
命中 且 token_ids 匹配 → 前缀缓存命中
→ 增加 ref_count,复用该物理块
→ num_cached_tokens += block_size
未命中 → 从 free_block_ids 取新块分配
将物理块 ID 追加到 seq.block_table
追加逻辑(may_append):
当 len(seq) % block_size == 1(刚进入新块):
→ 分配一个新物理块(旧块已满,其哈希已在之前被计算并注册)
当 len(seq) % block_size == 0(块刚好写满):
→ 计算并注册当前块的哈希,写入 hash_to_block_id
当 len(seq) % block_size 是其他值:
→ 块尚未写满,不做操作
判断能否追加(can_append):
def can_append(self, seq: Sequence) -> bool:
return len(self.free_block_ids) >= (len(seq) % self.block_size == 1)
仅在序列长度恰好是 block_size 的倍数加一(即刚进入新块的第一个 token)时需要申请新块,其余情况复用当前尾块。
3.5 模型执行器(ModelRunner)
文件: nanovllm/engine/model_runner.py
ModelRunner 是推理系统与 GPU 硬件交互的核心,负责所有 GPU 操作。
初始化流程:
1. dist.init_process_group("nccl", ...) # 初始化分布式通信
2. torch.cuda.set_device(rank) # 绑定 GPU
3. torch.set_default_dtype(hf_config.torch_dtype) # 设置精度(如 bfloat16)
4. torch.set_default_device("cuda") # 模型参数直接在 GPU 上分配
5. self.model = Qwen3ForCausalLM(hf_config) # 加载配置构建Qwen3模型
6. load_model(self.model, config.model) # Qwen3权重加载
7. warmup_model() # 空跑以探测峰值显存
8. allocate_kv_cache() # 根据剩余显存分配 KV Cache
9. capture_cudagraph() # 预捕获 CUDA Graph(若未禁用)
另外在初始化时还会创建共享内存块,如下所示:
if self.world_size > 1:
if rank == 0:
# 针对Rank0的主进程,先创建一块名为nanovllm的大小为1MB的共享内存,再等待所有工作进程执行到barrier
self.shm = SharedMemory(name="nanovllm", create=True, size=2**20)
dist.barrier()
else:
# 工作进程等待主进程完成创建执行到barrier后,连接到已经存在的共享内存
dist.barrier()
self.shm = SharedMemory(name="nanovllm")
self.loop() #工作进程的主循环,持续等待主进程的调用指令并读取执行对应的方法,直到接收到exit指令时退出循环
模型参数加载:
主要调用的是nanovllm/utils/loader.py中的load_model方法,packed_modules_mapping参见nanovllm/models/qwen3.py(187~193行)是作者定义的{原始Qwen3模型参数名 : (Qwen3并行处理后模型参数名 , 存储合并矩阵的子段名)}这样一个数据结构,这是为了适应通信优化后将Q/K/V计算的线性层合为一个大的线性层的参数拼接,同理还有解码的mlp中的升维降维线性层的合并,这能在推理时把多次小矩阵乘法合并成一次,这样在张量并行时只需要对一个矩阵做多卡的拆分,访存更友好。
def load_model(model: nn.Module, path: str):
'''
模型参数加载器:从指定路径加载模型权重文件,并将权重分配到模型的参数中。
根据模块是否在packed_modules_mapping键值中,判断是否需要进行权重重组和分片加载。对于没有packed_modules_mapping的模块,尝试加载weight_loader方法,若没有直接使用默认的权重加载器将权重加载到参数中
对于有packed_modules_mapping的模块则需要进行分片加载,根据映射关系将权重文件中的参数名转换为模型中对应的参数名,并调用相应的weight_loader方法进行加载。对于分片加载的参数,根据tp_rank和tp_size计算出当前进程需要加载的权重切片,并进行加载。
'''
packed_modules_mapping = getattr(model, "packed_modules_mapping", {})
for file in glob(os.path.join(path, "*.safetensors")):
with safe_open(file, "pt", "cpu") as f:
for weight_name in f.keys():
for k in packed_modules_mapping:
if k in weight_name:
v, shard_id = packed_modules_mapping[k]
param_name = weight_name.replace(k, v)
param = model.get_parameter(param_name)
weight_loader = getattr(param, "weight_loader")
weight_loader(param, f.get_tensor(weight_name), shard_id)
break
else:
param = model.get_parameter(weight_name)
weight_loader = getattr(param, "weight_loader", default_weight_loader)
weight_loader(param, f.get_tensor(weight_name))
显存探测与 KV Cache 分配:
在ModelRunner的初始化中会调用warmup_model和allocate_kv_cache两个方法,其中warmup_model方法会构造一个最大化负载的虚拟请求,以此来测算出模型在输入满负载条件下的预填充需要占用的显存。
def warmup_model(self):
'''
让CUDA完成JIT编译,构造一个最大化负载的虚拟请求批次,以触发最坏情况下的分配和计算路径
1、预热消除首请求延迟,确保第一个真实用户请求不会因为CUDA初始化或内核编译而变慢;
2、精确显存规划,通过模拟最大负载,提前分配并验证显存需求,防止在生产环境中因显存碎片或突发大请求导致OOM;
3、稳定性验证,在正式对客前,验证模型在极端配置下能否正常运行
'''
torch.cuda.empty_cache() # 释放GPU上未使用的缓存内存
torch.cuda.reset_peak_memory_stats() # 将当前 GPU 设备上记录的“自上次重置以来使用的最大显存量”清零
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() # 预热结束后,释放那些在预热过程中分配但随后不再需要的临时显存碎片,模型权重和KV cache会被保留或被管理
allocate_kv_cache方法通过公式计算出可用显存,然后根据模型参数计算与配置的cache块大小计算出单个cache所需的显存,二者相除得到可以分配的KV cache块的数量,预分配出这部分的KV cache块的连续显存。需要注意的是配置的KV cache块的大小kvcache_block_size为256,通常配置的都是远小于模型能输入序列的最大长度max_model_len=4096的,一个较长的序列的KV cache可能会由多个对应的cache块进行存储。如果分配的cache块太大,需要的显存呈几何量级增长,会导致这种输入序列独占cache块的形式有较多的显存浪费,块太小则会导致管理开销增大。
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)
# 每个KV块的字节数:2(K+V) × 层数 × block_size(token数) × kv_head数 × head_dim × 每元字节
# 注意计算kv cache大小时候使用的预估序列长度block_size(256)远小于模型可输入的最大序列长(4096),这是因为kv cache的分配是以块为单位的,过大的块大小会导致显存浪费和分配失败,而过小的块小会增加管理开销和降低效率
block_bytes = 2 * hf_config.num_hidden_layers * self.block_size * num_kv_heads *head_dim * hf_config.dtype.itemsize
# 可用显存 = 总显存 × 预设利用率 - 非本进程占用(used-current) - 模型权重(current) - 推理激活峰值(peak-current)
# 其中 used=总占用(含其他进程), peak=warmup峰值(含模型权重+激活), current=当前分配(模型权重); +current是为了从used和peak中扣除已包含的模型权重部分,避免重复减
config.num_kvcache_blocks = int(total * config.gpu_memory_utilization - used -peak + current) // block_bytes
assert config.num_kvcache_blocks > 0
# 预分配连续显存,避免后续推理过程中分配KV cache时频繁分配调用造成的分配延时和显存碎片问题
self.kv_cache = torch.empty(2, hf_config.num_hidden_layers, config.num_kvcache_blocks, self.block_size, num_kv_heads, head_dim)
通过先 warmup(空跑模型)再测量峰值显存,精确计算可分配的 KV Cache 块数,避免了静态保留的浪费。关于KV cache显存块的实际应用参考4.1 注意力机制(Attention)。
模型录制与推理:
这部分主要依赖torch.cuda.CUDAGraph(),通过虚拟一系列batch size的模型负载,依次通过模型调用这些虚拟的负载录制下CUDA内核执行序列,来消除Python解释器的开销、CUDA API多次调用的开销和大部分CPU-GPU同步等待。后续需要再次调用时只需要将对应batch size(找寻最接近的比当前处理批次大的录制size)的录制进行重放即可,将动态的由CPU调度的操作序列转化为静态的、预编译的执行流。
@torch.inference_mode()
def capture_cudagraph(self):
'''
预录制不同batchSize的decode CUDA kernel序列,将CPU的多次调度直接合并为单次提交给GPU
将运行时的CUDA API调用序列记录为静态执行图,以实现底层调度优化
'''
self.graph_bs = [1, 2, 4, 8] + list(range(16, max_bs + 1, 16))
self.graphs = {}
self.graph_pool = None
for bs in reversed(self.graph_bs):
graph = torch.cuda.CUDAGraph()
set_context(False, slot_mapping=slot_mapping[:bs], context_lens=context_len[: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()
@torch.inference_mode()
def run_model(self, input_ids, positions, is_prefill):
if is_prefill or enforce_eager or batch_size > 512:
return self.model.compute_logits(self.model(input_ids, positions))
else:
# 使用 CUDA Graph 加速 decode
graph = self.graphs[next(x for x in self.graph_bs if x >= batch_size)]
# 向 graph_vars 中填充数据 → replay → 返回输出
graph.replay() # CPU将预编译的指令流发送到GPU执行,跳过Python/C++运行时、CUDA Driver和动态图构建的开销
return self.model.compute_logits(graph_vars["outputs"][:bs])
预填充全局参数的计算:
预填充部分的参数计算代码有助于进一步理解全局变量参数的作用与KV cache块管理方式。计算过程中start_block是对应的是存储的cache块的起始块序号,对应的end_block是存储cache块的结束块序号,计算出的slot_mapping变量是预填充时token的kv数据存储在KV cache中的物理槽位。slot_mapping将本次需要处理的所有序列的对应槽位位置拼接起来供注意力计算时指导计算得到的最新kv存入对应的槽位。
def prepare_prefill(self, seqs: list[Sequence]):
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:
seqlen = len(seq)
start = min(seq.num_cached_tokens, seqlen - 1) # 已经缓存的token数量,作为本轮预填充的起始位置,确保每个token只被预填充一次
seqlen_q = seq.num_scheduled_tokens # 本轮需要做attention的token数量
seqlen_k = seqlen
end = start + seqlen_q
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)
# 在模型的warmup阶段,序列未设置block_table,则不会执行循环中后续的代码
if not seq.block_table: # warmup
continue
start_block = start // self.block_size # 第一个需要写入的块索引,block_size为预估的cache块能容纳的最大token数量,确保每个块被正确映射和管理
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 # block_table为序列写入占块的索引(0~num_kvcache_blocks-1),乘以块大小得到该块在KV cache中的起始位置
if i == start_block:
slot_start += start % self.block_size # 首位置需要考虑已经缓存的token造成的偏移
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)) # slot_mapping为分配后token需要写入cache中的实际物理位置,为后续的cache写入做准备
3.6 序列管理(Sequence)
文件: nanovllm/engine/sequence.py
Sequence 是单条推理请求的完整状态机:
class Sequence:
seq_id: int # 全局唯一 ID(类变量计数器)
status: SequenceStatus # WAITING / RUNNING / FINISHED
token_ids: list[int] # prompt + completion token 列表
num_tokens: int # 当前总 token 数
num_prompt_tokens: int # prompt 的 token 数(固定)
num_cached_tokens: int # 已通过前缀缓存命中的 token 数
block_table: list[int] # 物理 KV Cache 块 ID 列表
temperature: float # 采样温度
max_tokens: int # 最大生成 token 数
ignore_eos: bool # 是否忽略 EOS
关键属性计算:
num_completion_tokens = num_tokens - num_prompt_tokens
num_blocks = ceil(num_tokens / block_size)
last_block_num_tokens = num_tokens - (num_blocks - 1) * block_size
num_cached_blocks = num_cached_tokens // block_size
序列化优化(__getstate__ / __setstate__):
Sequence 实现了自定义序列化接口,在多进程(张量并行)传递时极度压缩数据量:
def __getstate__(self):
return (
num_tokens, num_prompt_tokens, num_cached_tokens, block_table,
token_ids if num_completion_tokens == 0 else last_token # prefill 传全量,decode 只传最后一个 token
)
这是一个精妙的优化:prefill 阶段需要传递所有 token ID;decode 阶段每步只需传递上一步生成的最后一个 token,大幅减少进程间数据拷贝量。
4. 神经网络层实现分析
该部分的模块依据是否进行参数按卡分散分为两类,模块参数分散的一类是需要重写模块的weight_loader方法,加载当前进程的GPU对应的切分部分参数。
4.1 注意力机制(Attention)
文件: nanovllm/layers/attention.py
Triton KV Cache 写入核(store_kvcache_kernel):
@triton.jit
def store_kvcache_kernel(key_ptr, key_stride, value_ptr, value_stride,
k_cache_ptr, v_cache_ptr, slot_mapping_ptr, D):
idx = tl.program_id(0) # 每个线程块处理一个 token
slot = tl.load(slot_mapping_ptr + idx) # 目标物理槽位
if slot == -1: return # warmup 阶段跳过
# 从连续 key/value 张量读取,写入分页 kv_cache 的指定槽位
key = tl.load(key_ptr + idx * key_stride + tl.arange(0, D))
tl.store(k_cache_ptr + slot * D + tl.arange(0, D), key)
# value 同理
这个自定义 Triton 核解决了非连续内存写入问题:Flash Attention 计算产生连续排列的 K/V,但需要写入到分页 KV Cache 中任意散乱的物理槽位。
Attention 前向传播分支:
def forward(self, q, k, v):
# 1. 写入 KV Cache
store_kvcache(k, v, k_cache, v_cache, context.slot_mapping)
if context.is_prefill:
if context.block_tables is not None: # 有前缀缓存命中
k, v = k_cache, v_cache # 使用缓存中的 K/V
o = flash_attn_varlen_func( # 变长序列 prefill
q, k, v,
cu_seqlens_q=..., cu_seqlens_k=...,
causal=True, block_table=context.block_tables
)
else: # decode
o = flash_attn_with_kvcache( # 单 token decode
q.unsqueeze(1), k_cache, v_cache,
cache_seqlens=context.context_lens,
block_table=context.block_tables,
causal=True
)
两个 Flash Attention API 的选择反映了 prefill 和 decode 的根本差异:
flash_attn_varlen_func:处理不等长序列的并行 prefillflash_attn_with_kvcache:高效的单步 decode,直接从分页 KV Cache 读取历史
抽取的Attention模块只需要计算Q/K/V向量的注意力输出,没有实际的参数,故无需进行参数的拆分。
4.2 张量并行线性层(Linear)
文件: nanovllm/layers/linear.py
项目实现了完整的张量并行线性层族,均继承自 LinearBase:
| 类名 | 并行策略 | AllReduce | 使用场景 |
|---|---|---|---|
ReplicatedLinear |
无(复制) | 否 | 不需要并行的层 |
ColumnParallelLinear |
列切分输出维 | 否 | FFN gate/up,attn q/k/v |
MergedColumnParallelLinear |
合并多个列切分 | 否 | gate_proj + up_proj 合并 |
QKVParallelLinear |
Q/K/V 分别切分列 | 否 | QKV 融合投影 |
RowParallelLinear |
行切分输入维 | 是 | FFN down,attn output |
权重加载机制(weight_loader):
每个参数对象都挂载了 weight_loader 方法,加载时自动处理 TP 分片:
class ColumnParallelLinear:
def weight_loader(self, param, loaded_weight):
shard_size = param.data.size(self.tp_dim) # 每个 rank 的分片大小
start_idx = self.tp_rank * shard_size
loaded_weight = loaded_weight.narrow(self.tp_dim, start_idx, shard_size)
param.data.copy_(loaded_weight)
RowParallelLinear 的 AllReduce:
def forward(self, x):
y = F.linear(x, self.weight, self.bias if self.tp_rank == 0 else None)
if self.tp_size > 1:
dist.all_reduce(y) # 聚合所有 rank 的部分和
return y
bias 只在 rank 0 添加,AllReduce 后自动正确(0 + 0 + … + bias = bias)。
class Qwen3Attention(nn.Module):
def forward(
self,
positions: torch.Tensor,
hidden_states: torch.Tensor,
) -> torch.Tensor:
qkv = self.qkv_proj(hidden_states)
q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
q = q.view(-1, self.num_heads, self.head_dim)
k = k.view(-1, self.num_kv_heads, self.head_dim)
v = v.view(-1, self.num_kv_heads, self.head_dim)
if not self.qkv_bias:
q = self.q_norm(q)
k = self.k_norm(k)
q, k = self.rotary_emb(positions, q, k)
o = self.attn(q, k, v)
output = self.o_proj(o.flatten(1, -1))
return output
在大模型架构中,通常是一个列并行线性层后下一个线性层使用行并行线性层,如上面示例中qkv_proj即为列并行线性层,而o_proj是行并行线性层。目的是列并行实际每张GPU上得到的仅仅是按特征维度的输出的部分切分,对特征维度的切分可以视作对注意力头维度的切分,多头机制天然是可以切分的,其后的行并行则是将原本分散在各卡上的输入能直接叠加得到结果,这样的设计可以将多GPU间的通信开销最小化。
4.3 归一化层(RMSNorm)
文件: nanovllm/layers/layernorm.py
项目实现了两种模式的 RMSNorm,均使用 @torch.compile 加速:
普通 RMSNorm(rms_forward):
@torch.compile
def rms_forward(self, x):
orig_dtype = x.dtype
x = x.float() # 升精度计算
var = x.pow(2).mean(dim=-1, keepdim=True)
x.mul_(torch.rsqrt(var + self.eps)) # 原地操作节省内存
return x.to(orig_dtype).mul_(self.weight)
残差合并 RMSNorm(add_rms_forward):
@torch.compile
def add_rms_forward(self, x, residual):
orig_dtype = x.dtype
x = x.float().add_(residual.float()) # 先加残差
residual = x.to(orig_dtype) # 保存残差供下层使用
var = x.pow(2).mean(dim=-1, keepdim=True)
x.mul_(torch.rsqrt(var + self.eps))
return x.to(orig_dtype).mul_(self.weight), residual
这是 vLLM 采用的经典优化:将残差加法与归一化融合,避免在 Decoder Layer 中额外存储中间张量,减少内存带宽消耗。
在 Qwen3DecoderLayer.forward 中,residual 作为独立变量在层间传递:
def forward(self, positions, hidden_states, residual):
hidden_states, residual = self.input_layernorm(hidden_states, residual)
hidden_states = self.self_attn(positions, hidden_states)
hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
hidden_states = self.mlp(hidden_states)
return hidden_states, residual
4.4 旋转位置编码(RoPE)
文件: nanovllm/layers/rotary_embedding.py
RoPE 实现采用预计算缓存策略:
def __init__(self, head_size, rotary_dim, max_position_embeddings, base):
inv_freq = 1.0 / (base ** (torch.arange(0, rotary_dim, 2) / rotary_dim))
t = torch.arange(max_position_embeddings)
freqs = torch.einsum("i,j -> ij", t, inv_freq)
cos = freqs.cos()
sin = freqs.sin()
cache = torch.cat((cos, sin), dim=-1).unsqueeze_(1)
self.register_buffer("cos_sin_cache", cache) # [max_pos, 1, head_size]
推理时只需按位置索引查表:
@torch.compile
def forward(self, positions, query, key):
cos_sin = self.cos_sin_cache[positions] # 查表
cos, sin = cos_sin.chunk(2, dim=-1)
query = apply_rotary_emb(query, cos, sin)
key = apply_rotary_emb(key, cos, sin)
return query, key
注意:当前 get_rope 函数使用 @lru_cache(1) 且断言 rope_scaling is None,意味着仅支持标准 RoPE,不支持 YaRN、LongRoPE 等扩展版本。
4.5 词表并行嵌入与 LM Head
文件: nanovllm/layers/embed_head.py
VocabParallelEmbedding:将词表均匀切分到各 GPU:
# 每个 rank 负责词表的一段 [vocab_start_idx, vocab_end_idx)
mask = (x >= self.vocab_start_idx) & (x < self.vocab_end_idx)
x = mask * (x - self.vocab_start_idx) # 越界 token 映射到 0
y = F.embedding(x, self.weight)
y = mask.unsqueeze(1) * y # 越界位置清零
dist.all_reduce(y) # 汇总各段结果
ParallelLMHead(继承 VocabParallelEmbedding):
def forward(self, x):
if context.is_prefill:
last_indices = context.cu_seqlens_q[1:] - 1
x = x[last_indices].contiguous() # 只取每条序列的最后一个 token
logits = F.linear(x, self.weight) # 复用 embedding 权重转置(当 tie_word_embeddings=True)
if self.tp_size > 1:
all_logits = [torch.empty_like(logits) for _ in range(self.tp_size)] if rank == 0 else None
dist.gather(logits, all_logits, 0)
logits = torch.cat(all_logits, -1) if rank == 0 else None
return logits
prefill 阶段通过 cu_seqlens_q 索引精确提取每条序列最后一个 token 的 hidden state,避免对整个序列计算 logits。
4.6 Token 采样器(Sampler)
文件: nanovllm/layers/sampler.py
@torch.compile
def forward(self, logits, temperatures):
logits = logits.float().div_(temperatures.unsqueeze(1)) # 温度缩放
probs = torch.softmax(logits, dim=-1)
# Gumbel-Max 采样:等价于 categorical 采样,但完全并行化
sample_tokens = probs.div_(
torch.empty_like(probs).exponential_(1).clamp_min_(1e-10)
).argmax(dim=-1)
return sample_tokens
Gumbel-Max 采样原理:
根据逆变换采样法,构建Gumbel分布与均匀分布时间的关系,设 u ∼ Uniform ( 0 , 1 ) u \sim \text{Uniform}(0,1) u∼Uniform(0,1),则 − ln ( − ln u ) ∼ Gumbel ( 0 , 1 ) -\ln(-\ln u) \sim \text{Gumbel}(0,1) −ln(−lnu)∼Gumbel(0,1)。又由于Gumbel分布的特性, A r g M a x ( ln P + G u m b e l N o i s e ) ArgMax(\ln P + Gumbel Noise) ArgMax(lnP+GumbelNoise)等价于从 P P P中进行采样,即计算 A r g M a x ( ln ( P − ln u ) ) ArgMax(\ln (\frac{P}{-\ln u})) ArgMax(ln(−lnuP)).
由于 Exponential ( 1 ) = − ln ( Uniform ( 0 , 1 ) ) \text{Exponential}(1) = -\ln(\text{Uniform}(0,1)) Exponential(1)=−ln(Uniform(0,1))和log函数在其定义域上的单调性,因此:
sample = arg max i ( p i Exponential i ) \text{sample} = \argmax_i \left(\frac{p_i}{\text{Exponential}_i}\right) sample=iargmax(Exponentialipi)
与标准的多项式采样等价,但可以通过 argmax 完全向量化并用 @torch.compile 融合,避免 CPU-GPU 同步。
5. 模型架构(Qwen3)
文件: nanovllm/models/qwen3.py
项目目前仅支持 Qwen3 系列模型,架构遵循标准 Decoder-Only Transformer:
层结构(Qwen3DecoderLayer):
输入
↓
RMSNorm(+ 残差合并)
↓
Qwen3Attention
├── QKVParallelLinear(合并 Q/K/V 投影)
├── 可选 q_norm / k_norm(RMSNorm,无 bias 时启用)← Qwen3 特有
├── RoPE
├── Flash Attention(prefill 或 decode 路径)
└── RowParallelLinear(输出投影)
↓
RMSNorm(+ 残差合并)
↓
Qwen3MLP
├── MergedColumnParallelLinear(gate + up 融合)
├── SiluAndMul(激活)
└── RowParallelLinear(down 投影)
↓
输出
Qwen3 特有设计:
-
QK Norm:当
attention_bias=False(即qkv_bias=False)时,对 Q 和 K 各应用独立的 RMSNorm,有助于稳定训练时的梯度:if not self.qkv_bias: self.q_norm = RMSNorm(self.head_dim, eps=rms_norm_eps) self.k_norm = RMSNorm(self.head_dim, eps=rms_norm_eps) -
权重绑定(
tie_word_embeddings):LM Head 与 Token Embedding 共享权重矩阵:if config.tie_word_embeddings: self.lm_head.weight.data = self.model.embed_tokens.weight.data -
Packed Modules Mapping:定义了 HuggingFace 原始权重名到 nano-vllm 合并权重名的映射:
packed_modules_mapping = { "q_proj": ("qkv_proj", "q"), "k_proj": ("qkv_proj", "k"), "v_proj": ("qkv_proj", "v"), "gate_proj": ("gate_up_proj", 0), "up_proj": ("gate_up_proj", 1), }这使得原本独立的 q/k/v 权重文件能够在加载时自动合并为
qkv_proj单一权重,减少矩阵乘法次数。
6. 关键优化技术综述
6.1 分页键值缓存(Paged KV Cache)
原理: 借鉴操作系统虚拟内存的分页思想,将 KV Cache 划分为固定大小的物理块(默认 256 tokens/块)。每条序列通过 block_table(类似页表)维护物理块 ID 列表,实现:
- 消除内存碎片:不同长度的序列共享物理块池,无需预留连续空间
- 按需分配:序列生成到需要新块时才分配,不预留未来空间
- 支持并发:不同序列的物理块可以任意交错排列
实现细节:
block_size = 256(可配置,需是 256 的倍数,源于 flash_attn_with_kvcache 对 paged attention 中 block_size 的底层对齐要求)- KV Cache 张量形状:
[2, num_layers, num_blocks, block_size, num_kv_heads, head_dim] - 通过
slot_mapping和自定义 Triton 核实现散射写入
6.2 前缀缓存(Prefix Caching)
原理: 对已填满的 KV Cache 块计算内容哈希,维护 hash → block_id 映射。新请求分配时,若发现哈希命中且内容匹配,直接复用已有物理块,无需重新计算注意力。
哈希链设计:
块 0 哈希 = xxh64(tokens[0:256])
块 1 哈希 = xxh64(block_0_hash || tokens[256:512])
块 2 哈希 = xxh64(block_1_hash || tokens[512:768])
...
前缀哈希的链式结构确保了哈希的全局唯一性:不同前缀的相同后缀块具有不同哈希,不会错误命中。
缓存命中后的处理:
prefill 阶段,若 cu_seqlens_k[-1] > cu_seqlens_q[-1](K 累积序列长度大于 Q 累积序列长度),说明有前缀缓存命中,此时:
- Attention 使用
block_tables参数从 KV Cache 读取历史 - 只对未命中的 tokens 计算 Attention(大幅减少计算量)
引用计数: 被多序列共享的块通过 ref_count 管理生命周期,只有 ref_count 降为 0 才真正释放。
6.3 连续批处理(Continuous Batching)
与传统静态批处理(等待所有序列完成后再处理下一批)不同,连续批处理在每个推理步骤后动态调整批次组成,允许:
- 新请求在任意步骤加入(prefill)
- 已完成请求立即退出并释放资源
nano-vllm 采用严格的两阶段策略:每次调度要么纯 prefill,要么纯 decode。这虽然不如 chunked-prefill 灵活,但实现更简单,且 decode 批次可以充分利用 CUDA Graph。
6.4 张量并行(Tensor Parallelism)
进程管理:
- 使用
torch.multiprocessing.spawn启动tensor_parallel_size - 1个子进程 - Rank 0 与主进程合并(避免额外进程开销)
- 通过 NCCL 后端进行 AllReduce 和 Gather 通信
通信点分析:
| 操作 | 通信类型 | 位置 |
|---|---|---|
RowParallelLinear.forward |
dist.all_reduce |
每个 attn/FFN 输出 |
VocabParallelEmbedding.forward |
dist.all_reduce |
Embedding 查表 |
ParallelLMHead.forward |
dist.gather |
最终 logits 汇聚到 rank 0 |
TP 度限制:
- 最大 8 卡(
assert 1 <= tensor_parallel_size <= 8) - Q 头数和 KV 头数均需整除 TP size
6.5 CUDA Graph 捕获
CUDA Graph 消除了 decode 阶段的 CPU-GPU 调度开销(kernel launch overhead),对小批量 decode 尤为有效。
捕获批次规格:
self.graph_bs = [1, 2, 4, 8] + list(range(16, max_bs + 1, 16))
# 示例:[1, 2, 4, 8, 16, 32, ..., 512]
共享内存池(self.graph_pool):
for bs in reversed(self.graph_bs): # 从大到小捕获
graph = torch.cuda.CUDAGraph()
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() # 首次捕获后获取共享内存池
从大到小捕获,并让所有 graph 共享同一个 CUDA 内存池,避免为每个 batch size 单独申请显存。
运行时 Padding:
实际 batch size 取最小满足条件的预捕获规格:
graph = self.graphs[next(x for x in self.graph_bs if x >= bs)]
graph_vars["input_ids"][:bs] = input_ids # 只填充实际数据
超出实际大小的部分(padding)不影响结果,因为输出只取 [:bs] 切片。
回退策略: batch_size > 512 时自动回退到 Eager 模式。
6.6 torch.compile 即时编译
以下模块使用 @torch.compile 编译:
| 模块 | 编译的函数 | 加速效果 |
|---|---|---|
SiluAndMul |
forward |
算子融合,减少中间张量 |
RMSNorm |
rms_forward、add_rms_forward |
融合归一化+残差计算 |
RotaryEmbedding |
forward |
融合查表+旋转运算 |
Sampler |
forward |
融合温度缩放+softmax+Gumbel 采样 |
torch.compile 使用 TorchInductor 后端,将多个 PyTorch 算子编译为单个 CUDA/Triton kernel,减少内存读写次数(kernel fusion)。
6.7 Flash Attention 集成
项目使用 flash-attn 库提供的两个接口:
Prefill: flash_attn_varlen_func
- 支持可变长度序列(通过
cu_seqlens传递累积序列长度) - 支持
block_table参数(前缀缓存命中时从 KV Cache 读取) - 因果掩码(
causal=True)
Decode: flash_attn_with_kvcache
- 专为单步 decode 优化,直接操作 paged KV Cache
cache_seqlens指定每条序列的历史长度block_table提供物理块到逻辑位置的映射
7. 进程间通信机制
文件: nanovllm/engine/model_runner.py(SharedMemory 相关)
多 GPU 张量并行时,需要将主进程的调度决策传递给 Worker 进程(rank > 0)。nano-vllm 使用 multiprocessing.shared_memory.SharedMemory 而非 pickle over pipe,以避免序列化瓶颈:
Rank 0(主进程) Rank 1,2,...(Worker 进程)
│ │
│ write_shm(method_name, *args) │
│ 1. pickle.dumps([method, args]) │
│ 2. 写入共享内存(前4字节=长度) │
│ 3. event.set() ───────────────────►│ event.wait()
│ │ read_shm()
│ │ 1. 读取长度
│ │ 2. pickle.loads(shm.buf)
│ │ event.clear()
│ call("run", seqs, is_prefill) ─────►│ call("run", seqs, is_prefill)
│ 执行本地 model.forward │ 执行本地 model.forward
│ NCCL AllReduce ↔ RowParallelLinear
│ NCCL AllReduce ↔ VocabParallelEmbedding
│ NCCL Gather ↔ ParallelLMHead
│ 接收 logits,采样 token_ids │ logits=None(不参与采样)
共享内存大小: 固定 1MB(size=2**20),对于传递 Sequence 的元数据(已做精简序列化)已绰绰有余。
同步机制: 每个 Worker 有一个 multiprocessing.Event,Rank 0 写完数据后 event.set() 通知所有 Worker,Worker 读取后 event.clear()。
8. 推理流水线全链路解析
以一次完整的 llm.generate(prompts, sampling_params) 调用为例,追踪请求从输入到输出的完整路径:
用户调用 generate(["Hello, world!"], SamplingParams(temperature=0.6, max_tokens=64))
│
├─ tokenizer.encode("Hello, world!") → [token_ids...]
├─ Sequence(token_ids, sampling_params) 构建序列对象
└─ scheduler.add(seq) → 加入 waiting 队列
while not is_finished():
│
├─ scheduler.schedule() ← 调度决策
│ │
│ ├─ [首次] Prefill:
│ │ ├─ block_manager.allocate(seq) ← 分配 KV Cache 块(尝试前缀缓存命中)
│ │ └─ 返回 ([seq], is_prefill=True)
│ │
│ └─ [后续] Decode:
│ ├─ block_manager.may_append(seq) ← 按需分配新 KV Cache 块
│ └─ 返回 ([seq, ...], is_prefill=False)
│
├─ model_runner.call("run", seqs, is_prefill)
│ │
│ ├─ [TP > 1] write_shm("run", seqs, is_prefill) → 通知 Workers
│ │
│ ├─ prepare_prefill(seqs) 或 prepare_decode(seqs)
│ │ ├─ 构建 input_ids, positions 张量(pin_memory → CUDA 非阻塞)
│ │ ├─ 构建 cu_seqlens_q/k, slot_mapping, block_tables
│ │ └─ set_context(...) ← 设置全局推理上下文
│ │
│ ├─ run_model(input_ids, positions, is_prefill)
│ │ ├─ [Prefill/Eager] model(input_ids, positions)
│ │ │ ├─ VocabParallelEmbedding(input_ids)
│ │ │ ├─ for layer in layers:
│ │ │ │ ├─ RMSNorm(hidden_states, residual)
│ │ │ │ ├─ Qwen3Attention:
│ │ │ │ │ ├─ QKVParallelLinear → q, k, v
│ │ │ │ │ ├─ [可选] q_norm, k_norm
│ │ │ │ │ ├─ RoPE(positions, q, k)
│ │ │ │ │ ├─ Attention.forward(q, k, v)
│ │ │ │ │ │ ├─ store_kvcache[Triton] → 写入 KV Cache
│ │ │ │ │ │ └─ flash_attn_varlen_func / flash_attn_with_kvcache
│ │ │ │ │ └─ RowParallelLinear → AllReduce
│ │ │ │ ├─ RMSNorm(hidden_states, residual)
│ │ │ │ └─ Qwen3MLP:
│ │ │ │ ├─ MergedColumnParallelLinear → gate_up
│ │ │ │ ├─ SiluAndMul[compiled]
│ │ │ │ └─ RowParallelLinear → AllReduce
│ │ │ └─ RMSNorm(hidden_states, residual)
│ │ │
│ │ └─ model.compute_logits(hidden_states)
│ │ ├─ ParallelLMHead: [Prefill] 取最后一个 token → logits
│ │ └─ [TP > 1] Gather 到 rank 0
│ │
│ ├─ sampler(logits, temperatures) ← Gumbel-Max 采样,仅 rank 0 执行
│ └─ reset_context()
│
├─ scheduler.postprocess(seqs, token_ids)
│ ├─ seq.append_token(token_id)
│ └─ if EOS or max_tokens: seq.status = FINISHED,释放 KV Cache
│
└─ 收集 FINISHED 序列的输出
outputs = tokenizer.decode(token_ids) ← 解码 token ID 为文本
9. 总结
核心技术贡献点:
- 完整的 Paged KV Cache 实现:包含物理块管理、引用计数、分配/释放/抢占全流程
- 基于内容哈希的前缀缓存:xxhash 链式哈希实现 O(1) 命中检查,支持多序列共享前缀块
- 高效的张量并行:基于 SharedMemory + Event 的轻量 IPC + NCCL 通信,支持 1-8 卡
- CUDA Graph + torch.compile 双重加速:decode 阶段几乎消除 CPU 开销
- 自定义 Triton KV Cache 写入核:解决了分页内存的散射写入问题
- 精简序列化:decode 阶段只传递最后一个 token,大幅降低进程间通信开销
- Gumbel-Max 向量化采样:完全 GPU 侧执行,无 CPU-GPU 同步
适用场景:
- 学习与研究:理解 vLLM 类推理引擎的核心原理
- 快速原型验证:在受限环境中快速测试新算法或新模型
- 定制化部署:在 Qwen3 基础上构建特定应用场景的推理服务
未来可能的扩展方向:
- 扩展模型支持(Llama、Mistral、DeepSeek 等主流架构)
- 添加 top-k/top-p 采样和 greedy 解码
- 增加 INT8/INT4 量化推理路径
- 提供异步/流式生成接口
本文档使用AI工具辅助撰写,作者已进行二次校对与完善。仅作为个人学习与技术分享使用,欢迎交流。
更多推荐
所有评论(0)