采样参数对解码开销的影响：temperature、top_p 与采样核路径

2024-06-13

在一组 8 卡 **** ****（单卡 64GB HBM，卡间互联 392 GB/s）上，用 **** 的基础软件栈配合 vLLM 0.5.0 的 v0 引擎起 OpenAI 兼容服务，入口是 python -m vllm.entrypoints.openai.api_server --model <model-path>。模型词表较大，压测并发 64。同样一条请求，把 top_p 从 1.0 调到 0.9、top_k 从 -1 调到 50 后，TPOT 出现可见上升；只改 temperature 则变化很小。原因在这一版采样是前向之后的一段独立的设备侧计算，作用在形状为 [num_seqs, vocab_size] 的 logits 上，开销随 num_seqs × vocab_size 增长、与模型参数量无关。top_p、top_k、min_p、各种惩罚项各自打开一趟针对整个词表维度的计算，而这一版的采样核没有定制算子，全部走 PyTorch。

解码一步里采样做什么

一步 decode 的顺序是：模型前向算出 hidden states，LM head 用 torch.matmul(hidden_states, embedding.t()) 得到词表维度的 logits，交给 Sampler.forward。Sampler.forward 这一步按固定次序做几件事：

_apply_min_tokens_penalty，处理 min_tokens；
_apply_penalties，处理 presence / frequency / repetition 三个惩罚，由 do_penalties 开关控制；
logits.div_(temperatures)，按序列各自的 temperature 缩放，元素级运算；
_apply_top_k_top_p，处理 top_k 与 top_p，由 do_top_p_top_k 开关控制；
_apply_min_p，处理 min_p，由 do_min_p 开关控制；
torch.softmax 与 torch.log_softmax，都指定 dtype=torch.float，各自产出一份 [num_seqs, vocab_size] 的 fp32 张量；
_sample，按采样类型选路产出下一个 token。

第 6 步是固定成本：无论参数怎么设，每步都要在 fp32 下算出概率与对数概率两份全词表张量。设本步参与采样的 token 数为 N、词表为 V，则这一步的材料开销是

probs    : N × V × 4 字节        # fp32
logprobs : N × V × 4 字节        # fp32

两份合计 8 × N × V 字节。词表 12.8 万、N = 64 时约 65 MB，每步都要读写一遍。它与模型大小无关，只跟 batch 与词表走，因此在词表大、并发高、模型偏小时占比更醒目。

哪些参数会额外打开一趟全词表计算

第 2、4、5 步是否执行，由三个布尔开关决定，开关在构造采样输入时对整批统计一次：

# vllm/model_executor/sampling_metadata.py（原文，0.5.0）
do_penalties = False
do_top_p_top_k = False
do_min_p = False
...
if not do_top_p_top_k and (top_p < 1.0 - _SAMPLING_EPS
                           or top_k != vocab_size):
    do_top_p_top_k = True
if not do_min_p and min_p > _SAMPLING_EPS:
    do_min_p = True
if not do_penalties and (abs(p) >= _SAMPLING_EPS
                         or abs(f) >= _SAMPLING_EPS
                         or abs(r - 1.0) >= _SAMPLING_EPS):
    do_penalties = True

三处判据各对应一类参数，默认值下都不触发：top_p = 1.0 且 top_k = -1（内部折算为 vocab_size）时不开 top_k/top_p；min_p = 0 时不开 min_p；三个惩罚项取 0、0、1.0 时不开 penalties。

参数	默认值	打开的开关	打开后新增的计算
`temperature`	1.0	无开关	`logits.div_` 一趟元素级缩放；`< 1e-5` 时采样类型转为 `GREEDY`
`top_p`	1.0	`do_top_p_top_k`	`_apply_top_k_top_p`，对词表维排序与掩码
`top_k`	-1	`do_top_p_top_k`	同上
`min_p`	0.0	`do_min_p`	`_apply_min_p`，一次 softmax 加阈值掩码
`presence_penalty`	0.0	`do_penalties`	`_apply_penalties`，逐 token 计数加全词表元素级运算
`frequency_penalty`	0.0	`do_penalties`	同上
`repetition_penalty`	1.0	`do_penalties`	同上

两个要点。其一，开关是整批统计的：只要批里有一条请求设了 top_p < 1.0，本步 do_top_p_top_k 即为真，排序在整个 batch 的 logits 上跑一遍，其余默认参数的序列也一起承担。混部负载里一条带参数的请求，会把这一趟计算摊到同批所有序列上。其二，top_k / top_p 路径要对词表维排序，min_p 路径是一次 softmax 加阈值，两者都开时排序仍然要跑。

采样类型与选路

SamplingType 是一个四值枚举，由参数推导：

# vllm/sampling_params.py（原文）
class SamplingType(IntEnum):
    GREEDY = 0
    RANDOM = 1
    RANDOM_SEED = 2
    BEAM = 3

# 推导（cached_property）
if self.use_beam_search:
    return SamplingType.BEAM
if self.temperature < _SAMPLING_EPS:      # _SAMPLING_EPS = 1e-5
    return SamplingType.GREEDY
if self.seed is not None:
    return SamplingType.RANDOM_SEED
return SamplingType.RANDOM

temperature < 1e-5 判为 GREEDY，选路走 _greedy_sample，用 argmax 取 token；seed 非空判为 RANDOM_SEED，走带固定种子的随机路径；use_beam_search 判为 BEAM，走 beam search 的候选扩展。三条路都建立在上文第 6 步那两份 fp32 全词表张量之上，GREEDY 省不掉 softmax，只改变取 token 的那一步算子。这一版的 _sample 固定调用 _sample_with_torch，Triton 采样核已实现但被注释掉（源码注明待其更快后再启用），因此没有一条能绕开 PyTorch 实现的后备路径。

best_of 与 n 的放大

best_of > 1 时，prompt 阶段一次产出 best_of 个候选，随后在输出处理阶段 fork 成多条序列：Sequence.fork 复制父序列，Scheduler.fork_seq 在块管理层同步分叉。n > 1 同理。代价有两层：采样侧每步要为更多序列算 token，解码侧多出来的序列各占一份 KV 并各走一次前向。best_of 还会改变调度行为：该请求的序列组被判为多序列组，显存不足时走换出路径而不是重算。

触发条件

词表大（十万量级）且并发高：第 6 步的两份 fp32 张量规模随之放大。
请求带非默认采样参数：top_p < 1、top_k != -1、min_p > 0 或任一惩罚项非默认。
混部负载里只有部分请求带参数：整批的开关被拉高，默认参数的请求也被裹进额外计算。
best_of > 1 或 n > 1：候选与 fork 放大解码与采样两侧的开销。

现象

把 top_p 调低或 top_k 收紧后 TPOT 上升，且并发越高上升越明显（开销随 N × V 走）。
只改 temperature（仍大于 1e-5）几乎看不到变化，因为它只多一趟元素级缩放。
只改 temperature = 0（转为 GREEDY）时 TPOT 略降，降幅来自取 token 的算子，不是来自 softmax。
混部负载里，把所有请求统一到同一组采样参数后 TPOT 回落，说明此前有请求拉高了整批开关。

定位步骤

记录请求侧的采样参数分布：分别统计 top_p、top_k、min_p、三个惩罚项非默认的请求占比。任一项大于零即会打开对应开关。
分段计时：固定请求集与并发，先跑全部默认参数（temperature = 1.0、top_p = 1.0、top_k = -1、min_p = 0、惩罚项取默认），再逐项打开，量每一步的 TPOT 变化。不要一次全开，否则分不清是排序、softmax 还是惩罚项贡献的。
算采样张量规模：用当前并发与词表算 8 × N × V 字节，与单步耗时对照，判断采样段是否已是可见量级。
查是否混部：若同一时刻既有带参数请求又有默认请求，确认整批开关被拉高；把两类请求分开压测可验证。
核对 best_of / n：从调用侧确认取值，> 1 时把 fork 出的序列数计入并发。

# 字段名示意：与采样开销有关的量
num_seqs_in_batch          # 本步参与采样的 token 数 N
vocab_size                 # 词表 V
do_top_p_top_k             # 批内是否有请求设置 top_p / top_k
do_min_p                   # 批内是否有请求设置 min_p
do_penalties               # 批内是否有请求设置惩罚项

处置

动作	作用	代价 / 前提
统一批内采样参数	避免个别请求拉高整批的 `do_top_p_top_k` / `do_penalties`	分散参数需求时须按参数分池或分离部署
用 `min_p` 替代 `top_p` / `top_k`	阈值路径是一次 softmax 加掩码，比词表维排序轻	截断语义不同，需重新评估输出质量
收紧 `top_k`	排序仍在整个词表上跑，收益有限	不改变排序规模，只改截断结果
需要多样性时降低 `best_of` / `n`	减少候选与 fork，同时降低采样与解码两侧开销	候选数下降，影响用途
温度取 0 走 `GREEDY`	选路换 argmax，省掉随机采样那一步	输出确定化，失去采样多样性
`--gpu-memory-utilization` 或并发调低	减小 `N`，直接压低采样段的绝对开销	峰值吞吐或并发下降

处置顺序：先确认批内参数是否混杂，混杂就先把参数统一（或分池），这一步常能直接消掉排序与惩罚项的整批开销；再看是否真需要固定种子或 top_p 这类需要整词表计算的参数；best_of / n 只在确实需要多候选时才留。

记录字段

# 一条采样开销记录（字段名示意）
sampling_params:
  temperature:                # 默认 1.0；< 1e-5 判为 GREEDY
  top_p:                      # 默认 1.0；< 1.0 打开 do_top_p_top_k
  top_k:                      # 默认 -1（折算为 vocab_size）；!= -1 打开 do_top_p_top_k
  min_p:                      # 默认 0.0；> 1e-5 打开 do_min_p
  presence_penalty:           # 默认 0.0
  frequency_penalty:          # 默认 0.0
  repetition_penalty:         # 默认 1.0
  seed:
  best_of:                    # 默认等于 n；> 1 触发 fork
  n:
  use_beam_search:            # 真值走 BEAM
batch:
  num_seqs:                   # N
  vocab_size:                 # V
  do_top_p_top_k:             # 批级开关
  do_min_p:
  do_penalties:
  param_mixed:                # 批内是否混杂非默认参数
workload:
  model_name:
  concurrency:
perf:
  tpot_default_params:        # 全默认参数下的 TPOT
  tpot_top_p_top_k_on:
  tpot_min_p_on:
  tpot_penalties_on:
  sampling_tensor_bytes:      # 估算 8 × N × V
status:                       # 采样段可见 / 采样段可忽略 / 受 fork 放大

小结

这一版采样是前向之后作用在全词表 logits 上的一段设备侧计算，固定成本是每步两份 [num_seqs, vocab_size] 的 fp32 概率与对数概率张量，规模 8 × N × V 字节，与模型参数量无关。top_p、top_k、min_p 与三个惩罚项各自由一个批级布尔开关控制是否额外打开一趟词表维计算：top_p < 1 或 top_k != -1 打开排序路径，min_p > 0 打开阈值路径，惩罚项非默认打开计数与掩码路径。开关对整批生效，混部负载里一条带参数的请求会把计算摊到同批所有序列上。temperature 只多做一趟元素级缩放，< 1e-5 时转为 GREEDY 换用 argmax 取 token。best_of > 1、n > 1 在采样与解码两侧同时放大，并让该请求进入多序列组的调度路径。