vLLM 推理流程详解

本文档详细说明当请求到来后，vLLM 的推理流程会经过哪些关键模块和脚本。

整体架构概览

vLLM 的推理流程主要分为以下几个阶段：

  1. API 服务器接收请求
  2. 输入处理（Tokenization）
  3. 请求调度（Scheduling）
  4. 模型执行（Model Execution）
  5. 采样（Sampling）
  6. 输出处理（Output Processing）
  7. KV Cache 管理

详细流程

1. API 服务器入口层

关键文件：

  • vllm/entrypoints/openai/api_server.py - OpenAI 兼容的 API 服务器
  • vllm/entrypoints/openai/serving_chat.py - Chat Completions 处理
  • vllm/entrypoints/openai/serving_completion.py - Completions 处理

流程：

HTTP 请求 → FastAPI 路由 → API Handler → AsyncLLM

关键代码位置：

@router.post(
    "/v1/chat/completions",
    dependencies=[Depends(validate_json_request)],
    ...
)
@with_cancellation
@load_aware_call
async def create_chat_completion(request: ChatCompletionRequest, raw_request: Request):
    ...
    handler = chat(raw_request)
    ...
    generator = await handler.create_chat_completion(request, raw_request)
    ...

说明：

  • 当客户端发送 HTTP 请求到 /v1/chat/completions 或 /v1/completions 时
  • FastAPI 路由将请求转发到对应的 handler（OpenAIServingChat 或 OpenAIServingCompletion）
  • Handler 调用 AsyncLLM.generate() 方法启动推理流程

2. AsyncLLM 层（异步引擎客户端）

关键文件：

  • vllm/v1/engine/async_llm.py - 异步 LLM 引擎客户端（这个是核心，关于调度器的执行和最后generation的输出均由它经手，为vllm盒子的出入口。）

流程：

AsyncLLM.add_request() → InputProcessor.process_inputs() → EngineCore.add_request()

关键代码位置：

async def add_request(
    self,
    request_id: str,
    prompt: EngineCoreRequest | PromptType,
    params: SamplingParams | PoolingParams,
    ...
) -> RequestOutputCollector:
    ...
    # 处理输入（tokenization）
    request = self.input_processor.process_inputs(
        request_id,
        prompt,
        params,
        ...
    )
    ...
    # 添加到 EngineCore
    await self.engine_core.add_request_async(request)

说明：

  • AsyncLLM 是异步引擎的客户端接口
  • 负责将请求转换为 EngineCoreRequest 并提交到 EngineCore
  • 创建 RequestOutputCollector 用于收集输出结果

3. 输入处理层（Input Processing）

关键文件：

  • vllm/v1/engine/input_processor.py - 输入处理器

主要功能：

  1. Tokenization（分词）：将文本转换为 token IDs
  2. 多模态处理：处理图像、视频等多模态输入
  3. 参数验证：验证采样参数、logprobs 等
  4. 结构化输出验证：验证 grammar、JSON schema 等

关键代码位置：

class InputProcessor:
    def __init__(
        self,
        vllm_config: VllmConfig,
        tokenizer: TokenizerLike | None,
        mm_registry: MultiModalRegistry = MULTIMODAL_REGISTRY,
    ) -> None:
        ...
        self.input_preprocessor = InputPreprocessor(
            self.model_config,
            tokenizer,
            mm_registry,
            ...
        )

说明：

  • InputProcessor 使用 InputPreprocessor 进行实际的 tokenization
  • 处理完成后生成 EngineCoreRequest 对象，包含：
    • prompt_token_ids: token ID 列表
    • sampling_params: 采样参数
    • lora_request: LoRA 适配器请求（如果有）

4. 引擎核心层（EngineCore）

关键文件：

  • vllm/v1/engine/core.py - 引擎核心，协调整个推理流程。即使Asyncllm的左膀右臂，更底层是用它来进行各类实际工作的执行。

主要职责：

  • 管理请求队列
  • 协调调度器和模型执行器
  • 处理输出结果

关键代码位置：

def step(self) -> tuple[dict[int, EngineCoreOutputs], bool]:
    """Schedule, execute, and make output."""
    
    if not self.scheduler.has_requests():
        return {}, False
    
    # 1. 调度：选择要处理的请求
    scheduler_output = self.scheduler.schedule()
    
    # 2. 执行模型
    future = self.model_executor.execute_model(scheduler_output, non_block=True)
    
    # 3. 获取 grammar bitmask（用于结构化输出）
    grammar_output = self.scheduler.get_grammar_bitmask(scheduler_output)
    
    # 4. 等待模型执行完成并采样
    model_output = future.result()
    if model_output is None:
        model_output = self.model_executor.sample_tokens(grammar_output)
    
    # 5. 处理中止请求
    self._process_aborts_queue()
    
    # 6. 更新调度器状态
    engine_core_outputs = self.scheduler.update_from_output(
        scheduler_output, model_output
    )
    
    return engine_core_outputs, scheduler_output.total_num_scheduled_tokens > 0

说明：

  • EngineCore.step() 是核心循环，每个 step 对应一次模型前向传播
  • 流程：调度 → 执行 → 采样 → 更新状态（一个Step对应的动作）

5. 调度器层（Scheduler）

关键文件：

  • vllm/v1/core/sched/scheduler.py - 请求调度器

主要功能：

  1. 请求选择：从等待队列中选择要处理的请求
  2. KV Cache 分配：为请求分配 KV cache blocks
  3. 批处理构建：构建批处理，决定每个请求处理多少 tokens
  4. 优先级管理：根据优先级调度请求
  5. 抢占（Preemption）：必要时抢占低优先级请求
5.1 批处理构建机制（Token Budget）

核心机制：

调度器使用 Token Budget（token 预算（批次层面的全局预算，所有的请求共享这个预算，每个请求的num_new_tokens不能超过剩余预算）） 机制来构建批处理。每个调度步骤都有一个固定的 max_num_scheduled_tokens（最大可调度 token 数），作为 token_budget。

构建流程：

  1. 优先调度 RUNNING 请求：

    • 遍历 self.running 队列中的每个请求
    • 计算每个请求需要处理的新 token 数：num_new_tokens = request.num_tokens_with_spec - request.num_computed_tokens（这里的spec指的是推测解码，调度器会尽量让num_computed_tokens追上num_tokens_with_spec）
    • 如果 num_new_tokens > token_budget，则限制为 token_budget（支持 chunked prefill）
    • 为每个请求分配 KV blocks（见下文）
    • 从 token_budget 中减去已分配的 token 数
    • 将请求加入 scheduled_running_reqs

  2. 调度 WAITING 请求：

    • 当 token_budget > 0 且 running 队列未满时
    • 从 waiting 队列取出请求
    • 检查前缀缓存：如果请求的前缀已缓存，直接复用缓存的 blocks
    • 计算需要的新 token 数（考虑已缓存的 tokens）
    • 如果 num_new_tokens <= token_budget，则调度该请求
    • 否则，如果启用了 chunked prefill，可以部分调度；否则跳过该请求

  3. 构建最终 Batch：

    • 所有被调度的请求组成一个 batch
    • SchedulerOutput 包含：
      • scheduled_running_reqs: 正在运行的请求列表
      • scheduled_new_reqs: 新加入的请求列表
      • num_scheduled_tokens: 每个请求要处理的 token 数（字典）
      • req_to_new_blocks: 每个请求新分配的 KV blocks

为什么不是每次只生成 1 个 token？

  • Prefill 阶段：需要处理整个 prompt（可能数百/数千 tokens），一次前向传播处理所有 prompt tokens
  • Chunked Prefill：长 prompt 被分成多个 chunks，每个 chunk 可能包含多个 tokens
  • Decode 阶段：虽然每次生成 1 个 token，但可能同时处理多个请求，形成批处理

关键代码位置：

# First, schedule the RUNNING requests.
req_index = 0
while req_index < len(self.running) and token_budget > 0:
    request = self.running[req_index]
    ...
    num_new_tokens = (
        request.num_tokens_with_spec
        + request.num_output_placeholders
        - request.num_computed_tokens
    )
    if 0 < self.scheduler_config.long_prefill_token_threshold < num_new_tokens:
        num_new_tokens = self.scheduler_config.long_prefill_token_threshold
    num_new_tokens = min(num_new_tokens, token_budget)
    ...
    token_budget -= num_new_tokens
    req_index += 1

# Next, schedule the WAITING requests.
if not preempted_reqs:
    while self.waiting and token_budget > 0:
        request = self.waiting.peek_request()
        ...
        num_new_tokens = min(num_new_tokens, token_budget)
        ...
5.2 请求状态转换

状态定义：

  • WAITING：等待调度
  • WAITING_FOR_REMOTE_KVS：等待远程 KV 传输（KV Connector 场景）
  • WAITING_FOR_FSM：等待 FSM 编译（结构化输出场景）
  • RUNNING：正在运行
  • PREEMPTED：被抢占（换出）

状态转换流程：

  • WAITING → RUNNING：请求被调度时
  • RUNNING → RUNNING：继续运行（每次 step 后仍可能被调度）
  • RUNNING → PREEMPTED：内存不足时被抢占
  • PREEMPTED → RUNNING：恢复运行时

为什么 RUNNING 请求可能仍然是 RUNNING？

每次 step 后，如果请求未完成且仍有 tokens 需要处理，会继续保持在 running 队列。只有完成或达到 max_tokens 时才会从 running 队列移除。

关键代码位置：

self.running.append(request)
...
request.status = RequestStatus.RUNNING
request.num_computed_tokens = num_computed_tokens
5.3 连续批处理实现

核心特点：

  • 每个 step 动态选择请求，不固定 batch
  • 新请求可以立即加入，完成请求立即移除
  • 通过 token_budget 控制批处理大小

实现机制：

  • 调度器维护 waiting 和 running 两个队列
  • 每个 step 从 running 队列优先调度，然后从 waiting 队列补充
  • 请求完成后自动从 running 队列移除
  • 新请求到达时立即加入 waiting 队列，等待下次调度

说明：

  • 调度器维护三个队列：
    • waiting: 等待调度的新请求
    • running: 正在运行的请求
    • swapped: 被换出的请求（内存不足时）
  • 返回 SchedulerOutput，包含：
    • 要处理的请求列表
    • 每个请求分配的 KV cache blocks
    • 每个请求要处理的 token 数量

6. KV Cache 管理（Block Pool）

关键文件：

  • vllm/v1/core/block_pool.py - KV Cache 块池管理
  • vllm/v1/core/kv_cache_manager.py - KV Cache 管理器
6.1 Block 分配策略

初始分配时机：

  • 新请求首次调度时（num_computed_tokens == 0），调用 allocate_slots()
  • 计算所需 blocks：num_blocks = ceil((num_prompt_tokens + max_tokens) / block_size)
  • 但实际分配是渐进式的：每次只分配当前 step 需要的 blocks

Prefill vs Decode 的分配策略：

  • Prefill 阶段：可能一次分配多个 blocks（如果 prompt 很长且未启用 chunked prefill）
  • Chunked Prefill：每次只分配一个 chunk 需要的 blocks
  • Decode 阶段：每次只分配 1 个 block（因为每次只生成 1 个 token）

剩余 Block 处理：

  • 如果分配的 block 未填满，剩余空间保留给后续 tokens
  • Block 填满后自动缓存（如果启用前缀缓存）
  • 请求完成时，所有 blocks 的 ref_cnt--，如果 ref_cnt == 0 则放回 free_block_queue

关键代码位置：

def allocate_slots(
    self,
    request: Request,
    num_new_tokens: int,
    ...
) -> KVCacheBlocks | None:
    """Add slots for a request with new tokens to append."""
    ...
    # The allocation has three stages:
    # - Free unnecessary blocks in `comp` and check
    #   if we have sufficient free blocks (return None if not).
    # - Handle prefix tokens (`comp + new_comp + ext_comp`):
    #     - Free unnecessary blocks (e.g. outside sliding window)
    #     - Allocate new blocks for `ext_comp` tokens inside
    #       sliding window
    # - Allocate new blocks for tokens to be computed (`new + lookahead`)
    ...

def get_new_blocks(self, num_blocks: int) -> list[KVCacheBlock]:
    """Get new blocks from the free block pool."""
    if num_blocks > self.get_num_free_blocks():
        raise ValueError(f"Cannot get {num_blocks} free blocks from the pool")
    
    ret: list[KVCacheBlock] = self.free_block_queue.popleft_n(num_blocks)
    
    for block in ret:
        assert block.ref_cnt == 0
        block.ref_cnt += 1  # 增加引用计数
        if self.metrics_collector:
            self.metrics_collector.on_block_allocated(block)
    return ret
6.2 前缀缓存机制

前缀缓存命中流程：

  • 调度新请求时，调用 get_computed_blocks() 查找前缀缓存
  • 通过 block_hashes 查找 BlockHashToBlockMap
  • 如果找到匹配的 blocks，直接复用（ref_cnt++）
  • num_computed_tokens 被设置为命中的 tokens 数量
  • 只对未命中的部分分配新 blocks

已计算 Tokens 的含义：

  • num_computed_tokens：已经完成前向传播并写入 KV cache 的 tokens 数量
  • 包括：
    • 本地计算的 tokens
    • 前缀缓存命中的 tokens
    • 外部计算的 tokens（KV connector）

关键代码位置：

def get_computed_blocks(self, request: Request) -> tuple[KVCacheBlocks, int]:
    """Get the computed (cached) blocks for the request."""
    if not self.enable_caching or request.skip_reading_prefix_cache:
        return self.empty_kv_cache_blocks, 0
    
    max_cache_hit_length = request.num_tokens - 1
    computed_blocks, num_new_computed_tokens = (
        self.coordinator.find_longest_cache_hit(
            request.block_hashes, max_cache_hit_length
        )
    )
    ...
    return self.create_kv_cache_blocks(computed_blocks), num_new_computed_tokens
6.3 多级缓存机制

缓存层级：

  1. 本地 KV Cache：GPU 内存中的 blocks（最快）
  2. 前缀缓存（Prefix Cache）：相同前缀的 blocks 可以共享
  3. 外部 KV Cache（KV Connector）：其他设备/节点的 KV cache

缓存查找顺序：

  1. 首先查找本地前缀缓存（BlockHashToBlockMap）
  2. 如果使用 KV Connector，查找外部缓存
  3. 如果都未命中，分配新 blocks 并计算

缓存更新：

  • Block 填满后自动缓存（计算 hash 并插入 BlockHashToBlockMap）
  • 使用引用计数管理共享 blocks

说明：

  • vLLM 使用 PagedAttention 机制，将 KV cache 分成固定大小的 blocks
  • 每个 block 可以存储多个 tokens 的 KV 值（通常 block_size = 16）
  • 支持前缀缓存：相同前缀可以共享 KV cache blocks

7. 模型执行层（Model Runner）

关键文件：

  • vllm/v1/worker/gpu_model_runner.py - GPU 模型执行器（V1 架构）
  • vllm/v1/worker/gpu/model_runner.py - GPU 模型执行器（简化版）
7.1 KV Cache 生成机制

前向传播中的 KV 写入：

  1. 准备阶段（ModelRunner）：

    • ModelRunner.prepare_input() 根据 SchedulerOutput 构建 InputBatch
    • InputBatch 包含：
      • input_ids: 当前 step 要处理的 token IDs
      • positions: 每个 token 在序列中的位置
      • attn_metadata: Attention 元数据，包含 block_tables（每个请求的 block 映射表）和 slot_mapping（每个 token 应该写入哪个 block 的哪个 slot）

  2. Attention 层执行：

    • 模型逐层执行，每层的 Attention.forward() 被调用
    • Attention 层接收 query、key、value 三个张量
    • 通过 ForwardContext 获取当前层的 KV cache 和 attention metadata

  3. KV Cache 写入（CUDA Kernel）：

    • Attention backend（如 Flash Attention、PagedAttention）调用底层 CUDA kernel
    • Kernel 根据 slot_mapping 确定每个 token 的 K/V 应该写入的位置： 
      slot = slot_mapping[token_idx]
block_id = slot // block_size
offset_in_block = slot % block_size
    • 将 key 和 value 张量写入对应的 block 和 offset
    • 写入操作由 reshape_and_cache kernel 完成

  4. Attention 计算：

    • 在写入新 K/V 的同时，从 KV cache 中读取历史 tokens 的 K/V
    • 使用 PagedAttention kernel 进行高效的 attention 计算
    • 计算 Q @ K^T 和 softmax(...) @ V

一次前向传播的完整流程：

Input Tokens → Embedding Layer
    ↓
Layer 1:
    → Attention.forward(query, key, value)
    → CUDA Kernel: reshape_and_cache(key, value, slot_mapping)
    → 写入 Layer 1 的 KV cache blocks
    → PagedAttention: 读取历史 K/V，计算 attention
    → FFN
    ↓
Layer 2:
    → Attention.forward(query, key, value)
    → CUDA Kernel: reshape_and_cache(key, value, slot_mapping)
    → 写入 Layer 2 的 KV cache blocks
    → PagedAttention: 读取历史 K/V，计算 attention
    → FFN
    ↓
...
    ↓
Output Layer → Logits

关键代码位置：

@staticmethod
def write_to_paged_cache(
    key: torch.Tensor,
    value: torch.Tensor,
    key_cache: torch.Tensor,
    value_cache: torch.Tensor,
    slot_mapping: torch.Tensor,
    kv_cache_dtype: str,
    k_scale: torch.Tensor,
    v_scale: torch.Tensor,
) -> None:
    ops.reshape_and_cache(
        key,
        value,
        key_cache,
        value_cache,
        slot_mapping.flatten(),
        kv_cache_dtype,
        k_scale,
        v_scale,
    )

def forward(
    self,
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    ...
) -> torch.Tensor:
    ...
    if self.use_direct_call:
        forward_context: ForwardContext = get_forward_context()
        attn_metadata = forward_context.attn_metadata
        self_kv_cache = self.kv_cache[forward_context.virtual_engine]
        self.impl.forward(
            self, query, key, value, self_kv_cache, attn_metadata, output=output
        )
    ...
7.2 Worker 与 ModelRunner

概念区分：

在 vLLM 的架构中，Worker 和 ModelRunner 是两个不同层次的概念：

  1. Worker（工作进程/设备实例）：

    • Worker 是进程或设备级别的概念
    • 在分布式场景中，每个 GPU 对应一个 worker 进程
    • Worker 负责：
      • 管理该设备上的模型实例
      • 接收来自 EngineCore 的调度指令
      • 协调该设备上的所有 ModelRunner 实例

  2. ModelRunner（模型执行器）：

    • ModelRunner 是单次前向传播的执行者
    • 一个 Worker 内部通常只有一个 ModelRunner 实例（如 GPUModelRunner）
    • ModelRunner 负责：
      • 准备输入数据（prepare_input()）
      • 调用模型的 forward() 方法
      • 管理 CUDA Graph、LoRA 适配器等
      • 与 KV cache 交互（通过 Attention metadata）

架构层次：

EngineCore (调度层)
    ↓
Executor (执行器接口)
    ↓
Worker Process (工作进程，每个 GPU 一个)
    ├── GPUModelRunner (模型执行器实例)
    │   ├── model (PyTorch 模型)
    │   ├── kv_cache (KV cache 存储)
    │   └── cudagraph_manager (CUDA Graph 管理器)

关键代码位置：

class GPUModelRunner(
    LoRAModelRunnerMixin, KVConnectorModelRunnerMixin, ECConnectorModelRunnerMixin
):
    """GPU Model Runner for V1 architecture."""

主要功能：

  1. 准备输入：构建 input_ids、positions、attention metadata
  2. 执行前向传播：调用模型进行推理
  3. 处理 LoRA：如果使用 LoRA，动态加载适配器
  4. CUDA Graph：使用 CUDA Graph 优化性能

关键代码位置：

@torch.inference_mode()
def execute_model(
    self,
    scheduler_output: SchedulerOutput,
    intermediate_tensors: Any | None = None,
    dummy_run: bool = False,
) -> ModelRunnerOutput | None:
    ...
    # 准备输入批次
    input_batch = self.prepare_input(scheduler_output)
    
    # 准备 attention metadata
    self.prepare_attn_metadata(input_batch)
    
    # 运行模型
    if cudagraph_mode == CUDAGraphMode.FULL:
        # 使用 CUDA Graph
        hidden_states = self.cudagraph_manager.run(...)
    else:
        # 正常 PyTorch 执行
        hidden_states = self.model(
            input_ids=input_batch.input_ids,
            positions=input_batch.positions,
        )
    ...

说明：

  • ModelRunner 负责实际的模型前向传播
  • 支持多种优化：
    • CUDA Graph：捕获和重放 CUDA kernel 序列
    • Flash Attention：高效的 attention 计算
    • 连续批处理（Continuous Batching）：动态批处理大小

8. Chunked Prefill 机制

关键文件：

  • vllm/v1/core/sched/scheduler.py - 调度器中的 chunked prefill 逻辑
  • vllm/v1/attention/backends/utils.py - Prefill chunk 分割工具

触发条件：

  • long_prefill_token_threshold：如果 num_new_tokens > threshold，则限制为 threshold
  • enable_chunked_prefill=True：必须显式启用（V1 中默认启用）

实现机制：

  • 长 prompt 被分成多个 chunks，每个 chunk 大小不超过 token_budget
  • 每个 chunk 独立调度，可以与其他 decode 请求混合
  • 调度策略：优先调度所有 decode 请求，然后调度 prefill chunks

优势：

  • 提高 GPU 利用率（compute-bound prefill + memory-bound decode）
  • 降低 ITL（Inter-Token Latency）
  • 允许长 prompt 请求立即开始处理，而不需要等待整个 prompt 处理完成

关键代码位置：

if 0 < self.scheduler_config.long_prefill_token_threshold < num_new_tokens:
    num_new_tokens = self.scheduler_config.long_prefill_token_threshold
num_new_tokens = min(num_new_tokens, token_budget)

# chunked prefill has to be enabled explicitly to allow
# pooling requests to be chunked
if (
    not self.scheduler_config.enable_chunked_prefill
    and num_new_tokens > token_budget
):
    # If chunked_prefill is disabled,
    # we can stop the scheduling here.
    break

num_new_tokens = min(num_new_tokens, token_budget)

def split_prefill_chunks(
    seq_lens_cpu: torch.Tensor, workspace_size: int, request_offset: int = 0
) -> list[tuple[int, int]]:
    """
    Split the prefill requests into chunks such that the total sequence length
    of each chunk is less than or equal to the workspace size.
    """
    chunk_bounds = []
    i, n = 0, len(seq_lens_cpu)
    assert torch.all(seq_lens_cpu <= workspace_size).item()

    while i < n:
        start, chunk_total = i, 0
        while i < n and (chunk_total + (s := seq_lens_cpu[i].item())) <= workspace_size:
            chunk_total += s
            i += 1
        chunk_bounds.append((start + request_offset, i + request_offset))
    return chunk_bounds

9. 采样层（Sampling）

关键文件：

  • vllm/v1/worker/gpu/model_runner.py - 采样逻辑

主要功能：

  1. Logits 处理：应用 temperature、top-p、top-k 等
  2. 采样：根据 logits 采样下一个 token
  3. 结构化输出：支持 grammar、JSON schema 等约束

关键代码位置：

@torch.inference_mode()
def sample_tokens(
    self,
    grammar_output: GrammarOutput | None,
) -> AsyncOutput | ModelRunnerOutput:
    ...

说明：

  • 从模型的 logits 输出中采样下一个 token
  • 支持多种采样策略：greedy、multinomial、beam search 等
  • 支持结构化输出约束（grammar、JSON schema）

10. 输出处理层（Output Processing）

关键文件：

  • vllm/v1/engine/output_processor.py - 输出处理器

主要功能：

  1. Token 解码：将 token IDs 解码为文本
  2. Logprobs 处理：处理 logprobs 信息
  3. 停止条件检查：检查是否达到停止条件（EOS、max_tokens、stop strings）
  4. 流式输出：支持流式返回结果

关键代码位置：

class RequestOutputCollector:
    """
    Collects streamed RequestOutputs per individual request,
    for hand-off to the consuming asyncio generate task.
    """
    ...

说明：

  • OutputProcessor 处理模型输出，生成 RequestOutput
  • 支持流式输出：每个新 token 都会生成一个 RequestOutput
  • 最终输出通过 RequestOutputCollector 返回给 API 层

11. KV Connector（可选，用于分布式场景）

关键文件：

  • vllm/v1/worker/kv_connector_model_runner_mixin.py - KV Connector 混入类
  • vllm/config/kv_transfer.py - KV Transfer 配置
11.1 功能说明

主要功能：

  1. KV Cache 传输：在多个 worker 之间传输 KV cache
  2. Offloading：将 KV cache offload 到 CPU 或其他设备
  3. 同步：同步分布式场景下的 KV cache 状态

关键代码位置：

class KVConnectorModelRunnerMixin:
    @staticmethod
    def maybe_setup_kv_connector(scheduler_output: "SchedulerOutput"):
        # Update KVConnector with the KVConnector metadata forward().
        if has_kv_transfer_group():
            kv_connector = get_kv_transfer_group()
            ...
            # Background KV cache transfers happen here.
            kv_connector.start_load_kv(get_forward_context())
    ...
11.2 启动配置

配置参数：

KV Connector 通过 --kv-transfer-config 参数配置，该参数接受 JSON 格式的配置：

{
  "kv_connector": "p2p-nccl",  // 或 "offloading"
  "kv_parallel_size": 2,        // 并行实例数量
  "kv_role": "kv_producer",     // 或 "kv_consumer", "kv_both"
  "kv_rank": 0,                 // 实例 rank（0 为 prefill，1 为 decode）
  "kv_ip": "127.0.0.1",         // KV connector IP
  "kv_port": 14579,             // KV connector 端口
  "kv_buffer_device": "cuda",   // 或 "cpu"
  "kv_buffer_size": 1000000000   // 缓冲区大小（字节）
}

启动命令示例：

  1. KV Connector (P2P)：
# Prefill 节点
python -m vllm.entrypoints.openai.api_server \
    --model <model> \
    --kv-transfer-config '{"kv_connector": "p2p-nccl", "kv_parallel_size": 2, "kv_role": "kv_producer", "kv_rank": 0, "kv_ip": "127.0.0.1", "kv_port": 14579}'

# Decode 节点
python -m vllm.entrypoints.openai.api_server \
    --model <model> \
    --kv-transfer-config '{"kv_connector": "p2p-nccl", "kv_parallel_size": 2, "kv_role": "kv_consumer", "kv_rank": 1, "kv_ip": "127.0.0.1", "kv_port": 14579}'
  1. KV Pool (Offloading)：
python -m vllm.entrypoints.openai.api_server \
    --model <model> \
    --kv-transfer-config '{"kv_connector": "offloading", "kv_parallel_size": 1, "kv_buffer_device": "cpu", "kv_buffer_size": 1000000000}'
  1. P/D 分离（Prefill/Decode Disaggregation）：
# Prefill 节点
python -m vllm.entrypoints.openai.api_server \
    --model <model> \
    --kv-transfer-config '{"kv_connector": "p2p-nccl", "kv_parallel_size": 2, "kv_role": "kv_producer", "kv_rank": 0}' \
    --enable-prefill-decoding-disaggregation

# Decode 节点
python -m vllm.entrypoints.openai.api_server \
    --model <model> \
    --kv-transfer-config '{"kv_connector": "p2p-nccl", "kv_parallel_size": 2, "kv_role": "kv_consumer", "kv_rank": 1}' \
    --enable-prefill-decoding-disaggregation

节点配置说明：

  • 每个节点需要指定相同的 kv_connector 和 kv_parallel_size
  • 节点类型由 kv_role 和 kv_rank 决定：
    • kv_producer + kv_rank=0：Prefill 节点
    • kv_consumer + kv_rank=1：Decode 节点
    • kv_both：混合节点（同时处理 prefill 和 decode）

关键代码位置：

@config
@dataclass
class KVTransferConfig:
    """Configuration for distributed KV cache transfer."""
    
    kv_connector: str | None = None
    """The KV connector for vLLM to transmit KV caches between vLLM instances."""
    
    kv_parallel_size: int = 1
    """The number of parallel instances for KV cache transfer. For
    P2pNcclConnector, this should be 2."""
    
    kv_role: KVRole | None = None
    """Whether this vLLM instance produces, consumes KV cache, or both."""
    
    kv_rank: int | None = None
    """The rank of this vLLM instance in the KV cache transfer."""
    ...

说明：

  • 在分布式推理或 KV cache offloading 场景中使用
  • 管理 KV cache 在不同设备/进程间的传输
  • 支持 P/D 分离架构，将 prefill 和 decode 分离到不同节点

完整流程图

┌─────────────────────────────────────────────────────────────┐
│ 1. API 服务器层                                              │
│    vllm/entrypoints/openai/api_server.py                    │
│    - 接收 HTTP 请求 (/v1/chat/completions)                  │
│    - 路由到 OpenAIServingChat                               │
└────────────────────┬────────────────────────────────────────┘
                     │
                     ▼
┌─────────────────────────────────────────────────────────────┐
│ 2. AsyncLLM 层                                              │
│    vllm/v1/engine/async_llm.py                              │
│    - add_request() 添加请求                                  │
│    - 创建 RequestOutputCollector                            │
└────────────────────┬────────────────────────────────────────┘
                     │
                     ▼
┌─────────────────────────────────────────────────────────────┐
│ 3. 输入处理层                                                │
│    vllm/v1/engine/input_processor.py                        │
│    - Tokenization（文本 → token IDs）                       │
│    - 多模态处理                                              │
│    - 参数验证                                                │
└────────────────────┬────────────────────────────────────────┘
                     │
                     ▼
┌─────────────────────────────────────────────────────────────┐
│ 4. EngineCore 层                                            │
│    vllm/v1/engine/core.py                                   │
│    - 管理请求队列                                            │
│    - 协调调度和执行                                          │
│    - step() 循环：调度 → 执行 → 采样 → 更新                  │
└────────────────────┬────────────────────────────────────────┘
                     │
         ┌───────────┴───────────┐
         │                       │
         ▼                       ▼
┌──────────────────┐    ┌──────────────────┐
│ 5. 调度器        │    │ 6. KV Cache 管理  │
│ Scheduler        │◄───┤ BlockPool         │
│ - Token Budget   │    │ - 分配 blocks     │
│ - 批处理构建     │    │ - 前缀缓存        │
│ - 状态转换       │    │ - 多级缓存        │
│ - Chunked Prefill│    │ - Block 回收      │
└────────┬─────────┘    └──────────────────┘
         │
         ▼
┌─────────────────────────────────────────────────────────────┐
│ 7. 模型执行层                                                │
│    vllm/v1/worker/gpu_model_runner.py                       │
│    - prepare_input() 准备输入                                │
│    - execute_model() 执行前向传播                            │
│    - KV Cache 写入（reshape_and_cache kernel）              │
│    - 支持 CUDA Graph、Flash Attention                       │
└────────────────────┬────────────────────────────────────────┘
                     │
                     ▼
┌─────────────────────────────────────────────────────────────┐
│ 8. Chunked Prefill（可选）                                   │
│    - 长 prompt 分块处理                                      │
│    - 与 decode 请求混合批处理                                │
└────────────────────┬────────────────────────────────────────┘
                     │
                     ▼
┌─────────────────────────────────────────────────────────────┐
│ 9. 采样层                                                    │
│    vllm/v1/worker/gpu_model_runner.py                       │
│    - sample_tokens() 采样下一个 token                        │
│    - 应用 temperature、top-p、grammar 等                     │
└────────────────────┬────────────────────────────────────────┘
                     │
                     ▼
┌─────────────────────────────────────────────────────────────┐
│ 10. 输出处理层                                               │
│     vllm/v1/engine/output_processor.py                     │
│     - 解码 token IDs → 文本                                  │
│     - 处理 logprobs                                         │
│     - 检查停止条件                                           │
│     - 生成 RequestOutput                                    │
└────────────────────┬────────────────────────────────────────┘
                     │
                     ▼
┌─────────────────────────────────────────────────────────────┐
│ 11. 返回结果                                                 │
│     - RequestOutputCollector 收集输出                        │
│     - 流式返回给 API 层                                      │
│     - 最终返回给客户端                                       │
└─────────────────────────────────────────────────────────────┘

关键数据结构

EngineCoreRequest

  • request_id: 请求 ID
  • prompt_token_ids: Prompt 的 token IDs
  • sampling_params: 采样参数
  • lora_request: LoRA 适配器请求（可选）

SchedulerOutput

  • scheduled_seq_groups: 被调度的序列组
  • num_scheduled_tokens: 调度的 token 总数
  • blocks_to_swap_in/out: 需要换入/换出的 KV cache blocks
  • attn_metadata: Attention 元数据

ModelRunnerOutput

  • sampler_output: 采样输出（token IDs、logprobs）
  • hidden_states: 隐藏状态（用于下一次迭代）
  • kv_connector_output: KV connector 输出（如果有）

RequestOutput

  • request_id: 请求 ID
  • outputs: 输出列表（每个可能包含多个候选）
  • finished: 是否完成
  • metrics: 性能指标

性能优化要点

  1. 连续批处理（Continuous Batching）：动态调整批处理大小，新请求可立即加入
  2. PagedAttention：高效的 KV cache 管理，使用固定大小的 blocks
  3. CUDA Graph：减少 kernel 启动开销
  4. Flash Attention：高效的 attention 计算
  5. 前缀缓存（Prefix Caching）：共享相同前缀的 KV cache，减少重复计算
  6. Chunked Prefill：长 prompt 分块处理，与 decode 请求混合批处理
  7. 多级缓存：本地缓存 + 前缀缓存 + 外部缓存（KV Connector）
  8. 异步调度（Async Scheduling）：重叠计算和通信

总结

vLLM 的推理流程是一个高度优化的流水线，从 HTTP 请求到最终输出，经过多个精心设计的模块：

  1. API 层：接收和路由请求
  2. 输入处理：Tokenization 和验证
  3. 调度：智能选择请求和分配资源
  4. 执行：高效的模型推理
  5. 采样：生成下一个 token
  6. 输出处理：解码和格式化结果

每个模块都有其特定的职责，共同构成了 vLLM 高性能推理系统。

Logo
2048 AI社区

有“AI”的1024 = 2048，欢迎大家加入2048 AI社区

更多推荐

cover

AI动画工具

avatar 2048 AI社区
cover

AI编辑器的兴起：如何用好AI编辑器解决实际问题？

avatar 2048 AI社区
cover

在线简历工具越来越多，10个简历制作免费软件实测排行，这几个更顺手

avatar 2048 AI社区