vLLM Worker 与 CacheEngine 源码精读

Executor 层级架构

vLLM 的执行层分为三层：Executor → Worker → ModelRunner。

graph TD E["LLMEngine"] GE["GPUExecutor
单 GPU"] DGE["DistributedGPUExecutor
多 GPU (Ray/MP)"] W["Worker
每个 GPU 一个"] CE["CacheEngine
管理 KV Cache"] MR["ModelRunner
执行前向推理"] E -->|"持有"| GE E -->|"或"| DGE GE -->|"持有 driver_worker"| W DGE -->|"持有多个"| W W -->|"持有"| CE W -->|"持有"| MR

GPUExecutor（gpu_executor.py L15-141）是单 GPU 场景的 Executor，核心逻辑：

gpu_executor.py — 初始化流程 L64-70

def _init_non_spec_worker(self):
    assert self.parallel_config.world_size == 1, (
        "GPUExecutor only supports single GPU.")
    self.driver_worker = self._create_worker()
    self.driver_worker.init_device()
    self.driver_worker.load_model()

Executor 是一个薄封装层，它的 execute_model、determine_num_available_blocks、initialize_cache 全部直接委托给 Worker。

Worker 初始化

Worker（worker.py L29-91）是每个 GPU 的核心管理者：

worker.py — Worker.__init__ L37-90

class Worker(WorkerBase):
    """A worker class that executes (a partition of) the model on a GPU.

    Each worker is associated with a single GPU. The worker is responsible for
    maintaining the KV cache and executing the model on the GPU. In case of
    distributed inference, each worker is assigned a partition of the model.
    """

    def __init__(self, model_config, parallel_config, scheduler_config,
                 device_config, cache_config, load_config,
                 local_rank, rank, distributed_init_method,
                 lora_config=None, vision_language_config=None,
                 is_driver_worker=False):
        # ... config 保存 ...
        if self.is_driver_worker:
            assert self.rank == 0, "The driver worker must have rank 0."

        self.model_runner = ModelRunner(
            model_config, parallel_config, scheduler_config,
            device_config, load_config=load_config,
            lora_config=self.lora_config,
            kv_cache_dtype=self.cache_config.cache_dtype,
            is_driver_worker=is_driver_worker,
            vision_language_config=vision_language_config,
        )
        # Uninitialized cache engine. Will be initialized by initialize_cache.
        self.cache_engine: CacheEngine
        self.gpu_cache: CacheList

关键属性：

属性	作用
`model_runner`	模型推理执行器，负责 forward pass
`cache_engine`	KV Cache 管理器，延迟初始化
`gpu_cache`	GPU 上的 KV Cache tensor 列表
`is_driver_worker`	是否为主 worker（rank 0），负责广播控制信号
`local_rank`	本机 GPU 编号

设备初始化与模型加载

worker.py — init_device L92-119

def init_device(self) -> None:
    if self.device_config.device.type == "cuda":
        # 防止 NCCL all_reduce 泄漏内存
        os.environ["TORCH_NCCL_AVOID_RECORD_STREAMS"] = "1"
        # Ray 设置的异步错误处理与 CUDA graph 冲突
        os.environ.pop("NCCL_ASYNC_ERROR_HANDLING", None)

        self.device = torch.device(f"cuda:{self.local_rank}")
        torch.cuda.set_device(self.device)
        _check_if_gpu_supports_dtype(self.model_config.dtype)
        torch.cuda.empty_cache()
        self.init_gpu_memory = torch.cuda.mem_get_info()[0]

    # Initialize the distributed environment.
    init_worker_distributed_environment(
        self.parallel_config, self.rank,
        self.distributed_init_method, self.local_rank)
    set_random_seed(self.model_config.seed)

init_device 完成三件事：

设置 CUDA 设备：绑定到指定的 GPU（local_rank）
记录初始显存：init_gpu_memory 用于后续 profiling 计算 peak memory
初始化分布式环境：NCCL 通信组、随机种子

load_model 直接委托给 ModelRunner：self.model_runner.load_model()。

显存 Profiling

determine_num_available_blocks（L128-175）是 vLLM 启动时的关键步骤——通过实际运行来测量模型的显存占用，然后计算剩余空间能放多少 KV Cache block：

worker.py — determine_num_available_blocks L128-175

@torch.inference_mode()
def determine_num_available_blocks(self) -> Tuple[int, int]:
    torch.cuda.empty_cache()

    # Execute a forward pass with dummy inputs to profile memory usage.
    self.model_runner.profile_run()

    torch.cuda.synchronize()
    free_gpu_memory, total_gpu_memory = torch.cuda.mem_get_info()

    peak_memory = self.init_gpu_memory - free_gpu_memory
    assert peak_memory > 0

    cache_block_size = self.get_cache_block_size_bytes()
    num_gpu_blocks = int(
        (total_gpu_memory * self.cache_config.gpu_memory_utilization -
         peak_memory) // cache_block_size)
    num_cpu_blocks = int(self.cache_config.swap_space_bytes //
                         cache_block_size)
    num_gpu_blocks = max(num_gpu_blocks, 0)
    num_cpu_blocks = max(num_cpu_blocks, 0)
    return num_gpu_blocks, num_cpu_blocks

计算公式：

可用显存 = total_gpu_memory × gpu_memory_utilization - peak_memory
num_gpu_blocks = 可用显存 // cache_block_size

为什么用 profile_run 而不是静态计算

模型的实际显存占用受很多因素影响：模型权重、优化器状态、CUDA 内部分配、activation 临时内存等。静态计算很难准确。通过 profile_run() 执行一次真实的前向传播（用 dummy 输入），可以精确测量 peak memory，确保 KV Cache 分配后不会 OOM。

CacheEngine 核心

CacheEngine（cache_engine.py L51-153）负责实际的 KV Cache 显存分配和管理：

cache_engine.py — CacheEngine.__init__ L59-92

class CacheEngine:
    """Manages the KV cache."""

    def __init__(self, cache_config, model_config, parallel_config):
        self.head_size = model_config.get_head_size()
        self.num_layers = model_config.get_num_layers(parallel_config)
        self.num_heads = model_config.get_num_kv_heads(parallel_config)
        self.block_size = cache_config.block_size
        self.num_gpu_blocks = cache_config.num_gpu_blocks
        self.num_cpu_blocks = cache_config.num_cpu_blocks

        if cache_config.cache_dtype == "auto":
            self.dtype = model_config.dtype
        else:
            self.dtype = STR_DTYPE_TO_TORCH_DTYPE[cache_config.cache_dtype]

        self.attn_backend = get_attn_backend(model_config.dtype)

        # Initialize the cache.
        self.gpu_cache = self._allocate_kv_cache(
            self.num_gpu_blocks, "cuda")
        self.cpu_cache = self._allocate_kv_cache(
            self.num_cpu_blocks, "cpu")

CacheEngine 在初始化时一次性分配好所有 GPU 和 CPU 的 KV Cache 内存。后续的 block 分配/释放只是在 Block Manager 的逻辑层面操作，不再涉及实际的 torch.cuda.malloc。

KV Cache 分配

cache_engine.py — _allocate_kv_cache L94-120

def _allocate_kv_cache(self, num_blocks, device):
    """Allocates KV cache on the specified device."""
    kv_cache_shape = self.attn_backend.get_kv_cache_shape(
        num_blocks, self.block_size, self.num_heads, self.head_size)
    pin_memory = is_pin_memory_available() if device == "cpu" else False

    kv_cache: CacheList = CacheList(
        num_const_cache=self.cache_config.num_const_cache,
        num_layers=self.num_linear_layers,
        cache_shape=(self.num_linear_heads, self.head_size, self.head_size),
        dtype=torch.float32,
        device=device,
    )
    for i in range(self.num_layers):
        kv_cache.append(
            torch.empty(kv_cache_shape, dtype=self.dtype,
                        pin_memory=pin_memory, device=device))
    return kv_cache

KV Cache 的数据结构：

外层：CacheList，长度 = num_layers，每个元素对应一层 Transformer
每层：一个大的 torch.Tensor，shape 由 attention backend 决定
典型 shape：[2, num_blocks, block_size, num_kv_heads, head_size]（2 表示 K 和 V）

Cache block size 的计算（L136-153）：

cache_block_size = num_layers × (key_block + value_block) × dtype_size
# 其中 key_block = block_size × num_kv_heads × head_size

CPU Cache 使用 pinned memory

CPU 侧的 KV Cache 使用 pin_memory=True（如果可用）。Pinned memory 不会被操作系统 swap 到磁盘，且支持异步的 CPU↔GPU 数据传输（cudaMemcpyAsync），这对 swap 操作的性能至关重要。

Swap 与 Copy 操作

CacheEngine 提供三个 KV Cache 操作接口：

cache_engine.py — swap & copy L122-133

def swap_in(self, src_to_dst: Dict[int, int]) -> None:
    for i in range(self.num_layers):
        self.attn_backend.swap_blocks(
            self.cpu_cache[i], self.gpu_cache[i], src_to_dst)

def swap_out(self, src_to_dst: Dict[int, int]) -> None:
    for i in range(self.num_layers):
        self.attn_backend.swap_blocks(
            self.gpu_cache[i], self.cpu_cache[i], src_to_dst)

def copy(self, src_to_dsts: Dict[int, List[int]]) -> None:
    self.attn_backend.copy_blocks(self.gpu_cache, src_to_dsts)

操作	方向	触发场景
`swap_in`	CPU → GPU	Scheduler 恢复被 swap out 的请求
`swap_out`	GPU → CPU	Scheduler 抢占请求，选择 SWAP 模式
`copy`	GPU → GPU	Copy-on-Write：beam search fork 后首次写入

这三个操作都委托给 attn_backend（如 FlashAttention 或 PagedAttention 的 CUDA kernel），在 GPU 上高效执行块级别的数据拷贝。src_to_dst 是 block ID 的映射关系，由 Scheduler 通过 Block Manager 计算得出。

execute_model 流程

Worker 的 execute_model 是每个推理 step 的入口。整体流程：

graph TD A["Worker.execute_model(request)"] --> B["cache_engine.swap_in
(blocks_to_swap_in)"] B --> C["cache_engine.swap_out
(blocks_to_swap_out)"] C --> D["cache_engine.copy
(blocks_to_copy)"] D --> E["model_runner.execute_model
(seq_group_metadata_list, kv_caches)"] E --> F["返回 SamplerOutput"]

执行顺序：

先执行 cache 操作：swap_in、swap_out、copy 必须在模型执行之前完成，确保 KV Cache 数据就绪
再执行模型前向：将 seq_group_metadata_list（调度器产出）和 gpu_cache（KV Cache tensors）传给 ModelRunner

Driver Worker 的特殊角色

在分布式场景中，只有 is_driver_worker=True 的 Worker（rank 0）会从 Engine 接收 ExecuteModelRequest，然后通过 broadcast_tensor_dict 将输入数据广播给其他 Worker。非 driver Worker 通过接收广播来获取输入。

分布式 Worker 协作

多 GPU 场景使用 DistributedGPUExecutor（基于 Ray 或 multiprocessing），每个 GPU 运行一个 Worker 进程。

协作模式：

Tensor Parallelism：模型的每一层被分割到多个 GPU 上。每个 Worker 持有模型的一个分片，通过 NCCL AllReduce 在层间同步
Pipeline Parallelism：模型的不同层分配到不同 GPU，数据在层间流水线式传递
KV Cache 独立：每个 Worker 的 CacheEngine 独立管理自己 GPU 上的 KV Cache，不需要跨 GPU 同步 cache 数据

初始化顺序（所有 Worker 同步执行）：

init_device() → 设置 CUDA 设备，初始化 NCCL
load_model() → 加载模型分片
determine_num_available_blocks() → profile 显存，取所有 Worker 的最小值
initialize_cache() → 分配 KV Cache

为什么取所有 Worker 的最小 block 数

不同 GPU 可能有不同的可用显存（例如 rank 0 的 GPU 可能被系统占用更多内存）。为了保证所有 Worker 的 KV Cache 大小一致（Scheduler 统一管理），必须取最小值。