Skip unnecessary penalizer (#1707)

2024-10-18 17:54:03 -07:00
parent bc12d4033f
commit 2bcfba1b08
7 changed files with 104 additions and 75 deletions
--- a/python/sglang/srt/mem_cache/memory_pool.py
+++ b/python/sglang/srt/mem_cache/memory_pool.py
@@ -90,7 +90,7 @@ class BaseTokenToKVPool:
        select_index = self.free_slots[:need_size]
        self.free_slots = self.free_slots[need_size:]

-        return select_index.to(self.device)
+        return select_index.to(self.device, non_blocking=True)

    def free(self, free_index: torch.Tensor):
        if self.is_not_in_free_group: