biondizzle/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

[5/N][torch.compile] torch.jit.script --> torch.compile (#10406)

Browse Source 
Signed-off-by: youkaichao <youkaichao@gmail.com>
This commit is contained in:
youkaichao
2024-11-18 07:20:06 -08:00
committed by GitHub
parent 4186be8111
commit 7851b45196
4 changed files with 6 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
vllm/model_executor/layers/vocab_parallel_embedding.py
View File

@@ -133,13 +133,13 @@ class VocabParallelEmbeddingShardIndices:
assert self.num_added_elements <= self.num_added_elements_padded
@torch.jit.script
@torch.compile(dynamic=True)
def get_masked_input_and_mask(
input_: torch.Tensor, org_vocab_start_index: int,
org_vocab_end_index: int, num_org_vocab_padding: int,
added_vocab_start_index: int,
added_vocab_end_index: int) -> Tuple[torch.Tensor, torch.Tensor]:
# torch.jit.script will fuse all of the pointwise ops below
# torch.compile will fuse all of the pointwise ops below
# into a single kernel, making it very fast
org_vocab_mask = (input_ >= org_vocab_start_index) & (input_ <
org_vocab_end_index)