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

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

Browse Source 
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
Baizhou Zhang
2026-04-19 21:35:35 -07:00
committed by GitHub
co-authored by Claude Opus 4.6
parent eb76aaba88
commit c304d0d64d
8 changed files with 148 additions and 402 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py
+1 -1
View File
@@ -807,7 +807,7 @@ class AscendAttnBackend(AttentionBackend):
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,
python/sglang/srt/layers/attention/nsa/nsa_indexer.py
+20 -17
View File
@@ -53,7 +53,6 @@ from sglang.srt.distributed import (
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,
)
@@ -61,6 +60,7 @@ 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
@@ -358,7 +358,7 @@ class Indexer(MultiPlatformOp):
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)
@@ -380,7 +380,7 @@ class Indexer(MultiPlatformOp):
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,
@@ -1231,17 +1231,17 @@ class Indexer(MultiPlatformOp):
)
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
@@ -1442,7 +1442,7 @@ class Indexer(MultiPlatformOp):
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),
@@ -1455,14 +1455,17 @@ class Indexer(MultiPlatformOp):
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
@@ -1517,7 +1520,7 @@ class Indexer(MultiPlatformOp):
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(
python/sglang/srt/layers/attention/nsa/utils.py
+2 -353
View File
@@ -1,24 +1,14 @@
# temp NSA debugging environ
from dataclasses import dataclass
from itertools import accumulate
from typing import TYPE_CHECKING, List, Tuple, Union
import torch
import torch.nn.functional as F
import triton
import triton.language as tl
from sglang.srt.distributed.device_communicators.pynccl_allocator import (
use_symmetric_memory,
)
from sglang.srt.layers.dp_attention import (
DpPaddingMode,
attn_cp_all_gather_into_tensor,
get_attention_cp_group,
get_attention_cp_rank,
get_attention_cp_size,
get_attention_dp_rank,
is_allocation_symmetric,
)
from sglang.srt.server_args import get_global_server_args
from sglang.srt.utils.common import ceil_align, ceil_div
@@ -135,27 +125,7 @@ def pad_nsa_cache_seqlens(forward_batch: "ForwardBatch", nsa_cache_seqlens):
return nsa_cache_seqlens
@dataclass
class NSAContextParallelMetadata:
split_list: List[int] = None
max_rank_len: List[int] = None
zigzag_index: List[int] = None
per_rank_actual_token: List[int] = None
reverse_split_len: List[int] = None
cp_reverse_index: List[int] = None
kv_len_prev: int = -1
kv_len_next: int = -1
actual_seq_q_prev: int = -1
actual_seq_q_next: int = -1
kv_len_prev_tensor: torch.Tensor = None
kv_len_next_tensor: torch.Tensor = None
actual_seq_q_prev_tensor: torch.Tensor = None
actual_seq_q_next_tensor: torch.Tensor = None
total_seq_lens: torch.Tensor = None
def can_cp_split(seq_len: int, cp_size: int, use_nsa: bool, forward_batch):
def can_nsa_cp_split(seq_len: int, cp_size: int, use_nsa: bool, forward_batch):
if is_nsa_prefill_cp_round_robin_split():
cur_cp_seq_len = seq_len // cp_size
assert (
@@ -179,42 +149,6 @@ def can_cp_split(seq_len: int, cp_size: int, use_nsa: bool, forward_batch):
return False
def cp_split_and_rebuild_data(forward_batch, input_: torch.Tensor):
if is_nsa_prefill_cp_round_robin_split():
cp_size = get_attention_cp_size()
assert (
input_.shape[0] % cp_size == 0
), f"Expect input shape 0 can divided by cp size, but got input shape {input_.shape}, cp size {cp_size}"
return nsa_cp_round_robin_split_data(input_)
input_list = list(
torch.split(input_, forward_batch.nsa_cp_metadata.split_list, dim=0)
)
result = torch.cat(
[input_list[i] for i in forward_batch.nsa_cp_metadata.zigzag_index], dim=0
).view(-1, input_.shape[-1])
return result
def cp_split_and_rebuild_position(forward_batch, positions: torch.Tensor):
if is_nsa_prefill_cp_round_robin_split():
cp_size = get_attention_cp_size()
assert positions.shape[0] % cp_size == 0, (
f"Expect positions shape 0 can divided by cp size, but got positions shape {positions.shape}, "
f"cp size {cp_size}"
)
return nsa_cp_round_robin_split_data(positions)
position_id_list = list(
torch.split(positions, forward_batch.nsa_cp_metadata.split_list, dim=-1)
)
positions = torch.cat(
[position_id_list[i] for i in forward_batch.nsa_cp_metadata.zigzag_index],
dim=-1,
)
return positions
@triton.jit
def nsa_cp_round_robin_split_q_seqs_kernel(
in_seqs_ptr,
@@ -285,295 +219,10 @@ def nsa_use_prefill_cp(forward_batch, nsa_enable_prefill_cp=None):
if nsa_enable_prefill_cp is None:
nsa_enable_prefill_cp = is_nsa_enable_prefill_cp()
if (
forward_batch.nsa_cp_metadata is not None
forward_batch.attn_cp_metadata is not None
and nsa_enable_prefill_cp
and forward_batch.forward_mode.is_context_parallel_extend()
):
return True
else:
return False
def cp_attn_tp_all_gather_reorganazied_into_tensor(
input_: torch.Tensor, total_len, attn_tp_size, forward_batch, stream_op
):
"""
Allgather communication for context_parallel(kv_cache, index_k, hidden_states).
This implementation mainly consists of three parts:
Step 1, padding the input shape to unify the shape for allgather communication (the shape must be the same).
Step 2, allgather communication(async).
Step 3, removing the padding and reassembling the data according to the actual tokens.
"""
# step1
max_len = (total_len + attn_tp_size - 1) // attn_tp_size
pad_size = max_len - input_.shape[0]
if pad_size > 0:
input_ = F.pad(input_, (0, 0, 0, pad_size), mode="constant", value=0)
with use_symmetric_memory(
get_attention_cp_group(), disabled=not is_allocation_symmetric()
):
input_tensor_all = torch.empty(
max_len * attn_tp_size,
input_.shape[1],
device=input_.device,
dtype=input_.dtype,
)
# step2
get_attention_cp_group().cp_all_gather_into_tensor_async(
input_tensor_all, input_, stream_op
)
# step3
outputs_list_max = list(
torch.split(input_tensor_all, forward_batch.nsa_cp_metadata.max_rank_len, dim=0)
)
outputs = torch.cat(
[
outputs_list_max[index][:per_rank_len]
for index, per_rank_len in enumerate(
forward_batch.nsa_cp_metadata.per_rank_actual_token
)
],
dim=0,
)
return outputs
def cp_all_gather_rerange_output(input_tensor, cp_size, forward_batch, stream):
"""
# for in-seq-split
| +-----------before allgather------------+|
| | dp_atten_tp0: block0, block7 |
| | dp_atten_tp1: block1, block6 |
| | dp_atten_tp2: block2, block5 |
| | dp_atten_tp3: block3, block4 |
|
| +----------before rerange---------------+|
| block0 | block7 | block1 | block6 | block2 | block5 | block3 | block4 |
|
| +--------------result-------------------+
| block0 | block1 | block2 | block3 | block4 | block5 | block6 | block7 |
| +-------------------------+
# for round-robin-split
| +-----------before allgather------------+|
| dp_atten_tp0: token0, token4, token8, token12, token16, ... |
| dp_atten_tp1: token1, token5, token9, token13, token17, ... |
| dp_atten_tp2: token2, token6, token10, token14, token18, ... |
| dp_atten_tp3: token3, token7, token11, token15, token19, ... |
|
| +--------------result-------------------+
| token0, token1, token2, token3, token4, token5, token6, token7, ...
| +-------------------------+
"""
if is_nsa_prefill_cp_round_robin_split():
with use_symmetric_memory(
get_attention_cp_group(), disabled=not is_allocation_symmetric()
):
output_tensor = input_tensor.new_empty(
(input_tensor.shape[0] * cp_size, *input_tensor.shape[1:]),
)
attn_cp_all_gather_into_tensor(
output_tensor,
input_tensor,
)
out_shape = output_tensor.shape
output_tensor = (
output_tensor.view(cp_size, -1, *out_shape[1:])
.transpose(0, 1)
.reshape(out_shape)
)
return output_tensor
bs_seq_len, hidden_size = input_tensor.shape
output_tensor = cp_attn_tp_all_gather_reorganazied_into_tensor(
input_tensor,
forward_batch.nsa_cp_metadata.total_seq_lens,
cp_size,
forward_batch,
stream,
)
outputs_list = list(
torch.split(
output_tensor, forward_batch.nsa_cp_metadata.reverse_split_len, dim=0
)
)
output_tensor = torch.cat(
[outputs_list[i] for i in forward_batch.nsa_cp_metadata.cp_reverse_index], dim=0
)
output_tensor = output_tensor.view(-1, hidden_size)
return output_tensor
def calculate_cp_seq_idx(cp_chunks_len, seqs_len):
"""Used to obtain the index of the seq corresponding
to each cp block in the forwardbatch, and the starting
and ending positions of the corresponding seq in the cp block"""
j = 0
tuple_len = [] # Only keep this result list
cumulative = {} # Used to track cumulative values for each index
for i in range(len(cp_chunks_len)):
current_dict = {}
current_tuples = []
c_val = cp_chunks_len[i]
while j < len(seqs_len):
s_val = seqs_len[j]
if s_val == c_val:
idx = j
current_dict[idx] = s_val
# Update cumulative value for this index
cumulative[idx] = cumulative.get(idx, 0) + s_val
j += 1
break
elif s_val > c_val:
idx = j
current_dict[idx] = c_val
# Update cumulative value for this index
cumulative[idx] = cumulative.get(idx, 0) + c_val
seqs_len[j] = s_val - c_val
break
else: # s_val < c_val
idx = j
current_dict[idx] = s_val
# Update cumulative value for this index
cumulative[idx] = cumulative.get(idx, 0) + s_val
c_val -= s_val
j += 1
# Build tuple: (index, historical cumulative, historical+current)
for idx, val in current_dict.items():
# Subtract current value to get historical cumulative
prev_cum = cumulative.get(idx, 0) - val
current_cum = prev_cum + val
current_tuples.append((idx, prev_cum, current_cum))
tuple_len.append(current_tuples)
return tuple_len
def prepare_input_dp_with_cp_dsa(
kv_len,
cp_rank,
cp_size,
seqs_len,
):
if is_nsa_prefill_cp_round_robin_split():
return True
"""prepare_input_dp_with_cp_dsa-zigzag index
Example (DP_ATTENT_TP == CP_SIZE == 4):
Description:
1. Start with a full-length request.
2. Split the request into multiple blocks (block0 to block7).
3. Rearrange these blocks to balance computational
load across different DP ranks.
4. Assign the rearranged blocks to different DP attention
time points (dp_atten_tp0 to dp_atten_tp3).
+---------------------------------+
| cp_split_tokens |
+---------------------------------+
| |
| request_with_full_length |
| | split (cp_size * 2) |
| +-------------------------+ |
| | block0 | block1 | block2 | block3 | block4 | block5 | block6 | block7 |
| +-------------------------+ |
| | rerange |
| +---------------------------------+
| | block0 | block7 | block1 | block6 | block2 | block5 | block3 | block4 |
| +---------------------------------+
| |
| +-------------------------+
| | dp_atten_tp0: block0, block7 |
| | dp_atten_tp1: block1, block6 |
| | dp_atten_tp2: block2, block5 |
| | dp_atten_tp3: block3, block4 |
| +-------------------------+
Why zigzag rearrange?
- Attention calculations must follow causal attention principles.
- Simply slicing by rank order can lead to computational load imbalance:
* First rank may focus on fewer historical key-value tokens (less computation)
* Last rank may focus on more tokens (more computation)
- To mitigate uneven load, the input hissenstate needs to be sliced by cp_size*2 and rearranged.
"""
# just support batch = 1
kv_len = torch.tensor(kv_len)
bs_per_cp_group = 1
kv_len_origin = kv_len
# get zigzag index
cp_segment_num = cp_size * 2
seq_per_batch = kv_len // cp_segment_num # seq_len for each batch and segment
split_list = seq_per_batch.repeat_interleave(cp_segment_num).int().tolist()
remainder = kv_len % (cp_segment_num)
if remainder > 0:
split_list[:remainder] = [x + 1 for x in split_list[:remainder]]
seq_max_rank_len = (kv_len + cp_size - 1) // cp_size
max_rank_len = seq_max_rank_len.repeat_interleave(cp_size).int().tolist()
zigzag_index = list(
range(cp_rank, cp_rank + bs_per_cp_group * cp_segment_num, cp_segment_num)
) + list(
range(
cp_segment_num - cp_rank - 1,
bs_per_cp_group * cp_segment_num,
cp_segment_num,
)
)
per_rank_actual_token = list(
split_list[i] + split_list[cp_size * 2 - i - 1] for i in range(cp_size)
)
reverse_split_len = [
element
for i in range(cp_size)
for element in (split_list[i], split_list[cp_size * 2 - i - 1])
]
# get zigzag reverse index
cp_reverse_index = []
for batch_id in range(bs_per_cp_group):
cp_reverse_index.extend(
list(range(batch_id, cp_segment_num * bs_per_cp_group, 2 * bs_per_cp_group))
+ list(
range(
(cp_segment_num - 1) * bs_per_cp_group + batch_id,
0,
-2 * bs_per_cp_group,
)
)
)
prefix_sum_list = list(accumulate(split_list))
# TODO Support multi-batch-cp-split, multi-batch-cp support has accuracy issues
# cp_seq_index = calculate_cp_seq_idx(split_list[:], seqs_len[:])
kv_len_prev = prefix_sum_list[cp_rank]
kv_len_next = prefix_sum_list[cp_size * 2 - cp_rank - 1]
actual_seq_q_prev = split_list[cp_rank]
actual_seq_q_next = split_list[cp_size * 2 - cp_rank - 1]
kv_len_prev_tensor = torch.tensor(kv_len_prev).to(device="cuda", dtype=torch.int32)
kv_len_next_tensor = torch.tensor(kv_len_next).to(device="cuda", dtype=torch.int32)
actual_seq_q_prev_tensor = torch.tensor(actual_seq_q_prev).to(
device="cuda", dtype=torch.int32
)
actual_seq_q_next_tensor = torch.tensor(actual_seq_q_next).to(
device="cuda", dtype=torch.int32
)
nsa_cp_metadata = NSAContextParallelMetadata(
split_list=split_list,
max_rank_len=max_rank_len,
zigzag_index=zigzag_index,
per_rank_actual_token=per_rank_actual_token,
reverse_split_len=reverse_split_len,
cp_reverse_index=cp_reverse_index,
kv_len_prev=kv_len_prev,
kv_len_next=kv_len_next,
actual_seq_q_prev=actual_seq_q_prev,
actual_seq_q_next=actual_seq_q_next,
kv_len_prev_tensor=kv_len_prev_tensor,
kv_len_next_tensor=kv_len_next_tensor,
actual_seq_q_prev_tensor=actual_seq_q_prev_tensor,
actual_seq_q_next_tensor=actual_seq_q_next_tensor,
total_seq_lens=kv_len_origin,
)
return nsa_cp_metadata
python/sglang/srt/layers/utils/cp_utils.py
+103 -12
View File
@@ -5,7 +5,15 @@ from typing import Callable, List
import torch
import torch.nn.functional as F
from sglang.srt.layers.dp_attention import get_attention_cp_group
from sglang.srt.distributed.device_communicators.pynccl_allocator import (
use_symmetric_memory,
)
from sglang.srt.layers.dp_attention import (
attn_cp_all_gather_into_tensor,
get_attention_cp_group,
get_attention_cp_size,
is_allocation_symmetric,
)
from sglang.srt.server_args import get_global_server_args
@@ -60,6 +68,18 @@ def can_cp_split(seq_len: int, cp_size: int, forward_batch):
def cp_split_and_rebuild_data(forward_batch, input_: torch.Tensor):
from sglang.srt.layers.attention.nsa.utils import (
is_nsa_prefill_cp_round_robin_split,
nsa_cp_round_robin_split_data,
)
if is_nsa_prefill_cp_round_robin_split():
cp_size = get_attention_cp_size()
assert (
input_.shape[0] % cp_size == 0
), f"Expect input shape 0 can divided by cp size, but got input shape {input_.shape}, cp size {cp_size}"
return nsa_cp_round_robin_split_data(input_)
input_list = list(
torch.split(input_, forward_batch.attn_cp_metadata.split_list, dim=0)
)
@@ -70,6 +90,19 @@ def cp_split_and_rebuild_data(forward_batch, input_: torch.Tensor):
def cp_split_and_rebuild_position(forward_batch, positions: torch.Tensor):
from sglang.srt.layers.attention.nsa.utils import (
is_nsa_prefill_cp_round_robin_split,
nsa_cp_round_robin_split_data,
)
if is_nsa_prefill_cp_round_robin_split():
cp_size = get_attention_cp_size()
assert positions.shape[0] % cp_size == 0, (
f"Expect positions shape 0 can divided by cp size, but got positions shape {positions.shape}, "
f"cp size {cp_size}"
)
return nsa_cp_round_robin_split_data(positions)
position_id_list = list(
torch.split(positions, forward_batch.attn_cp_metadata.split_list, dim=-1)
)
@@ -84,7 +117,11 @@ def cp_all_gather_reorganized_into_tensor(
input_tensor, total_len, cp_size, forward_batch, stream
):
"""
Allgather communication for context_parallel hidden_states.
Allgather communication for context_parallel(kv_cache, index_k, hidden_states).
This implementation mainly consists of three parts:
Step 1, padding the input shape to unify the shape for allgather communication (the shape must be the same).
Step 2, allgather communication(async).
Step 3, removing the padding and reassembling the data according to the actual tokens.
"""
# The input tensor should already be padded to the same length for allgather communication.
# No need to pad again.
@@ -95,12 +132,15 @@ def cp_all_gather_reorganized_into_tensor(
input_tensor = F.pad(
input_tensor, (0, 0, 0, pad_size), mode="constant", value=0
)
input_tensor_full = torch.empty(
max_len * cp_size,
input_tensor.shape[1],
device=input_tensor.device,
dtype=input_tensor.dtype,
)
with use_symmetric_memory(
get_attention_cp_group(), disabled=not is_allocation_symmetric()
):
input_tensor_full = torch.empty(
max_len * cp_size,
input_tensor.shape[1],
device=input_tensor.device,
dtype=input_tensor.dtype,
)
get_attention_cp_group().cp_all_gather_into_tensor_async(
input_tensor_full, input_tensor, stream
@@ -185,7 +225,40 @@ def cp_all_gather_rerange_output(input_tensor, cp_size, forward_batch, stream):
| +--------------result-------------------+
| block0 | block1 | block2 | block3 | block4 | block5 | block6 | block7 |
| +-------------------------+
# for round-robin-split
| +-----------before allgather------------+|
| dp_atten_tp0: token0, token4, token8, token12, token16, ... |
| dp_atten_tp1: token1, token5, token9, token13, token17, ... |
| dp_atten_tp2: token2, token6, token10, token14, token18, ... |
| dp_atten_tp3: token3, token7, token11, token15, token19, ... |
|
| +--------------result-------------------+
| token0, token1, token2, token3, token4, token5, token6, token7, ...
| +-------------------------+
"""
from sglang.srt.layers.attention.nsa.utils import (
is_nsa_prefill_cp_round_robin_split,
)
if is_nsa_prefill_cp_round_robin_split():
with use_symmetric_memory(
get_attention_cp_group(), disabled=not is_allocation_symmetric()
):
output_tensor = input_tensor.new_empty(
(input_tensor.shape[0] * cp_size, *input_tensor.shape[1:]),
)
attn_cp_all_gather_into_tensor(
output_tensor,
input_tensor,
)
out_shape = output_tensor.shape
output_tensor = (
output_tensor.view(cp_size, -1, *out_shape[1:])
.transpose(0, 1)
.reshape(out_shape)
)
return output_tensor
# TODO: Do we need to remove the padding here?
bs_seq_len, hidden_size = input_tensor.shape
@@ -321,6 +394,13 @@ def prepare_context_parallel_metadata(
cp_size,
seqs_len,
):
from sglang.srt.layers.attention.nsa.utils import (
is_nsa_prefill_cp_round_robin_split,
)
if is_nsa_prefill_cp_round_robin_split():
return ContextParallelMetadata()
"""prepare_input_dp_with_cp_dsa-zigzag index
Example (DP_ATTENT_TP == CP_SIZE == 4):
Description:
@@ -424,10 +504,21 @@ def prepare_context_parallel_metadata(
# TODO Support multi-batch-cp-split, multi-batch-cp support has accuracy issues
# Prefix offset is critical when radix cache hits (prefix_len > 0).
# These "cache_seqlens" values represent how many KV tokens are visible to
# each query segment during CP attention.
kv_len_prev = prefix_len + prefix_sum_list[cp_rank]
kv_len_next = prefix_len + prefix_sum_list[cp_size * 2 - cp_rank - 1]
# For non-NSA CP (e.g. qwen3-moe), consumers use these values directly as
# FlashAttention cache_seqlens, so the prefix must be baked in here.
# For NSA CP, `_get_topk_ragged_with_cp` re-adds the cached-prefix offset
# from (seq_lens_cpu - extend_seq_lens_cpu); baking prefix_len in here
# would silently drop it whenever the scheduler packs multiple requests
# into a single CP extend (len(seqs_len) != 1 -> prefix_len falls back
# to 0), corrupting the indexer's ke_offset on prefix-cache hits.
from sglang.srt.layers.attention.nsa.utils import is_nsa_enable_prefill_cp
if is_nsa_enable_prefill_cp():
kv_len_prev = prefix_sum_list[cp_rank]
kv_len_next = prefix_sum_list[cp_size * 2 - cp_rank - 1]
else:
kv_len_prev = prefix_len + prefix_sum_list[cp_rank]
kv_len_next = prefix_len + prefix_sum_list[cp_size * 2 - cp_rank - 1]
actual_seq_q_prev = split_list[cp_rank]
actual_seq_q_next = split_list[cp_size * 2 - cp_rank - 1]
# Flash Attention expects cache_seqlens to have shape (batch_size,), not scalar
python/sglang/srt/managers/schedule_batch.py
-2
View File
@@ -2151,8 +2151,6 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
# Clear context parallel metadata - CP is only for prefill, not decode
if hasattr(self, "attn_cp_metadata") and self.attn_cp_metadata is not None:
self.attn_cp_metadata = None
if hasattr(self, "nsa_cp_metadata") and self.nsa_cp_metadata is not None:
self.nsa_cp_metadata = None
if self.is_spec_v2:
# TODO(spec-v2): all spec v2 should go through this path
python/sglang/srt/model_executor/forward_batch_info.py
-3
View File
@@ -42,7 +42,6 @@ from sglang.srt.distributed.parallel_state import (
get_moe_expert_parallel_world_size,
get_tensor_model_parallel_world_size,
)
from sglang.srt.layers.attention.nsa.utils import NSAContextParallelMetadata
from sglang.srt.layers.dp_attention import (
DpPaddingMode,
get_attention_cp_size,
@@ -421,8 +420,6 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
dimensions: Optional[list[int]] = None
attn_cp_metadata: Optional[ContextParallelMetadata] = None
# Record the split metadata of the sequence number of NSA context parallels.
nsa_cp_metadata: Optional[NSAContextParallelMetadata] = None
# For hidden states before normal
return_hidden_states_before_norm: bool = False
python/sglang/srt/models/deepseek_nextn.py
+11 -7
View File
@@ -28,13 +28,9 @@ from sglang.srt.distributed import get_pp_group, get_tensor_model_parallel_world
from sglang.srt.environ import envs
from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder
from sglang.srt.layers.attention.nsa.utils import (
can_cp_split,
cp_all_gather_rerange_output,
cp_split_and_rebuild_data,
cp_split_and_rebuild_position,
can_nsa_cp_split,
is_nsa_enable_prefill_cp,
nsa_use_prefill_cp,
prepare_input_dp_with_cp_dsa,
)
from sglang.srt.layers.dp_attention import (
get_attention_cp_rank,
@@ -45,6 +41,12 @@ from sglang.srt.layers.layernorm import RMSNorm
from sglang.srt.layers.logits_processor import LogitsProcessor
from sglang.srt.layers.quantization import Fp8Config
from sglang.srt.layers.quantization.base_config import QuantizationConfig
from sglang.srt.layers.utils.cp_utils import (
cp_all_gather_rerange_output,
cp_split_and_rebuild_data,
cp_split_and_rebuild_position,
prepare_context_parallel_metadata,
)
from sglang.srt.layers.vocab_parallel_embedding import (
ParallelLMHead,
VocabParallelEmbedding,
@@ -268,8 +270,10 @@ class DeepseekV3ForCausalLMNextN(DeepseekV3ForCausalLM):
) -> torch.Tensor:
# TODO current just support prefill batch=1 and len(input_ids) > self.cp_size * 2
if self.nsa_enable_prefill_cp:
if can_cp_split(len(input_ids), self.cp_size, self.use_nsa, forward_batch):
forward_batch.nsa_cp_metadata = prepare_input_dp_with_cp_dsa(
if can_nsa_cp_split(
len(input_ids), self.cp_size, self.use_nsa, forward_batch
):
forward_batch.attn_cp_metadata = prepare_context_parallel_metadata(
len(input_ids),
self.cp_rank,
self.cp_size,
python/sglang/srt/models/deepseek_v2.py
+11 -7
View File
@@ -55,13 +55,9 @@ from sglang.srt.layers.activation import SiluAndMul
from sglang.srt.layers.amx_utils import PackWeightMethod
from sglang.srt.layers.attention.nsa.nsa_indexer import Indexer
from sglang.srt.layers.attention.nsa.utils import (
can_cp_split,
cp_all_gather_rerange_output,
cp_split_and_rebuild_data,
cp_split_and_rebuild_position,
can_nsa_cp_split,
is_nsa_enable_prefill_cp,
nsa_use_prefill_cp,
prepare_input_dp_with_cp_dsa,
)
from sglang.srt.layers.communicator import (
LayerCommunicator,
@@ -112,6 +108,12 @@ from sglang.srt.layers.quantization.fp8 import Fp8Config
from sglang.srt.layers.radix_attention import RadixAttention
from sglang.srt.layers.rotary_embedding import get_rope_wrapper
from sglang.srt.layers.utils import PPMissingLayer
from sglang.srt.layers.utils.cp_utils import (
cp_all_gather_rerange_output,
cp_split_and_rebuild_data,
cp_split_and_rebuild_position,
prepare_context_parallel_metadata,
)
from sglang.srt.layers.vocab_parallel_embedding import (
ParallelLMHead,
VocabParallelEmbedding,
@@ -2288,8 +2290,10 @@ class DeepseekV2ForCausalLM(nn.Module, DeepseekV2WeightLoaderMixin):
pp_proxy_tensors: Optional[PPProxyTensors] = None,
) -> torch.Tensor:
if self.nsa_enable_prefill_cp:
if can_cp_split(len(input_ids), self.cp_size, self.use_nsa, forward_batch):
forward_batch.nsa_cp_metadata = prepare_input_dp_with_cp_dsa(
if can_nsa_cp_split(
len(input_ids), self.cp_size, self.use_nsa, forward_batch
):
forward_batch.attn_cp_metadata = prepare_context_parallel_metadata(
len(input_ids),
self.cp_rank,
self.cp_size,