vLLM 源码完全拆解（九）：Sampler——从 Logits 到 Token 的采样管线与输出回传

为什么需要 SamplingMetadata#

在 vLLM 的一轮 step 中，batch 里可能同时存在十几个甚至几十个请求。每个请求有自己的 SamplingParams，包括：

1
class SamplingParams:
2
    temperature: float = 1.0        # 温度参数（0 表示 greedy）
3
    top_k: int = -1                # Top-K 过滤（-1 表示禁用）
4
    top_p: float = 1.0             # Top-P 过滤（1.0 表示禁用）
5
    min_p: float = 0.0             # Min-P 过滤
6
    frequency_penalty: float = 0.0 # 频率惩罚
7
    presence_penalty: float = 0.0  # 存在惩罚
8
    repetition_penalty: float = 1.0 # 重复惩罚
9
    seed: int | None = None        # 随机种子
10
    max_tokens: int = 16           # 最大生成 token 数
11
    stop: str | list[str] | None   # stop strings
12
    stop_token_ids: list[int] | None
13
    logprobs: int | None = None    # 返回多少个 logprobs
14
    ...

不同请求的这些参数几乎不可能完全一致。设想一个场景：请求 A（聊天补全，temperature=0.7，top_p=0.9），请求 B（代码补全，temperature=0，greedy），请求 C（创意写作，temperature=1.2，top_k=50）。Sampler 需要在 GPU 上一次性执行所有三个请求的采样，但每个请求又需要不同的 logits 处理方式。

直接把每一份 SamplingParams 传给 GPU Kernel 是不可行的——Kernel 无法接受 Python 对象的动态列表。vLLM 的解决方案是 SamplingMetadata：把所有请求的采样参数打包成一组等长张量（每请求一个标量，[batch_size] 形状），让 GPU Kernel 按索引读取对应请求的参数。

SamplingMetadata 的结构#

在 V1 中，SamplingMetadata（位于 vllm/v1/sample/metadata.py）的核心字段如下：

1
@dataclass
2
class SamplingMetadata:
3
    temperature: torch.Tensor       # [batch_size] float32
4
    top_k: torch.Tensor             # [batch_size] int32
5
    top_p: torch.Tensor             # [batch_size] float32
6
    min_p: torch.Tensor             # [batch_size] float32
7

8
    frequency_penalty: torch.Tensor # [batch_size] float32
9
    presence_penalty: torch.Tensor  # [batch_size] float32
10
    repetition_penalty: torch.Tensor # [batch_size] float32
11

12
    seed: torch.Tensor              # [batch_size] int64
13
    max_tokens: torch.Tensor        # [batch_size] int32
14

15
    all_greedy: bool                # 是否所有请求都是 greedy
16
    no_penalties: bool              # 是否所有请求都没有惩罚项
17

18
    # 输出规格
19
    output_token_ids: list[list[int]]  # 每个请求已经生成的 token ID（用于重复惩罚计算）
20

21
    # 惩罚相关的 bookkeeping
22
    prompt_token_ids: torch.Tensor | None  # prompt 的 token ID（已扁平化）
23
    num_prompt_tokens: torch.Tensor | None # 每个请求的 prompt token 数

注意 all_greedy 和 no_penalties 两个标志——它们是 vLLM 的快速路径优化。如果所有请求都是 temperature=0（greedy），Sampler 可以直接走 argmax，跳过随机数生成、top-k/top-p 排序等昂贵操作。如果没有任何惩罚项，就跳过频率/存在/重复惩罚的 logits 修正——在典型的生产场景中这非常常见。

如果所有请求的 temperature 为 0 且没有惩罚项，Sampler 可以只用一次高效的 argmax Kernel 完成全部采样。

SamplingMetadata 从哪来#

SamplingMetadata 的构造发生在 GPUModelRunner 的输入准备阶段。每次有新的调度输出时，InputBatch 中增删请求，SamplingMetadata 被更新为当前活跃请求的参数打包：

1
# 伪代码：SamplingMetadata 的更新逻辑
2
def _update_sampling_metadata(self, input_batch):
3
    req_ids = input_batch.req_ids  # 当前活跃的请求 ID 列表
4
    sampling_params = [self.requests[r].sampling_params for r in req_ids]
5

6
    self.sampling_metadata.temperature = tensor([
7
        sp.temperature for sp in sampling_params
8
    ])
9
    self.sampling_metadata.top_k = tensor([
10
        sp.top_k for sp in sampling_params
11
    ])
12
    self.sampling_metadata.all_greedy = all(
13
        sp.temperature == 0 for sp in sampling_params
14
    )
15
    # ... 其他字段类似

all_greedy 是引擎快速判断时常用的优化位。在生产中，大量请求（特别是代码补全、指令遵循）使用 temperature=0，这使得 argmax 路径成为默认热路径。

管线架构#

在 vllm/v1/sample/logits_processor/ 下，每种变换被实现为一个独立的类，都遵循统一的接口：

1
class LogitsProcessor(ABC):
2
    @abstractmethod
3
    def apply(self, logits: torch.Tensor,
4
              sampling_metadata: SamplingMetadata) -> torch.Tensor:
5
        ...

Sampler 内部维护一个有序的 Processor 列表，采样时按序调用：

1
# Sampler.sample() 中 LogitsProcessor 管线的调用
2
def _apply_logits_processors(self, logits, sampling_metadata):
3
    for processor in self.logits_processors:
4
        logits = processor.apply(logits, sampling_metadata)
5
    return logits

处理顺序经过精心选择，以保证数学正确性和数值稳定性。接下来逐个拆解每种核心变换。

2.1 Temperature（温度缩放）#

作用：把所有 logits 除以 temperature，控制 softmax 后分布的尖锐程度。

给定 temperature $T$ 和原始 logits $\ell \in \mathbb{R}^V$（$V$ 为词表大小），

当 $T \to 0$ 时，最大的 logit 被无穷放大——等效于 argmax（greedy decoding）。当 $T=1$ 时，保持原始 logits。当 $T > 1$ 时（如 $T=1.2$），logits 之间的差异被压缩，softmax 后的分布更平滑，采样更多样。

实现细节：temperature 缩放是最简单的变换，只是一个逐元素除法。但在批量异构场景中，不同请求需要不同的 temperature，而 logits 是 [num_tokens, vocab_size] 的连续张量——不能简单的一个 batch 一个除法。

vLLM 的解决方式是把每个请求的 temperature 广播到它在本轮中占用的 token 数：

1
# 伪代码：Temperature 的批量应用
2
def apply_temperature(logits, sampling_metadata):
3
    # temperatures_per_token: 从 sampling_metadata.temperature
4
    # 按 query_start_loc 扩展到每个 token
5
    temps = temperatures_per_token  # [num_tokens]
6
    logits.div_(temps.unsqueeze(-1))  # 广播到 vocab_size 维度
7
    return logits

当 sampling_metadata.all_greedy 为 True（所有 temperature=0）时，这一步骤被跳过——不是因为”等于 1”，而是因为 greedy 路径根本不需要经过 processor 管线，直接走 argmax。

为什么 $T=0$ 不能直接除：除零显然无效。vLLM 对 temperature=0 的处理是进入独立的 greedy 路径，完全不经过 temperature processor。

2.2 Top-K 过滤#

作用：只保留 logits 中最大的 K 个值，其余设为 $-\infty$（在 softmax 中等效于概率为 0）。

算法（对每个 token 位置独立执行）：

1. 找到 logits 中第 K 大的值（通过 torch.topk(logits, k).values[:, -1:]）
2. 构造 mask：logits < threshold
3. logits[mask] = -inf

为什么是 $-\infty$ 而不是 0：直接设 0 在 softmax 后仍然会获得非零概率——$e^0 = 1$，在所有 $e^{\ell_i}$ 的和中占据一定份额。只有 $\ell_i = -\infty$ 能让 $e^{-\infty} = 0$，彻底排除该 token。

批量异构 Top-K：不同请求有不同的 K 值。vLLM 需要为每个 token 位置取对应的 K，计算各自的阈值，然后单独 mask：

1
# 伪代码：批量异构 Top-K
2
def apply_top_k(logits, sampling_metadata):
3
    k_per_token = sampling_metadata.top_k_per_token()  # [num_tokens]
4
    # 在 CUDA/Triton kernel 中为每个 token 找它的第 K 大值
5
    thresholds = compute_per_token_topk_threshold(logits, k_per_token)
6
    mask = logits < thresholds.unsqueeze(-1)
7
    logits.masked_fill_(mask, float("-inf"))
8
    return logits

当某请求的 top_k <= 0 时，该 token 位置跳过 Top-K 过滤。

2.3 Top-P（Nucleus Sampling）#

作用：从概率最高的 token 开始累加，保留直到累积概率达到 p 的最小集合。由 Holtzman 等人于 2020 年在 ICLR 论文中提出。

算法（对每个 token 位置独立执行）：

1. 对 logits 做 softmax 得到概率分布
2. 按概率降序排列
3. 计算累积概率：$cumsum_i = \sum_{j=1}^{i} p_{\text{sorted}[j]}$
4. 保留满足 $cumsum_i \ge p$ 的最小的 $i$ 个 token
5. 其余 token 的 logits 设为 $-\infty$

与 Top-K 的关键区别：Top-K 是固定绝对数量的硬截断——无论分布是陡峭还是平坦，都只留 K 个。Top-P 是动态的自适应截断——如果头部 token 概率集中（如第一个 token 就占 0.85），可能只留 1-2 个；如果分布平坦，可能需要留几十个才能凑够 p=0.95。这使 Top-P 比 Top-K 更适应不同的采样场景。

Top-P 核采样概念图（来源：The Curious Case of Neural Text Degeneration，ICLR 2020）。不同上下文中的 softmax 概率分布差异巨大——左侧上下文中模型高度确定（前两个 token 就占绝大多数概率，Top-P 只需小核），右侧上下文则高度不确定（概率分布平坦，Top-P 需要大核才能覆盖相同累积概率）。这正是 Top-P 相比 Top-K 的优势：Top-P 的自适应核大小可以匹配上下文的不确定性程度，而不像 Top-K 一刀切地保留固定数量。

Tail Mass 示意图（来源：The Curious Case of Neural Text Degeneration）。绿色曲线标注 Top-K=50 时保留的概率质量分布，红色曲线标注被 Top-K 截掉的尾部分布（Tail Mass）。当模型的输出分布非常平坦时（红色面积很大），Top-K 的恒定截断会丢掉大量合理的候选 token——这些 token 本该有被采样到的机会。这就是 Top-P 用累积概率而非绝对排序数来决定截断边界的根本动机。

Top-P 在自然语言生成中的效果（来源：The Curious Case of Neural Text Degeneration）。当模型在若干候选 token 间存在不确定性时（如图中紫色节点所示），Top-P 保留足够多的分支让采样可以自然选择，而贪婪解码或低 Top-K 会过早确定路径，导致生成的文本陷入「重复-无意义」的模式。这张图直观解释了为什么 Top-P 核采样能产生比贪婪解码更自然、更多样的文本。

批量实现：GPU 上的高效实现需要 batch-level 排序——将所有 [num_tokens, vocab_size] 的 logits 降序排列，然后在最后一个维度上累积 softmax 概率，找每个位置的截止点。

2.4 Min-P#

作用：只保留概率不低于”最大概率的 p 倍”的 token。即，如果最高概率是 $p_{\max}$，则保留所有 $p_i \ge p_{\min} \cdot p_{\max}$ 的 token。

算法：

1. 对 logits 做 softmax 得到概率分布
2. 找到最大概率 $p_{\max} = \max_i p_i$
3. 构造 mask：$p_i < p_{\min} \cdot p_{\max}$
4. logits[mask] = -inf

为什么需要 Min-P：Top-K 和 Top-P 有各自的盲区。Top-K 在分布极陡时割掉可能有用的 token，在分布极平时保留太多噪声。Top-P 解决了”分布平坦”的问题，但有另一个问题：如果前 10 个 token 的概率加到一起都不够 p，Top-P 不得不保留大量极低概率的噪声 token。Min-P 的思路完全不同——不是从高往低累加，而是用最高概率做标杆，相对太低的全部砍掉。

当某请求的 min_p <= 0 时，该 token 位置跳过 Min-P。

2.5 惩罚项：频率惩罚、存在惩罚与重复惩罚#

这三项操作的目标都是抑制重复生成。模型在生成长文本时有一种倾向——在某些位置反复输出同样的词或短语。惩罚项通过修改 logits 来让已经出现过的 token 更难被再次选中。

频率惩罚（Frequency Penalty）：

其中 $\text{count}(i)$ 是 token $i$ 在已经生成的 token 中出现的次数，$\alpha_f$ 是频率惩罚系数。出现次数越多，惩罚越重。

存在惩罚（Presence Penalty）：

只要 token 至少出现过一次，就扣除固定惩罚值 $\alpha_p$。不同于频率惩罚的是，出现 100 次和出现 1 次的惩罚量相同。

重复惩罚（Repetition Penalty）：

与存在惩罚不同，重复惩罚是乘法而非减法——已出现 token 的 logits 被除以 $\alpha_r$（$\alpha_r > 1$ 时降低概率）。这导致惩罚量的语义是”概率缩小多少倍”而非”扣除多少分”，对大 logit token 和小 logit token 的惩罚更均匀。

批量实现中的 bookkeeping：计算每个 token 的出现次数需要在 GPU 上维护一个计数表。SamplingMetadata 中的 output_token_ids（每个请求已生成的 token 列表）是这项计算的数据源。vLLM 会将每个请求的已生成 token 做 bincount 得到出现次数，然后根据每个 token 位置上对应请求的计数表来修改对应行的 logits。

当 sampling_metadata.no_penalties 为 True 时，整个惩罚计算链被跳过。

处理器顺序与设计动机#

vLLM 的 LogitsProcessor 管线按以下顺序执行：

1. Penalties（频率/存在/重复惩罚）——在对数空间做加减/缩放，数值安全
2. Temperature——缩放整个 logits 向量的幅度
3. Min-P——需要先有 softmax 才知道 $p_{\max}$，但 Min-P 后仍可应用 Top-K/Top-P
4. Top-K——绝对值截断，不影响 softmax 的计算
5. Top-P——最后一步累积概率过滤

这个顺序的关键考量：惩罚项必须在 temperature 之前——如果先做 temperature 再做 penalty，$T \to 0$ 已经把 logits 推向 $+\infty$ 和 $-\infty$ 两极了，penalty 失去意义。Top-K 和 Top-P 放在最后是因为它们本质上是”softmax 后再过滤”——先让所有数值处理完成，再做最终的概率取舍。

3.1 Greedy（Temperature=0）路径#

当 sampling_metadata.all_greedy 为 True 时，所有请求都使用 temperature=0。此时不需要做 softmax——直接找 logits 中最大值的索引即可：

1
# 伪代码：Greedy 采样
2
if sampling_metadata.all_greedy:
3
    sampled = torch.argmax(logits, dim=-1)  # [num_tokens]

argmax 在 GPU 上是一个归约操作——每个 thread block 处理词表的一段，用共享内存做局部最大值比较，最后跨 block 汇总。

为什么 Greedy 不需要 softmax：softmax 是单调变换——不会改变最大值元素的位置。$\text{argmax}_i\ \ell_i = \text{argmax}_i\ \text{softmax}(\ell_i)$。跳过 softmax 省掉了 $V$ 次 $\exp$、$V$ 次加法、$V$ 次除法——对 128K 词表来说这是一笔可观的开销。

3.2 随机采样（Temperature > 0）路径#

当至少有一个请求的 temperature > 0 时，需要走随机采样路径。步骤：

1. Softmax：$\text{probs}_i = \exp(\ell_i) / \sum_j \exp(\ell_j)$
2. 从多项分布采样：PyTorch 的 multinomial 或自定义 CUDA Kernel

批量异构采样：经过 temperature processor 后，logits 已经被各请求对应的 temperature 缩放。同一个 batch 中 temperature=0.7 的请求和 temperature=1.0 的请求的 logits 已经各自由不同的 scale 处理过。所以采样阶段不需要再区分请求——直接对整个 [num_tokens, vocab_size] logits 做统一的 softmax + multinomial 即可。

vLLM 使用 torch.multinomial 作为主要的多项分布采样实现（在 PyTorch 中已高度优化），同时在部分专用后端（如 ROCm）中有自研 CUDA kernel 替代。

3.3 混合 batch 的处理#

当 all_greedy 为 False 但不是所有请求都需要随机采样时（部分 greedy + 部分随机），vLLM 的做法是将 batch 拆成两个子集分别处理：

1
# 伪代码：混合 batch 的采样
2
greedy_indices = sampling_metadata.greedy_indices  # temperature==0 的位置
3
random_indices = sampling_metadata.random_indices  # temperature>0 的位置
4

5
sampled = torch.empty(num_tokens, dtype=torch.int32, device=logits.device)
6
if len(greedy_indices) > 0:
7
    sampled[greedy_indices] = torch.argmax(logits[greedy_indices], dim=-1)
8
if len(random_indices) > 0:
9
    probs = torch.softmax(logits[random_indices], dim=-1)
10
    sampled[random_indices] = torch.multinomial(probs, num_samples=1).squeeze(-1)

请求级规格#

SamplingParams 中的 logprobs 字段有三种取值：

• None：不返回任何 logprobs（省计算）
• 0：只返回采样 token 的 log 概率
• k > 0：返回采样 token 的 log 概率 + 概率最高的 k 个备选 token 及各自的 log 概率和 token ID

计算方式#

采样完成后，logits 经过 softmax 再取 log：

1
# 伪代码：Logprobs 计算
2
logprobs = torch.log_softmax(logits, dim=-1)  # [num_tokens, vocab_size]

然后根据每个请求的需求提取不同级别的 logprobs：

1
# 每个 token 位置的采样 token 的 log 概率
2
sampled_logprobs = logprobs[range(num_tokens), sampled_token_ids]  # [num_tokens]
3

4
# 对需要 top-k logprobs 的 token 位置
5
topk_logprobs, topk_indices = torch.topk(logprobs, k=top_m, dim=-1)

工程考量：logprobs 的体积问题#

logprobs 的完整形状是 [num_tokens, vocab_size]，如果词表为 128K，FP32 下 256 个 token 就是 128 MB。大部分请求不要求返回全部 logprobs（甚至不要求返回任何 logprobs）。vLLM 的策略是分档处理：

• 如果没有任何请求需要 logprobs，完全跳过 softmax + log（greedy 路径）
• 如果有请求需要 logprobs，只计算需要的部分，并只提取需要的级别（采样 token 的 logprob 或 top-k logprobs）

异步 D2H 传输#

logprobs 的数据量不容忽视。对于 batch_size=64, top_m=20 的场景，每个 step 有 64×20×2（logprob+token_id）= 2560 个 32-bit 值 ≈ 10 KB——量不大，但如果每个 step 都同步 D2H 复制，会打乱 GPU 的流水线。

vLLM 使用异步 Stream 做 D2H 传输：

1
# 伪代码：异步 logprobs 回传
2
with torch.cuda.stream(self.sampler_output_copy_stream):
3
    sampled_token_ids_cpu = sampled_token_ids.to("cpu", non_blocking=True)
4
    if logprobs_tensors is not None:
5
        logprobs_cpu = logprobs_tensors.to("cpu", non_blocking=True)

non_blocking=True 保证 CPU 端在 GPU 传输完成前不等待，GPU 可以继续执行下一个 step。但这也意味着 EngineCore 在消费采样结果时需要确保传输已完成——通过 event 同步或 next step 的隐式 barrier 来保证。

在 EngineCore 中的消费#

在 EngineCore.step() 的主循环中，输出的消费发生在每步的开头：

1
# 伪代码：EngineCore.step() 的输出消费
2
def step(self):
3
    # 0. 处理上一步的输出（闭环关键步骤）
4
    self.scheduler.update_from_output(
5
        scheduler_output=self.last_scheduler_output,
6
        sampler_output=self.last_sampler_output,
7
    )
8

9
    # 1. 新一轮调度
10
    scheduler_output = self.scheduler.schedule()
11

12
    # 2. 执行模型
13
    model_output = self.model_executor.execute_model(
14
        scheduler_output,
15
    )
16
    self.last_sampler_output = model_output.sampler_output
17

18
    # ... 返回结果给前端

update_from_output() 做的事情：

1. 把采样的 token 写入请求的状态：request.output_tokens.append(sampled_token_id)
2. 检查 stop 条件：新 token 是否命中 stop_token_ids 或 stop_strs
3. 更新请求的完成状态：如果到达 max_tokens 或命中 stop，标记请求为已完成
4. 更新 KV Cache Manager 的请求状态：如果请求完成，释放其物理 KV blocks

这一步完成后，调度器就知道了：“上一步有哪些请求产生了哪些 token、哪些请求已经跑完了”。这直接影响下一轮调度时哪些请求继续被选择，哪些新请求可以被加入 batch。

请求完成时的收尾#

当一个请求生成了 max_tokens 个 token 或命中了 stop token/string，update_from_output() 将其标记为完成。完成的请求在下一轮 schedule() 中不会被分配新的 token 预算，其 KV blocks 被 free() 后重新加入空闲池，可供新请求分配。

这就是整个引擎控制环的闭合点——从采样结果的注入开始，到 KV 内存的回收结束，一整个 step 的状态迁移完成。