vLLM ModelRunner 与 Sampler 源码精读

ModelRunner（model_runner.py L106+）是 vLLM 推理的核心执行器，负责：

• 加载模型权重（load_model）
• 将 Scheduler 的调度结果转换为模型输入 tensor（prepare_input_tensors）
• 执行前向推理（execute_model）
• 管理 CUDA Graph 加速

class ModelRunner:
    def __init__(self, model_config, parallel_config, scheduler_config,
                 device_config, load_config, lora_config,
                 kv_cache_dtype="auto", is_driver_worker=False,
                 vision_language_config=None):
        # ...
        self.graph_runners: Dict[int, CUDAGraphRunner] = {}
        self.graph_memory_pool = None

        self.max_seq_len_to_capture = (
            self.model_config.max_seq_len_to_capture)
        self.pin_memory = is_pin_memory_available()
        self.kv_cache_dtype = kv_cache_dtype

BatchType 枚举定义了三种 batch 类型：

• PREFILL：全部是 prefill 请求
• DECODE：全部是 decode 请求
• MIXED：prefill 和 decode 混合（Chunked Prefill 场景）

ModelRunner 使用两个 NamedTuple 来组织 prefill 和 decode 的输入：

class PreparePromptMetadata(NamedTuple):
    input_tokens: List[int]       # 扁平化的 token ID 列表
    input_positions: List[int]    # 每个 token 的位置编码
    attn_metadata: Optional[AttentionMetadataPerStage]
    seq_lens: List[int]           # 每个序列的总长度
    query_lens: List[int]         # 每个序列本次 prefill 的 token 数
    lora_index_mapping: List[int]
    lora_prompt_mapping: List[int]
    lora_requests: Set[LoRARequest]
    multi_modal_input: Optional[torch.Tensor]  # 视觉输入
    slot_mapping: List[int]       # token → KV Cache slot 的映射

class PrepareDecodeMetadata(NamedTuple):
    input_tokens: List[int]       # 每个序列只有 1 个 token
    input_positions: List[int]    # 每个 token 的位置
    attn_metadata: Optional[AttentionMetadata]
    lora_index_mapping: List[int]
    lora_prompt_mapping: List[int]
    lora_requests: Set[LoRARequest]
    slot_mapping: List[int]       # token → KV Cache slot

_prepare_prompt 遍历所有 prefill 类型的 SequenceGroupMetadata，为每个序列组装输入：

1. Token IDs：从 seq_data.get_token_ids() 获取完整 token 序列，按 token_chunk_size 截取（支持 chunked prefill）
2. Positions：从 computed_len（已计算的 token 数）开始编号，支持 prefix caching 跳过已计算的前缀
3. Slot Mapping：通过 block_table 计算每个 token 对应的 KV Cache slot
4. Attention Metadata：包括 seq_lens、context_lens（prefix 长度）、block_tables 等

关键区别于 decode：prefill 的 input_tokens 包含多个 token（整个 prompt 或一个 chunk），而 decode 只有 1 个 token。

_prepare_decode 处理正在生成的序列。每个序列只输入最后一个 token（刚生成的），因为之前的 KV Cache 已经缓存：

• input_tokens：每个序列贡献 1 个 token
• input_positions：当前序列长度 - 1（最后一个位置）
• block_tables：完整的块表，供 PagedAttention kernel 查找历史 KV

Decode batch 的 input_tokens 长度 = 序列数，这使得 decode 的计算量远小于 prefill。但 attention 仍然需要访问所有历史 KV Cache（通过 block_tables）。

核心执行逻辑：

1. prepare_input_tensors：将 SequenceGroupMetadata 列表转换为模型需要的 tensor
2. 模型前向：调用 model.forward()，传入 input_ids、positions、kv_caches（来自 CacheEngine 的 GPU cache）和 attn_metadata
3. CUDA Graph 加速：对于纯 decode batch，如果 batch size 在预捕获的 graph 列表中，使用 CUDA Graph 避免 kernel launch 开销
4. 采样：Sampler 层根据 logits 和 SamplingMetadata 生成下一个 token

CUDA Graph 将一系列 GPU 操作捕获为一个图，后续执行时跳过 CPU 侧的 kernel launch 调度，显著降低延迟。

_BATCH_SIZE_ALIGNMENT = 8
# Capture graphs for token size 1, 2, 4, 8, 16, 24, 32, 40, ..., 256.
_BATCH_SIZES_TO_CAPTURE = [1, 2, 4] + [
    _BATCH_SIZE_ALIGNMENT * i for i in range(1, 33)
]

使用条件：

• 只在 decode 阶段使用（prefill 的 input 长度变化太大，无法预捕获）
• batch size 必须在 _BATCH_SIZES_TO_CAPTURE 列表中
• 序列长度不超过 max_seq_len_to_capture
• 不使用 LoRA 时效果最好

如果实际 batch size 不在预捕获列表中，会向上取整到最近的值（通过 _get_graph_batch_size），多余的位置用 padding 填充。

Sampler（sampler.py L26-120+）是模型的最后一层，将 logits 转换为生成的 token。

def forward(self, logits, sampling_metadata):
    """
    Args:
        logits: (num_tokens, vocab_size)
        sampling_metadata: Metadata for sampling.
    """
    # 1. Apply min_tokens penalty
    logits = _apply_min_tokens_penalty(logits, sampling_metadata)

    # 2. Prepare sampling tensors (pinned memory)
    (sampling_tensors, do_penalties,
     do_top_p_top_k, do_min_p) = SamplingTensors.from_sampling_metadata(
         sampling_metadata, vocab_size, logits.device, logits.dtype)

    # 3. Apply presence and frequency penalties
    if do_penalties:
        logits = _apply_penalties(logits, ...)

    # 4. Apply temperature scaling
    logits.div_(sampling_tensors.temperatures.unsqueeze_(dim=1))

    # 5. Apply top-k and top-p truncation
    if do_top_p_top_k:
        logits = _apply_top_k_top_p(logits, ...)

    # 6. Apply min-p truncation
    if do_min_p:
        logits = _apply_min_p(logits, ...)

    # 7. Compute probabilities and log probabilities
    probs = torch.softmax(logits, dim=-1, dtype=torch.float)
    logprobs = torch.log_softmax(logits, dim=-1, dtype=torch.float)

    # 8. Sample the next tokens
    sample_results = _sample(probs, logprobs, sampling_metadata, ...)

vLLM 支持三种基本采样方式（SamplingType）：

Penalty 参数

• presence_penalty：对已出现过的 token 施加固定惩罚，鼓励多样性
• frequency_penalty：惩罚与出现次数成正比，越频繁惩罚越大
• repetition_penalty：对已出现 token 的 logits 乘以惩罚因子
• temperature：控制分布的锐利程度。<1 更确定性，>1 更随机
• top_k：只保留概率最高的 k 个 token
• top_p：只保留累积概率达到 p 的最小 token 集合
• min_p：过滤掉概率低于 max_prob × min_p 的 token

Sampler 的输出是 SamplerOutput，包含每个序列组的采样结果：

SamplerOutput(
    outputs: List[CompletionSequenceGroupOutput(
        samples: List[SequenceOutput(
            parent_seq_id: int,
            output_token: int,
            logprobs: Dict[int, Logprob],
        )],
        prompt_logprobs: Optional[PromptLogprobs],
    )],
    sampled_token_probs: Optional[torch.Tensor],
    sampled_token_ids: Optional[torch.Tensor],
)

每个 SequenceOutput 包含：

• parent_seq_id：beam search 中的父序列 ID
• output_token：采样得到的 token ID
• logprobs：top-k log probabilities（如果请求了 logprobs）

这个 SamplerOutput 会返回给 LLMEngine.step()，引擎将 token append 到对应序列中，更新序列状态，然后进入下一轮调度。