nano-sampler

nano-vllm/nanovllm/layers/sampler.py · bb823b3 · L5–12 · 固定 commit 上游

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

5class Sampler(nn.Module):
6
7    @torch.compile
8    def forward(self, logits: torch.Tensor, temperatures: torch.Tensor):
9        logits = logits.float().div_(temperatures.unsqueeze(dim=1))
10        probs = torch.softmax(logits, dim=-1)
11        sample_tokens = probs.div_(torch.empty_like(probs).exponential_(1).clamp_min_(1e-10)).argmax(dim=-1)
12        return sample_tokens