本文介绍一种基于向量化操作的高效方法，替代原始 python 循环，实现对变长 token 序列的 logits 张量（如语言模型输出）进行分段索引与乘积-求和混合概率计算，显著提升推理效率。
本文介绍一种基于向量化操作的高效方法，替代原始 python 循环，实现对变长 token 序列的 logits 张量（如语言模型输出）进行分段索引与乘积-求和混合概率计算，显著提升推理效率。
在使用自回归语言模型（如 LLaMA、GPT 系列）进行 token-level 概率建模时，常需对一批变长 tokenized 单词（例如 subword 分词结果）计算其联合概率，并额外引入一组「新词起始 token ID」（如空格前缀 token）作为终止/切换条件。原始实现依赖逐样本 for 循环 + 条件判断，时间复杂度高、无法利用 GPU 并行性。以下提供一套完全向量化、数值稳定、且易于复用的 PyTorch 实现方案。
核心思路：分步向量化 + 掩码驱动索引
我们将目标概率拆解为两部分：
- 主序列概率：对每个单词的非填充 token，沿时间步索引 pxhs[i, t, token_id] 并连乘；
-
新词候选概率：在每个单词末尾对应的时间步（即首个
位置），对预定义的 new_word_token_ids 集合求 logits 的 softmax 概率和（或直接对 logit 求和后取 log，见下文说明）。
关键前提是：我们能提前获取每个单词的实际长度（即首个 pad token 的索引），记为 seq_lens: [n_words]。可通过 torch.argmax((token_list == pad_token_id).long(), dim=1) 安全获得（需处理全非 pad 场景，推荐用 torch.where + torch.min 更鲁棒）。
✅ 向量化实现（含完整可运行代码）
import torch
# 假设输入已准备就绪（实际中从 tokenizer 和 model.forward 获取）
pxhs = torch.rand((3, 4, 1000), requires_grad=False) # [n, seq_len+1, vocab]
word_token_list = torch.tensor([
[120, 1, 1], # pad_token_id = 1
[131, 132, 1],
[140, 141, 142],
])
pad_token_id = 1
new_word_token_ids = torch.tensor([0, 1, 2, 3, 5])
# Step 0: 构建 seq_lens —— 每个单词真实 token 数（即首个 pad 的索引）
mask_pad = (word_token_list == pad_token_id)
seq_lens = torch.where(
mask_pad.any(dim=1),
mask_pad.nonzero()[:, 1].reshape(-1, 1).max(dim=1)[0],
torch.full((word_token_list.size(0),), word_token_list.size(1))
)
n, max_seq_len_plus1, vocab_size = pxhs.shape
seq_idxs = torch.arange(n, device=pxhs.device)
# Step 1: 计算主序列 log-probabilities（向量化连乘 → 连加 log）
log_pxhs = pxhs.log() # 数值更稳定；若 pxhs 是 logits，应先 softmax 再 log，或用 log_softmax
log_p_x = torch.zeros(n, device=pxhs.device)
# 对每个时间步 t ∈ [0, max_len)，仅对未截断的样本累加
for t in range(max_seq_len_plus1 - 1): # 注意：主 token 最多占 seq_len 位，最后一位留给 new_word
valid_mask = (t <h3>⚠️ 注意事项与最佳实践</h3>
- 数值稳定性优先：始终在 log-space 中运算。避免 pxhs.exp().sum()，改用 torch.logsumexp(...)。
- seq_lens 必须准确：它决定了主 token 的边界和 new_word 的查询位置。若 word_token_list 无 pad，seq_lens 应设为 max_tokenization_length，而非越界。
- 维度对齐检查：确保 word_token_list.shape[0] == pxhs.shape[0]，且 seq_lens 长度一致；建议在函数开头加入 assert。
- 内存权衡：上述方案时间复杂度为 O(n × max_len)，空间为 O(n × vocab)。若 vocab_size 极大（如 128K），可改用稀疏索引：pxhs[seq_idxs, seq_lens][:, new_word_token_ids] 直接切片，再 sum(dim=1)。
- JIT 友好性：循环 for t in range(...) 在 TorchScript 中受限，生产环境建议用 torch.arange + torch.gather 重写为纯张量操作（需稍作 reshape）。
✅ 总结
该方案通过将“条件循环”转化为“掩码驱动的广播索引 + 分步聚合”，成功将原始 O(n·L) 的 CPU 循环提升至接近 GPU 原生速度，同时保持代码清晰与数值鲁棒性。适用于批量词级置信度评估、可控文本生成中的 token 筛选、以及子词合并概率建模等任务。只需提供 token_list、pxhs、pad_token_id 和 new_word_token_ids 四个输入，即可端到端向量化执行。
