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

[qwen 3.8 next] reuse old cuda stream instead of endlessly creating streams (#39474)

Browse Source
This commit is contained in:
Qiaolin Yu
2026-09-15 01:36:14 -07:00
committed by GitHub
parent 8565b11003
commit 2c0a70960c
1 changed files with 4 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/models/qwen4_exp.py
+4 -4
View File
@@ -1548,10 +1548,6 @@ class Qwen4ExpAttentionDecoderLayer(
# and writes QSA-private pool buffers. # and writes QSA-private pool buffers.
current_stream = torch.cuda.current_stream() current_stream = torch.cuda.current_stream()
self.alt_stream.wait_stream(current_stream) self.alt_stream.wait_stream(current_stream)
with torch.cuda.stream(self.alt_stream):
topk_indices = self._compute_qsa_topk_indices(
hidden_states, positions, forward_batch
)
q, k, v, gate = self._prepare_qkv_gate( q, k, v, gate = self._prepare_qkv_gate(
positions=positions, positions=positions,
@@ -1560,6 +1556,10 @@ class Qwen4ExpAttentionDecoderLayer(
) )
if overlap_indexer: if overlap_indexer:
with torch.cuda.stream(self.alt_stream):
topk_indices = self._compute_qsa_topk_indices(
hidden_states, positions, forward_batch
)
current_stream.wait_stream(self.alt_stream) current_stream.wait_stream(self.alt_stream)
# Allocated on alt_stream, consumed by attention on the current # Allocated on alt_stream, consumed by attention on the current
# stream; tell the caching allocator before alt_stream is reused. # stream; tell the caching allocator before alt_stream is reused.