| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import torch |
| 5 | |
| 6 | from vllm.triton_utils import triton |
| 7 | |
| 8 | |
| 9 | def set_triton_allocator(device: torch.device): |
| 10 | def alloc_fn(size: int, alignment: int, stream: int | None): |
| 11 | return torch.empty(size, device=device, dtype=torch.int8) |
| 12 | |
| 13 | triton.set_allocator(alloc_fn) |