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

Remove Sampler copy stream (#2209)

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