From 6ce02b95ad9daeb476d4c1c84dde3c4b8ce735ff Mon Sep 17 00:00:00 2001 From: Khoa Pham Date: Fri, 3 Jul 2026 13:53:16 -0700 Subject: [PATCH] fix: reorder DSA indexer dual-stream ops to avoid CUDA graph stream explosion (#30025) Co-authored-by: Khoa Pham <264503018+kpham-sgl@users.noreply.github.com> --- python/sglang/srt/environ.py | 2 +- .../srt/layers/attention/dsa/dsa_indexer.py | 15 +++++++-------- 2 files changed, 8 insertions(+), 9 deletions(-) diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index e6c526451..6cf331007 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -668,7 +668,7 @@ class Envs: SGLANG_USE_FUSED_METADATA_COPY = EnvBool(True) SGLANG_DSA_USE_FUSED_METADATA_GENERATION = EnvBool(True) SGLANG_DSA_TOPK_BROADCAST = EnvBool(False) - SGLANG_DISABLE_DSA_INDEXER_FUSION = EnvBool(True) + SGLANG_DISABLE_DSA_INDEXER_FUSION = EnvBool(False) # sgl-kernel SGLANG_SKIP_SGL_KERNEL_VERSION_CHECK = EnvBool(False) diff --git a/python/sglang/srt/layers/attention/dsa/dsa_indexer.py b/python/sglang/srt/layers/attention/dsa/dsa_indexer.py index 90cf90637..a57215e50 100644 --- a/python/sglang/srt/layers/attention/dsa/dsa_indexer.py +++ b/python/sglang/srt/layers/attention/dsa/dsa_indexer.py @@ -739,6 +739,13 @@ class Indexer(MultiPlatformOp): current_stream.wait_stream(self.alt_stream) self.alt_stream.wait_stream(current_stream) + q_fp8, weights = fused_q_indexer_rope_first_quant( + q.contiguous(), + weights_raw, + q_scale_gate, + self._indexer_cos_sin_cache, + positions, + ) with torch.cuda.stream(self.alt_stream): self._fused_k_prepare_and_store( key, @@ -749,14 +756,6 @@ class Indexer(MultiPlatformOp): out_cache_loc=out_cache_loc, ) - q_fp8, weights = fused_q_indexer_rope_first_quant( - q.contiguous(), - weights_raw, - q_scale_gate, - self._indexer_cos_sin_cache, - positions, - ) - current_stream.wait_stream(self.alt_stream) return q_fp8, weights