| 651 | triton.set_allocator(alloc_on_cuda) |
| 652 | |
| 653 | |
| 654 | def alloc_on_cuda(size: int, alignment: int, stream: int | None): |
| 655 | return torch.empty(size, device="cuda", dtype=torch.int8) |
| 656 | |
| 657 | |
| 658 | class Sampler(Protocol): |
| 659 | def prepare(self) -> "Sampler": |
| 660 | raise NotImplementedError() |
| 661 | |
| 662 | def sample(self, **kwargs) -> torch.Tensor: |
| 663 | raise NotImplementedError() |