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

This commit is contained in:
Qiaolin Yu
2026-09-15 01:36:14 -07:00
committed by GitHub
parent 8565b11003
commit 2c0a70960c
+4 -4
View File
@@ -1548,10 +1548,6 @@ class Qwen4ExpAttentionDecoderLayer(
# and writes QSA-private pool buffers.
current_stream = torch.cuda.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(
positions=positions,
@@ -1560,6 +1556,10 @@ class Qwen4ExpAttentionDecoderLayer(
)
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)
# Allocated on alt_stream, consumed by attention on the current
# stream; tell the caching allocator before alt_stream is reused.