289 lines
12 KiB
Python
289 lines
12 KiB
Python
from __future__ import annotations
|
|
|
|
import logging
|
|
import os
|
|
from typing import TYPE_CHECKING, Any
|
|
|
|
from sglang.srt.arg_groups.arg_utils import record_fields
|
|
from sglang.srt.arg_groups.overrides import (
|
|
declare_resolution,
|
|
model_config_of,
|
|
resolved_view,
|
|
resolving_view,
|
|
)
|
|
from sglang.srt.environ import envs
|
|
|
|
if TYPE_CHECKING:
|
|
from sglang.srt.server_args import ServerArgs
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def handle_pd_disaggregation(server_args: ServerArgs) -> None:
|
|
"""Validate and normalize PD-disaggregation server args."""
|
|
cfg = resolving_view(server_args)
|
|
|
|
# "mooncake_tcp" is mooncake with the TCP transport forced: set MC_FORCE_TCP
|
|
# so mooncake installs TcpTransport instead of RDMA, rewrite the backend to
|
|
# mooncake, and skip RDMA HCA selection. Must run before backend-name checks.
|
|
if cfg.disaggregation_transfer_backend == "mooncake_tcp":
|
|
os.environ.setdefault("MC_FORCE_TCP", "1")
|
|
declare_resolution(
|
|
server_args,
|
|
"handle_pd_disaggregation",
|
|
disaggregation_transfer_backend="mooncake",
|
|
)
|
|
declare_resolution(
|
|
server_args,
|
|
"handle_pd_disaggregation",
|
|
disaggregation_ib_device=None,
|
|
)
|
|
logger.info(
|
|
"disaggregation transfer backend 'mooncake_tcp' -> mooncake "
|
|
"with MC_FORCE_TCP=1 (TCP transport, no RDMA)"
|
|
)
|
|
|
|
if cfg.disaggregation_mode == "prefill" and cfg.dcp_size > 1:
|
|
logger.warning(
|
|
"DCP on a PD prefill server is supported when prefill and decode "
|
|
"use the same DCP layout, but it usually adds communication "
|
|
"overhead without improving prefill performance."
|
|
)
|
|
|
|
if cfg.disaggregation_mode == "decode" and cfg.dcp_size > 1:
|
|
# Fake transfer moves no KV and is only used for synthetic decode
|
|
# benchmarks, so it does not need the DCP relayout from Mooncake/NIXL.
|
|
if cfg.disaggregation_transfer_backend not in (
|
|
"mooncake",
|
|
"nixl",
|
|
"fake",
|
|
):
|
|
raise ValueError(
|
|
"PD decode DCP requires --disaggregation-transfer-backend "
|
|
"mooncake, nixl, or fake for synthetic benchmarking, got "
|
|
f"{cfg.disaggregation_transfer_backend!r}."
|
|
)
|
|
|
|
if cfg.disaggregation_mode == "decode":
|
|
if cfg.disaggregation_decode_enable_radix_cache:
|
|
if cfg.enable_hisparse:
|
|
raise ValueError(
|
|
"--disaggregation-decode-enable-radix-cache is incompatible "
|
|
"with --enable-hisparse"
|
|
)
|
|
if cfg.disaggregation_transfer_backend == "fake":
|
|
raise ValueError(
|
|
"--disaggregation-decode-enable-radix-cache is incompatible "
|
|
"with --disaggregation-transfer-backend fake"
|
|
)
|
|
if cfg.speculative_algorithm not in (None, "DSPARK"):
|
|
raise ValueError(
|
|
"--disaggregation-decode-enable-radix-cache is incompatible "
|
|
"with speculative decoding "
|
|
f"(--speculative-algorithm {cfg.speculative_algorithm})"
|
|
)
|
|
|
|
if resolved_view(server_args).enable_dp_attention:
|
|
logger.warning(
|
|
"EXPERIMENTAL: Decode radix cache with DP attention. "
|
|
"Requires prefix-aware DP rank routing for optimal cache hits."
|
|
)
|
|
declare_resolution(
|
|
server_args,
|
|
"handle_pd_disaggregation",
|
|
disable_radix_cache=False,
|
|
)
|
|
logger.warning("EXPERIMENTAL: Radix cache is enabled for decode server")
|
|
else:
|
|
declare_resolution(
|
|
server_args,
|
|
"handle_pd_disaggregation",
|
|
disable_radix_cache=True,
|
|
)
|
|
logger.warning("KV cache is forced as chunk cache for decode server")
|
|
|
|
# Default the number of *extra* decode req_to_token slots reserved for
|
|
# in-transfer (being-received-from-prefill) requests, on top of the
|
|
# max_running_requests-derived pool. Large batches get none; small
|
|
# per-worker batches reserve 2x the batch as cheap overlap headroom.
|
|
if cfg.disaggregation_decode_extra_slots is None:
|
|
extra_slots = 0
|
|
if cfg.max_running_requests is not None:
|
|
per_worker = cfg.max_running_requests // max(1, cfg.dp_size)
|
|
if per_worker <= 32:
|
|
extra_slots = per_worker * 2
|
|
declare_resolution(
|
|
server_args,
|
|
"handle_pd_disaggregation",
|
|
disaggregation_decode_extra_slots=extra_slots,
|
|
)
|
|
|
|
elif cfg.disaggregation_mode == "prefill":
|
|
assert cfg.disaggregation_transfer_backend != "fake", (
|
|
"Prefill server does not support 'fake' as the transfer backend"
|
|
)
|
|
|
|
if envs.SGLANG_RUST_SERVER.get():
|
|
_alias_bootstrap_port_to_api_port(server_args)
|
|
|
|
if cfg.disaggregation_mode in ("prefill", "decode"):
|
|
if (
|
|
envs.SGLANG_DISAGG_STAGING_BUFFER.get()
|
|
and cfg.disaggregation_transfer_backend not in ("mooncake", "nixl")
|
|
):
|
|
raise ValueError(
|
|
f"SGLANG_DISAGG_STAGING_BUFFER requires "
|
|
f"disaggregation_transfer_backend='mooncake' or 'nixl', "
|
|
f"got '{cfg.disaggregation_transfer_backend}'."
|
|
)
|
|
|
|
# Reject features whose role-specific state is not rebuilt on a flip.
|
|
if cfg.enable_pd_role_switch:
|
|
view = resolved_view(server_args)
|
|
unsupported = []
|
|
if view.enable_dp_attention:
|
|
unsupported.append("DP attention (--enable-dp-attention)")
|
|
if view.ep_size > 1:
|
|
unsupported.append(f"expert parallelism (--ep-size {view.ep_size})")
|
|
if view.moe_a2a_backend != "none":
|
|
unsupported.append(
|
|
f"MoE all-to-all (--moe-a2a-backend {view.moe_a2a_backend})"
|
|
)
|
|
if view.pp_size > 1:
|
|
unsupported.append(f"pipeline parallelism (--pp-size {view.pp_size})")
|
|
if view.dp_size > 1:
|
|
unsupported.append(f"data parallelism (--dp-size {view.dp_size})")
|
|
if view.dcp_size > 1:
|
|
unsupported.append(
|
|
f"decode context parallelism (--dcp-size {view.dcp_size})"
|
|
)
|
|
if view.speculative_algorithm is not None:
|
|
unsupported.append(
|
|
"speculative decoding "
|
|
f"(--speculative-algorithm {view.speculative_algorithm})"
|
|
)
|
|
if unsupported:
|
|
raise ValueError(
|
|
"--enable-pd-role-switch does not rebuild role-specific "
|
|
"state for the following features: "
|
|
+ ", ".join(unsupported)
|
|
+ ". Remove these options or drop --enable-pd-role-switch."
|
|
)
|
|
|
|
|
|
def _alias_bootstrap_port_to_api_port(server_args: ServerArgs) -> None:
|
|
"""Rust-server prefill serves the KV bootstrap registry on the api listener
|
|
itself, so the resolved bootstrap port must BE the api port — every internal
|
|
consumer (KVManager registration, PrefillBootstrapQueue) reads the resolved
|
|
field and agrees automatically. Decode is untouched: there the field names
|
|
the PREFILL side's bootstrap port and must stay as the operator set it.
|
|
"""
|
|
cfg = resolving_view(server_args)
|
|
default_port = next(
|
|
f.default
|
|
for f in record_fields(type(server_args))
|
|
if f.name == "disaggregation_bootstrap_port"
|
|
)
|
|
if cfg.disaggregation_bootstrap_port not in (
|
|
default_port,
|
|
cfg.port,
|
|
):
|
|
raise ValueError(
|
|
"SGLANG_RUST_SERVER serves the PD KV bootstrap registry on the api "
|
|
"port itself; --disaggregation-bootstrap-port "
|
|
f"{cfg.disaggregation_bootstrap_port} conflicts with --port "
|
|
f"{cfg.port}. Drop --disaggregation-bootstrap-port (decode "
|
|
"nodes and the PD router must then target the prefill api port)."
|
|
)
|
|
if cfg.disaggregation_bootstrap_port != cfg.port:
|
|
logger.info(
|
|
"SGLANG_RUST_SERVER: KV bootstrap registry is served on the api "
|
|
"port; disaggregation_bootstrap_port %d -> %d",
|
|
cfg.disaggregation_bootstrap_port,
|
|
cfg.port,
|
|
)
|
|
declare_resolution(
|
|
server_args,
|
|
"_alias_bootstrap_port_to_api_port",
|
|
disaggregation_bootstrap_port=cfg.port,
|
|
)
|
|
|
|
|
|
def handle_encoder_disaggregation(server_args: Any):
|
|
from sglang.srt.arg_groups.model_hook import handle_language_model_only
|
|
from sglang.srt.arg_groups.validation_hook import validate_ib_devices
|
|
from sglang.srt.server_args import resolve_encoder_transfer_backend
|
|
|
|
cfg = resolving_view(server_args)
|
|
handle_language_model_only(server_args)
|
|
if cfg.enable_prefix_mm_cache and not cfg.encoder_only:
|
|
raise ValueError(
|
|
"--enable-prefix-mm-cache requires --encoder-only to be enabled"
|
|
)
|
|
if cfg.encoder_only and cfg.language_only:
|
|
raise ValueError("Cannot set --encoder-only and --language-only together")
|
|
if cfg.encoder_only and not cfg.disaggregation_mode == "null":
|
|
raise ValueError(
|
|
"Cannot set --encoder-only and --disaggregation-mode prefill/decode together"
|
|
)
|
|
|
|
if cfg.language_only and len(cfg.encoder_urls) == 0:
|
|
logger.info(
|
|
"--language-only is set without --encoder-urls. Encoders are "
|
|
"expected to register dynamically via the "
|
|
"EncoderBootstrapServer."
|
|
)
|
|
|
|
# Validate IB devices when mooncake backend is used
|
|
if (
|
|
cfg.disaggregation_transfer_backend == "mooncake"
|
|
and cfg.disaggregation_mode in ("prefill", "decode")
|
|
) or cfg.encoder_transfer_backend == "mooncake":
|
|
declare_resolution(
|
|
server_args,
|
|
"_handle_encoder_disaggregation",
|
|
disaggregation_ib_device=validate_ib_devices(cfg.disaggregation_ib_device),
|
|
)
|
|
|
|
# Validate model type for encoder disaggregation
|
|
hf_config = model_config_of(server_args).hf_config
|
|
model_arch = hf_config.architectures[0]
|
|
if cfg.encoder_transfer_backend == "auto":
|
|
declare_resolution(
|
|
server_args,
|
|
"_handle_encoder_disaggregation",
|
|
encoder_transfer_backend=resolve_encoder_transfer_backend(
|
|
cfg.encoder_transfer_backend, model_arch, cfg.tp_size
|
|
),
|
|
)
|
|
if cfg.encoder_only or cfg.language_only:
|
|
logger.info(
|
|
"Encoder transfer backend auto-resolved to %s for %s at TP%d.",
|
|
cfg.encoder_transfer_backend,
|
|
model_arch,
|
|
cfg.tp_size,
|
|
)
|
|
if (cfg.encoder_only or cfg.language_only) and model_arch not in [
|
|
"Qwen2VLForConditionalGeneration",
|
|
"Qwen3VLForConditionalGeneration",
|
|
"Qwen2_5_VLForConditionalGeneration",
|
|
"Qwen3VLMoeForConditionalGeneration",
|
|
"Qwen3_5ForConditionalGeneration",
|
|
"Qwen3_5MoeForConditionalGeneration",
|
|
"InternS2PreviewForConditionalGeneration",
|
|
"Qwen3OmniMoeForConditionalGeneration",
|
|
"Qwen2AudioForConditionalGeneration",
|
|
"Qwen2_5OmniForConditionalGeneration",
|
|
"Dots3NoteForCausalLM",
|
|
"KimiVLForConditionalGeneration",
|
|
"KimiK25ForConditionalGeneration",
|
|
"KimiK3ForConditionalGeneration",
|
|
"MiMoV2ForCausalLM",
|
|
"Glm5NextForConditionalGeneration",
|
|
]:
|
|
raise ValueError(
|
|
f"Model type {model_arch} is not supported for encoder disaggregation. "
|
|
f"Supported architectures: Qwen2VL, Qwen3VL, Qwen3.5, InternS2, "
|
|
f"Qwen2Audio, Qwen2.5Omni, Dots3-Note, Kimi, MiMoV2, GLM5Next."
|
|
)
|