Fix minor style (#4460)

2025-03-15 21:51:12 -07:00
parent 158430473e
commit 2c4f5ccac1
3 changed files with 7 additions and 4 deletions
--- a/python/sglang/srt/mem_cache/memory_pool.py
+++ b/python/sglang/srt/mem_cache/memory_pool.py
@@ -340,6 +340,7 @@ class MHATokenToKVPool(KVCache):
            cache_v = cache_v.view(self.store_dtype)

        if self.capture_mode and cache_k.shape[0] < 4:
+            # Overlap the copy of K and V cache for small batch size
            current_stream = self.device_module.current_stream()
            self.alt_stream.wait_stream(current_stream)
            with self.device_module.stream(self.alt_stream):