louis_lifu
/
vllm
mirror of https://github.com/vllm-project/vllm
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Remove Sampler copy stream (#2209)

This commit is contained in:
Antoni Baum 2023-12-20 00:04:33 -08:00 committed by GitHub
parent 31bff69151
commit bd29cf3d3a
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 4 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
vllm/model_executor/layers/sampler.py
View File

@ -30,7 +30,6 @@ class Sampler(nn.Module):
def __init__(self, vocab_size: int) -> None:
super().__init__()
self.vocab_size = vocab_size
self._copy_stream: torch.cuda.Stream = torch.cuda.Stream()
def forward(
self,
@ -51,14 +50,10 @@ class Sampler(nn.Module):
# Apply logits processors (if any).
logits = _apply_logits_processors(logits, sampling_metadata)
# Prepare sampling tensors in another stream to overlap
# CPU<->GPU data transfer with GPU computation in forward pass.
with torch.cuda.stream(self._copy_stream):
(sampling_tensors, do_penalties, do_top_p_top_k,
do_min_p) = SamplingTensors.from_sampling_metadata(
sampling_metadata, vocab_size, logits.device, logits.dtype)
torch.cuda.current_stream().wait_stream(self._copy_stream)
# Prepare sampling tensors with pinned memory to avoid blocking.
(sampling_tensors, do_penalties, do_top_p_top_k,
do_min_p) = SamplingTensors.from_sampling_metadata(
sampling_metadata, vocab_size, logits.device, logits.dtype)
# Apply presence and frequency penalties.
if do_penalties: