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

[DSv32] Overlap indexer weights_proj during dual_stream decode (#16637)

Browse Source 
Co-authored-by: Ziang Li <ziangli@humansand.ai>
This commit is contained in:
Ziang Li
2026-01-09 21:06:44 -08:00
committed by GitHub
parent 32a569fb77
commit 20abaee26c
2 changed files with 64 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files

41
python/sglang/srt/layers/attention/nsa/nsa_indexer.py
View File

@@ -205,6 +205,12 @@ class Indexer(MultiPlatformOp):
else:
yield
@torch.compile(dynamic=True)
def _project_and_scale_head_gates(self, x: torch.Tensor):
weights, _ = self.weights_proj(x.float())
weights = weights * self.n_heads**-0.5
return weights
@torch.compile(dynamic=True)
def _get_logits_head_gate(self, x: torch.Tensor, q_scale: torch.Tensor):
weights, _ = self.weights_proj(x.float())
@@ -856,21 +862,36 @@ class Indexer(MultiPlatformOp):
return_indices,
)
query, key = self._get_q_k_bf16(
q_lora, x, positions, enable_dual_stream, forward_batch=forward_batch
)
if enable_dual_stream:
if enable_dual_stream and forward_batch.forward_mode.is_decode_or_idle():
current_stream = torch.cuda.current_stream()
self.alt_stream.wait_stream(current_stream)
q_fp8, q_scale = act_quant(query, self.block_size, self.scale_fmt)
weights = self._project_and_scale_head_gates(x)
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
else:
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)
query, key = self._get_q_k_bf16(
q_lora, x, positions, enable_dual_stream, forward_batch=forward_batch
)
if enable_dual_stream:
current_stream = torch.cuda.current_stream()
self.alt_stream.wait_stream(current_stream)
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)
current_stream.wait_stream(self.alt_stream)
else:
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)
weights = self._get_logits_head_gate(x, q_scale)
# k_fp8: (seq_len, head_dim) fp8_e4m3fn
# k_buffer: (num_total_tokens + page_size, head_dim) fp8_e4m3fn
@@ -885,8 +906,6 @@ class Indexer(MultiPlatformOp):
index_k_scale=k_scale,
)
weights = self._get_logits_head_gate(x, q_scale)
if is_cuda():
assert forward_batch.seq_lens_cpu is not None
if len(forward_batch.seq_lens_cpu) == 0:

45
python/sglang/srt/models/deepseek_v2.py
View File

@@ -1666,6 +1666,7 @@ class DeepseekV2AttentionMLA(nn.Module):
from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode
q_lora = None
topk_indices = None
if self.q_lora_rank is not None:
q, latent_cache = (
get_attn_tp_context()
@@ -1722,8 +1723,39 @@ class DeepseekV2AttentionMLA(nn.Module):
if self.use_nsa:
q_lora = q
k_nope = k_nope.unsqueeze(1)
q = self.q_b_proj(q)[0].view(-1, self.num_local_heads, self.qk_head_dim)
# overlap q_b_proj and indexer during decode
if (
self.alt_stream is not None
and get_is_capture_mode()
and forward_batch.forward_mode.is_decode_or_idle()
and q_lora is not None
):
current_stream = torch.cuda.current_stream()
self.alt_stream.wait_stream(current_stream)
with torch.cuda.stream(self.alt_stream):
k_nope = k_nope.unsqueeze(1)
q = self.q_b_proj(q)[0].view(
-1, self.num_local_heads, self.qk_head_dim
)
topk_indices = self.indexer(
x=hidden_states,
q_lora=q_lora,
positions=positions,
forward_batch=forward_batch,
layer_id=self.layer_id,
)
current_stream.wait_stream(self.alt_stream)
else:
k_nope = k_nope.unsqueeze(1)
q = self.q_b_proj(q)[0].view(-1, self.num_local_heads, self.qk_head_dim)
if q_lora is not None:
topk_indices = self.indexer(
x=hidden_states,
q_lora=q_lora,
positions=positions,
forward_batch=forward_batch,
layer_id=self.layer_id,
)
else:
q = self.q_proj(hidden_states)[0].view(
-1, self.num_local_heads, self.qk_head_dim
@@ -1818,15 +1850,6 @@ class DeepseekV2AttentionMLA(nn.Module):
k_nope, k_pe = self.rebuild_cp_kv_cache(
latent_cache, forward_batch, k_nope, k_pe
)
topk_indices = None
if q_lora is not None:
topk_indices = self.indexer(
x=hidden_states,
q_lora=q_lora,
positions=positions,
forward_batch=forward_batch,
layer_id=self.layer_id,
)
return (
q_pe,