From 649172879778646a8574525448734b95a55b6739 Mon Sep 17 00:00:00 2001 From: Baidu-AIAK Date: Tue, 24 Mar 2026 12:31:48 +0800 Subject: [PATCH] [Perf] Overlap NSA-CP key all-gather with query computation for DeepSeek-V3.2 (#20438) Co-authored-by: Shurui Jia <18817781975@163.com> Co-authored-by: Baidu-AIAK --- .../srt/layers/attention/nsa/nsa_indexer.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py index bff848c1b..7d1511963 100644 --- a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py +++ b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py @@ -329,6 +329,25 @@ class Indexer(MultiPlatformOp): with torch.cuda.stream(self.alt_stream): key = rotate_activation(key) current_stream.wait_stream(self.alt_stream) + elif ( + self.alt_stream is not None + and forward_batch.nsa_cp_metadata is not None + and self.nsa_enable_prefill_cp + ): + key = rotate_activation(key) + current_stream = torch.cuda.current_stream() + self.alt_stream.wait_stream(current_stream) + query = rotate_activation(query) + + with torch.cuda.stream(self.alt_stream): + key = cp_all_gather_rerange_output( + key.contiguous(), + self.cp_size, + forward_batch, + torch.cuda.current_stream(), + ) + current_stream.wait_stream(self.alt_stream) + return query, key else: query = rotate_activation(query) key = rotate_activation(key)