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

[DSv32] Overlap indexer qk projection and activation quant (#17688)

Browse Source
This commit is contained in:
Ziang Li
2026-01-28 11:46:49 +08:00
committed by GitHub
parent 52bca42870
commit a8dda2aa57
1 changed files with 4 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/nsa/nsa_indexer.py
+2 -2
View File
@@ -956,11 +956,11 @@ class Indexer(MultiPlatformOp):
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)
weights = self._project_and_scale_head_gates(x) weights = self._project_and_scale_head_gates(x)
with torch.cuda.stream(self.alt_stream):
query, key = self._get_q_k_bf16( query, key = self._get_q_k_bf16(
q_lora, x, positions, False, forward_batch=forward_batch q_lora, x, positions, enable_dual_stream, forward_batch=forward_batch
) )
q_fp8, q_scale = act_quant(query, self.block_size, self.scale_fmt) q_fp8, q_scale = act_quant(query, self.block_size, self.scale_fmt)
with torch.cuda.stream(self.alt_stream):
k_fp8, k_scale = act_quant(key, 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) current_stream.wait_stream(self.alt_stream)
weights = weights.unsqueeze(-1) * q_scale * self.softmax_scale weights = weights.unsqueeze(-1) * q_scale * self.softmax_scale