xinyun/vllm
1
0
Fork 0
mirror of https://git.datalinker.icu/vllm-project/vllm.git synced 2025-12-14 00:34:58 +08:00
Code Issues Packages Projects Releases Wiki Activity

[Misc]: optimize eager mode host time (#4196)

Browse Source 
Co-authored-by: xuhao <xuhao@cambricon.com>
This commit is contained in:
functionxu123 2024-05-31 13:14:50 +08:00 committed by GitHub
parent e9d3aa04f6
commit a377f0bd5e
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 5 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
vllm/utils.py
View File

@ -17,6 +17,7 @@ from typing import (Any, AsyncIterator, Awaitable, Callable, Dict, Generic,
Hashable, List, Optional, OrderedDict, Tuple, TypeVar, Hashable, List, Optional, OrderedDict, Tuple, TypeVar,
Union) Union)
import numpy as np
import psutil import psutil
import torch import torch
@ -501,11 +502,6 @@ def str_to_int_tuple(s: str) -> Tuple[int, ...]:
f"(e.g., 1, 2, 3). Given input: {s}") from e f"(e.g., 1, 2, 3). Given input: {s}") from e
def pad_to_max_length(x: List[int], max_len: int, pad: int) -> List[int]:
assert len(x) <= max_len
return x + [pad] * (max_len - len(x))
def make_tensor_with_pad( def make_tensor_with_pad(
x: List[List[int]], x: List[List[int]],
max_len: int, max_len: int,
@ -518,7 +514,10 @@ def make_tensor_with_pad(
The padding is applied to the end of each inner list until it reaches The padding is applied to the end of each inner list until it reaches
`max_len`. `max_len`.
""" """
padded_x = [pad_to_max_length(x_i, max_len, pad) for x_i in x] padded_x = np.zeros([len(x), max_len], dtype=np.int32) + pad
for ind, blocktb in enumerate(x):
assert len(blocktb) <= max_len
padded_x[ind, :len(blocktb)] = blocktb
return torch.tensor(padded_x, dtype=dtype, device=device) return torch.tensor(padded_x, dtype=dtype, device=device)