[CORE] Improvement in ranks code (#4718)

2024-05-12 20:47:47 -04:00 · 2024-05-12 20:47:47 -04:00 · a7be4d0072
commit a7be4d0072
parent a709e87a4f
1 changed files with 3 additions and 1 deletions
--- 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(