diff --git a/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py b/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py index e7acdee86e90..9d991652a075 100644 --- a/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py +++ b/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py @@ -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, diff --git a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py index 4c353f284f01..3aff6f5d776e 100644 --- a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py +++ b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py @@ -53,7 +53,6 @@ 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.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 @@ 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) @@ -380,7 +380,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, @@ -1231,17 +1231,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 @@ -1442,7 +1442,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), @@ -1455,14 +1455,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 @@ -1517,7 +1520,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( diff --git a/python/sglang/srt/layers/attention/nsa/utils.py b/python/sglang/srt/layers/attention/nsa/utils.py index 1e097983dee7..0d2c7ccdbdfb 100644 --- a/python/sglang/srt/layers/attention/nsa/utils.py +++ b/python/sglang/srt/layers/attention/nsa/utils.py @@ -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 diff --git a/python/sglang/srt/layers/utils/cp_utils.py b/python/sglang/srt/layers/utils/cp_utils.py index 698c118f9833..40aeb750b8f3 100644 --- a/python/sglang/srt/layers/utils/cp_utils.py +++ b/python/sglang/srt/layers/utils/cp_utils.py @@ -5,7 +5,15 @@ 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 diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 15995905ac3f..31d01a695184 100644 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -2151,8 +2151,6 @@ def prepare_for_decode(self): # 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 diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index 7e32a8356445..f8704040bc34 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/python/sglang/srt/model_executor/forward_batch_info.py @@ -42,7 +42,6 @@ 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 diff --git a/python/sglang/srt/models/deepseek_nextn.py b/python/sglang/srt/models/deepseek_nextn.py index f83db2d477a4..d147ad39dfa5 100644 --- a/python/sglang/srt/models/deepseek_nextn.py +++ b/python/sglang/srt/models/deepseek_nextn.py @@ -28,13 +28,9 @@ 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.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 @@ def forward( ) -> 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, diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index 7f71c837e0f6..8024ad25b6b3 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -55,13 +55,9 @@ 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.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 @@ def forward( 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,