[CORE] Improvement in ranks code (#4718)

a7be4d00 · Swapnil Parekh · GitHub · a709e87a · a7be4d00
Unverified Commit a7be4d00 authored May 12, 2024 by Swapnil Parekh Committed by GitHub May 12, 2024
Show whitespace changes
Inline Side-by-side

Showing with 3 additions and 1 deletion

vllm/model_executor/layers/sampler.py vllm/model_executor/layers/sampler.py +3 -1

No files found.
--- a/vllm/model_executor/layers/sampler.py
+++ b/vllm/model_executor/layers/sampler.py
@@ -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),
             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(