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

【NPU】【bugfix】accuracy fix when enable both nsa cp and prefixcache (#23268)

Browse Source
This commit is contained in:
cen121212
2026-04-28 09:08:28 +08:00
committed by GitHub
parent 9ffc0cc67e
commit b1e1fe8eee
2 changed files with 21 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/hardware_backend/npu/modules/deepseek_v2_attention_mla_npu.py
+3 -1
View File
@@ -359,7 +359,9 @@ def forward_dsa_prepare_npu(
if q_event is not None:
torch.npu.current_stream().wait_event(q_event)
else:
if fused_qkv_a_proj_out.shape[0] < 65535:
if fused_qkv_a_proj_out.shape[0] < 65535 and not nsa_use_prefill_cp(
forward_batch
):
q_lora, k_nope, k_pe = fused_split_qk_norm(
fused_qkv_a_proj_out,
m.q_a_layernorm,
python/sglang/srt/layers/attention/nsa/nsa_indexer.py
+18 -4
View File
@@ -1463,10 +1463,24 @@ class Indexer(MultiPlatformOp):
forward_batch.attn_cp_metadata.actual_seq_q_prev_tensor,
forward_batch.attn_cp_metadata.actual_seq_q_next_tensor,
)
forward_batch.attn_backend.forward_metadata.actual_seq_lengths_kv = (
forward_batch.attn_cp_metadata.kv_len_prev_tensor,
forward_batch.attn_cp_metadata.kv_len_next_tensor,
)
if sum(forward_batch.extend_prefix_lens_cpu) > 0:
total_kv_len_prev_tensor = (
forward_batch.attn_cp_metadata.kv_len_prev_tensor
+ forward_batch.extend_prefix_lens.squeeze()
)
total_kv_len_next_tensor = (
forward_batch.attn_cp_metadata.kv_len_next_tensor
+ forward_batch.extend_prefix_lens.squeeze()
)
forward_batch.attn_backend.forward_metadata.actual_seq_lengths_kv = (
total_kv_len_prev_tensor,
total_kv_len_next_tensor,
)
else:
forward_batch.attn_backend.forward_metadata.actual_seq_lengths_kv = (
forward_batch.attn_cp_metadata.kv_len_prev_tensor,
forward_batch.attn_cp_metadata.kv_len_next_tensor,
)
actual_seq_lengths_q = (
forward_batch.attn_backend.forward_metadata.actual_seq_lengths_q
)