| 30 | |
| 31 | |
| 32 | def prepare_inputs( |
| 33 | batch_size: int, |
| 34 | seq_len: int, |
| 35 | varlen: bool, |
| 36 | vocab_size: int, |
| 37 | device: torch.device |
| 38 | ): |
| 39 | if varlen: |
| 40 | tokens = torch.randint(high=vocab_size, size=(1, batch_size * seq_len), device=device) |
| 41 | offsets = torch.cat([ |
| 42 | torch.tensor([0], dtype=torch.long, device=device), |
| 43 | torch.randperm(batch_size * seq_len - 16, device=device)[:batch_size-1] + 16, |
| 44 | torch.tensor([batch_size * seq_len], dtype=torch.long, device=device) |
| 45 | ], 0).sort()[0] |
| 46 | else: |
| 47 | tokens = torch.randint(high=vocab_size, size=(batch_size, seq_len), device=device) |
| 48 | offsets = None |
| 49 | return tokens, offsets |
| 50 | |
| 51 | |
| 52 | def profile( |