vLLM 推理全流程深度解析

vLLM 使用 fastapi + uvicorn 启动服务。请求首先命中 vllm/entrypoints/openai/api_server.py 中的路由：

@router.post("/v1/chat/completions")
async def create_chat_completion(request: ChatCompletionRequest, raw_request: Request):
    # ...
    generator = await openai_serving_chat.create_chat_completion(
        request, raw_request)
    # ...

这里 openai_serving_chat 是 OpenAIServingChat 实例，负责将 OpenAI 格式请求转换为 vLLM 内部格式。

在 protocol.py 中，ChatCompletionRequest Pydantic 模型对字段做类型校验。关键字段包括：

• messages → 后续会被 tokenizer 的 apply_chat_template 处理
• temperature, top_p, top_k → 传给 SamplingParams
• max_tokens → 限制输出长度
• stream → 决定是返回 SSE stream 还是完整 JSON
• stop / stop_token_ids → 生成终止条件

vLLM 启动时已通过 TokenizerGroup 加载了 HuggingFace tokenizer。当请求到达时，调用链如下：

# vllm/entrypoints/openai/serving_chat.py
prompt = self.tokenizer.apply_chat_template(
    conversation=request.messages,
    tokenize=False,
    add_generation_prompt=True,
    chat_template=request.chat_template  # 用户可覆盖
)

以 Llama-2-chat 为例，apply_chat_template 会将 messages 展开为：

<s>[INST] <<SYS>>
You are a helpful assistant.
<</SYS>>

请解释 vLLM 的 PagedAttention 原理 [/INST]

注意 <s> 是 bos_token，[/INST] 是角色分隔符，末尾没有 answer 部分，因为 add_generation_prompt=True 只添加助手的起始标记，让模型继续生成。

# vllm/transformers_utils/tokenizer.py
input_ids = self.tokenizer.encode(prompt, add_special_tokens=True)

这一步的内部细节：

1. 如果 tokenizer 使用 BPE（如 GPT-2/LLaMA），文本先被预分词（regex 拆成单词/子词），然后查 merge table 逐步合并成 subword token。
2. 如果 tokenizer 使用 SentencePiece（如 Llama），文本直接被编码为 token id 序列，<s> 对应 id 1。
3. 得到的是一个 Python list[int]，例如 [1, 518, 25580, 29962, 3532, 14816, 29903, 6778, ...]

from vllm import SamplingParams
sampling_params = SamplingParams(
    temperature=0.7,
    top_p=1.0,
    top_k=-1,  # 默认禁用
    max_tokens=512,
    stop=[""],  # 默认 eos
    presence_penalty=0.0,
    frequency_penalty=0.0,
    repetition_penalty=1.0,
)

SamplingParams 还会根据 temperature 做数值校验，例如 temperature=0 时会自动进入 greedy decode 模式。

vLLM 内部使用 Sequence 对象表示一条序列（一个对话/生成轨迹）。对于单次生成（n=1）：

# vllm/sequence.py
seq = Sequence(
    seq_id=next(seq_counter),
    prompt=prompt,
    prompt_token_ids=input_ids,
    block_size=engine_config.block_size,  # 通常为 16
)

然后包装成 SequenceGroup：

seq_group = SequenceGroup(
    request_id=str(uuid.uuid4()),
    seqs=[seq],
    sampling_params=sampling_params,
    arrival_time=time.time(),
)

SequenceGroup 是调度（Scheduler）的基本单元。当用户设置 n=4（一次生成 4 个候选）时，这里会创建 4 个 Sequence 对象，它们共享同一个 prompt 的 KV Cache。

# vllm/engine/async_llm_engine.py
async def add_request(
    self,
    request_id: str,
    prompt: Optional[str],
    sampling_params: SamplingParams,
    prompt_token_ids: Optional[List[int]] = None,
    arrival_time: Optional[float] = None,
) -> AsyncStream:
    stream = AsyncStream()
    self._request_streams[request_id] = stream
    await self.engine.add_request_async(
        request_id, prompt, sampling_params, prompt_token_ids, arrival_time)
    return stream

AsyncStream 是一个异步迭代器，后续前端通过它逐 token 接收流式结果。真正的同步引擎 LLMEngine 把 seq_group 送入 Scheduler 的等待队列。

# vllm/core/scheduler.py
def schedule(self) -> Tuple[SchedulerOutputs, List[int]]:
    # 1. 检查是否可以做 prefix caching 的 prompt 复用（可选优化）
    # 2. 处理已完成的序列，释放 block
    # 3. 尝试将 waiting 队列的 seq_group 移入 running
    # 4. 尝试将 swapped 队列（被逐出到 CPU 的）移回 running
    # 5. 根据 max_num_seqs / max_num_batched_tokens 截断 batch
    # 6. 返回 SchedulerOutputs

新到达的请求初始状态为 SequenceStatus.WAITING。Scheduler 会尝试为它分配物理 KV block：

num_required_blocks = len(seq_group.prompt_token_ids) // block_size + 1

BlockManager 调用 allocate(seq_group)，从空闲 block 池中分配逻辑 block 到物理 block 的映射。

每个 Sequence 内部维护一个 block_table: List[int]，记录其 KV Cache 占用的物理块索引。例如：

# prompt 长度为 35，block_size=16
# 需要 3 个物理块
block_table = [12, 45, 89]

其中块 12 存 tokens 0-15 的 K/V，块 45 存 tokens 16-31，块 89 存 tokens 32-35（只使用了前 4 个 slot）。

这些 block 是全局显存池中的索引，vLLM 启动时根据 GPU 显存大小和 block_size 预分配了 num_gpu_blocks 个块。例如 A100 80GB 加载 7B 模型后，可能剩余约 70GB 用于 KV Cache，若 block_size=16, head_size=128, num_layers=32, num_kv_heads=32, fp16，则每块显存：

block_mem = 2 (K+V) * num_layers * num_kv_heads * head_size * block_size * sizeof(half)
          = 2 * 32 * 32 * 128 * 16 * 2 bytes ≈ 8.4 MB

70GB / 8.4MB ≈ 8500 个 block，可服务大量并发。

对于已经在 RUNNING 状态的序列，每个 decoding step 只产生 1 个新 token，因此 KV Cache 只需要新增 1 个 slot。Scheduler 检查当前物理块是否还有空 slot：

if 当前最后一个物理块还有空位:
    直接复用
else:
    从空闲池分配一个新的物理块，append 到 block_table

如果显存不足，Scheduler 会触发 Swapping：将部分序列的 KV block 逐出到 CPU 内存（cpu_allocator），或触发 Preemption：重新计算被踢出序列的 KV Cache（cost 高，尽量避免）。

最终，schedule() 返回：

@dataclass
class SchedulerOutputs:
    # 本轮要执行的 seq_groups（可能混合 prefill + decode）
    scheduled_seq_groups: List[SequenceGroup]
    # 被忽略的请求（如 prompt 过长）
    ignored_seq_groups: List[SequenceGroup]
    # 各 seq_group 的 token 数量统计
    num_batched_tokens: int
    # 本轮是新 prefill 还是纯 decode
    blocks_to_swap_in: List[Tuple[int, int]]   # CPU -> GPU
    blocks_to_swap_out: List[Tuple[int, int]]  # GPU -> CPU
    blocks_to_copy: List[Tuple[int, int]]      # 用于并行采样/Beam search 共享

这是连接调度层与执行层的关键数据结构。

# vllm/engine/llm_engine.py
def step(self) -> List[RequestOutput]:
    # 1. 调用 scheduler.schedule()
    seq_group_metadata_list, scheduler_outputs = self.scheduler.schedule()
    
    # 2. 如果为空，直接返回
    if not scheduler_outputs.scheduled_seq_groups:
        return []
    
    # 3. 执行模型推理
    output = self.model_executor.execute_model(
        seq_group_metadata_list,
        scheduler_outputs.blocks_to_swap_in,
        scheduler_outputs.blocks_to_swap_out,
        scheduler_outputs.blocks_to_copy,
    )
    
    # 4. 处理输出（采样、更新序列状态）
    request_outputs = self._process_model_outputs(output, scheduler_outputs)
    return request_outputs

vLLM 支持 TP（Tensor Parallelism，张量并行）和 PP（Pipeline Parallelism，流水线并行）。对于单卡或 TP，通常使用 GPUExecutor：

# vllm/executor/gpu_executor.py
def execute_model(self, seq_group_metadata_list, ...):
    # 如果是 TP，通过 torch.distributed 广播 input tensors 到各 rank
    # 各 Worker 的 model_runner 执行 forward
    return self.driver_worker.execute_model(seq_group_metadata_list, ...)

这是整个推理链路中最容易让人困惑的环节。Scheduler 输出的是 Python 对象，而 GPU 需要密集的 Tensor。ModelRunner 负责做这一转换：

# vllm/worker/model_runner.py
def prepare_input_tensors(
    self,
    seq_group_metadata_list: List[SequenceGroupMetadata],
) -> Tuple[torch.Tensor, torch.Tensor, InputMetadata, ...]:

核心逻辑拆解：

1. 收集 input token ids：对每个 seq_group，如果是 prefill 阶段，取全部 prompt_token_ids；如果是 decode 阶段，取最后一个生成的 token id。
2. 拼接为 1D input_ids tensor：vLLM 将所有序列的当前 step 要处理的 token 拼接成一个 1D LongTensor。例如 batch 中有 2 个 prefill（长度 35, 20）和 3 个 decode（各 1 个 token），则 input_ids.shape = [35+20+1+1+1] = [58]。
3. 构建 input_positions：记录每个 token 在各自序列中的位置索引，用于 RoPE / ALiBi 位置编码。
4. 构建 slot_mapping：一一映射每个 token 应该写入 KV Cache 的物理地址（block_number * block_size + offset）。这是 PagedAttention 的核心。
5. 构建 seq_lens / context_lens：记录每个序列的总长度（用于 attention mask / causal mask）。

假设本轮 batch 包含：

• Seq A: prefill, prompt_len=35, block_table=[12,45,89]
• Seq B: prefill, prompt_len=20, block_table=[7,22]
• Seq C: decode, gen_len=5（已生成5个token）, block_table=[3,8], 本轮输入 token 是第6个

input_ids       = [A0..A34 | B0..B19 | C5]          # shape [56]
positions       = [0..34   | 0..19   | 5]            # shape [56]
slot_mapping    = [12*16+0, ..., 12*16+15, 45*16+0, ...,
                   89*16+2, 7*16+0, ..., 8*16+3,
                   8*16+4]                           # shape [56]
seq_lens        = [35, 20, 6]                        # 各序列当前总长度
context_lens    = [0, 0, 5]  # prefill 时 context_len=0（无过去 KV），decode 时为已生成长度

# vllm/model_executor/models/llama.py
class LlamaForCausalLM(nn.Module):
    def forward(self, input_ids, positions, kv_caches, input_metadata):
        hidden_states = self.model(input_ids, positions, kv_caches, input_metadata)
        logits = self.lm_head(hidden_states)
        return logits

注意这里的 kv_caches 并不是传统意义上每个序列一个的 KV 张量，而是全局预分配的显存池。

def forward(self, input_ids, positions, kv_caches, input_metadata):
    hidden_states = self.embed_tokens(input_ids)   # [num_tokens, hidden_size]
    
    for i, layer in enumerate(self.layers):
        hidden_states = layer(
            hidden_states,
            positions,
            kv_caches[i],        # 第 i 层的全局 KV cache
            input_metadata,
        )
    hidden_states = self.norm(hidden_states)
    return hidden_states

def forward(self, hidden_states, positions, kv_cache, input_metadata):
    # 1. Self Attention (RMSNorm -> Attention -> Residual)
    residual = hidden_states
    hidden_states = self.input_layernorm(hidden_states)
    hidden_states = self.self_attn(
        hidden_states, positions, kv_cache, input_metadata)
    hidden_states = residual + hidden_states
    
    # 2. FFN (RMSNorm -> MLP -> Residual)
    residual = hidden_states
    hidden_states = self.post_attention_layernorm(hidden_states)
    hidden_states = self.mlp(hidden_states)
    hidden_states = residual + hidden_states
    return hidden_states

# vllm/model_executor/models/llama.py 中的 attention 实现
class LlamaAttention(nn.Module):
    def forward(self, hidden_states, positions, kv_cache, input_metadata):
        qkv = self.qkv_proj(hidden_states)  # [num_tokens, 3 * num_heads * head_size]
        q, k, v = qkv.split(...)
        
        # 使用 Rotary Position Embedding (RoPE)
        q, k = self.rotary_emb(positions, q, k)
        
        # 写入 KV Cache（通过自定义 CUDA kernel）
        # 这里不是简单的 tensor indexing，而是调用 paged_attention_ops.reshape_and_cache
        reshape_and_cache(k, v, kv_cache, input_metadata.slot_mapping)
        
        # 计算 attention（prefill 使用 xFormers/FlashAttention，decode 使用 PagedAttention）
        if input_metadata.is_prompt:  # prefill phase
            output = self.attn(q, k, v, input_metadata)
        else:  # decode phase
            output = self.attn(q, kv_cache, input_metadata)
        
        output = self.o_proj(output)
        return output

这是 vLLM 自定义 CUDA 算子之一，源码位于 csrc/cache_kernels.cu：

// 伪代码逻辑
__global__ void reshape_and_cache_kernel(
    const scalar_t* __restrict__ key,     // [num_tokens, num_kv_heads, head_size]
    const scalar_t* __restrict__ value,   // [num_tokens, num_kv_heads, head_size]
    scalar_t* __restrict__ key_cache,     // [num_blocks, block_size, num_kv_heads, head_size]
    scalar_t* __restrict__ value_cache,   // 同上
    const int64_t* __restrict__ slot_mapping, // [num_tokens]
) {
    const int64_t token_idx = blockIdx.x;
    const int64_t slot_idx = slot_mapping[token_idx];
    if (slot_idx < 0) return;  // 被忽略的 token
    
    const int64_t block_idx = slot_idx / block_size;
    const int64_t block_offset = slot_idx % block_size;
    
    // 将 key/value 写入对应 block 的对应 offset
    // key_cache 布局: [num_blocks, block_size, num_kv_heads, head_size]
    // 通过指针计算直接写入，无 Python overhead
}

关键理解： kv_cache 的 shape 是全局的 [num_gpu_blocks, block_size, num_kv_heads, head_size]，而不是按 batch 或 sequence 划分。这使得不同序列可以任意共享 block（如 beam search 中多个候选序列共享 prompt 的 block）。

对于 prefill 阶段，序列需要计算自注意力（所有 token 之间两两因果 mask）。vLLM 默认使用 xFormers memory_efficient_attention 或 FlashAttention：

from xformers.ops import memory_efficient_attention

# q/k/v shape: [batch_size, num_heads, seq_len, head_size]
# attention_mask: causal mask（下三角为 0，上三角为 -inf）
out = memory_efficient_attention(q, k, v, attn_bias=LowerTriangularMask())

FlashAttention 则通过定制的 CUDA kernel 在 SRAM 中完成 softmax 和 matmul，避免将巨大的注意力矩阵写出到 HBM（全局显存），从而大幅减少显存带宽压力。

Decode 阶段每个序列只输入 1 个新 token，但需要 attend 到之前所有 token 的 KV。若使用标准 Attention，需要将整个 KV Cache 从显存读入计算，带宽成为瓶颈（memory-bound）。

PagedAttention 的核心创新在于：它将 Attention 计算按 block 分块，利用共享内存和寄存器优化，直接从 block 化的 KV Cache 中读取。CUDA kernel 签名类似：

// csrc/paged_attention.cu
__global__ void paged_attention_v1_kernel(
    scalar_t* __restrict__ out,           // [num_seqs, num_heads, head_size]
    const scalar_t* __restrict__ q,       // [num_seqs, num_heads, head_size]
    const scalar_t* __restrict__ k_cache, // [num_blocks, block_size, num_kv_heads, head_size]
    const scalar_t* __restrict__ v_cache, // 同上
    const int* __restrict__ block_tables, // [num_seqs, max_num_blocks_per_seq]
    const int* __restrict__ context_lens, // [num_seqs]
    const int num_kv_heads,
    const float scale,
    ...
)

每个 warp / thread block 负责计算一个 head 的一个 sequence 的输出。它通过 block_tables 查找到该序列的物理块列表，然后逐块加载 K/V 到共享内存，做 dot-product 得到 attention score，online softmax 计算加权求和。

由于 decode 时 seq_len 很长（可能 4096+），但 q 只有一个 token，PagedAttention 的性能优势极大：它将访存从 O(seq_len) 的随机/分散访问优化为按 block 的连续突发读取（coalesced memory access），配合 warp 并行和 KV cache 量化（如 KV cache int8），吞吐可提升数倍。

# vllm/model_executor/layers/sampler.py
class Sampler(nn.Module):
    def forward(self, logits: torch.Tensor, sampling_metadata: SamplingMetadata):
        # logits: [num_tokens, vocab_size]
        # 注意：对于 prefill 阶段，通常只需要对最后一个位置的 token 采样
        # 对于 decode 阶段，每个 token 都对应一个要采样的 logits
        
        # 1. Apply penalties (presence, frequency, repetition)
        logits = self._apply_penalties(logits, sampling_metadata)
        
        # 2. Temperature scaling
        logits = self._apply_temperature(logits, sampling_metadata)
        
        # 3. Top-p / Top-k filtering
        logits = self._apply_top_p_top_k(logits, sampling_metadata)
        
        # 4. Sampling
        probs = torch.softmax(logits, dim=-1)
        samples = torch.multinomial(probs, num_samples=1)
        return samples

logits = logits / temperature

temperature → 0 时，logits 差异被放大，softmax 趋近于 argmax（greedy）；temperature → ∞ 时，分布趋近均匀。

• Top-K：对 logits 排序，只保留概率最高的 K 个，其余设为 -inf。
• Top-P：对排序后的概率做 cumsum，找到最小的集合使得概率和 ≥ p，其余丢弃。

vLLM 中通常使用 efficient GPU kernel 完成这些操作，避免将巨大的 vocab_size tensor 回传到 CPU。

# Presence Penalty: 如果 token 已在序列中出现过，直接减去 penalty 值
# Frequency Penalty: 根据出现次数成比例减去 penalty
# Repetition Penalty: 直接对 logits 除以 penalty（大于 1 时抑制重复）

实现上，Sampler 会读取每个序列已生成的 token set / token counter，对 logits 做 inplace 修改。

samples = torch.multinomial(probs, num_samples=1)  # [num_tokens, 1]
next_token_ids = samples.squeeze(-1)                # [num_tokens]

对于 prefill 阶段，vLLM 通常使用 _get_logprobs 返回 prompt 的每个 token 的 logprob，而真正的 next token 只取 prefill 序列最后一个位置的采样结果。

# vllm/engine/llm_engine.py 中 _process_model_outputs 的逻辑
for seq_group, outputs in zip(scheduled_seq_groups, output_by_seq_group):
    for seq, output in zip(seq_group.seqs, outputs):
        # 1. 将新 token 追加到序列
        seq.append_token_id(output.token_id, output.logprobs)
        
        # 2. 检查是否达到终止条件
        if seq.get_len() >= seq_group.sampling_params.max_tokens:
            seq.status = SequenceStatus.FINISHED_LENGTH_CAPPED
        elif output.token_id in seq_group.sampling_params.stop_token_ids:
            seq.status = SequenceStatus.FINISHED_STOPPED
        elif output.token_id == tokenizer.eos_token_id:
            seq.status = SequenceStatus.FINISHED_STOPPED
        elif seq.get_output_text().endswith(seq_group.sampling_params.stop_str):
            seq.status = SequenceStatus.FINISHED_STOPPED

由于 reshape_and_cache 已经在 forward 中将当前 step 的 K/V 写入了全局 KV Cache pool，这里不需要额外操作。Sequence 只需要记录自己的 block_table 和当前长度 data_len 即可。

# vllm/engine/async_llm_engine.py
for request_output in request_outputs:
    self._request_streams[request_output.request_id].put(request_output)

RequestOutput 包含：

• request_id
• outputs: 每个采样分支的 CompletionOutput（含生成的 text、token_ids、logprobs、finish_reason）

为了流式输出，vLLM 不能等所有 token 生成完再 decode，而是逐 token 增量解码。但 Byte-Pair Encoding（BPE）有一个棘手问题：一个 Unicode 字符可能被拆成多个 byte-level token，如果只解码当前 token，可能得到乱码（如 �）。

vLLM 的 detokenize_incrementally 解决如下：

# vllm/transformers_utils/detokenizer.py
def detokenize_incrementally(self, new_token_id: int) -> str:
    self.token_ids.append(new_token_id)
    # 先尝试解码整个序列
    new_text = self.tokenizer.decode(self.token_ids, skip_special_tokens=True)
    # 与上一次完整文本对比，取出增量部分
    incremental_text = new_text[len(self.prev_text):]
    self.prev_text = new_text
    return incremental_text

虽然看起来是重新 decode 整个序列，但由于 tokenizer 通常有 cache（如 SentencePiece 的 DecodeIds），且序列长度有限，开销极小。对于中文，增量文本可能一次输出 1-3 个汉字。

# vllm/entrypoints/openai/api_server.py
async def stream_generator():
    async for request_output in generator:
        for output in request_output.outputs:
            ret = {
                "choices": [{
                    "delta": {"content": output.text},
                    "finish_reason": output.finish_reason,
                }]
            }
            yield f"data: {json.dumps(ret)}\n\n"
    yield "data: [DONE]\n\n"

客户端每次收到 SSE 事件，将 delta.content 追加到界面，用户看到"打字机效果"。