Unified sampling interface.
(self, logits, top_k=None, top_p=None, temperature=1.0, deterministic=True)
| 40 | |
| 41 | @torch.inference_mode() |
| 42 | def sample(self, logits, top_k=None, top_p=None, temperature=1.0, deterministic=True): |
| 43 | """ |
| 44 | Unified sampling interface. |
| 45 | """ |
| 46 | logits = logits.contiguous() |
| 47 | if temperature != 1.0: |
| 48 | logits = logits / temperature |
| 49 | |
| 50 | if self.backend == constants.BackendLib.FLASHINFER.value: |
| 51 | from flashinfer.sampling import top_k_renorm_probs, top_p_sampling_from_probs |
| 52 | |
| 53 | logits = logits.float().contiguous() |
| 54 | probs = torch.softmax(logits, dim=-1) |
| 55 | |
| 56 | if top_k is None and top_p is None: |
| 57 | return torch.multinomial(probs, num_samples=1).view(-1) |
| 58 | |
| 59 | if top_k is not None: |
| 60 | probs = top_k_renorm_probs(probs, top_k) |
| 61 | |
| 62 | if top_p is not None: |
| 63 | return top_p_sampling_from_probs(probs, top_p, deterministic=deterministic) |
| 64 | |
| 65 | return torch.multinomial(probs, num_samples=1).view(-1) |
| 66 | |
| 67 | elif self.backend == constants.BackendLib.AITER.value: |
| 68 | # TODO: Connect to AITER's sampling operator |
| 69 | # return aiter.ops.sample(logits, ...) |
| 70 | pass |
| 71 | |
| 72 | # Fallback to native PyTorch sampling |
| 73 | if top_k is not None: |
| 74 | topk_values, _ = torch.topk(logits, top_k) |
| 75 | min_topk = topk_values[..., -1, None] |
| 76 | logits = torch.where(logits < min_topk, torch.full_like(logits, float("-inf")), logits) |
| 77 | probs = torch.softmax(logits, dim=-1) |
| 78 | return torch.multinomial(probs, num_samples=1).squeeze(-1) |
| 79 | |
| 80 | @torch.inference_mode() |
| 81 | def compute_logp(self, logits, token_ids): |