EngineX-Hygon/sglang
5
0
Fork 0
Code Issues Pull Requests Actions 7 Projects Releases Wiki Activity

Fix a regression introduced by overlapping KV cache writing (#4375)

Browse Source
This commit is contained in:
Lianmin Zheng
2025-03-13 03:49:05 -07:00
committed by GitHub
parent 6aaeb84872
commit 4fea040ca1
1 changed files with 1 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
python/sglang/srt/mem_cache/memory_pool.py
View File

@@ -326,7 +326,7 @@ class MHATokenToKVPool(KVCache):
cache_k = cache_k.view(self.store_dtype) cache_k = cache_k.view(self.store_dtype)
cache_v = cache_v.view(self.store_dtype) cache_v = cache_v.view(self.store_dtype)
if self.capture_mode: if self.capture_mode and cache_k.shape[0] < 4:
self.alt_stream.wait_stream(torch.cuda.current_stream()) self.alt_stream.wait_stream(torch.cuda.current_stream())
with torch.cuda.stream(self.alt_stream): with torch.cuda.stream(self.alt_stream):
self.k_buffer[layer_id][loc] = cache_k self.k_buffer[layer_id][loc] = cache_k