DSpark

@torch.inference_mode() #PyTorch 的推理模式
def generate_decoding_sample(
    *,
    target_model,
    input_ids: torch.Tensor,
    max_new_tokens: int,
    max_proposal_tokens: int,
    temperature: float,
    stop_token_ids: list[int] | None,
    init_context: Callable[..., Any],
    propose: Callable[..., DraftProposal],
    update: Callable[[Any, VerificationResult], None],
    post_verify: Callable[[DraftProposal, VerificationResult], None] | None = None,
) -> SimpleNamespace:
    # 通用的Speculative-decoding loop，任何Speculative-decoding算法都可以接入
    """Speculative-decoding loop.
    `init_context(initial_output, output_ids, position_ids, num_input_tokens)` builds the algorithm-specific state once after prefill.  //不同 draft 算法需要保存的东西不一样，这里就初始化对应的状态
    `propose(context, output_ids, position_ids, start, stop_token_ids)` returns the next DraftProposal.
    `update(context, verification)` advances the state when the loop continues.
    `post_verify` is an optional diagnostic hook called after every verification (used for confidence calibration).
    """
    """
    实现了统一流程：
    初始化
    提 proposal
    verify
    更新
    循环
    """
    assert max_proposal_tokens >= 1
    assert input_ids.size(0) == 1, "only bsz=1 is supported"

    device = input_ids.device
    num_input_tokens = input_ids.shape[1]
    max_length = num_input_tokens + int(max_new_tokens)

    output_ids = torch.empty(
        (1, max_length + max_proposal_tokens + 1),
        dtype=torch.long,
        device=device,
    ) # 显存优化操作：预先分配一块足够大的空白 Tensor 用来存放整个推理过程生成的所有 Token（包括最后一步投机可能越界的多余 Token），避免每次生成新字符时都去动态 torch.cat 导致显存碎片和性能开销。
    position_ids = torch.arange(output_ids.shape[1], device=device).unsqueeze(0)
    past_key_values_target = DynamicCache()
    # 做 prefill，先让 target 正常吐出一个 token
    output = target_model(
        input_ids=input_ids,
        position_ids=position_ids[:, :num_input_tokens],
        past_key_values=past_key_values_target,
        use_cache=True,
        output_hidden_states=True,
        logits_to_keep=1, #指定只保留最后几个位置的 Logits，这里是1
    )
    # 把prompt填入output_ids，然后从 target 的输出分布采样出第一个新 token
    output_ids[:, :num_input_tokens] = input_ids
    output_ids[:, num_input_tokens : num_input_tokens + 1] = sample_from_probs(
        logits_to_probs(output.logits, float(temperature))
    )
    # start 标记当前最新已被验证并采纳的 Token 的索引位置
    # （此时正好在 Prompt 结束、新生成 Token 的地方）。
    # 后面三个列表分别用于记录每一轮最终大模型的接纳长度、
    # 小模型的提议长度、小模型被接纳的纯草稿长度（用于后期统计效率指标）。
    # 其中 acceptance_lengths 和 accepted_draft_lengths 的区别是：
    # 无论小模型提议的草稿是对是错，大模型在验证时都会在小模型被接纳的草稿向后多预测一个 Token
    # （跑了一次的结果得用上啊），除非到结束符了。
    start = input_ids.shape[1]
    acceptance_lengths: list[int] = []
    proposal_lengths: list[int] = []
    accepted_draft_lengths: list[int] = []

    initial_token = output_ids[:, num_input_tokens : num_input_tokens + 1]
    if has_stop_token(initial_token, stop_token_ids):
        output_ids = output_ids[:, : num_input_tokens + 1]
        output_ids = trim_output_ids(output_ids, num_input_tokens, stop_token_ids)
        return SimpleNamespace(
            output_ids=output_ids,
            num_input_tokens=num_input_tokens,
            num_output_tokens=output_ids.shape[1] - num_input_tokens,
            acceptance_lengths=acceptance_lengths,
            proposal_lengths=proposal_lengths,
            accepted_draft_lengths=accepted_draft_lengths,
            verify_count=0,
        )
    # 初始化循环状态
    context = init_context(
        initial_output=output,
        output_ids=output_ids,
        position_ids=position_ids,
        num_input_tokens=num_input_tokens,
    )
    # Speculative Decoding 核心循环
    while start < max_length:
        # draft 进行提议
        proposal = propose(
            context=context,
            output_ids=output_ids,
            position_ids=position_ids,
            start=start,
            stop_token_ids=stop_token_ids,
        )
        # target 进行验证
        """
        这里的verify_draft_tokens的返回值为
        return VerificationResult(
        target_output=target_output,
        target_probs=target_probs,
        accept_prefix_mask=accept_prefix_mask,
        accepted_draft_tokens=accepted_draft_tokens,
        next_token=next_token,
        effective_proposal_length=effective_proposal_length,
        terminated_by_stop_token=terminated_by_stop_token,
        committed_tokens=committed_tokens,
    )
    """
        verification = verify_draft_tokens(
            target_model=target_model,
            proposal=proposal,
            position_ids=position_ids,
            start=start,
            past_key_values_target=past_key_values_target,
            temperature=temperature,
            max_proposal_tokens=max_proposal_tokens,
            current_token_ids=output_ids[:, start : start + 1],
            stop_token_ids=stop_token_ids,
        )
        # 如果有传入辅助性的诊断函数，在这步触发。
        if post_verify is not None:
            post_verify(proposal, verification)

        # 更新内存：把已经被大模型认可通过的小模型草稿片段，正式写回到 output_ids 序列中。
        proposal_lengths.append(int(verification.effective_proposal_length))
        accepted_draft_tokens = int(verification.accepted_draft_tokens)
        accepted_draft_lengths.append(accepted_draft_tokens)

        output_ids[:, start : start + accepted_draft_tokens + 1] = (
            proposal.verify_input_ids[:, : accepted_draft_tokens + 1]
        )
        # 判断是否碰到结束符，碰到就终止循环
        if verification.terminated_by_stop_token:
            acceptance_lengths.append(accepted_draft_tokens)
            start += accepted_draft_tokens
            past_key_values_target.crop(start)
            break
        # 没碰到就在被接收的draft后面，再补上target修正/产出的那个 next_token（这也是为什么每轮哪怕speculation草稿全错，大模型也必定能保底前进一步）
        output_ids[:, start + accepted_draft_tokens + 1] = verification.next_token
        new_token_ids = output_ids[:, start + 1 : start + accepted_draft_tokens + 2]
        acceptance_lengths.append(accepted_draft_tokens + 1)
        start += accepted_draft_tokens + 1
        past_key_values_target.crop(start) # KV 缓存裁剪：由于大模型在验证时连同小模型那些错误的草稿也一起计算了 KV 并塞进了缓存，此时必须根据最终实际接纳的正确长度start把后面错误的缓存裁掉
        update(context, verification)

        if has_stop_token(new_token_ids, stop_token_ids):
            break
    # 利用start将真正有效的Token序列截取出来
    output_ids = output_ids[:, : min(start + 1, max_length)]
    output_ids = trim_output_ids(output_ids, num_input_tokens, stop_token_ids)
    # 返回并封装
    return SimpleNamespace(
        output_ids=output_ids,
        num_input_tokens=num_input_tokens,
        num_output_tokens=output_ids.shape[1] - num_input_tokens,
        acceptance_lengths=acceptance_lengths,
        proposal_lengths=proposal_lengths,
        accepted_draft_lengths=accepted_draft_lengths,
        verify_count=len(proposal_lengths),
    )

# 结合小模型在每个格子的隐状态，以及小模型刚刚自己采样出来的草稿 Token，去算出一个“自信心分数矩阵”。用这个分数来评估小模型对自己刚写出来的草稿到底有多大的把握。
def _predict_confidence_logits(
    model: DSparkModel,
    *,
    proposal_hidden_states: torch.Tensor,
    draft_input_ids: torch.Tensor,
    sampled_tokens: torch.Tensor,
    block_size: int,
) -> torch.Tensor | None:
    # 先前一个token拼接上抛弃draft最后一个token的序列
    prev_token_ids = torch.cat(
        [draft_input_ids[:, :1], sampled_tokens[:, :-1]],
        dim=1,
    )
    # 上面的prev_token_ids 和 proposal_hidden_states一一对应的时候是错开一位的，比如proposal_hidden_states的第二个对应的是预测的第一个的token，这样就实现了前后关系对应，具体上面有例子
    confidence_pred = model.predict_confidence_step(
        proposal_hidden_states,
        prev_token_ids=prev_token_ids,
    )
    if confidence_pred is None:
        return None
    return confidence_pred.float().reshape(
        confidence_pred.shape[0],
        block_size,
        -1,
    )[:, :, 0]
def _confident_prefix_length(
    confidence_logits: torch.Tensor,
    *,
    block_size: int,
    threshold: float,
) -> int:
    if threshold <= 0.0:
        return int(block_size)
    below_threshold = confidence_logits.sigmoid() < threshold
    if not bool(below_threshold[0].any().item()):
        return int(block_size)
    return int(torch.nonzero(below_threshold[0], as_tuple=False)[0].item())

def build_dspark_proposal(
    model: DSparkModel,
    *,
    draft_input_ids: torch.Tensor,
    block_hidden: torch.Tensor,
    block_size: int,
    temperature: float,
    confidence_threshold: float,
) -> DSparkDraftProposal:
    assert draft_input_ids.size(0) == 1, "build_dspark_proposal requires batch_size=1"
    # draft 提议的每个位置的 hidden_states
    proposal_hidden_states = block_hidden[:, :block_size, :]
    base_draft_logits = model.compute_logits(proposal_hidden_states)
    sampled_tokens, draft_logits = model.sample_draft_tokens(
        base_draft_logits,
        first_prev_token_ids=draft_input_ids[:, 0],
        temperature=temperature,
        hidden_states=proposal_hidden_states,
    )
    # DSpark 的关键技术：动态置信度动态裁剪
    proposal_draft_tokens = int(block_size)
    confidence_logits = None
    if model.confidence_head is not None:
        # 如果有confidence_head，则调用_predict_confidence_logits
        confidence_logits = _predict_confidence_logits(
            model,
            proposal_hidden_states=proposal_hidden_states,
            draft_input_ids=draft_input_ids,
            sampled_tokens=sampled_tokens,
            block_size=block_size,
        )
        if confidence_logits is None:
            return _empty_dspark_proposal(draft_input_ids)
        # 将confidence_logits经过一个sigmoid，然后超过阈值的才被验证
        proposal_draft_tokens = _confident_prefix_length(
            confidence_logits,
            block_size=block_size,
            threshold=float(confidence_threshold),
        )

    if proposal_draft_tokens == 0:
        return _empty_dspark_proposal(draft_input_ids)

    verify_input_ids = torch.cat(
        [draft_input_ids[:, :1], sampled_tokens[:, :proposal_draft_tokens]],
        dim=1,
    )
    draft_probs = logits_to_probs(
        draft_logits[:, :proposal_draft_tokens, :],
        temperature,
    )
    return DSparkDraftProposal(
        draft_token_count=proposal_draft_tokens,
        verify_input_ids=verify_input_ids,
        draft_probs=draft_probs,
        confidence_logits=(
            confidence_logits[:, :proposal_draft_tokens]
            if confidence_logits is not None
            else None
        ),
    )

# 用一个很轻的线性 AcceptRatePredictor，根据每个 draft 位置的 hidden state，以及可选的前一个 token embedding，输出该位置的 confidence logit；推理时再经过 sigmoid，用来判断这一位是否值得继续送进 verify。
def predict_confidence_step(
    self,
    hidden_states,
    prev_token_ids=None,
):
    if self.confidence_head is None:
        return None
    if self.confidence_head_with_markov:
        prev_embeddings = self.markov_head.get_prev_embeddings(prev_token_ids)
        features = torch.cat([hidden_states, prev_embeddings], dim=-1)
        return self.confidence_head(features).float()
    return self.confidence_head(hidden_states).float()

def sample_draft_tokens(...):
    batch_size, proposal_len = base_logits.shape[:2]
    if proposal_len == 0:
        ...
    if self.markov_head is None:
        return sample_tokens(base_logits, temperature), base_logits
    return self.markov_head.sample_block_tokens(
        base_logits,
        first_prev_token_ids=first_prev_token_ids,
        hidden_states=hidden_states,
        temperature=temperature,
    )

sampled_tokens = []
corrected_logits = []
prev_token_ids = first_prev_token_ids.long()
for step_idx in range(proposal_len):
    step_hidden = hidden_states[:, step_idx, ...]
    # 用前一个位置对这一位做修正 最终 logits = 初稿 logits + 顺序修正 bias
    step_logits = self.apply_step_logits(
        base_logits[:, step_idx, :],
        token_ids=prev_token_ids,
        hidden_states=step_hidden,
    )
    corrected_logits.append(step_logits.unsqueeze(1))
    next_token_ids = sample_tokens(step_logits.unsqueeze(1), temperature=temperature).squeeze(1)
    sampled_tokens.append(next_token_ids)
    prev_token_ids = next_token_ids
return torch.stack(sampled_tokens, dim=1), torch.cat(corrected_logits, dim=1)

def verify_draft_tokens(
    *,
    target_model,
    proposal: DraftProposal,
    position_ids: torch.Tensor,
    start: int,
    past_key_values_target: DynamicCache,
    temperature: float,
    max_proposal_tokens: int,
    current_token_ids: torch.Tensor | None = None,
    stop_token_ids: list[int] | None = None,
) -> VerificationResult:
    """Verify draft tokens with the target model and rejection sampling."""
    if proposal.draft_token_count > max_proposal_tokens:
        raise ValueError(
            "DraftProposal.draft_token_count must not exceed "
            f"max_proposal_tokens={max_proposal_tokens}, "
            f"got {proposal.draft_token_count}."
        )
    if current_token_ids is not None and not torch.equal(
        proposal.verify_input_ids[:, :1],
        current_token_ids,
    ):
        raise ValueError(
            "DraftProposal.verify_input_ids must start with the current "
            "accepted token."
        )
    # 计算本次需要验证的总长度 verify_length（等于草稿长度 + 1 个基准 Token）。并从全局位置编码中切出对应这段区间的位置 ID verify_position_ids
    draft_token_count = int(proposal.draft_token_count)
    verify_length = draft_token_count + 1
    verify_position_ids = position_ids[:, start : start + verify_length]
    # 核心并行计算：把大模型上一轮最后一个正确 Token 和小模型新提议的K个草稿拼接成的整个序列，一次性全部喂给大模型。大模型在内部会并行计算并追加这一整段的KV缓存。
    target_output = target_model(
        input_ids=proposal.verify_input_ids,
        position_ids=verify_position_ids,
        past_key_values=past_key_values_target,
        use_cache=True,
        output_hidden_states=True,
    )
    if target_output.logits.ndim != 3:
        raise ValueError(
            "target model must return rank-3 logits [B, S, V], "
            f"got ndim={target_output.logits.ndim}."
        )
    target_probs = logits_to_probs(target_output.logits, float(temperature))
    if (
        draft_token_count > 0
        and proposal.draft_probs is not None
        and proposal.draft_probs.size(-1) != target_probs.size(-1)
    ):
        raise ValueError(
            "DraftProposal.draft_probs vocab size must match target logits, "
            f"got {proposal.draft_probs.size(-1)} and {target_probs.size(-1)}."
        )

    accept_prefix_mask = None
    if draft_token_count > 0:
        assert proposal.draft_probs is not None
        proposed_tokens = proposal.verify_input_ids[:, 1:]
        # 获取大模型对草稿选出的token id的预测概率
        selected_target_probs = gather_token_probs(
            target_probs[:, :-1, :],# target会在draft后多预测一个，这里要去掉
            proposed_tokens,
        )
        # 获取小模型对自身草稿选出的token id的预测概率
        selected_draft_probs = gather_token_probs(
            proposal.draft_probs,
            proposed_tokens,
        ).clamp_min(1e-8)
        # 计算接受概率（也就是上面提到的公式）
        accept_prob = torch.clamp(
            selected_target_probs / selected_draft_probs,
            max=1.0,
        )
        accept_mask = (torch.rand_like(accept_prob) < accept_prob).to(torch.int64) # 生成accept_prob大小的0~1随机浮点数，并和接受概率accept_prob比较
        accept_prefix_mask = accept_mask.cumprod(dim=1) # cumprod(dim=1) 是累乘（Cumulative Product），会沿着序列方向，让前面的结果和后面一直相乘，也就是只要有一个0后面就全是0。
        accepted_draft_tokens = int(accept_prefix_mask.sum(dim=1)[0].item())
    else:
        accepted_draft_tokens = 0
    # 检查大模型接受的草稿Token内部，是否包含了结束符
    effective_proposal_length = draft_token_count
    terminated_by_stop_token = False
    if stop_token_ids and accepted_draft_tokens > 0:
        accepted_slice = proposal.verify_input_ids[0, 1 : accepted_draft_tokens + 1]
        stop_tensor = torch.tensor(
            stop_token_ids,
            device=accepted_slice.device,
            dtype=accepted_slice.dtype,
        )
        eos_hits = torch.isin(accepted_slice, stop_tensor).nonzero(as_tuple=True)[0]
        if eos_hits.numel() > 0:
            eos_pos = int(eos_hits[0].item())
            accepted_draft_tokens = eos_pos + 1
            effective_proposal_length = eos_pos + 1
            terminated_by_stop_token = True

    if 0 < draft_token_count and accepted_draft_tokens < draft_token_count:
        assert proposal.draft_probs is not None
        # 如果草稿有错被拒绝了，利用Residual Sampling采样出next_token，Residual Sampling保证了分布一致
        next_token = sample_residual(
            target_probs[:, accepted_draft_tokens, :],
            proposal.draft_probs[:, accepted_draft_tokens, :],
        )
    else:
        # 如果都接受了，则直接采样下一个
        next_token = sample_from_probs(target_probs[:, -1:, :]).squeeze(1)

    committed_tokens = torch.cat(
        [
            proposal.verify_input_ids[:, 1 : accepted_draft_tokens + 1],
            next_token.unsqueeze(1),
        ],
        dim=1,
    )

    return VerificationResult(
        target_output=target_output,
        target_probs=target_probs,
        accept_prefix_mask=accept_prefix_mask,
        accepted_draft_tokens=accepted_draft_tokens,
        next_token=next_token,
        effective_proposal_length=effective_proposal_length,
        terminated_by_stop_token=terminated_by_stop_token,
        committed_tokens=committed_tokens,
    )