config: a parallel leaf with no live counterpart is read bare (#36620)

Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
Cheng Wan
2026-08-27 12:56:11 -07:00
committed by GitHub
co-authored by Claude Opus 5
parent 7c3b5a6732
commit ca1d7ed8e6
125 changed files with 343 additions and 405 deletions
+1 -1
View File
@@ -548,7 +548,7 @@ def _maybe_prepare_mlp_sync_batch(batch: ScheduleBatch, model_runner):
prepare_mlp_sync_batch_raw( prepare_mlp_sync_batch_raw(
batch, batch,
model_runner=model_runner, model_runner=model_runner,
dp_size=get_parallel().config.dp_size, dp_size=get_parallel().dp_size,
attn_tp_size=get_parallel().attn_tp_size, attn_tp_size=get_parallel().attn_tp_size,
attn_cp_size=model_runner.ps.attn_cp_size, attn_cp_size=model_runner.ps.attn_cp_size,
tp_group=model_runner.tp_group, tp_group=model_runner.tp_group,
+1 -1
View File
@@ -65,7 +65,7 @@ async def warm_up_compile(
disaggregation_mode: str, tokenizer_manager: TokenizerManager disaggregation_mode: str, tokenizer_manager: TokenizerManager
): ):
print("\nGenerate warm up request for compiling DeepGEMM...\n") print("\nGenerate warm up request for compiling DeepGEMM...\n")
dp_size = get_parallel().config.dp_size dp_size = get_parallel().dp_size
base_ids = [0, 1, 2, 3] base_ids = [0, 1, 2, 3]
sampling_params = { sampling_params = {
"temperature": 0.0, "temperature": 0.0,
@@ -787,7 +787,7 @@ class TboForwardBatchPreparer:
# TODO improve, e.g. unify w/ `init_raw` # TODO improve, e.g. unify w/ `init_raw`
if ( if (
get_parallel().config.moe_dense_tp_size == 1 get_parallel().moe_dense_tp_size == 1
and batch.global_dp_buffer_len is not None and batch.global_dp_buffer_len is not None
): ):
sum_len = end_token_index - start_token_index sum_len = end_token_index - start_token_index
+12 -16
View File
@@ -170,7 +170,7 @@ class CommonKVManager(BaseKVManager):
# for p/d multi node infer # for p/d multi node infer
self.bootstrap_host = get_serving().host self.bootstrap_host = get_serving().host
self.bootstrap_port = get_disagg().disaggregation_bootstrap_port self.bootstrap_port = get_disagg().disaggregation_bootstrap_port
self.dist_init_addr = get_parallel().config.dist_init_addr self.dist_init_addr = get_parallel().dist_init_addr
parallel = get_parallel() parallel = get_parallel()
self.attn_tp_size = parallel.attn_tp_size self.attn_tp_size = parallel.attn_tp_size
self.attn_tp_rank = parallel.attn_tp_rank self.attn_tp_rank = parallel.attn_tp_rank
@@ -181,9 +181,7 @@ class CommonKVManager(BaseKVManager):
self.attn_dp_size = get_attention_dp_size() self.attn_dp_size = get_attention_dp_size()
self.attn_dp_rank = get_attention_dp_rank() self.attn_dp_rank = get_attention_dp_rank()
self.system_dp_size = ( self.system_dp_size = (
1 1 if get_parallel().enable_dp_attention else get_parallel().dp_size
if get_parallel().config.enable_dp_attention
else get_parallel().config.dp_size
) )
self.system_dp_rank = ( self.system_dp_rank = (
self.kv_args.system_dp_rank if self.kv_args.system_dp_rank else 0 self.kv_args.system_dp_rank if self.kv_args.system_dp_rank else 0
@@ -192,8 +190,7 @@ class CommonKVManager(BaseKVManager):
self.pp_rank = self.kv_args.pp_rank self.pp_rank = self.kv_args.pp_rank
self.local_ip = get_local_ip_auto() self.local_ip = get_local_ip_auto()
cp_sharded_prefill = self.attn_cp_size > 1 and ( cp_sharded_prefill = self.attn_cp_size > 1 and (
self.is_hybrid_mla_backend self.is_hybrid_mla_backend or get_parallel().enable_dsa_cache_layer_split
or get_parallel().config.enable_dsa_cache_layer_split
) )
hybrid_decode_pulls_all_ranks = ( hybrid_decode_pulls_all_ranks = (
@@ -306,7 +303,7 @@ class CommonKVManager(BaseKVManager):
return ( return (
self.attn_cp_size > 1 self.attn_cp_size > 1
and self.attn_cp_rank != 0 and self.attn_cp_rank != 0
and not get_parallel().config.enable_dsa_cache_layer_split and not get_parallel().enable_dsa_cache_layer_split
) )
def requires_dcp_relayout(self, dst_dcp_size: int, dst_dcp_rank: int) -> bool: def requires_dcp_relayout(self, dst_dcp_size: int, dst_dcp_rank: int) -> bool:
@@ -751,7 +748,7 @@ class CommonKVManager(BaseKVManager):
`Connection refused`, and the leader's `prefill_port_table` ends `Connection refused`, and the leader's `prefill_port_table` ends
up missing rows. up missing rows.
""" """
if not self.dist_init_addr or get_parallel().config.nnodes == 1: if not self.dist_init_addr or get_parallel().nnodes == 1:
return local_port return local_port
if not (dist.is_available() and dist.is_initialized()): if not (dist.is_available() and dist.is_initialized()):
@@ -803,8 +800,8 @@ class CommonKVManager(BaseKVManager):
"rank_port": self.rank_port, "rank_port": self.rank_port,
"page_size": self.kv_args.page_size, "page_size": self.kv_args.page_size,
"kv_cache_dtype": self.kv_cache_dtype_str, "kv_cache_dtype": self.kv_cache_dtype_str,
"load_balance_method": get_parallel().config.load_balance_method, "load_balance_method": get_parallel().load_balance_method,
"enable_dsa_cache_layer_split": get_parallel().config.enable_dsa_cache_layer_split, "enable_dsa_cache_layer_split": get_parallel().enable_dsa_cache_layer_split,
# Self-register the HTTP API port so the decode can derive the PD # Self-register the HTTP API port so the decode can derive the PD
# retract rebootstrap /generate URL from bootstrap info instead of a # retract rebootstrap /generate URL from bootstrap info instead of a
# router-injected pd_rebootstrap_prefill_url. # router-injected pd_rebootstrap_prefill_url.
@@ -1186,12 +1183,11 @@ class CommonKVSender(BaseKVSender):
return return
self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Bootstrapping) self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Bootstrapping)
if get_parallel().config.dp_size > 1 and not req_has_disagg_prefill_dp_rank: if get_parallel().dp_size > 1 and not req_has_disagg_prefill_dp_rank:
if get_parallel().config.load_balance_method != "follow_bootstrap_room": if get_parallel().load_balance_method != "follow_bootstrap_room":
self._register_prefill_dp_rank() self._register_prefill_dp_rank()
elif ( elif (
self.kv_mgr.attn_dp_rank self.kv_mgr.attn_dp_rank != self.bootstrap_room % get_parallel().dp_size
!= self.bootstrap_room % get_parallel().config.dp_size
): ):
# follow_bootstrap_room was overridden by external routed_dp_rank # follow_bootstrap_room was overridden by external routed_dp_rank
if envs.SGLANG_DISAGGREGATION_FORCE_QUERY_PREFILL_DP_RANK.get(): if envs.SGLANG_DISAGGREGATION_FORCE_QUERY_PREFILL_DP_RANK.get():
@@ -1202,7 +1198,7 @@ class CommonKVSender(BaseKVSender):
f"follow_bootstrap_room conflict: dispatched to dp_rank " f"follow_bootstrap_room conflict: dispatched to dp_rank "
f"{self.kv_mgr.attn_dp_rank} but bootstrap_room " f"{self.kv_mgr.attn_dp_rank} but bootstrap_room "
f"{self.bootstrap_room} implies dp_rank " f"{self.bootstrap_room} implies dp_rank "
f"{self.bootstrap_room % get_parallel().config.dp_size}. " f"{self.bootstrap_room % get_parallel().dp_size}. "
f"Set SGLANG_DISAGGREGATION_FORCE_QUERY_PREFILL_DP_RANK=1 " f"Set SGLANG_DISAGGREGATION_FORCE_QUERY_PREFILL_DP_RANK=1 "
f"to allow mixed routing.", f"to allow mixed routing.",
) )
@@ -1276,7 +1272,7 @@ class CommonKVSender(BaseKVSender):
if ( if (
self.kv_mgr.enable_all_cp_ranks_for_transfer self.kv_mgr.enable_all_cp_ranks_for_transfer
and not get_parallel().config.enable_dsa_cache_layer_split and not get_parallel().enable_dsa_cache_layer_split
): ):
kv_indices, index_slice = filter_kv_indices_for_cp_rank( kv_indices, index_slice = filter_kv_indices_for_cp_rank(
self.kv_mgr, self.kv_mgr,
@@ -207,7 +207,7 @@ def launch_server(server_args: ServerArgs):
configure_logger(server_args, prefix=" encode_server") configure_logger(server_args, prefix=" encode_server")
publish(server_args, role="encoder") publish(server_args, role="encoder")
if get_parallel().config.dp_size > 1: if get_parallel().dp_size > 1:
dp_dispatcher = launch_dp_runtime(server_args) dp_dispatcher = launch_dp_runtime(server_args)
# runtime initializes multiprocess metrics before spawning; # runtime initializes multiprocess metrics before spawning;
# HTTP only exposes their endpoint. # HTTP only exposes their endpoint.
@@ -1492,10 +1492,10 @@ def launch_local_runtime(server_args: ServerArgs) -> EncoderRuntime:
This function owns backend construction only. HTTP/gRPC middleware, This function owns backend construction only. HTTP/gRPC middleware,
service registration, and network serving remain Transport concerns. service registration, and network serving remain Transport concerns.
""" """
if get_parallel().config.dp_size > 1: if get_parallel().dp_size > 1:
raise ValueError( raise ValueError(
"launch_local_runtime requires --dp-size 1; got " "launch_local_runtime requires --dp-size 1; got "
f"dp_size={get_parallel().config.dp_size}." f"dp_size={get_parallel().dp_size}."
) )
# Set up prometheus metrics. # Set up prometheus metrics.
@@ -1513,10 +1513,8 @@ def launch_local_runtime(server_args: ServerArgs) -> EncoderRuntime:
zmq_context = zmq.Context(10) zmq_context = zmq.Context(10)
ipc_path_prefix = random_uuid() ipc_path_prefix = random_uuid()
port_args = PortArgs.init_new(server_args) port_args = PortArgs.init_new(server_args)
if get_parallel().config.dist_init_addr: if get_parallel().dist_init_addr:
dist_init_method = NetworkAddress.parse( dist_init_method = NetworkAddress.parse(get_parallel().dist_init_addr).to_tcp()
get_parallel().config.dist_init_addr
).to_tcp()
else: else:
dist_init_method = NetworkAddress( dist_init_method = NetworkAddress(
get_serving().host or "127.0.0.1", port_args.nccl_port get_serving().host or "127.0.0.1", port_args.nccl_port
@@ -1572,12 +1570,12 @@ def launch_dp_runtime(server_args: ServerArgs) -> DPDispatcher:
HTTP uses this entry point today. gRPC can reuse it later without HTTP uses this entry point today. gRPC can reuse it later without
importing HTTP application state or Uvicorn. importing HTTP application state or Uvicorn.
""" """
if get_parallel().config.dp_size <= 1 or get_parallel().config.tp_size != 1: if get_parallel().dp_size <= 1 or get_parallel().config.tp_size != 1:
raise ValueError( raise ValueError(
"Encoder DP mode requires --dp-size > 1 and --tp-size 1; got " "Encoder DP mode requires --dp-size > 1 and --tp-size 1; got "
f"dp_size={get_parallel().config.dp_size}, tp_size={get_parallel().config.tp_size}." f"dp_size={get_parallel().dp_size}, tp_size={get_parallel().config.tp_size}."
) )
dp_size = get_parallel().config.dp_size dp_size = get_parallel().dp_size
logger.info(f"Launching encoder in DP mode: dp_size={dp_size}") logger.info(f"Launching encoder in DP mode: dp_size={dp_size}")
# DP mode: workers (subprocesses) write metrics to the shared multiproc dir; # DP mode: workers (subprocesses) write metrics to the shared multiproc dir;
+1 -1
View File
@@ -180,7 +180,7 @@ class PrefillBootstrapQueue:
"SGLANG_DISAGG_STAGING_BUFFER with pp_size > 1 is only " "SGLANG_DISAGG_STAGING_BUFFER with pp_size > 1 is only "
"supported by Mooncake." "supported by Mooncake."
) )
if get_parallel().config.enable_prefill_context_parallel: if get_parallel().enable_prefill_context_parallel:
# CP rewrites index_slice per rank, breaking the chunk grid. # CP rewrites index_slice per rank, breaking the chunk grid.
raise RuntimeError( raise RuntimeError(
"SGLANG_DISAGG_STAGING_BUFFER does not support " "SGLANG_DISAGG_STAGING_BUFFER does not support "
+2 -2
View File
@@ -129,7 +129,7 @@ def init_torch_distributed(
# included in later KV-cache sizing instead of appearing during capture. # included in later KV-cache sizing instead of appearing during capture.
if ( if (
device == "cuda" device == "cuda"
and get_parallel().config.enable_tp_lm_head_all_to_all and get_parallel().enable_tp_lm_head_all_to_all
and ps.tp_size > 1 and ps.tp_size > 1
): ):
_prewarm_tp_lm_head_all_to_all() _prewarm_tp_lm_head_all_to_all()
@@ -267,7 +267,7 @@ def _init_parallel_groups(
duplicate_attn_cp_group=( duplicate_attn_cp_group=(
is_hip() is_hip()
and server_args.enable_two_batch_overlap and server_args.enable_two_batch_overlap
and get_parallel().config.enable_dsa_prefill_context_parallel and get_parallel().enable_dsa_prefill_context_parallel
), ),
enable_symm_mem=get_exec().comm.enable_symm_mem, enable_symm_mem=get_exec().comm.enable_symm_mem,
recovered_rank=is_ep_joiner, recovered_rank=is_ep_joiner,
@@ -477,7 +477,7 @@ class MultimemAllGatherer:
# EP/mooncake setups, and keep multimem enabled. # EP/mooncake setups, and keep multimem enabled.
if ( if (
tp_group.world_size > 1 tp_group.world_size > 1
and get_parallel().config.nnodes > 1 and get_parallel().nnodes > 1
and not all(in_the_same_node_as(tp_group.cpu_group, source_rank=0)) and not all(in_the_same_node_as(tp_group.cpu_group, source_rank=0))
): ):
logger.warning( logger.warning(
+6 -7
View File
@@ -92,7 +92,7 @@ class ElasticEPStateManager:
if get_exec().moe.elastic_ep_backend is not None: if get_exec().moe.elastic_ep_backend is not None:
world_size = torch.distributed.get_world_size() world_size = torch.distributed.get_world_size()
active_rank_capacity = get_parallel().config.max_ep_size or world_size active_rank_capacity = get_parallel().max_ep_size or world_size
assert active_rank_capacity >= world_size, ( assert active_rank_capacity >= world_size, (
f"--max-ep-size ({active_rank_capacity}) must be >= " f"--max-ep-size ({active_rank_capacity}) must be >= "
f"world_size ({world_size})." f"world_size ({world_size})."
@@ -109,7 +109,7 @@ class ElasticEPStateManager:
if get_exec().moe.moe_a2a_backend == "nixl": if get_exec().moe.moe_a2a_backend == "nixl":
cls._on_scale = cls._on_scale_nixl cls._on_scale = cls._on_scale_nixl
inst.ep_join_rank_offset = get_parallel().config.ep_join_rank_offset inst.ep_join_rank_offset = get_parallel().ep_join_rank_offset
if server_args.is_ep_joiner: if server_args.is_ep_joiner:
cls._init_joiner_state(inst) cls._init_joiner_state(inst)
@@ -127,12 +127,11 @@ class ElasticEPStateManager:
if get_exec().moe.ep_join_mode == "scale": if get_exec().moe.ep_join_mode == "scale":
inst.effective_ep_size = ( inst.effective_ep_size = (
get_parallel().config.ep_join_rank_offset get_parallel().ep_join_rank_offset + get_parallel().config.tp_size
+ get_parallel().config.tp_size
) )
inst.original_ep_size = ( inst.original_ep_size = (
get_parallel().config.elastic_ep_initial_size get_parallel().elastic_ep_initial_size
or get_parallel().config.ep_join_rank_offset or get_parallel().ep_join_rank_offset
) )
inst.has_scaled = True inst.has_scaled = True
else: else:
@@ -317,7 +316,7 @@ def elastic_expanded_world_enabled() -> bool:
inst = ElasticEPStateManager.instance() inst = ElasticEPStateManager.instance()
if inst is None: if inst is None:
return False return False
if get_parallel().config.max_ep_size is None: if get_parallel().max_ep_size is None:
return False return False
active_target_size = inst.effective_ep_size active_target_size = inst.effective_ep_size
if inst.pending_ep_size is not None and inst.scale_phase in ( if inst.pending_ep_size is not None and inst.scale_phase in (
@@ -50,8 +50,8 @@ class ExpertBackupManager:
self.weight_pointer_map = {} self.weight_pointer_map = {}
self.transfer_engine = None self.transfer_engine = None
self.session_id = None self.session_id = None
self.engine_num = get_parallel().config.nnodes self.engine_num = get_parallel().nnodes
self.engine_rank = get_parallel().config.node_rank self.engine_rank = get_parallel().node_rank
self.expert_num = self.model_config.hf_config.n_routed_experts self.expert_num = self.model_config.hf_config.n_routed_experts
self.idmn = (self.expert_num // self.engine_num) * self.engine_rank self.idmn = (self.expert_num // self.engine_num) * self.engine_rank
self.idmx = (self.expert_num // self.engine_num) * (self.engine_rank + 1) self.idmx = (self.expert_num // self.engine_num) * (self.engine_rank + 1)
@@ -59,11 +59,11 @@ class ExpertBackupManager:
# Synchronization socket to avoid PUB/SUB slow joiner issues. # Synchronization socket to avoid PUB/SUB slow joiner issues.
self.recv_from_expert_backup_client = context.socket(zmq.PULL) self.recv_from_expert_backup_client = context.socket(zmq.PULL)
self.recv_from_expert_backup_client.bind( self.recv_from_expert_backup_client.bind(
f"tcp://{get_local_ip_auto()}:{PORT_BASE + get_parallel().config.node_rank * 2}" f"tcp://{get_local_ip_auto()}:{PORT_BASE + get_parallel().node_rank * 2}"
) )
self.send_to_expert_backup_client = context.socket(zmq.PUB) self.send_to_expert_backup_client = context.socket(zmq.PUB)
self.send_to_expert_backup_client.bind( self.send_to_expert_backup_client.bind(
f"tcp://{get_local_ip_auto()}:{PORT_BASE + get_parallel().config.node_rank * 2 + 1}" f"tcp://{get_local_ip_auto()}:{PORT_BASE + get_parallel().node_rank * 2 + 1}"
) )
self.backup_weights_from_disk() self.backup_weights_from_disk()
self.start_transfer_server() self.start_transfer_server()
+4 -8
View File
@@ -344,7 +344,7 @@ class Engine(EngineScoreMixin, EngineBase):
routed_dp_rank = data_parallel_rank routed_dp_rank = data_parallel_rank
if routed_dp_rank is not None: if routed_dp_rank is not None:
dp_size = get_parallel().config.dp_size dp_size = get_parallel().dp_size
if dp_size <= 1 and routed_dp_rank == 0: if dp_size <= 1 and routed_dp_rank == 0:
logger.debug( logger.debug(
f"routed_dp_rank={routed_dp_rank} is ignored because dp_size={dp_size}" f"routed_dp_rank={routed_dp_rank} is ignored because dp_size={dp_size}"
@@ -831,7 +831,7 @@ class Engine(EngineScoreMixin, EngineBase):
""" """
scheduler_procs = [] scheduler_procs = []
use_dp_controller = ( use_dp_controller = (
get_parallel().config.dp_size > 1 or get_exec().moe.ep_join_mode == "scale" get_parallel().dp_size > 1 or get_exec().moe.ep_join_mode == "scale"
) )
if not use_dp_controller: if not use_dp_controller:
@@ -1841,11 +1841,7 @@ def _compute_parallelism_ranks(
Called while the launcher is deciding what to spawn, so the sizes are the Called while the launcher is deciding what to spawn, so the sizes are the
configured ones -- the groups this is laying out do not exist yet. configured ones -- the groups this is laying out do not exist yet.
""" """
attn_dp_size = ( attn_dp_size = get_parallel().dp_size if get_parallel().enable_dp_attention else 1
get_parallel().config.dp_size
if get_parallel().config.enable_dp_attention
else 1
)
tp_size = server_args.tp_size tp_size = server_args.tp_size
attn_cp_size = get_parallel().config.attn_cp_size attn_cp_size = get_parallel().config.attn_cp_size
moe_dp_size = get_parallel().config.moe_dp_size moe_dp_size = get_parallel().config.moe_dp_size
@@ -1859,6 +1855,6 @@ def _compute_parallelism_ranks(
moe_ep_rank = ( moe_ep_rank = (
tp_rank tp_rank
% (tp_size // moe_dp_size) % (tp_size // moe_dp_size)
// (tp_size // moe_dp_size // get_parallel().config.ep_size) // (tp_size // moe_dp_size // get_parallel().ep_size)
) )
return attn_cp_rank, moe_dp_rank, moe_ep_rank return attn_cp_rank, moe_dp_rank, moe_ep_rank
+6 -10
View File
@@ -2180,7 +2180,7 @@ async def _send_disaggregation_warmup_requests(
return await asyncio.gather( return await asyncio.gather(
*( *(
send_request(session, dp_rank) send_request(session, dp_rank)
for dp_rank in range(get_parallel().config.dp_size) for dp_rank in range(get_parallel().dp_size)
) )
) )
@@ -2239,11 +2239,9 @@ def _execute_server_warmup(server_args: ServerArgs):
}, },
} }
if server_args.skip_tokenizer_init: if server_args.skip_tokenizer_init:
json_data["input_ids"] = [ json_data["input_ids"] = [[10, 11, 12] for _ in range(get_parallel().dp_size)]
[10, 11, 12] for _ in range(get_parallel().config.dp_size)
]
# TODO Workaround the bug that embedding errors for list of size 1 # TODO Workaround the bug that embedding errors for list of size 1
if get_parallel().config.dp_size == 1: if get_parallel().dp_size == 1:
json_data["input_ids"] = json_data["input_ids"][0] json_data["input_ids"] = json_data["input_ids"][0]
elif ( elif (
is_vlm is_vlm
@@ -2287,11 +2285,9 @@ def _execute_server_warmup(server_args: ServerArgs):
"temperature": 0.0, "temperature": 0.0,
} }
else: else:
json_data["text"] = [ json_data["text"] = ["The capital city of France is"] * get_parallel().dp_size
"The capital city of France is"
] * get_parallel().config.dp_size
# TODO Workaround the bug that embedding errors for list of size 1 # TODO Workaround the bug that embedding errors for list of size 1
if get_parallel().config.dp_size == 1: if get_parallel().dp_size == 1:
json_data["text"] = json_data["text"][0] json_data["text"] = json_data["text"][0]
# Config debug dumping # Config debug dumping
@@ -2332,7 +2328,7 @@ def _execute_server_warmup(server_args: ServerArgs):
if not failed_status_codes: if not failed_status_codes:
logger.info( logger.info(
"Disaggregation warmup requests completed for all %s DP ranks", "Disaggregation warmup requests completed for all %s DP ranks",
get_parallel().config.dp_size, get_parallel().dp_size,
) )
logger.info("End of disaggregation warmup") logger.info("End of disaggregation warmup")
else: else:
+2 -2
View File
@@ -148,8 +148,8 @@ async def get_loads(
"num_accelerators": _num_accelerators_per_dp_rank( "num_accelerators": _num_accelerators_per_dp_rank(
get_parallel().config.tp_size, get_parallel().config.tp_size,
get_parallel().config.pp_size, get_parallel().config.pp_size,
get_parallel().config.dp_size, get_parallel().dp_size,
get_parallel().config.enable_dp_attention, get_parallel().enable_dp_attention,
), ),
"loads": loads, "loads": loads,
} }
+2 -2
View File
@@ -157,7 +157,7 @@ class EPLBManager:
model=self._get_model(), model=self._get_model(),
new_expert_location_metadata=expert_location_metadata, new_expert_location_metadata=expert_location_metadata,
update_layer_ids=chunk_layer_ids, update_layer_ids=chunk_layer_ids,
nnodes=get_parallel().config.nnodes, nnodes=get_parallel().nnodes,
tp_rank=( tp_rank=(
self._elastic_global_rank() self._elastic_global_rank()
if is_post_scale_rebalance if is_post_scale_rebalance
@@ -221,7 +221,7 @@ class EPLBManager:
) )
def _elastic_global_rank(self) -> int: def _elastic_global_rank(self) -> int:
return self._ps.tp_rank + get_parallel().config.ep_join_rank_offset return self._ps.tp_rank + get_parallel().ep_join_rank_offset
def _check_rebalance_needed(self, average_utilization_rate_over_window): def _check_rebalance_needed(self, average_utilization_rate_over_window):
if average_utilization_rate_over_window is None: if average_utilization_rate_over_window is None:
+6 -9
View File
@@ -195,7 +195,7 @@ class ExpertLocationMetadata:
model_config_for_expert_location = common["model_config_for_expert_location"] model_config_for_expert_location = common["model_config_for_expert_location"]
num_physical_experts = common["num_physical_experts"] num_physical_experts = common["num_physical_experts"]
num_groups = model_config_for_expert_location.num_groups num_groups = model_config_for_expert_location.num_groups
num_nodes = 1 if use_flat_topology else get_parallel().config.nnodes num_nodes = 1 if use_flat_topology else get_parallel().nnodes
from sglang.srt.eplb import eplb_algorithms from sglang.srt.eplb import eplb_algorithms
@@ -238,15 +238,14 @@ class ExpertLocationMetadata:
+ get_exec().moe.ep_num_redundant_experts + get_exec().moe.ep_num_redundant_experts
) )
# elastic-EP scale-up rewrites ep_size on the published config # elastic-EP scale-up rewrites ep_size on the published config
ep_size = get_parallel().config.ep_size ep_size = get_parallel().ep_size
num_physical_experts = base_num_physical_experts num_physical_experts = base_num_physical_experts
initial_ep_size = get_parallel().config.elastic_ep_initial_size initial_ep_size = get_parallel().elastic_ep_initial_size
if initial_ep_size is not None: if initial_ep_size is not None:
if get_exec().moe.ep_join_mode == "scale": if get_exec().moe.ep_join_mode == "scale":
ep_size = max( ep_size = max(
ep_size, ep_size,
get_parallel().config.ep_join_rank_offset get_parallel().ep_join_rank_offset + get_parallel().config.tp_size,
+ get_parallel().config.tp_size,
) )
num_physical_experts, num_local_physical_experts = ( num_physical_experts, num_local_physical_experts = (
_compute_elastic_expert_layout( _compute_elastic_expert_layout(
@@ -570,7 +569,7 @@ def _compute_logical_to_all_physical_map(
num_local_gpu_physical_experts = num_physical_experts // ep_size num_local_gpu_physical_experts = num_physical_experts // ep_size
prefer_same_node = _prefer_same_node_experts() prefer_same_node = _prefer_same_node_experts()
num_gpus_per_node = ( num_gpus_per_node = (
get_parallel().config.ep_size // get_parallel().config.nnodes get_parallel().ep_size // get_parallel().nnodes
if prefer_same_node if prefer_same_node
else None else None
) )
@@ -634,9 +633,7 @@ def compute_logical_to_rank_dispatch_physical_map(
num_local_gpu_physical_experts = num_physical_experts // ep_size num_local_gpu_physical_experts = num_physical_experts // ep_size
prefer_same_node = _prefer_same_node_experts() prefer_same_node = _prefer_same_node_experts()
num_gpus_per_node = ( num_gpus_per_node = (
get_parallel().config.ep_size // get_parallel().config.nnodes get_parallel().ep_size // get_parallel().nnodes if prefer_same_node else None
if prefer_same_node
else None
) )
num_local_node_physical_experts = ( num_local_node_physical_experts = (
num_local_gpu_physical_experts * num_gpus_per_node num_local_gpu_physical_experts * num_gpus_per_node
@@ -106,7 +106,7 @@ def should_use_dsa_fused_topk(seed_dsa_topk_from_draft_extend: bool) -> bool:
def is_dsa_enable_prefill_cp(): def is_dsa_enable_prefill_cp():
if not envs.SGLANG_ENABLE_CP_V2.get(): if not envs.SGLANG_ENABLE_CP_V2.get():
return get_parallel().config.enable_dsa_prefill_context_parallel return get_parallel().enable_dsa_prefill_context_parallel
# Derive from the runtime CP topology + model arch rather than the legacy # Derive from the runtime CP topology + model arch rather than the legacy
# flag under CP-v2: DSA prefill CP is active when the CP group is on for a # flag under CP-v2: DSA prefill CP is active when the CP group is on for a
@@ -122,14 +122,14 @@ def is_dsa_enable_prefill_cp():
def is_dsa_prefill_cp_in_seq_split(): def is_dsa_prefill_cp_in_seq_split():
return ( return (
is_dsa_enable_prefill_cp() is_dsa_enable_prefill_cp()
and get_parallel().config.dsa_prefill_cp_mode == "in-seq-split" and get_parallel().dsa_prefill_cp_mode == "in-seq-split"
) )
def is_dsa_prefill_cp_round_robin_split(): def is_dsa_prefill_cp_round_robin_split():
return ( return (
is_dsa_enable_prefill_cp() is_dsa_enable_prefill_cp()
and get_parallel().config.dsa_prefill_cp_mode == "round-robin-split" and get_parallel().dsa_prefill_cp_mode == "round-robin-split"
) )
@@ -57,10 +57,7 @@ from sglang.kernels.ops.attention.flash_attention import (
def _should_disable_scheduler_metadata_precompute() -> bool: def _should_disable_scheduler_metadata_precompute() -> bool:
return bool( return bool(get_parallel().enable_prefill_cp or get_parallel().enable_dp_attention)
get_parallel().config.enable_prefill_cp
or get_parallel().config.enable_dp_attention
)
@dataclass @dataclass
+4 -4
View File
@@ -269,7 +269,7 @@ class AttnTpContext:
def init_context(self, q_lora_rank, is_dsa): def init_context(self, q_lora_rank, is_dsa):
self.is_dsa = is_dsa self.is_dsa = is_dsa
self.allow_input_scattered = ( self.allow_input_scattered = (
get_parallel().config.enable_attn_tp_input_scattered get_parallel().enable_attn_tp_input_scattered
and (_is_cuda or _is_npu) and (_is_cuda or _is_npu)
and q_lora_rank is not None and q_lora_rank is not None
and not is_dsa and not is_dsa
@@ -280,7 +280,7 @@ class AttnTpContext:
and not check_cuda_graph_backend(Phase.PREFILL, Backend.TC_PIECEWISE) and not check_cuda_graph_backend(Phase.PREFILL, Backend.TC_PIECEWISE)
and get_spec().speculative_algorithm != "EAGLE3" and get_spec().speculative_algorithm != "EAGLE3"
) )
if get_parallel().config.enable_attn_tp_input_scattered: if get_parallel().enable_attn_tp_input_scattered:
if not self.allow_input_scattered: if not self.allow_input_scattered:
logging.info( logging.info(
"attn_tp_input_scattered is not enabled while other conditions are not met" "attn_tp_input_scattered is not enabled while other conditions are not met"
@@ -438,11 +438,11 @@ class LayerScatterModes:
def enable_moe_dense_fully_dp(): def enable_moe_dense_fully_dp():
return get_parallel().config.moe_dense_tp_size == 1 return get_parallel().moe_dense_tp_size == 1
def enable_dwdp(): def enable_dwdp():
return get_parallel().config.dwdp_size > 1 return get_parallel().dwdp_size > 1
class LayerCommunicator: class LayerCommunicator:
@@ -51,7 +51,7 @@ class CpDecodeAttnTpContext:
"""Slices replicated attention weights across CP ranks during decode.""" """Slices replicated attention weights across CP ranks during decode."""
def __init__(self): def __init__(self):
enable_attn_tp = get_parallel().config.enable_cp_decode_attn_tp enable_attn_tp = get_parallel().enable_cp_decode_attn_tp
if enable_attn_tp and get_parallel().attn_cp_size > 1: if enable_attn_tp and get_parallel().attn_cp_size > 1:
self.decode_tp_rank = get_parallel().attn_cp_rank self.decode_tp_rank = get_parallel().attn_cp_rank
+1 -1
View File
@@ -63,7 +63,7 @@ def is_glm_dsa_cache_layer_split_enabled(model_runner: "ModelRunner") -> bool:
return ( return (
not model_runner.is_draft_worker not model_runner.is_draft_worker
and get_parallel().config.enable_dsa_cache_layer_split and get_parallel().enable_dsa_cache_layer_split
and model_runner.use_mla_backend and model_runner.use_mla_backend
and is_deepseek_dsa(model_runner.model_config.hf_config) and is_deepseek_dsa(model_runner.model_config.hf_config)
) )
+4 -7
View File
@@ -347,8 +347,8 @@ def initialize_dp_attention(
dp.max_len_with_idle = ( dp.max_len_with_idle = (
getattr(model_config.hf_config, "hybrid_override_pattern", None) is not None getattr(model_config.hf_config, "hybrid_override_pattern", None) is not None
) )
enable_dp_attention = get_parallel().config.enable_dp_attention enable_dp_attention = get_parallel().enable_dp_attention
dp_size = get_parallel().config.dp_size dp_size = get_parallel().dp_size
attn_cp_size = get_parallel().config.attn_cp_size attn_cp_size = get_parallel().config.attn_cp_size
dp.enabled = enable_dp_attention dp.enabled = enable_dp_attention
@@ -361,11 +361,8 @@ def initialize_dp_attention(
) )
_ATTN_DP_SIZE = dp_size if enable_dp_attention else 1 _ATTN_DP_SIZE = dp_size if enable_dp_attention else 1
if ( if get_exec().moe.elastic_ep_backend is not None and get_parallel().max_ep_size:
get_exec().moe.elastic_ep_backend is not None _ATTN_DP_RANK = tp_rank + get_parallel().ep_join_rank_offset
and get_parallel().config.max_ep_size
):
_ATTN_DP_RANK = tp_rank + get_parallel().config.ep_join_rank_offset
if server_args.is_ep_scale_joiner: if server_args.is_ep_scale_joiner:
dp.joiner_skip_all_gather = True dp.joiner_skip_all_gather = True
@@ -84,7 +84,7 @@ def resolve_flashinfer_allreduce_fusion_backend() -> Optional[str]:
backend = get_exec().comm.flashinfer_allreduce_fusion_backend backend = get_exec().comm.flashinfer_allreduce_fusion_backend
if backend is None: if backend is None:
return None return None
return _resolve_backend(backend, get_parallel().config.nnodes > 1) return _resolve_backend(backend, get_parallel().nnodes > 1)
if is_flashinfer_available(): if is_flashinfer_available():
+2 -4
View File
@@ -301,10 +301,8 @@ class LogitsProcessor(nn.Module):
self.config = config self.config = config
self.vocab_size = config.vocab_size self.vocab_size = config.vocab_size
self.logit_scale = logit_scale self.logit_scale = logit_scale
self.use_attn_tp_group = get_parallel().config.enable_dp_lm_head self.use_attn_tp_group = get_parallel().enable_dp_lm_head
self.use_tp_lm_head_all_to_all = ( self.use_tp_lm_head_all_to_all = get_parallel().enable_tp_lm_head_all_to_all
get_parallel().config.enable_tp_lm_head_all_to_all
)
self.use_fp32_lm_head = get_exec().features.enable_fp32_lm_head self.use_fp32_lm_head = get_exec().features.enable_fp32_lm_head
if self.use_attn_tp_group: if self.use_attn_tp_group:
self.attn_tp_size = get_parallel().attn_tp_size self.attn_tp_size = get_parallel().attn_tp_size
@@ -354,10 +354,10 @@ class FusedMoE(torch.nn.Module):
self._num_global_routed = num_experts - num_shared_slots self._num_global_routed = num_experts - num_shared_slots
if get_exec().moe.ep_join_mode == "scale": if get_exec().moe.ep_join_mode == "scale":
storage_ep_size = get_parallel().config.elastic_ep_initial_size storage_ep_size = get_parallel().elastic_ep_initial_size
assert storage_ep_size is not None assert storage_ep_size is not None
self._expert_storage_rank = ( self._expert_storage_rank = (
get_parallel().config.ep_join_rank_offset + self.moe_ep_rank get_parallel().ep_join_rank_offset + self.moe_ep_rank
) )
else: else:
storage_ep_size = self.moe_ep_size storage_ep_size = self.moe_ep_size
@@ -337,7 +337,7 @@ def ensure_cutedsl_wrapper(layer: torch.nn.Module) -> None:
else: else:
# Standard allgather path: the MoE sees up to dp_size local forwards # Standard allgather path: the MoE sees up to dp_size local forwards
# gathered together, so scale the per-rank forward bound by dp_size. # gathered together, so scale the per-rank forward bound by dp_size.
max_num_tokens = get_parallel().config.dp_size * cutedsl_moe_max_num_tokens() max_num_tokens = get_parallel().dp_size * cutedsl_moe_max_num_tokens()
top_k = layer.top_k if layer.top_k is not None else layer.moe_runner_config.top_k top_k = layer.top_k if layer.top_k is not None else layer.moe_runner_config.top_k
# inference_mode(False) ensures the wrapper's pre-allocated CUDA-graph # inference_mode(False) ensures the wrapper's pre-allocated CUDA-graph
# buffers are normal tensors. This call typically happens inside # buffers are normal tensors. This call typically happens inside
@@ -135,7 +135,7 @@ class NixlEPBuffer:
offset = ElasticEPStateManager.get_ep_join_rank_offset() offset = ElasticEPStateManager.get_ep_join_rank_offset()
global_rank = rank + offset global_rank = rank + offset
max_ep_size = get_parallel().config.max_ep_size or world_size max_ep_size = get_parallel().max_ep_size or world_size
nixl_max_ranks = max_ep_size nixl_max_ranks = max_ep_size
num_rdma_bytes = 0 num_rdma_bytes = 0
@@ -233,7 +233,7 @@ class _NixlEPDispatcherImplBase:
) )
self._active_world_size = dist.get_world_size(group) self._active_world_size = dist.get_world_size(group)
_max_ep = get_parallel().config.max_ep_size or self._active_world_size _max_ep = get_parallel().max_ep_size or self._active_world_size
self._mask_buffer = ( self._mask_buffer = (
torch.zeros(_max_ep, dtype=torch.int32, device="cuda") torch.zeros(_max_ep, dtype=torch.int32, device="cuda")
if self.active_ranks is not None if self.active_ranks is not None
@@ -155,7 +155,7 @@ class PplxAllToAllManager:
# pplx forces ep_size == world_size # pplx forces ep_size == world_size
# with pp_size == 1 (enforced in _ensure_nvshmem), so the EP group spans # with pp_size == 1 (enforced in _ensure_nvshmem), so the EP group spans
# a single node iff the whole job runs on one node. # a single node iff the whole job runs on one node.
is_internode = get_parallel().config.nnodes > 1 is_internode = get_parallel().nnodes > 1
if is_internode: if is_internode:
cls._all_to_all = AllToAll.internode( cls._all_to_all = AllToAll.internode(
+1 -1
View File
@@ -655,7 +655,7 @@ def should_skip_post_experts_all_reduce(*, is_tp_path: bool) -> bool:
""" """
if should_skip_mlp_all_reduce(): if should_skip_mlp_all_reduce():
return True return True
if get_parallel().config.dwdp_size > 1: if get_parallel().dwdp_size > 1:
return True return True
if should_use_dp_reduce_scatterv(): if should_use_dp_reduce_scatterv():
return True return True
+3 -3
View File
@@ -63,18 +63,18 @@ class ContextParallelMetadata:
def is_prefill_context_parallel_enabled(): def is_prefill_context_parallel_enabled():
return get_parallel().config.enable_prefill_context_parallel return get_parallel().enable_prefill_context_parallel
def is_prefill_cp_in_seq_split(): def is_prefill_cp_in_seq_split():
return ( return (
is_prefill_context_parallel_enabled() is_prefill_context_parallel_enabled()
and get_parallel().config.prefill_cp_mode == "in-seq-split" and get_parallel().prefill_cp_mode == "in-seq-split"
) )
def is_mla_prefill_cp_enabled() -> bool: def is_mla_prefill_cp_enabled() -> bool:
return get_parallel().config.enable_prefill_context_parallel and uses_mla_backend() return get_parallel().enable_prefill_context_parallel and uses_mla_backend()
def mla_use_prefill_cp(forward_batch, mla_enable_prefill_cp=None): def mla_use_prefill_cp(forward_batch, mla_enable_prefill_cp=None):
+1 -1
View File
@@ -99,7 +99,7 @@ class LoRAManager:
self.pending_lora_load_events = {} self.pending_lora_load_events = {}
self.eviction_policy = get_lora().lora_eviction_policy self.eviction_policy = get_lora().lora_eviction_policy
self.enable_dp_attention: bool = get_parallel().config.enable_dp_attention self.enable_dp_attention: bool = get_parallel().enable_dp_attention
self._experts_shared_outer_override: Optional[bool] = ( self._experts_shared_outer_override: Optional[bool] = (
get_lora().experts_shared_outer_loras get_lora().experts_shared_outer_loras
) )
@@ -148,12 +148,12 @@ class DataParallelController:
self.server_args = server_args self.server_args = server_args
self.port_args = port_args self.port_args = port_args
self.load_balance_method = LoadBalanceMethod.from_str( self.load_balance_method = LoadBalanceMethod.from_str(
get_parallel().config.load_balance_method get_parallel().load_balance_method
) )
self.run_scheduler_process_func = run_scheduler_process_func self.run_scheduler_process_func = run_scheduler_process_func
# Init inter-process communication # Init inter-process communication
self.context = zmq.Context(1 + get_parallel().config.dp_size) self.context = zmq.Context(1 + get_parallel().dp_size)
if server_args.node_rank == 0: if server_args.node_rank == 0:
self.recv_from_tokenizer = get_zmq_socket( self.recv_from_tokenizer = get_zmq_socket(
self.context, zmq.PULL, port_args.scheduler_input_ipc_name, False self.context, zmq.PULL, port_args.scheduler_input_ipc_name, False
@@ -173,8 +173,8 @@ class DataParallelController:
LoadBalanceMethod.TOTAL_TOKENS, LoadBalanceMethod.TOTAL_TOKENS,
) )
self.launch_dp_size: int = get_parallel().config.dp_size self.launch_dp_size: int = get_parallel().dp_size
self.max_dp_size: int = server_args.max_ep_size or get_parallel().config.dp_size self.max_dp_size: int = server_args.max_ep_size or get_parallel().dp_size
assert self.max_dp_size >= self.launch_dp_size, ( assert self.max_dp_size >= self.launch_dp_size, (
f"--max-ep-size ({self.max_dp_size}) must be >= " f"--max-ep-size ({self.max_dp_size}) must be >= "
f"--dp ({self.launch_dp_size})." f"--dp ({self.launch_dp_size})."
@@ -184,7 +184,7 @@ class DataParallelController:
self.max_dp_size - self.launch_dp_size self.max_dp_size - self.launch_dp_size
) )
self.dp_budget = DPBudget(get_parallel().config.dp_size) self.dp_budget = DPBudget(get_parallel().dp_size)
self.load_snapshot_reader = create_load_snapshot_reader( self.load_snapshot_reader = create_load_snapshot_reader(
port_args, port_args,
caller="DataParallelController", caller="DataParallelController",
@@ -201,16 +201,14 @@ class DataParallelController:
self._active_workers: List[int] = list(range(self.launch_dp_size)) self._active_workers: List[int] = list(range(self.launch_dp_size))
self._active_count_cache: int = self.launch_dp_size self._active_count_cache: int = self.launch_dp_size
if get_parallel().config.enable_dp_attention: if get_parallel().enable_dp_attention:
self.launch_dp_attention_schedulers(server_args, port_args) self.launch_dp_attention_schedulers(server_args, port_args)
# When local control broadcast is enabled, send control messages to # When local control broadcast is enabled, send control messages to
# every DP group leader (attn_tp_rank=0) so each leader broadcasts # every DP group leader (attn_tp_rank=0) so each leader broadcasts
# within its own attn_tp_group instead of the full tp_group. # within its own attn_tp_group instead of the full tp_group.
# Otherwise fall back to the original behaviour: send to only the # Otherwise fall back to the original behaviour: send to only the
# first leader, which then broadcasts over the full tp_group. # first leader, which then broadcasts over the full tp_group.
local_ctrl = ( local_ctrl = get_parallel().enable_dp_attention_local_control_broadcast
get_parallel().config.enable_dp_attention_local_control_broadcast
)
self.control_message_step = 1 if local_ctrl else server_args.tp_size self.control_message_step = 1 if local_ctrl else server_args.tp_size
else: else:
self.launch_dp_schedulers(server_args, port_args) self.launch_dp_schedulers(server_args, port_args)
@@ -374,7 +372,7 @@ class DataParallelController:
threads = [] threads = []
sockets = [] sockets = []
ready_events = [] ready_events = []
for dp_rank in range(get_parallel().config.dp_size): for dp_rank in range(get_parallel().dp_size):
tmp_port_args = PortArgs.init_new(server_args) tmp_port_args = PortArgs.init_new(server_args)
tmp_port_args.tokenizer_ipc_name = port_args.tokenizer_ipc_name tmp_port_args.tokenizer_ipc_name = port_args.tokenizer_ipc_name
tmp_port_args.detokenizer_ipc_name = port_args.detokenizer_ipc_name tmp_port_args.detokenizer_ipc_name = port_args.detokenizer_ipc_name
@@ -578,7 +576,7 @@ class DataParallelController:
bind_count = ( bind_count = (
self.max_dp_size self.max_dp_size
if server_args.elastic_ep_backend is not None if server_args.elastic_ep_backend is not None
else get_parallel().config.dp_size else get_parallel().dp_size
) )
for slot in range(bind_count): for slot in range(bind_count):
worker_port, worker_socket = get_zmq_socket_on_host( worker_port, worker_socket = get_zmq_socket_on_host(
@@ -608,7 +606,7 @@ class DataParallelController:
dp_rank: Optional[int], dp_rank: Optional[int],
worker_ports: Optional[List[int]] = None, worker_ports: Optional[List[int]] = None,
): ):
if not get_parallel().config.enable_dp_attention: if not get_parallel().enable_dp_attention:
logger.info(f"Launch DP{dp_rank} starting at GPU #{base_gpu_id}.") logger.info(f"Launch DP{dp_rank} starting at GPU #{base_gpu_id}.")
memory_saver_adapter = TorchMemorySaverAdapter.create( memory_saver_adapter = TorchMemorySaverAdapter.create(
@@ -642,13 +640,13 @@ class DataParallelController:
for tp_rank in tp_rank_range: for tp_rank in tp_rank_range:
rank_port_args = port_args rank_port_args = port_args
if get_parallel().config.enable_dp_attention: if get_parallel().enable_dp_attention:
# dp attention has different sharding logic # dp attention has different sharding logic
_, _, dp_rank, _ = compute_dp_attention_world_info( _, _, dp_rank, _ = compute_dp_attention_world_info(
get_parallel().config.enable_dp_attention, get_parallel().enable_dp_attention,
tp_rank, tp_rank,
server_args.tp_size, server_args.tp_size,
get_parallel().config.dp_size, get_parallel().dp_size,
get_parallel().config.attn_cp_size, get_parallel().config.attn_cp_size,
) )
# compute zmq ports for this dp rank # compute zmq ports for this dp rank
@@ -678,9 +676,7 @@ class DataParallelController:
+ (tp_rank % tp_size_per_node) * server_args.gpu_id_step + (tp_rank % tp_size_per_node) * server_args.gpu_id_step
) )
attn_dp_size = ( attn_dp_size = (
get_parallel().config.dp_size get_parallel().dp_size if get_parallel().enable_dp_attention else 1
if get_parallel().config.enable_dp_attention
else 1
) )
# Parallelism hierarchy (outermost to innermost): # Parallelism hierarchy (outermost to innermost):
@@ -703,7 +699,7 @@ class DataParallelController:
// ( // (
server_args.tp_size server_args.tp_size
// get_parallel().config.moe_dp_size // get_parallel().config.moe_dp_size
// get_parallel().config.ep_size // get_parallel().ep_size
) )
) )
+1 -2
View File
@@ -42,8 +42,7 @@ def maybe_create_ascend_config_store(transfer_backend: TransferBackend) -> None:
which the rust registry ports verbatim), leaving this store as the only which the rust registry ports verbatim), leaving this store as the only
``start_disagg_service`` duty left to perform.""" ``start_disagg_service`` duty left to perform."""
if not ( if not (
get_parallel().config.node_rank == 0 get_parallel().node_rank == 0 and transfer_backend == TransferBackend.ASCEND
and transfer_backend == TransferBackend.ASCEND
): ):
return return
try: try:
+7 -7
View File
@@ -71,7 +71,7 @@ def should_use_zmq() -> bool:
``SGLANG_LOAD_SNAPSHOT_USE_ZMQ`` forces zmq mode for testing. ``SGLANG_LOAD_SNAPSHOT_USE_ZMQ`` forces zmq mode for testing.
""" """
return ( return (
get_parallel().config.enable_dp_attention and get_parallel().config.nnodes > 1 get_parallel().enable_dp_attention and get_parallel().nnodes > 1
) or envs.SGLANG_LOAD_SNAPSHOT_USE_ZMQ.get() ) or envs.SGLANG_LOAD_SNAPSHOT_USE_ZMQ.get()
@@ -116,15 +116,15 @@ def zmq_reader_owner(caller: str) -> bool:
""" """
if not should_use_zmq(): if not should_use_zmq():
return False return False
if get_parallel().config.node_rank != 0: if get_parallel().node_rank != 0:
return False return False
if caller == "DataParallelController": if caller == "DataParallelController":
return ( return (
get_parallel().config.dp_size > 1 get_parallel().dp_size > 1
and get_parallel().config.load_balance_method.lower() in _LOAD_AWARE_METHODS and get_parallel().load_balance_method.lower() in _LOAD_AWARE_METHODS
) )
if get_parallel().config.dp_size > 1 and ( if get_parallel().dp_size > 1 and (
get_parallel().config.load_balance_method.lower() in _LOAD_AWARE_METHODS get_parallel().load_balance_method.lower() in _LOAD_AWARE_METHODS
): ):
return False return False
return caller == _tokenizer_load_snapshot_owner_caller() return caller == _tokenizer_load_snapshot_owner_caller()
@@ -658,7 +658,7 @@ def create_load_snapshot_reader(port_args, caller: str):
``"MultiTokenizerRouter"`` -- determines who binds the zmq PULL ``"MultiTokenizerRouter"`` -- determines who binds the zmq PULL
socket when zmq mode is active. socket when zmq mode is active.
""" """
dp_size = get_parallel().config.dp_size dp_size = get_parallel().dp_size
if zmq_reader_owner(caller): if zmq_reader_owner(caller):
return ZmqShmLoadSnapshotReader( return ZmqShmLoadSnapshotReader(
_zmq_addr_for(port_args), shm_path_for(port_args.instance_id), dp_size _zmq_addr_for(port_args), shm_path_for(port_args.instance_id), dp_size
@@ -108,7 +108,7 @@ class PrefillDelayer:
f"queue_trigger_enabled={self._queue_trigger_enabled}" f"queue_trigger_enabled={self._queue_trigger_enabled}"
) )
self.dp_size = dp_size self.dp_size = dp_size
self.enable_dp_attention = get_parallel().config.enable_dp_attention self.enable_dp_attention = get_parallel().enable_dp_attention
dp_size_dim = dp_size if self.enable_dp_attention else 1 dp_size_dim = dp_size if self.enable_dp_attention else 1
# Mirror scheduler_dp_attn_mixin's NCCL all-gather path: when the # Mirror scheduler_dp_attn_mixin's NCCL all-gather path: when the
+11 -11
View File
@@ -470,16 +470,16 @@ class Scheduler(
self.max_recv_per_poll = envs.SGLANG_SCHEDULER_MAX_RECV_PER_POLL.get() self.max_recv_per_poll = envs.SGLANG_SCHEDULER_MAX_RECV_PER_POLL.get()
self.max_new_tokens_limit = envs.SGLANG_MAX_NEW_TOKENS_LIMIT.get() self.max_new_tokens_limit = envs.SGLANG_MAX_NEW_TOKENS_LIMIT.get()
self.enable_hisparse = get_memory().enable_hisparse self.enable_hisparse = get_memory().enable_hisparse
self.enable_dp_attention = get_parallel().config.enable_dp_attention self.enable_dp_attention = get_parallel().enable_dp_attention
self.enable_unified_memory = get_memory().enable_unified_memory self.enable_unified_memory = get_memory().enable_unified_memory
# Distributed rank info # Distributed rank info
attn_tp_rank, attn_tp_size, attn_dp_rank, attn_dp_size = ( attn_tp_rank, attn_tp_size, attn_dp_rank, attn_dp_size = (
compute_dp_attention_world_info( compute_dp_attention_world_info(
get_parallel().config.enable_dp_attention, get_parallel().enable_dp_attention,
tp_rank, tp_rank,
get_parallel().config.tp_size, get_parallel().config.tp_size,
get_parallel().config.dp_size, get_parallel().dp_size,
get_parallel().config.attn_cp_size, get_parallel().config.attn_cp_size,
) )
) )
@@ -489,7 +489,7 @@ class Scheduler(
pp_rank=pp_rank, pp_rank=pp_rank,
pp_size=get_parallel().config.pp_size, pp_size=get_parallel().config.pp_size,
dp_rank=dp_rank, dp_rank=dp_rank,
dp_size=get_parallel().config.dp_size, dp_size=get_parallel().dp_size,
attn_tp_rank=attn_tp_rank, attn_tp_rank=attn_tp_rank,
attn_tp_size=attn_tp_size, attn_tp_size=attn_tp_size,
attn_cp_rank=attn_cp_rank, attn_cp_rank=attn_cp_rank,
@@ -499,7 +499,7 @@ class Scheduler(
attn_dp_rank=attn_dp_rank, attn_dp_rank=attn_dp_rank,
attn_dp_size=attn_dp_size, attn_dp_size=attn_dp_size,
moe_ep_rank=moe_ep_rank, moe_ep_rank=moe_ep_rank,
moe_ep_size=get_parallel().config.ep_size, moe_ep_size=get_parallel().ep_size,
moe_dp_rank=moe_dp_rank, moe_dp_rank=moe_dp_rank,
moe_dp_size=get_parallel().config.moe_dp_size, moe_dp_size=get_parallel().config.moe_dp_size,
gpu_id=gpu_id, gpu_id=gpu_id,
@@ -1088,7 +1088,7 @@ class Scheduler(
self.min_free_slots_delayer = MinFreeSlotsDelayer( self.min_free_slots_delayer = MinFreeSlotsDelayer(
min_free_slots=min_free_slots min_free_slots=min_free_slots
) )
if not get_parallel().config.pp_max_micro_batch_size: if not get_parallel().pp_max_micro_batch_size:
get_context().override( get_context().override(
"scheduler.pp_max_micro_batch_size_default", "scheduler.pp_max_micro_batch_size_default",
pp_max_micro_batch_size=max( pp_max_micro_batch_size=max(
@@ -1432,7 +1432,7 @@ class Scheduler(
gloo_group=self.attn_tp_cpu_group, gloo_group=self.attn_tp_cpu_group,
tp_rank=self.ps.tp_rank, tp_rank=self.ps.tp_rank,
tp_size=self.ps.tp_size, tp_size=self.ps.tp_size,
dp_size=get_parallel().config.dp_size, dp_size=get_parallel().dp_size,
gpu_id=self.ps.gpu_id, gpu_id=self.ps.gpu_id,
bootstrap_port=get_disagg().disaggregation_bootstrap_port, bootstrap_port=get_disagg().disaggregation_bootstrap_port,
max_total_num_tokens=self.max_total_num_tokens, max_total_num_tokens=self.max_total_num_tokens,
@@ -3291,7 +3291,7 @@ class Scheduler(
beam_width: Optional[int] = None, beam_width: Optional[int] = None,
running_batch: Optional[ScheduleBatch] = None, running_batch: Optional[ScheduleBatch] = None,
) -> int: ) -> int:
pp_budget = get_parallel().config.pp_max_micro_batch_size - running_bs pp_budget = get_parallel().pp_max_micro_batch_size - running_bs
available = self.req_to_token_pool.available_size() available = self.req_to_token_pool.available_size()
active_batch = running_batch or self.running_batch active_batch = running_batch or self.running_batch
@@ -5015,7 +5015,7 @@ class Scheduler(
old_ep_size = ElasticEPStateManager.get_effective_ep_size() old_ep_size = ElasticEPStateManager.get_effective_ep_size()
new_ep_size = recv_req.new_ep_size new_ep_size = recv_req.new_ep_size
max_ep_size = get_parallel().config.max_ep_size or old_ep_size max_ep_size = get_parallel().max_ep_size or old_ep_size
logger.debug( logger.debug(
"[Elastic EP][scale] request received: new_ep_size=%d " "[Elastic EP][scale] request received: new_ep_size=%d "
@@ -5285,7 +5285,7 @@ def configure_scheduler_process(
prefix += f" MOE_DP{moe_dp_rank}" prefix += f" MOE_DP{moe_dp_rank}"
if get_parallel().config.tp_size > 1: if get_parallel().config.tp_size > 1:
prefix += f" TP{shown_tp}" prefix += f" TP{shown_tp}"
if get_parallel().config.ep_size > 1: if get_parallel().ep_size > 1:
prefix += f" EP{shown_moe_ep}" prefix += f" EP{shown_moe_ep}"
# Config the process # Config the process
@@ -5301,7 +5301,7 @@ def configure_scheduler_process(
set_gpu_proc_affinity( set_gpu_proc_affinity(
get_parallel().config.pp_size, get_parallel().config.pp_size,
get_parallel().config.tp_size, get_parallel().config.tp_size,
get_parallel().config.nnodes, get_parallel().nnodes,
gpu_id, gpu_id,
) )
if not envs.SGLANG_NUMA_BIND_V2.get(): if not envs.SGLANG_NUMA_BIND_V2.get():
@@ -409,7 +409,7 @@ class SchedulerDPAttnAdapter:
return prepare_mlp_sync_batch_raw( return prepare_mlp_sync_batch_raw(
local_batch, local_batch,
model_runner=self.model_runner, model_runner=self.model_runner,
dp_size=get_parallel().config.dp_size, dp_size=get_parallel().dp_size,
attn_tp_size=self.ps.attn_tp_size, attn_tp_size=self.ps.attn_tp_size,
attn_cp_size=self.ps.attn_cp_size, attn_cp_size=self.ps.attn_cp_size,
tp_group=self.tp_group, tp_group=self.tp_group,
@@ -418,7 +418,7 @@ class SchedulerDPAttnAdapter:
require_mlp_tp_gather=require_mlp_tp_gather(), require_mlp_tp_gather=require_mlp_tp_gather(),
disable_overlap_schedule=get_schedule().disable_overlap_schedule, disable_overlap_schedule=get_schedule().disable_overlap_schedule,
offload_tags=self.offload_tags, offload_tags=self.offload_tags,
dwdp=get_parallel().config.dwdp_size > 1, dwdp=get_parallel().dwdp_size > 1,
) )
def maybe_prepare_mlp_sync_batch( def maybe_prepare_mlp_sync_batch(
@@ -34,7 +34,7 @@ class SchedulerRecvSkipper:
return ForwardMode.DECODE return ForwardMode.DECODE
def __init__(self): def __init__(self):
self._use_synced_mode = get_parallel().config.enable_dp_attention self._use_synced_mode = get_parallel().enable_dp_attention
self._counter = 0 self._counter = 0
self._threshold = get_schedule().scheduler_recv_interval self._threshold = get_schedule().scheduler_recv_interval
# All can be tuned if needed # All can be tuned if needed
@@ -151,7 +151,7 @@ class SchedulerRequestReceiver:
return recv_reqs return recv_reqs
def _broadcast_reqs_across_ranks(self, recv_reqs: Optional[List]) -> List: def _broadcast_reqs_across_ranks(self, recv_reqs: Optional[List]) -> List:
if get_parallel().config.enable_dp_attention: if get_parallel().enable_dp_attention:
if self.ps.attn_tp_rank == 0 and self.ps.attn_cp_rank == 0: if self.ps.attn_tp_rank == 0 and self.ps.attn_cp_rank == 0:
work_reqs, control_reqs = self._split_work_and_control_reqs(recv_reqs) work_reqs, control_reqs = self._split_work_and_control_reqs(recv_reqs)
else: else:
@@ -180,7 +180,7 @@ class SchedulerRequestReceiver:
# instead of the full tp_group. This avoids an expensive # instead of the full tp_group. This avoids an expensive
# all-ranks gloo sync. # all-ranks gloo sync.
_local_ctrl = ( _local_ctrl = (
get_parallel().config.enable_dp_attention_local_control_broadcast get_parallel().enable_dp_attention_local_control_broadcast
or is_ep_scale_joiner() or is_ep_scale_joiner()
) )
if _local_ctrl: if _local_ctrl:
@@ -258,7 +258,7 @@ class SchedulerRequestReceiver:
# peer ranks may still be unpickling ShmPointerMMData # peer ranks may still be unpickling ShmPointerMMData
# (-> shm_open). Synchronize the same CPU groups that carried # (-> shm_open). Synchronize the same CPU groups that carried
# SHM-backed work requests before materialize() unlinks them. # SHM-backed work requests before materialize() unlinks them.
if get_parallel().config.enable_dp_attention: if get_parallel().enable_dp_attention:
if self.ps.attn_tp_size > 1: if self.ps.attn_tp_size > 1:
barrier(group=self.attn_tp_cpu_group) barrier(group=self.attn_tp_cpu_group)
if self.ps.attn_cp_size > 1: if self.ps.attn_cp_size > 1:
@@ -128,7 +128,7 @@ class SchedulerPPMixin:
next_pp_outputs = None next_pp_outputs = None
next_batch_result = None next_batch_result = None
d2h_event = None d2h_event = None
if get_parallel().config.pp_async_batch_depth > 0: if get_parallel().pp_async_batch_depth > 0:
next_pp_outputs, next_batch_result, d2h_event = ( next_pp_outputs, next_batch_result, d2h_event = (
self._pp_commit_send_output_work_and_preprocess_output_tensors( self._pp_commit_send_output_work_and_preprocess_output_tensors(
next_first_rank_mb_id, next_first_rank_mb_id,
@@ -144,7 +144,7 @@ class SchedulerPPMixin:
self.mb_metadata, self.mb_metadata,
self.last_rank_comm_queue, self.last_rank_comm_queue,
) )
if get_parallel().config.pp_async_batch_depth == 0: if get_parallel().pp_async_batch_depth == 0:
next_pp_outputs, next_batch_result, d2h_event = ( next_pp_outputs, next_batch_result, d2h_event = (
self._pp_commit_send_output_work_and_preprocess_output_tensors( self._pp_commit_send_output_work_and_preprocess_output_tensors(
next_first_rank_mb_id, next_first_rank_mb_id,
@@ -274,7 +274,7 @@ class SchedulerPPMixin:
server_is_idle = False server_is_idle = False
pp_proxy_tensors = self._pp_recv_proxy_tensors() pp_proxy_tensors = self._pp_recv_proxy_tensors()
if get_parallel().config.pp_async_batch_depth > 0: if get_parallel().pp_async_batch_depth > 0:
next_pp_outputs, next_batch_result, d2h_event = ( next_pp_outputs, next_batch_result, d2h_event = (
self._pp_commit_send_output_work_and_preprocess_output_tensors( self._pp_commit_send_output_work_and_preprocess_output_tensors(
next_first_rank_mb_id, next_first_rank_mb_id,
@@ -292,7 +292,7 @@ class SchedulerPPMixin:
self.mb_metadata, self.mb_metadata,
self.last_rank_comm_queue, self.last_rank_comm_queue,
) )
if get_parallel().config.pp_async_batch_depth == 0: if get_parallel().pp_async_batch_depth == 0:
next_pp_outputs, next_batch_result, d2h_event = ( next_pp_outputs, next_batch_result, d2h_event = (
self._pp_commit_send_output_work_and_preprocess_output_tensors( self._pp_commit_send_output_work_and_preprocess_output_tensors(
next_first_rank_mb_id, next_first_rank_mb_id,
@@ -435,7 +435,7 @@ class SchedulerPPMixin:
pp_proxy_tensors = self._pp_recv_proxy_tensors() pp_proxy_tensors = self._pp_recv_proxy_tensors()
# early send output if possible # early send output if possible
if get_parallel().config.pp_async_batch_depth > 0: if get_parallel().pp_async_batch_depth > 0:
next_pp_outputs, next_batch_result, d2h_event = ( next_pp_outputs, next_batch_result, d2h_event = (
self._pp_commit_send_output_work_and_preprocess_output_tensors( self._pp_commit_send_output_work_and_preprocess_output_tensors(
next_first_rank_mb_id, next_first_rank_mb_id,
@@ -453,7 +453,7 @@ class SchedulerPPMixin:
self.last_rank_comm_queue, self.last_rank_comm_queue,
) )
if get_parallel().config.pp_async_batch_depth == 0: if get_parallel().pp_async_batch_depth == 0:
next_pp_outputs, next_batch_result, d2h_event = ( next_pp_outputs, next_batch_result, d2h_event = (
self._pp_commit_send_output_work_and_preprocess_output_tensors( self._pp_commit_send_output_work_and_preprocess_output_tensors(
next_first_rank_mb_id, next_first_rank_mb_id,
@@ -564,12 +564,10 @@ class SchedulerPPMixin:
self.on_idle() self.on_idle()
def init_pp_loop_state(self: Scheduler): def init_pp_loop_state(self: Scheduler):
self.pp_loop_size: int = ( self.pp_loop_size: int = self.ps.pp_size + get_parallel().pp_async_batch_depth
self.ps.pp_size + get_parallel().config.pp_async_batch_depth
)
# In CP mode, attention weights are duplicated, eliminating the need for the attention TP all-gather operation. # In CP mode, attention weights are duplicated, eliminating the need for the attention TP all-gather operation.
self.require_attn_tp_allgather = ( self.require_attn_tp_allgather = (
not get_parallel().config.enable_dsa_prefill_context_parallel not get_parallel().enable_dsa_prefill_context_parallel
) )
self.mbs = [None] * self.pp_loop_size self.mbs = [None] * self.pp_loop_size
self.last_mbs = [None] * self.pp_loop_size self.last_mbs = [None] * self.pp_loop_size
@@ -165,7 +165,7 @@ class TokenizerControlMixin:
mode = spec[2] if len(spec) > 2 else "queueing" mode = spec[2] if len(spec) > 2 else "queueing"
comm = FanOutCommunicator( comm = FanOutCommunicator(
self._dispatch_to_scheduler, self._dispatch_to_scheduler,
get_parallel().config.dp_size, get_parallel().dp_size,
mode, mode,
) )
setattr(self, f"{name}_communicator", comm) setattr(self, f"{name}_communicator", comm)
@@ -174,8 +174,8 @@ class TokenizerControlMixin:
def update_control_communicator_fan_out(self: TokenizerManager, worker_count: int): def update_control_communicator_fan_out(self: TokenizerManager, worker_count: int):
primary_group_control = ( primary_group_control = (
get_parallel().config.enable_dp_attention get_parallel().enable_dp_attention
and not get_parallel().config.enable_dp_attention_local_control_broadcast and not get_parallel().enable_dp_attention_local_control_broadcast
) )
if primary_group_control: if primary_group_control:
control_fan_out = ( control_fan_out = (
@@ -428,8 +428,7 @@ class TokenizerControlMixin:
) -> Tuple[bool, str]: ) -> Tuple[bool, str]:
self.auto_create_handle_loop() self.auto_create_handle_loop()
assert ( assert (
get_parallel().config.dp_size == 1 get_parallel().dp_size == 1 or get_parallel().enable_dp_attention
or get_parallel().config.enable_dp_attention
), "dp_size must be 1 or dp attention must be enabled for update weights from distributed" ), "dp_size must be 1 or dp attention must be enabled for update weights from distributed"
results = await self.init_weights_update_group_communicator(obj) results = await self.init_weights_update_group_communicator(obj)
@@ -442,8 +441,7 @@ class TokenizerControlMixin:
) -> Tuple[bool, str]: ) -> Tuple[bool, str]:
self.auto_create_handle_loop() self.auto_create_handle_loop()
assert ( assert (
get_parallel().config.dp_size == 1 get_parallel().dp_size == 1 or get_parallel().enable_dp_attention
or get_parallel().config.enable_dp_attention
), "dp_size must be 1 or dp attention must be enabled for destroy parameter update group" ), "dp_size must be 1 or dp attention must be enabled for destroy parameter update group"
results = await self.destroy_weights_update_group_communicator(obj) results = await self.destroy_weights_update_group_communicator(obj)
@@ -456,8 +454,7 @@ class TokenizerControlMixin:
) -> Tuple[bool, str]: ) -> Tuple[bool, str]:
self.auto_create_handle_loop() self.auto_create_handle_loop()
assert ( assert (
get_parallel().config.dp_size == 1 get_parallel().dp_size == 1 or get_parallel().enable_dp_attention
or get_parallel().config.enable_dp_attention
), "dp_size must be 1 or dp attention must be enabled for update weights from distributed" ), "dp_size must be 1 or dp attention must be enabled for update weights from distributed"
if obj.abort_all_requests: if obj.abort_all_requests:
@@ -490,7 +487,7 @@ class TokenizerControlMixin:
self.auto_create_handle_loop() self.auto_create_handle_loop()
# TODO: support DP # TODO: support DP
assert ( assert (
get_parallel().config.dp_size == 1 get_parallel().dp_size == 1
), "dp_size must be 1 for init_weights_send_group_for_remote_instance" ), "dp_size must be 1 for init_weights_send_group_for_remote_instance"
result = ( result = (
await self.init_weights_send_group_for_remote_instance_communicator(obj) await self.init_weights_send_group_for_remote_instance_communicator(obj)
@@ -505,7 +502,7 @@ class TokenizerControlMixin:
self.auto_create_handle_loop() self.auto_create_handle_loop()
# TODO: support DP # TODO: support DP
assert ( assert (
get_parallel().config.dp_size == 1 get_parallel().dp_size == 1
), "dp_size must be 1 for send_weights_to_remote_instance" ), "dp_size must be 1 for send_weights_to_remote_instance"
result = (await self.send_weights_to_remote_instance_communicator(obj))[0] result = (await self.send_weights_to_remote_instance_communicator(obj))[0]
return result.success, result.message return result.success, result.message
@@ -517,8 +514,7 @@ class TokenizerControlMixin:
) -> Tuple[bool, str]: ) -> Tuple[bool, str]:
self.auto_create_handle_loop() self.auto_create_handle_loop()
assert ( assert (
get_parallel().config.dp_size == 1 get_parallel().dp_size == 1 or get_parallel().enable_dp_attention
or get_parallel().config.enable_dp_attention
), "dp_size must be 1 or dp attention must be enabled for update weights from tensor" ), "dp_size must be 1 or dp attention must be enabled for update weights from tensor"
if obj.abort_all_requests: if obj.abort_all_requests:
@@ -556,8 +552,7 @@ class TokenizerControlMixin:
try: try:
# For now, we only support single data parallel instance # For now, we only support single data parallel instance
assert ( assert (
get_parallel().config.dp_size == 1 get_parallel().dp_size == 1 or get_parallel().enable_dp_attention
or get_parallel().config.enable_dp_attention
), "dp_size must be 1 or dp attention must be enabled for update weights from IPC" ), "dp_size must be 1 or dp attention must be enabled for update weights from IPC"
logger.info("Starting IPC weight update") logger.info("Starting IPC weight update")
@@ -620,8 +615,7 @@ class TokenizerControlMixin:
) )
assert ( assert (
get_parallel().config.dp_size == 1 get_parallel().dp_size == 1 or get_parallel().enable_dp_attention
or get_parallel().config.enable_dp_attention
), "dp_size must be 1 or dp attention must be enabled for dynamic lora loading" ), "dp_size must be 1 or dp attention must be enabled for dynamic lora loading"
logger.info( logger.info(
"Start load Lora adapter. Lora name=%s, path=%s", "Start load Lora adapter. Lora name=%s, path=%s",
@@ -699,8 +693,7 @@ class TokenizerControlMixin:
) )
assert ( assert (
get_parallel().config.dp_size == 1 get_parallel().dp_size == 1 or get_parallel().enable_dp_attention
or get_parallel().config.enable_dp_attention
), "dp_size must be 1 or dp attention must be enabled for dynamic lora loading" ), "dp_size must be 1 or dp attention must be enabled for dynamic lora loading"
logger.info( logger.info(
"Start load Lora adapter from tensors. Lora name=%s", "Start load Lora adapter from tensors. Lora name=%s",
@@ -780,8 +773,7 @@ class TokenizerControlMixin:
), "lora_name must be provided to unload LoRA adapter" ), "lora_name must be provided to unload LoRA adapter"
assert ( assert (
get_parallel().config.dp_size == 1 get_parallel().dp_size == 1 or get_parallel().enable_dp_attention
or get_parallel().config.enable_dp_attention
), "dp_size must be 1 or dp attention must be enabled for dynamic lora loading" ), "dp_size must be 1 or dp attention must be enabled for dynamic lora loading"
logger.info( logger.info(
"Start unload Lora adapter. Lora name=%s", "Start unload Lora adapter. Lora name=%s",
@@ -801,7 +793,7 @@ class TokenizerControlMixin:
self.auto_create_handle_loop() self.auto_create_handle_loop()
results = await self.get_weights_by_name_communicator(obj) results = await self.get_weights_by_name_communicator(obj)
all_parameters = [r.parameter for r in results] all_parameters = [r.parameter for r in results]
if get_parallel().config.dp_size == 1: if get_parallel().dp_size == 1:
return all_parameters[0] return all_parameters[0]
else: else:
return all_parameters return all_parameters
@@ -415,7 +415,7 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
self.server_args = server_args self.server_args = server_args
assert_published(server_args, role="tokenizer") assert_published(server_args, role="tokenizer")
self.startup_time: Optional[Dict[str, Any]] = None self.startup_time: Optional[Dict[str, Any]] = None
self.elastic_worker_count = get_parallel().config.dp_size self.elastic_worker_count = get_parallel().dp_size
self.elastic_pending_ep_size = None self.elastic_pending_ep_size = None
self.elastic_scale_phase = "idle" self.elastic_scale_phase = "idle"
self.elastic_last_error = None self.elastic_last_error = None
@@ -1550,7 +1550,7 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
return batch_size > 0 and ( return batch_size > 0 and (
get_serving().enable_tokenizer_batch_encode get_serving().enable_tokenizer_batch_encode
or ( or (
(not get_parallel().config.enable_dp_attention) (not get_parallel().enable_dp_attention)
and (not self._batch_has_text(batch_size, requests)) and (not self._batch_has_text(batch_size, requests))
) )
) )
@@ -279,9 +279,7 @@ def build_kv_cache(
), ),
is_eagle=spec_algorithm.is_eagle(), is_eagle=spec_algorithm.is_eagle(),
tp_cache_group=( tp_cache_group=(
attn_tp_cpu_group attn_tp_cpu_group if get_parallel().enable_dp_attention else tp_cpu_group
if get_parallel().config.enable_dp_attention
else tp_cpu_group
), ),
attn_cp_cache_group=attn_cp_cpu_group, attn_cp_cache_group=attn_cp_cpu_group,
attn_tp_cache_group=attn_tp_cpu_group, attn_tp_cache_group=attn_tp_cpu_group,
@@ -45,7 +45,7 @@ def ranks_per_host() -> int:
return 1 return 1
if world_group.world_size == 1: if world_group.world_size == 1:
return 1 return 1
return max(world_group.world_size // get_parallel().config.nnodes, 1) return max(world_group.world_size // get_parallel().nnodes, 1)
def host_memory_budget_bytes() -> int: def host_memory_budget_bytes() -> int:
@@ -610,7 +610,7 @@ class CPUGraphRunner:
model_runner.server_args.enable_profile_cuda_graph model_runner.server_args.enable_profile_cuda_graph
) )
self.tp_size = get_parallel().config.tp_size self.tp_size = get_parallel().config.tp_size
self.dp_size = get_parallel().config.dp_size self.dp_size = get_parallel().dp_size
self.pp_size = get_parallel().config.pp_size self.pp_size = get_parallel().config.pp_size
self.capture_forward_mode = ForwardMode.DECODE self.capture_forward_mode = ForwardMode.DECODE
@@ -486,19 +486,17 @@ class ModelRunner:
if not (get_exec().moe.elastic_ep_backend is not None and is_ep_scale_joiner()): if not (get_exec().moe.elastic_ep_backend is not None and is_ep_scale_joiner()):
return return
join_effective_ep_size = ( join_effective_ep_size = get_parallel().ep_join_rank_offset + self.ps.tp_size
get_parallel().config.ep_join_rank_offset + self.ps.tp_size
)
dist.barrier(group=self.tp_group.cpu_group) dist.barrier(group=self.tp_group.cpu_group)
if self.ps.tp_rank == 0: if self.ps.tp_rank == 0:
register_scale_cohort( register_scale_cohort(
get_parallel().config.ep_join_rank_offset, get_parallel().ep_join_rank_offset,
join_effective_ep_size, join_effective_ep_size,
) )
join_scale_process_group() join_scale_process_group()
get_context().override("elastic_ep.scale_join", ep_size=join_effective_ep_size) get_context().override("elastic_ep.scale_join", ep_size=join_effective_ep_size)
global_ep_rank = self.ps.tp_rank + get_parallel().config.ep_join_rank_offset global_ep_rank = self.ps.tp_rank + get_parallel().ep_join_rank_offset
broadcast_global_expert_location_metadata( broadcast_global_expert_location_metadata(
model_config=self.model_config, model_config=self.model_config,
moe_ep_rank=global_ep_rank, moe_ep_rank=global_ep_rank,
@@ -697,7 +695,7 @@ class ModelRunner:
if self.is_draft_worker: if self.is_draft_worker:
return return
expert_rank = self.ps.moe_ep_rank + ( expert_rank = self.ps.moe_ep_rank + (
get_parallel().config.ep_join_rank_offset if is_ep_scale_joiner() else 0 get_parallel().ep_join_rank_offset if is_ep_scale_joiner() else 0
) )
set_global_expert_location_metadata( set_global_expert_location_metadata(
compute_initial_expert_location_metadata( compute_initial_expert_location_metadata(
@@ -928,7 +926,7 @@ class ModelRunner:
device=self.device, device=self.device,
tp_group=( tp_group=(
self.attention_tp_group.cpu_group self.attention_tp_group.cpu_group
if get_parallel().config.enable_dp_attention if get_parallel().enable_dp_attention
else self.tp_group.cpu_group else self.tp_group.cpu_group
), ),
host_to_device_ratio=hisparse_cfg.host_to_device_ratio, host_to_device_ratio=hisparse_cfg.host_to_device_ratio,
@@ -964,7 +962,7 @@ class ModelRunner:
def post_capture_elastic_ep_recover(self): def post_capture_elastic_ep_recover(self):
join_process_groups() join_process_groups()
global_ep_rank = self.ps.tp_rank + get_parallel().config.ep_join_rank_offset global_ep_rank = self.ps.tp_rank + get_parallel().ep_join_rank_offset
broadcast_global_expert_location_metadata( broadcast_global_expert_location_metadata(
model_config=self.model_config, model_config=self.model_config,
moe_ep_rank=global_ep_rank, moe_ep_rank=global_ep_rank,
@@ -1004,7 +1002,7 @@ class ModelRunner:
self.decode_attn_backend = backends.decode_attn_backend self.decode_attn_backend = backends.decode_attn_backend
self.decode_attn_backend_group = backends.decode_attn_backend_group self.decode_attn_backend_group = backends.decode_attn_backend_group
if get_parallel().dcp_enabled and get_parallel().config.dcp_replicate_q_proj: if get_parallel().dcp_enabled and get_parallel().dcp_replicate_q_proj:
self._prepare_replicated_q_proj() self._prepare_replicated_q_proj()
def _prepare_replicated_q_proj(self) -> None: def _prepare_replicated_q_proj(self) -> None:
@@ -1282,7 +1280,7 @@ class ModelRunner:
def maybe_init_dwdp(self): def maybe_init_dwdp(self):
if self.is_draft_worker: if self.is_draft_worker:
return return
if get_parallel().config.dwdp_size <= 1: if get_parallel().dwdp_size <= 1:
return return
from sglang.srt.layers.moe.dwdp import DwdpManager from sglang.srt.layers.moe.dwdp import DwdpManager
@@ -1455,7 +1453,7 @@ class ModelRunner:
# rather than spawning additional processes, so dp_size must not be # rather than spawning additional processes, so dp_size must not be
# multiplied into the process count here (unlike regular DP, where # multiplied into the process count here (unlike regular DP, where
# dp_size * tp_size * pp_size is the true worker count). # dp_size * tp_size * pp_size is the true worker count).
dp_size = 1 if get_parallel().config.enable_dp_attention else self.ps.dp_size dp_size = 1 if get_parallel().enable_dp_attention else self.ps.dp_size
self.local_omp_cpuid = numa_utils.init_threads_binding( self.local_omp_cpuid = numa_utils.init_threads_binding(
numa_index=self.gpu_id, numa_index=self.gpu_id,
world_size=dp_size * self.ps.tp_size * self.ps.pp_size, world_size=dp_size * self.ps.tp_size * self.ps.pp_size,
@@ -1938,7 +1936,7 @@ class ModelRunner:
if added <= 0: if added <= 0:
return return
initial_ep_size = get_parallel().config.elastic_ep_initial_size initial_ep_size = get_parallel().elastic_ep_initial_size
assert initial_ep_size is not None assert initial_ep_size is not None
get_context().override("elastic_ep.scale", ep_size=effective_size) get_context().override("elastic_ep.scale", ep_size=effective_size)
@@ -1956,7 +1954,7 @@ class ModelRunner:
set_global_expert_location_metadata(new_metadata, allow_overwrite=True) set_global_expert_location_metadata(new_metadata, allow_overwrite=True)
def _elastic_global_rank(self) -> int: def _elastic_global_rank(self) -> int:
return self.ps.tp_rank + get_parallel().config.ep_join_rank_offset return self.ps.tp_rank + get_parallel().ep_join_rank_offset
def _rearm_eplb_after_elastic_scale(self) -> None: def _rearm_eplb_after_elastic_scale(self) -> None:
if self.eplb_manager is None: if self.eplb_manager is None:
@@ -78,13 +78,11 @@ class RemoteInstanceWeightTransporter:
""" """
import requests as http_requests import requests as http_requests
if get_parallel().config.dist_init_addr: if get_parallel().dist_init_addr:
# Multi-node: bootstrap server is on the head node (node_rank==0). # Multi-node: bootstrap server is on the head node (node_rank==0).
# Derive host from dist_init_addr (shared across all nodes). # Derive host from dist_init_addr (shared across all nodes).
bootstrap_host = ( bootstrap_host = (
NetworkAddress.parse(get_parallel().config.dist_init_addr) NetworkAddress.parse(get_parallel().dist_init_addr).resolved().host
.resolved()
.host
) )
else: else:
bootstrap_host = "127.0.0.1" bootstrap_host = "127.0.0.1"
@@ -123,8 +123,8 @@ class StartupWeightLoadOptions:
attn_cp_size=get_parallel().config.attn_cp_size, attn_cp_size=get_parallel().config.attn_cp_size,
dcp_size=get_parallel().config.dcp_size, dcp_size=get_parallel().config.dcp_size,
pp_size=get_parallel().config.pp_size, pp_size=get_parallel().config.pp_size,
dp_size=get_parallel().config.dp_size, dp_size=get_parallel().dp_size,
ep_size=get_parallel().config.ep_size, ep_size=get_parallel().ep_size,
cpu_offload_gb=get_exec().offload.cpu_offload_gb, cpu_offload_gb=get_exec().offload.cpu_offload_gb,
offload_group_size=get_exec().offload.offload_group_size, offload_group_size=get_exec().offload.offload_group_size,
enable_memory_saver=get_exec().features.enable_memory_saver, enable_memory_saver=get_exec().features.enable_memory_saver,
@@ -218,7 +218,7 @@ class BaseRunner(ABC):
self.device_module = torch.get_device_module(self.device) self.device_module = torch.get_device_module(self.device)
self.tp_size = get_parallel().config.tp_size self.tp_size = get_parallel().config.tp_size
# elastic-EP scale-up rewrites dp_size on the published config # elastic-EP scale-up rewrites dp_size on the published config
self.dp_size = get_parallel().config.dp_size self.dp_size = get_parallel().dp_size
self.pp_size = get_parallel().config.pp_size self.pp_size = get_parallel().config.pp_size
self.enable_pdmux = model_runner.server_args.enable_pdmux self.enable_pdmux = model_runner.server_args.enable_pdmux
self.return_hidden_states_mode = ( self.return_hidden_states_mode = (
@@ -289,7 +289,7 @@ class BaseRunner(ABC):
""" """
if ( if (
not get_parallel().dcp_enabled not get_parallel().dcp_enabled
or get_parallel().config.dcp_comm_backend != "fi_a2a" or get_parallel().dcp_comm_backend != "fi_a2a"
): ):
return return
@@ -348,7 +348,7 @@ class BaseRunner(ABC):
hidden_size=mr.model_config.hidden_size, hidden_size=mr.model_config.hidden_size,
vocab_size=mr.model_config.vocab_size, vocab_size=mr.model_config.vocab_size,
dtype=mr.model_config.dtype, dtype=mr.model_config.dtype,
dp_size=get_parallel().config.dp_size, dp_size=get_parallel().dp_size,
pp_size=get_parallel().config.pp_size, pp_size=get_parallel().config.pp_size,
is_encoder_decoder=mr.model_config.is_encoder_decoder, is_encoder_decoder=mr.model_config.is_encoder_decoder,
require_mlp_tp_gather=require_mlp_tp_gather(), require_mlp_tp_gather=require_mlp_tp_gather(),
@@ -541,7 +541,7 @@ class BaseRunner(ABC):
assert require_mlp_tp_gather_ or require_attn_tp_gather_ assert require_mlp_tp_gather_ or require_attn_tp_gather_
if require_mlp_tp_gather_: if require_mlp_tp_gather_:
global_num_tokens_cpu = [num_tokens] * get_parallel().config.dp_size global_num_tokens_cpu = [num_tokens] * get_parallel().dp_size
elif require_attn_tp_gather_: elif require_attn_tp_gather_:
global_num_tokens_cpu = [num_tokens] global_num_tokens_cpu = [num_tokens]
else: else:
@@ -242,7 +242,7 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
self.require_mlp_tp_gather or self.require_attn_tp_gather self.require_mlp_tp_gather or self.require_attn_tp_gather
) )
self.require_mlp_sync = ( self.require_mlp_sync = (
get_parallel().config.enable_dp_attention or self.require_gathered_buffer get_parallel().enable_dp_attention or self.require_gathered_buffer
) )
self.enable_two_batch_overlap = ( self.enable_two_batch_overlap = (
model_runner.server_args.enable_two_batch_overlap model_runner.server_args.enable_two_batch_overlap
@@ -144,7 +144,7 @@ class EagerRunner(BaseRunner):
encoder_lens_dtype=( encoder_lens_dtype=(
torch.int64 if torch.device(mr.device).type == "cpu" else torch.int32 torch.int64 if torch.device(mr.device).type == "cpu" else torch.int32
), ),
dp_size=get_parallel().config.dp_size, dp_size=get_parallel().dp_size,
) )
# Eager has no capture step, so warm up here (run-once via mr._kernel_warmed_up). # Eager has no capture step, so warm up here (run-once via mr._kernel_warmed_up).
self.warmup() self.warmup()
@@ -352,7 +352,7 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
self.moe_fusions = self.model_runner.moe_fusions self.moe_fusions = self.model_runner.moe_fusions
self.dsa_indexers = getattr(self.model_runner, "dsa_indexers", None) self.dsa_indexers = getattr(self.model_runner, "dsa_indexers", None)
self.dp_size = get_parallel().config.dp_size self.dp_size = get_parallel().dp_size
self.require_mlp_tp_gather = require_mlp_tp_gather() self.require_mlp_tp_gather = require_mlp_tp_gather()
self.require_attn_tp_gather = require_attn_tp_gather() self.require_attn_tp_gather = require_attn_tp_gather()
+2 -2
View File
@@ -1893,9 +1893,9 @@ class PreshardedModelLoader(DefaultModelLoader):
"dp": _safe(lambda: parallel.moe_dp_size), "dp": _safe(lambda: parallel.moe_dp_size),
"ep": _safe(lambda: parallel.moe_ep_size), "ep": _safe(lambda: parallel.moe_ep_size),
"pp": _safe(lambda: parallel.pp_size), "pp": _safe(lambda: parallel.pp_size),
"moe_dense_tp_size": parallel.config.moe_dense_tp_size, "moe_dense_tp_size": parallel.moe_dense_tp_size,
"moe_dp_size": get_parallel().config.moe_dp_size, "moe_dp_size": get_parallel().config.moe_dp_size,
"enable_dp_lm_head": parallel.config.enable_dp_lm_head, "enable_dp_lm_head": parallel.enable_dp_lm_head,
"enable_fp32_lm_head": get_exec().features.enable_fp32_lm_head, "enable_fp32_lm_head": get_exec().features.enable_fp32_lm_head,
"quantization": model_config.quantization, "quantization": model_config.quantization,
"model_dtype": str(model_config.dtype), "model_dtype": str(model_config.dtype),
+1 -1
View File
@@ -442,7 +442,7 @@ class ApertusForCausalLM(nn.Module):
config.hidden_size, config.hidden_size,
quant_config=quant_config, quant_config=quant_config,
prefix=add_prefix("lm_head", prefix), prefix=add_prefix("lm_head", prefix),
use_attn_tp_group=get_parallel().config.enable_dp_lm_head, use_attn_tp_group=get_parallel().enable_dp_lm_head,
) )
self.logits_processor = LogitsProcessor(config) self.logits_processor = LogitsProcessor(config)
self.pooler = Pooler(pooling_type=PoolingType.LAST, normalize=True) self.pooler = Pooler(pooling_type=PoolingType.LAST, normalize=True)
+1 -1
View File
@@ -405,7 +405,7 @@ class ArceeForCausalLM(nn.Module):
config.hidden_size, config.hidden_size,
quant_config=quant_config, quant_config=quant_config,
prefix=add_prefix("lm_head", prefix), prefix=add_prefix("lm_head", prefix),
use_attn_tp_group=get_parallel().config.enable_dp_lm_head, use_attn_tp_group=get_parallel().enable_dp_lm_head,
) )
self.logits_processor = LogitsProcessor(config) self.logits_processor = LogitsProcessor(config)
self.pooler = Pooler(pooling_type=PoolingType.LAST, normalize=True) self.pooler = Pooler(pooling_type=PoolingType.LAST, normalize=True)
+1 -1
View File
@@ -817,7 +817,7 @@ class BailingMoEForCausalLM(nn.Module):
config.hidden_size, config.hidden_size,
quant_config=quant_config, quant_config=quant_config,
prefix=add_prefix("lm_head", prefix), prefix=add_prefix("lm_head", prefix),
use_attn_tp_group=get_parallel().config.enable_dp_lm_head, use_attn_tp_group=get_parallel().enable_dp_lm_head,
) )
self.logits_processor = LogitsProcessor(config) self.logits_processor = LogitsProcessor(config)
@@ -1084,7 +1084,7 @@ class BailingMoELinearForCausalLM(nn.Module):
config.hidden_size, config.hidden_size,
params_dtype=torch.float32, params_dtype=torch.float32,
quant_config=quant_config, quant_config=quant_config,
use_attn_tp_group=get_parallel().config.enable_dp_lm_head, use_attn_tp_group=get_parallel().enable_dp_lm_head,
) )
) )
self.logits_processor = LogitsProcessor(config) self.logits_processor = LogitsProcessor(config)
@@ -241,7 +241,7 @@ class BailingMoeForCausalLMNextN(nn.Module):
config.hidden_size, config.hidden_size,
quant_config=quant_config, quant_config=quant_config,
prefix=add_prefix("model.shared_head.head", prefix), prefix=add_prefix("model.shared_head.head", prefix),
use_attn_tp_group=get_parallel().config.enable_dp_lm_head, use_attn_tp_group=get_parallel().enable_dp_lm_head,
) )
self.logits_processor = LogitsProcessor(config) self.logits_processor = LogitsProcessor(config)
if is_bailing_moe_v3: if is_bailing_moe_v3:
+1 -1
View File
@@ -1333,7 +1333,7 @@ class BailingMoeV3ForCausalLM(nn.Module):
# in the logits processor. Accuracy-neutral on ling-v3. # in the logits processor. Accuracy-neutral on ling-v3.
params_dtype=torch.bfloat16, params_dtype=torch.bfloat16,
quant_config=quant_config, quant_config=quant_config,
use_attn_tp_group=get_parallel().config.enable_dp_lm_head, use_attn_tp_group=get_parallel().enable_dp_lm_head,
) )
) )
self.logits_processor = LogitsProcessor(config) self.logits_processor = LogitsProcessor(config)
@@ -299,7 +299,7 @@ class DeepseekMLAForwardMixin:
# --dcp-replicate-q-proj: project full-head Q locally from pre-gathered # --dcp-replicate-q-proj: project full-head Q locally from pre-gathered
# weights and skip the per-layer Q all-gather (bf16 decode absorb only). # weights and skip the per-layer Q all-gather (bf16 decode absorb only).
q_replicate_active = ( q_replicate_active = (
get_parallel().config.dcp_replicate_q_proj get_parallel().dcp_replicate_q_proj
and is_dcp_mla_decode_phase(forward_batch) and is_dcp_mla_decode_phase(forward_batch)
and not self.use_deep_gemm_bmm and not self.use_deep_gemm_bmm
and self.w_kc_qrep is not None and self.w_kc_qrep is not None
@@ -779,7 +779,7 @@ class DeepseekMLAForwardMixin:
attn_output, self.num_local_heads attn_output, self.num_local_heads
) )
else: else:
dcp_comm_backend = get_parallel().config.dcp_comm_backend dcp_comm_backend = get_parallel().dcp_comm_backend
is_lse_base_on_e = is_mla_dcp_lse_base_on_e( is_lse_base_on_e = is_mla_dcp_lse_base_on_e(
self.current_attention_backend self.current_attention_backend
) )
@@ -331,7 +331,7 @@ class DeepseekMLARocmForwardMixin:
from sglang.srt.model_executor.runner import get_is_capture_mode from sglang.srt.model_executor.runner import get_is_capture_mode
q_replicate_active = ( q_replicate_active = (
get_parallel().config.dcp_replicate_q_proj get_parallel().dcp_replicate_q_proj
and is_dcp_mla_decode_phase(forward_batch) and is_dcp_mla_decode_phase(forward_batch)
and not self.use_deep_gemm_bmm and not self.use_deep_gemm_bmm
and self.w_kc_qrep is not None and self.w_kc_qrep is not None
@@ -778,7 +778,7 @@ class DeepseekMLARocmForwardMixin:
attn_output, self.num_local_heads attn_output, self.num_local_heads
) )
else: else:
dcp_comm_backend = get_parallel().config.dcp_comm_backend dcp_comm_backend = get_parallel().dcp_comm_backend
is_lse_base_on_e = is_mla_dcp_lse_base_on_e( is_lse_base_on_e = is_mla_dcp_lse_base_on_e(
self.current_attention_backend self.current_attention_backend
) )
+1 -1
View File
@@ -366,7 +366,7 @@ class DeepseekV3ForCausalLMNextN(DeepseekV3ForCausalLM):
config.hidden_size, config.hidden_size,
quant_config=quant_config, quant_config=quant_config,
prefix=add_prefix("model.shared_head.head", prefix), prefix=add_prefix("model.shared_head.head", prefix),
use_attn_tp_group=get_parallel().config.enable_dp_lm_head, use_attn_tp_group=get_parallel().enable_dp_lm_head,
) )
self.logits_processor = LogitsProcessor(config) self.logits_processor = LogitsProcessor(config)
+1 -1
View File
@@ -3007,7 +3007,7 @@ class DeepseekV2ForCausalLM(nn.Module, DeepseekV2WeightLoaderMixin):
config.hidden_size, config.hidden_size,
quant_config=quant_config, quant_config=quant_config,
prefix=add_prefix("lm_head", prefix), prefix=add_prefix("lm_head", prefix),
use_attn_tp_group=get_parallel().config.enable_dp_lm_head, use_attn_tp_group=get_parallel().enable_dp_lm_head,
) )
else: else:
# ranks other than the last rank will have a placeholder layer # ranks other than the last rank will have a placeholder layer
+1 -1
View File
@@ -3237,7 +3237,7 @@ class DeepseekV4ForCausalLM(nn.Module):
config.hidden_size, config.hidden_size,
quant_config=quant_config, quant_config=quant_config,
prefix=add_prefix("lm_head", prefix), prefix=add_prefix("lm_head", prefix),
use_attn_tp_group=get_parallel().config.enable_dp_lm_head, use_attn_tp_group=get_parallel().enable_dp_lm_head,
) )
else: else:
self.lm_head = PPMissingLayer() self.lm_head = PPMissingLayer()
@@ -758,7 +758,7 @@ class DeepseekV4ForCausalLMDSpark(nn.Module):
config.vocab_size, config.vocab_size,
config.hidden_size, config.hidden_size,
prefix=add_prefix("lm_head", prefix), prefix=add_prefix("lm_head", prefix),
use_attn_tp_group=get_parallel().config.enable_dp_lm_head, use_attn_tp_group=get_parallel().enable_dp_lm_head,
) )
else: else:
self.embed_tokens: Optional[nn.Module] = None self.embed_tokens: Optional[nn.Module] = None
@@ -250,7 +250,7 @@ class DeepseekV4ForCausalLMNextN(DeepseekV4ForCausalLM):
config.hidden_size, config.hidden_size,
quant_config=quant_config, quant_config=quant_config,
prefix=add_prefix("model.shared_head.head", prefix), prefix=add_prefix("model.shared_head.head", prefix),
use_attn_tp_group=get_parallel().config.enable_dp_lm_head, use_attn_tp_group=get_parallel().enable_dp_lm_head,
) )
self.logits_processor = LogitsProcessor(config) self.logits_processor = LogitsProcessor(config)
@@ -1876,7 +1876,7 @@ class Dots3LanguageModelForCausalLM(nn.Module):
config.hidden_size, config.hidden_size,
quant_config=quant_config, quant_config=quant_config,
prefix=add_prefix("lm_head", prefix), prefix=add_prefix("lm_head", prefix),
use_attn_tp_group=get_parallel().config.enable_dp_lm_head, use_attn_tp_group=get_parallel().enable_dp_lm_head,
) )
self.logits_processor = LogitsProcessor(config) self.logits_processor = LogitsProcessor(config)
@@ -167,7 +167,7 @@ class Dots3NoteForCausalLMNextN(Dots3LanguageModelForCausalLM):
config.hidden_size, config.hidden_size,
quant_config=quant_config, quant_config=quant_config,
prefix=add_prefix("model.shared_head.head", prefix), prefix=add_prefix("model.shared_head.head", prefix),
use_attn_tp_group=get_parallel().config.enable_dp_lm_head, use_attn_tp_group=get_parallel().enable_dp_lm_head,
) )
self.logits_processor = LogitsProcessor(config) self.logits_processor = LogitsProcessor(config)
self._mtp_loaded_embed = False self._mtp_loaded_embed = False
+1 -1
View File
@@ -439,7 +439,7 @@ class Exaone4ForCausalLM(nn.Module):
config.hidden_size, config.hidden_size,
quant_config=quant_config, quant_config=quant_config,
prefix=add_prefix("lm_head", prefix), prefix=add_prefix("lm_head", prefix),
use_attn_tp_group=get_parallel().config.enable_dp_lm_head, use_attn_tp_group=get_parallel().enable_dp_lm_head,
) )
self.logits_processor = LogitsProcessor(config) self.logits_processor = LogitsProcessor(config)
+1 -1
View File
@@ -643,7 +643,7 @@ class ExaoneMoEForCausalLM(nn.Module):
config.hidden_size, config.hidden_size,
quant_config=quant_config, quant_config=quant_config,
prefix=add_prefix("lm_head", prefix), prefix=add_prefix("lm_head", prefix),
use_attn_tp_group=get_parallel().config.enable_dp_lm_head, use_attn_tp_group=get_parallel().enable_dp_lm_head,
) )
self.logits_processor = LogitsProcessor(config) self.logits_processor = LogitsProcessor(config)
# For EAGLE3 support # For EAGLE3 support
+1 -1
View File
@@ -63,7 +63,7 @@ class ExaoneMoEForCausalLMMTP(ExaoneMoEForCausalLM):
config.hidden_size, config.hidden_size,
quant_config=quant_config, quant_config=quant_config,
prefix=add_prefix("lm_head", prefix), prefix=add_prefix("lm_head", prefix),
use_attn_tp_group=get_parallel().config.enable_dp_lm_head, use_attn_tp_group=get_parallel().enable_dp_lm_head,
) )
self.logits_processor = LogitsProcessor(config) self.logits_processor = LogitsProcessor(config)
+1 -1
View File
@@ -472,7 +472,7 @@ class FalconH1ForCausalLM(nn.Module):
quant_config=quant_config, quant_config=quant_config,
org_num_embeddings=config.vocab_size, org_num_embeddings=config.vocab_size,
prefix=add_prefix("lm_head", prefix), prefix=add_prefix("lm_head", prefix),
use_attn_tp_group=get_parallel().config.enable_dp_lm_head, use_attn_tp_group=get_parallel().enable_dp_lm_head,
) )
self.lm_head = self.lm_head.float() self.lm_head = self.lm_head.float()
self.lm_head_multiplier = config.lm_head_multiplier self.lm_head_multiplier = config.lm_head_multiplier
+1 -1
View File
@@ -1163,7 +1163,7 @@ class Glm4MoeForCausalLM(nn.Module):
config.hidden_size, config.hidden_size,
quant_config=quant_config, quant_config=quant_config,
prefix=add_prefix("lm_head", prefix), prefix=add_prefix("lm_head", prefix),
use_attn_tp_group=get_parallel().config.enable_dp_lm_head, use_attn_tp_group=get_parallel().enable_dp_lm_head,
) )
self.logits_processor = LogitsProcessor(config) self.logits_processor = LogitsProcessor(config)
+1 -1
View File
@@ -905,7 +905,7 @@ class Glm4MoeLiteForCausalLM(nn.Module, DeepseekV2WeightLoaderMixin):
config.hidden_size, config.hidden_size,
quant_config=quant_config, quant_config=quant_config,
prefix=add_prefix("lm_head", prefix), prefix=add_prefix("lm_head", prefix),
use_attn_tp_group=get_parallel().config.enable_dp_lm_head, use_attn_tp_group=get_parallel().enable_dp_lm_head,
) )
self.logits_processor = LogitsProcessor(config) self.logits_processor = LogitsProcessor(config)
@@ -159,7 +159,7 @@ class Glm4MoeLiteForCausalLMNextN(Glm4MoeLiteForCausalLM):
config.hidden_size, config.hidden_size,
quant_config=quant_config, quant_config=quant_config,
prefix=add_prefix("model.shared_head.head", prefix), prefix=add_prefix("model.shared_head.head", prefix),
use_attn_tp_group=get_parallel().config.enable_dp_lm_head, use_attn_tp_group=get_parallel().enable_dp_lm_head,
) )
self.logits_processor = LogitsProcessor(config) self.logits_processor = LogitsProcessor(config)
+1 -1
View File
@@ -144,7 +144,7 @@ class Glm4MoeForCausalLMNextN(Glm4MoeForCausalLM):
config.hidden_size, config.hidden_size,
quant_config=quant_config, quant_config=quant_config,
prefix=add_prefix("model.shared_head.head", prefix), prefix=add_prefix("model.shared_head.head", prefix),
use_attn_tp_group=get_parallel().config.enable_dp_lm_head, use_attn_tp_group=get_parallel().enable_dp_lm_head,
) )
self.logits_processor = LogitsProcessor(config) self.logits_processor = LogitsProcessor(config)
+1 -1
View File
@@ -70,7 +70,7 @@ class Glm4vMoeForConditionalGeneration(Glm4vForConditionalGeneration):
config.hidden_size, config.hidden_size,
quant_config=quant_config, quant_config=quant_config,
prefix=add_prefix("lm_head", prefix), prefix=add_prefix("lm_head", prefix),
use_attn_tp_group=get_parallel().config.enable_dp_lm_head, use_attn_tp_group=get_parallel().enable_dp_lm_head,
) )
else: else:
# ranks other than the last rank will have a placeholder layer # ranks other than the last rank will have a placeholder layer
+1 -1
View File
@@ -135,7 +135,7 @@ class GlmOcrForConditionalGenerationNextN(GlmOcrForConditionalGeneration):
config.hidden_size, config.hidden_size,
quant_config=quant_config, quant_config=quant_config,
prefix=add_prefix("model.shared_head.head", prefix), prefix=add_prefix("model.shared_head.head", prefix),
use_attn_tp_group=get_parallel().config.enable_dp_lm_head, use_attn_tp_group=get_parallel().enable_dp_lm_head,
) )
self.logits_processor = LogitsProcessor(config) self.logits_processor = LogitsProcessor(config)
+2 -2
View File
@@ -258,7 +258,7 @@ class GptOssSparseMoeBlock(nn.Module):
hidden_states: torch.Tensor, hidden_states: torch.Tensor,
forward_batch: Optional[ForwardBatch] = None, forward_batch: Optional[ForwardBatch] = None,
) -> torch.Tensor: ) -> torch.Tensor:
if get_parallel().config.dwdp_size > 1: if get_parallel().dwdp_size > 1:
return self.forward_dwdp(hidden_states) return self.forward_dwdp(hidden_states)
if not get_moe_a2a_backend().is_deepep(): if not get_moe_a2a_backend().is_deepep():
@@ -786,7 +786,7 @@ class GptOssForCausalLM(nn.Module):
config.hidden_size, config.hidden_size,
# quant_config=quant_config, # quant_config=quant_config,
prefix=add_prefix("lm_head", prefix), prefix=add_prefix("lm_head", prefix),
use_attn_tp_group=get_parallel().config.enable_dp_lm_head, use_attn_tp_group=get_parallel().enable_dp_lm_head,
) )
self.logits_processor = LogitsProcessor(config) self.logits_processor = LogitsProcessor(config)
self.capture_aux_hidden_states = False self.capture_aux_hidden_states = False
+4 -4
View File
@@ -296,7 +296,7 @@ class KimiK3MLP(nn.Module):
# but allow the NPU launcher to retain the proven attention-TP layout # but allow the NPU launcher to retain the proven attention-TP layout
# without a device-type branch in shared model code. # without a device-type branch in shared model code.
self._dense_attn_tp = ( self._dense_attn_tp = (
get_parallel().config.enable_dense_mlp_attn_tp get_parallel().enable_dense_mlp_attn_tp
and is_dp_attention_enabled() and is_dp_attention_enabled()
and tp_rank is None and tp_rank is None
and tp_size is None and tp_size is None
@@ -554,13 +554,13 @@ class KimiK3MoE(nn.Module):
# a TP-sharded partial sum could never be reduced across ranks that # a TP-sharded partial sum could never be reduced across ranks that
# hold different tokens. # hold different tokens.
self._shared_experts_tp1 = ( self._shared_experts_tp1 = (
self._ep_a2a and not get_parallel().config.enable_shared_experts_attn_tp self._ep_a2a and not get_parallel().enable_shared_experts_attn_tp
) )
# NPU compatibility mode keeps DeepEP's DP-local token dispatch but # NPU compatibility mode keeps DeepEP's DP-local token dispatch but
# uses the original TP-sharded shared MLP. Gather only that branch's # uses the original TP-sharded shared MLP. Gather only that branch's
# inputs, then reduce-scatter its output back to the DP-local rows. # inputs, then reduce-scatter its output back to the DP-local rows.
self._shared_experts_attn_tp_comm = ( self._shared_experts_attn_tp_comm = (
get_parallel().config.enable_shared_experts_attn_tp get_parallel().enable_shared_experts_attn_tp
and self._ep_a2a and self._ep_a2a
and self._dp_attention and self._dp_attention
and get_parallel().attn_tp_size > 1 and get_parallel().attn_tp_size > 1
@@ -2870,7 +2870,7 @@ class KimiK3LinearForCausalLM(nn.Module):
config.hidden_size, config.hidden_size,
quant_config=quant_config, quant_config=quant_config,
prefix=maybe_prefix(prefix, "lm_head"), prefix=maybe_prefix(prefix, "lm_head"),
use_attn_tp_group=get_parallel().config.enable_dp_lm_head, use_attn_tp_group=get_parallel().enable_dp_lm_head,
) )
else: else:
self.lm_head = PPMissingLayer() self.lm_head = PPMissingLayer()
+1 -1
View File
@@ -667,7 +667,7 @@ class LagunaForCausalLM(nn.Module):
config.hidden_size, config.hidden_size,
quant_config=quant_config, quant_config=quant_config,
prefix=add_prefix("lm_head", prefix), prefix=add_prefix("lm_head", prefix),
use_attn_tp_group=get_parallel().config.enable_dp_lm_head, use_attn_tp_group=get_parallel().enable_dp_lm_head,
) )
else: else:
self.lm_head = PPMissingLayer() self.lm_head = PPMissingLayer()
+1 -1
View File
@@ -824,7 +824,7 @@ class LLaDA2MoeModelLM(nn.Module):
config.hidden_size, config.hidden_size,
quant_config=quant_config, quant_config=quant_config,
prefix=add_prefix("lm_head", prefix), prefix=add_prefix("lm_head", prefix),
use_attn_tp_group=get_parallel().config.enable_dp_lm_head, use_attn_tp_group=get_parallel().enable_dp_lm_head,
) )
self.logits_processor = LogitsProcessor(config, return_full_logits=True) self.logits_processor = LogitsProcessor(config, return_full_logits=True)
+1 -1
View File
@@ -536,7 +536,7 @@ class LlamaForCausalLM(nn.Module):
config.hidden_size, config.hidden_size,
quant_config=quant_config, quant_config=quant_config,
prefix=add_prefix("lm_head", prefix), prefix=add_prefix("lm_head", prefix),
use_attn_tp_group=get_parallel().config.enable_dp_lm_head, use_attn_tp_group=get_parallel().enable_dp_lm_head,
) )
self.logits_processor = LogitsProcessor(config) self.logits_processor = LogitsProcessor(config)
self.pooler = Pooler(pooling_type=PoolingType.LAST, normalize=True) self.pooler = Pooler(pooling_type=PoolingType.LAST, normalize=True)
+1 -1
View File
@@ -721,7 +721,7 @@ class LongcatFlashForCausalLM(nn.Module):
config.hidden_size, config.hidden_size,
quant_config=quant_config, quant_config=quant_config,
prefix=add_prefix("lm_head", prefix), prefix=add_prefix("lm_head", prefix),
use_attn_tp_group=get_parallel().config.enable_dp_lm_head, use_attn_tp_group=get_parallel().enable_dp_lm_head,
) )
self.logits_processor = LogitsProcessor(config) self.logits_processor = LogitsProcessor(config)
self.capture_aux_hidden_states = False self.capture_aux_hidden_states = False
+1 -1
View File
@@ -520,7 +520,7 @@ class MellumForCausalLM(Qwen3MoeForCausalLM):
cfg.hidden_size, cfg.hidden_size,
quant_config=quant_config, quant_config=quant_config,
prefix=add_prefix("lm_head", prefix), prefix=add_prefix("lm_head", prefix),
use_attn_tp_group=get_parallel().config.enable_dp_lm_head, use_attn_tp_group=get_parallel().enable_dp_lm_head,
) )
self.logits_processor = LogitsProcessor(cfg) self.logits_processor = LogitsProcessor(cfg)
self.capture_aux_hidden_states = False self.capture_aux_hidden_states = False
+1 -1
View File
@@ -1187,7 +1187,7 @@ class MiMoV2ForCausalLM(nn.Module, AudioEncoderMixin):
config.hidden_size, config.hidden_size,
quant_config=quant_config, quant_config=quant_config,
prefix=add_prefix("lm_head", prefix), prefix=add_prefix("lm_head", prefix),
use_attn_tp_group=get_parallel().config.enable_dp_lm_head, use_attn_tp_group=get_parallel().enable_dp_lm_head,
) )
else: else:
self.lm_head = PPMissingLayer() self.lm_head = PPMissingLayer()
+1 -1
View File
@@ -259,7 +259,7 @@ class MiMoV2MTP(MiMoV2ForCausalLM):
config.hidden_size, config.hidden_size,
quant_config=quant_config, quant_config=quant_config,
prefix=add_prefix("lm_head", prefix), prefix=add_prefix("lm_head", prefix),
use_attn_tp_group=get_parallel().config.enable_dp_lm_head, use_attn_tp_group=get_parallel().enable_dp_lm_head,
) )
self.logits_processor = LogitsProcessor(config) self.logits_processor = LogitsProcessor(config)
+1 -1
View File
@@ -1574,7 +1574,7 @@ class MiniMaxM3SparseForCausalLM(nn.Module):
config.hidden_size, config.hidden_size,
quant_config=quant_config, quant_config=quant_config,
prefix=add_prefix("lm_head", prefix), prefix=add_prefix("lm_head", prefix),
use_attn_tp_group=get_parallel().config.enable_dp_lm_head, use_attn_tp_group=get_parallel().enable_dp_lm_head,
) )
self.logits_processor = LogitsProcessor(config) self.logits_processor = LogitsProcessor(config)
+1 -1
View File
@@ -123,7 +123,7 @@ class MiniMaxM3SparseForConditionalGeneration(nn.Module):
text_config.hidden_size, text_config.hidden_size,
quant_config=quant_config, quant_config=quant_config,
prefix=add_prefix("language_model.lm_head", prefix), prefix=add_prefix("language_model.lm_head", prefix),
use_attn_tp_group=get_parallel().config.enable_dp_lm_head, use_attn_tp_group=get_parallel().enable_dp_lm_head,
) )
else: else:
self.lm_head = PPMissingLayer() self.lm_head = PPMissingLayer()
+1 -1
View File
@@ -964,7 +964,7 @@ class NemotronHForCausalLM(nn.Module):
else lora_config.lora_vocab_padding_size else lora_config.lora_vocab_padding_size
), ),
quant_config=quant_config, quant_config=quant_config,
use_attn_tp_group=get_parallel().config.enable_dp_lm_head, use_attn_tp_group=get_parallel().enable_dp_lm_head,
prefix=add_prefix("lm_head", prefix), prefix=add_prefix("lm_head", prefix),
) )
else: else:
+1 -1
View File
@@ -339,7 +339,7 @@ class NemotronHForCausalLMMTP(NemotronHForCausalLM):
self.config.hidden_size, self.config.hidden_size,
quant_config=quant_config, quant_config=quant_config,
prefix=add_prefix("lm_head", prefix), prefix=add_prefix("lm_head", prefix),
use_attn_tp_group=get_parallel().config.enable_dp_lm_head, use_attn_tp_group=get_parallel().enable_dp_lm_head,
) )
self.logits_processor = LogitsProcessor(config) self.logits_processor = LogitsProcessor(config)
+1 -1
View File
@@ -1119,7 +1119,7 @@ class Qwen2MoeForCausalLM(nn.Module):
config.hidden_size, config.hidden_size,
quant_config=quant_config, quant_config=quant_config,
prefix=add_prefix("lm_head", prefix), prefix=add_prefix("lm_head", prefix),
use_attn_tp_group=get_parallel().config.enable_dp_lm_head, use_attn_tp_group=get_parallel().enable_dp_lm_head,
) )
self.logits_processor = LogitsProcessor(config) self.logits_processor = LogitsProcessor(config)
# For EAGLE3 support # For EAGLE3 support
+1 -1
View File
@@ -492,7 +492,7 @@ class Qwen3ForCausalLM(nn.Module):
config.vocab_size, config.vocab_size,
config.hidden_size, config.hidden_size,
quant_config=quant_config, quant_config=quant_config,
use_attn_tp_group=get_parallel().config.enable_dp_lm_head, use_attn_tp_group=get_parallel().enable_dp_lm_head,
prefix=add_prefix("lm_head", prefix), prefix=add_prefix("lm_head", prefix),
) )
else: else:
+1 -1
View File
@@ -80,7 +80,7 @@ class Qwen3_5ForCausalLM(nn.Module):
quant_config=quant_config, quant_config=quant_config,
org_num_embeddings=config.vocab_size, org_num_embeddings=config.vocab_size,
prefix=add_prefix("lm_head", prefix), prefix=add_prefix("lm_head", prefix),
use_attn_tp_group=get_parallel().config.enable_dp_lm_head, use_attn_tp_group=get_parallel().enable_dp_lm_head,
) )
else: else:
self.lm_head = PPMissingLayer() self.lm_head = PPMissingLayer()
+1 -1
View File
@@ -961,7 +961,7 @@ class Qwen3MoeForCausalLM(nn.Module):
config.hidden_size, config.hidden_size,
quant_config=quant_config, quant_config=quant_config,
prefix=add_prefix("lm_head", prefix), prefix=add_prefix("lm_head", prefix),
use_attn_tp_group=get_parallel().config.enable_dp_lm_head, use_attn_tp_group=get_parallel().enable_dp_lm_head,
) )
self.logits_processor = LogitsProcessor(config) self.logits_processor = LogitsProcessor(config)
self.capture_aux_hidden_states = False self.capture_aux_hidden_states = False
+1 -1
View File
@@ -63,7 +63,7 @@ class Qwen3MoeForCausalLMMTP(Qwen3MoeForCausalLM):
config.hidden_size, config.hidden_size,
quant_config=quant_config, quant_config=quant_config,
prefix=add_prefix("lm_head", prefix), prefix=add_prefix("lm_head", prefix),
use_attn_tp_group=get_parallel().config.enable_dp_lm_head, use_attn_tp_group=get_parallel().enable_dp_lm_head,
) )
self.logits_processor = LogitsProcessor(config) self.logits_processor = LogitsProcessor(config)
+1 -1
View File
@@ -1027,7 +1027,7 @@ class Qwen3NextForCausalLM(nn.Module):
quant_config=quant_config, quant_config=quant_config,
org_num_embeddings=config.vocab_size, org_num_embeddings=config.vocab_size,
prefix=add_prefix("lm_head", prefix), prefix=add_prefix("lm_head", prefix),
use_attn_tp_group=get_parallel().config.enable_dp_lm_head, use_attn_tp_group=get_parallel().enable_dp_lm_head,
) )
self.logits_processor = LogitsProcessor(config) self.logits_processor = LogitsProcessor(config)
# For EAGLE3 support # For EAGLE3 support
+1 -1
View File
@@ -80,7 +80,7 @@ class Qwen3NextForCausalLMMTP(Qwen3NextForCausalLM):
config.hidden_size, config.hidden_size,
quant_config=quant_config, quant_config=quant_config,
prefix=add_prefix("model.shared_head.head", prefix), prefix=add_prefix("model.shared_head.head", prefix),
use_attn_tp_group=get_parallel().config.enable_dp_lm_head, use_attn_tp_group=get_parallel().enable_dp_lm_head,
) )
self.logits_processor = LogitsProcessor(config) self.logits_processor = LogitsProcessor(config)
# Mirror Qwen3NextForCausalLM.__init__'s shared-expert fusion setup so # Mirror Qwen3NextForCausalLM.__init__'s shared-expert fusion setup so
+1 -1
View File
@@ -1316,7 +1316,7 @@ class Qwen3VLForConditionalGeneration(nn.Module):
self.config.vocab_size, self.config.vocab_size,
self.config.hidden_size, self.config.hidden_size,
quant_config=quant_config, quant_config=quant_config,
use_attn_tp_group=get_parallel().config.enable_dp_lm_head, use_attn_tp_group=get_parallel().enable_dp_lm_head,
prefix=add_prefix("lm_head", prefix), prefix=add_prefix("lm_head", prefix),
) )
else: else:
+1 -1
View File
@@ -1220,7 +1220,7 @@ class SarvamMLAForCausalLM(nn.Module):
config.hidden_size, config.hidden_size,
quant_config=quant_config, quant_config=quant_config,
prefix=add_prefix("lm_head", prefix), prefix=add_prefix("lm_head", prefix),
use_attn_tp_group=get_parallel().config.enable_dp_lm_head, use_attn_tp_group=get_parallel().enable_dp_lm_head,
) )
self.logits_processor = LogitsProcessor(config) self.logits_processor = LogitsProcessor(config)

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