deepinfer-confidence

deepinfer/deepinfer/model_executor/models/mtp/dflash.py · a679bb8 · L46–79 · 固定 commit 上游

这是固定版本源码的只读摘录；可在线阅读，也可通过上游链接查看原文件。

46class ConfidenceHead(nn.Module):
47    """Predict the acceptance probability of each draft token."""
48
49    def __init__(
50        self,
51        n_embed: int,
52        markov_rank: int,
53        num_steps: int,
54        weight: torch.nn.Parameter,
55    ) -> None:
56        super().__init__()
57        self.proj = ReplicatedLinear(
58            input_size=n_embed + markov_rank,
59            output_size=1,
60            weight=weight,
61        )
62        temperature = envs.DEEPINFER_DFLASH_CONFIDENCE_HEAD_TEMPERATURE
63        temperatures = temperature if isinstance(temperature, tuple) else (temperature,) * num_steps
64        self.temperature = tuple(max(value, SAMPLING_EPS) for value in temperatures)
65
66    def forward(
67        self,
68        hidden: torch.Tensor,
69        markov_emb: torch.Tensor | None,
70        out: torch.Tensor | None = None,
71        step_i: int = 0,
72    ) -> torch.Tensor:
73        return confidence_head_fused(
74            hidden,
75            weight=self.proj.weight,
76            temperature=self.temperature[step_i],
77            markov_emb=markov_emb,
78            out=out,
79        )