Simulates standard PyTorch sampling logic (Top-K -> Top-P -> Softmax -> Multinomial)
(logits, top_k=None, top_p=None, temperature=1.0)
| 13 | |
| 14 | def native_sampling(logits, top_k=None, top_p=None, temperature=1.0): |
| 15 | """ |
| 16 | Simulates standard PyTorch sampling logic (Top-K -> Top-P -> Softmax -> Multinomial) |
| 17 | """ |
| 18 | if temperature != 1.0: |
| 19 | logits = logits / temperature |
| 20 | |
| 21 | logits = logits.float() |
| 22 | |
| 23 | if top_k is not None: |
| 24 | v, _ = torch.topk(logits, min(top_k, logits.size(-1))) |
| 25 | logits[logits < v[:, [-1]]] = float("-inf") |
| 26 | |
| 27 | if top_p is not None: |
| 28 | sorted_logits, sorted_indices = torch.sort(logits, descending=True) |
| 29 | cumulative_probs = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1) |
| 30 | sorted_indices_to_remove = cumulative_probs > top_p |
| 31 | sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone() |
| 32 | sorted_indices_to_remove[..., 0] = 0 |
| 33 | |
| 34 | for i in range(logits.size(0)): |
| 35 | indices_to_remove = sorted_indices[i][sorted_indices_to_remove[i]] |
| 36 | logits[i, indices_to_remove] = float("-inf") |
| 37 | |
| 38 | probs = torch.softmax(logits, dim=-1) |
| 39 | return torch.multinomial(probs, num_samples=1) |
| 40 | |
| 41 | |
| 42 | def run_benchmark(args, return_data: bool = False): |
| 43 | device = device_ctx.device |
no outgoing calls
no test coverage detected