From a8dda2aa5727a7dfa54fa7323d334cb6e0aa69e0 Mon Sep 17 00:00:00 2001 From: Ziang Li <106564213+zianglih@users.noreply.github.com> Date: Tue, 27 Jan 2026 19:46:49 -0800 Subject: [PATCH] [DSv32] Overlap indexer qk projection and activation quant (#17688) --- python/sglang/srt/layers/attention/nsa/nsa_indexer.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py index 7d7b26d62..659236033 100644 --- a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py +++ b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py @@ -956,11 +956,11 @@ class Indexer(MultiPlatformOp): current_stream = torch.cuda.current_stream() self.alt_stream.wait_stream(current_stream) weights = self._project_and_scale_head_gates(x) + query, key = self._get_q_k_bf16( + q_lora, x, positions, enable_dual_stream, forward_batch=forward_batch + ) + q_fp8, q_scale = act_quant(query, self.block_size, self.scale_fmt) with torch.cuda.stream(self.alt_stream): - query, key = self._get_q_k_bf16( - q_lora, x, positions, False, forward_batch=forward_batch - ) - q_fp8, q_scale = act_quant(query, self.block_size, self.scale_fmt) k_fp8, k_scale = act_quant(key, self.block_size, self.scale_fmt) current_stream.wait_stream(self.alt_stream) weights = weights.unsqueeze(-1) * q_scale * self.softmax_scale