flashinfer.sampling.chain_speculative_sampling

flashinfer.sampling.chain_speculative_sampling(draft_probs, draft_token_ids, target_probs, maybe_output_accepted_token_num: Tensor | None = None, maybe_output_emitted_draft_token_num: Tensor | None = None, deterministic: bool = True, generator: Generator | None = None, seed: int | None = None, offset: int | None = None) → Tensor

用于序列生成的推测采样的融合 GPU 内核（论文《Accelerating Large Language Model Decoding with Speculative Sampling》中提出），其中草稿模型为每个请求生成一系列（链）的 token。

参数:

  • draft_probs (torch.Tensor) – 草稿模型生成的词汇表概率。形状：(batch_size, num_speculate_tokens, vocab_size)

  • draft_token_ids (torch.Tensor) – 草稿模型生成的 token 索引。形状：(batch_size, num_speculate_tokens)

  • target_probs (torch.Tensor) – 目标模型生成的词汇表概率。与输入 draft_probs 相比，目标模型的概率在末尾有一个额外的槽，因为目标模型将生成比草稿模型多一个 token。形状：(batch_size, num_speculate_tokens + 1, vocab_size)

  • maybe_output_accepted_token_num (Optional[torch.Tensor]) – 如果每个 token 独立考虑，则可以接受的 token 数量。此指标不考虑拒绝采样将在不满足概率要求 r < p/q 的第一个 token 时停止的事实。它仅评估草稿模型和目标模型的对齐情况。形状：(batch_size) 如果指定，则将接受的 token 数量将在此张量中原地添加。默认值为 None。

  • maybe_output_emitted_draft_token_num (Optional[torch.Tensor]) – 对于每个请求，最终发出的草稿 token 数量。不包括奖励 token。（因此，给定请求采样的 token 总数是 output_emitted_draft_token_num + 1）。形状：(batch_size) 如果指定，则将发出的 token 数量将在此张量中原地添加。默认值为 None。

  • deterministic (bool) – 是否使用确定性内核实现，默认值为 True。

  • generator (Optional[torch.Generator]) – 操作的随机数生成器。

  • seed (Optional[int]) – 采样操作期间用于 rng 的种子值。

  • offset (Optional[int]) – 采样操作期间用于 rng 的偏移值。

返回值:

  • output_token_ids (torch.Tensor) – 由目标模型验证的输出 token 索引，拒绝的样本用 -1 填充。与输入 draft_token_ids 相比，输出张量在末尾有一个额外的 token 索引用于最终 token，如果所有先前的 token 都被接受，则将从目标模型的概率中采样另一个“奖励”token。形状：(batch_size, num_speculate_tokens + 1)

  • output_accepted_token_num (torch.Tensor) – 如果每个 token 独立考虑，则可以接受的 token 数量。此指标不考虑拒绝采样将在不满足概率要求 r < p/q 的第一个 token 时停止的事实。它仅评估草稿模型和目标模型的对齐情况。形状：(batch_size)

  • output_emitted_draft_token_num (torch.Tensor) – 对于每个请求，最终发出的草稿 token 数量。不包括奖励 token。（因此，给定请求采样的 token 总数是 output_emitted_draft_token_num + 1）。形状：(batch_size)

示例

>>> import torch
>>> import flashinfer
>>> torch.manual_seed(42)
>>> batch_size = 1
>>> num_speculate_tokens = 2
>>> vocab_size = 4
>>> draft_probs = torch.tensor([[[0.1, 0.2, 0.3, 0.4], [0.2, 0.3, 0.4, 0.1]]]).to(0)
>>> # token 2 was sampled from draft model for the first token, and
>>> # token 1 was sampled from draft model for the second token
>>> draft_token_ids = torch.tensor([[2, 1]], dtype=torch.int32).to(0)
>>> target_probs = torch.tensor([[[0.0, 0.1, 0.6, 0.3], [1.0, 0.0, 0.0, 0.0], [0.7, 0.1, 0.1, 0.1]]]).to(0)
>>> output_token_ids, output_accepted_token_num, output_emitted_draft_token_num =\
...     flashinfer.sampling.chain_speculative_sampling(
...         draft_probs, draft_token_ids, target_probs)
>>> # the first token is accepted, the second token is rejected and sampled from the difference
>>> # between the target model and the draft model, the third token is padded with -1
>>> output_token_ids
tensor([[ 2,  0, -1]], device='cuda:0', dtype=torch.int32)
>>> output_accepted_token_num
tensor([1], device='cuda:0')
>>> output_emitted_draft_token_num
tensor([1], device='cuda:0')