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

Read process groups through the runtime context (#40068)

Browse Source
This commit is contained in:
Cheng Wan
2026-09-18 17:40:32 -07:00
committed by GitHub
parent 5931fd60ee
commit afe71f4b9e
165 changed files with 569 additions and 608 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/ops/layernorm/mhc.py
+7 -4
View File
@@ -13,12 +13,11 @@ from sglang.kernels.jit.utils import is_arch_support_pdl
from sglang.srt.distributed.device_communicators.pynccl_allocator import ( from sglang.srt.distributed.device_communicators.pynccl_allocator import (
use_symmetric_memory, use_symmetric_memory,
) )
from sglang.srt.distributed.parallel_state import get_tp_group
from sglang.srt.environ import envs from sglang.srt.environ import envs
from sglang.srt.layers.attention.dsa.utils import is_dsa_prefill_cp_interleave from sglang.srt.layers.attention.dsa.utils import is_dsa_prefill_cp_interleave
from sglang.srt.layers.dp_attention import is_allocation_symmetric from sglang.srt.layers.dp_attention import is_allocation_symmetric
from sglang.srt.layers.utils.common import strict_contiguous from sglang.srt.layers.utils.common import strict_contiguous
from sglang.srt.runtime_context import get_platform from sglang.srt.runtime_context import get_parallel, get_platform
from sglang.srt.utils.common import is_gfx1250_supported from sglang.srt.utils.common import is_gfx1250_supported
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@@ -1035,7 +1034,9 @@ def mhc_pre(
# NCCL symmetric path: the Triton inplace MoE runner writes the expert # NCCL symmetric path: the Triton inplace MoE runner writes the expert
# output back into this buffer, so a symmetric input yields a symmetric # output back into this buffer, so a symmetric input yields a symmetric
# all-reduce input. # all-reduce input.
with use_symmetric_memory(get_tp_group(), disabled=not is_allocation_symmetric()): with use_symmetric_memory(
get_parallel().tp_group, disabled=not is_allocation_symmetric()
):
layer_input = torch.empty( layer_input = torch.empty(
num_tokens, hidden_size, dtype=torch.bfloat16, device=residual.device num_tokens, hidden_size, dtype=torch.bfloat16, device=residual.device
) )
@@ -1697,7 +1698,9 @@ def mhc_fused_post_pre(
# layer_input_cur is the post-norm activation fed into the MoE; allocate it # layer_input_cur is the post-norm activation fed into the MoE; allocate it
# in the symmetric memory pool so the Triton inplace MoE runner yields a # in the symmetric memory pool so the Triton inplace MoE runner yields a
# symmetric all-reduce input (see _mhc_pre_impl). # symmetric all-reduce input (see _mhc_pre_impl).
with use_symmetric_memory(get_tp_group(), disabled=not is_allocation_symmetric()): with use_symmetric_memory(
get_parallel().tp_group, disabled=not is_allocation_symmetric()
):
layer_input_cur = torch.empty( layer_input_cur = torch.empty(
num_tokens, num_tokens,
hidden_size, hidden_size,
python/sglang/srt/disaggregation/ascend/conn.py
+2 -2
View File
@@ -15,7 +15,7 @@ from sglang.srt.disaggregation.mooncake.conn import (
MooncakeKVReceiver, MooncakeKVReceiver,
MooncakeKVSender, MooncakeKVSender,
) )
from sglang.srt.distributed import get_pp_group from sglang.srt.runtime_context import get_parallel
from sglang.srt.utils.network import get_local_ip_auto from sglang.srt.utils.network import get_local_ip_auto
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@@ -130,7 +130,7 @@ class AscendKVManager(MooncakeKVManager):
else: else:
sliced_dst_kv_ptrs = [] sliced_dst_kv_ptrs = []
start_layer = self.kv_args.prefill_start_layer start_layer = self.kv_args.prefill_start_layer
transfer_draft_kv = get_pp_group().is_last_rank and draft_kv_layers transfer_draft_kv = get_parallel().pp_group.is_last_rank and draft_kv_layers
if transfer_draft_kv: if transfer_draft_kv:
end_layer = start_layer + src_layers - draft_kv_layers end_layer = start_layer + src_layers - draft_kv_layers
else: else:
python/sglang/srt/disaggregation/ascend/transfer_engine.py
+5 -6
View File
@@ -57,10 +57,7 @@ class AscendTransferEngine(MooncakeTransferEngine):
self.session_id = NetworkAddress(self.hostname, rpc_port).to_host_port_str() self.session_id = NetworkAddress(self.hostname, rpc_port).to_host_port_str()
def initialize(self) -> None: def initialize(self) -> None:
from sglang.srt.distributed.parallel_state import ( from sglang.srt.runtime_context import get_parallel
get_world_group,
get_world_size,
)
transfer_protocol = self._get_transfer_protocol() transfer_protocol = self._get_transfer_protocol()
if transfer_protocol == "device_rdma": if transfer_protocol == "device_rdma":
@@ -68,10 +65,12 @@ class AscendTransferEngine(MooncakeTransferEngine):
# through all_gather to avoid conflicts with rdma initialization. # through all_gather to avoid conflicts with rdma initialization.
tmp_tensor = torch.zeros(1, device="npu") tmp_tensor = torch.zeros(1, device="npu")
output_tensor_list = [ output_tensor_list = [
torch.empty_like(tmp_tensor) for _ in range(get_world_size()) torch.empty_like(tmp_tensor) for _ in range(get_parallel().world_size)
] ]
torch.distributed.all_gather( torch.distributed.all_gather(
output_tensor_list, tmp_tensor, group=get_world_group().device_group output_tensor_list,
tmp_tensor,
group=get_parallel().world_group.device_group,
) )
trans_op_type = self._resolve_trans_op_type(transfer_protocol) trans_op_type = self._resolve_trans_op_type(transfer_protocol)
python/sglang/srt/disaggregation/common/conn.py
+3 -4
View File
@@ -31,7 +31,6 @@ from sglang.srt.disaggregation.utils import (
filter_kv_indices_for_cp_rank, filter_kv_indices_for_cp_rank,
get_dsv41_spec_layout, get_dsv41_spec_layout,
) )
from sglang.srt.distributed import get_pp_group, get_world_group
from sglang.srt.environ import envs from sglang.srt.environ import envs
from sglang.srt.runtime_context import ( from sglang.srt.runtime_context import (
get_disagg, get_disagg,
@@ -259,7 +258,7 @@ class CommonKVManager(BaseKVManager):
self._deferred_ack_targets: Dict[int, Tuple[str, int]] = {} self._deferred_ack_targets: Dict[int, Tuple[str, int]] = {}
self.req_to_decode_prefix_len: Dict[int, int] = {} self.req_to_decode_prefix_len: Dict[int, int] = {}
self.decode_kv_args_table = {} self.decode_kv_args_table = {}
self.pp_group = get_pp_group() self.pp_group = get_parallel().pp_group
# If a timeout happens on the prefill side, it means prefill instances # If a timeout happens on the prefill side, it means prefill instances
# fail to receive the KV indices from the decode instance of this request. # fail to receive the KV indices from the decode instance of this request.
# These timeout requests should be aborted to release the tree cache. # These timeout requests should be aborted to release the tree cache.
@@ -1059,7 +1058,7 @@ class CommonKVManager(BaseKVManager):
"multi-node prefill mode." "multi-node prefill mode."
) )
world_group = get_world_group() world_group = get_parallel().world_group
synced_port = world_group.broadcast_object(local_port, src=0) synced_port = world_group.broadcast_object(local_port, src=0)
if synced_port != local_port: if synced_port != local_port:
logger.info( logger.info(
@@ -1111,7 +1110,7 @@ class CommonKVManager(BaseKVManager):
} }
if envs.SGLANG_RUST_SERVER.get() and self.attn_dp_size > 1: if envs.SGLANG_RUST_SERVER.get() and self.attn_dp_size > 1:
topology_rows = get_world_group().all_gather_object(payload) topology_rows = get_parallel().world_group.all_gather_object(payload)
# Every scheduler contributes a topology row. Only the scheduler # Every scheduler contributes a topology row. Only the scheduler
# ranks that own a Rust listener populate their local registry. # ranks that own a Rust listener populate their local registry.
if self.kv_args.rust_http_port is None: if self.kv_args.rust_http_port is None:
python/sglang/srt/disaggregation/encoder/server.py
+2 -3
View File
@@ -35,7 +35,6 @@ from sglang.srt.disaggregation.encoder.receiver import (
from sglang.srt.distributed.parallel_state import ( from sglang.srt.distributed.parallel_state import (
get_default_distributed_backend, get_default_distributed_backend,
get_mooncake_transfer_engine, get_mooncake_transfer_engine,
get_tp_group,
init_distributed_environment, init_distributed_environment,
initialize_model_parallel, initialize_model_parallel,
) )
@@ -654,7 +653,7 @@ class MMEncoder:
get_parallel().tp_size, get_parallel().tp_size,
embedding_store=embedding_store, embedding_store=embedding_store,
hidden_dims=self._embedding_dims, hidden_dims=self._embedding_dims,
tp_group=get_tp_group().cpu_group, tp_group=get_parallel().tp_group.cpu_group,
all_rank_get=False, all_rank_get=False,
dtype=self._embedding_dtype, dtype=self._embedding_dtype,
) )
@@ -1275,7 +1274,7 @@ class MMEncoder:
layout_digest: tuple[int, int], layout_digest: tuple[int, int],
) -> List[torch.Tensor]: ) -> List[torch.Tensor]:
"""Raise the same preparation error on every TP rank.""" """Raise the same preparation error on every TP rank."""
tp_group = get_tp_group() tp_group = get_parallel().tp_group
error_code = ( error_code = (
int( int(
local_error.code local_error.code
python/sglang/srt/disaggregation/prefill.py
+3 -2
View File
@@ -62,7 +62,6 @@ from sglang.srt.disaggregation.utils import (
prepare_abort, prepare_abort,
setup_state_kv_args, setup_state_kv_args,
) )
from sglang.srt.distributed import get_pp_group
from sglang.srt.environ import envs from sglang.srt.environ import envs
from sglang.srt.managers.schedule_batch import ( from sglang.srt.managers.schedule_batch import (
FINISH_ABORT, FINISH_ABORT,
@@ -89,6 +88,7 @@ from sglang.srt.observability.scheduler_stage_metrics import (
) )
from sglang.srt.runtime_context import ( from sglang.srt.runtime_context import (
get_disagg, get_disagg,
get_parallel,
get_schedule, get_schedule,
) )
from sglang.srt.utils import is_npu from sglang.srt.utils import is_npu
@@ -266,7 +266,8 @@ class PrefillBootstrapQueue:
draft_kv_pool = ( draft_kv_pool = (
self.draft_token_to_kv_pool self.draft_token_to_kv_pool
if transfer_draft_cache and (not _is_npu or get_pp_group().is_last_rank) if transfer_draft_cache
and (not _is_npu or get_parallel().pp_group.is_last_rank)
else None else None
) )
num_draft_entries = 0 num_draft_entries = 0
python/sglang/srt/elastic_ep/elastic_ep.py
+2 -2
View File
@@ -7,7 +7,7 @@ from typing import TYPE_CHECKING, Callable, Iterator, List, Optional
import torch import torch
from sglang.srt.distributed import get_world_group, parallel_state from sglang.srt.distributed import parallel_state
from sglang.srt.distributed.utils import get_global_tcp_store from sglang.srt.distributed.utils import get_global_tcp_store
from sglang.srt.eplb.expert_location import broadcast_global_expert_location_metadata from sglang.srt.eplb.expert_location import broadcast_global_expert_location_metadata
from sglang.srt.runtime_context import ( from sglang.srt.runtime_context import (
@@ -446,7 +446,7 @@ def join_process_groups() -> None:
def get_healthy_expert_location_src_rank( def get_healthy_expert_location_src_rank(
*, invoked_in_elastic_ep_rejoin_path: bool *, invoked_in_elastic_ep_rejoin_path: bool
) -> int: ) -> int:
world_group = get_world_group() world_group = get_parallel().world_group
# NOTE: do not key off `self.server_args.elastic_ep_rejoin` here. # NOTE: do not key off `self.server_args.elastic_ep_rejoin` here.
# A rank that was started as a rejoin rank may later act as a healthy # A rank that was started as a rejoin rank may later act as a healthy
# rank in a subsequent recovery cycle. # rank in a subsequent recovery cycle.
python/sglang/srt/elastic_ep/expert_backup_client.py
+4 -8
View File
@@ -7,10 +7,6 @@ from typing import Any, Callable
import torch import torch
import zmq import zmq
from sglang.srt.distributed.parallel_state import (
get_world_group,
get_world_size,
)
from sglang.srt.environ import envs from sglang.srt.environ import envs
from sglang.srt.eplb.expert_location import get_global_expert_location_metadata from sglang.srt.eplb.expert_location import get_global_expert_location_metadata
from sglang.srt.managers.io_struct import UpdateExpertBackupReq, sock_recv, sock_send from sglang.srt.managers.io_struct import UpdateExpertBackupReq, sock_recv, sock_send
@@ -54,23 +50,23 @@ class ExpertBackupClient:
self.buffer_size = 0 self.buffer_size = 0
self.use_backup = False self.use_backup = False
local_ip = get_local_ip_auto() local_ip = get_local_ip_auto()
all_ips = [None] * get_world_size() all_ips = [None] * get_parallel().world_size
torch.distributed.all_gather_object( torch.distributed.all_gather_object(
all_ips, local_ip, group=get_world_group().cpu_group all_ips, local_ip, group=get_parallel().world_group.cpu_group
) )
logger.info(f"all_ips: {all_ips}") logger.info(f"all_ips: {all_ips}")
for i in range(self.engine_num): for i in range(self.engine_num):
self.recv_list[i] = context.socket(zmq.SUB) self.recv_list[i] = context.socket(zmq.SUB)
self.recv_list[i].connect( self.recv_list[i].connect(
f"tcp://{all_ips[i * get_world_size() // get_parallel().nnodes]}:{PORT_BASE + i * 2 + 1}" f"tcp://{all_ips[i * get_parallel().world_size // get_parallel().nnodes]}:{PORT_BASE + i * 2 + 1}"
) )
self.recv_list[i].setsockopt(zmq.SUBSCRIBE, b"") self.recv_list[i].setsockopt(zmq.SUBSCRIBE, b"")
# Synchronization channel to notify the manager when this client is ready. # Synchronization channel to notify the manager when this client is ready.
self.ready_sockets[i] = context.socket(zmq.PUSH) self.ready_sockets[i] = context.socket(zmq.PUSH)
self.ready_sockets[i].connect( self.ready_sockets[i].connect(
f"tcp://{all_ips[i * get_world_size() // get_parallel().nnodes]}:{PORT_BASE + i * 2}" f"tcp://{all_ips[i * get_parallel().world_size // get_parallel().nnodes]}:{PORT_BASE + i * 2}"
) )
sock_send(self.ready_sockets[i], UpdateExpertBackupReq()) sock_send(self.ready_sockets[i], UpdateExpertBackupReq())
python/sglang/srt/hardware_backend/musa/attention/flashattention_backend.py
+3 -3
View File
@@ -8,7 +8,7 @@ from flash_attn_interface import flash_attn_varlen_func
from flash_attn_interface import flash_attn_with_kvcache as mate_flash_attn_with_kvcache from flash_attn_interface import flash_attn_with_kvcache as mate_flash_attn_with_kvcache
from flash_attn_interface import get_scheduler_metadata from flash_attn_interface import get_scheduler_metadata
from sglang.srt.distributed import get_pp_group, get_pp_indices from sglang.srt.distributed import get_pp_indices
from sglang.srt.environ import envs from sglang.srt.environ import envs
from sglang.srt.hardware_backend.musa.layers.utils.cp_utils import ( from sglang.srt.hardware_backend.musa.layers.utils.cp_utils import (
musa_cp_attn_forward_extend as cp_attn_forward_extend, musa_cp_attn_forward_extend as cp_attn_forward_extend,
@@ -23,7 +23,7 @@ from sglang.srt.layers.utils.cp_utils import (
cp_allgather_and_save_kv_cache, cp_allgather_and_save_kv_cache,
) )
from sglang.srt.mem_cache.memory_pool import KVWriteLoc from sglang.srt.mem_cache.memory_pool import KVWriteLoc
from sglang.srt.runtime_context import get_schedule from sglang.srt.runtime_context import get_parallel, get_schedule
if TYPE_CHECKING: if TYPE_CHECKING:
from sglang.srt.layers.radix_attention import RadixAttention from sglang.srt.layers.radix_attention import RadixAttention
@@ -61,7 +61,7 @@ def _compute_scheduler_metadata(
# Determine if scheduler metadata should be updated # Determine if scheduler metadata should be updated
should_update = True should_update = True
pp_group = get_pp_group() pp_group = get_parallel().pp_group
pp_rank = pp_group.rank_in_group pp_rank = pp_group.rank_in_group
start_layer_id, _ = get_pp_indices( start_layer_id, _ = get_pp_indices(
backend.num_hidden_layers, pp_group.rank_in_group, pp_group.world_size backend.num_hidden_layers, pp_group.rank_in_group, pp_group.world_size
python/sglang/srt/hardware_backend/npu/moe/fuseep.py
+2 -3
View File
@@ -12,12 +12,11 @@ from typing import TYPE_CHECKING
import torch import torch
from sglang.srt.distributed import get_moe_ep_group
from sglang.srt.environ import envs from sglang.srt.environ import envs
from sglang.srt.hardware_backend.npu.utils import npu_format_cast from sglang.srt.hardware_backend.npu.utils import npu_format_cast
from sglang.srt.layers.moe.token_dispatcher.deepep import DeepEPBuffer from sglang.srt.layers.moe.token_dispatcher.deepep import DeepEPBuffer
from sglang.srt.layers.moe.utils import DeepEPMode from sglang.srt.layers.moe.utils import DeepEPMode
from sglang.srt.runtime_context import get_exec from sglang.srt.runtime_context import get_exec, get_parallel
if TYPE_CHECKING: if TYPE_CHECKING:
from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE
@@ -30,7 +29,7 @@ _PARAMS_BYTES = 2 # bf16 — Ascend's Dispatch & Combine does not support fp16
def _get_fuseep_buffer(layer: FusedMoE): def _get_fuseep_buffer(layer: FusedMoE):
DeepEPBuffer.set_dispatch_mode_as_low_latency() DeepEPBuffer.set_dispatch_mode_as_low_latency()
return DeepEPBuffer.get_deepep_buffer( return DeepEPBuffer.get_deepep_buffer(
get_moe_ep_group().device_group, get_parallel().moe_ep_group.device_group,
layer.hidden_size, layer.hidden_size,
_PARAMS_BYTES, _PARAMS_BYTES,
DeepEPMode.LOW_LATENCY, DeepEPMode.LOW_LATENCY,
python/sglang/srt/layers/attention/dsa/dsa_indexer.py
+4 -6
View File
@@ -112,10 +112,6 @@ if _is_xpu:
if _use_aiter: if _use_aiter:
from aiter.ops.cache import indexer_k_quant_and_cache from aiter.ops.cache import indexer_k_quant_and_cache
from sglang.srt.distributed import (
get_attn_tp_group,
)
from sglang.srt.distributed.parallel_state import get_pp_group
from sglang.srt.layers import deep_gemm_wrapper from sglang.srt.layers import deep_gemm_wrapper
from sglang.srt.layers.cp.base import get_cp_strategy from sglang.srt.layers.cp.base import get_cp_strategy
from sglang.srt.layers.cp.utils import is_cp_active from sglang.srt.layers.cp.utils import is_cp_active
@@ -157,7 +153,7 @@ if _is_cuda:
def _broadcast_indexer_topk_from_rank0_impl(topk_indices: torch.Tensor) -> None: def _broadcast_indexer_topk_from_rank0_impl(topk_indices: torch.Tensor) -> None:
group = get_attn_tp_group() group = get_parallel().attn_tp_group
if group.world_size == 1: if group.world_size == 1:
return return
@@ -260,7 +256,9 @@ class Indexer(DSANPUIndexerMixin, BaseFusedOp):
self.sm_count = deep_gemm.get_num_sms() self.sm_count = deep_gemm.get_num_sms()
self.half_device_sm_count = ceil_align(self.sm_count // 2, 8) self.half_device_sm_count = ceil_align(self.sm_count // 2, 8)
pp_size = get_parallel().pp_size pp_size = get_parallel().pp_size
self.logits_with_pp_recv = pp_size > 1 and not get_pp_group().is_last_rank self.logits_with_pp_recv = (
pp_size > 1 and not get_parallel().pp_group.is_last_rank
)
else: else:
self.logits_with_pp_recv = False self.logits_with_pp_recv = False
python/sglang/srt/layers/communicator.py
+8 -9
View File
@@ -23,7 +23,6 @@ import torch
from sglang.srt.distributed import ( from sglang.srt.distributed import (
attention_tensor_model_parallel_all_reduce, attention_tensor_model_parallel_all_reduce,
attention_tensor_model_parallel_quant_all_reduce, attention_tensor_model_parallel_quant_all_reduce,
get_tp_group,
tensor_model_parallel_all_reduce, tensor_model_parallel_all_reduce,
) )
from sglang.srt.distributed.device_communicators.pynccl_allocator import ( from sglang.srt.distributed.device_communicators.pynccl_allocator import (
@@ -266,7 +265,7 @@ class AttentionInputs:
def tp_all_gather_hidden_states(self, hidden_states, forward_batch): def tp_all_gather_hidden_states(self, hidden_states, forward_batch):
total_tokens = forward_batch.input_ids.shape[0] total_tokens = forward_batch.input_ids.shape[0]
output = hidden_states.new_empty((total_tokens, hidden_states.shape[-1])) output = hidden_states.new_empty((total_tokens, hidden_states.shape[-1]))
get_tp_group().all_gather_into_tensor(output, hidden_states) get_parallel().tp_group.all_gather_into_tensor(output, hidden_states)
return output return output
def fetch_qkv_latent(self): def fetch_qkv_latent(self):
@@ -510,7 +509,7 @@ def tp_reduce_scatter(
) )
local_tokens = hidden_states.shape[0] // context.tp_size local_tokens = hidden_states.shape[0] // context.tp_size
output = hidden_states.new_empty(local_tokens, *hidden_states.shape[1:]) output = hidden_states.new_empty(local_tokens, *hidden_states.shape[1:])
get_tp_group().reduce_scatter_tensor(output, hidden_states) get_parallel().tp_group.reduce_scatter_tensor(output, hidden_states)
if residual is not None: if residual is not None:
residual = residual.tensor_split(context.tp_size)[context.tp_rank] residual = residual.tensor_split(context.tp_size)[context.tp_rank]
return output, residual return output, residual
@@ -1075,7 +1074,7 @@ class CommunicateSimpleFn:
gathered_hidden_states = [] gathered_hidden_states = []
for local_hidden_states in hidden_states: for local_hidden_states in hidden_states:
with use_symmetric_memory( with use_symmetric_memory(
get_tp_group(), get_parallel().tp_group,
disabled=not is_allocation_symmetric(), disabled=not is_allocation_symmetric(),
): ):
output = torch.empty( output = torch.empty(
@@ -1270,7 +1269,7 @@ class CommunicateWithAllReduceAndLayerNormFn:
hidden_states hidden_states
) )
with use_symmetric_memory( with use_symmetric_memory(
get_tp_group(), get_parallel().tp_group,
disabled=not is_allocation_symmetric(), disabled=not is_allocation_symmetric(),
): ):
hidden_states, residual = layernorm(hidden_states, residual) hidden_states, residual = layernorm(hidden_states, residual)
@@ -1278,7 +1277,7 @@ class CommunicateWithAllReduceAndLayerNormFn:
hidden_states += residual hidden_states += residual
hidden_states, local_hidden_states = ( hidden_states, local_hidden_states = (
get_global_dp_buffer(get_tp_group()), get_global_dp_buffer(get_parallel().tp_group),
hidden_states, hidden_states,
) )
if use_layer_norm_before_gather: if use_layer_norm_before_gather:
@@ -1510,7 +1509,7 @@ class CommunicateSummableTensorPairFn:
allow_reduce_scatter: bool = False, allow_reduce_scatter: bool = False,
): ):
if get_parallel().tp_size == get_parallel().attn_dp_size: if get_parallel().tp_size == get_parallel().attn_dp_size:
group = get_tp_group() group = get_parallel().tp_group
else: else:
group = get_parallel().attn_tp_group group = get_parallel().attn_tp_group
hidden_states, global_hidden_states = ( hidden_states, global_hidden_states = (
@@ -1518,7 +1517,7 @@ class CommunicateSummableTensorPairFn:
hidden_states, hidden_states,
) )
if should_use_dp_reduce_scatterv(): if should_use_dp_reduce_scatterv():
get_tp_group().reduce_scatterv( get_parallel().tp_group.reduce_scatterv(
global_hidden_states, global_hidden_states,
output=hidden_states, output=hidden_states,
sizes=get_dp_global_num_tokens(), sizes=get_dp_global_num_tokens(),
@@ -1600,7 +1599,7 @@ class CommunicateSummableTensorPairFn:
# DP scatter (if DP attention is enabled) # DP scatter (if DP attention is enabled)
if context.attn_dp_size > 1: if context.attn_dp_size > 1:
if get_parallel().tp_size == get_parallel().attn_dp_size: if get_parallel().tp_size == get_parallel().attn_dp_size:
group = get_tp_group() group = get_parallel().tp_group
else: else:
group = get_parallel().attn_tp_group group = get_parallel().attn_tp_group
hidden_states_output, global_hidden_states = ( hidden_states_output, global_hidden_states = (
python/sglang/srt/layers/communicator_mhc.py
+9 -9
View File
@@ -18,7 +18,6 @@ from typing import Callable, Optional
import torch import torch
from sglang.kernels.ops.layernorm.mhc import hc_contract, hc_expand from sglang.kernels.ops.layernorm.mhc import hc_contract, hc_expand
from sglang.srt.distributed import get_tp_group
from sglang.srt.distributed.communication_op import ( from sglang.srt.distributed.communication_op import (
attention_tensor_model_parallel_all_reduce, attention_tensor_model_parallel_all_reduce,
) )
@@ -50,6 +49,7 @@ from sglang.srt.layers.dp_attention import (
) )
from sglang.srt.layers.moe import should_use_dp_reduce_scatterv from sglang.srt.layers.moe import should_use_dp_reduce_scatterv
from sglang.srt.model_executor.forward_batch_info import ForwardBatch from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.runtime_context import get_parallel
def tp_all_gather_hidden_states(hidden_states, forward_batch): def tp_all_gather_hidden_states(hidden_states, forward_batch):
@@ -58,7 +58,7 @@ def tp_all_gather_hidden_states(hidden_states, forward_batch):
) )
total_tokens = forward_batch.input_ids.shape[0] total_tokens = forward_batch.input_ids.shape[0]
output = hidden_states.new_empty((total_tokens, hidden_states.shape[-1])) output = hidden_states.new_empty((total_tokens, hidden_states.shape[-1]))
get_tp_group().all_gather_into_tensor(output, hidden_states) get_parallel().tp_group.all_gather_into_tensor(output, hidden_states)
return output return output
@@ -199,7 +199,7 @@ class MHCCommunicateWithAllReduceAndLayerNormFn(CommunicateWithAllReduceAndLayer
return hidden_states, hidden_states return hidden_states, hidden_states
scatter_states = hidden_states.tensor_split(context.tp_size)[context.tp_rank] scatter_states = hidden_states.tensor_split(context.tp_size)[context.tp_rank]
get_tp_group().reduce_scatter_tensor(scatter_states, hidden_states) get_parallel().tp_group.reduce_scatter_tensor(scatter_states, hidden_states)
scatter_states, residual = mhc.attn_to_mlp( scatter_states, residual = mhc.attn_to_mlp(
scatter_states, residual, out_norm=layernorm scatter_states, residual, out_norm=layernorm
@@ -238,7 +238,7 @@ class MHCCommunicateWithAllReduceAndLayerNormFn(CommunicateWithAllReduceAndLayer
if context.attn_dp_size != 1: if context.attn_dp_size != 1:
if hidden_states.shape[0] != 0: if hidden_states.shape[0] != 0:
with use_symmetric_memory( with use_symmetric_memory(
get_tp_group(), get_parallel().tp_group,
disabled=not is_allocation_symmetric(), disabled=not is_allocation_symmetric(),
): ):
hidden_states, residual = mhc.attn_to_mlp( hidden_states, residual = mhc.attn_to_mlp(
@@ -248,7 +248,7 @@ class MHCCommunicateWithAllReduceAndLayerNormFn(CommunicateWithAllReduceAndLayer
hidden_states, residual = mhc.attn_to_mlp(hidden_states, residual) hidden_states, residual = mhc.attn_to_mlp(hidden_states, residual)
hidden_states, local_hidden_states = ( hidden_states, local_hidden_states = (
get_global_dp_buffer(get_tp_group()), get_global_dp_buffer(get_parallel().tp_group),
hidden_states, hidden_states,
) )
dp_gather_replicate(hidden_states, local_hidden_states, forward_batch) dp_gather_replicate(hidden_states, local_hidden_states, forward_batch)
@@ -305,7 +305,7 @@ class MHCCommunicateSummableTensorPairFn(CommunicateSummableTensorPairFn):
hidden_states = local_states.new_empty( hidden_states = local_states.new_empty(
local_states.shape[0] * context.tp_size, *local_states.shape[1:] local_states.shape[0] * context.tp_size, *local_states.shape[1:]
) )
get_tp_group().all_gather_into_tensor(hidden_states, local_states) get_parallel().tp_group.all_gather_into_tensor(hidden_states, local_states)
return hidden_states, None return hidden_states, None
@@ -322,13 +322,13 @@ class MHCCommunicateSummableTensorPairFn(CommunicateSummableTensorPairFn):
**kwargs, **kwargs,
): ):
hidden_states, global_hidden_states = ( hidden_states, global_hidden_states = (
get_local_dp_buffer_mhc(get_tp_group(), 1), get_local_dp_buffer_mhc(get_parallel().tp_group, 1),
hidden_states, hidden_states,
) )
# MoE skips its post-expert all-reduce with reduce_scatterv, so this # MoE skips its post-expert all-reduce with reduce_scatterv, so this
# scatter must reduce while combining local-expert partial sums. # scatter must reduce while combining local-expert partial sums.
if should_use_dp_reduce_scatterv(): if should_use_dp_reduce_scatterv():
get_tp_group().reduce_scatterv( get_parallel().tp_group.reduce_scatterv(
global_hidden_states, global_hidden_states,
output=hidden_states, output=hidden_states,
sizes=get_dp_global_num_tokens(), sizes=get_dp_global_num_tokens(),
@@ -362,7 +362,7 @@ class MHCCommunicateSummableTensorPairFn(CommunicateSummableTensorPairFn):
hidden_states, local_hidden_states = ( hidden_states, local_hidden_states = (
get_local_dp_buffer_mhc( get_local_dp_buffer_mhc(
get_tp_group(), 1 if is_last_layer else mhc.hc_mult get_parallel().tp_group, 1 if is_last_layer else mhc.hc_mult
), ),
hidden_states, hidden_states,
) )
python/sglang/srt/layers/dp_attention.py
+34 -29
View File
@@ -15,16 +15,11 @@ from sglang.srt.arg_groups.model_override_base import (
) )
from sglang.srt.distributed import ( from sglang.srt.distributed import (
GroupCoordinator, GroupCoordinator,
get_attn_cp_group,
get_attn_tensor_model_parallel_rank,
get_attn_tensor_model_parallel_world_size, get_attn_tensor_model_parallel_world_size,
get_attn_tp_group,
) )
from sglang.srt.distributed import get_moe_dp_group as _get_moe_dp_group from sglang.srt.distributed import get_moe_dp_group as _get_moe_dp_group
from sglang.srt.distributed import ( from sglang.srt.distributed import (
get_tensor_model_parallel_rank,
get_tensor_model_parallel_world_size, get_tensor_model_parallel_world_size,
get_tp_group,
tensor_model_parallel_all_reduce, tensor_model_parallel_all_reduce,
) )
from sglang.srt.distributed.device_communicators.pynccl_allocator import ( from sglang.srt.distributed.device_communicators.pynccl_allocator import (
@@ -380,7 +375,7 @@ def initialize_dp_attention(
dp.enabled = enable_dp_attention dp.enabled = enable_dp_attention
tp_rank = get_tensor_model_parallel_rank() tp_rank = get_parallel().tp_rank
tp_size = get_tensor_model_parallel_world_size() tp_size = get_tensor_model_parallel_world_size()
_, _, attn_dp_rank, attn_dp_size = compute_dp_attention_world_info( _, _, attn_dp_rank, attn_dp_size = compute_dp_attention_world_info(
@@ -503,9 +498,7 @@ def _dp_gather_via_all_reduce(
assert local_tokens.is_contiguous() assert local_tokens.is_contiguous()
assert global_tokens.is_contiguous() assert global_tokens.is_contiguous()
if local_tokens.shape[0] > 0 and ( if local_tokens.shape[0] > 0 and (is_partial or get_parallel().attn_tp_rank == 0):
is_partial or get_attn_tensor_model_parallel_rank() == 0
):
assert local_tokens.untyped_storage() is not global_tokens.untyped_storage(), ( assert local_tokens.untyped_storage() is not global_tokens.untyped_storage(), (
"aliasing between global_tokens and local_tokens not allowed" "aliasing between global_tokens and local_tokens not allowed"
) )
@@ -527,7 +520,9 @@ def _dp_gather_via_all_reduce(
): ):
from sglang.srt.distributed.parallel_state import inplace_all_reduce from sglang.srt.distributed.parallel_state import inplace_all_reduce
inplace_all_reduce(global_tokens, group_name=get_tp_group().unique_name) inplace_all_reduce(
global_tokens, group_name=get_parallel().tp_group.unique_name
)
else: else:
global_tokens[:] = tensor_model_parallel_all_reduce(global_tokens) global_tokens[:] = tensor_model_parallel_all_reduce(global_tokens)
@@ -549,16 +544,18 @@ def _dp_gather_via_all_gather(
group=torch.distributed.group.WORLD, group=torch.distributed.group.WORLD,
) )
else: else:
get_tp_group().all_gather_into_tensor(global_tokens, local_tokens) get_parallel().tp_group.all_gather_into_tensor(global_tokens, local_tokens)
return return
if not is_partial: if not is_partial:
if get_attn_tensor_model_parallel_rank() != 0: if get_parallel().attn_tp_rank != 0:
local_tokens.fill_(0) local_tokens.fill_(0)
scattered_local_tokens = local_tokens.tensor_split( scattered_local_tokens = local_tokens.tensor_split(
get_attn_tensor_model_parallel_world_size() get_attn_tensor_model_parallel_world_size()
)[get_attn_tensor_model_parallel_rank()] )[get_parallel().attn_tp_rank]
get_attn_tp_group().reduce_scatter_tensor(scattered_local_tokens, local_tokens) get_parallel().attn_tp_group.reduce_scatter_tensor(
scattered_local_tokens, local_tokens
)
if use_world: if use_world:
torch.distributed.all_gather_into_tensor( torch.distributed.all_gather_into_tensor(
global_tokens, global_tokens,
@@ -566,7 +563,9 @@ def _dp_gather_via_all_gather(
group=torch.distributed.group.WORLD, group=torch.distributed.group.WORLD,
) )
else: else:
get_tp_group().all_gather_into_tensor(global_tokens, scattered_local_tokens) get_parallel().tp_group.all_gather_into_tensor(
global_tokens, scattered_local_tokens
)
# Variable-length DP-MoE gather (reference https://github.com/ROCm/ATOM/pull/930): instead of padding every # Variable-length DP-MoE gather (reference https://github.com/ROCm/ATOM/pull/930): instead of padding every
@@ -696,7 +695,7 @@ def _dp_gather_via_all_gatherv_fp8(
local_real.contiguous(), _DP_GATHER_FP8_GROUP local_real.contiguous(), _DP_GATHER_FP8_GROUP
) )
gq, gs = _get_dp_gather_fp8_bufs(rows, hidden, global_tokens.device) gq, gs = _get_dp_gather_fp8_bufs(rows, hidden, global_tokens.device)
tp_group = get_tp_group() tp_group = get_parallel().tp_group
tp_group.all_gatherv(q.view(torch.uint8), sizes=sizes, output=gq) tp_group.all_gatherv(q.view(torch.uint8), sizes=sizes, output=gq)
tp_group.all_gatherv(s, sizes=sizes, output=gs) tp_group.all_gatherv(s, sizes=sizes, output=gs)
_dequant_per_token_group_fp8_kernel[(rows,)]( _dequant_per_token_group_fp8_kernel[(rows,)](
@@ -785,7 +784,7 @@ def _dp_gather_via_all_gatherv(
): ):
_dp_gather_via_all_gatherv_fp8(global_tokens, local_real, sizes) _dp_gather_via_all_gatherv_fp8(global_tokens, local_real, sizes)
return return
get_tp_group().all_gatherv(local_real, sizes=sizes, output=global_tokens) get_parallel().tp_group.all_gatherv(local_real, sizes=sizes, output=global_tokens)
def _note_dp_gather_in_prefill_graph() -> None: def _note_dp_gather_in_prefill_graph() -> None:
@@ -882,16 +881,18 @@ def dp_reduce_scatter_tensor(output: torch.Tensor, input: torch.Tensor):
# the default reduce-scatter path if per-rank sizes are unavailable. # the default reduce-scatter path if per-rank sizes are unavailable.
sizes = get_dp_global_num_tokens() sizes = get_dp_global_num_tokens()
if sizes is not None: if sizes is not None:
get_tp_group().reduce_scatterv(input, output=output, sizes=sizes) get_parallel().tp_group.reduce_scatterv(input, output=output, sizes=sizes)
return return
if get_tensor_model_parallel_world_size() == get_parallel().attn_dp_size: if get_tensor_model_parallel_world_size() == get_parallel().attn_dp_size:
get_tp_group().reduce_scatter_tensor(output, input) get_parallel().tp_group.reduce_scatter_tensor(output, input)
else: else:
scattered_local_tokens = input.tensor_split( scattered_local_tokens = input.tensor_split(
get_tensor_model_parallel_world_size() get_tensor_model_parallel_world_size()
)[get_tensor_model_parallel_rank()] )[get_parallel().tp_rank]
get_tp_group().reduce_scatter_tensor(scattered_local_tokens, input) get_parallel().tp_group.reduce_scatter_tensor(scattered_local_tokens, input)
get_attn_tp_group().all_gather_into_tensor(output, scattered_local_tokens) get_parallel().attn_tp_group.all_gather_into_tensor(
output, scattered_local_tokens
)
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
@@ -989,29 +990,31 @@ def dp_reduce_scatterv_async(
ev = _tbo_event(event_key) ev = _tbo_event(event_key)
with torch.cuda.stream(comm): with torch.cuda.stream(comm):
comm.wait_stream(compute) comm.wait_stream(compute)
get_tp_group().reduce_scatterv(global_tokens, output=output_local, sizes=sizes) get_parallel().tp_group.reduce_scatterv(
global_tokens, output=output_local, sizes=sizes
)
ev.record(comm) ev.record(comm)
return ev return ev
def attn_tp_reduce_scatter_tensor(output: torch.Tensor, input: torch.Tensor): def attn_tp_reduce_scatter_tensor(output: torch.Tensor, input: torch.Tensor):
return get_attn_tp_group().reduce_scatter_tensor(output, input) return get_parallel().attn_tp_group.reduce_scatter_tensor(output, input)
def attn_cp_reduce_scatter_tensor(output: torch.Tensor, input: torch.Tensor): def attn_cp_reduce_scatter_tensor(output: torch.Tensor, input: torch.Tensor):
return get_attn_cp_group().reduce_scatter_tensor(output, input) return get_parallel().attn_cp_group.reduce_scatter_tensor(output, input)
def attn_tp_all_reduce(input: torch.Tensor): def attn_tp_all_reduce(input: torch.Tensor):
return get_attn_tp_group().all_reduce(input) return get_parallel().attn_tp_group.all_reduce(input)
def attn_tp_all_gather_into_tensor(output: torch.Tensor, input: torch.Tensor): def attn_tp_all_gather_into_tensor(output: torch.Tensor, input: torch.Tensor):
return get_attn_tp_group().all_gather_into_tensor(output, input) return get_parallel().attn_tp_group.all_gather_into_tensor(output, input)
def attn_cp_all_gather_into_tensor(output: torch.Tensor, input: torch.Tensor): def attn_cp_all_gather_into_tensor(output: torch.Tensor, input: torch.Tensor):
return get_attn_cp_group().all_gather_into_tensor(output, input) return get_parallel().attn_cp_group.all_gather_into_tensor(output, input)
def get_moe_cp_group() -> GroupCoordinator: def get_moe_cp_group() -> GroupCoordinator:
@@ -1041,4 +1044,6 @@ def moe_cp_all_gather_into_tensor(output: torch.Tensor, input: torch.Tensor):
def attn_tp_all_gather(output_list: List[torch.Tensor], input: torch.Tensor): def attn_tp_all_gather(output_list: List[torch.Tensor], input: torch.Tensor):
return get_attn_tp_group().all_gather(input, output_tensor_list=output_list) return get_parallel().attn_tp_group.all_gather(
input, output_tensor_list=output_list
)
python/sglang/srt/layers/engram.py
+1 -2
View File
@@ -33,7 +33,6 @@ from sglang.kernels.ops.embeddings.engram_hash import (
engram_hash_ids_and_commit, engram_hash_ids_and_commit,
) )
from sglang.srt.distributed import tensor_model_parallel_all_reduce from sglang.srt.distributed import tensor_model_parallel_all_reduce
from sglang.srt.distributed.parallel_state import get_tp_group
from sglang.srt.environ import envs from sglang.srt.environ import envs
from sglang.srt.layers.dp_attention import ( from sglang.srt.layers.dp_attention import (
attn_cp_all_gather_into_tensor, attn_cp_all_gather_into_tensor,
@@ -708,7 +707,7 @@ class EngramEmbedding(nn.Module):
layout, layout,
max(1, w_bytes + s_bytes), # mmap requires storage even for an empty shard. max(1, w_bytes + s_bytes), # mmap requires storage even for an empty shard.
f"sglang_engram_{layer_id}", f"sglang_engram_{layer_id}",
get_tp_group(), get_parallel().tp_group,
) )
raw = self.host_table.bytes[: w_bytes + s_bytes] raw = self.host_table.bytes[: w_bytes + s_bytes]
weight = raw[:w_bytes].view(torch.float8_e4m3fn).view(n, dim) weight = raw[:w_bytes].view(torch.float8_e4m3fn).view(n, dim)
python/sglang/srt/layers/flashinfer_comm_fusion.py
+6 -12
View File
@@ -6,12 +6,6 @@ import torch
import torch.distributed as dist import torch.distributed as dist
from torch.distributed import ProcessGroup from torch.distributed import ProcessGroup
from sglang.srt.distributed import (
get_attn_tp_group,
get_moe_ep_group,
get_moe_tp_group,
get_tp_group,
)
from sglang.srt.distributed.parallel_state import in_the_same_node_as from sglang.srt.distributed.parallel_state import in_the_same_node_as
from sglang.srt.runtime_context import ( from sglang.srt.runtime_context import (
get_exec, get_exec,
@@ -335,7 +329,7 @@ def _preflight_check_workspace_memory(
group = cpu_group group = cpu_group
if group is None: if group is None:
tp_group = get_tp_group() tp_group = get_parallel().tp_group
if tp_group.world_size <= 1: if tp_group.world_size <= 1:
return True return True
group = tp_group.cpu_group group = tp_group.cpu_group
@@ -670,12 +664,12 @@ def resolve_fusion_group(*, use_attn_tp_group: bool):
parallel = get_parallel() parallel = get_parallel()
if use_attn_tp_group: if use_attn_tp_group:
return parallel.attn_tp_size, parallel.attn_tp_rank, get_attn_tp_group() return parallel.attn_tp_size, parallel.attn_tp_rank, parallel.attn_tp_group
if can_merge_post_experts_all_reduce(): if can_merge_post_experts_all_reduce():
return parallel.tp_size, parallel.tp_rank, get_tp_group() return parallel.tp_size, parallel.tp_rank, parallel.tp_group
if parallel.moe_ep_size > 1: if parallel.moe_ep_size > 1:
return parallel.moe_ep_size, parallel.moe_ep_rank, get_moe_ep_group() return parallel.moe_ep_size, parallel.moe_ep_rank, parallel.moe_ep_group
return parallel.moe_tp_size, parallel.moe_tp_rank, get_moe_tp_group() return parallel.moe_tp_size, parallel.moe_tp_rank, parallel.moe_tp_group
def _sync_allreduce_unavailable_across_tp(): def _sync_allreduce_unavailable_across_tp():
@@ -691,7 +685,7 @@ def _sync_allreduce_unavailable_across_tp():
try: try:
import torch.distributed as dist import torch.distributed as dist
tp_group = get_tp_group() tp_group = get_parallel().tp_group
if tp_group.world_size <= 1: if tp_group.world_size <= 1:
return return
flag = torch.tensor( flag = torch.tensor(
python/sglang/srt/layers/flashinfer_mnnvl_cutedsl.py
+2 -2
View File
@@ -308,10 +308,10 @@ def get_flashinfer_mnnvl_cutedsl_ar_fusion(
assert max_m is not None assert max_m is not None
assert rms_epsilon is not None assert rms_epsilon is not None
assert weight_bias is not None assert weight_bias is not None
from sglang.srt.distributed.parallel_state import get_tp_group from sglang.srt.runtime_context import get_parallel
device = torch.device("cuda", torch.cuda.current_device()) device = torch.device("cuda", torch.cuda.current_device())
process_group = get_tp_group().device_group process_group = get_parallel().tp_group.device_group
domain = ( domain = (
int(hidden_size), int(hidden_size),
int(top_k), int(top_k),
python/sglang/srt/layers/k3_ar_fusion.py
+4 -5
View File
@@ -25,7 +25,6 @@ import torch
import sglang.srt.runtime_context as ctx import sglang.srt.runtime_context as ctx
from sglang.kernels.jit.utils import cache_once from sglang.kernels.jit.utils import cache_once
from sglang.srt.environ import envs from sglang.srt.environ import envs
from sglang.srt.runtime_context import get_parallel
if TYPE_CHECKING: if TYPE_CHECKING:
from sglang.srt.distributed.device_communicators.custom_all_reduce_v2 import ( from sglang.srt.distributed.device_communicators.custom_all_reduce_v2 import (
@@ -84,11 +83,11 @@ def _get_state() -> Optional[_State]:
from sglang.srt.distributed.device_communicators.custom_all_reduce_v2 import ( from sglang.srt.distributed.device_communicators.custom_all_reduce_v2 import (
CustomAllReduceV2, CustomAllReduceV2,
) )
from sglang.srt.distributed.parallel_state import get_tp_group from sglang.srt.runtime_context import get_parallel
if get_parallel().tp_size <= 1: if get_parallel().tp_size <= 1:
return None return None
group = get_tp_group() group = get_parallel().tp_group
comm = group.ca_comm comm = group.ca_comm
if ( if (
not isinstance(comm, CustomAllReduceV2) not isinstance(comm, CustomAllReduceV2)
@@ -198,9 +197,9 @@ def symm_buffer(
its own. Each name belongs to one group, so the name alone identifies it. its own. Each name belongs to one group, so the name alone identifies it.
""" """
if group_name is None: if group_name is None:
from sglang.srt.distributed.parallel_state import get_tp_group from sglang.srt.runtime_context import get_parallel
group_name = get_tp_group().cpu_group.group_name group_name = get_parallel().tp_group.cpu_group.group_name
buf: _Buffer = ctx.get_buffer( buf: _Buffer = ctx.get_buffer(
f"k3_symm:{name}", lambda: _create_buffer(name, width, dtype, group_name) f"k3_symm:{name}", lambda: _create_buffer(name, width, dtype, group_name)
) )
python/sglang/srt/layers/k3_gemm_ar.py
+2 -2
View File
@@ -56,7 +56,7 @@ def maybe_wrap_o_proj(o_proj: RowParallelLinear) -> None:
if not _init(): if not _init():
return return
from sglang.kernels.ops.kimi_k3 import gemm_ar as mod from sglang.kernels.ops.kimi_k3 import gemm_ar as mod
from sglang.srt.distributed.parallel_state import get_tp_group from sglang.srt.runtime_context import get_parallel
parallel = get_parallel() parallel = get_parallel()
world_size = parallel.tp_size world_size = parallel.tp_size
@@ -73,7 +73,7 @@ def maybe_wrap_o_proj(o_proj: RowParallelLinear) -> None:
mod.init( mod.init(
world_size=world_size, world_size=world_size,
rank=parallel.tp_rank, rank=parallel.tp_rank,
group=get_tp_group().cpu_group, group=get_parallel().tp_group.cpu_group,
k=weight.shape[1], k=weight.shape[1],
) )
# per-K compile + base-address stash, pre-capture # per-K compile + base-address stash, pre-capture
python/sglang/srt/layers/layernorm_sp.py
+5 -6
View File
@@ -39,7 +39,6 @@ from typing import Optional
import torch import torch
from sglang.srt.distributed import get_tp_group
from sglang.srt.runtime_context import ( from sglang.srt.runtime_context import (
get_flags, get_flags,
get_forward, get_forward,
@@ -112,7 +111,7 @@ def sp_entry_scatter(hidden_states: torch.Tensor) -> torch.Tensor:
""" """
num_tokens = hidden_states.shape[0] num_tokens = hidden_states.shape[0]
set_sp_num_tokens(num_tokens) set_sp_num_tokens(num_tokens)
tp_group = get_tp_group() tp_group = get_parallel().tp_group
tp_size = tp_group.world_size tp_size = tp_group.world_size
if tp_size == 1: if tp_size == 1:
return hidden_states return hidden_states
@@ -127,7 +126,7 @@ def sp_entry_scatter(hidden_states: torch.Tensor) -> torch.Tensor:
def sp_exit_gather(hidden_states: torch.Tensor, num_tokens: int) -> torch.Tensor: def sp_exit_gather(hidden_states: torch.Tensor, num_tokens: int) -> torch.Tensor:
"""g: all-gather the per-rank shards back to the full sequence along dim 0, """g: all-gather the per-rank shards back to the full sequence along dim 0,
then narrow to ``num_tokens`` (dropping the entry-scatter padding).""" then narrow to ``num_tokens`` (dropping the entry-scatter padding)."""
tp_group = get_tp_group() tp_group = get_parallel().tp_group
tp_size = tp_group.world_size tp_size = tp_group.world_size
if tp_size == 1: if tp_size == 1:
return hidden_states[:num_tokens] return hidden_states[:num_tokens]
@@ -207,7 +206,7 @@ def column_parallel_g_matmul(
""" """
num_tokens = sp_num_tokens() num_tokens = sp_num_tokens()
if sp_fused_matmul_eligible(linear): if sp_fused_matmul_eligible(linear):
group_name = get_tp_group().device_group.group_name group_name = get_parallel().tp_group.device_group.group_name
_, mm_outputs = torch.ops.symm_mem.fused_all_gather_matmul( _, mm_outputs = torch.ops.symm_mem.fused_all_gather_matmul(
input_parallel.contiguous(), input_parallel.contiguous(),
[linear.weight.t()], [linear.weight.t()],
@@ -235,7 +234,7 @@ def row_parallel_gbar_matmul(linear, input_: torch.Tensor, bias) -> torch.Tensor
if padded != num_tokens: if padded != num_tokens:
x = torch.nn.functional.pad(x, (0, 0, 0, padded - num_tokens)) x = torch.nn.functional.pad(x, (0, 0, 0, padded - num_tokens))
if sp_fused_matmul_eligible(linear): if sp_fused_matmul_eligible(linear):
group_name = get_tp_group().device_group.group_name group_name = get_parallel().tp_group.device_group.group_name
return torch.ops.symm_mem.fused_matmul_reduce_scatter( return torch.ops.symm_mem.fused_matmul_reduce_scatter(
x, x,
linear.weight.t(), linear.weight.t(),
@@ -245,5 +244,5 @@ def row_parallel_gbar_matmul(linear, input_: torch.Tensor, bias) -> torch.Tensor
) )
full = linear.quant_method.apply(linear, x, bias) full = linear.quant_method.apply(linear, x, bias)
output = full.new_empty((padded // tp_size, *full.shape[1:])) output = full.new_empty((padded // tp_size, *full.shape[1:]))
get_tp_group().reduce_scatter_tensor(output, full) get_parallel().tp_group.reduce_scatter_tensor(output, full)
return output return output
python/sglang/srt/layers/linear.py
+1 -2
View File
@@ -15,7 +15,6 @@ from torch.nn.parameter import Parameter, UninitializedParameter
from sglang.kernels.kernel_api_logging import wrap_method_with_debug_kernel_once from sglang.kernels.kernel_api_logging import wrap_method_with_debug_kernel_once
from sglang.srt.distributed import ( from sglang.srt.distributed import (
divide, divide,
get_tp_group,
split_tensor_along_last_dim, split_tensor_along_last_dim,
tensor_model_parallel_all_gather, tensor_model_parallel_all_gather,
tensor_model_parallel_all_reduce, tensor_model_parallel_all_reduce,
@@ -1648,7 +1647,7 @@ class RowParallelLinear(LinearBase):
symm_ctx = use_symmetric_memory(get_parallel().attn_tp_group) symm_ctx = use_symmetric_memory(get_parallel().attn_tp_group)
else: else:
symm_ctx = use_symmetric_memory( symm_ctx = use_symmetric_memory(
get_tp_group(), disabled=not is_allocation_symmetric() get_parallel().tp_group, disabled=not is_allocation_symmetric()
) )
with symm_ctx: with symm_ctx:
if output_tensor is None: if output_tensor is None:
python/sglang/srt/layers/logits_processor.py
+9 -3
View File
@@ -26,7 +26,6 @@ from sglang.kernels.ops.activation.softcap import (
softcap_inplace_logits as fused_softcap, softcap_inplace_logits as fused_softcap,
) )
from sglang.srt.beam_search.logits_capture import BeamLogitsCapture from sglang.srt.beam_search.logits_capture import BeamLogitsCapture
from sglang.srt.distributed import get_attn_tp_group, get_tp_group
from sglang.srt.distributed.device_communicators import triton_symm_mem_ag from sglang.srt.distributed.device_communicators import triton_symm_mem_ag
from sglang.srt.environ import envs from sglang.srt.environ import envs
from sglang.srt.layers import layernorm_sp from sglang.srt.layers import layernorm_sp
@@ -463,7 +462,12 @@ class LogitsProcessor(nn.Module):
self.do_tensor_parallel_all_gather self.do_tensor_parallel_all_gather
and not self.do_tensor_parallel_all_gather_dp_attn and not self.do_tensor_parallel_all_gather_dp_attn
): ):
group = get_attn_tp_group() if self.use_attn_tp_group else get_tp_group() parallel = get_parallel()
group = (
parallel.attn_tp_group
if self.use_attn_tp_group
else parallel.tp_group
)
chunking_group = group.cpu_group chunking_group = group.cpu_group
self.input_logprob_processor = InputLogprobProcessor( self.input_logprob_processor = InputLogprobProcessor(
self.vocab_size, chunking_group=chunking_group self.vocab_size, chunking_group=chunking_group
@@ -1061,7 +1065,9 @@ class LogitsProcessor(nn.Module):
"""Exchange only the row block owned by each destination DP rank.""" """Exchange only the row block owned by each destination DP rank."""
logits = logits.contiguous() logits = logits.contiguous()
all_to_all_output = torch.empty_like(logits) all_to_all_output = torch.empty_like(logits)
get_tp_group().all_to_all_single(all_to_all_output.view(-1), logits.view(-1)) get_parallel().tp_group.all_to_all_single(
all_to_all_output.view(-1), logits.view(-1)
)
return _reassemble_tp_lm_head_all_to_all_output( return _reassemble_tp_lm_head_all_to_all_output(
all_to_all_output, get_parallel().tp_size all_to_all_output, get_parallel().tp_size
) )
python/sglang/srt/layers/moe/fused_moe_triton/layer.py
+6 -8
View File
@@ -14,8 +14,6 @@ from sglang.srt.batch_overlap.single_batch_overlap import DownGemmOverlapArgs
from sglang.srt.batch_overlap.two_batch_overlap import MaybeTboDeepEPDispatcher from sglang.srt.batch_overlap.two_batch_overlap import MaybeTboDeepEPDispatcher
from sglang.srt.configs.moe_model_registry import model_requires_fp32_silu_mul from sglang.srt.configs.moe_model_registry import model_requires_fp32_silu_mul
from sglang.srt.distributed import ( from sglang.srt.distributed import (
get_moe_ep_group,
get_tp_group,
tensor_model_parallel_all_reduce, tensor_model_parallel_all_reduce,
) )
from sglang.srt.distributed.device_communicators.pynccl_allocator import ( from sglang.srt.distributed.device_communicators.pynccl_allocator import (
@@ -150,13 +148,13 @@ def _maybe_copy_weight_view_before_h2d(
def _get_deepep_comm_group(a2a_backend): def _get_deepep_comm_group(a2a_backend):
group = get_tp_group().device_group group = get_parallel().tp_group.device_group
if a2a_backend.is_mori(): if a2a_backend.is_mori():
group = get_tp_group() group = get_parallel().tp_group
elif _is_npu: elif _is_npu:
group = get_moe_ep_group().device_group group = get_parallel().moe_ep_group.device_group
return group return group
@@ -204,7 +202,7 @@ def create_moe_dispatcher(
_deepep_v2_experts_are_fp8(quant_method) _deepep_v2_experts_are_fp8(quant_method)
) )
return DeepEPv2Dispatcher( return DeepEPv2Dispatcher(
group=get_tp_group().device_group, group=get_parallel().tp_group.device_group,
router_topk=moe_runner_config.top_k, router_topk=moe_runner_config.top_k,
num_experts=moe_runner_config.num_experts, num_experts=moe_runner_config.num_experts,
num_local_experts=moe_runner_config.num_local_experts, num_local_experts=moe_runner_config.num_local_experts,
@@ -219,7 +217,7 @@ def create_moe_dispatcher(
) )
elif a2a_backend.is_flashinfer(): elif a2a_backend.is_flashinfer():
return FlashinferDispatcher( return FlashinferDispatcher(
group=get_tp_group().device_group, group=get_parallel().tp_group.device_group,
router_topk=moe_runner_config.top_k, router_topk=moe_runner_config.top_k,
num_experts=moe_runner_config.num_experts, num_experts=moe_runner_config.num_experts,
num_local_experts=moe_runner_config.num_local_experts, num_local_experts=moe_runner_config.num_local_experts,
@@ -1583,7 +1581,7 @@ class FusedMoE(torch.nn.Module):
dwdp_mgr.record_compute_and_prefetch_next(self.layer_id) dwdp_mgr.record_compute_and_prefetch_next(self.layer_id)
with use_symmetric_memory( with use_symmetric_memory(
get_tp_group(), disabled=not is_allocation_symmetric() get_parallel().tp_group, disabled=not is_allocation_symmetric()
): ):
final_hidden_states = self.dispatcher.combine(combine_input=combine_input) final_hidden_states = self.dispatcher.combine(combine_input=combine_input)
python/sglang/srt/layers/moe/mega_moe.py
+2 -2
View File
@@ -190,7 +190,7 @@ def _run_mega_routed(
) -> torch.Tensor: ) -> torch.Tensor:
import deep_gemm import deep_gemm
from sglang.srt.distributed.parallel_state import get_moe_ep_group from sglang.srt.runtime_context import get_parallel
hidden_size = moe.config.hidden_size hidden_size = moe.config.hidden_size
@@ -216,7 +216,7 @@ def _run_mega_routed(
topk_ids = None topk_ids = None
topk_weights = None topk_weights = None
ep_group = get_moe_ep_group().device_group ep_group = get_parallel().moe_ep_group.device_group
num_experts = moe.experts.num_experts num_experts = moe.experts.num_experts
top_k = moe.config.num_experts_per_tok + moe.num_fused_shared_experts top_k = moe.config.num_experts_per_tok + moe.num_fused_shared_experts
intermediate_size = moe.config.moe_intermediate_size intermediate_size = moe.config.moe_intermediate_size
python/sglang/srt/layers/moe/moe_runner/deep_gemm.py
+7 -6
View File
@@ -17,7 +17,6 @@ from sglang.kernels.ops.quantization.per_token_group_quant import per_token_grou
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
from sglang.srt.distributed import get_tp_group
from sglang.srt.distributed.device_communicators.pynccl_allocator import ( from sglang.srt.distributed.device_communicators.pynccl_allocator import (
use_symmetric_memory, use_symmetric_memory,
) )
@@ -596,7 +595,7 @@ class DeepGemmRunnerCore(MoeRunnerCore):
# symmetric path. Only this final output enters the pool; intermediate # symmetric path. Only this final output enters the pool; intermediate
# buffers stay on the default allocator to bound pool occupancy. # buffers stay on the default allocator to bound pool occupancy.
with use_symmetric_memory( with use_symmetric_memory(
get_tp_group(), disabled=not is_allocation_symmetric() get_parallel().tp_group, disabled=not is_allocation_symmetric()
): ):
down_output = torch.empty( down_output = torch.empty(
(all_tokens, K), (all_tokens, K),
@@ -678,7 +677,7 @@ class DeepGemmRunnerCore(MoeRunnerCore):
# GroupGemm-2: (M, N/2) (E, K, N/2) -> (M, K) # GroupGemm-2: (M, N/2) (E, K, N/2) -> (M, K)
with use_symmetric_memory( with use_symmetric_memory(
get_tp_group(), disabled=not is_allocation_symmetric() get_parallel().tp_group, disabled=not is_allocation_symmetric()
): ):
down_output = torch.empty( down_output = torch.empty(
(all_tokens, K), (all_tokens, K),
@@ -855,7 +854,7 @@ class DeepGemmRunnerCore(MoeRunnerCore):
activation_scale_width=down_input_scale.shape[-1], activation_scale_width=down_input_scale.shape[-1],
) )
with use_symmetric_memory( with use_symmetric_memory(
get_tp_group(), disabled=not is_allocation_symmetric() get_parallel().tp_group, disabled=not is_allocation_symmetric()
): ):
down_output = torch.empty( down_output = torch.empty(
(num_groups, m, n), device=hidden_states_device, dtype=torch.bfloat16 (num_groups, m, n), device=hidden_states_device, dtype=torch.bfloat16
@@ -948,7 +947,7 @@ class DeepGemmRunnerCore(MoeRunnerCore):
n = w2_weight.shape[1] n = w2_weight.shape[1]
with use_symmetric_memory( with use_symmetric_memory(
get_tp_group(), disabled=not is_allocation_symmetric() get_parallel().tp_group, disabled=not is_allocation_symmetric()
): ):
down_output = torch.empty( down_output = torch.empty(
(num_groups, m, n), device=hidden_states_device, dtype=torch.bfloat16 (num_groups, m, n), device=hidden_states_device, dtype=torch.bfloat16
@@ -1259,7 +1258,9 @@ def post_permute_deep_gemm_to_standard(
src2dst = running_state["src2dst"] src2dst = running_state["src2dst"]
with use_symmetric_memory(get_tp_group(), disabled=not is_allocation_symmetric()): with use_symmetric_memory(
get_parallel().tp_group, disabled=not is_allocation_symmetric()
):
output = torch.empty( output = torch.empty(
hidden_states_shape, dtype=hidden_states_dtype, device=hidden_states_device hidden_states_shape, dtype=hidden_states_dtype, device=hidden_states_device
) )
python/sglang/srt/layers/moe/moe_runner/flashinfer_cutlass.py
+5 -3
View File
@@ -14,7 +14,6 @@ from typing import TYPE_CHECKING, Optional
import torch import torch
from sglang.kernels.ops.quantization.fp8_kernel import scaled_fp8_quant from sglang.kernels.ops.quantization.fp8_kernel import scaled_fp8_quant
from sglang.srt.distributed import get_tp_group
from sglang.srt.distributed.device_communicators.pynccl_allocator import ( from sglang.srt.distributed.device_communicators.pynccl_allocator import (
use_symmetric_memory, use_symmetric_memory,
) )
@@ -25,6 +24,7 @@ from sglang.srt.layers.moe.moe_runner.base import (
MoeRunnerConfig, MoeRunnerConfig,
register_fused_func, register_fused_func,
) )
from sglang.srt.runtime_context import get_parallel
from sglang.srt.utils import is_flashinfer_available from sglang.srt.utils import is_flashinfer_available
from sglang.srt.utils.common import next_power_of_2 from sglang.srt.utils.common import next_power_of_2
@@ -206,7 +206,7 @@ def _run_flashinfer_cutlass(
if output is None: if output is None:
with use_symmetric_memory( with use_symmetric_memory(
get_tp_group(), disabled=not is_allocation_symmetric() get_parallel().tp_group, disabled=not is_allocation_symmetric()
): ):
output = torch.empty( output = torch.empty(
x.shape[0], x.shape[0],
@@ -435,7 +435,9 @@ def _fused_experts_flashinfer_mxfp4_cutlass(
# new keyword at all on the existing W4A16/MXFP8 paths, so those paths keep # new keyword at all on the existing W4A16/MXFP8 paths, so those paths keep
# working with SGLang's currently pinned release. # working with SGLang's currently pinned release.
humming_kwargs = {"use_wfp4afp8_humming": True} if use_wfp4afp8_humming else {} humming_kwargs = {"use_wfp4afp8_humming": True} if use_wfp4afp8_humming else {}
with use_symmetric_memory(get_tp_group(), disabled=not is_allocation_symmetric()): with use_symmetric_memory(
get_parallel().tp_group, disabled=not is_allocation_symmetric()
):
out = torch.empty(x.shape[0], out_hidden, dtype=output_dtype, device=x.device) out = torch.empty(x.shape[0], out_hidden, dtype=output_dtype, device=x.device)
flashinfer_cutlass_fused_moe( flashinfer_cutlass_fused_moe(
python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py
+10 -6
View File
@@ -15,7 +15,6 @@ from sglang.kernels.ops.quantization.fp8_kernel import (
) )
# Import to register custom ops for torch.compile compatibility # Import to register custom ops for torch.compile compatibility
from sglang.srt.distributed import get_tp_group
from sglang.srt.distributed.device_communicators.pynccl_allocator import ( from sglang.srt.distributed.device_communicators.pynccl_allocator import (
is_symmetric_memory_enabled, is_symmetric_memory_enabled,
is_tensor_in_symmetric_mempool, is_tensor_in_symmetric_mempool,
@@ -34,6 +33,7 @@ from sglang.srt.layers.moe.moe_runner.base import (
register_fused_func, register_fused_func,
) )
from sglang.srt.layers.utils import copy_or_rebind_param from sglang.srt.layers.utils import copy_or_rebind_param
from sglang.srt.runtime_context import get_parallel
from sglang.srt.utils.common import ( from sglang.srt.utils.common import (
is_flashinfer_available, is_flashinfer_available,
next_power_of_2, next_power_of_2,
@@ -817,7 +817,7 @@ def fused_experts_none_to_flashinfer_trtllm_fp8(
# The deferred path returns FlashInfer's permuted/padded GEMM2 # The deferred path returns FlashInfer's permuted/padded GEMM2
# materialization and must not allocate the ordinary final output. # materialization and must not allocate the ordinary final output.
with use_symmetric_memory( with use_symmetric_memory(
get_tp_group(), disabled=not is_allocation_symmetric() get_parallel().tp_group, disabled=not is_allocation_symmetric()
): ):
symm_output = torch.empty( symm_output = torch.empty(
hidden_states.shape[0], hidden_states.shape[0],
@@ -943,7 +943,7 @@ def fused_experts_none_to_flashinfer_trtllm_fp8(
# Allocate output inside symmetric memory context # Allocate output inside symmetric memory context
with use_symmetric_memory( with use_symmetric_memory(
get_tp_group(), disabled=not is_allocation_symmetric() get_parallel().tp_group, disabled=not is_allocation_symmetric()
): ):
symm_output = torch.empty( symm_output = torch.empty(
hidden_states.shape[0], hidden_states.shape[0],
@@ -1083,7 +1083,7 @@ def _fused_experts_flashinfer_mxfp4_sm100_trtllm_gen(
) )
if symm_output is None: if symm_output is None:
with use_symmetric_memory( with use_symmetric_memory(
get_tp_group(), disabled=not is_allocation_symmetric() get_parallel().tp_group, disabled=not is_allocation_symmetric()
): ):
symm_output = torch.empty( symm_output = torch.empty(
num_tokens, num_tokens,
@@ -1401,7 +1401,9 @@ def fused_experts_none_to_flashinfer_trtllm_fp4(
): ):
symm_output = _provided symm_output = _provided
else: else:
with use_symmetric_memory(get_tp_group(), disabled=not _symm_required): with use_symmetric_memory(
get_parallel().tp_group, disabled=not _symm_required
):
symm_output = torch.empty( symm_output = torch.empty(
num_tokens, num_tokens,
hidden_size, hidden_size,
@@ -1567,7 +1569,9 @@ def fused_experts_none_to_flashinfer_trtllm_bf16(
hidden_states = dispatch_output.hidden_states hidden_states = dispatch_output.hidden_states
topk_output = dispatch_output.topk_output topk_output = dispatch_output.topk_output
with use_symmetric_memory(get_tp_group(), disabled=not is_allocation_symmetric()): with use_symmetric_memory(
get_parallel().tp_group, disabled=not is_allocation_symmetric()
):
if use_routed_topk: if use_routed_topk:
assert runner_config.top_k is not None, ( assert runner_config.top_k is not None, (
"runner_config.top_k is required for flashinfer_trtllm_routed." "runner_config.top_k is required for flashinfer_trtllm_routed."
python/sglang/srt/layers/moe/moe_runner/triton_utils/fused_moe.py
+2 -3
View File
@@ -22,14 +22,13 @@ from sglang.kernels.ops.moe.fused_moe_triton_kernels import (
support_tensor_descriptor, support_tensor_descriptor,
) )
from sglang.srt.batch_invariant_ops import is_batch_invariant_mode_enabled from sglang.srt.batch_invariant_ops import is_batch_invariant_mode_enabled
from sglang.srt.distributed import get_tp_group
from sglang.srt.distributed.device_communicators.pynccl_allocator import ( from sglang.srt.distributed.device_communicators.pynccl_allocator import (
use_symmetric_memory, use_symmetric_memory,
) )
from sglang.srt.layers.dp_attention import is_allocation_symmetric from sglang.srt.layers.dp_attention import is_allocation_symmetric
from sglang.srt.layers.moe.moe_runner import MoeRunnerConfig from sglang.srt.layers.moe.moe_runner import MoeRunnerConfig
from sglang.srt.layers.moe.utils import get_moe_padding_size, get_moe_runner_backend from sglang.srt.layers.moe.utils import get_moe_padding_size, get_moe_runner_backend
from sglang.srt.runtime_context import get_exec from sglang.srt.runtime_context import get_exec, get_parallel
from sglang.srt.utils import ( from sglang.srt.utils import (
cpu_has_amx_support, cpu_has_amx_support,
get_bool_env_var, get_bool_env_var,
@@ -578,7 +577,7 @@ def _fused_moe_kernel_sequence(
# symmetric path. Only this output enters the pool; the intermediate caches # symmetric path. Only this output enters the pool; the intermediate caches
# below stay on the default allocator to bound pool occupancy. # below stay on the default allocator to bound pool occupancy.
with use_symmetric_memory( with use_symmetric_memory(
get_tp_group(), disabled=not is_allocation_symmetric() get_parallel().tp_group, disabled=not is_allocation_symmetric()
): ):
out_hidden_states = torch.empty_like(hidden_states) out_hidden_states = torch.empty_like(hidden_states)
python/sglang/srt/layers/moe/token_dispatcher/deepep.py
+2 -2
View File
@@ -7,7 +7,6 @@ from contextlib import nullcontext
from dataclasses import dataclass from dataclasses import dataclass
from typing import TYPE_CHECKING, List, NamedTuple, Optional, Tuple, Union from typing import TYPE_CHECKING, List, NamedTuple, Optional, Tuple, Union
from sglang.srt.distributed.parallel_state import get_tp_group
from sglang.srt.environ import envs from sglang.srt.environ import envs
from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder
from sglang.srt.layers import deep_gemm_wrapper from sglang.srt.layers import deep_gemm_wrapper
@@ -29,6 +28,7 @@ from sglang.srt.layers.moe.utils import (
get_deepep_output_dtype, get_deepep_output_dtype,
is_tbo_enabled, is_tbo_enabled,
) )
from sglang.srt.runtime_context import get_parallel
from sglang.srt.utils import ( from sglang.srt.utils import (
get_bool_env_var, get_bool_env_var,
get_cuda_version, get_cuda_version,
@@ -100,7 +100,7 @@ def _deepep_precompile_tp_barrier() -> None:
# To avoid this, we use torch.distributed's barrier during the compile stage. # To avoid this, we use torch.distributed's barrier during the compile stage.
# We apply this barrier only in the compile stage to prevent extra all-reduce overhead at runtime. # We apply this barrier only in the compile stage to prevent extra all-reduce overhead at runtime.
if envs.SGLANG_IN_DEEPGEMM_PRECOMPILE_STAGE.get(): if envs.SGLANG_IN_DEEPGEMM_PRECOMPILE_STAGE.get():
get_tp_group().barrier() get_parallel().tp_group.barrier()
class DeepEPPDispatchHooks(DispatcherBaseHooks): class DeepEPPDispatchHooks(DispatcherBaseHooks):
python/sglang/srt/layers/moe/token_dispatcher/standard.py
+4 -7
View File
@@ -4,9 +4,6 @@ from typing import TYPE_CHECKING, NamedTuple, Optional, Tuple
import torch import torch
from sglang.srt.distributed import (
get_tp_group,
)
from sglang.srt.distributed.device_communicators.pynccl_allocator import ( from sglang.srt.distributed.device_communicators.pynccl_allocator import (
use_symmetric_memory, use_symmetric_memory,
) )
@@ -153,7 +150,7 @@ class StandardDispatcher(BaseDispatcher):
# Quantize before comm, swizzle after. # Quantize before comm, swizzle after.
with use_symmetric_memory( with use_symmetric_memory(
get_tp_group(), disabled=not is_allocation_symmetric() get_parallel().tp_group, disabled=not is_allocation_symmetric()
): ):
if hidden_states.shape[0] > 0: if hidden_states.shape[0] > 0:
x, x_sf = fp4_quantize_flashinfer( x, x_sf = fp4_quantize_flashinfer(
@@ -167,7 +164,7 @@ class StandardDispatcher(BaseDispatcher):
x_sf = torch.zeros( x_sf = torch.zeros(
0, x_col // 16, dtype=torch.uint8, device=hidden_states.device 0, x_col // 16, dtype=torch.uint8, device=hidden_states.device
) )
topk_weights, topk_ids, x, x_sf = get_tp_group().all_gatherv( topk_weights, topk_ids, x, x_sf = get_parallel().tp_group.all_gatherv(
[topk_weights, topk_ids, x, x_sf], sizes=get_dp_global_num_tokens() [topk_weights, topk_ids, x, x_sf], sizes=get_dp_global_num_tokens()
) )
# TODO: fuse into cutlass moe # TODO: fuse into cutlass moe
@@ -251,10 +248,10 @@ class StandardDispatcher(BaseDispatcher):
(hidden_states,) = combine_input (hidden_states,) = combine_input
if should_use_flashinfer_cutlass_moe_fp4_allgather(): if should_use_flashinfer_cutlass_moe_fp4_allgather():
hidden_states, global_hidden_states = ( hidden_states, global_hidden_states = (
get_local_dp_buffer(get_tp_group()), get_local_dp_buffer(get_parallel().tp_group),
hidden_states, hidden_states,
) )
get_tp_group().reduce_scatterv( get_parallel().tp_group.reduce_scatterv(
global_hidden_states, global_hidden_states,
output=hidden_states, output=hidden_states,
sizes=get_dp_global_num_tokens(), sizes=get_dp_global_num_tokens(),
python/sglang/srt/layers/moe/topk.py
+2 -5
View File
@@ -87,9 +87,6 @@ except ImportError:
from sglang.kernels.fused_op import BaseFusedOp from sglang.kernels.fused_op import BaseFusedOp
from sglang.kernels.ops.attention.dsv4 import mask_topk_ids from sglang.kernels.ops.attention.dsv4 import mask_topk_ids
from sglang.srt.distributed import (
get_tp_group,
)
from sglang.srt.distributed.device_communicators.pynccl_allocator import ( from sglang.srt.distributed.device_communicators.pynccl_allocator import (
use_symmetric_memory, use_symmetric_memory,
) )
@@ -693,7 +690,7 @@ class TopK(BaseFusedOp):
else: else:
self.topk_config.torch_native = False self.topk_config.torch_native = False
with use_symmetric_memory( with use_symmetric_memory(
get_tp_group(), disabled=not is_allocation_symmetric() get_parallel().tp_group, disabled=not is_allocation_symmetric()
): ):
topk_output = select_experts( topk_output = select_experts(
hidden_states=hidden_states, hidden_states=hidden_states,
@@ -784,7 +781,7 @@ class TopK(BaseFusedOp):
) )
topk = self.topk_config.top_k - self.topk_config.num_fused_shared_experts topk = self.topk_config.top_k - self.topk_config.num_fused_shared_experts
with use_symmetric_memory( with use_symmetric_memory(
get_tp_group(), disabled=not is_allocation_symmetric() get_parallel().tp_group, disabled=not is_allocation_symmetric()
): ):
topk_weights = torch.empty((0, topk), dtype=torch.float32, device=device) topk_weights = torch.empty((0, topk), dtype=torch.float32, device=device)
topk_ids = torch.full((0, topk), -1, dtype=torch.int32, device=device) topk_ids = torch.full((0, topk), -1, dtype=torch.int32, device=device)
python/sglang/srt/layers/moe/waterfill.py
+2 -2
View File
@@ -167,12 +167,12 @@ class WaterfillBalancer:
local_routed_counts: Tensor, num_tokens: int local_routed_counts: Tensor, num_tokens: int
) -> Tuple[Tensor, Tensor]: ) -> Tuple[Tensor, Tensor]:
"""Aggregate dynamic load with SGLang EP communication.""" """Aggregate dynamic load with SGLang EP communication."""
from sglang.srt.distributed import get_moe_ep_group
from sglang.srt.distributed.communication_op import ( from sglang.srt.distributed.communication_op import (
moe_expert_parallel_all_reduce, moe_expert_parallel_all_reduce,
) )
from sglang.srt.runtime_context import get_parallel
group = get_moe_ep_group() group = get_parallel().moe_ep_group
world = group.world_size world = group.world_size
buf = torch.zeros( buf = torch.zeros(
world * 2, dtype=torch.int64, device=local_routed_counts.device world * 2, dtype=torch.int64, device=local_routed_counts.device
python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_mxint4_moe.py
+1 -2
View File
@@ -6,7 +6,6 @@ from typing import TYPE_CHECKING
import torch import torch
from compressed_tensors import CompressionFormat from compressed_tensors import CompressionFormat
from sglang.srt.distributed import get_tp_group
from sglang.srt.distributed.device_communicators.pynccl_allocator import ( from sglang.srt.distributed.device_communicators.pynccl_allocator import (
use_symmetric_memory, use_symmetric_memory,
) )
@@ -325,7 +324,7 @@ class CompressedTensorsMxInt4MoE(CompressedTensorsMoEScheme):
) )
with use_symmetric_memory( with use_symmetric_memory(
get_tp_group(), disabled=not is_allocation_symmetric() get_parallel().tp_group, disabled=not is_allocation_symmetric()
): ):
num_tokens = x.shape[0] num_tokens = x.shape[0]
hidden_size = x.shape[-1] hidden_size = x.shape[-1]
python/sglang/srt/layers/quantization/fp8.py
+1 -2
View File
@@ -18,7 +18,6 @@ from sglang.kernels.ops.quantization.fp8_kernel import (
per_token_group_quant_fp8, per_token_group_quant_fp8,
scaled_fp8_quant, scaled_fp8_quant,
) )
from sglang.srt.distributed import get_tp_group
from sglang.srt.distributed.device_communicators.pynccl_allocator import ( from sglang.srt.distributed.device_communicators.pynccl_allocator import (
use_symmetric_memory, use_symmetric_memory,
) )
@@ -2902,7 +2901,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
from sglang.srt.layers.moe.cutlass_moe import cutlass_fused_experts_fp8 from sglang.srt.layers.moe.cutlass_moe import cutlass_fused_experts_fp8
with use_symmetric_memory( with use_symmetric_memory(
get_tp_group(), disabled=not is_allocation_symmetric() get_parallel().tp_group, disabled=not is_allocation_symmetric()
): ):
symm_output = torch.empty_like(x) symm_output = torch.empty_like(x)
python/sglang/srt/layers/quantization/moe_wna16.py
+2 -2
View File
@@ -9,7 +9,6 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional
import numpy as np import numpy as np
import torch import torch
from sglang.srt.distributed.parallel_state import get_tp_group
from sglang.srt.layers.moe import MoeRunner, MoeRunnerBackend, MoeRunnerConfig from sglang.srt.layers.moe import MoeRunner, MoeRunnerBackend, MoeRunnerConfig
from sglang.srt.layers.moe.moe_runner.triton import TritonMoeQuantInfo from sglang.srt.layers.moe.moe_runner.triton import TritonMoeQuantInfo
from sglang.srt.layers.quantization.awq import AWQConfig from sglang.srt.layers.quantization.awq import AWQConfig
@@ -453,7 +452,8 @@ class MoeWNA16Method(FusedMoEMethodBase):
if not layer.quant_config.has_zp and "qzeros" in weight_name: if not layer.quant_config.has_zp and "qzeros" in weight_name:
return return
device = get_tp_group().device tp_group = get_parallel().tp_group
device = tp_group.device
tp_rank = get_parallel().tp_rank tp_rank = get_parallel().tp_rank
loaded_weight = loaded_weight.to(device) loaded_weight = loaded_weight.to(device)
shard_size = layer.intermediate_size_per_partition shard_size = layer.intermediate_size_per_partition
python/sglang/srt/layers/quantization/mxfp4_flashinfer_trtllm_moe.py
+4 -4
View File
@@ -7,7 +7,6 @@ import torch
from torch.nn import Module from torch.nn import Module
from torch.nn.parameter import Parameter from torch.nn.parameter import Parameter
from sglang.srt.distributed import get_tp_group
from sglang.srt.distributed.device_communicators.pynccl_allocator import ( from sglang.srt.distributed.device_communicators.pynccl_allocator import (
use_symmetric_memory, use_symmetric_memory,
) )
@@ -15,6 +14,7 @@ from sglang.srt.layers.dp_attention import is_allocation_symmetric
from sglang.srt.layers.moe.utils import RoutingMethodType from sglang.srt.layers.moe.utils import RoutingMethodType
from sglang.srt.runtime_context import ( from sglang.srt.runtime_context import (
get_exec, get_exec,
get_parallel,
get_platform, get_platform,
) )
from sglang.srt.utils import ( from sglang.srt.utils import (
@@ -418,7 +418,7 @@ class Mxfp4FlashinferTrtllmMoEMethod:
symm_output = None symm_output = None
if not defer_finalize: if not defer_finalize:
with use_symmetric_memory( with use_symmetric_memory(
get_tp_group(), disabled=not is_allocation_symmetric() get_parallel().tp_group, disabled=not is_allocation_symmetric()
): ):
out_hidden_size = ( out_hidden_size = (
x_quant.shape[-1] * 2 x_quant.shape[-1] * 2
@@ -530,7 +530,7 @@ def _fused_finalize_all_reduce_comm_world_size() -> Optional[int]:
CustomAllReduceV2, CustomAllReduceV2,
) )
ca_comm = get_tp_group().ca_comm ca_comm = get_parallel().tp_group.ca_comm
if isinstance(ca_comm, CustomAllReduceV2) and not ca_comm.disabled: if isinstance(ca_comm, CustomAllReduceV2) and not ca_comm.disabled:
all_reduce_fusion.register_comm(ca_comm.obj) all_reduce_fusion.register_comm(ca_comm.obj)
_fused_finalize_all_reduce_world_size = ca_comm.world_size _fused_finalize_all_reduce_world_size = ca_comm.world_size
@@ -560,7 +560,7 @@ def should_use_fuse_finalize_all_reduce(
if not all_reduce_fusion.valid_cluster_sizes(hidden_dim): if not all_reduce_fusion.valid_cluster_sizes(hidden_dim):
return False return False
tp_group = get_tp_group() tp_group = get_parallel().tp_group
if _fused_finalize_all_reduce_comm_world_size() != tp_group.world_size: if _fused_finalize_all_reduce_comm_world_size() != tp_group.world_size:
return False return False
# one push phase counter per row (the plane has num_sm of them) # one push phase counter per row (the plane has num_sm of them)
python/sglang/srt/layers/sampler.py
+1 -2
View File
@@ -7,7 +7,6 @@ import torch.distributed as dist
from torch import nn from torch import nn
from sglang.kernels.ops.sampling.murmur_hash import murmur_hash32 from sglang.kernels.ops.sampling.murmur_hash import murmur_hash32
from sglang.srt.distributed import get_tp_group
from sglang.srt.environ import envs from sglang.srt.environ import envs
from sglang.srt.layers.dp_attention import ( from sglang.srt.layers.dp_attention import (
is_dp_attention_enabled, is_dp_attention_enabled,
@@ -109,7 +108,7 @@ def _select_sampling_mask_rows(
class Sampler(nn.Module): class Sampler(nn.Module):
def __init__(self): def __init__(self):
super().__init__() super().__init__()
self.tp_sync_group = get_tp_group().device_group self.tp_sync_group = get_parallel().tp_group.device_group
self.cp_sync_group = None self.cp_sync_group = None
if is_dp_attention_enabled(): if is_dp_attention_enabled():
self.tp_sync_group = get_parallel().attn_tp_group.device_group self.tp_sync_group = get_parallel().attn_tp_group.device_group
python/sglang/srt/layers/vocab_parallel_embedding.py
+1 -2
View File
@@ -14,7 +14,6 @@ from sglang.kernels.ops.embeddings.vocab_parallel_embedding import (
) )
from sglang.srt.distributed import ( from sglang.srt.distributed import (
divide, divide,
get_tp_group,
tensor_model_parallel_all_reduce, tensor_model_parallel_all_reduce,
) )
from sglang.srt.distributed.device_communicators.pynccl_allocator import ( from sglang.srt.distributed.device_communicators.pynccl_allocator import (
@@ -537,7 +536,7 @@ class VocabParallelEmbedding(torch.nn.Module):
in-place fill deliberately stay outside the pool. in-place fill deliberately stay outside the pool.
""" """
symm_alloc = use_symmetric_memory( symm_alloc = use_symmetric_memory(
get_tp_group(), disabled=not is_allocation_symmetric() get_parallel().tp_group, disabled=not is_allocation_symmetric()
) )
if self.tp_size == 1: if self.tp_size == 1:
with symm_alloc: with symm_alloc:
python/sglang/srt/lora/mem_pool.py
+1 -2
View File
@@ -17,7 +17,6 @@ import torch
from sglang.srt.distributed import ( from sglang.srt.distributed import (
divide, divide,
get_pp_group,
) )
from sglang.srt.environ import envs from sglang.srt.environ import envs
from sglang.srt.lora.eviction_policy import get_eviction_policy from sglang.srt.lora.eviction_policy import get_eviction_policy
@@ -1457,7 +1456,7 @@ class LoRAMemoryPool:
# Non-last PP stages do not own lm_head, so adapters can # Non-last PP stages do not own lm_head, so adapters can
# legitimately contain lm_head LoRA weights with no local # legitimately contain lm_head LoRA weights with no local
# module to load them into, otherwise we should have been able to load this weight. # module to load them into, otherwise we should have been able to load this weight.
assert not get_pp_group().is_last_rank, ( assert not get_parallel().pp_group.is_last_rank, (
f"Failed to load lm_head LoRA weight: {name}, this is only expected to happen on non-last PP stages." f"Failed to load lm_head LoRA weight: {name}, this is only expected to happen on non-last PP stages."
) )
continue continue
python/sglang/srt/lora/trtllm_lora_temp/lora_dispatch.py
+11 -5
View File
@@ -19,11 +19,11 @@ from typing import TYPE_CHECKING
import torch import torch
from sglang.kernels.ops.quantization.fp8_kernel import per_token_group_quant_fp8 from sglang.kernels.ops.quantization.fp8_kernel import per_token_group_quant_fp8
from sglang.srt.distributed import get_tp_group
from sglang.srt.distributed.device_communicators.pynccl_allocator import ( from sglang.srt.distributed.device_communicators.pynccl_allocator import (
use_symmetric_memory, use_symmetric_memory,
) )
from sglang.srt.layers.dp_attention import is_allocation_symmetric from sglang.srt.layers.dp_attention import is_allocation_symmetric
from sglang.srt.runtime_context import get_parallel
from sglang.srt.utils.common import next_power_of_2 from sglang.srt.utils.common import next_power_of_2
if TYPE_CHECKING: if TYPE_CHECKING:
@@ -167,7 +167,7 @@ def fused_experts_none_to_experimental_sgl_trtllm_fp8_lora(
direct_down_output = None direct_down_output = None
if use_virtual_lora_store: if use_virtual_lora_store:
with use_symmetric_memory( with use_symmetric_memory(
get_tp_group(), disabled=not is_allocation_symmetric() get_parallel().tp_group, disabled=not is_allocation_symmetric()
): ):
direct_down_output = torch.empty( direct_down_output = torch.empty(
hidden_states.shape[0], hidden_states.shape[0],
@@ -277,7 +277,9 @@ def fused_experts_none_to_experimental_sgl_trtllm_fp8_lora(
topk_ids, topk_ids,
) )
with use_symmetric_memory(get_tp_group(), disabled=not is_allocation_symmetric()): with use_symmetric_memory(
get_parallel().tp_group, disabled=not is_allocation_symmetric()
):
output = torch.empty( output = torch.empty(
hidden_states.shape[0], hidden_states.shape[0],
hidden_states.shape[1], hidden_states.shape[1],
@@ -401,7 +403,9 @@ def fused_experts_none_to_experimental_sgl_trtllm_bf16_lora(
elif routing_method_type == RoutingMethodType.DeepSeekV3: elif routing_method_type == RoutingMethodType.DeepSeekV3:
routing_method_type = RoutingMethodType.TopK routing_method_type = RoutingMethodType.TopK
with use_symmetric_memory(get_tp_group(), disabled=not is_allocation_symmetric()): with use_symmetric_memory(
get_parallel().tp_group, disabled=not is_allocation_symmetric()
):
direct_down_output = torch.empty( direct_down_output = torch.empty(
hidden_states.shape[0], hidden_states.shape[0],
hidden_states.shape[1], hidden_states.shape[1],
@@ -557,7 +561,9 @@ def fused_experts_none_to_experimental_sgl_trtllm_fp4_lora(
topk_weights=topk_weights, topk_weights=topk_weights,
) )
with use_symmetric_memory(get_tp_group(), disabled=not is_allocation_symmetric()): with use_symmetric_memory(
get_parallel().tp_group, disabled=not is_allocation_symmetric()
):
direct_down_output = torch.empty( direct_down_output = torch.empty(
hidden_states.shape[0], hidden_states.shape[0],
hidden_states.shape[1], hidden_states.shape[1],
python/sglang/srt/lora/trtllm_lora_temp/moe_overlap.py
+12 -6
View File
@@ -62,7 +62,6 @@ def fused_experts_none_to_experimental_sgl_trtllm_fp8_lora_two_stream(
merged_experts_fused_moe_lora_add, merged_experts_fused_moe_lora_add,
) )
from sglang.kernels.ops.quantization.fp8_kernel import per_token_group_quant_fp8 from sglang.kernels.ops.quantization.fp8_kernel import per_token_group_quant_fp8
from sglang.srt.distributed import get_tp_group
from sglang.srt.distributed.device_communicators.pynccl_allocator import ( from sglang.srt.distributed.device_communicators.pynccl_allocator import (
use_symmetric_memory, use_symmetric_memory,
) )
@@ -73,6 +72,7 @@ def fused_experts_none_to_experimental_sgl_trtllm_fp8_lora_two_stream(
from sglang.srt.lora.trtllm_lora_temp.shared_add_overlap import ( from sglang.srt.lora.trtllm_lora_temp.shared_add_overlap import (
maybe_overlap_staged_shared_add, maybe_overlap_staged_shared_add,
) )
from sglang.srt.runtime_context import get_parallel
from sglang.srt.utils.common import next_power_of_2 from sglang.srt.utils.common import next_power_of_2
assert runner_config.activation == "silu" and runner_config.is_gated, ( assert runner_config.activation == "silu" and runner_config.is_gated, (
@@ -195,7 +195,9 @@ def fused_experts_none_to_experimental_sgl_trtllm_fp8_lora_two_stream(
topk_weights=topk_weights, topk_weights=topk_weights,
) )
with use_symmetric_memory(get_tp_group(), disabled=not is_allocation_symmetric()): with use_symmetric_memory(
get_parallel().tp_group, disabled=not is_allocation_symmetric()
):
direct_down_output = torch.empty( direct_down_output = torch.empty(
hidden_states.shape[0], hidden_states.shape[0],
hidden_states.shape[1], hidden_states.shape[1],
@@ -365,13 +367,13 @@ def fused_experts_none_to_experimental_sgl_trtllm_fp4_lora_two_stream(
from sglang.kernels.ops.moe.trtllm_lora_temp.virtual_experts import ( from sglang.kernels.ops.moe.trtllm_lora_temp.virtual_experts import (
merged_experts_fused_moe_lora_add, merged_experts_fused_moe_lora_add,
) )
from sglang.srt.distributed import get_tp_group
from sglang.srt.distributed.device_communicators.pynccl_allocator import ( from sglang.srt.distributed.device_communicators.pynccl_allocator import (
use_symmetric_memory, use_symmetric_memory,
) )
from sglang.srt.layers.dp_attention import is_allocation_symmetric from sglang.srt.layers.dp_attention import is_allocation_symmetric
from sglang.srt.layers.moe.token_dispatcher.standard import StandardCombineInput from sglang.srt.layers.moe.token_dispatcher.standard import StandardCombineInput
from sglang.srt.layers.moe.topk import TopKOutputChecker from sglang.srt.layers.moe.topk import TopKOutputChecker
from sglang.srt.runtime_context import get_parallel
assert runner_config.activation == "silu" and runner_config.is_gated, ( assert runner_config.activation == "silu" and runner_config.is_gated, (
"experimental_sgl_trtllm NVFP4 LoRA currently supports the gated SwiGLU path only." "experimental_sgl_trtllm NVFP4 LoRA currently supports the gated SwiGLU path only."
@@ -464,7 +466,9 @@ def fused_experts_none_to_experimental_sgl_trtllm_fp4_lora_two_stream(
topk_ids=topk_ids, topk_ids=topk_ids,
topk_weights=topk_weights, topk_weights=topk_weights,
) )
with use_symmetric_memory(get_tp_group(), disabled=not is_allocation_symmetric()): with use_symmetric_memory(
get_parallel().tp_group, disabled=not is_allocation_symmetric()
):
direct_down_output = torch.empty( direct_down_output = torch.empty(
hidden_states.shape[0], hidden_states.shape[0],
hidden_states.shape[1], hidden_states.shape[1],
@@ -611,7 +615,6 @@ def fused_experts_none_to_experimental_sgl_trtllm_bf16_lora_two_stream(
from sglang.kernels.ops.moe.trtllm_lora_temp.virtual_experts import ( from sglang.kernels.ops.moe.trtllm_lora_temp.virtual_experts import (
merged_experts_fused_moe_lora_add, merged_experts_fused_moe_lora_add,
) )
from sglang.srt.distributed import get_tp_group
from sglang.srt.distributed.device_communicators.pynccl_allocator import ( from sglang.srt.distributed.device_communicators.pynccl_allocator import (
use_symmetric_memory, use_symmetric_memory,
) )
@@ -622,6 +625,7 @@ def fused_experts_none_to_experimental_sgl_trtllm_bf16_lora_two_stream(
from sglang.srt.layers.moe.token_dispatcher.standard import StandardCombineInput from sglang.srt.layers.moe.token_dispatcher.standard import StandardCombineInput
from sglang.srt.layers.moe.topk import TopKOutputChecker from sglang.srt.layers.moe.topk import TopKOutputChecker
from sglang.srt.layers.moe.utils import RoutingMethodType from sglang.srt.layers.moe.utils import RoutingMethodType
from sglang.srt.runtime_context import get_parallel
assert runner_config.activation == "silu" and runner_config.is_gated, ( assert runner_config.activation == "silu" and runner_config.is_gated, (
"experimental_sgl_trtllm BF16 LoRA currently supports the gated SwiGLU path only." "experimental_sgl_trtllm BF16 LoRA currently supports the gated SwiGLU path only."
@@ -720,7 +724,9 @@ def fused_experts_none_to_experimental_sgl_trtllm_bf16_lora_two_stream(
elif routing_method_type == RoutingMethodType.DeepSeekV3: elif routing_method_type == RoutingMethodType.DeepSeekV3:
routing_method_type = RoutingMethodType.TopK routing_method_type = RoutingMethodType.TopK
with use_symmetric_memory(get_tp_group(), disabled=not is_allocation_symmetric()): with use_symmetric_memory(
get_parallel().tp_group, disabled=not is_allocation_symmetric()
):
direct_down_output = torch.empty( direct_down_output = torch.empty(
hidden_states.shape[0], hidden_states.shape[0],
hidden_states.shape[1], hidden_states.shape[1],
python/sglang/srt/lora/trtllm_lora_temp/sgl_fp8_moe.py
+3 -3
View File
@@ -29,7 +29,6 @@ def fused_experts_fp8_sgl(
from sglang.kernels.ops.moe.trtllm_lora_temp.topk_pack import fused_pack_topk from sglang.kernels.ops.moe.trtllm_lora_temp.topk_pack import fused_pack_topk
from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import ( from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import (
get_tp_group,
is_allocation_symmetric, is_allocation_symmetric,
next_power_of_2, next_power_of_2,
per_token_group_quant_fp8, per_token_group_quant_fp8,
@@ -46,6 +45,7 @@ def fused_experts_fp8_sgl(
from sglang.srt.lora.trtllm_lora_temp.experimental_sgl_trtllm_moe import ( from sglang.srt.lora.trtllm_lora_temp.experimental_sgl_trtllm_moe import (
sgl_trtllm_fp8_block_scale_routed_moe_wrapper as trtllm_fp8_block_scale_routed_moe_wrapper, sgl_trtllm_fp8_block_scale_routed_moe_wrapper as trtllm_fp8_block_scale_routed_moe_wrapper,
) )
from sglang.srt.runtime_context import get_parallel
_SUPPORTED_FP8_ACTIVATIONS = {"silu", "relu2"} _SUPPORTED_FP8_ACTIVATIONS = {"silu", "relu2"}
assert runner_config.activation in _SUPPORTED_FP8_ACTIVATIONS, ( assert runner_config.activation in _SUPPORTED_FP8_ACTIVATIONS, (
@@ -98,7 +98,7 @@ def fused_experts_fp8_sgl(
# Allocate output inside symmetric memory context # Allocate output inside symmetric memory context
with use_symmetric_memory( with use_symmetric_memory(
get_tp_group(), disabled=not is_allocation_symmetric() get_parallel().tp_group, disabled=not is_allocation_symmetric()
): ):
symm_output = torch.empty( symm_output = torch.empty(
hidden_states.shape[0], hidden_states.shape[0],
@@ -198,7 +198,7 @@ def fused_experts_fp8_sgl(
# Allocate output inside symmetric memory context # Allocate output inside symmetric memory context
with use_symmetric_memory( with use_symmetric_memory(
get_tp_group(), disabled=not is_allocation_symmetric() get_parallel().tp_group, disabled=not is_allocation_symmetric()
): ):
symm_output = torch.empty( symm_output = torch.empty(
hidden_states.shape[0], hidden_states.shape[0],
python/sglang/srt/managers/scheduler.py
+3 -5
View File
@@ -99,10 +99,8 @@ from sglang.srt.disaggregation.utils import (
prepare_abort, prepare_abort,
unified_memory_disagg_move_gate, unified_memory_disagg_move_gate,
) )
from sglang.srt.distributed import get_pp_group, get_world_group
from sglang.srt.distributed.parallel_state import ( from sglang.srt.distributed.parallel_state import (
abort_distributed_environment, abort_distributed_environment,
get_tp_group,
) )
from sglang.srt.distributed.parallel_state_wrapper import ParallelState from sglang.srt.distributed.parallel_state_wrapper import ParallelState
from sglang.srt.dllm.mixin.scheduler import SchedulerDllmMixin from sglang.srt.dllm.mixin.scheduler import SchedulerDllmMixin
@@ -1213,14 +1211,14 @@ class Scheduler(
), ),
) )
self.tp_group = get_tp_group() self.tp_group = get_parallel().tp_group
self.tp_cpu_group = self.tp_group.cpu_group self.tp_cpu_group = self.tp_group.cpu_group
self.attn_tp_group = get_parallel().attn_tp_group self.attn_tp_group = get_parallel().attn_tp_group
self.attn_tp_cpu_group = self.attn_tp_group.cpu_group self.attn_tp_cpu_group = self.attn_tp_group.cpu_group
self.attn_cp_group = get_parallel().attn_cp_group self.attn_cp_group = get_parallel().attn_cp_group
self.attn_cp_cpu_group = self.attn_cp_group.cpu_group self.attn_cp_cpu_group = self.attn_cp_group.cpu_group
self.pp_group = get_pp_group() self.pp_group = get_parallel().pp_group
self.world_group = get_world_group() self.world_group = get_parallel().world_group
# NOTE: dp_tp_* are request/data-plane coordination groups (not tensor collectives). # NOTE: dp_tp_* are request/data-plane coordination groups (not tensor collectives).
# When DP attention is enabled, scope to the attention-TP group; otherwise use # When DP attention is enabled, scope to the attention-TP group; otherwise use
python/sglang/srt/managers/scheduler_components/dp_attn.py
+4 -5
View File
@@ -7,7 +7,6 @@ import torch
from sglang.srt.batch_overlap.two_batch_overlap import TboDPAttentionPreparer from sglang.srt.batch_overlap.two_batch_overlap import TboDPAttentionPreparer
from sglang.srt.configs.model_config import ModelConfig from sglang.srt.configs.model_config import ModelConfig
from sglang.srt.distributed.parallel_state import get_tp_group
from sglang.srt.distributed.parallel_state_wrapper import ParallelState from sglang.srt.distributed.parallel_state_wrapper import ParallelState
from sglang.srt.environ import envs from sglang.srt.environ import envs
from sglang.srt.layers.cp.utils import get_cp_strategy from sglang.srt.layers.cp.utils import get_cp_strategy
@@ -192,9 +191,9 @@ class MLPSyncBatchInfo:
) )
num_ranks_in_tp_info = tp_info.shape[0] num_ranks_in_tp_info = tp_info.shape[0]
if device == "cpu": if device == "cpu":
tp_active_ranks = get_tp_group().active_ranks_cpu tp_active_ranks = get_parallel().tp_group.active_ranks_cpu
else: else:
tp_active_ranks = get_tp_group().active_ranks tp_active_ranks = get_parallel().tp_group.active_ranks
if tp_active_ranks.shape[0] < num_ranks_in_tp_info: if tp_active_ranks.shape[0] < num_ranks_in_tp_info:
tp_active_ranks = torch.ones( tp_active_ranks = torch.ones(
num_ranks_in_tp_info, num_ranks_in_tp_info,
@@ -432,9 +431,9 @@ def prepare_mlp_sync_batch_raw(
tbo_preparer = TboDPAttentionPreparer() tbo_preparer = TboDPAttentionPreparer()
use_world_group = world_dp_gather_enabled() use_world_group = world_dp_gather_enabled()
if use_world_group: if use_world_group:
from sglang.srt.distributed.parallel_state import get_world_group from sglang.srt.runtime_context import get_parallel
world = get_world_group() world = get_parallel().world_group
group = torch.distributed.group.WORLD group = torch.distributed.group.WORLD
device = world.device device = world.device
elif len(offload_tags) == 0 and ( elif len(offload_tags) == 0 and (
python/sglang/srt/managers/tp_worker.py
+2 -3
View File
@@ -22,7 +22,6 @@ from typing import TYPE_CHECKING, List, Optional, Tuple
import torch import torch
from sglang.srt.beam_search.logits_capture import capture_pre_sample_logits from sglang.srt.beam_search.logits_capture import capture_pre_sample_logits
from sglang.srt.distributed import get_pp_group, get_world_group
from sglang.srt.distributed.parallel_state_wrapper import ParallelState from sglang.srt.distributed.parallel_state_wrapper import ParallelState
from sglang.srt.environ import envs from sglang.srt.environ import envs
from sglang.srt.managers.io_struct import ( from sglang.srt.managers.io_struct import (
@@ -388,8 +387,8 @@ class TpModelWorker(BaseTpWorker):
self.device = self.model_runner.device self.device = self.model_runner.device
# Init nccl groups # Init nccl groups
self.pp_group = get_pp_group() self.pp_group = get_parallel().pp_group
self.world_group = get_world_group() self.world_group = get_parallel().world_group
# Sync random seed across TP workers. # Sync random seed across TP workers.
# Elastic joiners and last-stage-only draft workers cannot enter the WORLD # Elastic joiners and last-stage-only draft workers cannot enter the WORLD
python/sglang/srt/mem_cache/kv_cache_configurator.py
+3 -4
View File
@@ -28,7 +28,6 @@ from sglang.srt.configs.model_config import (
is_deepseek_v4, is_deepseek_v4,
is_minimax_sparse, is_minimax_sparse,
) )
from sglang.srt.distributed.parallel_state import get_world_group
from sglang.srt.distributed.utils import get_pp_indices from sglang.srt.distributed.utils import get_pp_indices
from sglang.srt.environ import envs from sglang.srt.environ import envs
from sglang.srt.layers.quantization.fp4_kv_cache_quant_method import ( from sglang.srt.layers.quantization.fp4_kv_cache_quant_method import (
@@ -2196,8 +2195,8 @@ class KVCacheConfigurator:
available_gpu_memory = get_available_gpu_memory( available_gpu_memory = get_available_gpu_memory(
self.device, self.device,
self.gpu_id, self.gpu_id,
distributed=get_world_group().world_size > 1, distributed=get_parallel().world_group.world_size > 1,
cpu_group=get_world_group().cpu_group, cpu_group=get_parallel().world_group.cpu_group,
) )
slack_gb = pre_model_load_memory * (1 - get_schedule().mem_fraction_static) slack_gb = pre_model_load_memory * (1 - get_schedule().mem_fraction_static)
@@ -2292,7 +2291,7 @@ class KVCacheConfigurator:
torch.distributed.all_reduce( torch.distributed.all_reduce(
tensor, tensor,
op=torch.distributed.ReduceOp.MIN, op=torch.distributed.ReduceOp.MIN,
group=get_world_group().cpu_group, group=get_parallel().world_group.cpu_group,
) )
token_capacity = tensor.item() token_capacity = tensor.item()
python/sglang/srt/mem_cache/pool_host/base.py
+3 -4
View File
@@ -9,7 +9,6 @@ from typing import Optional
import psutil import psutil
import torch import torch
from sglang.srt.distributed.parallel_state import get_world_group
from sglang.srt.mem_cache.memory_pool import KVCache from sglang.srt.mem_cache.memory_pool import KVCache
from sglang.srt.mem_cache.pool_host.common import ( from sglang.srt.mem_cache.pool_host.common import (
_cuda_host_unregister, _cuda_host_unregister,
@@ -40,7 +39,7 @@ def ranks_per_host() -> int:
if not (torch.distributed.is_available() and torch.distributed.is_initialized()): if not (torch.distributed.is_available() and torch.distributed.is_initialized()):
return 1 return 1
try: try:
world_group = get_world_group() world_group = get_parallel().world_group
except AssertionError: except AssertionError:
return 1 return 1
if world_group.world_size == 1: if world_group.world_size == 1:
@@ -75,9 +74,9 @@ def sync_fixed_hicache_size(size: int, host_size: int) -> int:
return size return size
try: try:
from sglang.srt.distributed.parallel_state import get_pp_group from sglang.srt.runtime_context import get_parallel
pp_group = get_pp_group() pp_group = get_parallel().pp_group
except AssertionError: except AssertionError:
return size return size
python/sglang/srt/mem_cache/storage/flexkv/__init__.py
+7 -12
View File
@@ -22,19 +22,14 @@ def _flexkv_factory(ctx):
"""Build a :class:`FlexKVRadixCache` from a ``TreeCacheBuildContext``. """Build a :class:`FlexKVRadixCache` from a ``TreeCacheBuildContext``.
``TreeCacheBuildContext`` carries TP rank/size and the TP group ``TreeCacheBuildContext`` carries TP rank/size and the TP group
coordinator, but not PP/CP. We pick those up from the global coordinator, but not PP/CP. We pick those up from ``get_parallel()``;
accessors in :mod:`sglang.srt.distributed.parallel_state`; FlexKV FlexKV needs them to fan out lookup/store decisions across the full
needs them to fan out lookup/store decisions across the full TP × CP TP × CP × PP topology.
× PP topology.
""" """
from sglang.srt.distributed.parallel_state import (
get_attn_cp_group,
get_attn_tp_group,
get_pp_group,
)
from sglang.srt.mem_cache.storage.flexkv.flexkv_radix_cache import ( from sglang.srt.mem_cache.storage.flexkv.flexkv_radix_cache import (
FlexKVRadixCache, FlexKVRadixCache,
) )
from sglang.srt.runtime_context import get_parallel
server_args = ctx.server_args server_args = ctx.server_args
@@ -42,15 +37,15 @@ def _flexkv_factory(ctx):
# the regular TP group when attn DP is off — that's fine, the # the regular TP group when attn DP is off — that's fine, the
# connector treats size-1 groups as no-ops. # connector treats size-1 groups as no-ops.
try: try:
pp_group = get_pp_group() pp_group = get_parallel().pp_group
except (RuntimeError, AssertionError): except (RuntimeError, AssertionError):
pp_group = None pp_group = None
try: try:
attn_tp_group = get_attn_tp_group() attn_tp_group = get_parallel().attn_tp_group
except (RuntimeError, AssertionError): except (RuntimeError, AssertionError):
attn_tp_group = ctx.tp_group attn_tp_group = ctx.tp_group
try: try:
attn_cp_group = get_attn_cp_group() attn_cp_group = get_parallel().attn_cp_group
except (RuntimeError, AssertionError): except (RuntimeError, AssertionError):
attn_cp_group = None attn_cp_group = None
python/sglang/srt/mem_cache/storage/flexkv/flexkv_comm.py
+2 -2
View File
@@ -36,7 +36,7 @@ from typing import Any, Dict, List
import torch import torch
import torch.distributed as dist import torch.distributed as dist
from sglang.srt.distributed.parallel_state import get_world_group from sglang.srt.runtime_context import get_parallel
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@@ -174,7 +174,7 @@ class FlexKVComm:
self.pp_size > 1 or self.attn_tp_size > 1 or self.attn_cp_size > 1 self.pp_size > 1 or self.attn_tp_size > 1 or self.attn_cp_size > 1
) )
self._world_cpu_group = get_world_group().cpu_group self._world_cpu_group = get_parallel().world_group.cpu_group
self.pp_group = ( self.pp_group = (
self.pp_cpu_group self.pp_cpu_group
python/sglang/srt/model_executor/model_runner_components/attention_backend_setup.py
+4 -4
View File
@@ -5,13 +5,13 @@ from typing import TYPE_CHECKING, Optional
import msgspec import msgspec
from sglang.srt.distributed import get_world_group
from sglang.srt.environ import envs from sglang.srt.environ import envs
from sglang.srt.layers.attention.attention_registry import ( from sglang.srt.layers.attention.attention_registry import (
ATTENTION_BACKENDS, ATTENTION_BACKENDS,
attn_backend_wrapper, attn_backend_wrapper,
) )
from sglang.srt.layers.attention.tbo_backend import TboAttnBackend from sglang.srt.layers.attention.tbo_backend import TboAttnBackend
from sglang.srt.runtime_context import get_parallel
from sglang.srt.utils import init_cublas from sglang.srt.utils import init_cublas
if TYPE_CHECKING: if TYPE_CHECKING:
@@ -145,9 +145,9 @@ def build_attention_backends(*, model_runner: ModelRunner) -> AttentionBackends:
lazy_init_zbal_gva_mem( lazy_init_zbal_gva_mem(
model_runner.device, model_runner.device,
model_runner.gpu_id, model_runner.gpu_id,
get_world_group().rank_in_group, get_parallel().world_group.rank_in_group,
get_world_group().world_size, get_parallel().world_group.world_size,
get_world_group().cpu_group, get_parallel().world_group.cpu_group,
) )
# Record resolved per-mode backends on the backend for model dispatch. # Record resolved per-mode backends on the backend for model dispatch.
python/sglang/srt/model_executor/model_runner_components/cuda_graph_setup.py
+1 -2
View File
@@ -8,7 +8,6 @@ from typing import TYPE_CHECKING, Any, Optional
import msgspec import msgspec
from sglang.srt.configs.model_config import ModelImpl from sglang.srt.configs.model_config import ModelImpl
from sglang.srt.distributed import get_world_group
from sglang.srt.distributed.device_communicators.pynccl_allocator import ( from sglang.srt.distributed.device_communicators.pynccl_allocator import (
prealloc_symmetric_memory_pool, prealloc_symmetric_memory_pool,
) )
@@ -224,7 +223,7 @@ def refresh_deep_gemm_layout_memory_budget(
set_masked_standard_layout_memory_budget, set_masked_standard_layout_memory_budget,
) )
world_group = get_world_group() world_group = get_parallel().world_group
available_memory_gb = get_available_gpu_memory( available_memory_gb = get_available_gpu_memory(
model_runner.device, model_runner.device,
model_runner.gpu_id, model_runner.gpu_id,
python/sglang/srt/model_executor/model_runner_components/kv_pool_runtime.py
+3 -3
View File
@@ -8,7 +8,6 @@ import torch
from sglang.srt.arg_groups.overrides import post_capture_kv_sizing_planned from sglang.srt.arg_groups.overrides import post_capture_kv_sizing_planned
from sglang.srt.configs.hybrid_arch import mambaish_config from sglang.srt.configs.hybrid_arch import mambaish_config
from sglang.srt.distributed import get_world_group
from sglang.srt.mem_cache.kv_cache_configurator import mm_runtime_reservation_gb from sglang.srt.mem_cache.kv_cache_configurator import mm_runtime_reservation_gb
from sglang.srt.model_executor.cuda_graph_config import Backend from sglang.srt.model_executor.cuda_graph_config import Backend
from sglang.srt.model_executor.runner_utils.pool import graph_pool_borrow_enabled from sglang.srt.model_executor.runner_utils.pool import graph_pool_borrow_enabled
@@ -17,6 +16,7 @@ from sglang.srt.runtime_context import (
get_disagg, get_disagg,
get_exec, get_exec,
get_mm, get_mm,
get_parallel,
pre_capture_activation_reserve_mb, pre_capture_activation_reserve_mb,
) )
from sglang.srt.utils.common import get_available_gpu_memory, get_device_memory_capacity from sglang.srt.utils.common import get_available_gpu_memory, get_device_memory_capacity
@@ -59,8 +59,8 @@ def compute_post_capture_kv_resize(
free_gb = get_available_gpu_memory( free_gb = get_available_gpu_memory(
model_runner.device, model_runner.device,
model_runner.gpu_id, model_runner.gpu_id,
distributed=get_world_group().world_size > 1, distributed=get_parallel().world_group.world_size > 1,
cpu_group=get_world_group().cpu_group, cpu_group=get_parallel().world_group.cpu_group,
) )
headroom_gb = model_runner.pre_model_load_memory * ( headroom_gb = model_runner.pre_model_load_memory * (
1 - model_runner.mem_fraction_static 1 - model_runner.mem_fraction_static
python/sglang/srt/model_executor/model_runner_components/load_model_utils.py
+3 -3
View File
@@ -21,7 +21,6 @@ from sglang.srt.constants import GPU_MEMORY_TYPE_WEIGHTS
from sglang.srt.debug_utils.tensor_dump_forward_hook import ( from sglang.srt.debug_utils.tensor_dump_forward_hook import (
register_forward_hook_for_model, register_forward_hook_for_model,
) )
from sglang.srt.distributed import get_tp_group
from sglang.srt.distributed.parallel_state import monkey_patch_vllm_parallel_state from sglang.srt.distributed.parallel_state import monkey_patch_vllm_parallel_state
from sglang.srt.model_loader.loader import get_model_loader from sglang.srt.model_loader.loader import get_model_loader
from sglang.srt.model_loader.remote_instance_weight_loader_utils import ( from sglang.srt.model_loader.remote_instance_weight_loader_utils import (
@@ -34,6 +33,7 @@ from sglang.srt.runtime_context import (
get_exec, get_exec,
get_model, get_model,
get_observability, get_observability,
get_parallel,
) )
from sglang.srt.utils.common import is_npu from sglang.srt.utils.common import is_npu
from sglang.srt.utils.network import NetworkAddress from sglang.srt.utils.network import NetworkAddress
@@ -373,12 +373,12 @@ def dist_barrier_after_load(
if elastic_ep_backend == "mooncake": if elastic_ep_backend == "mooncake":
# Mooncake does not support `monitored_barrier` # Mooncake does not support `monitored_barrier`
if not is_ep_joiner: if not is_ep_joiner:
dist.barrier(group=get_tp_group().cpu_group) dist.barrier(group=get_parallel().tp_group.cpu_group)
else: else:
# Handle the case where some ranks do not finish loading. # Handle the case where some ranks do not finish loading.
try: try:
dist.monitored_barrier( dist.monitored_barrier(
group=get_tp_group().cpu_group, group=get_parallel().tp_group.cpu_group,
timeout=datetime.timedelta(seconds=UNBALANCED_MODEL_LOADING_TIMEOUT_S), timeout=datetime.timedelta(seconds=UNBALANCED_MODEL_LOADING_TIMEOUT_S),
wait_all_ranks=True, wait_all_ranks=True,
) )
python/sglang/srt/model_executor/model_runner_components/moe_ep_setup.py
+2 -2
View File
@@ -75,7 +75,7 @@ def prepare_moe_topk(
def init_lplb_solvers(*, model_config: ModelConfig) -> None: def init_lplb_solvers(*, model_config: ModelConfig) -> None:
"""Initialize per-layer LPLB solvers from current expert location metadata.""" """Initialize per-layer LPLB solvers from current expert location metadata."""
from sglang.srt.distributed import get_moe_ep_group from sglang.srt.runtime_context import get_parallel
# Gate: refuse LP for non-DeepSeek MoE families whose empty-token paths # Gate: refuse LP for non-DeepSeek MoE families whose empty-token paths
# don't participate in the EP all-reduce (would deadlock under DP- # don't participate in the EP all-reduce (would deadlock under DP-
@@ -88,7 +88,7 @@ def init_lplb_solvers(*, model_config: ModelConfig) -> None:
if metadata is None: if metadata is None:
return return
clear_global_lplb_solvers() clear_global_lplb_solvers()
ep_group = get_moe_ep_group() ep_group = get_parallel().moe_ep_group
for lid in range(metadata.num_layers): for lid in range(metadata.num_layers):
solver = LPLBSolver( solver = LPLBSolver(
phy2log=metadata.physical_to_logical_map[lid], phy2log=metadata.physical_to_logical_map[lid],
python/sglang/srt/model_loader/loader.py
+6 -6
View File
@@ -2046,9 +2046,9 @@ class PreshardedModelLoader(DefaultModelLoader):
cls, local_sig: Optional[str] cls, local_sig: Optional[str]
) -> Optional[str]: ) -> Optional[str]:
try: try:
from sglang.srt.distributed import get_world_group from sglang.srt.runtime_context import get_parallel
group = get_world_group() group = get_parallel().world_group
if group.world_size <= 1: if group.world_size <= 1:
return local_sig return local_sig
all_sigs = group.all_gather_object(local_sig) all_sigs = group.all_gather_object(local_sig)
@@ -2068,20 +2068,20 @@ class PreshardedModelLoader(DefaultModelLoader):
@staticmethod @staticmethod
def _world_rank_and_size() -> Tuple[int, int]: def _world_rank_and_size() -> Tuple[int, int]:
from sglang.srt.distributed import get_world_group from sglang.srt.runtime_context import get_parallel
try: try:
g = get_world_group() g = get_parallel().world_group
return g.rank_in_group, g.world_size return g.rank_in_group, g.world_size
except (AssertionError, AttributeError): except (AssertionError, AttributeError):
return 0, 1 return 0, 1
@staticmethod @staticmethod
def _world_barrier() -> None: def _world_barrier() -> None:
from sglang.srt.distributed import get_world_group from sglang.srt.runtime_context import get_parallel
try: try:
get_world_group().barrier() get_parallel().world_group.barrier()
except (AssertionError, AttributeError): except (AssertionError, AttributeError):
pass pass
python/sglang/srt/model_loader/weight_utils.py
+1 -2
View File
@@ -44,7 +44,6 @@ from sglang.srt.configs.model_config import (
ModelConfig, ModelConfig,
is_qwen3_5_mtp_draft, is_qwen3_5_mtp_draft,
) )
from sglang.srt.distributed import get_world_group
from sglang.srt.layers.quantization import QuantizationConfig, get_quantization_config from sglang.srt.layers.quantization import QuantizationConfig, get_quantization_config
from sglang.srt.layers.quantization.fp8 import Fp8Config from sglang.srt.layers.quantization.fp8 import Fp8Config
from sglang.srt.layers.quantization.modelopt_quant import ( from sglang.srt.layers.quantization.modelopt_quant import (
@@ -1007,7 +1006,7 @@ def _prefetch_all_checkpoints(
# full checkpoint into its own page cache. Global rank would split files # full checkpoint into its own page cache. Global rank would split files
# across nodes, but page cache is not shared across nodes. # across nodes, but page cache is not shared across nodes.
if torch.distributed.is_initialized(): if torch.distributed.is_initialized():
world_group = get_world_group() world_group = get_parallel().world_group
local_rank = world_group.local_rank local_rank = world_group.local_rank
local_world_size = world_group.local_size or world_group.world_size local_world_size = world_group.local_size or world_group.world_size
else: else:
python/sglang/srt/models/apertus.py
+2 -5
View File
@@ -26,9 +26,6 @@ import torch
from torch import nn from torch import nn
from transformers import ApertusConfig from transformers import ApertusConfig
from sglang.srt.distributed import (
get_pp_group,
)
from sglang.srt.layers.activation import XIELU from sglang.srt.layers.activation import XIELU
from sglang.srt.layers.layernorm import RMSNorm from sglang.srt.layers.layernorm import RMSNorm
from sglang.srt.layers.linear import ( from sglang.srt.layers.linear import (
@@ -298,7 +295,7 @@ class ApertusModel(nn.Module):
self.padding_idx = config.pad_token_id self.padding_idx = config.pad_token_id
self.vocab_size = config.vocab_size self.vocab_size = config.vocab_size
self.org_vocab_size = config.vocab_size self.org_vocab_size = config.vocab_size
self.pp_group = get_pp_group() self.pp_group = get_parallel().pp_group
if self.pp_group.is_first_rank: if self.pp_group.is_first_rank:
self.embed_tokens = VocabParallelEmbedding( self.embed_tokens = VocabParallelEmbedding(
config.vocab_size, config.vocab_size,
@@ -430,7 +427,7 @@ class ApertusForCausalLM(nn.Module):
prefix: str = "", prefix: str = "",
) -> None: ) -> None:
super().__init__() super().__init__()
self.pp_group = get_pp_group() self.pp_group = get_parallel().pp_group
self.config = config self.config = config
self.quant_config = quant_config self.quant_config = quant_config
self.model = self._init_model(config, quant_config, add_prefix("model", prefix)) self.model = self._init_model(config, quant_config, add_prefix("model", prefix))
python/sglang/srt/models/arcee.py
+2 -5
View File
@@ -20,9 +20,6 @@ import torch
from torch import nn from torch import nn
from transformers import LlamaConfig from transformers import LlamaConfig
from sglang.srt.distributed import (
get_pp_group,
)
from sglang.srt.layers.activation import get_act_fn from sglang.srt.layers.activation import get_act_fn
from sglang.srt.layers.layernorm import RMSNorm from sglang.srt.layers.layernorm import RMSNorm
from sglang.srt.layers.linear import ( from sglang.srt.layers.linear import (
@@ -272,7 +269,7 @@ class ArceeModel(nn.Module):
self.config = config self.config = config
self.padding_idx = config.pad_token_id self.padding_idx = config.pad_token_id
self.vocab_size = config.vocab_size self.vocab_size = config.vocab_size
self.pp_group = get_pp_group() self.pp_group = get_parallel().pp_group
if self.pp_group.is_first_rank: if self.pp_group.is_first_rank:
self.embed_tokens = VocabParallelEmbedding( self.embed_tokens = VocabParallelEmbedding(
config.vocab_size, config.vocab_size,
@@ -395,7 +392,7 @@ class ArceeForCausalLM(nn.Module):
prefix: str = "", prefix: str = "",
) -> None: ) -> None:
super().__init__() super().__init__()
self.pp_group = get_pp_group() self.pp_group = get_parallel().pp_group
self.config = config self.config = config
self.quant_config = quant_config self.quant_config = quant_config
self.model = self._init_model(config, quant_config, add_prefix("model", prefix)) self.model = self._init_model(config, quant_config, add_prefix("model", prefix))
python/sglang/srt/models/bailing_mm.py
+2 -3
View File
@@ -21,7 +21,6 @@ import torch.nn as nn
import torch.nn.functional as F import torch.nn.functional as F
from transformers import PretrainedConfig from transformers import PretrainedConfig
from sglang.srt.distributed import get_pp_group
from sglang.srt.layers.quantization.base_config import QuantizationConfig from sglang.srt.layers.quantization.base_config import QuantizationConfig
from sglang.srt.layers.utils import PPMissingLayer from sglang.srt.layers.utils import PPMissingLayer
from sglang.srt.managers.mm_utils import ( from sglang.srt.managers.mm_utils import (
@@ -37,7 +36,7 @@ from sglang.srt.model_loader.weight_utils import default_weight_loader
from sglang.srt.models.bailing_moe import BailingMoeV2ForCausalLM from sglang.srt.models.bailing_moe import BailingMoeV2ForCausalLM
from sglang.srt.models.qwen2_5_vl import Qwen2_5_VisionTransformer from sglang.srt.models.qwen2_5_vl import Qwen2_5_VisionTransformer
from sglang.srt.multimodal.mm_utils import materialize_multimodal_features from sglang.srt.multimodal.mm_utils import materialize_multimodal_features
from sglang.srt.runtime_context import get_mm from sglang.srt.runtime_context import get_mm, get_parallel
from sglang.srt.utils import add_prefix from sglang.srt.utils import add_prefix
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@@ -53,7 +52,7 @@ class BailingMMNativeForConditionalGeneration(nn.Module):
prefix: str = "", prefix: str = "",
) -> None: ) -> None:
super().__init__() super().__init__()
self.pp_group = get_pp_group() self.pp_group = get_parallel().pp_group
self.config = config self.config = config
self.quant_config = quant_config self.quant_config = quant_config
self.use_data_parallel = get_mm().mm_enable_dp_encoder self.use_data_parallel = get_mm().mm_enable_dp_encoder
python/sglang/srt/models/bailing_mm_v3.py
+2 -3
View File
@@ -22,7 +22,6 @@ import torch.nn.functional as F
from transformers import PretrainedConfig from transformers import PretrainedConfig
from sglang.srt.configs.bailing_hybrid import is_bailing_multi_gate_enabled from sglang.srt.configs.bailing_hybrid import is_bailing_multi_gate_enabled
from sglang.srt.distributed import get_pp_group
from sglang.srt.layers.quantization.base_config import QuantizationConfig from sglang.srt.layers.quantization.base_config import QuantizationConfig
from sglang.srt.layers.utils import PPMissingLayer from sglang.srt.layers.utils import PPMissingLayer
from sglang.srt.managers.mm_utils import ( from sglang.srt.managers.mm_utils import (
@@ -40,7 +39,7 @@ from sglang.srt.models.bailing_moe_v3 import (
) )
from sglang.srt.models.qwen3_vl import Qwen3VLMoeVisionModel from sglang.srt.models.qwen3_vl import Qwen3VLMoeVisionModel
from sglang.srt.multimodal.mm_utils import materialize_multimodal_features from sglang.srt.multimodal.mm_utils import materialize_multimodal_features
from sglang.srt.runtime_context import get_mm from sglang.srt.runtime_context import get_mm, get_parallel
from sglang.srt.utils import add_prefix from sglang.srt.utils import add_prefix
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@@ -75,7 +74,7 @@ class BailingMoeV3VLForConditionalGeneration(nn.Module):
prefix: str = "", prefix: str = "",
) -> None: ) -> None:
super().__init__() super().__init__()
self.pp_group = get_pp_group() self.pp_group = get_parallel().pp_group
self.config = config self.config = config
self.quant_config = quant_config self.quant_config = quant_config
self.norm_query_embeds = getattr(config, "norm_query_embeds", False) self.norm_query_embeds = getattr(config, "norm_query_embeds", False)
python/sglang/srt/models/bailing_moe.py
+3 -5
View File
@@ -28,8 +28,6 @@ from torch import nn
from transformers import PretrainedConfig from transformers import PretrainedConfig
from sglang.srt.distributed import ( from sglang.srt.distributed import (
get_pp_group,
parallel_state,
tensor_model_parallel_all_reduce, tensor_model_parallel_all_reduce,
) )
from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder
@@ -331,7 +329,7 @@ class BailingMoESparseMoeBlock(nn.Module):
self.ep_size = get_parallel().tp_size self.ep_size = get_parallel().tp_size
self.deepep_dispatcher = DeepEPDispatcher( self.deepep_dispatcher = DeepEPDispatcher(
group=parallel_state.get_tp_group().device_group, group=get_parallel().tp_group.device_group,
router_topk=self.top_k, router_topk=self.top_k,
permute_fusion=True, permute_fusion=True,
num_experts=self.num_experts, num_experts=self.num_experts,
@@ -790,7 +788,7 @@ class BailingMoEModel(nn.Module):
prefix: str = "", prefix: str = "",
): ):
super().__init__() super().__init__()
self.pp_group = get_pp_group() self.pp_group = get_parallel().pp_group
self.config = config self.config = config
self.vocab_size = config.vocab_size self.vocab_size = config.vocab_size
self.embed_dim = config.hidden_size self.embed_dim = config.hidden_size
@@ -895,7 +893,7 @@ class BailingMoEForCausalLM(nn.Module):
prefix: str = "", prefix: str = "",
): ):
super().__init__() super().__init__()
self.pp_group = get_pp_group() self.pp_group = get_parallel().pp_group
self.config = config self.config = config
self.quant_config = quant_config self.quant_config = quant_config
alt_stream = get_stream("alt") if _is_cuda else None alt_stream = get_stream("alt") if _is_cuda else None
python/sglang/srt/models/bailing_moe_linear.py
+2 -3
View File
@@ -13,7 +13,6 @@ from sglang.kernels.ops.attention.fla.layernorm_gated import RMSNorm as RMSNormG
from sglang.kernels.ops.attention.fla.layernorm_gated import layernorm_fn from sglang.kernels.ops.attention.fla.layernorm_gated import layernorm_fn
from sglang.kernels.ops.quantization.fp8_kernel import is_fp8_fnuz from sglang.kernels.ops.quantization.fp8_kernel import is_fp8_fnuz
from sglang.srt.distributed import ( from sglang.srt.distributed import (
get_pp_group,
tensor_model_parallel_all_reduce, tensor_model_parallel_all_reduce,
) )
from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder
@@ -923,7 +922,7 @@ class BailingMoELinearModel(nn.Module):
prefix: str = "", prefix: str = "",
) -> None: ) -> None:
super().__init__() super().__init__()
self.pp_group = get_pp_group() self.pp_group = get_parallel().pp_group
self.config = config self.config = config
self.vocab_size = config.vocab_size self.vocab_size = config.vocab_size
self.embed_dim = config.hidden_size self.embed_dim = config.hidden_size
@@ -1069,7 +1068,7 @@ class BailingMoELinearForCausalLM(nn.Module):
prefix: str = "", prefix: str = "",
) -> None: ) -> None:
super().__init__() super().__init__()
self.pp_group = get_pp_group() self.pp_group = get_parallel().pp_group
self.config = config self.config = config
self.quant_config = quant_config self.quant_config = quant_config
self.model = BailingMoELinearModel( self.model = BailingMoELinearModel(
python/sglang/srt/models/bailing_moe_v3.py
+2 -3
View File
@@ -20,7 +20,6 @@ from sglang.kernels.ops.quantization.fp8_kernel import (
from sglang.srt.configs import KimiLinearConfig from sglang.srt.configs import KimiLinearConfig
from sglang.srt.configs.bailing_hybrid import is_bailing_multi_gate_enabled from sglang.srt.configs.bailing_hybrid import is_bailing_multi_gate_enabled
from sglang.srt.distributed import ( from sglang.srt.distributed import (
get_pp_group,
moe_expert_parallel_all_reduce, moe_expert_parallel_all_reduce,
moe_tensor_model_parallel_all_reduce, moe_tensor_model_parallel_all_reduce,
) )
@@ -1263,7 +1262,7 @@ class BailingMoELinearModel(nn.Module):
num_fused_shared_experts: int = 0, num_fused_shared_experts: int = 0,
) -> None: ) -> None:
super().__init__() super().__init__()
self.pp_group = get_pp_group() self.pp_group = get_parallel().pp_group
self.config = config self.config = config
self.vocab_size = config.vocab_size self.vocab_size = config.vocab_size
self.embed_dim = config.hidden_size self.embed_dim = config.hidden_size
@@ -1436,7 +1435,7 @@ class BailingMoeV3ForCausalLM(nn.Module):
) -> None: ) -> None:
super().__init__() super().__init__()
self.pp_group = get_pp_group() self.pp_group = get_parallel().pp_group
self.config = config self.config = config
self.quant_config = quant_config self.quant_config = quant_config
self.tp_size = get_parallel().tp_size self.tp_size = get_parallel().tp_size
python/sglang/srt/models/deepseek_nextn.py
+1 -2
View File
@@ -25,7 +25,6 @@ from torch import nn
from transformers import PretrainedConfig from transformers import PretrainedConfig
from sglang.kernels.ops.layernorm.fused_eh_norm import fused_eh_norm from sglang.kernels.ops.layernorm.fused_eh_norm import fused_eh_norm
from sglang.srt.distributed import get_pp_group
from sglang.srt.environ import envs from sglang.srt.environ import envs
from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder
from sglang.srt.layers.attention.index_topk_share import IndexTopKShareState from sglang.srt.layers.attention.index_topk_share import IndexTopKShareState
@@ -281,7 +280,7 @@ class DeepseekV3ForCausalLMNextN(DeepseekV3ForCausalLM):
self.tp_size = get_parallel().tp_size self.tp_size = get_parallel().tp_size
self.quant_config = quant_config self.quant_config = quant_config
# if not set, model load will be broken in DeepseekV3ForCausalLM load_weights() # if not set, model load will be broken in DeepseekV3ForCausalLM load_weights()
self.pp_group = get_pp_group() self.pp_group = get_parallel().pp_group
self.determine_num_fused_shared_experts() self.determine_num_fused_shared_experts()
nextn_quant_config = self._resolve_nextn_quant_config(config, quant_config) nextn_quant_config = self._resolve_nextn_quant_config(config, quant_config)
python/sglang/srt/models/deepseek_v2.py
+3 -3
View File
@@ -52,7 +52,7 @@ from sglang.srt.configs.model_config import (
is_deepseek_dsa, is_deepseek_dsa,
is_glm_moe_dsa, is_glm_moe_dsa,
) )
from sglang.srt.distributed import divide, get_pp_group from sglang.srt.distributed import divide
from sglang.srt.environ import envs from sglang.srt.environ import envs
from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder
from sglang.srt.eplb.expert_location import ModelConfigForExpertLocation from sglang.srt.eplb.expert_location import ModelConfigForExpertLocation
@@ -2775,7 +2775,7 @@ class DeepseekV2Model(nn.Module):
self.padding_id = config.pad_token_id self.padding_id = config.pad_token_id
self.vocab_size = config.vocab_size self.vocab_size = config.vocab_size
self.first_k_dense_replace = config.first_k_dense_replace self.first_k_dense_replace = config.first_k_dense_replace
self.pp_group = get_pp_group() self.pp_group = get_parallel().pp_group
if self.pp_group.is_first_rank or (_is_npu and self.pp_group.is_last_rank): if self.pp_group.is_first_rank or (_is_npu and self.pp_group.is_last_rank):
self.embed_tokens = VocabParallelEmbedding( self.embed_tokens = VocabParallelEmbedding(
@@ -3097,7 +3097,7 @@ class DeepseekV2ForCausalLM(nn.Module, DeepseekV2WeightLoaderMixin):
if quant_config is not None: if quant_config is not None:
quant_config.update_packed_modules_mapping(self.packed_modules_mapping) quant_config.update_packed_modules_mapping(self.packed_modules_mapping)
self.pp_group = get_pp_group() self.pp_group = get_parallel().pp_group
self.config = config self.config = config
self.tp_size = get_parallel().tp_size self.tp_size = get_parallel().tp_size
self.quant_config = quant_config self.quant_config = quant_config
python/sglang/srt/models/deepseek_v4.py
+9 -13
View File
@@ -42,10 +42,6 @@ from sglang.kernels.ops.quantization.fp8_kernel import (
) )
from sglang.srt.compilation.compilation_config import register_split_op from sglang.srt.compilation.compilation_config import register_split_op
from sglang.srt.configs.deepseek_v4 import DeepSeekV4Config from sglang.srt.configs.deepseek_v4 import DeepSeekV4Config
from sglang.srt.distributed import (
get_pp_group,
get_tp_group,
)
from sglang.srt.distributed.device_communicators.pynccl_allocator import ( from sglang.srt.distributed.device_communicators.pynccl_allocator import (
use_symmetric_memory, use_symmetric_memory,
) )
@@ -2872,7 +2868,7 @@ class DeepseekV4DecoderLayer(nn.Module):
# all-reduce input. Gated by is_allocation_symmetric() (mirrors the # all-reduce input. Gated by is_allocation_symmetric() (mirrors the
# TileLang path in _mhc_pre_impl / mhc_fused_post_pre). # TileLang path in _mhc_pre_impl / mhc_fused_post_pre).
with use_symmetric_memory( with use_symmetric_memory(
get_tp_group(), disabled=not is_allocation_symmetric() get_parallel().tp_group, disabled=not is_allocation_symmetric()
): ):
y = hc_combine(x_flat, pre.squeeze(1), self.hc_mult, dtype) y = hc_combine(x_flat, pre.squeeze(1), self.hc_mult, dtype)
return y, post.squeeze(1), comb.squeeze(1), False return y, post.squeeze(1), comb.squeeze(1), False
@@ -3611,7 +3607,7 @@ class DeepseekV4DecoderLayer(nn.Module):
) )
elif _use_tp_moe_gather: elif _use_tp_moe_gather:
hidden_states, local_hidden_states = ( hidden_states, local_hidden_states = (
get_global_dp_buffer(get_tp_group()), get_global_dp_buffer(get_parallel().tp_group),
hidden_states, hidden_states,
) )
if _do_shared_local and local_hidden_states.shape[0] > 0: if _do_shared_local and local_hidden_states.shape[0] > 0:
@@ -3661,7 +3657,7 @@ class DeepseekV4DecoderLayer(nn.Module):
hidden_states = dsa_cp_reduce_scatter_hidden_states(hidden_states) hidden_states = dsa_cp_reduce_scatter_hidden_states(hidden_states)
elif _use_tp_moe_gather: elif _use_tp_moe_gather:
hidden_states, global_hidden_states = ( hidden_states, global_hidden_states = (
get_local_dp_buffer(get_tp_group()), get_local_dp_buffer(get_parallel().tp_group),
hidden_states, hidden_states,
) )
if should_use_dp_reduce_scatterv() or _use_reduce_scatterv: if should_use_dp_reduce_scatterv() or _use_reduce_scatterv:
@@ -3669,7 +3665,7 @@ class DeepseekV4DecoderLayer(nn.Module):
# each rank its own token slice, in one op. Correct because the # each rank its own token slice, in one op. Correct because the
# MoE-internal all_reduce was skipped (mlp_reduce_scatter above). # MoE-internal all_reduce was skipped (mlp_reduce_scatter above).
# This is the symmetric inverse of the all_gatherv gather. # This is the symmetric inverse of the all_gatherv gather.
get_tp_group().reduce_scatterv( get_parallel().tp_group.reduce_scatterv(
global_hidden_states, global_hidden_states,
output=hidden_states, output=hidden_states,
sizes=get_dp_global_num_tokens(), sizes=get_dp_global_num_tokens(),
@@ -3987,7 +3983,7 @@ class DeepseekV4Model(nn.Module):
) -> None: ) -> None:
super().__init__() super().__init__()
self.config = config self.config = config
self.pp_group = get_pp_group() self.pp_group = get_parallel().pp_group
self.hidden_size = config.hidden_size self.hidden_size = config.hidden_size
if self.pp_group.is_first_rank: if self.pp_group.is_first_rank:
embedding_quant_config = ( embedding_quant_config = (
@@ -4376,7 +4372,7 @@ class DeepseekV4Model(nn.Module):
# once across DP ranks, then populate each child's global_num_tokens + # once across DP ranks, then populate each child's global_num_tokens +
# global_dp_buffer_len so the gatherv/reduce_scatterv buffers size correctly. # global_dp_buffer_len so the gatherv/reduce_scatterv buffers size correctly.
if get_moe_a2a_backend().is_none() and get_parallel().attn_dp_size > 1: if get_moe_a2a_backend().is_none() and get_parallel().attn_dp_size > 1:
tp_group = get_tp_group() tp_group = get_parallel().tp_group
world = tp_group.world_size world = tp_group.world_size
children = forward_batch.tbo_children children = forward_batch.tbo_children
local_lens = torch.tensor( local_lens = torch.tensor(
@@ -4620,7 +4616,7 @@ class DeepseekV4ForCausalLM(nn.Module):
if config.model_type == "deepseek_v41" and config.vision_n_layers > 0: if config.model_type == "deepseek_v41" and config.vision_n_layers > 0:
if ( if (
get_parallel().attn_cp_size != 1 get_parallel().attn_cp_size != 1
or get_pp_group().world_size != 1 or get_parallel().pp_group.world_size != 1
or not get_moe_a2a_backend().is_none() or not get_moe_a2a_backend().is_none()
): ):
raise ValueError( raise ValueError(
@@ -4636,7 +4632,7 @@ class DeepseekV4ForCausalLM(nn.Module):
self.model = DeepseekV4Model( self.model = DeepseekV4Model(
config, quant_config, prefix=add_prefix("model", prefix) config, quant_config, prefix=add_prefix("model", prefix)
) )
self.pp_group = get_pp_group() self.pp_group = get_parallel().pp_group
if self.pp_group.is_last_rank: if self.pp_group.is_last_rank:
if self.pp_group.world_size == 1 and config.tie_word_embeddings: if self.pp_group.world_size == 1 and config.tie_word_embeddings:
self.lm_head = self.model.embed_tokens self.lm_head = self.model.embed_tokens
@@ -5108,7 +5104,7 @@ class DeepseekV4ForCausalLM(nn.Module):
compile_secs = time.perf_counter() - tic compile_secs = time.perf_counter() - tic
# Runs before init_memory_pool(); don't let transients skew pool sizing. # Runs before init_memory_pool(); don't let transients skew pool sizing.
torch.cuda.empty_cache() torch.cuda.empty_cache()
get_tp_group().barrier() get_parallel().tp_group.barrier()
logger.info( logger.info(
"DeepSeek V4 MHC prewarm at load: compile %.1fs, rank sync +%.1fs", "DeepSeek V4 MHC prewarm at load: compile %.1fs, rank sync +%.1fs",
compile_secs, compile_secs,
python/sglang/srt/models/deepseek_v4_nextn.py
+1 -2
View File
@@ -6,7 +6,6 @@ import torch.nn.functional as F
from torch import nn from torch import nn
from transformers import PretrainedConfig from transformers import PretrainedConfig
from sglang.srt.distributed import get_pp_group
from sglang.srt.hardware_backend.npu.dsv4.dsv4_rope import prime_rope_cos_sin from sglang.srt.hardware_backend.npu.dsv4.dsv4_rope import prime_rope_cos_sin
from sglang.srt.layers.attention.dsa.utils import ( from sglang.srt.layers.attention.dsa.utils import (
dsa_use_prefill_cp, dsa_use_prefill_cp,
@@ -219,7 +218,7 @@ class DeepseekV4ForCausalLMNextN(DeepseekV4ForCausalLM):
nn.Module.__init__(self) nn.Module.__init__(self)
self.config = config self.config = config
self.tp_size = get_parallel().tp_size self.tp_size = get_parallel().tp_size
self.pp_group = get_pp_group() self.pp_group = get_parallel().pp_group
self.quant_config = quant_config self.quant_config = quant_config
self.wo_a_fp8 = wo_a_fp8_gemm_enabled(quant_config) self.wo_a_fp8 = wo_a_fp8_gemm_enabled(quant_config)
self.determine_num_fused_shared_experts() self.determine_num_fused_shared_experts()
python/sglang/srt/models/dots3_common/modeling.py
+6 -8
View File
@@ -45,8 +45,6 @@ from sglang.srt.batch_overlap.two_batch_overlap import (
) )
from sglang.srt.configs.dots3 import Dots3Config from sglang.srt.configs.dots3 import Dots3Config
from sglang.srt.distributed import ( from sglang.srt.distributed import (
get_pp_group,
parallel_state,
tensor_model_parallel_all_reduce, tensor_model_parallel_all_reduce,
) )
from sglang.srt.distributed.device_communicators.pynccl_allocator import ( from sglang.srt.distributed.device_communicators.pynccl_allocator import (
@@ -391,7 +389,7 @@ class Dots3MoE(nn.Module):
) )
self.deepep_dispatcher = MaybeTboDeepEPDispatcher( self.deepep_dispatcher = MaybeTboDeepEPDispatcher(
group=parallel_state.get_tp_group().device_group, group=get_parallel().tp_group.device_group,
router_topk=self.top_k, router_topk=self.top_k,
permute_fusion=True, permute_fusion=True,
num_experts=self.num_experts, num_experts=self.num_experts,
@@ -459,7 +457,7 @@ class Dots3MoE(nn.Module):
final_hidden_states = self.experts(hidden_states, topk_output) final_hidden_states = self.experts(hidden_states, topk_output)
current_stream.wait_stream(self.alt_stream) current_stream.wait_stream(self.alt_stream)
with use_symmetric_memory(parallel_state.get_tp_group()) as sm: with use_symmetric_memory(get_parallel().tp_group) as sm:
final_hidden_states_out = torch.empty_like(final_hidden_states) final_hidden_states_out = torch.empty_like(final_hidden_states)
torch.add(final_hidden_states, shared_output, out=final_hidden_states_out) torch.add(final_hidden_states, shared_output, out=final_hidden_states_out)
@@ -491,7 +489,7 @@ class Dots3MoE(nn.Module):
final_hidden_states = self.experts(hidden_states, topk_output) final_hidden_states = self.experts(hidden_states, topk_output)
if shared_output is not None: if shared_output is not None:
with use_symmetric_memory(parallel_state.get_tp_group()) as sm: with use_symmetric_memory(get_parallel().tp_group) as sm:
final_hidden_states_out = torch.empty_like(final_hidden_states) final_hidden_states_out = torch.empty_like(final_hidden_states)
torch.add(final_hidden_states, shared_output, out=final_hidden_states_out) torch.add(final_hidden_states, shared_output, out=final_hidden_states_out)
final_hidden_states = final_hidden_states_out final_hidden_states = final_hidden_states_out
@@ -1731,7 +1729,7 @@ class Dots3Model(nn.Module):
super().__init__() super().__init__()
_require_cuda() _require_cuda()
self.first_k_dense_replace = config.first_k_dense_replace self.first_k_dense_replace = config.first_k_dense_replace
self.pp_group = get_pp_group() self.pp_group = get_parallel().pp_group
if self.pp_group.is_first_rank: if self.pp_group.is_first_rank:
self.embed_tokens = VocabParallelEmbedding( self.embed_tokens = VocabParallelEmbedding(
@@ -1865,7 +1863,7 @@ class Dots3LanguageModelForCausalLM(nn.Module):
"g_proj", "g_proj",
] ]
self.pp_group = get_pp_group() self.pp_group = get_parallel().pp_group
self.config = config self.config = config
self.tp_size = get_parallel().tp_size self.tp_size = get_parallel().tp_size
self.quant_config = quant_config self.quant_config = quant_config
@@ -2605,7 +2603,7 @@ class DotsNoteOmniThinkerForConditionalGeneration(nn.Module):
) )
self.config = config self.config = config
self.pp_group = get_pp_group() self.pp_group = get_parallel().pp_group
model_dir = Path(config._name_or_path) model_dir = Path(config._name_or_path)
self.language_model = Dots3LanguageModelForCausalLM( self.language_model = Dots3LanguageModelForCausalLM(
config, config,
python/sglang/srt/models/dots3_common/nextn.py
+1 -2
View File
@@ -7,7 +7,6 @@ import torch
from torch import nn from torch import nn
from transformers import PretrainedConfig from transformers import PretrainedConfig
from sglang.srt.distributed import get_pp_group
from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder
from sglang.srt.layers.dp_attention import is_dp_attention_enabled from sglang.srt.layers.dp_attention import is_dp_attention_enabled
from sglang.srt.layers.layernorm import RMSNorm from sglang.srt.layers.layernorm import RMSNorm
@@ -148,7 +147,7 @@ class Dots3NoteForCausalLMNextN(Dots3LanguageModelForCausalLM):
self.config = config self.config = config
self.tp_size = get_parallel().tp_size self.tp_size = get_parallel().tp_size
self.quant_config = quant_config self.quant_config = quant_config
self.pp_group = get_pp_group() self.pp_group = get_parallel().pp_group
self.fuse_qkv_a_g_proj = True self.fuse_qkv_a_g_proj = True
self.packed_modules_mapping = { self.packed_modules_mapping = {
"fused_qkv_a_g_proj_with_mqa": [ "fused_qkv_a_g_proj_with_mqa": [
python/sglang/srt/models/dots_vlm.py
+2 -2
View File
@@ -23,7 +23,6 @@ import torch
from torch import nn from torch import nn
from sglang.srt.configs.dots_vlm import DotsVLMConfig from sglang.srt.configs.dots_vlm import DotsVLMConfig
from sglang.srt.distributed import get_pp_group
from sglang.srt.layers.quantization.base_config import QuantizationConfig from sglang.srt.layers.quantization.base_config import QuantizationConfig
from sglang.srt.managers.mm_utils import ( from sglang.srt.managers.mm_utils import (
MultiModalityDataPaddingPatternMultimodalTokens, MultiModalityDataPaddingPatternMultimodalTokens,
@@ -33,6 +32,7 @@ from sglang.srt.managers.schedule_batch import MultimodalDataItem, MultimodalInp
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors
from sglang.srt.model_loader.weight_utils import default_weight_loader from sglang.srt.model_loader.weight_utils import default_weight_loader
from sglang.srt.models.deepseek_v2 import DeepseekV2ForCausalLM from sglang.srt.models.deepseek_v2 import DeepseekV2ForCausalLM
from sglang.srt.runtime_context import get_parallel
from .dots_vlm_vit import DotsVisionTransformer from .dots_vlm_vit import DotsVisionTransformer
@@ -56,7 +56,7 @@ class DotsVLMForCausalLM(nn.Module):
self.config = config self.config = config
self.image_token_id = config.im_span_id self.image_token_id = config.im_span_id
self.video_token_id = config.video_span_id self.video_token_id = config.video_span_id
self.pp_group = get_pp_group() self.pp_group = get_parallel().pp_group
if not config.encoder_only: if not config.encoder_only:
self.language_model = DeepseekV2ForCausalLM( self.language_model = DeepseekV2ForCausalLM(
python/sglang/srt/models/ernie45_moe_vl.py
+1 -2
View File
@@ -23,7 +23,6 @@ from torch import nn
from transformers import PretrainedConfig from transformers import PretrainedConfig
from sglang.srt.distributed import ( from sglang.srt.distributed import (
get_pp_group,
tensor_model_parallel_all_reduce, tensor_model_parallel_all_reduce,
) )
from sglang.srt.layers.dp_attention import is_dp_attention_enabled from sglang.srt.layers.dp_attention import is_dp_attention_enabled
@@ -472,7 +471,7 @@ class Ernie4_5_VLMoeModel(nn.Module):
) -> None: ) -> None:
super().__init__() super().__init__()
self.config = config self.config = config
self.pp_group = get_pp_group() self.pp_group = get_parallel().pp_group
if self.pp_group.is_first_rank: if self.pp_group.is_first_rank:
self.embed_tokens = VocabParallelEmbedding( self.embed_tokens = VocabParallelEmbedding(
python/sglang/srt/models/exaone4.py
+2 -3
View File
@@ -5,7 +5,6 @@ import torch
from torch import nn from torch import nn
from transformers import Exaone4Config from transformers import Exaone4Config
from sglang.srt.distributed import get_pp_group
from sglang.srt.layers.activation import SiluAndMul from sglang.srt.layers.activation import SiluAndMul
from sglang.srt.layers.layernorm import RMSNorm from sglang.srt.layers.layernorm import RMSNorm
from sglang.srt.layers.linear import ( from sglang.srt.layers.linear import (
@@ -310,7 +309,7 @@ class Exaone4Model(nn.Module):
self.config = config self.config = config
self.quant_config = quant_config self.quant_config = quant_config
self.vocab_size = config.vocab_size self.vocab_size = config.vocab_size
self.pp_group = get_pp_group() self.pp_group = get_parallel().pp_group
if self.pp_group.is_first_rank: if self.pp_group.is_first_rank:
self.embed_tokens = VocabParallelEmbedding( self.embed_tokens = VocabParallelEmbedding(
config.vocab_size, config.vocab_size,
@@ -423,7 +422,7 @@ class Exaone4ForCausalLM(nn.Module):
prefix: str = "", prefix: str = "",
): ):
super().__init__() super().__init__()
self.pp_group = get_pp_group() self.pp_group = get_parallel().pp_group
self.config = config self.config = config
self.quant_config = quant_config self.quant_config = quant_config
python/sglang/srt/models/exaone_moe.py
+3 -4
View File
@@ -25,7 +25,6 @@ from torch import nn
from transformers import PretrainedConfig from transformers import PretrainedConfig
from sglang.srt.distributed import ( from sglang.srt.distributed import (
get_pp_group,
tensor_model_parallel_all_reduce, tensor_model_parallel_all_reduce,
) )
from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder
@@ -536,7 +535,7 @@ class ExaoneMoEModel(nn.Module):
self.config = config self.config = config
self.padding_idx = config.pad_token_id self.padding_idx = config.pad_token_id
self.vocab_size = config.vocab_size self.vocab_size = config.vocab_size
self.pp_group = get_pp_group() self.pp_group = get_parallel().pp_group
if self.pp_group.is_first_rank: if self.pp_group.is_first_rank:
self.embed_tokens = VocabParallelEmbedding( self.embed_tokens = VocabParallelEmbedding(
@@ -624,7 +623,7 @@ class ExaoneMoEForCausalLM(nn.Module):
prefix: str = "", prefix: str = "",
) -> None: ) -> None:
super().__init__() super().__init__()
self.pp_group = get_pp_group() self.pp_group = get_parallel().pp_group
self.config = config self.config = config
self.quant_config = quant_config self.quant_config = quant_config
alt_stream = get_stream("alt") if _is_cuda else None alt_stream = get_stream("alt") if _is_cuda else None
@@ -856,7 +855,7 @@ class ExaoneMoEForCausalLM(nn.Module):
) )
def set_eagle3_layers_to_capture(self, layer_ids: Optional[list[int]] = None): def set_eagle3_layers_to_capture(self, layer_ids: Optional[list[int]] = None):
if not get_pp_group().is_last_rank: if not get_parallel().pp_group.is_last_rank:
return return
self.capture_aux_hidden_states = True self.capture_aux_hidden_states = True
python/sglang/srt/models/exaone_moe_mtp.py
+1 -2
View File
@@ -23,7 +23,6 @@ import torch
from torch import nn from torch import nn
from transformers import PretrainedConfig from transformers import PretrainedConfig
from sglang.srt.distributed import get_pp_group
from sglang.srt.layers.layernorm import RMSNorm from sglang.srt.layers.layernorm import RMSNorm
from sglang.srt.layers.logits_processor import LogitsProcessor from sglang.srt.layers.logits_processor import LogitsProcessor
from sglang.srt.layers.quantization.base_config import QuantizationConfig from sglang.srt.layers.quantization.base_config import QuantizationConfig
@@ -48,7 +47,7 @@ class ExaoneMoEForCausalLMMTP(ExaoneMoEForCausalLM):
config.num_hidden_layers = 1 config.num_hidden_layers = 1
self.tp_size = get_parallel().tp_size self.tp_size = get_parallel().tp_size
self.quant_config = quant_config self.quant_config = quant_config
self.pp_group = get_pp_group() self.pp_group = get_parallel().pp_group
self.fc = nn.Linear(2 * config.hidden_size, config.hidden_size, bias=False) self.fc = nn.Linear(2 * config.hidden_size, config.hidden_size, bias=False)
self.pre_fc_norm_embedding = RMSNorm( self.pre_fc_norm_embedding = RMSNorm(
python/sglang/srt/models/falcon_h1.py
+1 -2
View File
@@ -5,7 +5,6 @@ import torch
from torch import nn from torch import nn
from sglang.srt.configs.falcon_h1 import FalconH1Config from sglang.srt.configs.falcon_h1 import FalconH1Config
from sglang.srt.distributed import get_pp_group
from sglang.srt.layers.activation import SiluAndMul from sglang.srt.layers.activation import SiluAndMul
from sglang.srt.layers.attention.hybrid_linear_attn_backend import ( from sglang.srt.layers.attention.hybrid_linear_attn_backend import (
HybridLinearAttnBackend, HybridLinearAttnBackend,
@@ -457,7 +456,7 @@ class FalconH1ForCausalLM(nn.Module):
) -> None: ) -> None:
super().__init__() super().__init__()
self.config = config self.config = config
self.pp_group = get_pp_group() self.pp_group = get_parallel().pp_group
assert self.pp_group.is_first_rank and self.pp_group.is_last_rank assert self.pp_group.is_first_rank and self.pp_group.is_last_rank
self.quant_config = quant_config self.quant_config = quant_config
self.model = FalconH1Model( self.model = FalconH1Model(
python/sglang/srt/models/gemma4_causal.py
+2 -5
View File
@@ -31,9 +31,6 @@ from sglang.kernels.ops.layernorm.gemma4_fused_ops import (
gemma_rmsnorm_residual_scalar, gemma_rmsnorm_residual_scalar,
gemma_routing_post_topk, gemma_routing_post_topk,
) )
from sglang.srt.distributed import (
get_pp_group,
)
from sglang.srt.layers.layernorm import Gemma4RMSNorm, RMSNorm from sglang.srt.layers.layernorm import Gemma4RMSNorm, RMSNorm
from sglang.srt.layers.linear import ( from sglang.srt.layers.linear import (
QKVParallelLinear, QKVParallelLinear,
@@ -782,7 +779,7 @@ class Gemma4TextModel(PreTrainedModel):
self.quant_config = quant_config self.quant_config = quant_config
self.vocab_size = config.vocab_size self.vocab_size = config.vocab_size
self.padding_idx = getattr(config, "pad_token_id", None) self.padding_idx = getattr(config, "pad_token_id", None)
self.pp_group = get_pp_group() self.pp_group = get_parallel().pp_group
# Token / per-layer embedding tables and the per-layer projection only # Token / per-layer embedding tables and the per-layer projection only
# produce activations consumed at the model entry, so they live on the # produce activations consumed at the model entry, so they live on the
@@ -1090,7 +1087,7 @@ class Gemma4ForCausalLM(PreTrainedModel):
prefix: str = "", prefix: str = "",
) -> None: ) -> None:
super().__init__(config=config) super().__init__(config=config)
self.pp_group = get_pp_group() self.pp_group = get_parallel().pp_group
self.config = config self.config = config
self.quant_config = quant_config self.quant_config = quant_config
python/sglang/srt/models/gemma4_mm.py
+2 -2
View File
@@ -28,7 +28,6 @@ from transformers import (
PreTrainedModel, PreTrainedModel,
) )
from sglang.srt.distributed import get_pp_group
from sglang.srt.environ import envs from sglang.srt.environ import envs
from sglang.srt.layers.attention.triton_backend import TritonAttnBackend from sglang.srt.layers.attention.triton_backend import TritonAttnBackend
from sglang.srt.layers.layernorm import Gemma4RMSNorm from sglang.srt.layers.layernorm import Gemma4RMSNorm
@@ -65,6 +64,7 @@ from sglang.srt.models.gemma4_causal import (
pp_filter_load_weight, pp_filter_load_weight,
) )
from sglang.srt.models.gemma4_vision import Gemma4VisionEncoder from sglang.srt.models.gemma4_vision import Gemma4VisionEncoder
from sglang.srt.runtime_context import get_parallel
from sglang.srt.utils import add_prefix, cpu_has_amx_support, is_cpu from sglang.srt.utils import add_prefix, cpu_has_amx_support, is_cpu
from sglang.srt.utils.hf_transformers_utils import get_processor from sglang.srt.utils.hf_transformers_utils import get_processor
@@ -187,7 +187,7 @@ class Gemma4ForConditionalGeneration(PreTrainedModel):
prefix: str = "", prefix: str = "",
) -> None: ) -> None:
super().__init__(config=config) super().__init__(config=config)
self.pp_group = get_pp_group() self.pp_group = get_parallel().pp_group
self.config = config self.config = config
self.quant_config = quant_config self.quant_config = quant_config
python/sglang/srt/models/gemma4_mtp.py
+2 -2
View File
@@ -21,7 +21,6 @@ import torch
from torch import nn from torch import nn
from transformers import PretrainedConfig, PreTrainedModel from transformers import PretrainedConfig, PreTrainedModel
from sglang.srt.distributed import get_pp_group
from sglang.srt.layers.linear import ReplicatedLinear from sglang.srt.layers.linear import ReplicatedLinear
from sglang.srt.layers.logits_processor import ( from sglang.srt.layers.logits_processor import (
LogitsMetadata, LogitsMetadata,
@@ -32,6 +31,7 @@ from sglang.srt.layers.quantization.base_config import QuantizationConfig
from sglang.srt.mem_cache.memory_pool import KVCache from sglang.srt.mem_cache.memory_pool import KVCache
from sglang.srt.model_executor.forward_batch_info import ForwardBatch from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.models.gemma4_causal import Gemma4ForCausalLM, Gemma4TextModel from sglang.srt.models.gemma4_causal import Gemma4ForCausalLM, Gemma4TextModel
from sglang.srt.runtime_context import get_parallel
from sglang.srt.speculative.frozen_kv_mtp_info import FrozenKVMTPContext from sglang.srt.speculative.frozen_kv_mtp_info import FrozenKVMTPContext
from sglang.srt.utils import add_prefix from sglang.srt.utils import add_prefix
@@ -73,7 +73,7 @@ class Gemma4AssistantForCausalLM(Gemma4ForCausalLM):
self.assistant_config = config self.assistant_config = config
self.config = text_config self.config = text_config
self.quant_config = quant_config self.quant_config = quant_config
self.pp_group = get_pp_group() self.pp_group = get_parallel().pp_group
self.vocab_size = text_config.vocab_size self.vocab_size = text_config.vocab_size
self.hidden_size = text_config.hidden_size self.hidden_size = text_config.hidden_size
python/sglang/srt/models/gemma4_unified.py
+2 -2
View File
@@ -40,7 +40,6 @@ import torch
from torch import nn from torch import nn
from transformers import PreTrainedModel from transformers import PreTrainedModel
from sglang.srt.distributed import get_pp_group
from sglang.srt.layers.layernorm import Gemma4RMSNorm from sglang.srt.layers.layernorm import Gemma4RMSNorm
from sglang.srt.layers.logits_processor import LogitsProcessor, LogitsProcessorOutput from sglang.srt.layers.logits_processor import LogitsProcessor, LogitsProcessorOutput
from sglang.srt.layers.quantization.base_config import QuantizationConfig from sglang.srt.layers.quantization.base_config import QuantizationConfig
@@ -53,6 +52,7 @@ from sglang.srt.managers.schedule_batch import (
from sglang.srt.model_loader.weight_utils import default_weight_loader from sglang.srt.model_loader.weight_utils import default_weight_loader
from sglang.srt.models.gemma4_causal import Gemma4TextModel, pp_filter_load_weight from sglang.srt.models.gemma4_causal import Gemma4TextModel, pp_filter_load_weight
from sglang.srt.models.gemma4_mm import Gemma4ForConditionalGeneration from sglang.srt.models.gemma4_mm import Gemma4ForConditionalGeneration
from sglang.srt.runtime_context import get_parallel
from sglang.srt.utils import add_prefix from sglang.srt.utils import add_prefix
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@@ -144,7 +144,7 @@ class Gemma4UnifiedForConditionalGeneration(Gemma4ForConditionalGeneration):
# Skip Gemma4ForConditionalGeneration.__init__ (it builds the SigLIP / # Skip Gemma4ForConditionalGeneration.__init__ (it builds the SigLIP /
# conformer towers we do not have) and initialise the HF base directly. # conformer towers we do not have) and initialise the HF base directly.
PreTrainedModel.__init__(self, config=config) PreTrainedModel.__init__(self, config=config)
self.pp_group = get_pp_group() self.pp_group = get_parallel().pp_group
self.config = config self.config = config
self.quant_config = quant_config self.quant_config = quant_config
python/sglang/srt/models/glm4.py
+2 -5
View File
@@ -23,9 +23,6 @@ from typing import Any, Dict, Iterable, Optional, Tuple, Union
import torch import torch
from torch import nn from torch import nn
from sglang.srt.distributed import (
get_pp_group,
)
from sglang.srt.layers.activation import SiluAndMul from sglang.srt.layers.activation import SiluAndMul
from sglang.srt.layers.dp_attention import is_dp_attention_enabled from sglang.srt.layers.dp_attention import is_dp_attention_enabled
from sglang.srt.layers.layernorm import RMSNorm from sglang.srt.layers.layernorm import RMSNorm
@@ -298,7 +295,7 @@ class Glm4Model(nn.Module):
self.config = config self.config = config
self.padding_idx = config.pad_token_id self.padding_idx = config.pad_token_id
self.vocab_size = config.vocab_size self.vocab_size = config.vocab_size
self.pp_group = get_pp_group() self.pp_group = get_parallel().pp_group
if self.pp_group.is_first_rank: if self.pp_group.is_first_rank:
self.embed_tokens = VocabParallelEmbedding( self.embed_tokens = VocabParallelEmbedding(
@@ -420,7 +417,7 @@ class Glm4ForCausalLM(nn.Module):
prefix: str = "", prefix: str = "",
) -> None: ) -> None:
super().__init__() super().__init__()
self.pp_group = get_pp_group() self.pp_group = get_parallel().pp_group
self.config = config self.config = config
self.quant_config = quant_config self.quant_config = quant_config
self.model = Glm4Model( self.model = Glm4Model(
python/sglang/srt/models/glm4_moe.py
+3 -5
View File
@@ -28,9 +28,7 @@ from sglang.kernels.ops.quantization.fp8_kernel import is_fp8_fnuz
from sglang.srt.batch_overlap.single_batch_overlap import SboFlags from sglang.srt.batch_overlap.single_batch_overlap import SboFlags
from sglang.srt.batch_overlap.two_batch_overlap import model_forward_maybe_tbo from sglang.srt.batch_overlap.two_batch_overlap import model_forward_maybe_tbo
from sglang.srt.distributed import ( from sglang.srt.distributed import (
get_pp_group,
get_pp_indices, get_pp_indices,
parallel_state,
tensor_model_parallel_all_reduce, tensor_model_parallel_all_reduce,
) )
from sglang.srt.distributed.device_communicators.pynccl_allocator import ( from sglang.srt.distributed.device_communicators.pynccl_allocator import (
@@ -646,7 +644,7 @@ class Glm4MoeSparseMoeBlock(nn.Module):
final_hidden_states *= self.routed_scaling_factor final_hidden_states *= self.routed_scaling_factor
if shared_output is not None: if shared_output is not None:
with use_symmetric_memory( with use_symmetric_memory(
parallel_state.get_tp_group(), disabled=not is_allocation_symmetric() get_parallel().tp_group, disabled=not is_allocation_symmetric()
): ):
final_hidden_states_out = torch.empty_like(final_hidden_states) final_hidden_states_out = torch.empty_like(final_hidden_states)
torch.add(final_hidden_states, shared_output, out=final_hidden_states_out) torch.add(final_hidden_states, shared_output, out=final_hidden_states_out)
@@ -1059,7 +1057,7 @@ class Glm4MoeModel(nn.Module):
prefix: str = "", prefix: str = "",
): ):
super().__init__() super().__init__()
self.pp_group = get_pp_group() self.pp_group = get_parallel().pp_group
self.config = config self.config = config
self.vocab_size = config.vocab_size self.vocab_size = config.vocab_size
self.first_k_dense_replace = config.first_k_dense_replace self.first_k_dense_replace = config.first_k_dense_replace
@@ -1185,7 +1183,7 @@ class Glm4MoeForCausalLM(nn.Module):
prefix: str = "", prefix: str = "",
) -> None: ) -> None:
nn.Module.__init__(self) nn.Module.__init__(self)
self.pp_group = get_pp_group() self.pp_group = get_parallel().pp_group
self.config = config self.config = config
self.tp_size = get_parallel().tp_size self.tp_size = get_parallel().tp_size
self.quant_config = quant_config self.quant_config = quant_config
python/sglang/srt/models/glm4_moe_lite.py
+3 -5
View File
@@ -26,8 +26,6 @@ from transformers import PretrainedConfig
from sglang.srt.batch_overlap.single_batch_overlap import SboFlags from sglang.srt.batch_overlap.single_batch_overlap import SboFlags
from sglang.srt.batch_overlap.two_batch_overlap import model_forward_maybe_tbo from sglang.srt.batch_overlap.two_batch_overlap import model_forward_maybe_tbo
from sglang.srt.distributed import ( from sglang.srt.distributed import (
get_pp_group,
parallel_state,
tensor_model_parallel_all_reduce, tensor_model_parallel_all_reduce,
) )
from sglang.srt.distributed.device_communicators.pynccl_allocator import ( from sglang.srt.distributed.device_communicators.pynccl_allocator import (
@@ -365,7 +363,7 @@ class Glm4MoeLiteSparseMoeBlock(nn.Module):
final_hidden_states *= self.routed_scaling_factor final_hidden_states *= self.routed_scaling_factor
if shared_output is not None: if shared_output is not None:
with use_symmetric_memory( with use_symmetric_memory(
parallel_state.get_tp_group(), disabled=not is_allocation_symmetric() get_parallel().tp_group, disabled=not is_allocation_symmetric()
): ):
final_hidden_states_out = torch.empty_like(final_hidden_states) final_hidden_states_out = torch.empty_like(final_hidden_states)
torch.add(final_hidden_states, shared_output, out=final_hidden_states_out) torch.add(final_hidden_states, shared_output, out=final_hidden_states_out)
@@ -760,7 +758,7 @@ class Glm4MoeLiteModel(nn.Module):
self.padding_id = config.pad_token_id self.padding_id = config.pad_token_id
self.vocab_size = config.vocab_size self.vocab_size = config.vocab_size
self.first_k_dense_replace = config.first_k_dense_replace self.first_k_dense_replace = config.first_k_dense_replace
self.pp_group = get_pp_group() self.pp_group = get_parallel().pp_group
if self.pp_group.is_first_rank: if self.pp_group.is_first_rank:
self.embed_tokens = VocabParallelEmbedding( self.embed_tokens = VocabParallelEmbedding(
@@ -892,7 +890,7 @@ class Glm4MoeLiteForCausalLM(nn.Module, DeepseekV2WeightLoaderMixin):
self.config = config self.config = config
self.tp_size = get_parallel().tp_size self.tp_size = get_parallel().tp_size
self.quant_config = quant_config self.quant_config = quant_config
self.pp_group = get_pp_group() self.pp_group = get_parallel().pp_group
self.determine_num_fused_shared_experts() self.determine_num_fused_shared_experts()
self.model = Glm4MoeLiteModel( self.model = Glm4MoeLiteModel(
config, quant_config, prefix=add_prefix("model", prefix) config, quant_config, prefix=add_prefix("model", prefix)
python/sglang/srt/models/glm4v.py
+1 -2
View File
@@ -27,7 +27,6 @@ import torch.nn.functional as F
from einops import rearrange from einops import rearrange
from transformers.models.glm4v.configuration_glm4v import Glm4vConfig, Glm4vVisionConfig from transformers.models.glm4v.configuration_glm4v import Glm4vConfig, Glm4vVisionConfig
from sglang.srt.distributed.parallel_state import get_pp_group
from sglang.srt.layers.activation import SiluAndMul from sglang.srt.layers.activation import SiluAndMul
from sglang.srt.layers.attention import vision_utils from sglang.srt.layers.attention import vision_utils
from sglang.srt.layers.attention.vision import ( from sglang.srt.layers.attention.vision import (
@@ -556,7 +555,7 @@ class Glm4vForConditionalGeneration(nn.Module):
) -> None: ) -> None:
super().__init__() super().__init__()
self.pp_group = get_pp_group() self.pp_group = get_parallel().pp_group
self.config = config self.config = config
self.use_data_parallel = get_mm().mm_enable_dp_encoder self.use_data_parallel = get_mm().mm_enable_dp_encoder
vision_utils.update_vit_attn_dummy_heads_config(self.config) vision_utils.update_vit_attn_dummy_heads_config(self.config)
python/sglang/srt/models/glm4v_moe.py
+1 -2
View File
@@ -6,7 +6,6 @@ import torch
import torch.nn as nn import torch.nn as nn
from transformers.models.glm4v_moe.configuration_glm4v_moe import Glm4vMoeConfig from transformers.models.glm4v_moe.configuration_glm4v_moe import Glm4vMoeConfig
from sglang.srt.distributed.parallel_state import get_pp_group
from sglang.srt.layers.attention import vision_utils from sglang.srt.layers.attention import vision_utils
from sglang.srt.layers.logits_processor import LogitsProcessor from sglang.srt.layers.logits_processor import LogitsProcessor
from sglang.srt.layers.moe import get_moe_a2a_backend from sglang.srt.layers.moe import get_moe_a2a_backend
@@ -40,7 +39,7 @@ class Glm4vMoeForConditionalGeneration(Glm4vForConditionalGeneration):
) -> None: ) -> None:
nn.Module.__init__(self) nn.Module.__init__(self)
self.pp_group = get_pp_group() self.pp_group = get_parallel().pp_group
self.config = config self.config = config
self.use_data_parallel = get_mm().mm_enable_dp_encoder self.use_data_parallel = get_mm().mm_enable_dp_encoder
vision_utils.update_vit_attn_dummy_heads_config(self.config) vision_utils.update_vit_attn_dummy_heads_config(self.config)
python/sglang/srt/models/glm5_next.py
+2 -3
View File
@@ -16,7 +16,6 @@ from sglang.srt.batch_overlap.two_batch_overlap import (
) )
from sglang.srt.configs.glm5_next import Glm5NextConfig, Glm5NextTextConfig from sglang.srt.configs.glm5_next import Glm5NextConfig, Glm5NextTextConfig
from sglang.srt.configs.model_config import is_deepseek_dsa from sglang.srt.configs.model_config import is_deepseek_dsa
from sglang.srt.distributed.parallel_state import get_pp_group
from sglang.srt.distributed.utils import divide from sglang.srt.distributed.utils import divide
from sglang.srt.environ import envs from sglang.srt.environ import envs
from sglang.srt.eplb.expert_distribution import ( from sglang.srt.eplb.expert_distribution import (
@@ -861,7 +860,7 @@ class Glm5NextModel(nn.Module):
self.padding_id = config.pad_token_id self.padding_id = config.pad_token_id
self.vocab_size = config.vocab_size self.vocab_size = config.vocab_size
self.first_k_dense_replace = config.first_k_dense_replace self.first_k_dense_replace = config.first_k_dense_replace
self.pp_group = get_pp_group() self.pp_group = get_parallel().pp_group
if self.pp_group.is_first_rank: if self.pp_group.is_first_rank:
self.embed_tokens = VocabParallelEmbedding( self.embed_tokens = VocabParallelEmbedding(
@@ -1123,7 +1122,7 @@ class Glm5NextForConditionalGeneration(nn.Module):
and getattr(text_config, "q_lora_rank", None) is not None and getattr(text_config, "q_lora_rank", None) is not None
) )
self.pp_group = get_pp_group() self.pp_group = get_parallel().pp_group
self.config = text_config self.config = text_config
self.tp_size = get_parallel().tp_size self.tp_size = get_parallel().tp_size
self.quant_config = quant_config self.quant_config = quant_config
python/sglang/srt/models/glm_ocr.py
+2 -3
View File
@@ -29,7 +29,6 @@ from transformers.models.glm_ocr.configuration_glm_ocr import (
GlmOcrVisionConfig, GlmOcrVisionConfig,
) )
from sglang.srt.distributed.parallel_state import get_pp_group
from sglang.srt.layers.attention import vision_utils from sglang.srt.layers.attention import vision_utils
from sglang.srt.layers.attention.vision import ( from sglang.srt.layers.attention.vision import (
VisionAttention, VisionAttention,
@@ -53,7 +52,7 @@ from sglang.srt.models.glm4v import (
Glm4vVisionModel, Glm4vVisionModel,
Glm4vVisionPatchEmbed, Glm4vVisionPatchEmbed,
) )
from sglang.srt.runtime_context import get_mm from sglang.srt.runtime_context import get_mm, get_parallel
from sglang.srt.utils import add_prefix from sglang.srt.utils import add_prefix
from sglang.srt.utils.hf_transformers_utils import get_processor from sglang.srt.utils.hf_transformers_utils import get_processor
@@ -285,7 +284,7 @@ class GlmOcrForConditionalGeneration(Glm4vForConditionalGeneration):
) -> None: ) -> None:
super().__init__(config, quant_config, prefix) super().__init__(config, quant_config, prefix)
self.pp_group = get_pp_group() self.pp_group = get_parallel().pp_group
self.config = config self.config = config
self.use_data_parallel = get_mm().mm_enable_dp_encoder self.use_data_parallel = get_mm().mm_enable_dp_encoder
self.visual = GlmOcrVisionModel( self.visual = GlmOcrVisionModel(
python/sglang/srt/models/gpt_oss.py
+2 -3
View File
@@ -28,7 +28,6 @@ from transformers import PretrainedConfig
from sglang.kernels.jit.utils import is_arch_support_pdl from sglang.kernels.jit.utils import is_arch_support_pdl
from sglang.srt.distributed import ( from sglang.srt.distributed import (
get_pp_group,
tensor_model_parallel_all_reduce, tensor_model_parallel_all_reduce,
) )
from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder
@@ -680,7 +679,7 @@ class GptOssModel(nn.Module):
super().__init__() super().__init__()
self.padding_idx = config.pad_token_id self.padding_idx = config.pad_token_id
self.vocab_size = config.vocab_size self.vocab_size = config.vocab_size
self.pp_group = get_pp_group() self.pp_group = get_parallel().pp_group
if _is_npu: if _is_npu:
config.hidden_act = "npu_swiglu_oai" config.hidden_act = "npu_swiglu_oai"
@@ -790,7 +789,7 @@ class GptOssForCausalLM(nn.Module):
prefix: str = "", prefix: str = "",
) -> None: ) -> None:
super().__init__() super().__init__()
self.pp_group = get_pp_group() self.pp_group = get_parallel().pp_group
self.config = config self.config = config
self.quant_config = quant_config self.quant_config = quant_config
self.model = GptOssModel( self.model = GptOssModel(
python/sglang/srt/models/granitemoehybrid.py
+2 -3
View File
@@ -4,7 +4,6 @@ import torch
from torch import nn from torch import nn
from sglang.srt.configs.granitemoehybrid import GraniteMoeHybridConfig from sglang.srt.configs.granitemoehybrid import GraniteMoeHybridConfig
from sglang.srt.distributed import get_pp_group
from sglang.srt.layers.attention.hybrid_linear_attn_backend import ( from sglang.srt.layers.attention.hybrid_linear_attn_backend import (
HybridLinearAttnBackend, HybridLinearAttnBackend,
Mamba2AttnBackend, Mamba2AttnBackend,
@@ -327,7 +326,7 @@ class GraniteMoeHybridModel(nn.Module):
self.vocab_size = config.vocab_size self.vocab_size = config.vocab_size
self.pp_group = get_pp_group() self.pp_group = get_parallel().pp_group
if self.pp_group.is_first_rank: if self.pp_group.is_first_rank:
self.embed_tokens = VocabParallelEmbedding( self.embed_tokens = VocabParallelEmbedding(
@@ -443,7 +442,7 @@ class GraniteMoeHybridForCausalLM(
super().__init__() super().__init__()
self.capture_aux_hidden_states = False self.capture_aux_hidden_states = False
self.pp_group = get_pp_group() self.pp_group = get_parallel().pp_group
self.quant_config = quant_config self.quant_config = quant_config
self.config = config self.config = config
python/sglang/srt/models/hunyuan_v4.py
+2 -3
View File
@@ -6,7 +6,6 @@ import torch
from torch import nn from torch import nn
from transformers import PretrainedConfig from transformers import PretrainedConfig
from sglang.srt.distributed import get_pp_group
from sglang.srt.layers.attention.index_topk_share import IndexTopKShareState from sglang.srt.layers.attention.index_topk_share import IndexTopKShareState
from sglang.srt.layers.communicator import AttentionInputs, get_attn_tp_context from sglang.srt.layers.communicator import AttentionInputs, get_attn_tp_context
from sglang.srt.layers.layernorm import RMSNorm from sglang.srt.layers.layernorm import RMSNorm
@@ -614,7 +613,7 @@ class HYV4DecoderLayer(nn.Module):
class HYV4Model(nn.Module): class HYV4Model(nn.Module):
def __init__(self, config, quant_config=None, prefix=""): def __init__(self, config, quant_config=None, prefix=""):
super().__init__() super().__init__()
if get_pp_group().world_size != 1: if get_parallel().pp_group.world_size != 1:
raise ValueError("HYV4 pipeline parallelism is not supported") raise ValueError("HYV4 pipeline parallelism is not supported")
self.config = config self.config = config
self.start_layer = 0 self.start_layer = 0
@@ -675,7 +674,7 @@ class HYV4ForCausalLM(nn.Module, DeepseekV2WeightLoaderMixin):
super().__init__() super().__init__()
self.config = config self.config = config
self.quant_config = quant_config self.quant_config = quant_config
self.pp_group = get_pp_group() self.pp_group = get_parallel().pp_group
self.model = HYV4Model(config, quant_config, f"{prefix}.model") self.model = HYV4Model(config, quant_config, f"{prefix}.model")
self.num_fused_shared_experts = max( self.num_fused_shared_experts = max(
( (
python/sglang/srt/models/hunyuan_v4_nextn.py
+1 -2
View File
@@ -4,7 +4,6 @@ from typing import Iterable, Tuple
import torch import torch
from torch import nn from torch import nn
from sglang.srt.distributed import get_pp_group
from sglang.srt.layers.attention.index_topk_share import IndexTopKShareState from sglang.srt.layers.attention.index_topk_share import IndexTopKShareState
from sglang.srt.layers.communicator import AttentionInputs, get_attn_tp_context from sglang.srt.layers.communicator import AttentionInputs, get_attn_tp_context
from sglang.srt.layers.layernorm import RMSNorm from sglang.srt.layers.layernorm import RMSNorm
@@ -170,7 +169,7 @@ class HYV4ForCausalLMNextN(nn.Module, DeepseekV2WeightLoaderMixin):
super().__init__() super().__init__()
self.config = config self.config = config
self.quant_config = quant_config self.quant_config = quant_config
self.pp_group = get_pp_group() self.pp_group = get_parallel().pp_group
nextn_quant_config = _mtp_quant_config(quant_config) nextn_quant_config = _mtp_quant_config(quant_config)
self.model = HYV4ModelNextN( self.model = HYV4ModelNextN(
config, nextn_quant_config, prefix=f"{prefix}.model" config, nextn_quant_config, prefix=f"{prefix}.model"
python/sglang/srt/models/interns2_mobius.py
+2 -2
View File
@@ -10,7 +10,7 @@ from sglang.srt.configs.interns2_mobius import (
InternS2MobiusConfig, InternS2MobiusConfig,
InternS2MobiusTextConfig, InternS2MobiusTextConfig,
) )
from sglang.srt.distributed import get_pp_group, tensor_model_parallel_all_reduce from sglang.srt.distributed import tensor_model_parallel_all_reduce
from sglang.srt.layers.communicator import LayerCommunicator, LayerScatterModes from sglang.srt.layers.communicator import LayerCommunicator, LayerScatterModes
from sglang.srt.layers.dp_attention import is_dp_attention_enabled from sglang.srt.layers.dp_attention import is_dp_attention_enabled
from sglang.srt.layers.layernorm import GemmaRMSNorm from sglang.srt.layers.layernorm import GemmaRMSNorm
@@ -695,7 +695,7 @@ class InternS2MobiusForCausalLM(Qwen3_5ForCausalLM):
nn.Module.__init__(self) nn.Module.__init__(self)
self.config = config self.config = config
self.hidden_size = config.hidden_size self.hidden_size = config.hidden_size
self.pp_group = get_pp_group() self.pp_group = get_parallel().pp_group
if self.pp_group.world_size != 1: if self.pp_group.world_size != 1:
raise ValueError( raise ValueError(
"Intern-S2-Mobius baseline does not support pipeline parallelism" "Intern-S2-Mobius baseline does not support pipeline parallelism"
python/sglang/srt/models/kimi_k25_eagle3.py
+2 -2
View File
@@ -24,7 +24,6 @@ import torch
from torch import nn from torch import nn
from transformers import PretrainedConfig from transformers import PretrainedConfig
from sglang.srt.distributed import get_pp_group
from sglang.srt.distributed.device_communicators import triton_symm_mem_ag from sglang.srt.distributed.device_communicators import triton_symm_mem_ag
from sglang.srt.layers.communicator import AttentionInputs, get_attn_tp_context from sglang.srt.layers.communicator import AttentionInputs, get_attn_tp_context
from sglang.srt.layers.layernorm import RMSNorm from sglang.srt.layers.layernorm import RMSNorm
@@ -39,6 +38,7 @@ from sglang.srt.layers.vocab_parallel_embedding import (
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors
from sglang.srt.model_loader.weight_utils import default_weight_loader from sglang.srt.model_loader.weight_utils import default_weight_loader
from sglang.srt.models.deepseek_v2 import DeepseekV2AttentionMLA, DeepseekV2MLP from sglang.srt.models.deepseek_v2 import DeepseekV2AttentionMLA, DeepseekV2MLP
from sglang.srt.runtime_context import get_parallel
from sglang.srt.utils import BumpAllocator, add_prefix from sglang.srt.utils import BumpAllocator, add_prefix
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@@ -336,7 +336,7 @@ class Eagle3DeepseekV2ForCausalLM(nn.Module):
) )
quant_config = None quant_config = None
self.quant_config = quant_config self.quant_config = quant_config
self.pp_group = get_pp_group() self.pp_group = get_parallel().pp_group
self.model = Eagle3MLAModel( self.model = Eagle3MLAModel(
config, quant_config=quant_config, prefix=add_prefix("model", prefix) config, quant_config=quant_config, prefix=add_prefix("model", prefix)
python/sglang/srt/models/kimi_k3.py
+9 -11
View File
@@ -21,9 +21,7 @@ from sglang.srt.configs.kimi_k3 import KimiK3Config
from sglang.srt.configs.kimi_linear import KimiLinearConfig from sglang.srt.configs.kimi_linear import KimiLinearConfig
from sglang.srt.distributed import ( from sglang.srt.distributed import (
divide, divide,
get_pp_group,
get_shared_experts_tp_group, get_shared_experts_tp_group,
get_tp_group,
tensor_model_parallel_all_reduce, tensor_model_parallel_all_reduce,
) )
from sglang.srt.distributed.device_communicators.pynccl_allocator import ( from sglang.srt.distributed.device_communicators.pynccl_allocator import (
@@ -257,7 +255,7 @@ def _dp_local_buffer_group():
CommunicateSummableTensorPairFn._scatter_hidden_states).""" CommunicateSummableTensorPairFn._scatter_hidden_states)."""
parallel = get_parallel() parallel = get_parallel()
if parallel.tp_size == parallel.attn_dp_size: if parallel.tp_size == parallel.attn_dp_size:
return get_tp_group() return get_parallel().tp_group
return parallel.attn_tp_group return parallel.attn_tp_group
@@ -361,7 +359,7 @@ class KimiK3MLP(nn.Module):
) )
if use_dp: if use_dp:
local_hidden_states = hidden_states local_hidden_states = hidden_states
hidden_states = get_global_dp_buffer(get_tp_group()) hidden_states = get_global_dp_buffer(get_parallel().tp_group)
dp_gather_replicate(hidden_states, local_hidden_states, forward_batch) dp_gather_replicate(hidden_states, local_hidden_states, forward_batch)
gate_up, _ = self.gate_up_proj(hidden_states) gate_up, _ = self.gate_up_proj(hidden_states)
hidden_states = self.act_fn(gate_up) hidden_states = self.act_fn(gate_up)
@@ -796,12 +794,12 @@ class KimiK3MoE(nn.Module):
import deep_gemm import deep_gemm
from sglang.kernels.ops.attention.dsv4 import mega_moe_pre_dispatch from sglang.kernels.ops.attention.dsv4 import mega_moe_pre_dispatch
from sglang.srt.distributed.parallel_state import get_moe_ep_group
from sglang.srt.environ import envs from sglang.srt.environ import envs
from sglang.srt.layers.moe.mega_moe import ( from sglang.srt.layers.moe.mega_moe import (
_configure_mega_moe_deep_gemm_num_sms, _configure_mega_moe_deep_gemm_num_sms,
_get_mega_moe_symm_buffer, _get_mega_moe_symm_buffer,
) )
from sglang.srt.runtime_context import get_parallel
# In SP-MoE mode (KimiK3DecoderLayer reduce-scatters the o_proj # In SP-MoE mode (KimiK3DecoderLayer reduce-scatters the o_proj
# output) the incoming rows are already this rank's token shard, so # output) the incoming rows are already this rank's token shard, so
@@ -819,7 +817,7 @@ class KimiK3MoE(nn.Module):
f"the env var to cover the per-rank rows" f"the env var to cover the per-rank rows"
) )
buf = _get_mega_moe_symm_buffer( buf = _get_mega_moe_symm_buffer(
get_moe_ep_group().device_group, get_parallel().moe_ep_group.device_group,
num_experts=self.experts.num_experts, num_experts=self.experts.num_experts,
num_max_tokens_per_rank=num_max_tokens_per_rank, num_max_tokens_per_rank=num_max_tokens_per_rank,
num_topk=self._mega_top_k, num_topk=self._mega_top_k,
@@ -1389,7 +1387,7 @@ class KimiK3MoE(nn.Module):
).view(-1) ).view(-1)
else: else:
with use_symmetric_memory( with use_symmetric_memory(
get_tp_group(), disabled=not is_allocation_symmetric() get_parallel().tp_group, disabled=not is_allocation_symmetric()
): ):
buf = hidden_states.new_empty(latent_numel + num_tokens * hidden_size) buf = hidden_states.new_empty(latent_numel + num_tokens * hidden_size)
@@ -1505,7 +1503,7 @@ class KimiK3MoE(nn.Module):
use_dp = self._dp_attention and forward_batch is not None and not self._ep_a2a use_dp = self._dp_attention and forward_batch is not None and not self._ep_a2a
if use_dp: if use_dp:
local_hidden_states = hidden_states local_hidden_states = hidden_states
hidden_states = get_global_dp_buffer(get_tp_group()) hidden_states = get_global_dp_buffer(get_parallel().tp_group)
dp_gather_replicate(hidden_states, local_hidden_states, forward_batch) dp_gather_replicate(hidden_states, local_hidden_states, forward_batch)
dp_prefix_sum, prefix_sum = prefix_sum, None dp_prefix_sum, prefix_sum = prefix_sum, None
if hidden_states.shape[0] > 0 and self._eligible_for_fused_front: if hidden_states.shape[0] > 0 and self._eligible_for_fused_front:
@@ -2919,7 +2917,7 @@ class KimiK3LinearModel(nn.Module):
): ):
super().__init__() super().__init__()
self.config = config self.config = config
self.pp_group = get_pp_group() self.pp_group = get_parallel().pp_group
self.dspark_layers_to_capture: Optional[list[int]] = None self.dspark_layers_to_capture: Optional[list[int]] = None
self._dp_attention = is_dp_attention_enabled() self._dp_attention = is_dp_attention_enabled()
self._trim_padded_attn = require_mlp_sync() self._trim_padded_attn = require_mlp_sync()
@@ -2990,7 +2988,7 @@ class KimiK3LinearModel(nn.Module):
inputs_embeds: torch.Tensor | None = None, inputs_embeds: torch.Tensor | None = None,
pp_proxy_tensors: Optional[PPProxyTensors] = None, pp_proxy_tensors: Optional[PPProxyTensors] = None,
) -> torch.Tensor: ) -> torch.Tensor:
if get_pp_group().is_first_rank: if get_parallel().pp_group.is_first_rank:
if inputs_embeds is not None: if inputs_embeds is not None:
hidden_states = inputs_embeds hidden_states = inputs_embeds
else: else:
@@ -3188,7 +3186,7 @@ class KimiK3LinearForCausalLM(nn.Module):
self.model = KimiK3LinearModel( self.model = KimiK3LinearModel(
config, quant_config, prefix=maybe_prefix(prefix, "model") config, quant_config, prefix=maybe_prefix(prefix, "model")
) )
self.pp_group = get_pp_group() self.pp_group = get_parallel().pp_group
if self.pp_group.is_last_rank: if self.pp_group.is_last_rank:
self.lm_head = ParallelLMHead( self.lm_head = ParallelLMHead(
config.vocab_size, config.vocab_size,
python/sglang/srt/models/kimi_linear.py
+3 -4
View File
@@ -12,7 +12,6 @@ from sglang.kernels.ops.attention.fla.fused_norm_gate import FusedRMSNormGated
from sglang.srt.configs.kimi_linear import KimiLinearConfig from sglang.srt.configs.kimi_linear import KimiLinearConfig
from sglang.srt.distributed import ( from sglang.srt.distributed import (
divide, divide,
get_pp_group,
tensor_model_parallel_all_reduce, tensor_model_parallel_all_reduce,
) )
from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder
@@ -653,7 +652,7 @@ class KimiLinearModel(nn.Module):
self.padding_idx = config.pad_token_id self.padding_idx = config.pad_token_id
self.vocab_size = config.vocab_size self.vocab_size = config.vocab_size
self.pp_group = get_pp_group() self.pp_group = get_parallel().pp_group
self.dspark_layers_to_capture: Optional[list[int]] = None self.dspark_layers_to_capture: Optional[list[int]] = None
if self.pp_group.is_first_rank: if self.pp_group.is_first_rank:
@@ -699,7 +698,7 @@ class KimiLinearModel(nn.Module):
inputs_embeds: torch.Tensor | None = None, inputs_embeds: torch.Tensor | None = None,
pp_proxy_tensors: Optional[PPProxyTensors] = None, pp_proxy_tensors: Optional[PPProxyTensors] = None,
) -> torch.Tensor: ) -> torch.Tensor:
if get_pp_group().is_first_rank: if get_parallel().pp_group.is_first_rank:
if inputs_embeds is not None: if inputs_embeds is not None:
hidden_states = inputs_embeds hidden_states = inputs_embeds
else: else:
@@ -771,7 +770,7 @@ class KimiLinearForCausalLM(nn.Module):
) )
self.start_layer = self.model.start_layer self.start_layer = self.model.start_layer
self.end_layer = self.model.end_layer self.end_layer = self.model.end_layer
self.pp_group = get_pp_group() self.pp_group = get_parallel().pp_group
if self.pp_group.is_last_rank: if self.pp_group.is_last_rank:
self.lm_head = ParallelLMHead( self.lm_head = ParallelLMHead(
self.config.vocab_size, self.config.vocab_size,
python/sglang/srt/models/laguna.py
+2 -3
View File
@@ -18,7 +18,6 @@ from torch import nn
from sglang.srt.configs.laguna import LagunaConfig, normalize_gating from sglang.srt.configs.laguna import LagunaConfig, normalize_gating
from sglang.srt.distributed import ( from sglang.srt.distributed import (
get_pp_group,
tensor_model_parallel_all_reduce, tensor_model_parallel_all_reduce,
) )
from sglang.srt.environ import envs from sglang.srt.environ import envs
@@ -534,7 +533,7 @@ class LagunaModel(nn.Module):
self.config = config self.config = config
self.padding_idx = getattr(config, "pad_token_id", None) self.padding_idx = getattr(config, "pad_token_id", None)
self.vocab_size = config.vocab_size self.vocab_size = config.vocab_size
self.pp_group = get_pp_group() self.pp_group = get_parallel().pp_group
if self.pp_group.is_first_rank: if self.pp_group.is_first_rank:
self.embed_tokens = VocabParallelEmbedding( self.embed_tokens = VocabParallelEmbedding(
@@ -656,7 +655,7 @@ class LagunaForCausalLM(nn.Module):
prefix: str = "", prefix: str = "",
) -> None: ) -> None:
super().__init__() super().__init__()
self.pp_group = get_pp_group() self.pp_group = get_parallel().pp_group
self.config = config self.config = config
self.model = LagunaModel( self.model = LagunaModel(
config, quant_config=quant_config, prefix=add_prefix("model", prefix) config, quant_config=quant_config, prefix=add_prefix("model", prefix)
python/sglang/srt/models/lfm2.py
+1 -2
View File
@@ -22,7 +22,6 @@ from sglang.kernels.ops.mamba.causal_conv1d_triton import (
causal_conv1d_update as causal_conv1d_update_triton, causal_conv1d_update as causal_conv1d_update_triton,
) )
from sglang.srt.configs.lfm2 import Lfm2Config from sglang.srt.configs.lfm2 import Lfm2Config
from sglang.srt.distributed import get_pp_group
from sglang.srt.layers.attention.mamba.causal_conv1d import ( from sglang.srt.layers.attention.mamba.causal_conv1d import (
causal_conv1d_fn, causal_conv1d_fn,
causal_conv1d_update, causal_conv1d_update,
@@ -688,7 +687,7 @@ class Lfm2ForCausalLM(nn.Module):
) -> None: ) -> None:
super().__init__() super().__init__()
self.config = config self.config = config
self.pp_group = get_pp_group() self.pp_group = get_parallel().pp_group
assert self.pp_group.is_first_rank and self.pp_group.is_last_rank assert self.pp_group.is_first_rank and self.pp_group.is_last_rank
self.quant_config = quant_config self.quant_config = quant_config
python/sglang/srt/models/lfm2_moe.py
+1 -2
View File
@@ -23,7 +23,6 @@ from sglang.kernels.ops.mamba.lfm_short_conv import (
fused_lfm_short_conv_prefill, fused_lfm_short_conv_prefill,
) )
from sglang.srt.configs.lfm2_moe import Lfm2MoeConfig from sglang.srt.configs.lfm2_moe import Lfm2MoeConfig
from sglang.srt.distributed import get_pp_group
from sglang.srt.layers.activation import SiluAndMul from sglang.srt.layers.activation import SiluAndMul
from sglang.srt.layers.attention.mamba.causal_conv1d import ( from sglang.srt.layers.attention.mamba.causal_conv1d import (
causal_conv1d_fn, causal_conv1d_fn,
@@ -598,7 +597,7 @@ class Lfm2MoeForCausalLM(nn.Module):
) -> None: ) -> None:
super().__init__() super().__init__()
self.config = config self.config = config
self.pp_group = get_pp_group() self.pp_group = get_parallel().pp_group
assert self.pp_group.is_first_rank and self.pp_group.is_last_rank assert self.pp_group.is_first_rank and self.pp_group.is_last_rank
self.quant_config = quant_config self.quant_config = quant_config
python/sglang/srt/models/llada2.py
+3 -5
View File
@@ -28,8 +28,6 @@ from torch import nn
from transformers import PretrainedConfig from transformers import PretrainedConfig
from sglang.srt.distributed import ( from sglang.srt.distributed import (
get_pp_group,
parallel_state,
tensor_model_parallel_all_reduce, tensor_model_parallel_all_reduce,
) )
from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder
@@ -309,7 +307,7 @@ class LLaDA2MoeSparseMoeBlock(nn.Module):
self.ep_size = get_parallel().tp_size self.ep_size = get_parallel().tp_size
self.deepep_dispatcher = DeepEPDispatcher( self.deepep_dispatcher = DeepEPDispatcher(
group=parallel_state.get_tp_group().device_group, group=get_parallel().tp_group.device_group,
router_topk=self.top_k, router_topk=self.top_k,
permute_fusion=True, permute_fusion=True,
num_experts=self.num_experts, num_experts=self.num_experts,
@@ -717,7 +715,7 @@ class LLaDA2MoeModel(nn.Module):
prefix: str = "", prefix: str = "",
): ):
super().__init__() super().__init__()
self.pp_group = get_pp_group() self.pp_group = get_parallel().pp_group
self.config = config self.config = config
self.vocab_size = config.vocab_size self.vocab_size = config.vocab_size
self.embed_dim = config.hidden_size self.embed_dim = config.hidden_size
@@ -804,7 +802,7 @@ class LLaDA2MoeModelLM(nn.Module):
prefix: str = "", prefix: str = "",
): ):
super().__init__() super().__init__()
self.pp_group = get_pp_group() self.pp_group = get_parallel().pp_group
self.config = config self.config = config
self.quant_config = quant_config self.quant_config = quant_config
alt_stream = get_stream("alt") if _is_cuda else None alt_stream = get_stream("alt") if _is_cuda else None
python/sglang/srt/models/llama.py
+2 -3
View File
@@ -26,7 +26,6 @@ from torch import nn
from transformers import LlamaConfig from transformers import LlamaConfig
from sglang.srt.distributed import ( from sglang.srt.distributed import (
get_pp_group,
get_pp_indices, get_pp_indices,
) )
from sglang.srt.layers.activation import SiluAndMul from sglang.srt.layers.activation import SiluAndMul
@@ -379,7 +378,7 @@ class LlamaModel(nn.Module):
self.config = config self.config = config
self.padding_idx = config.pad_token_id self.padding_idx = config.pad_token_id
self.vocab_size = config.vocab_size self.vocab_size = config.vocab_size
self.pp_group = get_pp_group() self.pp_group = get_parallel().pp_group
if self.pp_group.is_first_rank: if self.pp_group.is_first_rank:
self.embed_tokens = VocabParallelEmbedding( self.embed_tokens = VocabParallelEmbedding(
config.vocab_size, config.vocab_size,
@@ -521,7 +520,7 @@ class LlamaForCausalLM(nn.Module):
prefix: str = "", prefix: str = "",
) -> None: ) -> None:
super().__init__() super().__init__()
self.pp_group = get_pp_group() self.pp_group = get_parallel().pp_group
self.config = config self.config = config
self.quant_config = quant_config self.quant_config = quant_config
self.model = self._init_model(config, quant_config, add_prefix("model", prefix)) self.model = self._init_model(config, quant_config, add_prefix("model", prefix))
python/sglang/srt/models/llama_eagle.py
+2 -2
View File
@@ -25,7 +25,6 @@ import torch
from torch import nn from torch import nn
from transformers import LlamaConfig from transformers import LlamaConfig
from sglang.srt.distributed import get_pp_group
from sglang.srt.layers.logits_processor import LogitsProcessor from sglang.srt.layers.logits_processor import LogitsProcessor
from sglang.srt.layers.quantization.base_config import QuantizationConfig from sglang.srt.layers.quantization.base_config import QuantizationConfig
from sglang.srt.layers.vocab_parallel_embedding import ( from sglang.srt.layers.vocab_parallel_embedding import (
@@ -34,6 +33,7 @@ from sglang.srt.layers.vocab_parallel_embedding import (
) )
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors
from sglang.srt.models.llama import LlamaDecoderLayer, LlamaForCausalLM from sglang.srt.models.llama import LlamaDecoderLayer, LlamaForCausalLM
from sglang.srt.runtime_context import get_parallel
class LlamaDecoderLayer(LlamaDecoderLayer): class LlamaDecoderLayer(LlamaDecoderLayer):
@@ -120,7 +120,7 @@ class LlamaForCausalLMEagle(LlamaForCausalLM):
nn.Module.__init__(self) nn.Module.__init__(self)
self.config = config self.config = config
self.quant_config = quant_config self.quant_config = quant_config
self.pp_group = get_pp_group() self.pp_group = get_parallel().pp_group
self.model = LlamaModel( self.model = LlamaModel(
config, quant_config=quant_config, prefix=add_prefix("model", prefix) config, quant_config=quant_config, prefix=add_prefix("model", prefix)
) )

Some files were not shown because too many files have changed in this diff Show More