Skip to content

[Refactor] Deduplicate NSA utils.py into cp_utils.py for context parallel #22914

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Fridge003 merged 7 commits into from
Apr 20, 2026
Merged

[Refactor] Deduplicate NSA utils.py into cp_utils.py for context parallel #22914

Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
6d3ee63
[Refactor] Deduplicate NSA utils.py into cp_utils.py for context para…
Fridge003 Apr 15, 2026
bb069d2
some fixes
Fridge003 Apr 16, 2026
93ccfc7
Fix prefix_len double-count in _get_topk_ragged_with_cp
Fridge003 Apr 17, 2026
4b0ee1a
Keep extend-only kv_len_prev/next so NSA CP handles multi-batch extend
Fridge003 Apr 18, 2026
f351936
Gate extend-only kv_len_prev/next on NSA; keep prefix-baked path for …
Fridge003 Apr 18, 2026
67f3914
Merge branch 'main' into dedup-nsa-utils-cp-utils
Fridge003 Apr 20, 2026
63d3e85
Merge branch 'main' into dedup-nsa-utils-cp-utils
Fridge003 Apr 20, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -807,7 +807,7 @@ def forward_sparse(
if (
is_prefill
and is_nsa_enable_prefill_cp()
and forward_batch.nsa_cp_metadata is not None
and forward_batch.attn_cp_metadata is not None
):
attn_out = self.do_cp_balance_attn(
q_nope,
Expand Down
46 changes: 24 additions & 22 deletions python/sglang/srt/layers/attention/nsa/nsa_indexer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -51,14 +51,14 @@
from sglang.srt.distributed.parallel_state import get_pp_group
from sglang.srt.layers import deep_gemm_wrapper
from sglang.srt.layers.attention.nsa.utils import (
cp_all_gather_rerange_output,
is_nsa_enable_prefill_cp,
is_nsa_prefill_cp_in_seq_split,
)
from sglang.srt.layers.communicator import ScatterMode
from sglang.srt.layers.linear import ReplicatedLinear
from sglang.srt.layers.quantization.base_config import QuantizationConfig
from sglang.srt.layers.rotary_embedding import get_rope_wrapper
from sglang.srt.layers.utils.cp_utils import cp_all_gather_rerange_output
from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.server_args import get_global_server_args
Expand Down Expand Up @@ -355,7 +355,7 @@ def _get_q_k_bf16(
current_stream.wait_stream(self.alt_stream)
elif (
self.alt_stream is not None
and forward_batch.nsa_cp_metadata is not None
and forward_batch.attn_cp_metadata is not None
and self.nsa_enable_prefill_cp
):
key = rotate_activation(key)
Expand All @@ -377,7 +377,7 @@ def _get_q_k_bf16(
key = rotate_activation(key)

# allgather+rerrange
if forward_batch.nsa_cp_metadata is not None and self.nsa_enable_prefill_cp:
if forward_batch.attn_cp_metadata is not None and self.nsa_enable_prefill_cp:
key = cp_all_gather_rerange_output(
key.contiguous(),
self.cp_size,
Expand Down Expand Up @@ -858,11 +858,10 @@ def _get_topk_ragged_with_cp(
batch_idx_list=batch_idx_list,
)
else:
kv_len = (
forward_batch.seq_lens_cpu[0].item()
- forward_batch.extend_seq_lens_cpu[0]
+ kv_len
)
# kv_len already includes the prefix offset
# (prepare_context_parallel_metadata sums prefix_len into kv_len_prev/next).
# Do NOT add (seq_lens_cpu - extend_seq_lens_cpu) again here, or the prefix
# would be double-counted and get_index_k_continuous reads past the block table.
Comment thread
Fridge003 marked this conversation as resolved.
Outdated
k_fp8 = forward_batch.token_to_kv_pool.get_index_k_continuous(
layer_id,
kv_len,
Expand Down Expand Up @@ -1212,17 +1211,17 @@ def forward_cuda(
)
else:
if (
forward_batch.nsa_cp_metadata is not None
forward_batch.attn_cp_metadata is not None
and is_nsa_prefill_cp_in_seq_split()
):
kv_len_prev = forward_batch.nsa_cp_metadata.kv_len_prev
kv_len_next = forward_batch.nsa_cp_metadata.kv_len_next
actual_seq_q_prev = forward_batch.nsa_cp_metadata.actual_seq_q_prev
actual_seq_q_next = forward_batch.nsa_cp_metadata.actual_seq_q_next
kv_len_prev = forward_batch.attn_cp_metadata.kv_len_prev
kv_len_next = forward_batch.attn_cp_metadata.kv_len_next
actual_seq_q_prev = forward_batch.attn_cp_metadata.actual_seq_q_prev
actual_seq_q_next = forward_batch.attn_cp_metadata.actual_seq_q_next

# TODO support mutil-batch
# cp_batch_seq_index_prev = forward_batch.nsa_cp_metadata["cp_batch_seq_index_prev"]
# cp_batch_seq_index_next = forward_batch.nsa_cp_metadata["cp_batch_seq_index_next"]
# cp_batch_seq_index_prev = forward_batch.attn_cp_metadata["cp_batch_seq_index_prev"]
# cp_batch_seq_index_next = forward_batch.attn_cp_metadata["cp_batch_seq_index_next"]
# TODO prev, next, combined into a single call
q_fp8_prev, q_fp8_next = torch.split(
q_fp8, (q_fp8.shape[0] + 1) // 2, dim=0
Expand Down Expand Up @@ -1423,7 +1422,7 @@ def forward_npu(
if (
is_prefill
and self.nsa_enable_prefill_cp
and forward_batch.nsa_cp_metadata is not None
and forward_batch.attn_cp_metadata is not None
):
k = cp_all_gather_rerange_output(
k.contiguous().view(-1, self.head_dim),
Expand All @@ -1436,14 +1435,17 @@ def forward_npu(
layer_id, forward_batch.out_cache_loc, k
)
if is_prefill:
if self.nsa_enable_prefill_cp and forward_batch.nsa_cp_metadata is not None:
if (
self.nsa_enable_prefill_cp
and forward_batch.attn_cp_metadata is not None
):
forward_batch.attn_backend.forward_metadata.actual_seq_lengths_q = (
forward_batch.nsa_cp_metadata.actual_seq_q_prev_tensor,
forward_batch.nsa_cp_metadata.actual_seq_q_next_tensor,
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.nsa_cp_metadata.kv_len_prev_tensor,
forward_batch.nsa_cp_metadata.kv_len_next_tensor,
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
Expand Down Expand Up @@ -1498,7 +1500,7 @@ def forward_npu(
if (
is_prefill
and self.nsa_enable_prefill_cp
and forward_batch.nsa_cp_metadata is not None
and forward_batch.attn_cp_metadata is not None
):
block_table = block_table[: actual_seq_lengths_q[0].numel()]
topk_indices = self.do_npu_cp_balance_indexer(
Expand Down
Loading
Loading