深入 SGLang Context Parallel：源码实现详解

# python/sglang/srt/server_args.py
attn_cp_size: int = 1
moe_dp_size: int = 1
 
# Context parallelism used in the long sequence prefill phase of DeepSeek v3.2
enable_nsa_prefill_context_parallel: bool = False
nsa_prefill_cp_mode: str = "round-robin-split"
 
# Context parallelism
enable_prefill_context_parallel: bool = False
prefill_cp_mode: str = "in-seq-split"

# python/sglang/srt/server_args.py
attn_cp_size: int = 1
moe_dp_size: int = 1
 
# Context parallelism used in the long sequence prefill phase of DeepSeek v3.2
enable_nsa_prefill_context_parallel: bool = False
nsa_prefill_cp_mode: str = "round-robin-split"
 
# Context parallelism
enable_prefill_context_parallel: bool = False
prefill_cp_mode: str = "in-seq-split"

--attention-context-parallel-size / --attn-cp-size
--moe-data-parallel-size / --moe-dp-size
--enable-nsa-prefill-context-parallel
--nsa-prefill-cp-mode     # round-robin-split 或 in-seq-split
--enable-prefill-context-parallel
--prefill-cp-mode         # 当前只有 in-seq-split

--attention-context-parallel-size / --attn-cp-size
--moe-data-parallel-size / --moe-dp-size
--enable-nsa-prefill-context-parallel
--nsa-prefill-cp-mode     # round-robin-split 或 in-seq-split
--enable-prefill-context-parallel
--prefill-cp-mode         # 当前只有 in-seq-split

args.tp_size = args.tensor_parallel_size
args.pp_size = args.pipeline_parallel_size
args.attn_cp_size = args.attention_context_parallel_size
args.moe_dp_size = args.moe_data_parallel_size
args.dp_size = args.data_parallel_size
args.ep_size = args.expert_parallel_size

args.tp_size = args.tensor_parallel_size
args.pp_size = args.pipeline_parallel_size
args.attn_cp_size = args.attention_context_parallel_size
args.moe_dp_size = args.moe_data_parallel_size
args.dp_size = args.data_parallel_size
args.ep_size = args.expert_parallel_size

if self.enable_nsa_prefill_context_parallel:
    logger.warning(
        "Context parallel feature is still under experiment. It has only been verified on Hopper platform."
    )
    if self.nsa_prefill_cp_mode == "in-seq-split":
        self.enable_dp_attention = True
        self.moe_dense_tp_size = 1
        self.moe_a2a_backend = "deepep"
        self.ep_size = self.tp_size
        logger.warning(
            "For in-seq split mode, we have the following restrictions: moe_dense_tp_size == 1, moe_a2a_backend == deepep, ep_size == tp_size, batch_size == 1"
        )
    else:
        self.enable_dp_attention = True
        self.moe_dense_tp_size = 1
        assert self.dp_size == 1, "For round-robin split mode, dp attention is not supported."
    assert self.tp_size == 8, (
        "Current multi-machine CP support suffers from precision issues. "
        "So context parallel only support Single machine(tp_size == 8)"
    )
    self.attn_cp_size = self.tp_size // self.dp_size

if self.enable_nsa_prefill_context_parallel:
    logger.warning(
        "Context parallel feature is still under experiment. It has only been verified on Hopper platform."
    )
    if self.nsa_prefill_cp_mode == "in-seq-split":
        self.enable_dp_attention = True
        self.moe_dense_tp_size = 1
        self.moe_a2a_backend = "deepep"
        self.ep_size = self.tp_size
        logger.warning(
            "For in-seq split mode, we have the following restrictions: moe_dense_tp_size == 1, moe_a2a_backend == deepep, ep_size == tp_size, batch_size == 1"
        )
    else:
        self.enable_dp_attention = True
        self.moe_dense_tp_size = 1
        assert self.dp_size == 1, "For round-robin split mode, dp attention is not supported."
    assert self.tp_size == 8, (
        "Current multi-machine CP support suffers from precision issues. "
        "So context parallel only support Single machine(tp_size == 8)"
    )
    self.attn_cp_size = self.tp_size // self.dp_size

if self.attn_cp_size > 1:
    assert self.tp_size % self.attn_cp_size == 0
    assert self.tp_size % (self.dp_size * self.attn_cp_size) == 0
    assert not self.enable_aiter_allreduce_fusion
 
if self.moe_dp_size > 1:
    assert self.tp_size % self.moe_dp_size == 0
    assert self.ep_size * self.moe_dp_size <= self.tp_size
    assert self.pp_size == 1
    if self.ep_size > 1:
        assert self.ep_size * self.moe_dp_size == self.tp_size
    assert not self.enable_aiter_allreduce_fusion
 
if self.attn_cp_size != self.moe_dp_size:
    assert self.moe_dp_size == 1

if self.attn_cp_size > 1:
    assert self.tp_size % self.attn_cp_size == 0
    assert self.tp_size % (self.dp_size * self.attn_cp_size) == 0
    assert not self.enable_aiter_allreduce_fusion
 
if self.moe_dp_size > 1:
    assert self.tp_size % self.moe_dp_size == 0
    assert self.ep_size * self.moe_dp_size <= self.tp_size
    assert self.pp_size == 1
    if self.ep_size > 1:
        assert self.ep_size * self.moe_dp_size == self.tp_size
    assert not self.enable_aiter_allreduce_fusion
 
if self.attn_cp_size != self.moe_dp_size:
    assert self.moe_dp_size == 1

def _compute_parallelism_ranks(server_args: ServerArgs, tp_rank: int):
    attn_dp_size = server_args.dp_size if server_args.enable_dp_attention else 1
 
    # Parallelism hierarchy (outermost to innermost):
    # - Attention: Global(TP) -> DP -> ATTN_CP -> ATTN_TP (innermost)
    # - MoE: Global(TP) -> MOE_DP -> EP -> MOE_TP (innermost)
    attn_tp_size = server_args.tp_size // attn_dp_size // server_args.attn_cp_size
    attn_cp_rank = (tp_rank // attn_tp_size) % server_args.attn_cp_size
    moe_dp_rank = tp_rank // (server_args.tp_size // server_args.moe_dp_size)
    moe_ep_rank = (
        tp_rank
        % (server_args.tp_size // server_args.moe_dp_size)
        // (server_args.tp_size // server_args.moe_dp_size // server_args.ep_size)
    )
    return attn_cp_rank, moe_dp_rank, moe_ep_rank

def _compute_parallelism_ranks(server_args: ServerArgs, tp_rank: int):
    attn_dp_size = server_args.dp_size if server_args.enable_dp_attention else 1
 
    # Parallelism hierarchy (outermost to innermost):
    # - Attention: Global(TP) -> DP -> ATTN_CP -> ATTN_TP (innermost)
    # - MoE: Global(TP) -> MOE_DP -> EP -> MOE_TP (innermost)
    attn_tp_size = server_args.tp_size // attn_dp_size // server_args.attn_cp_size
    attn_cp_rank = (tp_rank // attn_tp_size) % server_args.attn_cp_size
    moe_dp_rank = tp_rank // (server_args.tp_size // server_args.moe_dp_size)
    moe_ep_rank = (
        tp_rank
        % (server_args.tp_size // server_args.moe_dp_size)
        // (server_args.tp_size // server_args.moe_dp_size // server_args.ep_size)
    )
    return attn_cp_rank, moe_dp_rank, moe_ep_rank

attn_dp_size = dp_size if enable_dp_attention else 1
attn_tp_size = tp_size // attn_dp_size // attn_cp_size
attn_tp_rank = tp_rank % attn_tp_size
attn_dp_rank = tp_rank // (attn_tp_size * attn_cp_size)

attn_dp_size = dp_size if enable_dp_attention else 1
attn_tp_size = tp_size // attn_dp_size // attn_cp_size
attn_tp_rank = tp_rank % attn_tp_size
attn_dp_rank = tp_rank // (attn_tp_size * attn_cp_size)

tp_rank = (attn_dp_rank * attn_cp_size + attn_cp_rank) * attn_tp_size + attn_tp_rank

tp_rank = (attn_dp_rank * attn_cp_size + attn_cp_rank) * attn_tp_size + attn_tp_rank

attn_tp_size = 8 / 1 / 2 = 4
tp_rank:       0 1 2 3 4 5 6 7
attn_tp_rank: 0 1 2 3 0 1 2 3
attn_cp_rank: 0 0 0 0 1 1 1 1
ATTN_CP 组:   [0,4], [1,5], [2,6], [3,7]
ATTN_TP 组:   [0,1,2,3], [4,5,6,7]

attn_tp_size = 8 / 1 / 2 = 4
tp_rank:       0 1 2 3 4 5 6 7
attn_tp_rank: 0 1 2 3 0 1 2 3
attn_cp_rank: 0 0 0 0 1 1 1 1
ATTN_CP 组:   [0,4], [1,5], [2,6], [3,7]
ATTN_TP 组:   [0,1,2,3], [4,5,6,7]

attn_tp_size = 8 / 2 / 4 = 1
tp_rank:       0 1 2 3 | 4 5 6 7
attn_dp_rank:  0 0 0 0 | 1 1 1 1
attn_cp_rank:  0 1 2 3 | 0 1 2 3
attn_tp_rank:  0 0 0 0 | 0 0 0 0
ATTN_CP 组:   [0,1,2,3], [4,5,6,7]
ATTN_TP 组:   每个 rank 单独成组

attn_tp_size = 8 / 2 / 4 = 1
tp_rank:       0 1 2 3 | 4 5 6 7
attn_dp_rank:  0 0 0 0 | 1 1 1 1
attn_cp_rank:  0 1 2 3 | 0 1 2 3
attn_tp_rank:  0 0 0 0 | 0 0 0 0
ATTN_CP 组:   [0,1,2,3], [4,5,6,7]
ATTN_TP 组:   每个 rank 单独成组

attn_dp_size = attention_data_parallel_size
attn_cp_size = attention_context_model_parallel_size
attn_tp_size = tensor_model_parallel_size // attn_cp_size // attn_dp_size
 
if attn_cp_size == tensor_model_parallel_size:
    _ATTN_CP = _TP
else:
    group_ranks = []
    for tp_group_idx in range(num_tensor_model_parallel_groups):
        for dp_idx in range(attn_dp_size):
            for attn_tp_idx in range(attn_tp_size):
                st = (
                    tp_group_idx * tensor_model_parallel_size
                    + dp_idx * attn_tp_size * attn_cp_size
                    + attn_tp_idx
                )
                en = (
                    tp_group_idx * tensor_model_parallel_size
                    + (dp_idx + 1) * attn_tp_size * attn_cp_size
                    + attn_tp_idx
                )
                ranks = list(range(st, en, attn_tp_size))
                group_ranks.append(ranks)
    _ATTN_CP = init_model_parallel_group(..., group_name="attn_cp")

attn_dp_size = attention_data_parallel_size
attn_cp_size = attention_context_model_parallel_size
attn_tp_size = tensor_model_parallel_size // attn_cp_size // attn_dp_size
 
if attn_cp_size == tensor_model_parallel_size:
    _ATTN_CP = _TP
else:
    group_ranks = []
    for tp_group_idx in range(num_tensor_model_parallel_groups):
        for dp_idx in range(attn_dp_size):
            for attn_tp_idx in range(attn_tp_size):
                st = (
                    tp_group_idx * tensor_model_parallel_size
                    + dp_idx * attn_tp_size * attn_cp_size
                    + attn_tp_idx
                )
                en = (
                    tp_group_idx * tensor_model_parallel_size
                    + (dp_idx + 1) * attn_tp_size * attn_cp_size
                    + attn_tp_idx
                )
                ranks = list(range(st, en, attn_tp_size))
                group_ranks.append(ranks)
    _ATTN_CP = init_model_parallel_group(..., group_name="attn_cp")

for tp_group_idx in range(num_tensor_model_parallel_groups):
    for cp_dp_combined_idx in range(attn_cp_size * attn_dp_size):
        st = tp_group_idx * tensor_model_parallel_size + cp_dp_combined_idx * attn_tp_size
        en = tp_group_idx * tensor_model_parallel_size + (cp_dp_combined_idx + 1) * attn_tp_size
        ranks = list(range(st, en))
        group_ranks.append(ranks)

for tp_group_idx in range(num_tensor_model_parallel_groups):
    for cp_dp_combined_idx in range(attn_cp_size * attn_dp_size):
        st = tp_group_idx * tensor_model_parallel_size + cp_dp_combined_idx * attn_tp_size
        en = tp_group_idx * tensor_model_parallel_size + (cp_dp_combined_idx + 1) * attn_tp_size
        ranks = list(range(st, en))
        group_ranks.append(ranks)

if attn_cp_size > moe_dp_size:
    # When moe_dp_size < attn_cp_size, CP ranks must share tokens before MoE.
    # The MOE_DP group includes these CP partners, so the existing DP
    # allgather/scatter handles the token sharing.
    _MOE_DP = _ATTN_CP

if attn_cp_size > moe_dp_size:
    # When moe_dp_size < attn_cp_size, CP ranks must share tokens before MoE.
    # The MOE_DP group includes these CP partners, so the existing DP
    # allgather/scatter handles the token sharing.
    _MOE_DP = _ATTN_CP

def get_attention_cp_group() -> GroupCoordinator:
    return get_attn_cp_group()
 
def get_attention_cp_rank() -> int:
    return get_attn_context_model_parallel_rank()
 
def get_attention_cp_size() -> int:
    return get_attn_context_model_parallel_world_size()
 
def attn_cp_all_gather_into_tensor(output: torch.Tensor, input: torch.Tensor):
    return get_attention_cp_group().all_gather_into_tensor(output, input)
 
def attn_cp_reduce_scatter_tensor(output: torch.Tensor, input: torch.Tensor):
    return get_attention_cp_group().reduce_scatter_tensor(output, input)

def get_attention_cp_group() -> GroupCoordinator:
    return get_attn_cp_group()
 
def get_attention_cp_rank() -> int:
    return get_attn_context_model_parallel_rank()
 
def get_attention_cp_size() -> int:
    return get_attn_context_model_parallel_world_size()
 
def attn_cp_all_gather_into_tensor(output: torch.Tensor, input: torch.Tensor):
    return get_attention_cp_group().all_gather_into_tensor(output, input)
 
def attn_cp_reduce_scatter_tensor(output: torch.Tensor, input: torch.Tensor):
    return get_attention_cp_group().reduce_scatter_tensor(output, input)

def cp_all_gather_into_tensor_async(self, output, input, stream):
    pynccl_comm = self.pynccl_comm
    if pynccl_comm is None or pynccl_comm.disabled:
        self.all_gather_into_tensor(output, input)
    else:
        pynccl_comm.cp_all_gather_into_tensor(output, input, stream=stream)

def cp_all_gather_into_tensor_async(self, output, input, stream):
    pynccl_comm = self.pynccl_comm
    if pynccl_comm is None or pynccl_comm.disabled:
        self.all_gather_into_tensor(output, input)
    else:
        pynccl_comm.cp_all_gather_into_tensor(output, input, stream=stream)

def cp_all_gather_into_tensor(self, output_tensor, input_tensor, stream, sizes=None):
    assert input_tensor.device == self.device
    self.nccl.ncclAllGather(
        buffer_type(input_tensor.data_ptr()),
        buffer_type(output_tensor.data_ptr()),
        input_tensor.numel(),
        ncclDataTypeEnum.from_torch(input_tensor.dtype),
        self.comm,
        cudaStream_t(stream.cuda_stream),
    )

def cp_all_gather_into_tensor(self, output_tensor, input_tensor, stream, sizes=None):
    assert input_tensor.device == self.device
    self.nccl.ncclAllGather(
        buffer_type(input_tensor.data_ptr()),
        buffer_type(output_tensor.data_ptr()),
        input_tensor.numel(),
        ncclDataTypeEnum.from_torch(input_tensor.dtype),
        self.comm,
        cudaStream_t(stream.cuda_stream),
    )

self.attn_cp_rank = attn_cp_rank
self.attn_cp_size = server_args.attn_cp_size
self.attn_tp_rank, self.attn_tp_size, self.attn_dp_rank = (
    compute_dp_attention_world_info(
        server_args.enable_dp_attention,
        self.tp_rank,
        self.tp_size,
        self.dp_size,
        self.attn_cp_size,
    )
)

self.attn_cp_rank = attn_cp_rank
self.attn_cp_size = server_args.attn_cp_size
self.attn_tp_rank, self.attn_tp_size, self.attn_dp_rank = (
    compute_dp_attention_world_info(
        server_args.enable_dp_attention,
        self.tp_rank,
        self.tp_size,
        self.dp_size,
        self.attn_cp_size,
    )
)

if self.server_args.enable_dp_attention:
    if self.attn_tp_rank == 0 and self.attn_cp_rank == 0:
        work_reqs, control_reqs = self._split_work_and_control_reqs(recv_reqs)
    else:
        work_reqs = None
        control_reqs = None
 
    if self.attn_tp_size != 1:
        work_reqs = broadcast_pyobj(..., self.attn_tp_cpu_group, src=self.attn_tp_group.ranks[0])
 
    if self.attn_cp_size != 1:
        work_reqs = broadcast_pyobj(..., self.attn_cp_cpu_group, src=self.attn_cp_group.ranks[0])

if self.server_args.enable_dp_attention:
    if self.attn_tp_rank == 0 and self.attn_cp_rank == 0:
        work_reqs, control_reqs = self._split_work_and_control_reqs(recv_reqs)
    else:
        work_reqs = None
        control_reqs = None
 
    if self.attn_tp_size != 1:
        work_reqs = broadcast_pyobj(..., self.attn_tp_cpu_group, src=self.attn_tp_group.ranks[0])
 
    if self.attn_cp_size != 1:
        work_reqs = broadcast_pyobj(..., self.attn_cp_cpu_group, src=self.attn_cp_group.ranks[0])

def is_context_parallel_extend(self, include_draft_extend_v2: bool = False):
    return (
        self == ForwardMode.EXTEND
        or self == ForwardMode.MIXED
        or (
            self == ForwardMode.DRAFT_EXTEND_V2
            if include_draft_extend_v2
            else False
        )
    )

def is_context_parallel_extend(self, include_draft_extend_v2: bool = False):
    return (
        self == ForwardMode.EXTEND
        or self == ForwardMode.MIXED
        or (
            self == ForwardMode.DRAFT_EXTEND_V2
            if include_draft_extend_v2
            else False
        )
    )

attn_cp_metadata: Optional[ContextParallelMetadata] = None

attn_cp_metadata: Optional[ContextParallelMetadata] = None

def prepare_for_decode(self):
    self.forward_mode = ForwardMode.DECODE
    ...
    # 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

def prepare_for_decode(self):
    self.forward_mode = ForwardMode.DECODE
    ...
    # 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

global_num_tokens = self.global_num_tokens_cpu
attn_tp_size = get_attention_tp_size()
 
for i in range(sync_group_size):
    global_num_tokens[i] = ceil_align(global_num_tokens[i], attn_tp_size)
 
attn_cp_size = get_attention_cp_size()
for i in range(sync_group_size):
    global_num_tokens[i] = ceil_align(global_num_tokens[i], attn_cp_size)

global_num_tokens = self.global_num_tokens_cpu
attn_tp_size = get_attention_tp_size()
 
for i in range(sync_group_size):
    global_num_tokens[i] = ceil_align(global_num_tokens[i], attn_tp_size)
 
attn_cp_size = get_attention_cp_size()
for i in range(sync_group_size):
    global_num_tokens[i] = ceil_align(global_num_tokens[i], attn_cp_size)

@dataclass
class ContextParallelMetadata:
    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
 
    # metadata for attention
    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

@dataclass
class ContextParallelMetadata:
    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
 
    # metadata for attention
    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

kv_len = torch.tensor(kv_len)
cp_segment_num = cp_size * 2
seq_per_batch = kv_len // cp_segment_num
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 = [
    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])
]

kv_len = torch.tensor(kv_len)
cp_segment_num = cp_size * 2
seq_per_batch = kv_len // cp_segment_num
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 = [
    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])
]

原始顺序: block0 block1 block2 block3 block4 block5 block6 block7
zigzag:  block0 block7 block1 block6 block2 block5 block3 block4

rank0: block0 + block7
rank1: block1 + block6
rank2: block2 + block5
rank3: block3 + block4

原始顺序: block0 block1 block2 block3 block4 block5 block6 block7
zigzag:  block0 block7 block1 block6 block2 block5 block3 block4

rank0: block0 + block7
rank1: block1 + block6
rank2: block2 + block5
rank3: block3 + block4

prefix_len = 0
try:
    if seqs_len is not None and len(seqs_len) == 1:
        prefix_len = int(seqs_len[0]) - int(kv_len_origin.item())
        if prefix_len < 0:
            prefix_len = 0
except Exception:
    prefix_len = 0

prefix_len = 0
try:
    if seqs_len is not None and len(seqs_len) == 1:
        prefix_len = int(seqs_len[0]) - int(kv_len_origin.item())
        if prefix_len < 0:
            prefix_len = 0
except Exception:
    prefix_len = 0

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]

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]

def cp_split_and_rebuild_data(forward_batch, input_: torch.Tensor):
    input_list = list(
        torch.split(input_, forward_batch.attn_cp_metadata.split_list, dim=0)
    )
    result = torch.cat(
        [input_list[i] for i in forward_batch.attn_cp_metadata.zigzag_index], dim=0
    ).view(-1, input_.shape[-1])
    return result
 
def cp_split_and_rebuild_position(forward_batch, positions: torch.Tensor):
    position_id_list = list(
        torch.split(positions, forward_batch.attn_cp_metadata.split_list, dim=-1)
    )
    positions = torch.cat(
        [position_id_list[i] for i in forward_batch.attn_cp_metadata.zigzag_index],
        dim=-1,
    )
    return positions

def cp_split_and_rebuild_data(forward_batch, input_: torch.Tensor):
    input_list = list(
        torch.split(input_, forward_batch.attn_cp_metadata.split_list, dim=0)
    )
    result = torch.cat(
        [input_list[i] for i in forward_batch.attn_cp_metadata.zigzag_index], dim=0
    ).view(-1, input_.shape[-1])
    return result
 
def cp_split_and_rebuild_position(forward_batch, positions: torch.Tensor):
    position_id_list = list(
        torch.split(positions, forward_batch.attn_cp_metadata.split_list, dim=-1)
    )
    positions = torch.cat(
        [position_id_list[i] for i in forward_batch.attn_cp_metadata.zigzag_index],
        dim=-1,
    )
    return positions

def nsa_cp_round_robin_split_data(input_):
    cp_size = get_attention_cp_size()
    cp_rank = get_attention_cp_rank()
    if isinstance(input_, (tuple, list)):
        indices = range(cp_rank, len(input_), cp_size)
        return input_[indices]
 
    tokens = len(input_)
    if tokens % cp_size != 0:
        cur_len = tokens // cp_size + (tokens % cp_size > cp_rank)
        if cur_len == 0:
            return input_.new_empty(0, *input_.shape[1:])
        indices = torch.arange(cp_rank, tokens, cp_size, device=input_.device)
        return input_[indices]
 
    return input_.view(-1, cp_size, *input_.shape[1:])[:, cp_rank].contiguous()

def nsa_cp_round_robin_split_data(input_):
    cp_size = get_attention_cp_size()
    cp_rank = get_attention_cp_rank()
    if isinstance(input_, (tuple, list)):
        indices = range(cp_rank, len(input_), cp_size)
        return input_[indices]
 
    tokens = len(input_)
    if tokens % cp_size != 0:
        cur_len = tokens // cp_size + (tokens % cp_size > cp_rank)
        if cur_len == 0:
            return input_.new_empty(0, *input_.shape[1:])
        indices = torch.arange(cp_rank, tokens, cp_size, device=input_.device)
        return input_[indices]
 
    return input_.view(-1, cp_size, *input_.shape[1:])[:, cp_rank].contiguous()

rank0: token0, token4, token8,  ...
rank1: token1, token5, token9,  ...
rank2: token2, token6, token10, ...
rank3: token3, token7, token11, ...

rank0: token0, token4, token8,  ...
rank1: token1, token5, token9,  ...
rank2: token2, token6, token10, ...
rank3: token3, token7, token11, ...

extra_seq = 0
q_seqs = []
for bs, cur_len in enumerate(extend_seqs):
    cur_len += extra_seq
    cur_seq = cur_len // cp_size + int(cur_len % cp_size > cp_rank)
    q_seqs.append(cur_seq)
    extra_seq = cur_len - cur_seq * cp_size
bs_idx = [i for i, x in enumerate(q_seqs) if x > 0]
q_seqs = [q_len for q_len in q_seqs if q_len > 0]

extra_seq = 0
q_seqs = []
for bs, cur_len in enumerate(extend_seqs):
    cur_len += extra_seq
    cur_seq = cur_len // cp_size + int(cur_len % cp_size > cp_rank)
    q_seqs.append(cur_seq)
    extra_seq = cur_len - cur_seq * cp_size
bs_idx = [i for i, x in enumerate(q_seqs) if x > 0]
q_seqs = [q_len for q_len in q_seqs if q_len > 0]

if is_prefill_context_parallel_enabled():
    if can_cp_split(len(input_ids), self.attn_cp_size, forward_batch):
        forward_batch.attn_cp_metadata = prepare_context_parallel_metadata(
            len(input_ids),
            self.attn_cp_rank,
            self.attn_cp_size,
            forward_batch.seq_lens_cpu.tolist(),
        )

if is_prefill_context_parallel_enabled():
    if can_cp_split(len(input_ids), self.attn_cp_size, forward_batch):
        forward_batch.attn_cp_metadata = prepare_context_parallel_metadata(
            len(input_ids),
            self.attn_cp_rank,
            self.attn_cp_size,
            forward_batch.seq_lens_cpu.tolist(),
        )

cur_cp_seq_len = seq_len // (cp_size * 2)
return (
    cur_cp_seq_len != 0
    and cp_size > 1
    and forward_batch.forward_mode.is_context_parallel_extend()
    and is_prefill_context_parallel_enabled()
    and forward_batch.seq_lens_cpu.shape[0] == 1
)

cur_cp_seq_len = seq_len // (cp_size * 2)
return (
    cur_cp_seq_len != 0
    and cp_size > 1
    and forward_batch.forward_mode.is_context_parallel_extend()
    and is_prefill_context_parallel_enabled()
    and forward_batch.seq_lens_cpu.shape[0] == 1
)

if (
    is_prefill_context_parallel_enabled()
    and forward_batch.forward_mode.is_context_parallel_extend()
    and forward_batch.attn_cp_metadata is not None
):
    if self.pp_group.is_first_rank:
        hidden_states = cp_split_and_rebuild_data(forward_batch, hidden_states)
    positions = cp_split_and_rebuild_position(forward_batch, positions)

if (
    is_prefill_context_parallel_enabled()
    and forward_batch.forward_mode.is_context_parallel_extend()
    and forward_batch.attn_cp_metadata is not None
):
    if self.pp_group.is_first_rank:
        hidden_states = cp_split_and_rebuild_data(forward_batch, hidden_states)
    positions = cp_split_and_rebuild_position(forward_batch, positions)

if (
    self.pp_group.is_last_rank
    and is_prefill_context_parallel_enabled()
    and forward_batch.forward_mode.is_context_parallel_extend()
    and forward_batch.attn_cp_metadata is not None
):
    hidden_states = cp_all_gather_rerange_output(
        hidden_states,
        self.attn_cp_size,
        forward_batch,
        torch.cuda.current_stream(),
    )

if (
    self.pp_group.is_last_rank
    and is_prefill_context_parallel_enabled()
    and forward_batch.forward_mode.is_context_parallel_extend()
    and forward_batch.attn_cp_metadata is not None
):
    hidden_states = cp_all_gather_rerange_output(
        hidden_states,
        self.attn_cp_size,
        forward_batch,
        torch.cuda.current_stream(),
    )

if self.nsa_enable_prefill_cp:
    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,
            forward_batch.seq_lens_cpu.tolist(),
        )
 
with get_attn_tp_context().maybe_input_scattered(forward_batch):
    hidden_states = self.model(
        input_ids, positions, forward_batch, input_embeds, pp_proxy_tensors
    )

if self.nsa_enable_prefill_cp:
    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,
            forward_batch.seq_lens_cpu.tolist(),
        )
 
with get_attn_tp_context().maybe_input_scattered(forward_batch):
    hidden_states = self.model(
        input_ids, positions, forward_batch, input_embeds, pp_proxy_tensors
    )

if nsa_use_prefill_cp(forward_batch):
    if self.pp_group.is_first_rank:
        hidden_states = cp_split_and_rebuild_data(forward_batch, hidden_states)
    positions = cp_split_and_rebuild_position(forward_batch, positions)
 
...
 
if self.pp_group.is_last_rank and nsa_use_prefill_cp(forward_batch):
    hidden_states = cp_all_gather_rerange_output(
        hidden_states,
        self.cp_size,
        forward_batch,
        torch.cuda.current_stream(),
    )

if nsa_use_prefill_cp(forward_batch):
    if self.pp_group.is_first_rank:
        hidden_states = cp_split_and_rebuild_data(forward_batch, hidden_states)
    positions = cp_split_and_rebuild_position(forward_batch, positions)
 
...
 
if self.pp_group.is_last_rank and nsa_use_prefill_cp(forward_batch):
    hidden_states = cp_all_gather_rerange_output(
        hidden_states,
        self.cp_size,
        forward_batch,
        torch.cuda.current_stream(),
    )

if self.nsa_enable_prefill_cp:
    self.layer_communicator = NSACPLayerCommunicator(
        layer_scatter_modes=self.layer_scatter_modes,
        input_layernorm=self.input_layernorm,
        post_attention_layernorm=self.post_attention_layernorm,
        allow_reduce_scatter=True,
        is_last_layer=(is_nextn or (self.layer_id == self.config.num_hidden_layers - 1)),
        qkv_latent_func=self.self_attn.prepare_qkv_latent,
    )
else:
    self.layer_communicator = LayerCommunicator(...)

if self.nsa_enable_prefill_cp:
    self.layer_communicator = NSACPLayerCommunicator(
        layer_scatter_modes=self.layer_scatter_modes,
        input_layernorm=self.input_layernorm,
        post_attention_layernorm=self.post_attention_layernorm,
        allow_reduce_scatter=True,
        is_last_layer=(is_nextn or (self.layer_id == self.config.num_hidden_layers - 1)),
        qkv_latent_func=self.self_attn.prepare_qkv_latent,
    )
else:
    self.layer_communicator = LayerCommunicator(...)

max_len = (total_len + cp_size - 1) // cp_size
pad_size = max_len - input_tensor.shape[0]
if pad_size > 0:
    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,
)
 
get_attention_cp_group().cp_all_gather_into_tensor_async(
    input_tensor_full, input_tensor, stream
)
 
outputs_list_max = list(
    torch.split(input_tensor_full, forward_batch.attn_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.attn_cp_metadata.per_rank_actual_token
        )
    ],
    dim=0,
)

max_len = (total_len + cp_size - 1) // cp_size
pad_size = max_len - input_tensor.shape[0]
if pad_size > 0:
    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,
)
 
get_attention_cp_group().cp_all_gather_into_tensor_async(
    input_tensor_full, input_tensor, stream
)
 
outputs_list_max = list(
    torch.split(input_tensor_full, forward_batch.attn_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.attn_cp_metadata.per_rank_actual_token
        )
    ],
    dim=0,
)

input_tensor: [T_rank, num_heads, head_dim]
output_tensor: [T_full, num_heads, head_dim]

input_tensor: [T_rank, num_heads, head_dim]
output_tensor: [T_full, num_heads, head_dim]

padding = [0, 0] * (input_tensor.ndim - 1) + [0, pad_size]
input_tensor = F.pad(input_tensor, padding, mode="constant", value=0)

padding = [0, 0] * (input_tensor.ndim - 1) + [0, pad_size]
input_tensor = F.pad(input_tensor, padding, mode="constant", value=0)

def cp_allgather_and_save_kv_cache(forward_batch, layer, k, v, cp_size):
    cache_loc = (
        forward_batch.out_cache_loc
        if not layer.is_cross_attention
        else forward_batch.encoder_out_cache_loc
    )
 
    k = k.contiguous()
    v = v.contiguous()
 
    key_cache_full = cp_all_gather_rerange_kv_cache(
        k, cp_size, forward_batch, torch.cuda.current_stream()
    )
    value_cache_full = cp_all_gather_rerange_kv_cache(
        v, cp_size, forward_batch, torch.cuda.current_stream()
    )
 
    forward_batch.token_to_kv_pool.set_kv_buffer(
        layer,
        cache_loc,
        key_cache_full,
        value_cache_full,
        layer.k_scale,
        layer.v_scale,
    )

def cp_allgather_and_save_kv_cache(forward_batch, layer, k, v, cp_size):
    cache_loc = (
        forward_batch.out_cache_loc
        if not layer.is_cross_attention
        else forward_batch.encoder_out_cache_loc
    )
 
    k = k.contiguous()
    v = v.contiguous()
 
    key_cache_full = cp_all_gather_rerange_kv_cache(
        k, cp_size, forward_batch, torch.cuda.current_stream()
    )
    value_cache_full = cp_all_gather_rerange_kv_cache(
        v, cp_size, forward_batch, torch.cuda.current_stream()
    )
 
    forward_batch.token_to_kv_pool.set_kv_buffer(
        layer,
        cache_loc,
        key_cache_full,
        value_cache_full,
        layer.k_scale,
        layer.v_scale,
    )

k / v local:      [T_rank, tp_k_or_v_head_num, head_dim]
key_cache_full:   [T_full, tp_k_head_num, head_dim]
value_cache_full: [T_full, tp_v_head_num, v_head_dim]
cache_loc:        out_cache_loc 对应本次 extend 的完整 token cache 位置

k / v local:      [T_rank, tp_k_or_v_head_num, head_dim]
key_cache_full:   [T_full, tp_k_head_num, head_dim]
value_cache_full: [T_full, tp_v_head_num, v_head_dim]
cache_loc:        out_cache_loc 对应本次 extend 的完整 token cache 位置

is_cp_mode = (
    forward_batch.forward_mode.is_context_parallel_extend()
    and forward_batch.attn_cp_metadata is not None
    and self.attn_cp_size > 1
)
 
if save_kv_cache and not is_cp_mode and not self.fa_skip_kv_cache:
    token_to_kv_pool.set_kv_buffer(...)
if is_cp_mode:
    cp_allgather_and_save_kv_cache(forward_batch, layer, k, v, self.attn_cp_size)

is_cp_mode = (
    forward_batch.forward_mode.is_context_parallel_extend()
    and forward_batch.attn_cp_metadata is not None
    and self.attn_cp_size > 1
)
 
if save_kv_cache and not is_cp_mode and not self.fa_skip_kv_cache:
    token_to_kv_pool.set_kv_buffer(...)
if is_cp_mode:
    cp_allgather_and_save_kv_cache(forward_batch, layer, k, v, self.attn_cp_size)

key_cache, value_cache = forward_batch.token_to_kv_pool.get_kv_buffer(layer.layer_id)
key_cache = key_cache.view(-1, self.page_size, layer.tp_k_head_num, layer.head_dim)
value_cache = value_cache.view(-1, self.page_size, layer.tp_v_head_num, layer.v_head_dim)

key_cache, value_cache = forward_batch.token_to_kv_pool.get_kv_buffer(layer.layer_id)
key_cache = key_cache.view(-1, self.page_size, layer.tp_k_head_num, layer.head_dim)
value_cache = value_cache.view(-1, self.page_size, layer.tp_v_head_num, layer.v_head_dim)

def _fa_cp_attn(q_chunk, cu_seqlens_q_cp, cache_seqlens_cp, max_seqlen_q_cp):
    return flash_attn_with_kvcache(
        q=q_chunk,
        k_cache=key_cache,
        v_cache=value_cache,
        page_table=page_table,
        cache_seqlens=cache_seqlens_cp,
        cu_seqlens_q=cu_seqlens_q_cp,
        cu_seqlens_k_new=cu_seqlens_k if not use_local_attn else None,
        max_seqlen_q=max_seqlen_q_cp,
        softmax_scale=layer.scaling,
        causal=False if use_cascade_attn else causal,
        window_size=window_size,
        softcap=layer.logit_cap,
        k_descale=k_descale,
        v_descale=v_descale,
        return_softmax_lse=use_cascade_attn,
        num_splits=self.num_splits,
        ver=self.fa_impl_ver,
        **kwargs,
    )
 
result = cp_attn_forward_extend(
    forward_batch,
    q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim),
    self.device,
    _fa_cp_attn,
)

def _fa_cp_attn(q_chunk, cu_seqlens_q_cp, cache_seqlens_cp, max_seqlen_q_cp):
    return flash_attn_with_kvcache(
        q=q_chunk,
        k_cache=key_cache,
        v_cache=value_cache,
        page_table=page_table,
        cache_seqlens=cache_seqlens_cp,
        cu_seqlens_q=cu_seqlens_q_cp,
        cu_seqlens_k_new=cu_seqlens_k if not use_local_attn else None,
        max_seqlen_q=max_seqlen_q_cp,
        softmax_scale=layer.scaling,
        causal=False if use_cascade_attn else causal,
        window_size=window_size,
        softcap=layer.logit_cap,
        k_descale=k_descale,
        v_descale=v_descale,
        return_softmax_lse=use_cascade_attn,
        num_splits=self.num_splits,
        ver=self.fa_impl_ver,
        **kwargs,
    )
 
result = cp_attn_forward_extend(
    forward_batch,
    q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim),
    self.device,
    _fa_cp_attn,
)

q_prev, q_next = torch.chunk(q, 2, dim=0)
 
cu_seqlens_q_prev = torch.tensor([0, cp_meta.actual_seq_q_prev], device=device, dtype=torch.int32)
result_prev = attn_fn(q_prev, cu_seqlens_q_prev, cp_meta.kv_len_prev_tensor, cp_meta.actual_seq_q_prev)
 
cu_seqlens_q_next = torch.tensor([0, cp_meta.actual_seq_q_next], device=device, dtype=torch.int32)
result_next = attn_fn(q_next, cu_seqlens_q_next, cp_meta.kv_len_next_tensor, cp_meta.actual_seq_q_next)
 
return torch.concat([result_prev, result_next], dim=0)

q_prev, q_next = torch.chunk(q, 2, dim=0)
 
cu_seqlens_q_prev = torch.tensor([0, cp_meta.actual_seq_q_prev], device=device, dtype=torch.int32)
result_prev = attn_fn(q_prev, cu_seqlens_q_prev, cp_meta.kv_len_prev_tensor, cp_meta.actual_seq_q_prev)
 
cu_seqlens_q_next = torch.tensor([0, cp_meta.actual_seq_q_next], device=device, dtype=torch.int32)
result_next = attn_fn(q_next, cu_seqlens_q_next, cp_meta.kv_len_next_tensor, cp_meta.actual_seq_q_next)
 
return torch.concat([result_prev, result_next], dim=0)

k_flat = k.contiguous().reshape(k.shape[0], -1)  # [S_local, k_feat]
v_flat = v.contiguous().reshape(v.shape[0], -1)  # [S_local, v_feat]
k_feat_size = k_flat.shape[-1]
kv_flat = torch.cat([k_flat, v_flat], dim=-1)    # [S_local, k_feat + v_feat]
 
kv_full = cp_all_gather_rerange_kv_cache(
    kv_flat, cp_size, forward_batch, get_current_device_stream_fast()
)  # [S_full, k_feat + v_feat]
 
key_cache_full = kv_full[..., :k_feat_size].reshape(-1, *k_tail)
value_cache_full = kv_full[..., k_feat_size:].reshape(-1, *v_tail)

k_flat = k.contiguous().reshape(k.shape[0], -1)  # [S_local, k_feat]
v_flat = v.contiguous().reshape(v.shape[0], -1)  # [S_local, v_feat]
k_feat_size = k_flat.shape[-1]
kv_flat = torch.cat([k_flat, v_flat], dim=-1)    # [S_local, k_feat + v_feat]
 
kv_full = cp_all_gather_rerange_kv_cache(
    kv_flat, cp_size, forward_batch, get_current_device_stream_fast()
)  # [S_full, k_feat + v_feat]
 
key_cache_full = kv_full[..., :k_feat_size].reshape(-1, *k_tail)
value_cache_full = kv_full[..., k_feat_size:].reshape(-1, *v_tail)

q_prev, q_next = torch.chunk(q, 2, dim=0)
q_prev = q_prev.contiguous().reshape(-1, layer.tp_q_head_num, layer.qk_head_dim)
q_next = q_next.contiguous().reshape(-1, layer.tp_q_head_num, layer.qk_head_dim)

q_prev, q_next = torch.chunk(q, 2, dim=0)
q_prev = q_prev.contiguous().reshape(-1, layer.tp_q_head_num, layer.qk_head_dim)
q_next = q_next.contiguous().reshape(-1, layer.tp_q_head_num, layer.qk_head_dim)

output_tensor = cp_all_gather_reorganized_into_tensor(
    input_tensor,
    forward_batch.attn_cp_metadata.total_seq_lens,
    cp_size,
    forward_batch,
    stream,
)
outputs_list = list(
    torch.split(
        output_tensor, forward_batch.attn_cp_metadata.reverse_split_len, dim=0
    )
)
output_tensor = torch.cat(
    [outputs_list[i] for i in forward_batch.attn_cp_metadata.cp_reverse_index],
    dim=0,
)
output_tensor = output_tensor.view(-1, hidden_size)

output_tensor = cp_all_gather_reorganized_into_tensor(
    input_tensor,
    forward_batch.attn_cp_metadata.total_seq_lens,
    cp_size,
    forward_batch,
    stream,
)
outputs_list = list(
    torch.split(
        output_tensor, forward_batch.attn_cp_metadata.reverse_split_len, dim=0
    )
)
output_tensor = torch.cat(
    [outputs_list[i] for i in forward_batch.attn_cp_metadata.cp_reverse_index],
    dim=0,
)
output_tensor = output_tensor.view(-1, hidden_size)

all-gather 后: block0 block7 block1 block6 block2 block5 block3 block4
切段后:        [0]    [7]    [1]    [6]    [2]    [5]    [3]    [4]
reverse 后:    block0 block1 block2 block3 block4 block5 block6 block7

all-gather 后: block0 block7 block1 block6 block2 block5 block3 block4
切段后:        [0]    [7]    [1]    [6]    [2]    [5]    [3]    [4]
reverse 后:    block0 block1 block2 block3 block4 block5 block6 block7

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)
)

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)
)

rank0 tokens: token0 token4 token8
rank1 tokens: token1 token5 token9
rank2 tokens: token2 token6 token10
rank3 tokens: token3 token7 token11

rank0 tokens: token0 token4 token8
rank1 tokens: token1 token5 token9
rank2 tokens: token2 token6 token10
rank3 tokens: token3 token7 token11

token0 token1 token2 token3 token4 token5 token6 token7 ...

token0 token1 token2 token3 token4 token5 token6 token7 ...

def is_nsa_enable_prefill_cp():
    return get_global_server_args().enable_nsa_prefill_context_parallel
 
def is_nsa_prefill_cp_in_seq_split():
    return (
        is_nsa_enable_prefill_cp()
        and get_global_server_args().nsa_prefill_cp_mode == "in-seq-split"
    )
 
def is_nsa_prefill_cp_round_robin_split():
    return (
        is_nsa_enable_prefill_cp()
        and get_global_server_args().nsa_prefill_cp_mode == "round-robin-split"
    )

def is_nsa_enable_prefill_cp():
    return get_global_server_args().enable_nsa_prefill_context_parallel
 
def is_nsa_prefill_cp_in_seq_split():
    return (
        is_nsa_enable_prefill_cp()
        and get_global_server_args().nsa_prefill_cp_mode == "in-seq-split"
    )
 
def is_nsa_prefill_cp_round_robin_split():
    return (
        is_nsa_enable_prefill_cp()
        and get_global_server_args().nsa_prefill_cp_mode == "round-robin-split"
    )

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()
    return (
        forward_batch.attn_cp_metadata is not None
        and nsa_enable_prefill_cp
        and forward_batch.forward_mode.is_context_parallel_extend()
    )

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()
    return (
        forward_batch.attn_cp_metadata is not None
        and nsa_enable_prefill_cp
        and forward_batch.forward_mode.is_context_parallel_extend()
    )

if is_nsa_prefill_cp_round_robin_split():
    cur_cp_seq_len = seq_len // cp_size
    assert seq_len % cp_size == 0
else:
    cur_cp_seq_len = seq_len // (cp_size * 2)
 
return (
    cur_cp_seq_len != 0
    and cp_size > 1
    and use_nsa
    and forward_batch.forward_mode.is_context_parallel_extend()
    and is_nsa_enable_prefill_cp()
    and sum(forward_batch.extend_seq_lens_cpu) >= cp_size
)

if is_nsa_prefill_cp_round_robin_split():
    cur_cp_seq_len = seq_len // cp_size
    assert seq_len % cp_size == 0
else:
    cur_cp_seq_len = seq_len // (cp_size * 2)
 
return (
    cur_cp_seq_len != 0
    and cp_size > 1
    and use_nsa
    and forward_batch.forward_mode.is_context_parallel_extend()
    and is_nsa_enable_prefill_cp()
    and sum(forward_batch.extend_seq_lens_cpu) >= cp_size
)

seqlens_expanded = torch.cat(
    [
        torch.arange(
            kv_len - qo_len + 1,
            kv_len + 1,
            dtype=torch.int32,
            device=device,
        )
        for qo_len, kv_len in zip(
            forward_batch.extend_seq_lens_cpu,
            forward_batch.seq_lens_cpu.tolist(),
            strict=True,
        )
    ]
)

seqlens_expanded = torch.cat(
    [
        torch.arange(
            kv_len - qo_len + 1,
            kv_len + 1,
            dtype=torch.int32,
            device=device,
        )
        for qo_len, kv_len in zip(
            forward_batch.extend_seq_lens_cpu,
            forward_batch.seq_lens_cpu.tolist(),
            strict=True,
        )
    ]
)

if can_nsa_prefill_cp_round_robin_split(forward_batch):
    seqlens_expanded = nsa_cp_round_robin_split_data(seqlens_expanded)
    extend_seq_lens_cpu, extend_seq_lens, bs_idx_cpu, bs_idx = (
        nsa_cp_round_robin_split_q_seqs(
            extend_seq_lens_cpu, extend_seq_lens
        )
    )
    indexer_seq_lens_cpu = indexer_seq_lens_cpu[bs_idx_cpu]
    indexer_seq_lens = indexer_seq_lens[bs_idx]
    cache_seqlens_int32 = cache_seqlens_int32[bs_idx]
    cu_seqlens_k = compute_cu_seqlens(cache_seqlens_int32)
    max_seqlen_k = (
        int(indexer_seq_lens_cpu.max().item() + draft_token_num)
        if len(indexer_seq_lens_cpu) != 0
        else 0
    )
    page_table = page_table[bs_idx, :max_seqlen_k]

if can_nsa_prefill_cp_round_robin_split(forward_batch):
    seqlens_expanded = nsa_cp_round_robin_split_data(seqlens_expanded)
    extend_seq_lens_cpu, extend_seq_lens, bs_idx_cpu, bs_idx = (
        nsa_cp_round_robin_split_q_seqs(
            extend_seq_lens_cpu, extend_seq_lens
        )
    )
    indexer_seq_lens_cpu = indexer_seq_lens_cpu[bs_idx_cpu]
    indexer_seq_lens = indexer_seq_lens[bs_idx]
    cache_seqlens_int32 = cache_seqlens_int32[bs_idx]
    cu_seqlens_k = compute_cu_seqlens(cache_seqlens_int32)
    max_seqlen_k = (
        int(indexer_seq_lens_cpu.max().item() + draft_token_num)
        if len(indexer_seq_lens_cpu) != 0
        else 0
    )
    page_table = page_table[bs_idx, :max_seqlen_k]

if bs_idx is not None:
    assert can_nsa_prefill_cp_round_robin_split(forward_batch)
    ks = nsa_cp_round_robin_split_data(ks)
    ke = nsa_cp_round_robin_split_data(ke)
    token_to_batch_idx = nsa_cp_round_robin_split_data(token_to_batch_idx)

if bs_idx is not None:
    assert can_nsa_prefill_cp_round_robin_split(forward_batch)
    ks = nsa_cp_round_robin_split_data(ks)
    ke = nsa_cp_round_robin_split_data(ke)
    token_to_batch_idx = nsa_cp_round_robin_split_data(token_to_batch_idx)

self.use_mha = (
    ...
    and (not is_nsa_enable_prefill_cp())  # CP not enabled
    and (forward_batch.hisparse_coordinator is None)
)

self.use_mha = (
    ...
    and (not is_nsa_enable_prefill_cp())  # CP not enabled
    and (forward_batch.hisparse_coordinator is None)
)

if nsa_use_prefill_cp(forward_batch):
    # support allgather+rerrange
    k_nope, k_pe = self.rebuild_cp_kv_cache(
        latent_cache, forward_batch, k_nope, k_pe
    )

if nsa_use_prefill_cp(forward_batch):
    # support allgather+rerrange
    k_nope, k_pe = self.rebuild_cp_kv_cache(
        latent_cache, forward_batch, k_nope, k_pe
    )

latent_cache[..., : self.kv_lora_rank] = k_nope.squeeze(1)
latent_cache[..., self.kv_lora_rank :] = k_pe.squeeze(1)
latent_cache_output = cp_all_gather_rerange_output(
    latent_cache.contiguous(),
    self.cp_size,
    forward_batch,
    torch.cuda.current_stream(),
)
k_nope = latent_cache_output[..., : self.kv_lora_rank].unsqueeze(1)
k_pe = latent_cache_output[..., self.kv_lora_rank :].unsqueeze(1)

latent_cache[..., : self.kv_lora_rank] = k_nope.squeeze(1)
latent_cache[..., self.kv_lora_rank :] = k_pe.squeeze(1)
latent_cache_output = cp_all_gather_rerange_output(
    latent_cache.contiguous(),
    self.cp_size,
    forward_batch,
    torch.cuda.current_stream(),
)
k_nope = latent_cache_output[..., : self.kv_lora_rank].unsqueeze(1)
k_pe = latent_cache_output[..., self.kv_lora_rank :].unsqueeze(1)

class ScatterMode(Enum):
    """
    SCATTERED: [a, b, c, d]
    TP_ATTN_FULL: [ab, ab, cd, cd]
    FULL: [abcd, abcd, abcd, abcd]
    MOE_FULL: full within the MoE group (cp_per_moe CP chunks), used when moe_dp_size < attn_cp_size
    """

class ScatterMode(Enum):
    """
    SCATTERED: [a, b, c, d]
    TP_ATTN_FULL: [ab, ab, cd, cd]
    FULL: [abcd, abcd, abcd, abcd]
    MOE_FULL: full within the MoE group (cp_per_moe CP chunks), used when moe_dp_size < attn_cp_size
    """

process_group_sizes = {
    ScatterMode.SCATTERED: 1,
    ScatterMode.TP_ATTN_FULL: attn_tp_size,
    # With context parallel enabled, we should exclude the attn_cp_size from the total tp_size
    ScatterMode.FULL: tp_size // attn_cp_size,
    ScatterMode.MOE_FULL: tp_size // (attn_cp_size // moe_cp_size),
}

process_group_sizes = {
    ScatterMode.SCATTERED: 1,
    ScatterMode.TP_ATTN_FULL: attn_tp_size,
    # With context parallel enabled, we should exclude the attn_cp_size from the total tp_size
    ScatterMode.FULL: tp_size // attn_cp_size,
    ScatterMode.MOE_FULL: tp_size // (attn_cp_size // moe_cp_size),
}

moe_cp_size = get_moe_cp_size()
if (
    moe_cp_size > 1
    and hidden_states.shape[0] > 0
    and forward_batch.forward_mode.is_context_parallel_extend()
    and forward_batch.attn_cp_metadata is not None
):
    per_rank_tokens = forward_batch.attn_cp_metadata.per_rank_actual_token
    max_tokens = max(per_rank_tokens)
    pad_size = max_tokens - hidden_states.shape[0]
    if pad_size > 0:
        hidden_states = torch.nn.functional.pad(
            hidden_states, [0, 0, 0, pad_size]
        )
 
    output = torch.empty(
        (max_tokens * moe_cp_size, hidden_states.shape[1]),
        dtype=hidden_states.dtype,
        device=hidden_states.device,
    )
    moe_cp_all_gather_into_tensor(output, hidden_states)
    hidden_states = output

moe_cp_size = get_moe_cp_size()
if (
    moe_cp_size > 1
    and hidden_states.shape[0] > 0
    and forward_batch.forward_mode.is_context_parallel_extend()
    and forward_batch.attn_cp_metadata is not None
):
    per_rank_tokens = forward_batch.attn_cp_metadata.per_rank_actual_token
    max_tokens = max(per_rank_tokens)
    pad_size = max_tokens - hidden_states.shape[0]
    if pad_size > 0:
        hidden_states = torch.nn.functional.pad(
            hidden_states, [0, 0, 0, pad_size]
        )
 
    output = torch.empty(
        (max_tokens * moe_cp_size, hidden_states.shape[1]),
        dtype=hidden_states.dtype,
        device=hidden_states.device,
    )
    moe_cp_all_gather_into_tensor(output, hidden_states)
    hidden_states = output

self._communicate_simple_fn = NSACPCommunicateSimpleFn.get_fn(
    input_mode=ScatterMode.SCATTERED,
    output_mode=ScatterMode.SCATTERED,
    context=self._context,
)
self._communicate_with_all_reduce_and_layer_norm_fn = NSACPCommunicateWithAllReduceAndLayerNormFn.get_fn(
    hidden_states_input_mode=ScatterMode.SCATTERED,
    residual_input_mode=ScatterMode.SCATTERED,
    hidden_states_output_mode=self.layer_scatter_modes.mlp_mode,
    residual_output_mode=ScatterMode.SCATTERED,
    context=self._context,
)

self._communicate_simple_fn = NSACPCommunicateSimpleFn.get_fn(
    input_mode=ScatterMode.SCATTERED,
    output_mode=ScatterMode.SCATTERED,
    context=self._context,
)
self._communicate_with_all_reduce_and_layer_norm_fn = NSACPCommunicateWithAllReduceAndLayerNormFn.get_fn(
    hidden_states_input_mode=ScatterMode.SCATTERED,
    residual_input_mode=ScatterMode.SCATTERED,
    hidden_states_output_mode=self.layer_scatter_modes.mlp_mode,
    residual_output_mode=ScatterMode.SCATTERED,
    context=self._context,
)

if nsa_use_prefill_cp(forward_batch):
    assert context.attn_dp_size == 1
    hidden_states, local_hidden_states = (
        get_local_dp_buffer(),
        hidden_states,
    )
    attn_cp_all_gather_into_tensor(
        hidden_states,
        local_hidden_states,
    )

if nsa_use_prefill_cp(forward_batch):
    assert context.attn_dp_size == 1
    hidden_states, local_hidden_states = (
        get_local_dp_buffer(),
        hidden_states,
    )
    attn_cp_all_gather_into_tensor(
        hidden_states,
        local_hidden_states,
    )

if nsa_use_prefill_cp(forward_batch):
    assert context.attn_dp_size == 1
    input_hidden_states = hidden_states
    hidden_states = hidden_states.tensor_split(context.attn_cp_size)[
        context.attn_cp_rank
    ]
    attn_cp_reduce_scatter_tensor(hidden_states, input_hidden_states)

if nsa_use_prefill_cp(forward_batch):
    assert context.attn_dp_size == 1
    input_hidden_states = hidden_states
    hidden_states = hidden_states.tensor_split(context.attn_cp_size)[
        context.attn_cp_rank
    ]
    attn_cp_reduce_scatter_tensor(hidden_states, input_hidden_states)

self.attn_tp_size = get_attention_tp_size()
self.attn_tp_rank = get_attention_tp_rank()
self.attn_cp_size = get_attention_cp_size()
self.attn_cp_rank = get_attention_cp_rank()
self.attn_dp_size = get_attention_dp_size()
self.attn_dp_rank = get_attention_dp_rank()

self.attn_tp_size = get_attention_tp_size()
self.attn_tp_rank = get_attention_tp_rank()
self.attn_cp_size = get_attention_cp_size()
self.attn_cp_rank = get_attention_cp_rank()
self.attn_dp_size = get_attention_dp_size()
self.attn_dp_rank = get_attention_dp_rank()

payload = {
    "attn_tp_size": self.attn_tp_size,
    "attn_tp_rank": self.attn_tp_rank,
    "attn_cp_size": self.attn_cp_size,
    "attn_cp_rank": self.attn_cp_rank,
    "attn_dp_size": self.attn_dp_size,
    "attn_dp_rank": self.attn_dp_rank,
    ...
}

payload = {
    "attn_tp_size": self.attn_tp_size,
    "attn_tp_rank": self.attn_tp_rank,
    "attn_cp_size": self.attn_cp_size,
    "attn_cp_rank": self.attn_cp_rank,
    "attn_dp_size": self.attn_dp_size,
    "attn_dp_rank": self.attn_dp_rank,
    ...
}

dp_group_table = self.prefill_port_table.setdefault(dp_group, {})
cp_group_table = dp_group_table.setdefault(attn_cp_rank, {})
tp_group_table = cp_group_table.setdefault(attn_tp_rank, {})
tp_group_table[pp_rank] = PrefillRankInfo(...)

dp_group_table = self.prefill_port_table.setdefault(dp_group, {})
cp_group_table = dp_group_table.setdefault(attn_cp_rank, {})
tp_group_table = cp_group_table.setdefault(attn_tp_rank, {})
tp_group_table[pp_rank] = PrefillRankInfo(...)

assert self.attn_cp_size == 1, (
    f"Decode cp size ({self.attn_cp_size}) should be equal to 1",
)

assert self.attn_cp_size == 1, (
    f"Decode cp size ({self.attn_cp_size}) should be equal to 1",
)

target_cp_ranks = list(range(info.attn_cp_size))
if not self.enable_all_cp_ranks_for_transfer:
    # Only retrieve from prefill CP rank 0 when not using all ranks
    target_cp_ranks = target_cp_ranks[:1]
else:
    required_prefill_response_num *= info.attn_cp_size // self.attn_cp_size

target_cp_ranks = list(range(info.attn_cp_size))
if not self.enable_all_cp_ranks_for_transfer:
    # Only retrieve from prefill CP rank 0 when not using all ranks
    target_cp_ranks = target_cp_ranks[:1]
else:
    required_prefill_response_num *= info.attn_cp_size // self.attn_cp_size

if self.kv_mgr.enable_all_cp_ranks_for_transfer:
    kv_indices, index_slice = filter_kv_indices_for_cp_rank(
        self.kv_mgr,
        kv_indices,
        index_slice,
    )
elif self.kv_mgr.is_dummy_cp_rank:
    if not is_last_chunk:
        return
    else:
        self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Success)
        return

if self.kv_mgr.enable_all_cp_ranks_for_transfer:
    kv_indices, index_slice = filter_kv_indices_for_cp_rank(
        self.kv_mgr,
        kv_indices,
        index_slice,
    )
elif self.kv_mgr.is_dummy_cp_rank:
    if not is_last_chunk:
        return
    else:
        self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Success)
        return

base = total_pages // cp_size
rem = total_pages % cp_size
 
if rem == 0:
    local_start = cp_rank * base
    local_end = local_start + base
else:
    local_start = cp_rank * base + min(cp_rank, rem)
    n_pages = base + (1 if cp_rank < rem else 0)
    local_end = local_start + n_pages
 
start_page = first_page + local_start
end_page = first_page + local_end
mask = (page_indices >= start_page) & (page_indices < end_page)

base = total_pages // cp_size
rem = total_pages % cp_size
 
if rem == 0:
    local_start = cp_rank * base
    local_end = local_start + base
else:
    local_start = cp_rank * base + min(cp_rank, rem)
    n_pages = base + (1 if cp_rank < rem else 0)
    local_end = local_start + n_pages
 
start_page = first_page + local_start
end_page = first_page + local_end
mask = (page_indices >= start_page) & (page_indices < end_page)

polls = poll_and_all_reduce_attn_cp_tp_group(
    [req.disagg_kv_sender for req in self.disagg_prefill_inflight_queue],
    self.attn_cp_cpu_group,
    self.attn_tp_cpu_group,
)

polls = poll_and_all_reduce_attn_cp_tp_group(
    [req.disagg_kv_sender for req in self.disagg_prefill_inflight_queue],
    self.attn_cp_cpu_group,
    self.attn_tp_cpu_group,
)

polls = poll_and_all_reduce(pollers, attn_tp_cpu_group)
tensor_to_reduce = torch.tensor(polls, dtype=torch.uint8, device="cpu")
dist.all_reduce(tensor_to_reduce, op=dist.ReduceOp.MIN, group=attn_cp_cpu_group)

polls = poll_and_all_reduce(pollers, attn_tp_cpu_group)
tensor_to_reduce = torch.tensor(polls, dtype=torch.uint8, device="cpu")
dist.all_reduce(tensor_to_reduce, op=dist.ReduceOp.MIN, group=attn_cp_cpu_group)

params = CacheInitParams(
    ...
    tp_cache_group=(
        self.attn_tp_cpu_group
        if self.server_args.enable_dp_attention
        else self.tp_cpu_group
    ),
    attn_cp_cache_group=self.attn_cp_cpu_group,
    attn_tp_cache_group=self.attn_tp_cpu_group,
    ...
)

params = CacheInitParams(
    ...
    tp_cache_group=(
        self.attn_tp_cpu_group
        if self.server_args.enable_dp_attention
        else self.tp_cpu_group
    ),
    attn_cp_cache_group=self.attn_cp_cpu_group,
    attn_tp_cache_group=self.attn_tp_cpu_group,
    ...
)

if self.attn_cp_group is not None:
    return (
        torch.distributed.get_rank(group=self.attn_cp_group),
        torch.distributed.get_world_size(group=self.attn_cp_group),
    )
return 0, 1

if self.attn_cp_group is not None:
    return (
        torch.distributed.get_rank(group=self.attn_cp_group),
        torch.distributed.get_world_size(group=self.attn_cp_group),
    )
return 0, 1

return HiCacheStorageConfig(
    tp_rank=self.tp_rank,
    tp_size=self.tp_size,
    pp_rank=self.pp_rank,
    pp_size=self.pp_size,
    attn_cp_rank=attn_cp_rank,
    attn_cp_size=attn_cp_size,
    ...
)

return HiCacheStorageConfig(
    tp_rank=self.tp_rank,
    tp_size=self.tp_size,
    pp_rank=self.pp_rank,
    pp_size=self.pp_size,
    attn_cp_rank=attn_cp_rank,
    attn_cp_size=attn_cp_size,
    ...
)

if self.attn_cp_size > 1:
    self.disable_piecewise_cuda_graph = True

if self.attn_cp_size > 1:
    self.disable_piecewise_cuda_graph = True

if require_gathered_buffer(self.model_runner.server_args):
    mul_base = self.attn_tp_size
    attn_cp_size = get_attention_cp_size()
    if mul_base % attn_cp_size != 0:
        mul_base *= attn_cp_size
    filtered = [n for n in self.capture_num_tokens if n % mul_base == 0]

if require_gathered_buffer(self.model_runner.server_args):
    mul_base = self.attn_tp_size
    attn_cp_size = get_attention_cp_size()
    if mul_base % attn_cp_size != 0:
        mul_base *= attn_cp_size
    filtered = [n for n in self.capture_num_tokens if n % mul_base == 0]

return (
    _is_cuda
    and ...
    and not is_prefill_context_parallel_enabled()
) or (_is_hip and not is_prefill_context_parallel_enabled())

return (
    _is_cuda
    and ...
    and not is_prefill_context_parallel_enabled()
) or (_is_hip and not is_prefill_context_parallel_enabled())

# python/sglang/srt/entrypoints/engine.py
attn_dp_size = server_args.dp_size if server_args.enable_dp_attention else 1
 
# Parallelism hierarchy (outermost to innermost):
# - Attention: Global(TP) -> DP -> ATTN_CP -> ATTN_TP (innermost)
# - MoE: Global(TP) -> MOE_DP -> EP -> MOE_TP (innermost)
attn_tp_size = server_args.tp_size // attn_dp_size // server_args.attn_cp_size
attn_cp_rank = (tp_rank // attn_tp_size) % server_args.attn_cp_size
moe_dp_rank = tp_rank // (server_args.tp_size // server_args.moe_dp_size)

# python/sglang/srt/entrypoints/engine.py
attn_dp_size = server_args.dp_size if server_args.enable_dp_attention else 1
 
# Parallelism hierarchy (outermost to innermost):
# - Attention: Global(TP) -> DP -> ATTN_CP -> ATTN_TP (innermost)
# - MoE: Global(TP) -> MOE_DP -> EP -> MOE_TP (innermost)
attn_tp_size = server_args.tp_size // attn_dp_size // server_args.attn_cp_size
attn_cp_rank = (tp_rank // attn_tp_size) % server_args.attn_cp_size
moe_dp_rank = tp_rank // (server_args.tp_size // server_args.moe_dp_size)

# python/sglang/srt/layers/dp_attention.py
attn_dp_size = dp_size if enable_dp_attention else 1
attn_tp_size = tp_size // attn_dp_size // attn_cp_size
attn_tp_rank = tp_rank % attn_tp_size
 
if not enable_dp_attention:
    attn_dp_rank = 0
else:
    # Rank layout is (dp, cp, tp) where tp is the fastest-changing dim:
    # tp_rank = (attn_dp_rank * attn_cp_size + attn_cp_rank) * attn_tp_size + attn_tp_rank
    attn_dp_rank = tp_rank // (attn_tp_size * attn_cp_size)

# python/sglang/srt/layers/dp_attention.py
attn_dp_size = dp_size if enable_dp_attention else 1
attn_tp_size = tp_size // attn_dp_size // attn_cp_size
attn_tp_rank = tp_rank % attn_tp_size
 
if not enable_dp_attention:
    attn_dp_rank = 0
else:
    # Rank layout is (dp, cp, tp) where tp is the fastest-changing dim:
    # tp_rank = (attn_dp_rank * attn_cp_size + attn_cp_rank) * attn_tp_size + attn_tp_rank
    attn_dp_rank = tp_rank // (attn_tp_size * attn_cp_size)