Improve _prune_hidden_states micro-benchmark (#707)

28873a27 · Aman Gupta Karmani · GitHub · 0080d832 · 28873a27
Unverified Commit 28873a27 authored Aug 31, 2023 by Aman Gupta Karmani Committed by GitHub Aug 31, 2023
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 1 deletion

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

No files found.
--- a/vllm/model_executor/layers/sampler.py
+++ b/vllm/model_executor/layers/sampler.py
@@ -100,7 +100,8 @@ def _prune_hidden_states(
        start_idx += prompt_len
    last_token_indicies.extend(
        range(start_idx, start_idx + input_metadata.num_generation_tokens))
-    return hidden_states[last_token_indicies]
+    return hidden_states.index_select(
+        0, torch.tensor(last_token_indicies, device=hidden_states.device))


 def _get_penalties(