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

[CORE] Improvement in ranks code (#4718)

Browse Source
This commit is contained in:
Swapnil Parekh 2024-05-12 20:47:47 -04:00 committed by GitHub
parent a709e87a4f
commit a7be4d0072
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 3 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -681,7 +681,9 @@ def _get_ranks(x: torch.Tensor, indices: torch.Tensor) -> torch.Tensor:
""" """
vals = x[torch.arange(0, len(x), device=x.device, dtype=indices.dtype), vals = x[torch.arange(0, len(x), device=x.device, dtype=indices.dtype),
indices] indices]
return (x > vals[:, None]).long().sum(1).add_(1) result = (x > vals[:, None])
del vals
return result.sum(1).add_(1)
def _get_logprobs( def _get_logprobs(