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

[NPU][Bug fix] context parallel bug fix (#19820)

Browse Source
This commit is contained in:
liupeng374
2026-03-10 14:43:49 +08:00
committed by GitHub
parent 5f717913a0
commit 9b2e5526fb
1 changed files with 2 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
python/sglang/srt/hardware_backend/npu/modules/deepseek_v2_attention_mla_npu.py
View File

@@ -12,7 +12,6 @@ from sglang.srt.hardware_backend.npu.attention.mla_preprocess import (
)
from sglang.srt.layers.attention.nsa.nsa_indexer import scattered_to_tp_attn_full
from sglang.srt.layers.attention.nsa.utils import (
cp_split_and_rebuild_position,
nsa_use_prefill_cp,
)
from sglang.srt.layers.communicator import ScatterMode, get_attn_tp_context
@@ -224,12 +223,9 @@ def forward_mla_prepare_npu(
q_nope_out = q_nope_out.transpose(0, 1)
if nsa_use_prefill_cp(forward_batch, m.nsa_enable_prefill_cp):
positions = cp_split_and_rebuild_position(forward_batch, positions)
q_pe, k_pe = m.rotary_emb(positions, q_pe, k_pe)
if nsa_use_prefill_cp(forward_batch, m.nsa_enable_prefill_cp):
if nsa_use_prefill_cp(forward_batch):
# support allgather+rerrange
k_nope, k_pe = m.rebuild_cp_kv_cache(
latent_cache, forward_batch, k_nope, k_pe
@@ -367,12 +363,9 @@ def forward_dsa_prepare_npu(
q_nope_out = q_nope_out.transpose(0, 1)
if nsa_use_prefill_cp(forward_batch, m.nsa_enable_prefill_cp):
positions = cp_split_and_rebuild_position(forward_batch, positions)
q_pe, k_pe = m.rotary_emb(positions, q_pe, k_pe)
if nsa_use_prefill_cp(forward_batch, m.nsa_enable_prefill_cp):
if nsa_use_prefill_cp(forward_batch):
# support allgather+rerrange
k_nope, k_pe = m.rebuild_cp_kv_cache(
latent_cache, forward_batch, k_nope, k_pe