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:
co-authored by
Claude Opus 5
parent
7c3b5a6732
commit
ca1d7ed8e6
@@ -548,7 +548,7 @@ def _maybe_prepare_mlp_sync_batch(batch: ScheduleBatch, model_runner):
|
||||
prepare_mlp_sync_batch_raw(
|
||||
batch,
|
||||
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_cp_size=model_runner.ps.attn_cp_size,
|
||||
tp_group=model_runner.tp_group,
|
||||
|
||||
@@ -65,7 +65,7 @@ async def warm_up_compile(
|
||||
disaggregation_mode: str, tokenizer_manager: TokenizerManager
|
||||
):
|
||||
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]
|
||||
sampling_params = {
|
||||
"temperature": 0.0,
|
||||
|
||||
@@ -787,7 +787,7 @@ class TboForwardBatchPreparer:
|
||||
|
||||
# TODO improve, e.g. unify w/ `init_raw`
|
||||
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
|
||||
):
|
||||
sum_len = end_token_index - start_token_index
|
||||
|
||||
@@ -170,7 +170,7 @@ class CommonKVManager(BaseKVManager):
|
||||
# for p/d multi node infer
|
||||
self.bootstrap_host = get_serving().host
|
||||
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()
|
||||
self.attn_tp_size = parallel.attn_tp_size
|
||||
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_rank = get_attention_dp_rank()
|
||||
self.system_dp_size = (
|
||||
1
|
||||
if get_parallel().config.enable_dp_attention
|
||||
else get_parallel().config.dp_size
|
||||
1 if get_parallel().enable_dp_attention else get_parallel().dp_size
|
||||
)
|
||||
self.system_dp_rank = (
|
||||
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.local_ip = get_local_ip_auto()
|
||||
cp_sharded_prefill = self.attn_cp_size > 1 and (
|
||||
self.is_hybrid_mla_backend
|
||||
or get_parallel().config.enable_dsa_cache_layer_split
|
||||
self.is_hybrid_mla_backend or get_parallel().enable_dsa_cache_layer_split
|
||||
)
|
||||
|
||||
hybrid_decode_pulls_all_ranks = (
|
||||
@@ -306,7 +303,7 @@ class CommonKVManager(BaseKVManager):
|
||||
return (
|
||||
self.attn_cp_size > 1
|
||||
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:
|
||||
@@ -751,7 +748,7 @@ class CommonKVManager(BaseKVManager):
|
||||
`Connection refused`, and the leader's `prefill_port_table` ends
|
||||
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
|
||||
|
||||
if not (dist.is_available() and dist.is_initialized()):
|
||||
@@ -803,8 +800,8 @@ class CommonKVManager(BaseKVManager):
|
||||
"rank_port": self.rank_port,
|
||||
"page_size": self.kv_args.page_size,
|
||||
"kv_cache_dtype": self.kv_cache_dtype_str,
|
||||
"load_balance_method": get_parallel().config.load_balance_method,
|
||||
"enable_dsa_cache_layer_split": get_parallel().config.enable_dsa_cache_layer_split,
|
||||
"load_balance_method": get_parallel().load_balance_method,
|
||||
"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
|
||||
# retract rebootstrap /generate URL from bootstrap info instead of a
|
||||
# router-injected pd_rebootstrap_prefill_url.
|
||||
@@ -1186,12 +1183,11 @@ class CommonKVSender(BaseKVSender):
|
||||
return
|
||||
|
||||
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().config.load_balance_method != "follow_bootstrap_room":
|
||||
if get_parallel().dp_size > 1 and not req_has_disagg_prefill_dp_rank:
|
||||
if get_parallel().load_balance_method != "follow_bootstrap_room":
|
||||
self._register_prefill_dp_rank()
|
||||
elif (
|
||||
self.kv_mgr.attn_dp_rank
|
||||
!= self.bootstrap_room % get_parallel().config.dp_size
|
||||
self.kv_mgr.attn_dp_rank != self.bootstrap_room % get_parallel().dp_size
|
||||
):
|
||||
# follow_bootstrap_room was overridden by external routed_dp_rank
|
||||
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"{self.kv_mgr.attn_dp_rank} but bootstrap_room "
|
||||
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"to allow mixed routing.",
|
||||
)
|
||||
@@ -1276,7 +1272,7 @@ class CommonKVSender(BaseKVSender):
|
||||
|
||||
if (
|
||||
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(
|
||||
self.kv_mgr,
|
||||
|
||||
@@ -207,7 +207,7 @@ def launch_server(server_args: ServerArgs):
|
||||
|
||||
configure_logger(server_args, prefix=" encode_server")
|
||||
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)
|
||||
# runtime initializes multiprocess metrics before spawning;
|
||||
# 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,
|
||||
service registration, and network serving remain Transport concerns.
|
||||
"""
|
||||
if get_parallel().config.dp_size > 1:
|
||||
if get_parallel().dp_size > 1:
|
||||
raise ValueError(
|
||||
"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.
|
||||
@@ -1513,10 +1513,8 @@ def launch_local_runtime(server_args: ServerArgs) -> EncoderRuntime:
|
||||
zmq_context = zmq.Context(10)
|
||||
ipc_path_prefix = random_uuid()
|
||||
port_args = PortArgs.init_new(server_args)
|
||||
if get_parallel().config.dist_init_addr:
|
||||
dist_init_method = NetworkAddress.parse(
|
||||
get_parallel().config.dist_init_addr
|
||||
).to_tcp()
|
||||
if get_parallel().dist_init_addr:
|
||||
dist_init_method = NetworkAddress.parse(get_parallel().dist_init_addr).to_tcp()
|
||||
else:
|
||||
dist_init_method = NetworkAddress(
|
||||
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
|
||||
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(
|
||||
"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}")
|
||||
|
||||
# DP mode: workers (subprocesses) write metrics to the shared multiproc dir;
|
||||
|
||||
@@ -180,7 +180,7 @@ class PrefillBootstrapQueue:
|
||||
"SGLANG_DISAGG_STAGING_BUFFER with pp_size > 1 is only "
|
||||
"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.
|
||||
raise RuntimeError(
|
||||
"SGLANG_DISAGG_STAGING_BUFFER does not support "
|
||||
|
||||
@@ -129,7 +129,7 @@ def init_torch_distributed(
|
||||
# included in later KV-cache sizing instead of appearing during capture.
|
||||
if (
|
||||
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
|
||||
):
|
||||
_prewarm_tp_lm_head_all_to_all()
|
||||
@@ -267,7 +267,7 @@ def _init_parallel_groups(
|
||||
duplicate_attn_cp_group=(
|
||||
is_hip()
|
||||
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,
|
||||
recovered_rank=is_ep_joiner,
|
||||
|
||||
@@ -477,7 +477,7 @@ class MultimemAllGatherer:
|
||||
# EP/mooncake setups, and keep multimem enabled.
|
||||
if (
|
||||
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))
|
||||
):
|
||||
logger.warning(
|
||||
|
||||
@@ -92,7 +92,7 @@ class ElasticEPStateManager:
|
||||
|
||||
if get_exec().moe.elastic_ep_backend is not None:
|
||||
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, (
|
||||
f"--max-ep-size ({active_rank_capacity}) must be >= "
|
||||
f"world_size ({world_size})."
|
||||
@@ -109,7 +109,7 @@ class ElasticEPStateManager:
|
||||
if get_exec().moe.moe_a2a_backend == "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:
|
||||
cls._init_joiner_state(inst)
|
||||
|
||||
@@ -127,12 +127,11 @@ class ElasticEPStateManager:
|
||||
|
||||
if get_exec().moe.ep_join_mode == "scale":
|
||||
inst.effective_ep_size = (
|
||||
get_parallel().config.ep_join_rank_offset
|
||||
+ get_parallel().config.tp_size
|
||||
get_parallel().ep_join_rank_offset + get_parallel().config.tp_size
|
||||
)
|
||||
inst.original_ep_size = (
|
||||
get_parallel().config.elastic_ep_initial_size
|
||||
or get_parallel().config.ep_join_rank_offset
|
||||
get_parallel().elastic_ep_initial_size
|
||||
or get_parallel().ep_join_rank_offset
|
||||
)
|
||||
inst.has_scaled = True
|
||||
else:
|
||||
@@ -317,7 +316,7 @@ def elastic_expanded_world_enabled() -> bool:
|
||||
inst = ElasticEPStateManager.instance()
|
||||
if inst is None:
|
||||
return False
|
||||
if get_parallel().config.max_ep_size is None:
|
||||
if get_parallel().max_ep_size is None:
|
||||
return False
|
||||
active_target_size = inst.effective_ep_size
|
||||
if inst.pending_ep_size is not None and inst.scale_phase in (
|
||||
|
||||
@@ -50,8 +50,8 @@ class ExpertBackupManager:
|
||||
self.weight_pointer_map = {}
|
||||
self.transfer_engine = None
|
||||
self.session_id = None
|
||||
self.engine_num = get_parallel().config.nnodes
|
||||
self.engine_rank = get_parallel().config.node_rank
|
||||
self.engine_num = get_parallel().nnodes
|
||||
self.engine_rank = get_parallel().node_rank
|
||||
self.expert_num = self.model_config.hf_config.n_routed_experts
|
||||
self.idmn = (self.expert_num // self.engine_num) * self.engine_rank
|
||||
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.
|
||||
self.recv_from_expert_backup_client = context.socket(zmq.PULL)
|
||||
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.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.start_transfer_server()
|
||||
|
||||
@@ -344,7 +344,7 @@ class Engine(EngineScoreMixin, EngineBase):
|
||||
routed_dp_rank = data_parallel_rank
|
||||
|
||||
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:
|
||||
logger.debug(
|
||||
f"routed_dp_rank={routed_dp_rank} is ignored because dp_size={dp_size}"
|
||||
@@ -831,7 +831,7 @@ class Engine(EngineScoreMixin, EngineBase):
|
||||
"""
|
||||
scheduler_procs = []
|
||||
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:
|
||||
@@ -1841,11 +1841,7 @@ def _compute_parallelism_ranks(
|
||||
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.
|
||||
"""
|
||||
attn_dp_size = (
|
||||
get_parallel().config.dp_size
|
||||
if get_parallel().config.enable_dp_attention
|
||||
else 1
|
||||
)
|
||||
attn_dp_size = get_parallel().dp_size if get_parallel().enable_dp_attention else 1
|
||||
tp_size = server_args.tp_size
|
||||
attn_cp_size = get_parallel().config.attn_cp_size
|
||||
moe_dp_size = get_parallel().config.moe_dp_size
|
||||
@@ -1859,6 +1855,6 @@ def _compute_parallelism_ranks(
|
||||
moe_ep_rank = (
|
||||
tp_rank
|
||||
% (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
|
||||
|
||||
@@ -2180,7 +2180,7 @@ async def _send_disaggregation_warmup_requests(
|
||||
return await asyncio.gather(
|
||||
*(
|
||||
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:
|
||||
json_data["input_ids"] = [
|
||||
[10, 11, 12] for _ in range(get_parallel().config.dp_size)
|
||||
]
|
||||
json_data["input_ids"] = [[10, 11, 12] for _ in range(get_parallel().dp_size)]
|
||||
# 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]
|
||||
elif (
|
||||
is_vlm
|
||||
@@ -2287,11 +2285,9 @@ def _execute_server_warmup(server_args: ServerArgs):
|
||||
"temperature": 0.0,
|
||||
}
|
||||
else:
|
||||
json_data["text"] = [
|
||||
"The capital city of France is"
|
||||
] * get_parallel().config.dp_size
|
||||
json_data["text"] = ["The capital city of France is"] * get_parallel().dp_size
|
||||
# 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]
|
||||
|
||||
# Config debug dumping
|
||||
@@ -2332,7 +2328,7 @@ def _execute_server_warmup(server_args: ServerArgs):
|
||||
if not failed_status_codes:
|
||||
logger.info(
|
||||
"Disaggregation warmup requests completed for all %s DP ranks",
|
||||
get_parallel().config.dp_size,
|
||||
get_parallel().dp_size,
|
||||
)
|
||||
logger.info("End of disaggregation warmup")
|
||||
else:
|
||||
|
||||
@@ -148,8 +148,8 @@ async def get_loads(
|
||||
"num_accelerators": _num_accelerators_per_dp_rank(
|
||||
get_parallel().config.tp_size,
|
||||
get_parallel().config.pp_size,
|
||||
get_parallel().config.dp_size,
|
||||
get_parallel().config.enable_dp_attention,
|
||||
get_parallel().dp_size,
|
||||
get_parallel().enable_dp_attention,
|
||||
),
|
||||
"loads": loads,
|
||||
}
|
||||
|
||||
@@ -157,7 +157,7 @@ class EPLBManager:
|
||||
model=self._get_model(),
|
||||
new_expert_location_metadata=expert_location_metadata,
|
||||
update_layer_ids=chunk_layer_ids,
|
||||
nnodes=get_parallel().config.nnodes,
|
||||
nnodes=get_parallel().nnodes,
|
||||
tp_rank=(
|
||||
self._elastic_global_rank()
|
||||
if is_post_scale_rebalance
|
||||
@@ -221,7 +221,7 @@ class EPLBManager:
|
||||
)
|
||||
|
||||
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):
|
||||
if average_utilization_rate_over_window is None:
|
||||
|
||||
@@ -195,7 +195,7 @@ class ExpertLocationMetadata:
|
||||
model_config_for_expert_location = common["model_config_for_expert_location"]
|
||||
num_physical_experts = common["num_physical_experts"]
|
||||
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
|
||||
|
||||
@@ -238,15 +238,14 @@ class ExpertLocationMetadata:
|
||||
+ get_exec().moe.ep_num_redundant_experts
|
||||
)
|
||||
# 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
|
||||
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 get_exec().moe.ep_join_mode == "scale":
|
||||
ep_size = max(
|
||||
ep_size,
|
||||
get_parallel().config.ep_join_rank_offset
|
||||
+ get_parallel().config.tp_size,
|
||||
get_parallel().ep_join_rank_offset + get_parallel().config.tp_size,
|
||||
)
|
||||
num_physical_experts, num_local_physical_experts = (
|
||||
_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
|
||||
prefer_same_node = _prefer_same_node_experts()
|
||||
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
|
||||
)
|
||||
@@ -634,9 +633,7 @@ def compute_logical_to_rank_dispatch_physical_map(
|
||||
num_local_gpu_physical_experts = num_physical_experts // ep_size
|
||||
prefer_same_node = _prefer_same_node_experts()
|
||||
num_gpus_per_node = (
|
||||
get_parallel().config.ep_size // get_parallel().config.nnodes
|
||||
if prefer_same_node
|
||||
else None
|
||||
get_parallel().ep_size // get_parallel().nnodes if prefer_same_node else None
|
||||
)
|
||||
num_local_node_physical_experts = (
|
||||
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():
|
||||
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
|
||||
# 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():
|
||||
return (
|
||||
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():
|
||||
return (
|
||||
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:
|
||||
return bool(
|
||||
get_parallel().config.enable_prefill_cp
|
||||
or get_parallel().config.enable_dp_attention
|
||||
)
|
||||
return bool(get_parallel().enable_prefill_cp or get_parallel().enable_dp_attention)
|
||||
|
||||
|
||||
@dataclass
|
||||
|
||||
@@ -269,7 +269,7 @@ class AttnTpContext:
|
||||
def init_context(self, q_lora_rank, is_dsa):
|
||||
self.is_dsa = is_dsa
|
||||
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 q_lora_rank is not None
|
||||
and not is_dsa
|
||||
@@ -280,7 +280,7 @@ class AttnTpContext:
|
||||
and not check_cuda_graph_backend(Phase.PREFILL, Backend.TC_PIECEWISE)
|
||||
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:
|
||||
logging.info(
|
||||
"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():
|
||||
return get_parallel().config.moe_dense_tp_size == 1
|
||||
return get_parallel().moe_dense_tp_size == 1
|
||||
|
||||
|
||||
def enable_dwdp():
|
||||
return get_parallel().config.dwdp_size > 1
|
||||
return get_parallel().dwdp_size > 1
|
||||
|
||||
|
||||
class LayerCommunicator:
|
||||
|
||||
@@ -51,7 +51,7 @@ class CpDecodeAttnTpContext:
|
||||
"""Slices replicated attention weights across CP ranks during decode."""
|
||||
|
||||
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:
|
||||
self.decode_tp_rank = get_parallel().attn_cp_rank
|
||||
|
||||
@@ -63,7 +63,7 @@ def is_glm_dsa_cache_layer_split_enabled(model_runner: "ModelRunner") -> bool:
|
||||
|
||||
return (
|
||||
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 is_deepseek_dsa(model_runner.model_config.hf_config)
|
||||
)
|
||||
|
||||
@@ -347,8 +347,8 @@ def initialize_dp_attention(
|
||||
dp.max_len_with_idle = (
|
||||
getattr(model_config.hf_config, "hybrid_override_pattern", None) is not None
|
||||
)
|
||||
enable_dp_attention = get_parallel().config.enable_dp_attention
|
||||
dp_size = get_parallel().config.dp_size
|
||||
enable_dp_attention = get_parallel().enable_dp_attention
|
||||
dp_size = get_parallel().dp_size
|
||||
attn_cp_size = get_parallel().config.attn_cp_size
|
||||
|
||||
dp.enabled = enable_dp_attention
|
||||
@@ -361,11 +361,8 @@ def initialize_dp_attention(
|
||||
)
|
||||
_ATTN_DP_SIZE = dp_size if enable_dp_attention else 1
|
||||
|
||||
if (
|
||||
get_exec().moe.elastic_ep_backend is not None
|
||||
and get_parallel().config.max_ep_size
|
||||
):
|
||||
_ATTN_DP_RANK = tp_rank + get_parallel().config.ep_join_rank_offset
|
||||
if get_exec().moe.elastic_ep_backend is not None and get_parallel().max_ep_size:
|
||||
_ATTN_DP_RANK = tp_rank + get_parallel().ep_join_rank_offset
|
||||
if server_args.is_ep_scale_joiner:
|
||||
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
|
||||
if backend is None:
|
||||
return None
|
||||
return _resolve_backend(backend, get_parallel().config.nnodes > 1)
|
||||
return _resolve_backend(backend, get_parallel().nnodes > 1)
|
||||
|
||||
|
||||
if is_flashinfer_available():
|
||||
|
||||
@@ -301,10 +301,8 @@ class LogitsProcessor(nn.Module):
|
||||
self.config = config
|
||||
self.vocab_size = config.vocab_size
|
||||
self.logit_scale = logit_scale
|
||||
self.use_attn_tp_group = get_parallel().config.enable_dp_lm_head
|
||||
self.use_tp_lm_head_all_to_all = (
|
||||
get_parallel().config.enable_tp_lm_head_all_to_all
|
||||
)
|
||||
self.use_attn_tp_group = get_parallel().enable_dp_lm_head
|
||||
self.use_tp_lm_head_all_to_all = get_parallel().enable_tp_lm_head_all_to_all
|
||||
self.use_fp32_lm_head = get_exec().features.enable_fp32_lm_head
|
||||
if self.use_attn_tp_group:
|
||||
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
|
||||
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
|
||||
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:
|
||||
storage_ep_size = self.moe_ep_size
|
||||
|
||||
@@ -337,7 +337,7 @@ def ensure_cutedsl_wrapper(layer: torch.nn.Module) -> None:
|
||||
else:
|
||||
# Standard allgather path: the MoE sees up to dp_size local forwards
|
||||
# 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
|
||||
# inference_mode(False) ensures the wrapper's pre-allocated CUDA-graph
|
||||
# buffers are normal tensors. This call typically happens inside
|
||||
|
||||
@@ -135,7 +135,7 @@ class NixlEPBuffer:
|
||||
offset = ElasticEPStateManager.get_ep_join_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
|
||||
|
||||
num_rdma_bytes = 0
|
||||
@@ -233,7 +233,7 @@ class _NixlEPDispatcherImplBase:
|
||||
)
|
||||
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 = (
|
||||
torch.zeros(_max_ep, dtype=torch.int32, device="cuda")
|
||||
if self.active_ranks is not None
|
||||
|
||||
@@ -155,7 +155,7 @@ class PplxAllToAllManager:
|
||||
# pplx forces ep_size == world_size
|
||||
# with pp_size == 1 (enforced in _ensure_nvshmem), so the EP group spans
|
||||
# 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:
|
||||
cls._all_to_all = AllToAll.internode(
|
||||
|
||||
@@ -655,7 +655,7 @@ def should_skip_post_experts_all_reduce(*, is_tp_path: bool) -> bool:
|
||||
"""
|
||||
if should_skip_mlp_all_reduce():
|
||||
return True
|
||||
if get_parallel().config.dwdp_size > 1:
|
||||
if get_parallel().dwdp_size > 1:
|
||||
return True
|
||||
if should_use_dp_reduce_scatterv():
|
||||
return True
|
||||
|
||||
@@ -63,18 +63,18 @@ class ContextParallelMetadata:
|
||||
|
||||
|
||||
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():
|
||||
return (
|
||||
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:
|
||||
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):
|
||||
|
||||
@@ -99,7 +99,7 @@ class LoRAManager:
|
||||
self.pending_lora_load_events = {}
|
||||
|
||||
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] = (
|
||||
get_lora().experts_shared_outer_loras
|
||||
)
|
||||
|
||||
@@ -148,12 +148,12 @@ class DataParallelController:
|
||||
self.server_args = server_args
|
||||
self.port_args = port_args
|
||||
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
|
||||
|
||||
# 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:
|
||||
self.recv_from_tokenizer = get_zmq_socket(
|
||||
self.context, zmq.PULL, port_args.scheduler_input_ipc_name, False
|
||||
@@ -173,8 +173,8 @@ class DataParallelController:
|
||||
LoadBalanceMethod.TOTAL_TOKENS,
|
||||
)
|
||||
|
||||
self.launch_dp_size: int = get_parallel().config.dp_size
|
||||
self.max_dp_size: int = server_args.max_ep_size or 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().dp_size
|
||||
assert self.max_dp_size >= self.launch_dp_size, (
|
||||
f"--max-ep-size ({self.max_dp_size}) must be >= "
|
||||
f"--dp ({self.launch_dp_size})."
|
||||
@@ -184,7 +184,7 @@ class DataParallelController:
|
||||
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(
|
||||
port_args,
|
||||
caller="DataParallelController",
|
||||
@@ -201,16 +201,14 @@ class DataParallelController:
|
||||
self._active_workers: List[int] = list(range(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)
|
||||
# When local control broadcast is enabled, send control messages to
|
||||
# every DP group leader (attn_tp_rank=0) so each leader broadcasts
|
||||
# within its own attn_tp_group instead of the full tp_group.
|
||||
# Otherwise fall back to the original behaviour: send to only the
|
||||
# first leader, which then broadcasts over the full tp_group.
|
||||
local_ctrl = (
|
||||
get_parallel().config.enable_dp_attention_local_control_broadcast
|
||||
)
|
||||
local_ctrl = get_parallel().enable_dp_attention_local_control_broadcast
|
||||
self.control_message_step = 1 if local_ctrl else server_args.tp_size
|
||||
else:
|
||||
self.launch_dp_schedulers(server_args, port_args)
|
||||
@@ -374,7 +372,7 @@ class DataParallelController:
|
||||
threads = []
|
||||
sockets = []
|
||||
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.tokenizer_ipc_name = port_args.tokenizer_ipc_name
|
||||
tmp_port_args.detokenizer_ipc_name = port_args.detokenizer_ipc_name
|
||||
@@ -578,7 +576,7 @@ class DataParallelController:
|
||||
bind_count = (
|
||||
self.max_dp_size
|
||||
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):
|
||||
worker_port, worker_socket = get_zmq_socket_on_host(
|
||||
@@ -608,7 +606,7 @@ class DataParallelController:
|
||||
dp_rank: Optional[int],
|
||||
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}.")
|
||||
|
||||
memory_saver_adapter = TorchMemorySaverAdapter.create(
|
||||
@@ -642,13 +640,13 @@ class DataParallelController:
|
||||
for tp_rank in tp_rank_range:
|
||||
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_rank, _ = compute_dp_attention_world_info(
|
||||
get_parallel().config.enable_dp_attention,
|
||||
get_parallel().enable_dp_attention,
|
||||
tp_rank,
|
||||
server_args.tp_size,
|
||||
get_parallel().config.dp_size,
|
||||
get_parallel().dp_size,
|
||||
get_parallel().config.attn_cp_size,
|
||||
)
|
||||
# compute zmq ports for this dp rank
|
||||
@@ -678,9 +676,7 @@ class DataParallelController:
|
||||
+ (tp_rank % tp_size_per_node) * server_args.gpu_id_step
|
||||
)
|
||||
attn_dp_size = (
|
||||
get_parallel().config.dp_size
|
||||
if get_parallel().config.enable_dp_attention
|
||||
else 1
|
||||
get_parallel().dp_size if get_parallel().enable_dp_attention else 1
|
||||
)
|
||||
|
||||
# Parallelism hierarchy (outermost to innermost):
|
||||
@@ -703,7 +699,7 @@ class DataParallelController:
|
||||
// (
|
||||
server_args.tp_size
|
||||
// get_parallel().config.moe_dp_size
|
||||
// get_parallel().config.ep_size
|
||||
// get_parallel().ep_size
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
@@ -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
|
||||
``start_disagg_service`` duty left to perform."""
|
||||
if not (
|
||||
get_parallel().config.node_rank == 0
|
||||
and transfer_backend == TransferBackend.ASCEND
|
||||
get_parallel().node_rank == 0 and transfer_backend == TransferBackend.ASCEND
|
||||
):
|
||||
return
|
||||
try:
|
||||
|
||||
@@ -71,7 +71,7 @@ def should_use_zmq() -> bool:
|
||||
``SGLANG_LOAD_SNAPSHOT_USE_ZMQ`` forces zmq mode for testing.
|
||||
"""
|
||||
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()
|
||||
|
||||
|
||||
@@ -116,15 +116,15 @@ def zmq_reader_owner(caller: str) -> bool:
|
||||
"""
|
||||
if not should_use_zmq():
|
||||
return False
|
||||
if get_parallel().config.node_rank != 0:
|
||||
if get_parallel().node_rank != 0:
|
||||
return False
|
||||
if caller == "DataParallelController":
|
||||
return (
|
||||
get_parallel().config.dp_size > 1
|
||||
and get_parallel().config.load_balance_method.lower() in _LOAD_AWARE_METHODS
|
||||
get_parallel().dp_size > 1
|
||||
and get_parallel().load_balance_method.lower() in _LOAD_AWARE_METHODS
|
||||
)
|
||||
if get_parallel().config.dp_size > 1 and (
|
||||
get_parallel().config.load_balance_method.lower() in _LOAD_AWARE_METHODS
|
||||
if get_parallel().dp_size > 1 and (
|
||||
get_parallel().load_balance_method.lower() in _LOAD_AWARE_METHODS
|
||||
):
|
||||
return False
|
||||
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
|
||||
socket when zmq mode is active.
|
||||
"""
|
||||
dp_size = get_parallel().config.dp_size
|
||||
dp_size = get_parallel().dp_size
|
||||
if zmq_reader_owner(caller):
|
||||
return ZmqShmLoadSnapshotReader(
|
||||
_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}"
|
||||
)
|
||||
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
|
||||
|
||||
# Mirror scheduler_dp_attn_mixin's NCCL all-gather path: when the
|
||||
|
||||
@@ -470,16 +470,16 @@ class Scheduler(
|
||||
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.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
|
||||
|
||||
# Distributed rank info
|
||||
attn_tp_rank, attn_tp_size, attn_dp_rank, attn_dp_size = (
|
||||
compute_dp_attention_world_info(
|
||||
get_parallel().config.enable_dp_attention,
|
||||
get_parallel().enable_dp_attention,
|
||||
tp_rank,
|
||||
get_parallel().config.tp_size,
|
||||
get_parallel().config.dp_size,
|
||||
get_parallel().dp_size,
|
||||
get_parallel().config.attn_cp_size,
|
||||
)
|
||||
)
|
||||
@@ -489,7 +489,7 @@ class Scheduler(
|
||||
pp_rank=pp_rank,
|
||||
pp_size=get_parallel().config.pp_size,
|
||||
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_size=attn_tp_size,
|
||||
attn_cp_rank=attn_cp_rank,
|
||||
@@ -499,7 +499,7 @@ class Scheduler(
|
||||
attn_dp_rank=attn_dp_rank,
|
||||
attn_dp_size=attn_dp_size,
|
||||
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_size=get_parallel().config.moe_dp_size,
|
||||
gpu_id=gpu_id,
|
||||
@@ -1088,7 +1088,7 @@ class Scheduler(
|
||||
self.min_free_slots_delayer = MinFreeSlotsDelayer(
|
||||
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(
|
||||
"scheduler.pp_max_micro_batch_size_default",
|
||||
pp_max_micro_batch_size=max(
|
||||
@@ -1432,7 +1432,7 @@ class Scheduler(
|
||||
gloo_group=self.attn_tp_cpu_group,
|
||||
tp_rank=self.ps.tp_rank,
|
||||
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,
|
||||
bootstrap_port=get_disagg().disaggregation_bootstrap_port,
|
||||
max_total_num_tokens=self.max_total_num_tokens,
|
||||
@@ -3291,7 +3291,7 @@ class Scheduler(
|
||||
beam_width: Optional[int] = None,
|
||||
running_batch: Optional[ScheduleBatch] = None,
|
||||
) -> 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()
|
||||
|
||||
active_batch = running_batch or self.running_batch
|
||||
@@ -5015,7 +5015,7 @@ class Scheduler(
|
||||
|
||||
old_ep_size = ElasticEPStateManager.get_effective_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(
|
||||
"[Elastic EP][scale] request received: new_ep_size=%d "
|
||||
@@ -5285,7 +5285,7 @@ def configure_scheduler_process(
|
||||
prefix += f" MOE_DP{moe_dp_rank}"
|
||||
if get_parallel().config.tp_size > 1:
|
||||
prefix += f" TP{shown_tp}"
|
||||
if get_parallel().config.ep_size > 1:
|
||||
if get_parallel().ep_size > 1:
|
||||
prefix += f" EP{shown_moe_ep}"
|
||||
|
||||
# Config the process
|
||||
@@ -5301,7 +5301,7 @@ def configure_scheduler_process(
|
||||
set_gpu_proc_affinity(
|
||||
get_parallel().config.pp_size,
|
||||
get_parallel().config.tp_size,
|
||||
get_parallel().config.nnodes,
|
||||
get_parallel().nnodes,
|
||||
gpu_id,
|
||||
)
|
||||
if not envs.SGLANG_NUMA_BIND_V2.get():
|
||||
|
||||
@@ -409,7 +409,7 @@ class SchedulerDPAttnAdapter:
|
||||
return prepare_mlp_sync_batch_raw(
|
||||
local_batch,
|
||||
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_cp_size=self.ps.attn_cp_size,
|
||||
tp_group=self.tp_group,
|
||||
@@ -418,7 +418,7 @@ class SchedulerDPAttnAdapter:
|
||||
require_mlp_tp_gather=require_mlp_tp_gather(),
|
||||
disable_overlap_schedule=get_schedule().disable_overlap_schedule,
|
||||
offload_tags=self.offload_tags,
|
||||
dwdp=get_parallel().config.dwdp_size > 1,
|
||||
dwdp=get_parallel().dwdp_size > 1,
|
||||
)
|
||||
|
||||
def maybe_prepare_mlp_sync_batch(
|
||||
|
||||
@@ -34,7 +34,7 @@ class SchedulerRecvSkipper:
|
||||
return ForwardMode.DECODE
|
||||
|
||||
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._threshold = get_schedule().scheduler_recv_interval
|
||||
# All can be tuned if needed
|
||||
|
||||
@@ -151,7 +151,7 @@ class SchedulerRequestReceiver:
|
||||
return recv_reqs
|
||||
|
||||
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:
|
||||
work_reqs, control_reqs = self._split_work_and_control_reqs(recv_reqs)
|
||||
else:
|
||||
@@ -180,7 +180,7 @@ class SchedulerRequestReceiver:
|
||||
# instead of the full tp_group. This avoids an expensive
|
||||
# all-ranks gloo sync.
|
||||
_local_ctrl = (
|
||||
get_parallel().config.enable_dp_attention_local_control_broadcast
|
||||
get_parallel().enable_dp_attention_local_control_broadcast
|
||||
or is_ep_scale_joiner()
|
||||
)
|
||||
if _local_ctrl:
|
||||
@@ -258,7 +258,7 @@ class SchedulerRequestReceiver:
|
||||
# peer ranks may still be unpickling ShmPointerMMData
|
||||
# (-> shm_open). Synchronize the same CPU groups that carried
|
||||
# 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:
|
||||
barrier(group=self.attn_tp_cpu_group)
|
||||
if self.ps.attn_cp_size > 1:
|
||||
|
||||
@@ -128,7 +128,7 @@ class SchedulerPPMixin:
|
||||
next_pp_outputs = None
|
||||
next_batch_result = 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 = (
|
||||
self._pp_commit_send_output_work_and_preprocess_output_tensors(
|
||||
next_first_rank_mb_id,
|
||||
@@ -144,7 +144,7 @@ class SchedulerPPMixin:
|
||||
self.mb_metadata,
|
||||
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 = (
|
||||
self._pp_commit_send_output_work_and_preprocess_output_tensors(
|
||||
next_first_rank_mb_id,
|
||||
@@ -274,7 +274,7 @@ class SchedulerPPMixin:
|
||||
server_is_idle = False
|
||||
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 = (
|
||||
self._pp_commit_send_output_work_and_preprocess_output_tensors(
|
||||
next_first_rank_mb_id,
|
||||
@@ -292,7 +292,7 @@ class SchedulerPPMixin:
|
||||
self.mb_metadata,
|
||||
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 = (
|
||||
self._pp_commit_send_output_work_and_preprocess_output_tensors(
|
||||
next_first_rank_mb_id,
|
||||
@@ -435,7 +435,7 @@ class SchedulerPPMixin:
|
||||
pp_proxy_tensors = self._pp_recv_proxy_tensors()
|
||||
|
||||
# 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 = (
|
||||
self._pp_commit_send_output_work_and_preprocess_output_tensors(
|
||||
next_first_rank_mb_id,
|
||||
@@ -453,7 +453,7 @@ class SchedulerPPMixin:
|
||||
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 = (
|
||||
self._pp_commit_send_output_work_and_preprocess_output_tensors(
|
||||
next_first_rank_mb_id,
|
||||
@@ -564,12 +564,10 @@ class SchedulerPPMixin:
|
||||
self.on_idle()
|
||||
|
||||
def init_pp_loop_state(self: Scheduler):
|
||||
self.pp_loop_size: int = (
|
||||
self.ps.pp_size + get_parallel().config.pp_async_batch_depth
|
||||
)
|
||||
self.pp_loop_size: int = self.ps.pp_size + get_parallel().pp_async_batch_depth
|
||||
# In CP mode, attention weights are duplicated, eliminating the need for the attention TP all-gather operation.
|
||||
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.last_mbs = [None] * self.pp_loop_size
|
||||
|
||||
@@ -165,7 +165,7 @@ class TokenizerControlMixin:
|
||||
mode = spec[2] if len(spec) > 2 else "queueing"
|
||||
comm = FanOutCommunicator(
|
||||
self._dispatch_to_scheduler,
|
||||
get_parallel().config.dp_size,
|
||||
get_parallel().dp_size,
|
||||
mode,
|
||||
)
|
||||
setattr(self, f"{name}_communicator", comm)
|
||||
@@ -174,8 +174,8 @@ class TokenizerControlMixin:
|
||||
|
||||
def update_control_communicator_fan_out(self: TokenizerManager, worker_count: int):
|
||||
primary_group_control = (
|
||||
get_parallel().config.enable_dp_attention
|
||||
and not get_parallel().config.enable_dp_attention_local_control_broadcast
|
||||
get_parallel().enable_dp_attention
|
||||
and not get_parallel().enable_dp_attention_local_control_broadcast
|
||||
)
|
||||
if primary_group_control:
|
||||
control_fan_out = (
|
||||
@@ -428,8 +428,7 @@ class TokenizerControlMixin:
|
||||
) -> Tuple[bool, str]:
|
||||
self.auto_create_handle_loop()
|
||||
assert (
|
||||
get_parallel().config.dp_size == 1
|
||||
or get_parallel().config.enable_dp_attention
|
||||
get_parallel().dp_size == 1 or get_parallel().enable_dp_attention
|
||||
), "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)
|
||||
@@ -442,8 +441,7 @@ class TokenizerControlMixin:
|
||||
) -> Tuple[bool, str]:
|
||||
self.auto_create_handle_loop()
|
||||
assert (
|
||||
get_parallel().config.dp_size == 1
|
||||
or get_parallel().config.enable_dp_attention
|
||||
get_parallel().dp_size == 1 or get_parallel().enable_dp_attention
|
||||
), "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)
|
||||
@@ -456,8 +454,7 @@ class TokenizerControlMixin:
|
||||
) -> Tuple[bool, str]:
|
||||
self.auto_create_handle_loop()
|
||||
assert (
|
||||
get_parallel().config.dp_size == 1
|
||||
or get_parallel().config.enable_dp_attention
|
||||
get_parallel().dp_size == 1 or get_parallel().enable_dp_attention
|
||||
), "dp_size must be 1 or dp attention must be enabled for update weights from distributed"
|
||||
|
||||
if obj.abort_all_requests:
|
||||
@@ -490,7 +487,7 @@ class TokenizerControlMixin:
|
||||
self.auto_create_handle_loop()
|
||||
# TODO: support DP
|
||||
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"
|
||||
result = (
|
||||
await self.init_weights_send_group_for_remote_instance_communicator(obj)
|
||||
@@ -505,7 +502,7 @@ class TokenizerControlMixin:
|
||||
self.auto_create_handle_loop()
|
||||
# TODO: support DP
|
||||
assert (
|
||||
get_parallel().config.dp_size == 1
|
||||
get_parallel().dp_size == 1
|
||||
), "dp_size must be 1 for send_weights_to_remote_instance"
|
||||
result = (await self.send_weights_to_remote_instance_communicator(obj))[0]
|
||||
return result.success, result.message
|
||||
@@ -517,8 +514,7 @@ class TokenizerControlMixin:
|
||||
) -> Tuple[bool, str]:
|
||||
self.auto_create_handle_loop()
|
||||
assert (
|
||||
get_parallel().config.dp_size == 1
|
||||
or get_parallel().config.enable_dp_attention
|
||||
get_parallel().dp_size == 1 or get_parallel().enable_dp_attention
|
||||
), "dp_size must be 1 or dp attention must be enabled for update weights from tensor"
|
||||
|
||||
if obj.abort_all_requests:
|
||||
@@ -556,8 +552,7 @@ class TokenizerControlMixin:
|
||||
try:
|
||||
# For now, we only support single data parallel instance
|
||||
assert (
|
||||
get_parallel().config.dp_size == 1
|
||||
or get_parallel().config.enable_dp_attention
|
||||
get_parallel().dp_size == 1 or get_parallel().enable_dp_attention
|
||||
), "dp_size must be 1 or dp attention must be enabled for update weights from IPC"
|
||||
logger.info("Starting IPC weight update")
|
||||
|
||||
@@ -620,8 +615,7 @@ class TokenizerControlMixin:
|
||||
)
|
||||
|
||||
assert (
|
||||
get_parallel().config.dp_size == 1
|
||||
or get_parallel().config.enable_dp_attention
|
||||
get_parallel().dp_size == 1 or get_parallel().enable_dp_attention
|
||||
), "dp_size must be 1 or dp attention must be enabled for dynamic lora loading"
|
||||
logger.info(
|
||||
"Start load Lora adapter. Lora name=%s, path=%s",
|
||||
@@ -699,8 +693,7 @@ class TokenizerControlMixin:
|
||||
)
|
||||
|
||||
assert (
|
||||
get_parallel().config.dp_size == 1
|
||||
or get_parallel().config.enable_dp_attention
|
||||
get_parallel().dp_size == 1 or get_parallel().enable_dp_attention
|
||||
), "dp_size must be 1 or dp attention must be enabled for dynamic lora loading"
|
||||
logger.info(
|
||||
"Start load Lora adapter from tensors. Lora name=%s",
|
||||
@@ -780,8 +773,7 @@ class TokenizerControlMixin:
|
||||
), "lora_name must be provided to unload LoRA adapter"
|
||||
|
||||
assert (
|
||||
get_parallel().config.dp_size == 1
|
||||
or get_parallel().config.enable_dp_attention
|
||||
get_parallel().dp_size == 1 or get_parallel().enable_dp_attention
|
||||
), "dp_size must be 1 or dp attention must be enabled for dynamic lora loading"
|
||||
logger.info(
|
||||
"Start unload Lora adapter. Lora name=%s",
|
||||
@@ -801,7 +793,7 @@ class TokenizerControlMixin:
|
||||
self.auto_create_handle_loop()
|
||||
results = await self.get_weights_by_name_communicator(obj)
|
||||
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]
|
||||
else:
|
||||
return all_parameters
|
||||
|
||||
@@ -415,7 +415,7 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
|
||||
self.server_args = server_args
|
||||
assert_published(server_args, role="tokenizer")
|
||||
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_scale_phase = "idle"
|
||||
self.elastic_last_error = None
|
||||
@@ -1550,7 +1550,7 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
|
||||
return batch_size > 0 and (
|
||||
get_serving().enable_tokenizer_batch_encode
|
||||
or (
|
||||
(not get_parallel().config.enable_dp_attention)
|
||||
(not get_parallel().enable_dp_attention)
|
||||
and (not self._batch_has_text(batch_size, requests))
|
||||
)
|
||||
)
|
||||
|
||||
@@ -279,9 +279,7 @@ def build_kv_cache(
|
||||
),
|
||||
is_eagle=spec_algorithm.is_eagle(),
|
||||
tp_cache_group=(
|
||||
attn_tp_cpu_group
|
||||
if get_parallel().config.enable_dp_attention
|
||||
else tp_cpu_group
|
||||
attn_tp_cpu_group if get_parallel().enable_dp_attention else tp_cpu_group
|
||||
),
|
||||
attn_cp_cache_group=attn_cp_cpu_group,
|
||||
attn_tp_cache_group=attn_tp_cpu_group,
|
||||
|
||||
@@ -45,7 +45,7 @@ def ranks_per_host() -> int:
|
||||
return 1
|
||||
if world_group.world_size == 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:
|
||||
|
||||
@@ -610,7 +610,7 @@ class CPUGraphRunner:
|
||||
model_runner.server_args.enable_profile_cuda_graph
|
||||
)
|
||||
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.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()):
|
||||
return
|
||||
|
||||
join_effective_ep_size = (
|
||||
get_parallel().config.ep_join_rank_offset + self.ps.tp_size
|
||||
)
|
||||
join_effective_ep_size = get_parallel().ep_join_rank_offset + self.ps.tp_size
|
||||
dist.barrier(group=self.tp_group.cpu_group)
|
||||
if self.ps.tp_rank == 0:
|
||||
register_scale_cohort(
|
||||
get_parallel().config.ep_join_rank_offset,
|
||||
get_parallel().ep_join_rank_offset,
|
||||
join_effective_ep_size,
|
||||
)
|
||||
join_scale_process_group()
|
||||
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(
|
||||
model_config=self.model_config,
|
||||
moe_ep_rank=global_ep_rank,
|
||||
@@ -697,7 +695,7 @@ class ModelRunner:
|
||||
if self.is_draft_worker:
|
||||
return
|
||||
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(
|
||||
compute_initial_expert_location_metadata(
|
||||
@@ -928,7 +926,7 @@ class ModelRunner:
|
||||
device=self.device,
|
||||
tp_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
|
||||
),
|
||||
host_to_device_ratio=hisparse_cfg.host_to_device_ratio,
|
||||
@@ -964,7 +962,7 @@ class ModelRunner:
|
||||
def post_capture_elastic_ep_recover(self):
|
||||
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(
|
||||
model_config=self.model_config,
|
||||
moe_ep_rank=global_ep_rank,
|
||||
@@ -1004,7 +1002,7 @@ class ModelRunner:
|
||||
self.decode_attn_backend = backends.decode_attn_backend
|
||||
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()
|
||||
|
||||
def _prepare_replicated_q_proj(self) -> None:
|
||||
@@ -1282,7 +1280,7 @@ class ModelRunner:
|
||||
def maybe_init_dwdp(self):
|
||||
if self.is_draft_worker:
|
||||
return
|
||||
if get_parallel().config.dwdp_size <= 1:
|
||||
if get_parallel().dwdp_size <= 1:
|
||||
return
|
||||
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
|
||||
# multiplied into the process count here (unlike regular DP, where
|
||||
# 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(
|
||||
numa_index=self.gpu_id,
|
||||
world_size=dp_size * self.ps.tp_size * self.ps.pp_size,
|
||||
@@ -1938,7 +1936,7 @@ class ModelRunner:
|
||||
if added <= 0:
|
||||
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
|
||||
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)
|
||||
|
||||
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:
|
||||
if self.eplb_manager is None:
|
||||
|
||||
+2
-4
@@ -78,13 +78,11 @@ class RemoteInstanceWeightTransporter:
|
||||
"""
|
||||
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).
|
||||
# Derive host from dist_init_addr (shared across all nodes).
|
||||
bootstrap_host = (
|
||||
NetworkAddress.parse(get_parallel().config.dist_init_addr)
|
||||
.resolved()
|
||||
.host
|
||||
NetworkAddress.parse(get_parallel().dist_init_addr).resolved().host
|
||||
)
|
||||
else:
|
||||
bootstrap_host = "127.0.0.1"
|
||||
|
||||
@@ -123,8 +123,8 @@ class StartupWeightLoadOptions:
|
||||
attn_cp_size=get_parallel().config.attn_cp_size,
|
||||
dcp_size=get_parallel().config.dcp_size,
|
||||
pp_size=get_parallel().config.pp_size,
|
||||
dp_size=get_parallel().config.dp_size,
|
||||
ep_size=get_parallel().config.ep_size,
|
||||
dp_size=get_parallel().dp_size,
|
||||
ep_size=get_parallel().ep_size,
|
||||
cpu_offload_gb=get_exec().offload.cpu_offload_gb,
|
||||
offload_group_size=get_exec().offload.offload_group_size,
|
||||
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.tp_size = get_parallel().config.tp_size
|
||||
# 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.enable_pdmux = model_runner.server_args.enable_pdmux
|
||||
self.return_hidden_states_mode = (
|
||||
@@ -289,7 +289,7 @@ class BaseRunner(ABC):
|
||||
"""
|
||||
if (
|
||||
not get_parallel().dcp_enabled
|
||||
or get_parallel().config.dcp_comm_backend != "fi_a2a"
|
||||
or get_parallel().dcp_comm_backend != "fi_a2a"
|
||||
):
|
||||
return
|
||||
|
||||
@@ -348,7 +348,7 @@ class BaseRunner(ABC):
|
||||
hidden_size=mr.model_config.hidden_size,
|
||||
vocab_size=mr.model_config.vocab_size,
|
||||
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,
|
||||
is_encoder_decoder=mr.model_config.is_encoder_decoder,
|
||||
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_
|
||||
|
||||
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_:
|
||||
global_num_tokens_cpu = [num_tokens]
|
||||
else:
|
||||
|
||||
@@ -242,7 +242,7 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
|
||||
self.require_mlp_tp_gather or self.require_attn_tp_gather
|
||||
)
|
||||
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 = (
|
||||
model_runner.server_args.enable_two_batch_overlap
|
||||
|
||||
@@ -144,7 +144,7 @@ class EagerRunner(BaseRunner):
|
||||
encoder_lens_dtype=(
|
||||
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).
|
||||
self.warmup()
|
||||
|
||||
@@ -352,7 +352,7 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
|
||||
self.moe_fusions = self.model_runner.moe_fusions
|
||||
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_attn_tp_gather = require_attn_tp_gather()
|
||||
|
||||
|
||||
@@ -1893,9 +1893,9 @@ class PreshardedModelLoader(DefaultModelLoader):
|
||||
"dp": _safe(lambda: parallel.moe_dp_size),
|
||||
"ep": _safe(lambda: parallel.moe_ep_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,
|
||||
"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,
|
||||
"quantization": model_config.quantization,
|
||||
"model_dtype": str(model_config.dtype),
|
||||
|
||||
@@ -442,7 +442,7 @@ class ApertusForCausalLM(nn.Module):
|
||||
config.hidden_size,
|
||||
quant_config=quant_config,
|
||||
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.pooler = Pooler(pooling_type=PoolingType.LAST, normalize=True)
|
||||
|
||||
@@ -405,7 +405,7 @@ class ArceeForCausalLM(nn.Module):
|
||||
config.hidden_size,
|
||||
quant_config=quant_config,
|
||||
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.pooler = Pooler(pooling_type=PoolingType.LAST, normalize=True)
|
||||
|
||||
@@ -817,7 +817,7 @@ class BailingMoEForCausalLM(nn.Module):
|
||||
config.hidden_size,
|
||||
quant_config=quant_config,
|
||||
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)
|
||||
|
||||
|
||||
@@ -1084,7 +1084,7 @@ class BailingMoELinearForCausalLM(nn.Module):
|
||||
config.hidden_size,
|
||||
params_dtype=torch.float32,
|
||||
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)
|
||||
|
||||
@@ -241,7 +241,7 @@ class BailingMoeForCausalLMNextN(nn.Module):
|
||||
config.hidden_size,
|
||||
quant_config=quant_config,
|
||||
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)
|
||||
if is_bailing_moe_v3:
|
||||
|
||||
@@ -1333,7 +1333,7 @@ class BailingMoeV3ForCausalLM(nn.Module):
|
||||
# in the logits processor. Accuracy-neutral on ling-v3.
|
||||
params_dtype=torch.bfloat16,
|
||||
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)
|
||||
|
||||
@@ -299,7 +299,7 @@ class DeepseekMLAForwardMixin:
|
||||
# --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).
|
||||
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 not self.use_deep_gemm_bmm
|
||||
and self.w_kc_qrep is not None
|
||||
@@ -779,7 +779,7 @@ class DeepseekMLAForwardMixin:
|
||||
attn_output, self.num_local_heads
|
||||
)
|
||||
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(
|
||||
self.current_attention_backend
|
||||
)
|
||||
|
||||
+2
-2
@@ -331,7 +331,7 @@ class DeepseekMLARocmForwardMixin:
|
||||
from sglang.srt.model_executor.runner import get_is_capture_mode
|
||||
|
||||
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 not self.use_deep_gemm_bmm
|
||||
and self.w_kc_qrep is not None
|
||||
@@ -778,7 +778,7 @@ class DeepseekMLARocmForwardMixin:
|
||||
attn_output, self.num_local_heads
|
||||
)
|
||||
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(
|
||||
self.current_attention_backend
|
||||
)
|
||||
|
||||
@@ -366,7 +366,7 @@ class DeepseekV3ForCausalLMNextN(DeepseekV3ForCausalLM):
|
||||
config.hidden_size,
|
||||
quant_config=quant_config,
|
||||
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)
|
||||
|
||||
|
||||
@@ -3007,7 +3007,7 @@ class DeepseekV2ForCausalLM(nn.Module, DeepseekV2WeightLoaderMixin):
|
||||
config.hidden_size,
|
||||
quant_config=quant_config,
|
||||
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:
|
||||
# ranks other than the last rank will have a placeholder layer
|
||||
|
||||
@@ -3237,7 +3237,7 @@ class DeepseekV4ForCausalLM(nn.Module):
|
||||
config.hidden_size,
|
||||
quant_config=quant_config,
|
||||
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:
|
||||
self.lm_head = PPMissingLayer()
|
||||
|
||||
@@ -758,7 +758,7 @@ class DeepseekV4ForCausalLMDSpark(nn.Module):
|
||||
config.vocab_size,
|
||||
config.hidden_size,
|
||||
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:
|
||||
self.embed_tokens: Optional[nn.Module] = None
|
||||
|
||||
@@ -250,7 +250,7 @@ class DeepseekV4ForCausalLMNextN(DeepseekV4ForCausalLM):
|
||||
config.hidden_size,
|
||||
quant_config=quant_config,
|
||||
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)
|
||||
|
||||
|
||||
@@ -1876,7 +1876,7 @@ class Dots3LanguageModelForCausalLM(nn.Module):
|
||||
config.hidden_size,
|
||||
quant_config=quant_config,
|
||||
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)
|
||||
|
||||
|
||||
@@ -167,7 +167,7 @@ class Dots3NoteForCausalLMNextN(Dots3LanguageModelForCausalLM):
|
||||
config.hidden_size,
|
||||
quant_config=quant_config,
|
||||
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._mtp_loaded_embed = False
|
||||
|
||||
@@ -439,7 +439,7 @@ class Exaone4ForCausalLM(nn.Module):
|
||||
config.hidden_size,
|
||||
quant_config=quant_config,
|
||||
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)
|
||||
|
||||
@@ -643,7 +643,7 @@ class ExaoneMoEForCausalLM(nn.Module):
|
||||
config.hidden_size,
|
||||
quant_config=quant_config,
|
||||
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)
|
||||
# For EAGLE3 support
|
||||
|
||||
@@ -63,7 +63,7 @@ class ExaoneMoEForCausalLMMTP(ExaoneMoEForCausalLM):
|
||||
config.hidden_size,
|
||||
quant_config=quant_config,
|
||||
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)
|
||||
|
||||
|
||||
@@ -472,7 +472,7 @@ class FalconH1ForCausalLM(nn.Module):
|
||||
quant_config=quant_config,
|
||||
org_num_embeddings=config.vocab_size,
|
||||
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_multiplier = config.lm_head_multiplier
|
||||
|
||||
@@ -1163,7 +1163,7 @@ class Glm4MoeForCausalLM(nn.Module):
|
||||
config.hidden_size,
|
||||
quant_config=quant_config,
|
||||
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)
|
||||
|
||||
|
||||
@@ -905,7 +905,7 @@ class Glm4MoeLiteForCausalLM(nn.Module, DeepseekV2WeightLoaderMixin):
|
||||
config.hidden_size,
|
||||
quant_config=quant_config,
|
||||
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)
|
||||
|
||||
|
||||
@@ -159,7 +159,7 @@ class Glm4MoeLiteForCausalLMNextN(Glm4MoeLiteForCausalLM):
|
||||
config.hidden_size,
|
||||
quant_config=quant_config,
|
||||
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)
|
||||
|
||||
|
||||
@@ -144,7 +144,7 @@ class Glm4MoeForCausalLMNextN(Glm4MoeForCausalLM):
|
||||
config.hidden_size,
|
||||
quant_config=quant_config,
|
||||
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)
|
||||
|
||||
|
||||
@@ -70,7 +70,7 @@ class Glm4vMoeForConditionalGeneration(Glm4vForConditionalGeneration):
|
||||
config.hidden_size,
|
||||
quant_config=quant_config,
|
||||
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:
|
||||
# ranks other than the last rank will have a placeholder layer
|
||||
|
||||
@@ -135,7 +135,7 @@ class GlmOcrForConditionalGenerationNextN(GlmOcrForConditionalGeneration):
|
||||
config.hidden_size,
|
||||
quant_config=quant_config,
|
||||
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)
|
||||
|
||||
|
||||
@@ -258,7 +258,7 @@ class GptOssSparseMoeBlock(nn.Module):
|
||||
hidden_states: torch.Tensor,
|
||||
forward_batch: Optional[ForwardBatch] = None,
|
||||
) -> torch.Tensor:
|
||||
if get_parallel().config.dwdp_size > 1:
|
||||
if get_parallel().dwdp_size > 1:
|
||||
return self.forward_dwdp(hidden_states)
|
||||
|
||||
if not get_moe_a2a_backend().is_deepep():
|
||||
@@ -786,7 +786,7 @@ class GptOssForCausalLM(nn.Module):
|
||||
config.hidden_size,
|
||||
# quant_config=quant_config,
|
||||
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.capture_aux_hidden_states = False
|
||||
|
||||
@@ -296,7 +296,7 @@ class KimiK3MLP(nn.Module):
|
||||
# but allow the NPU launcher to retain the proven attention-TP layout
|
||||
# without a device-type branch in shared model code.
|
||||
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 tp_rank 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
|
||||
# hold different tokens.
|
||||
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
|
||||
# uses the original TP-sharded shared MLP. Gather only that branch's
|
||||
# inputs, then reduce-scatter its output back to the DP-local rows.
|
||||
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._dp_attention
|
||||
and get_parallel().attn_tp_size > 1
|
||||
@@ -2870,7 +2870,7 @@ class KimiK3LinearForCausalLM(nn.Module):
|
||||
config.hidden_size,
|
||||
quant_config=quant_config,
|
||||
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:
|
||||
self.lm_head = PPMissingLayer()
|
||||
|
||||
@@ -667,7 +667,7 @@ class LagunaForCausalLM(nn.Module):
|
||||
config.hidden_size,
|
||||
quant_config=quant_config,
|
||||
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:
|
||||
self.lm_head = PPMissingLayer()
|
||||
|
||||
@@ -824,7 +824,7 @@ class LLaDA2MoeModelLM(nn.Module):
|
||||
config.hidden_size,
|
||||
quant_config=quant_config,
|
||||
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)
|
||||
|
||||
|
||||
@@ -536,7 +536,7 @@ class LlamaForCausalLM(nn.Module):
|
||||
config.hidden_size,
|
||||
quant_config=quant_config,
|
||||
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.pooler = Pooler(pooling_type=PoolingType.LAST, normalize=True)
|
||||
|
||||
@@ -721,7 +721,7 @@ class LongcatFlashForCausalLM(nn.Module):
|
||||
config.hidden_size,
|
||||
quant_config=quant_config,
|
||||
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.capture_aux_hidden_states = False
|
||||
|
||||
@@ -520,7 +520,7 @@ class MellumForCausalLM(Qwen3MoeForCausalLM):
|
||||
cfg.hidden_size,
|
||||
quant_config=quant_config,
|
||||
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.capture_aux_hidden_states = False
|
||||
|
||||
@@ -1187,7 +1187,7 @@ class MiMoV2ForCausalLM(nn.Module, AudioEncoderMixin):
|
||||
config.hidden_size,
|
||||
quant_config=quant_config,
|
||||
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:
|
||||
self.lm_head = PPMissingLayer()
|
||||
|
||||
@@ -259,7 +259,7 @@ class MiMoV2MTP(MiMoV2ForCausalLM):
|
||||
config.hidden_size,
|
||||
quant_config=quant_config,
|
||||
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)
|
||||
|
||||
|
||||
@@ -1574,7 +1574,7 @@ class MiniMaxM3SparseForCausalLM(nn.Module):
|
||||
config.hidden_size,
|
||||
quant_config=quant_config,
|
||||
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)
|
||||
|
||||
@@ -123,7 +123,7 @@ class MiniMaxM3SparseForConditionalGeneration(nn.Module):
|
||||
text_config.hidden_size,
|
||||
quant_config=quant_config,
|
||||
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:
|
||||
self.lm_head = PPMissingLayer()
|
||||
|
||||
@@ -964,7 +964,7 @@ class NemotronHForCausalLM(nn.Module):
|
||||
else lora_config.lora_vocab_padding_size
|
||||
),
|
||||
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),
|
||||
)
|
||||
else:
|
||||
|
||||
@@ -339,7 +339,7 @@ class NemotronHForCausalLMMTP(NemotronHForCausalLM):
|
||||
self.config.hidden_size,
|
||||
quant_config=quant_config,
|
||||
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)
|
||||
|
||||
@@ -1119,7 +1119,7 @@ class Qwen2MoeForCausalLM(nn.Module):
|
||||
config.hidden_size,
|
||||
quant_config=quant_config,
|
||||
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)
|
||||
# For EAGLE3 support
|
||||
|
||||
@@ -492,7 +492,7 @@ class Qwen3ForCausalLM(nn.Module):
|
||||
config.vocab_size,
|
||||
config.hidden_size,
|
||||
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),
|
||||
)
|
||||
else:
|
||||
|
||||
@@ -80,7 +80,7 @@ class Qwen3_5ForCausalLM(nn.Module):
|
||||
quant_config=quant_config,
|
||||
org_num_embeddings=config.vocab_size,
|
||||
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:
|
||||
self.lm_head = PPMissingLayer()
|
||||
|
||||
@@ -961,7 +961,7 @@ class Qwen3MoeForCausalLM(nn.Module):
|
||||
config.hidden_size,
|
||||
quant_config=quant_config,
|
||||
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.capture_aux_hidden_states = False
|
||||
|
||||
@@ -63,7 +63,7 @@ class Qwen3MoeForCausalLMMTP(Qwen3MoeForCausalLM):
|
||||
config.hidden_size,
|
||||
quant_config=quant_config,
|
||||
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)
|
||||
|
||||
|
||||
@@ -1027,7 +1027,7 @@ class Qwen3NextForCausalLM(nn.Module):
|
||||
quant_config=quant_config,
|
||||
org_num_embeddings=config.vocab_size,
|
||||
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)
|
||||
# For EAGLE3 support
|
||||
|
||||
@@ -80,7 +80,7 @@ class Qwen3NextForCausalLMMTP(Qwen3NextForCausalLM):
|
||||
config.hidden_size,
|
||||
quant_config=quant_config,
|
||||
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)
|
||||
# Mirror Qwen3NextForCausalLM.__init__'s shared-expert fusion setup so
|
||||
|
||||
@@ -1316,7 +1316,7 @@ class Qwen3VLForConditionalGeneration(nn.Module):
|
||||
self.config.vocab_size,
|
||||
self.config.hidden_size,
|
||||
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),
|
||||
)
|
||||
else:
|
||||
|
||||
@@ -1220,7 +1220,7 @@ class SarvamMLAForCausalLM(nn.Module):
|
||||
config.hidden_size,
|
||||
quant_config=quant_config,
|
||||
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)
|
||||
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user