minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Fix DeepSeek V4 multistream QKV buffer lifetime (#36547)

Browse Source
This commit is contained in:
Aurick Qiao
2026-08-28 10:10:36 +08:00
committed by GitHub
parent cbd0271574
commit 84a6f51bd6
1 changed files with 2 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/models/deepseek_v4.py
+2 -4
View File
@@ -1059,8 +1059,6 @@ class MQALayer(MqaAttentionBase):
x_linear, positions, forward_batch, attn_backend, qkv_a=qkv_a x_linear, positions, forward_batch, attn_backend, qkv_a=qkv_a
) )
del qkv_a
if self.compressor is not None: if self.compressor is not None:
with torch.cuda.stream(stream_compressor): with torch.cuda.stream(stream_compressor):
attn_backend.forward_core_compressor( attn_backend.forward_core_compressor(
@@ -1071,6 +1069,7 @@ class MQALayer(MqaAttentionBase):
current_stream.wait_stream(stream_kv) current_stream.wait_stream(stream_kv)
current_stream.wait_stream(stream_compressor) current_stream.wait_stream(stream_compressor)
current_stream.wait_stream(stream_indexer) current_stream.wait_stream(stream_indexer)
del qkv_a
return q return q
@@ -1153,8 +1152,6 @@ class MQALayer(MqaAttentionBase):
q_out.copy_(q) q_out.copy_(q)
q.record_stream(stream_q) q.record_stream(stream_q)
del qkv_a
# Indexer + compressor: serial on current. # Indexer + compressor: serial on current.
if self.indexer is not None: if self.indexer is not None:
self.indexer( self.indexer(
@@ -1174,6 +1171,7 @@ class MQALayer(MqaAttentionBase):
# Join stream_kv + stream_q before downstream attention. # Join stream_kv + stream_q before downstream attention.
current_stream.wait_stream(stream_kv) current_stream.wait_stream(stream_kv)
current_stream.wait_stream(stream_q) current_stream.wait_stream(stream_q)
del qkv_a
return q return q
def _forward_prepare_multi_stream_hip( def _forward_prepare_multi_stream_hip(