Upgrade to vllm 0.17.0 corex v4.1 overlay
This commit is contained in:
13
vllm/triton_utils/allocation.py
Normal file
13
vllm/triton_utils/allocation.py
Normal file
@@ -0,0 +1,13 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import torch
|
||||
|
||||
from vllm.triton_utils import triton
|
||||
|
||||
|
||||
def set_triton_allocator(device: torch.device):
|
||||
def alloc_fn(size: int, alignment: int, stream: int | None):
|
||||
return torch.empty(size, device=device, dtype=torch.int8)
|
||||
|
||||
triton.set_allocator(alloc_fn)
|
||||
Reference in New Issue
Block a user