# SPDX-License-Identifier: Apache-2.0 """Server-argument resolution for context- and decode-context parallelism.""" from __future__ import annotations import logging import os from typing import Any from sglang.srt.arg_groups.overrides import ( declare_resolution, resolved_view, resolving_view, should_report_expert_balancedness, ) from sglang.srt.connector import ConnectorType from sglang.srt.environ import envs from sglang.srt.model_executor.cuda_graph_config import Backend, Phase, with_phase from sglang.srt.utils.common import is_cuda, parse_connector_type logger = logging.getLogger(__name__) def handle_context_parallelism(server_args: Any): from sglang.srt.arg_groups.overrides import model_config_of cfg = resolving_view(server_args) if parse_connector_type(cfg.model_path) != ConnectorType.INSTANCE: from sglang.srt.configs.model_config import is_deepseek_dsa from sglang.srt.layers.cp.utils import CP_V2_DEFAULT_MODEL_CLASSES model_config = model_config_of(server_args) hf_config = model_config.hf_config model_arch = hf_config.architectures[0] if model_arch in CP_V2_DEFAULT_MODEL_CLASSES: is_dsa_default_model = is_deepseek_dsa(hf_config) # DSA CP-v2 currently supports only the interleave strategy. enable_default_cp_v2 = not is_dsa_default_model or ( cfg.enable_prefill_cp and cfg.cp_strategy == "interleave" ) if enable_default_cp_v2 and not envs.SGLANG_ENABLE_CP_V2.is_set(): envs.SGLANG_ENABLE_CP_V2.set(True) if ( cfg.enable_prefill_cp and model_arch in ("MiMoV2ForCausalLM", "MiMoV2FlashForCausalLM") and envs.SGLANG_ENABLE_CP_V2.get() ): if cfg.cp_strategy != "zigzag": raise ValueError("MiMo V2 CP-v2 only supports --cp-strategy zigzag.") if ( model_config.is_multimodal and not cfg.language_only and not cfg.language_model_only ): raise ValueError( "MiMo V2 CP-v2 only supports text inference; add " "--language-only." ) if cfg.enable_prefill_cp and cfg.cp_strategy is None: raise ValueError( "--cp-strategy must be set when --enable-prefill-cp is enabled." ) if cfg.enable_prefill_context_parallel and cfg.enable_dsa_prefill_context_parallel: raise ValueError( "--enable-prefill-context-parallel and " "--enable-nsa-prefill-context-parallel are mutually " "exclusive. Use --enable-nsa-prefill-context-parallel for " "DeepSeek V3.2 (NSA) models and " "--enable-prefill-context-parallel for MLA-based models " "(DeepSeek V3/R1, Kimi K2.5) or MHA/GQA-based models." ) view = resolved_view(server_args) if view.attn_cp_size > 1: # The tp_size is the world size, not the real tensor parallel size assert ( cfg.tp_size % view.attn_cp_size == 0 ), "tp_size must be divisible by attn_cp_size" assert ( cfg.tp_size % (cfg.dp_size * view.attn_cp_size) == 0 ), "tp_size must be divisible by dp_size * attn_cp_size" assert ( not cfg.enable_aiter_allreduce_fusion ), "Aiter allreduce fusion is not supported with context parallelism" if cfg.moe_dp_size > 1: # The tp_size is the world size, not the real tensor parallel size assert ( cfg.tp_size % cfg.moe_dp_size == 0 ), "tp_size must be divisible by moe_dp_size" assert ( view.ep_size * cfg.moe_dp_size <= cfg.tp_size ), "ep_size * moe_dp_size must be less than or equal to tp_size" assert cfg.pp_size == 1, "PP is not supported with context parallelism" if view.ep_size > 1: assert ( view.ep_size * cfg.moe_dp_size == cfg.tp_size ), "ep_size * moe_dp_size must be equal to tp_size" assert ( not cfg.enable_aiter_allreduce_fusion ), "Aiter allreduce fusion is not supported with context parallelism" if view.attn_cp_size != cfg.moe_dp_size: assert ( cfg.moe_dp_size == 1 ), "attn_cp_size != moe_dp_size is only supported when moe_dp_size == 1" from sglang.srt.layers.cp.base import init_cp_strategy init_cp_strategy( enable_prefill_cp=bool(cfg.enable_prefill_cp), cp_size=cfg.attn_cp_size, cp_strategy=cfg.cp_strategy, ) def handle_dcp_validation(server_args: Any): cfg = resolving_view(server_args) if cfg.dcp_size < 1: raise ValueError( "Decode context parallel size (--dcp-size / " "--decode-context-parallel-size) must be >= 1, but got " f"dcp_size={cfg.dcp_size}." ) if cfg.dcp_comm_backend in ("a2a", "fi_a2a") and cfg.dcp_size <= 1: raise ValueError( f"--dcp-comm-backend {cfg.dcp_comm_backend} only affects the " "decode context-parallel attention reduction and therefore " "requires --dcp-size / --decode-context-parallel-size > 1, but " f"got dcp_size={cfg.dcp_size}." ) if cfg.dcp_comm_backend == "fi_a2a" and not is_cuda(): raise ValueError( "--dcp-comm-backend fi_a2a delegates the exchange to FlashInfer's " "MNNVL All-to-All kernel, which requires an NVIDIA CUDA platform " "with SM90+ and MNNVL fabric memory (e.g. GB200 NVL72). The " "authoritative fabric probe runs at model-runner init; use 'a2a' " "or 'ag_rs' on clusters without MNNVL." ) if cfg.dcp_replicate_q_proj: if cfg.dcp_size <= 1: raise ValueError("--dcp-replicate-q-proj requires --dcp-size > 1.") if cfg.dcp_comm_backend not in ("a2a", "fi_a2a"): raise ValueError( "--dcp-replicate-q-proj only applies to the a2a/fi_a2a DCP " "communication backend (it removes the head-dim Q all-gather); " f"got --dcp-comm-backend={cfg.dcp_comm_backend}." ) def handle_data_parallelism(server_args: Any): # The dp_size==1 resets moved to the resolution pipeline # (arg_groups/overrides.py: _data_parallelism_defaults). from sglang.srt.arg_groups.cuda_graph_hook import ( generate_prefill_cuda_graph_batch_sizes, ) cfg = resolving_view(server_args) from sglang.srt.arg_groups.overrides import ( _data_parallelism_defaults, run_post_process_pass, ) run_post_process_pass(server_args, _data_parallelism_defaults) if cfg.mm_enable_dp_encoder: if cfg.tp_size == 1: logger.warning( "--mm-enable-dp-encoder is enabled with TP=1, so the encoder " "has no data-parallel work to distribute. Disable it unless " "you need to validate this configuration." ) else: logger.info( "--mm-enable-dp-encoder is enabled across TP=%d. It replicates " "the vision encoder and distributes image work across ranks; " "this is most useful when high-resolution or multi-image ViT " "prefill is a material part of TTFT. Measure against the default " "for small-image workloads because replication and aggregation " "can increase memory use and overhead.", cfg.tp_size, ) if resolved_view(server_args).enable_dp_attention: declare_resolution( server_args, "_handle_data_parallelism", schedule_conservativeness=cfg.schedule_conservativeness * 0.3, ) assert cfg.tp_size % cfg.dp_size == 0 original_chunked_prefill_size = cfg.chunked_prefill_size declare_resolution( server_args, "_handle_data_parallelism", chunked_prefill_size=cfg.chunked_prefill_size // cfg.dp_size, ) logger.warning( f"DP attention is enabled. chunked prefill size is adjusted " f"from {original_chunked_prefill_size} to {cfg.chunked_prefill_size}." ) # The prefill CUDA graph max_bs was derived from the pre-DP-division # chunked_prefill_size in _handle_gpu_memory_settings (which runs # before this handler). Re-clamp it (and the captured shape list) to # the per-DP-rank chunked_prefill_size so breakable CUDA graph # capture never exceeds the MoE all-to-all's max_num_tokens budget, # which is also sized from the DP-adjusted chunked_prefill_size. prefill_cfg = cfg.cuda_graph_config.prefill if ( prefill_cfg.backend != Backend.DISABLED and prefill_cfg.max_bs is not None and prefill_cfg.max_bs > cfg.chunked_prefill_size and (Phase.PREFILL, "max_bs") not in server_args._cuda_graph_config_locked ): clamped = {"max_bs": cfg.chunked_prefill_size} if (Phase.PREFILL, "bs") not in server_args._cuda_graph_config_locked: clamped["bs"] = generate_prefill_cuda_graph_batch_sizes( clamped["max_bs"] ) declare_resolution( server_args, "_handle_data_parallelism", cuda_graph_config=with_phase( cfg.cuda_graph_config, Phase.PREFILL, **clamped ), ) # Resolve the phase-aware TP LM-head default before validating the # resulting DP/TP LM-head configuration. from sglang.srt.arg_groups.overrides import ( _dp_lm_head_validation, _tp_lm_head_all_to_all_default, ) run_post_process_pass(server_args, _tp_lm_head_all_to_all_default) run_post_process_pass(server_args, _dp_lm_head_validation) def handle_dwdp(server_args: Any): cfg = resolving_view(server_args) if cfg.dwdp_size <= 1: return assert ( cfg.dwdp_size >= 2 ), f"dwdp_size must be >= 2 when enabled, got {cfg.dwdp_size}" assert ( cfg.dwdp_size == cfg.tp_size ), f"dwdp_size ({cfg.dwdp_size}) must equal tp_size ({cfg.tp_size})" assert cfg.disaggregation_mode in ( "null", "prefill", ), "DWDP requires --disaggregation-mode null or prefill" assert ( not cfg.enable_eplb ), "EPLB dynamic migration conflicts with static DWDP partitioning" assert ( cfg.speculative_algorithm is None ), "DWDP does not support speculative decoding (MTP/draft workers)" assert cfg.pp_size == 1, "DWDP requires pp_size == 1" assert ( not cfg.enable_two_batch_overlap ), "DWDP's prefetch event protocol does not support two-batch overlap" if cfg.disaggregation_mode == "null": logger.warning( "DWDP with --disaggregation-mode null: decode steps re-fetch all " "remote expert weights every step, which is slow. DWDP is " "recommended only with --disaggregation-mode prefill." ) declare_resolution( server_args, "_handle_dwdp", dp_size=cfg.dwdp_size, ) declare_resolution( server_args, "_handle_dwdp", enable_dp_attention=True, ) declare_resolution( server_args, "_handle_dwdp", enable_dp_attention_local_control_broadcast=True ) declare_resolution( server_args, "_handle_dwdp", enable_dp_lm_head=True, ) declare_resolution( server_args, "_handle_dwdp", moe_dense_tp_size=1, ) declare_resolution( server_args, "_handle_dwdp", ep_size=cfg.dwdp_size, ) declare_resolution( server_args, "_handle_dwdp", moe_dp_size=1, ) declare_resolution( server_args, "_handle_dwdp", moe_a2a_backend="none", ) envs.SGLANG_SCHEDULER_SKIP_ALL_GATHER.set(True) declare_resolution( server_args, "_handle_dwdp", disable_cuda_graph=True, ) logger.info( f"DWDP enabled: dwdp_size={cfg.dwdp_size}, " f"auto-forced dp_size={cfg.dp_size}, ep_size={cfg.dwdp_size}, " f"moe_dense_tp_size=1, moe_a2a_backend=none, " f"dp_attention_local_control_broadcast=True, " f"enable_dp_lm_head=True, SCHEDULER_SKIP_ALL_GATHER=True, " f"disable_cuda_graph=True" ) def handle_elastic_ep(server_args: Any): from sglang.srt.arg_groups.validation_hook import validate_ib_devices cfg = resolving_view(server_args) if cfg.elastic_ep_rejoin: if cfg.ep_join_mode is None: logger.warning( "--elastic-ep-rejoin is deprecated, use --elastic-ep-join-mode recover instead." ) declare_resolution( server_args, "_handle_elastic_ep", ep_join_mode="recover", ) else: assert cfg.ep_join_mode == "recover", ( "--elastic-ep-rejoin (deprecated) conflicts with " f"--elastic-ep-join-mode {cfg.ep_join_mode}." ) if cfg.elastic_ep_backend is not None: if cfg.enable_eplb: if cfg.eplb_algorithm == "auto": declare_resolution( server_args, "_handle_elastic_ep", eplb_algorithm="elasticity_aware", ) assert cfg.eplb_algorithm in [ "elasticity_aware", "elasticity_aware_hierarchical", ], "Elastic EP requires eplb_algorithm to be set to 'auto' or 'elasticity_aware(_hierarchical)'." assert cfg.pp_size == 1, "PP size should be set to 1 under elastic EP" if cfg.elastic_ep_backend == "mooncake": declare_resolution( server_args, "_handle_elastic_ep", mooncake_ib_device=validate_ib_devices(cfg.mooncake_ib_device), ) if cfg.ep_join_mode is not None: assert ( cfg.elastic_ep_backend is not None ), "--elastic-ep-join-mode requires --elastic-ep-backend to be set." if cfg.ep_join_mode == "scale": assert cfg.node_rank == 1, ( "Elastic EP scale-up requires one joining TP group at " f"--node-rank 1 (got {cfg.node_rank})." ) assert cfg.ep_join_rank_offset > 0, ( "Elastic EP scale joiners require " "--elastic-ep-join-rank-offset set to the current " "effective EP size." ) if cfg.ep_join_rank_offset != 0: assert cfg.ep_join_mode == "scale", ( "--elastic-ep-join-rank-offset is only valid with " "--elastic-ep-join-mode scale." ) assert cfg.ep_join_rank_offset >= 0, "elastic EP join rank offset must be >= 0." if cfg.max_ep_size is not None: assert ( cfg.elastic_ep_backend is not None ), "--max-ep-size requires --elastic-ep-backend to be set." assert cfg.max_ep_size > 0, "--max-ep-size must be a positive integer." scaling_active = ( cfg.elastic_ep_backend is not None and cfg.max_ep_size is not None and cfg.max_ep_size > cfg.tp_size ) if cfg.elastic_ep_initial_size is not None: assert scaling_active, ( "--elastic-ep-initial-size is only valid for an Elastic EP " "deployment with --max-ep-size larger than its local TP size." ) if scaling_active: resolved = resolved_view(server_args) assert ( cfg.elastic_ep_scale_timeout > 0 ), "--elastic-ep-scale-timeout must be greater than zero." assert cfg.tokenizer_worker_num == 1, ( "Elastic EP runtime scale-up currently requires " "--tokenizer-worker-num 1." ) assert ( not cfg.use_ray ), "Elastic EP runtime scale-up does not support --use-ray." assert not cfg.enable_elastic_expert_backup, ( "Elastic EP runtime scale-up does not support " "--enable-elastic-expert-backup." ) declare_resolution( server_args, "_handle_elastic_ep", enable_dp_attention_local_control_broadcast=True, ) if cfg.ep_join_mode == "scale": assert cfg.elastic_ep_initial_size is not None, ( "Elastic EP scale joiners require --elastic-ep-initial-size " "set to the primary deployment's launch-time EP size." ) assert cfg.elastic_ep_initial_size <= cfg.ep_join_rank_offset, ( "--elastic-ep-initial-size cannot exceed the current EP size " f"(initial={cfg.elastic_ep_initial_size}, " f"current={cfg.ep_join_rank_offset})." ) join_target = cfg.ep_join_rank_offset + cfg.tp_size assert join_target <= cfg.max_ep_size, ( "Elastic EP joining group exceeds --max-ep-size " f"(join_target={join_target}, max_ep_size={cfg.max_ep_size})." ) if cfg.tp_size == 1: assert cfg.moe_dense_tp_size == 1, ( "A single-rank Elastic EP joining group requires " "--moe-dense-tp-size 1." ) else: if cfg.elastic_ep_initial_size is None: declare_resolution( server_args, "_handle_elastic_ep", elastic_ep_initial_size=cfg.tp_size, ) assert cfg.elastic_ep_initial_size == cfg.tp_size, ( "The primary --elastic-ep-initial-size must equal its " f"launch-time TP size ({cfg.tp_size})." ) assert cfg.elastic_ep_initial_size > 0 assert cfg.load_balance_method == "round_robin", ( "Elastic EP scale-up requires --load-balance-method round_robin; " "load-aware methods " "require global-rank load snapshots after scale " f"(got {cfg.load_balance_method})." ) assert cfg.elastic_ep_backend == "mooncake", ( "Elastic EP runtime scale-up requires --elastic-ep-backend " f"mooncake (got elastic_ep_backend={cfg.elastic_ep_backend})." ) assert cfg.pp_size == 1, ( "Elastic EP scale-up requires --pp-size 1 " f"(got pp_size={cfg.pp_size}); WORLD must not span PP stages." ) decode_cuda_graph_disabled = ( cfg.cuda_graph_config.decode.backend == Backend.DISABLED ) prefill_cuda_graph_disabled = ( cfg.cuda_graph_config.prefill.backend == Backend.DISABLED ) assert decode_cuda_graph_disabled and prefill_cuda_graph_disabled, ( "Elastic EP runtime scale-up requires decode and prefill CUDA " "graphs to be disabled." ) assert resolved.enable_dp_attention, ( "Elastic EP scale-up requires --enable-dp-attention; without it " "the TP group is not equivalent to WORLD and the post-scale " "collective path is invalid." ) assert resolved.enable_dp_lm_head, ( "Elastic EP scale-up requires --enable-dp-lm-head so output " "projection does not depend on the joining group's TP size." ) assert resolved.attn_cp_size == 1, ( "Elastic EP scale-up requires --attn-cp-size 1 " f"(got attn_cp_size={resolved.attn_cp_size})." ) assert cfg.moe_dp_size == 1, ( "Elastic EP scale-up requires --moe-dp-size 1 " f"(got moe_dp_size={cfg.moe_dp_size})." ) assert resolved.ep_size == cfg.tp_size, ( "Elastic EP scale-up requires ep_size == tp_size " f"(got ep_size={resolved.ep_size}, tp_size={cfg.tp_size}); EP, TP " "and the attention DP group must all coincide with WORLD." ) assert cfg.dp_size == cfg.tp_size, ( "Elastic EP scale-up requires dp_size == tp_size " f"(got dp_size={cfg.dp_size}, tp_size={cfg.tp_size})." ) assert resolved.moe_a2a_backend == "nixl", ( "Elastic EP scale-up requires --moe-a2a-backend nixl " f"(got moe_a2a_backend={resolved.moe_a2a_backend})." ) def handle_eplb_and_dispatch(server_args: Any): cfg = resolving_view(server_args) if cfg.enable_eplb and (cfg.expert_distribution_recorder_mode is None): declare_resolution( server_args, "_handle_eplb_and_dispatch", expert_distribution_recorder_mode="stat", ) logger.warning( "EPLB is enabled. The expert_distribution_recorder_mode is automatically set." ) # Without an a2a backend all EP ranks run the MoE over the same tokens and # sum their partial outputs, so the pick has to agree across ranks. needs_rank_invariant_dispatch = resolved_view(server_args).moe_a2a_backend == "none" if (cfg.enable_eplb or (cfg.init_expert_location != "trivial")) and ( cfg.ep_dispatch_algorithm is None ): declare_resolution( server_args, "_handle_eplb_and_dispatch", ep_dispatch_algorithm=( "dynamic" if needs_rank_invariant_dispatch else "static" ), ) # `dynamic` / `fake` switch to the row-index pick; `static` reads a # per-rank table and `lp` samples inside its kernel. if needs_rank_invariant_dispatch and cfg.ep_dispatch_algorithm in ( "static", "lp", ): raise ValueError( f"--ep-dispatch-algorithm {cfg.ep_dispatch_algorithm} picks a " "different physical replica per rank, which only holds up when an " "a2a backend routes each token to a single rank. Use " "--ep-dispatch-algorithm dynamic with --moe-a2a-backend none." ) if cfg.enable_eplb and cfg.ep_join_mode != "scale": assert resolved_view(server_args).ep_size > 1 def handle_legacy_cp_arguments(server_args: Any): cfg = resolving_view(server_args) legacy_mode_to_strategy = { "in-seq-split": "zigzag", "round-robin-split": "interleave", } strategy_to_legacy_mode = { "zigzag": "in-seq-split", "interleave": "round-robin-split", } if cfg.enable_prefill_context_parallel or cfg.enable_dsa_prefill_context_parallel: declare_resolution( server_args, "_handle_legacy_cp_arguments", enable_prefill_cp=True, ) if cfg.enable_prefill_context_parallel and cfg.cp_strategy is None: declare_resolution( server_args, "_handle_legacy_cp_arguments", cp_strategy=legacy_mode_to_strategy[cfg.prefill_cp_mode], ) if cfg.enable_dsa_prefill_context_parallel and cfg.cp_strategy is None: declare_resolution( server_args, "_handle_legacy_cp_arguments", cp_strategy=legacy_mode_to_strategy[cfg.dsa_prefill_cp_mode], ) if cfg.enable_prefill_context_parallel and cfg.enable_dsa_prefill_context_parallel: return if not cfg.enable_prefill_cp or cfg.cp_strategy is None: return mode = strategy_to_legacy_mode[cfg.cp_strategy] use_dsa_legacy_aliases = cfg.enable_dsa_prefill_context_parallel or getattr( resolved_view(server_args), "attention_backend", None ) in ("dsa", "dsv4") if use_dsa_legacy_aliases: declare_resolution( server_args, "_handle_legacy_cp_arguments", enable_dsa_prefill_context_parallel=True, ) declare_resolution( server_args, "_handle_legacy_cp_arguments", enable_prefill_context_parallel=False, ) else: declare_resolution( server_args, "_handle_legacy_cp_arguments", enable_prefill_context_parallel=True, ) declare_resolution( server_args, "_handle_legacy_cp_arguments", dsa_prefill_cp_mode=mode, ) declare_resolution( server_args, "_handle_legacy_cp_arguments", prefill_cp_mode=mode, ) def handle_expert_distribution_metrics(server_args: Any): cfg = resolving_view(server_args) if "SGLANG_ENABLE_EPLB_BALANCEDNESS_METRIC" in os.environ: raise ValueError( "SGLANG_ENABLE_EPLB_BALANCEDNESS_METRIC is no longer supported. Use " "--expert-balancedness-report-mode with one of: off, server_log, " "prometheus, both." ) if should_report_expert_balancedness(server_args) and ( cfg.expert_distribution_recorder_mode is None ): declare_resolution( server_args, "_handle_expert_distribution_metrics", expert_distribution_recorder_mode="stat", ) if cfg.expert_distribution_recorder_buffer_size is None: if (x := cfg.eplb_rebalance_num_iterations) is not None: declare_resolution( server_args, "_handle_expert_distribution_metrics", expert_distribution_recorder_buffer_size=x, ) elif cfg.expert_distribution_recorder_mode is not None: declare_resolution( server_args, "_handle_expert_distribution_metrics", expert_distribution_recorder_buffer_size=1000, )