推理引擎架构深度解析：从内存管理到计算调度的全栈优化设计
一、推理引擎的性能瓶颈：框架开销不容忽视
大模型推理的性能，模型本身只占一部分，推理引擎的架构设计同样关键。拿 LLaMA-7B 来说，用原生 PyTorch 跑，延迟大概在 150ms/token；换成 llama.cpp 能压到 30ms/token；要是用 vLLM 在高并发下，吞吐量甚至是 PyTorch 的 10 倍。这些差距不是硬件或模型带来的，而是推理引擎在内存管理、计算调度和并发控制上的架构差异。
推理引擎的核心任务很明确：在 GPU 显存有限的情况下，既要最大化吞吐量（tokens/s），又要最小化首 token 延迟（TTFT）。但这俩目标经常打架------高吞吐需要大 batch size，大 batch 又意味着显存占用和计算延迟增加。架构设计的本质，就是在这些约束里找平衡点。
二、推理引擎的核心架构与数据流转
推理引擎可以拆成五个核心子系统，数据在它们之间流转的方式直接决定了端到端的性能。
请求调度器是引擎的"大脑"。连续批处理（Continuous Batching）能把不同时间到的请求动态合到一个 batch 里，不用等所有请求都跑完再开下一轮。这能大幅提升 GPU 利用率------传统静态批处理里，短请求跑完了 GPU 还得等长请求，算力就浪费了。
KV Cache 管理器是引擎的"心脏"。自回归生成时，每生成一个 token 都得重新算前面所有 token 的 Key 和 Value，KV Cache 就是把这些中间结果存起来避免重算。但 KV Cache 占的显存会随序列长度线性增长，是显存瓶颈所在。PagedAttention（vLLM 的核心）把 KV Cache 分页存，按需分配释放，消除了预分配带来的浪费。
算子执行引擎负责张量计算。FlashAttention 用分块计算（Tiling）减少 HBM 访问，把注意力计算从内存带宽瓶颈转成计算瓶颈。算子融合把 QKV 投影、注意力计算和输出投影合到一个内核里，省去了中间张量的显存读写。
内存分配器管 GPU 显存生命周期。显存池化预分配大块显存，按需切分给不同请求，避免频繁 cudaMalloc/cudaFree 和显存碎片。CPU 卸载在显存不够时把不活跃的 KV Cache 挪到主机内存，需要时再拉回 GPU。
采样与解码器把模型输出的 logits 转成 token。支持多种采样策略（贪心、Top-K、Top-P、温度缩放），还要处理停止词检测和特殊 token。
三、推理引擎核心组件的代码实现
下面这段代码展示了连续批处理调度器和 PagedAttention 风格的 KV Cache 管理器的核心逻辑。
python
from dataclasses import dataclass, field
from typing import Optional
from enum import Enum
import time
class RequestStatus(Enum):
"""请求状态枚举"""
WAITING = "waiting" # 等待调度
RUNNING = "running" # 正在推理
PREEMPTED = "preempted" # 被抢占（显存不足）
COMPLETED = "completed" # 推理完成
@dataclass
class InferenceRequest:
"""推理请求"""
request_id: int
prompt_tokens: list[int] # 输入 token 序列
max_tokens: int # 最大生成长度
generated_tokens: list[int] = field(default_factory=list)
status: RequestStatus = RequestStatus.WAITING
arrival_time: float = 0.0 # 请求到达时间
priority: int = 0 # 优先级（首 token 请求优先）
@dataclass
class KVBlock:
"""KV Cache 分页块：固定大小的存储单元"""
block_id: int
size: int = 16 # 每块存储 16 个 token 的 KV
ref_count: int = 0 # 引用计数（支持前缀共享）
class PagedKVCache:
"""
PagedAttention 风格的 KV Cache 管理器
将 KV Cache 分页存储，按需分配和释放，消除预分配浪费
"""
def __init__(self, total_blocks: int, block_size: int = 16):
self.block_size = block_size
# 空闲块池：所有可分配的 KV 块
self.free_blocks: list[KVBlock] = [
KVBlock(block_id=i, size=block_size) for i in range(total_blocks)
]
# 每个请求占用的块列表
self.request_blocks: dict[int, list[KVBlock]] = {}
# 物理块表：块ID → 实际数据存储位置
self.block_table: dict[int, Optional[bytes]] = {}
def allocate(self, request_id: int, num_tokens: int) -> bool:
"""
为请求分配 KV Cache 块
返回 True 表示分配成功，False 表示显存不足
"""
num_blocks_needed = (num_tokens + self.block_size - 1) // self.block_size
num_blocks_existing = len(self.request_blocks.get(request_id, []))
num_blocks_to_alloc = num_blocks_needed - num_blocks_existing
if num_blocks_to_alloc <= 0:
return True # 已有足够块
if len(self.free_blocks) < num_blocks_to_alloc:
return False # 显存不足
# 从空闲池中取出所需块数
if request_id not in self.request_blocks:
self.request_blocks[request_id] = []
for _ in range(num_blocks_to_alloc):
block = self.free_blocks.pop()
block.ref_count += 1
self.request_blocks[request_id].append(block)
return True
def free(self, request_id: int) -> int:
"""
释放请求占用的所有 KV Cache 块
返回释放的块数量
"""
if request_id not in self.request_blocks:
return 0
freed = 0
for block in self.request_blocks[request_id]:
block.ref_count -= 1
if block.ref_count == 0:
# 引用计数归零，块可被复用
self.free_blocks.append(block)
freed += 1
del self.request_blocks[request_id]
return freed
def get_num_free_blocks(self) -> int:
return len(self.free_blocks)
class ContinuousBatchScheduler:
"""
连续批处理调度器
动态合并请求到同一 batch，短请求完成后立即释放资源
"""
def __init__(self, kv_cache: PagedKVCache, max_batch_size: int = 32):
self.kv_cache = kv_cache
self.max_batch_size = max_batch_size
self.waiting_queue: list[InferenceRequest] = []
self.running_batch: list[InferenceRequest] = []
def add_request(self, request: InferenceRequest) -> None:
"""添加新请求到等待队列"""
request.arrival_time = time.time()
request.priority = 1 # 新请求优先调度以降低首 token 延迟
self.waiting_queue.append(request)
def schedule(self) -> list[InferenceRequest]:
"""
调度一轮推理：从等待队列中选择请求加入运行批次
考虑显存容量和 batch size 限制
"""
# 先处理已完成的请求，释放 KV Cache
completed = [r for r in self.running_batch
if len(r.generated_tokens) >= r.max_tokens]
for req in completed:
req.status = RequestStatus.COMPLETED
self.kv_cache.free(req.request_id)
self.running_batch = [r for r in self.running_batch
if r.status == RequestStatus.RUNNING]
# 从等待队列中补充请求
still_waiting = []
for req in self.waiting_queue:
if len(self.running_batch) >= self.max_batch_size:
still_waiting.append(req)
continue
# 尝试为新请求分配 KV Cache
total_tokens = len(req.prompt_tokens) + 1 # 至少生成1个token
if self.kv_cache.allocate(req.request_id, total_tokens):
req.status = RequestStatus.RUNNING
self.running_batch.append(req)
else:
# 显存不足，请求保留在等待队列
still_waiting.append(req)
self.waiting_queue = still_waiting
return self.running_batch
def preempt_if_needed(self) -> Optional[InferenceRequest]:
"""
抢占机制：当显存严重不足时，换出最低优先级的运行请求
被抢占的请求的 KV Cache 被释放，请求回到等待队列
"""
if self.kv_cache.get_num_free_blocks() > 0:
return None
# 找到优先级最低的运行请求（已生成最多 token 的请求）
victim = min(self.running_batch,
key=lambda r: r.priority * 1000 - len(r.generated_tokens))
# 释放其 KV Cache
self.kv_cache.free(victim.request_id)
victim.status = RequestStatus.PREEMPTED
self.running_batch.remove(victim)
# 重新加入等待队列，降低优先级防止饥饿
victim.priority = 0
self.waiting_queue.append(victim)
return victim代码里 PagedKVCache 实现了分页 KV Cache 管理，按需分配释放显存块，用引用计数支持前缀共享。ContinuousBatchScheduler 实现了连续批处理调度，动态合并请求，显存不够时触发抢占。
四、推理引擎架构的权衡与适用场景
PagedAttention 的碎片开销：分页存储引入了块表查找的间接开销，块大小选得不好会影响内部碎片率。块太小块表太大、查找慢；块太大内部碎片又严重。vLLM 实验发现，块大小设为 16 个 token 时大多数场景下平衡最好。
连续批处理的延迟抖动：新请求加入 batch 会增加单步推理的计算量，正在跑的请求延迟就会增加。对延迟敏感的场景（比如对话 AI），得限制最大 batch size 或者设个延迟预算。
CPU 卸载的带宽瓶颈：KV Cache 在 GPU 和 CPU 之间传，受 PCIe 带宽限制。PCIe 4.0 x16 理论带宽大概 32GB/s，一个 7B 模型的 KV Cache 可能有数 GB，卸载加载的延迟能到几百毫秒。CPU 卸载只适合吞吐优先、对延迟不敏感的场景。
适用边界：这架构适合大语言模型的在线推理服务，尤其是多用户并发场景。单用户离线推理（比如批量文本生成）用简单的顺序推理更高效，不需要复杂的调度和显存管理。小模型（< 1B 参数）KV Cache 显存占用不是瓶颈，PagedAttention 收益有限。
五、总结
推理引擎架构的核心是在显存约束下最大化吞吐量并控制延迟。PagedAttention 解决了 KV Cache 的显存浪费，连续批处理提升了 GPU 利用率，抢占机制保证了显存不足时的系统稳定性。落地建议：优先实现连续批处理调度，这是吞吐量提升最明显的优化；KV Cache 管理用分页方案，块大小设 16 个 token；显存不足时通过抢占释放资源，别直接拒绝新请求。架构选择得看实际场景的延迟和吞吐需求，没有通用的最优方案。