huggingface/text-generation-inference
5
0
Fork 0
mirror of https://github.com/huggingface/text-generation-inference.git synced 2025-09-12 12:54:52 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: allocate tmp based on sgmv kernel if available

Browse Source
This commit is contained in:
drbh 2024-07-31 22:15:52 +00:00
parent 7451041ecd
commit e4a0bf3b71
1 changed files with 8 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
server/text_generation_server/utils/sgmv.py
View File

@ -151,13 +151,17 @@ def get_tmp_expand_size(size: int) -> int:
def get_tmp_tensors(
nsegments: int, lora_rank: int, device: torch.device
) -> Tuple[torch.Tensor, torch.Tensor]:
if use_cutlass_shrink(lora_rank) and has_sgmv():
use_cutlass = use_cutlass_shrink(lora_rank) and has_sgmv()
has_sgmv_available = has_sgmv()
if use_cutlass:
tmp = get_tmp_tensor_for_size(nsegments, device)
return tmp, tmp
elif has_sgmv_available:
return get_tmp_tensor(device), get_tmp_tensor_for_size(nsegments, device)
else:
tmp_shrink = get_tmp_tensor(device)
tmp_expand = get_tmp_tensor_for_size_no_kernels(nsegments, device)
return tmp_shrink, tmp_expand
tmp = get_tmp_tensor_for_size(nsegments, device)
return tmp, tmp
def lora_a_sgmv_cutlass(