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(
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,
+1 -1
View File
@@ -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
+12 -16
View File
@@ -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;
+1 -1
View File
@@ -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 "
+2 -2
View File
@@ -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(
+6 -7
View File
@@ -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()
+4 -8
View File
@@ -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
+6 -10
View File
@@ -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:
+2 -2
View File
@@ -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,
}
+2 -2
View File
@@ -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:
+6 -9
View File
@@ -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
+4 -4
View File
@@ -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
+1 -1
View File
@@ -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)
)
+4 -7
View File
@@ -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():
+2 -4
View File
@@ -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(
+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():
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
+3 -3
View File
@@ -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):
+1 -1
View File
@@ -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
)
)
+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
``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:
+7 -7
View File
@@ -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
+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_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:
@@ -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()
+2 -2
View File
@@ -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),
+1 -1
View File
@@ -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)
+1 -1
View File
@@ -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)
+1 -1
View File
@@ -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:
+1 -1
View File
@@ -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
)
@@ -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
)
+1 -1
View File
@@ -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)
+1 -1
View File
@@ -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
+1 -1
View File
@@ -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
+1 -1
View File
@@ -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)
+1 -1
View File
@@ -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
+1 -1
View File
@@ -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)
+1 -1
View File
@@ -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
+1 -1
View File
@@ -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)
+1 -1
View File
@@ -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)
+1 -1
View File
@@ -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)
+1 -1
View File
@@ -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
+1 -1
View File
@@ -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)
+2 -2
View File
@@ -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
+4 -4
View File
@@ -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()
+1 -1
View File
@@ -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()
+1 -1
View File
@@ -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)
+1 -1
View File
@@ -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)
+1 -1
View File
@@ -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
+1 -1
View File
@@ -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
+1 -1
View File
@@ -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()
+1 -1
View File
@@ -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)
+1 -1
View File
@@ -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)
+1 -1
View File
@@ -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()
+1 -1
View File
@@ -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:
+1 -1
View File
@@ -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)
+1 -1
View File
@@ -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
+1 -1
View File
@@ -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:
+1 -1
View File
@@ -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()
+1 -1
View File
@@ -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
+1 -1
View File
@@ -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)
+1 -1
View File
@@ -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
+1 -1
View File
@@ -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
+1 -1
View File
@@ -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:
+1 -1
View File
@@ -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