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

refactor context parallel state (#17213)

Browse Source 
Co-authored-by: Shunkang <182541032+Shunkangz@users.noreply.github.co>
This commit is contained in:
dongjiyingdjy
2026-02-13 23:18:17 +08:00
committed by GitHub
co-authored by Shunkang
parent 0012d6a4eb
commit 8b4c364960
27 changed files with 847 additions and 118 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/distributed/parallel_state.py
+219 -16
View File
@@ -1351,6 +1351,7 @@ def init_model_parallel_group(
group_ranks: List[List[int]],
local_rank: int,
backend: str,
use_pynccl: Optional[bool] = None,
use_custom_allreduce: Optional[bool] = None,
use_message_queue_broadcaster: bool = False,
group_name: Optional[str] = None,
@@ -1368,7 +1369,11 @@ def init_model_parallel_group(
group_ranks=group_ranks,
local_rank=local_rank,
torch_distributed_backend=backend,
use_pynccl=not (_is_npu or _is_xpu or backend == "mooncake"),
use_pynccl=(
not (_is_npu or _is_xpu or backend == "mooncake")
if use_pynccl is None
else use_pynccl
),
use_pymscclpp=use_mscclpp_allreduce,
use_custom_allreduce=use_custom_allreduce,
use_torch_symm_mem_all_reduce=use_torch_symm_mem_allreduce,
@@ -1382,6 +1387,8 @@ def init_model_parallel_group(
_TP: Optional[GroupCoordinator] = None
_ATTN_TP: Optional[GroupCoordinator] = None
_ATTN_CP: Optional[GroupCoordinator] = None
# duplicate GroupCoordinator for prefill in PD-Multiplexing
_PDMUX_PREFILL_TP_GROUP: Optional[GroupCoordinator] = None
@@ -1404,10 +1411,30 @@ def get_tp_group() -> GroupCoordinator:
return _TP
def get_attn_tp_group() -> GroupCoordinator:
assert (
_ATTN_TP is not None
), "attention tensor model parallel group is not initialized"
return _ATTN_TP
def get_attn_cp_group() -> GroupCoordinator:
assert (
_ATTN_CP is not None
), "attention context model parallel group is not initialized"
return _ATTN_CP
_MOE_DP: Optional[GroupCoordinator] = None
_MOE_EP: Optional[GroupCoordinator] = None
_MOE_TP: Optional[GroupCoordinator] = None
def get_moe_dp_group() -> GroupCoordinator:
assert _MOE_DP is not None, "moe data parallel group is not initialized"
return _MOE_DP
def get_moe_ep_group() -> GroupCoordinator:
assert _MOE_EP is not None, "expert model parallel group is not initialized"
return _MOE_EP
@@ -1558,6 +1585,9 @@ def initialize_model_parallel(
tensor_model_parallel_size: int = 1,
expert_model_parallel_size: int = 1,
pipeline_model_parallel_size: int = 1,
attention_data_parallel_size: int = 1,
attention_context_model_parallel_size: int = 1,
moe_data_model_parallel_size: int = 1,
backend: Optional[str] = None,
duplicate_tp_group: bool = False,
) -> None:
@@ -1567,8 +1597,16 @@ def initialize_model_parallel(
Arguments:
tensor_model_parallel_size: number of GPUs used for tensor model
parallelism.
expert_model_parallel_size: number of GPUs used for expert model
parallelism.
pipeline_model_parallel_size: number of GPUs used for pipeline model
parallelism.
attention_data_parallel_size: number of GPUs used for attention data
parallelism.
attention_context_model_parallel_size: number of GPUs used for attention context
parallelism.
moe_data_model_parallel_size: number of GPUs used for moe data
parallelism.
Let's say we have a total of 8 GPUs denoted by g0 ... g7 and we
use 2 GPUs to parallelize the model tensor, and 4 GPUs to parallelize
@@ -1578,6 +1616,20 @@ def initialize_model_parallel(
[g0, g1], [g2, g3], [g4, g5], [g6, g7]
2 pipeline model-parallel groups:
[g0, g2, g4, g6], [g1, g3, g5, g7]
Let's say we use 2 GPUs for attention context parallelism (attn_cp_size=2) and 4 GPUs for
attention tensor parallelism (attn_tp_size=4). As for MoE part, we use 2 GPUs for moe data
parallelism (moe_dp_size=2) and 4 GPUs for moe expert parallelism (moe_ep_size=4). The present
function will create the following groups:
2 tensor model-parallel groups:
[g0, g1, g2, g3], [g4, g5, g6, g7]
4 attention context-parallel groups:
[g0, g4], [g1, g5], [g2, g6], [g3, g7]
2 moe expert-parallel groups:
[g0, g1, g2, g3], [g4, g5, g6, g7]
4 moe data-parallel groups:
[g0, g4], [g1, g5], [g2, g6], [g3, g7]
Note that for efficiency, the caller should make sure adjacent ranks
are on the same DGX box. For example if we are using 2 DGX-1 boxes
with a total of 16 GPUs, rank 0 to 7 belong to the first box and
@@ -1600,9 +1652,12 @@ def initialize_model_parallel(
global _TP
assert _TP is None, "tensor model parallel group is already initialized"
group_ranks = []
for i in range(num_tensor_model_parallel_groups):
for tp_group_idx in range(num_tensor_model_parallel_groups):
ranks = list(
range(i * tensor_model_parallel_size, (i + 1) * tensor_model_parallel_size)
range(
tp_group_idx * tensor_model_parallel_size,
(tp_group_idx + 1) * tensor_model_parallel_size,
)
)
group_ranks.append(ranks)
@@ -1637,8 +1692,98 @@ def initialize_model_parallel(
_TP.pynccl_comm.disabled = False
_PDMUX_PREFILL_TP_GROUP.pynccl_comm.disabled = False
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
global _ATTN_CP
assert (
_ATTN_CP is None
), "attention context model parallel group is already initialized"
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_ranks,
get_world_group().local_rank,
backend,
group_name="attn_cp",
)
from sglang.srt.layers.sampler import SYNC_TOKEN_IDS_ACROSS_TP
global _ATTN_TP
assert (
_ATTN_TP is None
), "attention tensor model parallel group is already initialized"
if attn_tp_size == tensor_model_parallel_size:
_ATTN_TP = _TP
else:
group_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)
_ATTN_TP = init_model_parallel_group(
group_ranks,
get_world_group().local_rank,
backend,
use_pynccl=SYNC_TOKEN_IDS_ACROSS_TP,
use_mscclpp_allreduce=False,
use_custom_allreduce=False,
use_torch_symm_mem_allreduce=False,
group_name="attention_tp",
)
moe_ep_size = expert_model_parallel_size
moe_tp_size = tensor_model_parallel_size // moe_ep_size
moe_dp_size = moe_data_model_parallel_size
moe_tp_size = tensor_model_parallel_size // moe_ep_size // moe_dp_size
global _MOE_DP
assert _MOE_DP is None, "moe data parallel group is already initialized"
# gpus_per_pp_stage = tensor_model_parallel_size * attention_context_model_parallel_size
if moe_dp_size == tensor_model_parallel_size:
_MOE_DP = _TP
else:
group_ranks = []
for tp_group_idx in range(num_tensor_model_parallel_groups):
for tp_ep_combined_idx in range(moe_tp_size * moe_ep_size):
st = tp_group_idx * tensor_model_parallel_size + tp_ep_combined_idx
en = (
tp_group_idx + 1
) * tensor_model_parallel_size + tp_ep_combined_idx
ranks = list(range(st, en, moe_tp_size * moe_ep_size))
group_ranks.append(ranks)
_MOE_DP = init_model_parallel_group(
group_ranks,
get_world_group().local_rank,
backend,
group_name="moe_dp",
)
global _MOE_EP
assert _MOE_EP is None, "expert model parallel group is already initialized"
@@ -1647,12 +1792,17 @@ def initialize_model_parallel(
else:
# TODO(ch-wan): use split_group to save memory
group_ranks = []
for i in range(num_tensor_model_parallel_groups):
for j in range(moe_tp_size):
st = i * tensor_model_parallel_size + j
en = (i + 1) * tensor_model_parallel_size + j
ranks = list(range(st, en, moe_tp_size))
group_ranks.append(ranks)
for tp_group_idx in range(num_tensor_model_parallel_groups):
for moe_dp_idx in range(moe_dp_size):
for moe_tp_idx in range(moe_tp_size):
st = (
tp_group_idx * tensor_model_parallel_size
+ moe_dp_idx * moe_ep_size * moe_tp_size
+ moe_tp_idx
)
en = st + moe_ep_size * moe_tp_size
ranks = list(range(st, en, moe_tp_size))
group_ranks.append(ranks)
_MOE_EP = init_model_parallel_group(
group_ranks,
get_world_group().local_rank,
@@ -1667,10 +1817,16 @@ def initialize_model_parallel(
else:
# TODO(ch-wan): use split_group to save memory
group_ranks = []
for i in range(num_tensor_model_parallel_groups):
for j in range(moe_ep_size):
st = i * tensor_model_parallel_size + j * moe_tp_size
en = i * tensor_model_parallel_size + (j + 1) * moe_tp_size
for tp_group_idx in range(num_tensor_model_parallel_groups):
for ep_dp_combined_idx in range(moe_ep_size * moe_dp_size):
st = (
tp_group_idx * tensor_model_parallel_size
+ ep_dp_combined_idx * moe_tp_size
)
en = (
tp_group_idx * tensor_model_parallel_size
+ (ep_dp_combined_idx + 1) * moe_tp_size
)
ranks = list(range(st, en))
group_ranks.append(ranks)
_MOE_TP = init_model_parallel_group(
@@ -1685,8 +1841,10 @@ def initialize_model_parallel(
global _PP
assert _PP is None, "pipeline model parallel group is already initialized"
group_ranks = []
for i in range(num_pipeline_model_parallel_groups):
ranks = list(range(i, world_size, num_pipeline_model_parallel_groups))
for pp_group_idx in range(num_pipeline_model_parallel_groups):
ranks = list(
range(pp_group_idx, world_size, num_pipeline_model_parallel_groups)
)
group_ranks.append(ranks)
# pipeline parallel does not need custom allreduce
_PP = init_model_parallel_group(
@@ -1833,6 +1991,28 @@ def get_tensor_model_parallel_rank():
return get_tp_group().rank_in_group
# ATTN_TP
def get_attn_tensor_model_parallel_world_size():
"""Return world size for the attention tensor model parallel group."""
return get_attn_tp_group().world_size
def get_attn_tensor_model_parallel_rank():
"""Return my rank for the attention tensor model parallel group."""
return get_attn_tp_group().rank_in_group
# ATTN_CP
def get_attn_context_model_parallel_world_size():
"""Return world size for the attention context model parallel group."""
return get_attn_cp_group().world_size
def get_attn_context_model_parallel_rank():
"""Return my rank for the attention context model parallel group."""
return get_attn_cp_group().rank_in_group
def get_pipeline_model_parallel_world_size():
"""Return world size for the pipeline model parallel group."""
return get_pp_group().world_size
@@ -1843,6 +2023,18 @@ def get_pipeline_model_parallel_rank():
return get_pp_group().rank_in_group
# MOE_DP
def get_moe_data_parallel_world_size():
"""Return world size for the moe data parallel group."""
return get_moe_dp_group().world_size
def get_moe_data_parallel_rank():
"""Return my rank for the moe data parallel group."""
return get_moe_dp_group().rank_in_group
# MOE_EP
def get_moe_expert_parallel_world_size():
"""Return world size for the moe expert parallel group."""
return get_moe_ep_group().world_size
@@ -1853,6 +2045,7 @@ def get_moe_expert_parallel_rank():
return get_moe_ep_group().rank_in_group
# MOE_TP
def get_moe_tensor_parallel_world_size():
"""Return world size for the moe tensor parallel group."""
return get_moe_tp_group().world_size
@@ -1885,6 +2078,16 @@ def destroy_model_parallel():
_MOE_TP.destroy()
_MOE_TP = None
global _ATTN_CP
if _ATTN_CP:
_ATTN_CP.destroy()
_ATTN_CP = None
global _MOE_DP
if _MOE_DP:
_MOE_DP.destroy()
_MOE_DP = None
global _PDMUX_PREFILL_TP_GROUP
if _PDMUX_PREFILL_TP_GROUP: # type: ignore[union-attr]
_PDMUX_PREFILL_TP_GROUP.destroy()
python/sglang/srt/entrypoints/engine.py
+25 -1
View File
@@ -938,7 +938,29 @@ def _launch_scheduler_processes(
+ ((pp_rank % pp_size_per_node) * tp_size_per_node)
+ (tp_rank % tp_size_per_node) * server_args.gpu_id_step
)
moe_ep_rank = tp_rank // (server_args.tp_size // server_args.ep_size)
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
)
)
with maybe_reindex_device_id(gpu_id) as gpu_id:
proc = mp.Process(
@@ -948,6 +970,8 @@ def _launch_scheduler_processes(
port_args,
gpu_id,
tp_rank,
attn_cp_rank,
moe_dp_rank,
moe_ep_rank,
pp_rank,
None,
python/sglang/srt/layers/attention/nsa/nsa_indexer.py
+6 -3
View File
@@ -28,6 +28,10 @@ if is_npu():
import torch_npu
from sglang.srt.hardware_backend.npu.utils import get_indexer_weight_stream
from sglang.srt.distributed import (
get_attn_context_model_parallel_rank,
get_attn_context_model_parallel_world_size,
)
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 (
@@ -35,7 +39,6 @@ from sglang.srt.layers.attention.nsa.utils import (
is_nsa_enable_prefill_cp,
is_nsa_prefill_cp_in_seq_split,
)
from sglang.srt.layers.dp_attention import get_attention_tp_rank, get_attention_tp_size
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
@@ -162,8 +165,8 @@ class Indexer(MultiPlatformOp):
self.alt_stream = alt_stream
self.nsa_enable_prefill_cp = is_nsa_enable_prefill_cp()
if self.nsa_enable_prefill_cp:
self.cp_size = get_attention_tp_size()
self.cp_rank = get_attention_tp_rank()
self.cp_size = get_attn_context_model_parallel_world_size()
self.cp_rank = get_attn_context_model_parallel_rank()
else:
self.cp_size = None
self.cp_rank = None
python/sglang/srt/layers/attention/nsa/utils.py
+22 -21
View File
@@ -13,11 +13,11 @@ from sglang.srt.distributed.device_communicators.pynccl_allocator import (
)
from sglang.srt.layers.dp_attention import (
DpPaddingMode,
attn_tp_all_gather_into_tensor,
attn_cp_all_gather_into_tensor,
get_attention_cp_group,
get_attention_cp_rank,
get_attention_cp_size,
get_attention_dp_rank,
get_attention_tp_group,
get_attention_tp_rank,
get_attention_tp_size,
is_allocation_symmetric,
)
from sglang.srt.server_args import get_global_server_args
@@ -52,7 +52,7 @@ def is_nsa_prefill_cp_round_robin_split():
def can_nsa_prefill_cp_round_robin_split(forward_batch: "ForwardBatch"):
if not forward_batch.forward_mode.is_context_parallel_extend():
return False
cp_size = get_attention_tp_size()
cp_size = get_attention_cp_size()
seq_len = sum(forward_batch.extend_seq_lens_cpu)
return is_nsa_prefill_cp_round_robin_split() and seq_len > 0 and cp_size > 1
@@ -70,8 +70,8 @@ def nsa_cp_round_robin_split_data(input_: Union[torch.Tensor, List]):
| dp_atten_tp3: token3, token7, token11, token15, token19, ... |
| +-------------------------+
"""
cp_size = get_attention_tp_size()
cp_rank = get_attention_tp_rank()
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]
@@ -93,9 +93,9 @@ def cal_padded_tokens(forward_batch: "ForwardBatch"):
# calculate the actual token length after padding when attn_tp_size > 1 or in the MAX_LEN padding mode.
global_num_tokens = forward_batch.global_num_tokens_cpu.copy()
sync_group_size = len(global_num_tokens)
attn_tp_size = get_attention_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_tp_size)
global_num_tokens[i] = ceil_align(global_num_tokens[i], attn_cp_size)
dp_padding_mode = DpPaddingMode.get_dp_padding_mode(
forward_batch.is_extend_in_batch, global_num_tokens
)
@@ -106,12 +106,13 @@ def cal_padded_tokens(forward_batch: "ForwardBatch"):
else:
tokens = global_num_tokens[0]
if can_nsa_prefill_cp_round_robin_split(forward_batch):
tokens = ceil_div(tokens, attn_tp_size)
tokens = ceil_div(tokens, attn_cp_size)
return tokens
def pad_nsa_cache_seqlens(forward_batch: "ForwardBatch", nsa_cache_seqlens):
if forward_batch.global_num_tokens_cpu is None:
attn_cp_size = get_attention_cp_size()
if attn_cp_size == 1 or not can_nsa_prefill_cp_round_robin_split(forward_batch):
return nsa_cache_seqlens
tokens = cal_padded_tokens(forward_batch)
pad_len = tokens - nsa_cache_seqlens.shape[0]
@@ -170,7 +171,7 @@ def can_cp_split(seq_len: int, cp_size: int, use_nsa: bool, forward_batch):
def cp_split_and_rebuild_data(forward_batch, input_: torch.Tensor):
if is_nsa_prefill_cp_round_robin_split():
cp_size = get_attention_tp_size()
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}"
@@ -187,7 +188,7 @@ def cp_split_and_rebuild_data(forward_batch, input_: torch.Tensor):
def cp_split_and_rebuild_position(forward_batch, positions: torch.Tensor):
if is_nsa_prefill_cp_round_robin_split():
cp_size = get_attention_tp_size()
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}"
@@ -227,8 +228,8 @@ def nsa_cp_round_robin_split_q_seqs_kernel(
def nsa_cp_round_robin_split_q_seqs_cpu(extend_seqs):
cp_size = get_attention_tp_size()
cp_rank = get_attention_tp_rank()
cp_size = get_attention_cp_size()
cp_rank = get_attention_cp_rank()
extra_seq = 0
q_seqs = []
for bs, cur_len in enumerate(extend_seqs):
@@ -253,8 +254,8 @@ def nsa_cp_round_robin_split_q_seqs(
bs_idx_cpu(List) and bs_idx(torch.Tensor): marks which sequences are ultimately selected,
i.e., those with a partitioned length greater than zero.
"""
cp_size = get_attention_tp_size()
cp_rank = get_attention_tp_rank()
cp_size = get_attention_cp_size()
cp_rank = get_attention_cp_rank()
# len(ret_q_lens_cpu) == len(bs_idx_cpu)
ret_q_lens_cpu, bs_idx_cpu = nsa_cp_round_robin_split_q_seqs_cpu(extend_seqs_cpu)
ret_q_lens = torch.empty(
@@ -299,7 +300,7 @@ def cp_attn_tp_all_gather_reorganazied_into_tensor(
if pad_size > 0:
input_ = F.pad(input_, (0, 0, 0, pad_size), mode="constant", value=0)
with use_symmetric_memory(
get_attention_tp_group(), disabled=not is_allocation_symmetric()
get_attention_cp_group(), disabled=not is_allocation_symmetric()
):
input_tensor_all = torch.empty(
max_len * attn_tp_size,
@@ -308,7 +309,7 @@ def cp_attn_tp_all_gather_reorganazied_into_tensor(
dtype=input_.dtype,
)
# step2
get_attention_tp_group().cp_all_gather_into_tensor_async(
get_attention_cp_group().cp_all_gather_into_tensor_async(
input_tensor_all, input_, stream_op
)
# step3
@@ -356,12 +357,12 @@ def cp_all_gather_rerange_output(input_tensor, cp_size, forward_batch, stream):
"""
if is_nsa_prefill_cp_round_robin_split():
with use_symmetric_memory(
get_attention_tp_group(), disabled=not is_allocation_symmetric()
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_tp_all_gather_into_tensor(
attn_cp_all_gather_into_tensor(
output_tensor,
input_tensor,
)
python/sglang/srt/layers/communicator.py
+8
View File
@@ -39,6 +39,8 @@ from sglang.srt.layers.dp_attention import (
dp_gather_partial,
dp_reduce_scatter_tensor,
dp_scatter,
get_attention_cp_rank,
get_attention_cp_size,
get_attention_dp_size,
get_attention_tp_rank,
get_attention_tp_size,
@@ -611,6 +613,8 @@ class CommunicateContext:
attn_tp_rank: int
attn_tp_size: int
attn_dp_size: int
attn_cp_rank: int
attn_cp_size: int
tp_size: int
cache = None
tp_rank: int
@@ -623,6 +627,8 @@ class CommunicateContext:
attn_tp_rank = get_attention_tp_rank()
attn_tp_size = get_attention_tp_size()
attn_dp_size = get_attention_dp_size()
attn_cp_size = get_attention_cp_size()
attn_cp_rank = get_attention_cp_rank()
tp_size = get_tensor_model_parallel_world_size()
tp_rank = get_tensor_model_parallel_rank()
process_group_sizes = {
@@ -636,6 +642,8 @@ class CommunicateContext:
attn_tp_rank=attn_tp_rank,
attn_tp_size=attn_tp_size,
attn_dp_size=attn_dp_size,
attn_cp_rank=attn_cp_rank,
attn_cp_size=attn_cp_size,
tp_size=tp_size,
tp_rank=tp_rank,
)
python/sglang/srt/layers/communicator_nsa_cp.py
+6 -6
View File
@@ -32,8 +32,8 @@ from sglang.srt.layers.communicator import (
ScatterMode,
)
from sglang.srt.layers.dp_attention import (
attn_tp_all_gather_into_tensor,
attn_tp_reduce_scatter_tensor,
attn_cp_all_gather_into_tensor,
attn_cp_reduce_scatter_tensor,
get_local_dp_buffer,
)
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
@@ -157,7 +157,7 @@ class NSACPCommunicateWithAllReduceAndLayerNormFn(
get_local_dp_buffer(),
hidden_states,
)
attn_tp_all_gather_into_tensor(
attn_cp_all_gather_into_tensor(
hidden_states,
local_hidden_states,
)
@@ -203,8 +203,8 @@ class NSACPCommunicateSummableTensorPairFn(CommunicateSummableTensorPairFn):
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_tp_size)[
context.attn_tp_rank
hidden_states = hidden_states.tensor_split(context.attn_cp_size)[
context.attn_cp_rank
]
attn_tp_reduce_scatter_tensor(hidden_states, input_hidden_states)
attn_cp_reduce_scatter_tensor(hidden_states, input_hidden_states)
return hidden_states, residual
python/sglang/srt/layers/dp_attention.py
+48 -47
View File
@@ -12,6 +12,12 @@ import triton.language as tl
from sglang.srt.distributed import (
GroupCoordinator,
get_attn_context_model_parallel_rank,
get_attn_context_model_parallel_world_size,
get_attn_cp_group,
get_attn_tensor_model_parallel_rank,
get_attn_tensor_model_parallel_world_size,
get_attn_tp_group,
get_tensor_model_parallel_rank,
get_tensor_model_parallel_world_size,
get_tp_group,
@@ -31,9 +37,6 @@ logger = logging.getLogger(__name__)
if TYPE_CHECKING:
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
_ATTN_TP_GROUP: Optional[GroupCoordinator] = None
_ATTN_TP_RANK: Optional[int] = None
_ATTN_TP_SIZE: Optional[int] = None
_ATTN_DP_RANK: Optional[int] = None
_ATTN_DP_SIZE: Optional[int] = None
_LOCAL_ATTN_DP_SIZE: Optional[int] = None
@@ -224,14 +227,20 @@ def is_dp_max_padding() -> bool:
return _DpGatheredBufferWrapper.is_dp_max_padding()
def compute_dp_attention_world_info(enable_dp_attention, tp_rank, tp_size, dp_size):
if not enable_dp_attention:
return tp_rank, tp_size, 0
attn_tp_size = tp_size // dp_size
attn_dp_rank = tp_rank // attn_tp_size
def compute_dp_attention_world_info(
enable_dp_attention, tp_rank, tp_size, dp_size, attn_cp_size: int = 1
):
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 = ((cp_rank * dp_size) + dp_rank) * attn_tp_size + attn_tp_rank
attn_dp_rank = tp_rank // (attn_tp_size * attn_cp_size)
return attn_tp_rank, attn_tp_size, attn_dp_rank
@@ -256,23 +265,20 @@ def initialize_dp_attention(
server_args: ServerArgs,
model_config: ModelConfig,
):
global _ATTN_TP_GROUP, _ATTN_TP_RANK, _ATTN_TP_SIZE, _ATTN_DP_RANK, _ATTN_DP_SIZE
global _ATTN_DP_RANK, _ATTN_DP_SIZE
global _LOCAL_ATTN_DP_SIZE, _LOCAL_ATTN_DP_RANK, _ENABLE_DP_ATTENTION_FLAG
from sglang.srt.layers.sampler import SYNC_TOKEN_IDS_ACROSS_TP
enable_dp_attention = server_args.enable_dp_attention
tp_size = server_args.tp_size
dp_size = server_args.dp_size
moe_dense_tp_size = server_args.moe_dense_tp_size
pp_size = server_args.pp_size
tp_rank = get_tensor_model_parallel_rank()
attn_cp_size = server_args.attn_cp_size
_ENABLE_DP_ATTENTION_FLAG = enable_dp_attention
_ATTN_TP_RANK, _ATTN_TP_SIZE, _ATTN_DP_RANK = compute_dp_attention_world_info(
enable_dp_attention, tp_rank, tp_size, dp_size
tp_rank = get_tensor_model_parallel_rank()
tp_size = get_tensor_model_parallel_world_size()
_, _, _ATTN_DP_RANK = compute_dp_attention_world_info(
enable_dp_attention, tp_rank, tp_size, dp_size, attn_cp_size
)
_, _, _LOCAL_ATTN_DP_RANK = compute_dp_attention_local_info(
enable_dp_attention, tp_rank, tp_size, dp_size, moe_dense_tp_size
@@ -288,28 +294,6 @@ def initialize_dp_attention(
_ATTN_DP_SIZE = 1
_LOCAL_ATTN_DP_SIZE = 1
tp_group = get_tp_group()
# Trick to solve circular references
from sglang.srt.layers.attention.nsa.utils import is_nsa_enable_prefill_cp
use_pynccl = True if is_nsa_enable_prefill_cp() else SYNC_TOKEN_IDS_ACROSS_TP
_ATTN_TP_GROUP = GroupCoordinator(
[
list(range(head, head + _ATTN_TP_SIZE))
for head in range(0, pp_size * tp_size, _ATTN_TP_SIZE)
],
tp_group.local_rank,
torch.distributed.get_backend(tp_group.device_group),
use_pynccl=use_pynccl,
use_pymscclpp=False,
use_custom_allreduce=False,
use_torch_symm_mem_all_reduce=False,
use_hpu_communicator=False,
use_xpu_communicator=False,
use_npu_communicator=False,
group_name="attention_tp",
)
_DpGatheredBufferWrapper.set_metadata(
hidden_size=model_config.hidden_size,
dtype=model_config.dtype,
@@ -326,18 +310,27 @@ def is_allocation_symmetric() -> bool:
def get_attention_tp_group() -> GroupCoordinator:
assert _ATTN_TP_GROUP is not None, "dp attention not initialized!"
return _ATTN_TP_GROUP
return get_attn_tp_group()
def get_attention_tp_rank() -> int:
assert _ATTN_TP_RANK is not None, "dp attention not initialized!"
return _ATTN_TP_RANK
return get_attn_tensor_model_parallel_rank()
def get_attention_tp_size() -> int:
assert _ATTN_TP_SIZE is not None, "dp attention not initialized!"
return _ATTN_TP_SIZE
return get_attn_tensor_model_parallel_world_size()
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 get_attention_dp_rank() -> int:
@@ -564,6 +557,10 @@ def attn_tp_reduce_scatter_tensor(output: torch.Tensor, input: torch.Tensor):
return get_attention_tp_group().reduce_scatter_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 attn_tp_all_reduce(input: torch.Tensor):
return get_attention_tp_group().all_reduce(input)
@@ -572,5 +569,9 @@ def attn_tp_all_gather_into_tensor(output: torch.Tensor, input: torch.Tensor):
return get_attention_tp_group().all_gather_into_tensor(output, input)
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_tp_all_gather(output_list: List[torch.Tensor], input: torch.Tensor):
return get_attention_tp_group().all_gather(input, output_tensor_list=output_list)
python/sglang/srt/managers/data_parallel_controller.py
+29 -1
View File
@@ -418,6 +418,8 @@ class DataParallelController:
tp_size_per_node * (server_args.node_rank % nnodes_per_tp_group + 1),
)
attn_cp_rank = 0
moe_dp_rank = 0
for pp_rank in pp_rank_range:
for tp_rank in tp_rank_range:
rank_port_args = port_args
@@ -429,6 +431,7 @@ class DataParallelController:
tp_rank,
server_args.tp_size,
server_args.dp_size,
server_args.attn_cp_size,
)
# compute zmq ports for this dp rank
rank_port_args = PortArgs.init_new(
@@ -445,7 +448,30 @@ class DataParallelController:
+ ((pp_rank % pp_size_per_node) * tp_size_per_node)
+ (tp_rank % tp_size_per_node) * server_args.gpu_id_step
)
moe_ep_rank = tp_rank // (server_args.tp_size // server_args.ep_size)
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
)
)
with self.env_lock, maybe_reindex_device_id(gpu_id) as gpu_id:
proc = mp.Process(
target=self.run_scheduler_process_func,
@@ -454,6 +480,8 @@ class DataParallelController:
rank_port_args,
gpu_id,
tp_rank,
attn_cp_rank,
moe_dp_rank,
moe_ep_rank,
pp_rank,
dp_rank,
python/sglang/srt/managers/scheduler.py
+35 -4
View File
@@ -63,6 +63,7 @@ from sglang.srt.environ import envs
from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder
from sglang.srt.layers.dp_attention import (
compute_dp_attention_world_info,
get_attention_cp_group,
get_attention_tp_group,
)
from sglang.srt.layers.moe import initialize_moe_config
@@ -267,6 +268,8 @@ class Scheduler(
tp_rank: int,
moe_ep_rank: int,
pp_rank: int,
attn_cp_rank: int,
moe_dp_rank: int,
dp_rank: Optional[int],
):
self.is_initializing = True
@@ -277,6 +280,10 @@ class Scheduler(
self.tp_rank = tp_rank
self.moe_ep_rank = moe_ep_rank
self.pp_rank = pp_rank
self.attn_cp_rank = attn_cp_rank
self.attn_cp_size = server_args.attn_cp_size
self.moe_dp_rank = moe_dp_rank
self.moe_dp_size = server_args.moe_dp_size
self.dp_rank = dp_rank
self.tp_size = server_args.tp_size
self.moe_ep_size = server_args.ep_size
@@ -322,6 +329,7 @@ class Scheduler(
self.tp_rank,
self.tp_size,
self.dp_size,
self.attn_cp_size,
)
)
@@ -405,7 +413,7 @@ class Scheduler(
context = zmq.Context(2)
self.idle_sleeper = None
if self.pp_rank == 0 and self.attn_tp_rank == 0:
if self.pp_rank == 0 and self.attn_tp_rank == 0 and self.attn_cp_rank == 0:
self.recv_from_tokenizer = get_zmq_socket(
context, zmq.PULL, port_args.scheduler_input_ipc_name, False
)
@@ -506,6 +514,8 @@ class Scheduler(
tp_rank=self.tp_rank,
moe_ep_rank=self.moe_ep_rank,
pp_rank=self.pp_rank,
attn_cp_rank=self.attn_cp_rank,
moe_dp_rank=self.moe_dp_rank,
dp_rank=self.dp_rank,
nccl_port=self.nccl_port,
)
@@ -524,6 +534,8 @@ class Scheduler(
nccl_port=self.nccl_port,
target_worker=self.tp_worker,
dp_rank=self.dp_rank,
attn_cp_rank=self.attn_cp_rank,
moe_dp_rank=self.moe_dp_rank,
)
if self.server_args.speculative_draft_load_format is not None:
@@ -571,6 +583,8 @@ class Scheduler(
self.tp_cpu_group = self.tp_group.cpu_group
self.attn_tp_group = get_attention_tp_group()
self.attn_tp_cpu_group = self.attn_tp_group.cpu_group
self.attn_cp_group = get_attention_cp_group()
self.attn_cp_cpu_group = self.attn_cp_group.cpu_group
self.pp_group = get_pp_group()
self.world_group = get_world_group()
@@ -1201,7 +1215,7 @@ class Scheduler(
return []
if self.pp_rank == 0:
if self.attn_tp_rank == 0:
if self.attn_tp_rank == 0 and self.attn_cp_rank == 0:
recv_reqs = []
while True:
@@ -1225,7 +1239,7 @@ class Scheduler(
else:
recv_reqs = None
else:
if self.attn_tp_rank == 0:
if self.attn_tp_rank == 0 and self.attn_cp_rank == 0:
dp_offset = self.attn_dp_rank * self.attn_tp_size
recv_reqs = point_to_point_pyobj(
[],
@@ -1241,7 +1255,7 @@ class Scheduler(
recv_reqs = self.input_blocker.handle(recv_reqs)
if self.server_args.enable_dp_attention:
if self.attn_tp_rank == 0:
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
@@ -1254,6 +1268,15 @@ class Scheduler(
self.attn_tp_cpu_group,
src=self.attn_tp_group.ranks[0],
)
if self.attn_cp_size != 1:
work_reqs = broadcast_pyobj(
work_reqs,
self.attn_cp_group.rank,
self.attn_cp_cpu_group,
src=self.attn_cp_group.ranks[0],
)
if self.tp_size != 1:
control_reqs = broadcast_pyobj(
control_reqs,
@@ -3003,6 +3026,8 @@ def run_scheduler_process(
port_args: PortArgs,
gpu_id: int,
tp_rank: int,
attn_cp_rank: int,
moe_dp_rank: int,
moe_ep_rank: int,
pp_rank: int,
dp_rank: Optional[int],
@@ -3017,6 +3042,10 @@ def run_scheduler_process(
prefix += f" DP{dp_rank}"
if server_args.pp_size > 1:
prefix += f" PP{pp_rank}"
if server_args.attn_cp_size > 1:
prefix += f" ATTN_CP{attn_cp_rank}"
if server_args.moe_dp_size > 1:
prefix += f" MOE_DP{moe_dp_rank}"
if server_args.tp_size > 1:
prefix += f" TP{tp_rank}"
if server_args.ep_size > 1:
@@ -3061,6 +3090,8 @@ def run_scheduler_process(
tp_rank,
moe_ep_rank,
pp_rank,
attn_cp_rank,
moe_dp_rank,
dp_rank,
)
result_dict = {
python/sglang/srt/managers/scheduler_dp_attn_mixin.py
+6 -2
View File
@@ -25,6 +25,7 @@ _ENABLE_METRICS_DP_ATTENTION = envs.SGLANG_ENABLE_METRICS_DP_ATTENTION.get()
class MLPSyncBatchInfo:
dp_size: int
tp_size: int
cp_size: int
num_tokens: int
num_tokens_for_logprob: int
@@ -72,7 +73,7 @@ class MLPSyncBatchInfo:
def all_gather(self, device, group: torch.distributed.ProcessGroup):
local_info_tensor = self._get_local_tensor(device=device)
global_info_tensor = torch.empty(
(self.dp_size, self.tp_size, 6),
(self.dp_size, self.tp_size * self.cp_size, 6),
dtype=torch.int64,
device=device,
)
@@ -88,7 +89,7 @@ class MLPSyncBatchInfo:
tp_active_ranks = get_tp_group().active_ranks
# Set fallback values for inactive ranks
tp_info = global_info_tensor.view(self.dp_size * self.tp_size, 6)
tp_info = global_info_tensor.view(self.dp_size * self.tp_size * self.cp_size, 6)
tp_info[tp_active_ranks == 0] = self._get_fallback_tensor(device=device)
tp0_info = global_info_tensor[:, 0, :]
@@ -129,6 +130,7 @@ def prepare_mlp_sync_batch_raw(
local_batch: ScheduleBatch,
dp_size: int,
attn_tp_size: int,
attn_cp_size: int,
tp_group: GroupCoordinator,
get_idle_batch: Callable[[], ScheduleBatch],
disable_cuda_graph: bool,
@@ -185,6 +187,7 @@ def prepare_mlp_sync_batch_raw(
mlp_sync_info = MLPSyncBatchInfo(
dp_size=dp_size,
tp_size=attn_tp_size,
cp_size=attn_cp_size,
num_tokens=num_tokens,
num_tokens_for_logprob=num_tokens_for_logprob,
can_cuda_graph=can_cuda_graph,
@@ -226,6 +229,7 @@ class SchedulerDPAttnMixin:
local_batch,
dp_size=self.server_args.dp_size,
attn_tp_size=self.attn_tp_size,
attn_cp_size=self.attn_cp_size,
tp_group=self.tp_group,
get_idle_batch=self.get_idle_batch,
disable_cuda_graph=self.server_args.disable_cuda_graph,
python/sglang/srt/managers/tp_worker.py
+4
View File
@@ -212,6 +212,8 @@ class TpModelWorker(BaseTpWorker):
tp_rank: int,
moe_ep_rank: int,
pp_rank: int,
attn_cp_rank: int,
moe_dp_rank: int,
dp_rank: Optional[int],
nccl_port: int,
is_draft_worker: bool = False,
@@ -234,6 +236,8 @@ class TpModelWorker(BaseTpWorker):
self.is_multi_layer_eagle = is_multi_layer_eagle
self.req_to_token_pool = req_to_token_pool
self.token_to_kv_pool_allocator = token_to_kv_pool_allocator
self.attn_cp_rank = attn_cp_rank
self.moe_dp_rank = moe_dp_rank
# MTP model runners
self.model_runner_list: List[ModelRunner] = []
python/sglang/srt/model_executor/cuda_graph_runner.py
+4
View File
@@ -43,6 +43,7 @@ from sglang.srt.dllm.config import DllmConfig
from sglang.srt.layers.attention.nsa.utils import is_nsa_enable_prefill_cp
from sglang.srt.layers.dp_attention import (
DpPaddingMode,
get_attention_cp_size,
get_attention_tp_rank,
get_attention_tp_size,
set_dp_buffer_len,
@@ -204,6 +205,9 @@ def get_batch_sizes_to_capture(model_runner: ModelRunner, num_tokens_per_bs=1):
if require_gathered_buffer(server_args):
mul_base *= get_attention_tp_size()
if mul_base % get_attention_cp_size() != 0:
mul_base *= get_attention_cp_size()
# Model input token count = bs * num_tokens_per_bs; must be a multiple of attn_tp_size.
capture_bs = [bs for bs in capture_bs if bs * num_tokens_per_bs % mul_base == 0]
python/sglang/srt/model_executor/forward_batch_info.py
+6
View File
@@ -45,6 +45,7 @@ from sglang.srt.distributed.parallel_state import (
from sglang.srt.layers.attention.nsa.utils import NSAContextParallelMetadata
from sglang.srt.layers.dp_attention import (
DpPaddingMode,
get_attention_cp_size,
get_attention_dp_rank,
get_attention_tp_rank,
get_attention_tp_size,
@@ -749,6 +750,11 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
# there is no reduce-scatter in LM logprob, so we do not need to adjust the padded length for logprob
global_num_tokens[i] = ceil_align(global_num_tokens[i], attn_tp_size)
# make sure that each rank has the same number of tokens to do collective communication.
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)
dp_padding_mode = DpPaddingMode.get_dp_padding_mode(
self.is_extend_in_batch, global_num_tokens
)
python/sglang/srt/model_executor/model_runner.py
+11 -3
View File
@@ -290,6 +290,8 @@ class ModelRunner(ModelRunnerKVCacheMixin):
nccl_port: int,
server_args: ServerArgs,
dp_rank: Optional[int] = None,
attn_cp_rank: Optional[int] = None,
moe_dp_rank: Optional[int] = None,
is_draft_worker: bool = False,
req_to_token_pool: Optional[ReqToTokenPool] = None,
token_to_kv_pool_allocator: Optional[BaseTokenToKVPoolAllocator] = None,
@@ -303,9 +305,13 @@ class ModelRunner(ModelRunnerKVCacheMixin):
self.tp_size = tp_size
self.moe_ep_rank = moe_ep_rank
self.moe_ep_size = moe_ep_size
self.dp_size = server_args.dp_size
self.dp_size = server_args.dp_size if server_args.enable_dp_attention else 1
self.pp_rank = pp_rank
self.pp_size = pp_size
self.attn_cp_rank = attn_cp_rank
self.attn_cp_size = server_args.attn_cp_size
self.moe_dp_rank = moe_dp_rank
self.moe_dp_size = server_args.moe_dp_size
self.model_config = model_config
self.dist_port = nccl_port
self.server_args = server_args
@@ -586,8 +592,7 @@ class ModelRunner(ModelRunnerKVCacheMixin):
(
self.max_total_num_tokens // 2
if server_args.max_running_requests is None
else server_args.max_running_requests
// (server_args.dp_size if server_args.enable_dp_attention else 1)
else server_args.max_running_requests // (self.dp_size)
),
self.req_to_token_pool.size,
)
@@ -797,8 +802,11 @@ class ModelRunner(ModelRunnerKVCacheMixin):
)
initialize_model_parallel(
tensor_model_parallel_size=self.tp_size,
attention_data_parallel_size=self.dp_size,
pipeline_model_parallel_size=self.pp_size,
expert_model_parallel_size=self.moe_ep_size,
attention_context_model_parallel_size=self.attn_cp_size,
moe_data_model_parallel_size=self.moe_dp_size,
duplicate_tp_group=self.server_args.enable_pdmux,
)
initialize_dp_attention(
python/sglang/srt/models/deepseek_v2.py
+6 -6
View File
@@ -72,6 +72,8 @@ from sglang.srt.layers.communicator import (
)
from sglang.srt.layers.communicator_nsa_cp import NSACPLayerCommunicator
from sglang.srt.layers.dp_attention import (
get_attention_cp_rank,
get_attention_cp_size,
get_attention_tp_rank,
get_attention_tp_size,
is_dp_attention_enabled,
@@ -1097,9 +1099,7 @@ class DeepseekV2AttentionMLA(nn.Module, DeepseekMHAForwardMixin):
assert self.use_nsa, "CP currently only supports deepseek v3.2 model"
# cp reuse the attn_tp comm group but need to duplicate the weights
if self.nsa_enable_prefill_cp and self.use_nsa:
attn_tp_rank = 0
attn_tp_size = 1
self.cp_size = get_attention_tp_size()
self.cp_size = get_attention_cp_size()
self.num_heads = num_heads
assert num_heads % attn_tp_size == 0
self.num_local_heads = num_heads // attn_tp_size
@@ -2512,7 +2512,7 @@ class DeepseekV2Model(nn.Module):
self.pp_group = get_pp_group()
self.nsa_enable_prefill_cp = is_nsa_enable_prefill_cp()
if self.nsa_enable_prefill_cp:
self.cp_size = get_attention_tp_size()
self.cp_size = get_attention_cp_size()
else:
self.cp_size = None
@@ -2827,8 +2827,8 @@ class DeepseekV2ForCausalLM(nn.Module, DeepseekV2WeightLoaderMixin):
self.nsa_enable_prefill_cp = is_nsa_enable_prefill_cp()
if self.nsa_enable_prefill_cp:
self.cp_rank = get_attention_tp_rank()
self.cp_size = get_attention_tp_size()
self.cp_rank = get_attention_cp_rank()
self.cp_size = get_attention_cp_size()
else:
self.cp_rank = self.cp_size = None
python/sglang/srt/server_args.py
+48
View File
@@ -417,6 +417,9 @@ class ServerArgs:
dp_size: int = 1
load_balance_method: str = "auto"
attn_cp_size: int = 1
moe_dp_size: int = 1
# Multi-node distributed serving
dist_init_addr: Optional[str] = None
nnodes: int = 1
@@ -745,6 +748,9 @@ class ServerArgs:
# Handle data parallelism.
self._handle_data_parallelism()
# Handle context parallelism.
self._handle_context_parallelism()
# Handle MoE configurations.
self._handle_moe_kernel_config()
self._handle_a2a_moe()
@@ -2042,6 +2048,32 @@ class ServerArgs:
if self.grammar_backend is None:
self.grammar_backend = "xgrammar"
def _handle_context_parallelism(self):
if self.attn_cp_size > 1:
# The tp_size is the world size, not the real tensor parallel size
assert (
self.tp_size % self.attn_cp_size == 0
), "tp_size must be divisible by attn_cp_size"
assert (
self.tp_size % (self.dp_size * self.attn_cp_size) == 0
), "tp_size must be divisible by dp_size * attn_cp_size"
assert self.pp_size == 1, "PP is not supported with context parallelism"
if self.moe_dp_size > 1:
# The tp_size is the world size, not the real tensor parallel size
assert (
self.tp_size % self.moe_dp_size == 0
), "tp_size must be divisible by moe_dp_size"
assert (
self.ep_size * self.moe_dp_size <= self.tp_size
), "ep_size * moe_dp_size must be less than or equal to tp_size"
assert self.pp_size == 1, "PP is not supported with context parallelism"
if self.ep_size > 1:
assert (
self.ep_size * self.moe_dp_size == self.tp_size
), "ep_size * moe_dp_size must be equal to tp_size"
def _handle_data_parallelism(self):
if self.dp_size == 1:
self.enable_dp_attention = False
@@ -3241,6 +3273,20 @@ class ServerArgs:
default=ServerArgs.tp_size,
help="The tensor parallelism size.",
)
parser.add_argument(
"--attention-context-parallel-size",
"--attn-cp-size",
type=int,
default=ServerArgs.attn_cp_size,
help="The attention context parallelism size.",
)
parser.add_argument(
"--moe-data-parallel-size",
"--moe-dp-size",
type=int,
default=ServerArgs.moe_dp_size,
help="The moe data parallelism size.",
)
parser.add_argument(
"--pipeline-parallel-size",
"--pp-size",
@@ -4989,6 +5035,8 @@ class ServerArgs:
def from_cli_args(cls, args: argparse.Namespace):
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
python/sglang/srt/speculative/eagle_worker.py
+4
View File
@@ -84,6 +84,8 @@ class EAGLEWorker(TpModelWorker):
tp_rank: int,
dp_rank: Optional[int],
moe_ep_rank: int,
attn_cp_rank: int,
moe_dp_rank: int,
nccl_port: int,
target_worker: TpModelWorker,
):
@@ -144,6 +146,8 @@ class EAGLEWorker(TpModelWorker):
pp_rank=0, # FIXME
dp_rank=dp_rank,
moe_ep_rank=moe_ep_rank,
attn_cp_rank=attn_cp_rank,
moe_dp_rank=moe_dp_rank,
nccl_port=nccl_port,
is_draft_worker=True,
req_to_token_pool=self.req_to_token_pool,
python/sglang/srt/speculative/eagle_worker_v2.py
+17 -1
View File
@@ -86,6 +86,8 @@ class EagleDraftWorker(BaseDraftWorker):
tp_rank: int,
dp_rank: int,
moe_ep_rank: int,
attn_cp_rank: int,
moe_dp_rank: int,
nccl_port: int,
target_worker: TpModelWorker,
):
@@ -97,6 +99,8 @@ class EagleDraftWorker(BaseDraftWorker):
self.moe_ep_rank = moe_ep_rank
self.nccl_port = nccl_port
self.target_worker = target_worker
self.attn_cp_rank = attn_cp_rank
self.moe_dp_rank = moe_dp_rank
# Args for easy access
self.device = server_args.device
@@ -134,6 +138,8 @@ class EagleDraftWorker(BaseDraftWorker):
pp_rank=0, # FIXME
dp_rank=dp_rank,
moe_ep_rank=moe_ep_rank,
attn_cp_rank=attn_cp_rank,
moe_dp_rank=moe_dp_rank,
nccl_port=nccl_port,
is_draft_worker=True,
req_to_token_pool=self.req_to_token_pool,
@@ -582,6 +588,8 @@ class EAGLEWorkerV2(BaseSpecWorker):
tp_rank: int,
dp_rank: Optional[int],
moe_ep_rank: int,
attn_cp_rank: int,
moe_dp_rank: int,
nccl_port: int,
target_worker: TpModelWorker,
):
@@ -608,7 +616,15 @@ class EAGLEWorkerV2(BaseSpecWorker):
server_args.context_length = target_worker.model_runner.model_config.context_len
self._draft_worker = EagleDraftWorker(
server_args, gpu_id, tp_rank, dp_rank, moe_ep_rank, nccl_port, target_worker
server_args,
gpu_id,
tp_rank,
dp_rank,
moe_ep_rank,
attn_cp_rank,
moe_dp_rank,
nccl_port,
target_worker,
)
# Some dummy tensors
python/sglang/srt/speculative/multi_layer_eagle_worker.py
+4
View File
@@ -76,6 +76,8 @@ class MultiLayerEagleWorker(TpModelWorker):
tp_rank: int,
dp_rank: Optional[int],
moe_ep_rank: int,
attn_cp_rank: int,
moe_dp_rank: int,
nccl_port: int,
target_worker: TpModelWorker,
):
@@ -135,6 +137,8 @@ class MultiLayerEagleWorker(TpModelWorker):
pp_rank=0, # FIXME
dp_rank=dp_rank,
moe_ep_rank=moe_ep_rank,
attn_cp_rank=attn_cp_rank,
moe_dp_rank=moe_dp_rank,
nccl_port=nccl_port,
is_draft_worker=True,
req_to_token_pool=self.req_to_token_pool,
python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py
+15 -1
View File
@@ -70,6 +70,8 @@ class MultiLayerEagleDraftWorker(BaseDraftWorker):
tp_rank: int,
dp_rank: int,
moe_ep_rank: int,
attn_cp_rank: int,
moe_dp_rank: int,
nccl_port: int,
target_worker: TpModelWorker,
):
@@ -117,6 +119,8 @@ class MultiLayerEagleDraftWorker(BaseDraftWorker):
pp_rank=0, # FIXME
dp_rank=dp_rank,
moe_ep_rank=moe_ep_rank,
attn_cp_rank=attn_cp_rank,
moe_dp_rank=moe_dp_rank,
nccl_port=nccl_port,
is_draft_worker=True,
req_to_token_pool=self.req_to_token_pool,
@@ -532,6 +536,8 @@ class MultiLayerEagleWorkerV2(BaseSpecWorker):
tp_rank: int,
dp_rank: Optional[int],
moe_ep_rank: int,
attn_cp_rank: int,
moe_dp_rank: int,
nccl_port: int,
target_worker: TpModelWorker,
):
@@ -557,7 +563,15 @@ class MultiLayerEagleWorkerV2(BaseSpecWorker):
server_args.context_length = target_worker.model_runner.model_config.context_len
self._draft_worker = MultiLayerEagleDraftWorker(
server_args, gpu_id, tp_rank, dp_rank, moe_ep_rank, nccl_port, target_worker
server_args,
gpu_id,
tp_rank,
dp_rank,
moe_ep_rank,
attn_cp_rank,
moe_dp_rank,
nccl_port,
target_worker,
)
# Some dummy tensors
python/sglang/srt/speculative/ngram_worker.py
+2
View File
@@ -30,6 +30,8 @@ class NGRAMWorker:
tp_rank: int,
dp_rank: Optional[int],
moe_ep_rank: int,
attn_cp_rank: int,
moe_dp_rank: int,
nccl_port: int,
target_worker: TpModelWorker,
):
python/sglang/srt/speculative/standalone_worker.py
+4
View File
@@ -30,6 +30,8 @@ class StandaloneWorker(EAGLEWorker):
tp_rank: int,
dp_rank: Optional[int],
moe_ep_rank: int,
attn_cp_rank: int,
moe_dp_rank: int,
nccl_port: int,
target_worker: TpModelWorker,
):
@@ -79,6 +81,8 @@ class StandaloneWorker(EAGLEWorker):
pp_rank=0, # FIXME
dp_rank=dp_rank,
moe_ep_rank=moe_ep_rank,
attn_cp_rank=attn_cp_rank,
moe_dp_rank=moe_dp_rank,
nccl_port=nccl_port,
is_draft_worker=True,
req_to_token_pool=self.req_to_token_pool,
python/sglang/srt/speculative/standalone_worker_v2.py
+17 -1
View File
@@ -42,6 +42,8 @@ class StandaloneDraftWorker(EagleDraftWorker):
tp_rank: int,
dp_rank: int,
moe_ep_rank: int,
attn_cp_rank: int,
moe_dp_rank: int,
nccl_port: int,
target_worker: TpModelWorker,
):
@@ -53,6 +55,8 @@ class StandaloneDraftWorker(EagleDraftWorker):
self.moe_ep_rank = moe_ep_rank
self.nccl_port = nccl_port
self.target_worker = target_worker
self.attn_cp_rank = attn_cp_rank
self.moe_dp_rank = moe_dp_rank
# Args for easy access
self.device = server_args.device
@@ -89,6 +93,8 @@ class StandaloneDraftWorker(EagleDraftWorker):
pp_rank=0, # FIXME
dp_rank=dp_rank,
moe_ep_rank=moe_ep_rank,
attn_cp_rank=attn_cp_rank,
moe_dp_rank=moe_dp_rank,
nccl_port=nccl_port,
is_draft_worker=True,
req_to_token_pool=self.req_to_token_pool,
@@ -131,6 +137,8 @@ class StandaloneWorkerV2(EAGLEWorkerV2):
tp_rank: int,
dp_rank: Optional[int],
moe_ep_rank: int,
attn_cp_rank: int,
moe_dp_rank: int,
nccl_port: int,
target_worker: TpModelWorker,
):
@@ -157,7 +165,15 @@ class StandaloneWorkerV2(EAGLEWorkerV2):
# Create our custom draft worker that doesn't share embeddings/lm_head
self._draft_worker = StandaloneDraftWorker(
server_args, gpu_id, tp_rank, dp_rank, moe_ep_rank, nccl_port, target_worker
server_args,
gpu_id,
tp_rank,
dp_rank,
moe_ep_rank,
attn_cp_rank,
moe_dp_rank,
nccl_port,
target_worker,
)
# Some dummy tensors