minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[rust-server] PD disaggregation support (#33125)

Browse Source 
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Rain Jiang <rain-jiang@outlook.com>
This commit is contained in:
Kan Wu
2026-08-02 18:32:59 -07:00
committed by GitHub
co-authored by Claude Fable 5 Cursor Rain Jiang
parent c844244da5
commit 5d2dbb35a6
15 changed files with 1230 additions and 49 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/arg_groups/pd_disaggregation_hook.py
+37
View File
@@ -1,5 +1,6 @@
from __future__ import annotations
import dataclasses
import logging
import os
from typing import TYPE_CHECKING
@@ -101,6 +102,9 @@ def handle_pd_disaggregation(server_args: ServerArgs) -> None:
server_args.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 server_args.disaggregation_mode in ("prefill", "decode"):
if (
envs.SGLANG_DISAGG_STAGING_BUFFER.get()
@@ -111,3 +115,36 @@ def handle_pd_disaggregation(server_args: ServerArgs) -> None:
f"disaggregation_transfer_backend='mooncake' or 'nixl', "
f"got '{server_args.disaggregation_transfer_backend}'."
)
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.
"""
default_port = next(
f.default
for f in dataclasses.fields(server_args)
if f.name == "disaggregation_bootstrap_port"
)
if server_args.disaggregation_bootstrap_port not in (
default_port,
server_args.port,
):
raise ValueError(
"SGLANG_RUST_SERVER serves the PD KV bootstrap registry on the api "
"port itself; --disaggregation-bootstrap-port "
f"{server_args.disaggregation_bootstrap_port} conflicts with --port "
f"{server_args.port}. Drop --disaggregation-bootstrap-port (decode "
"nodes and the PD router must then target the prefill api port)."
)
if server_args.disaggregation_bootstrap_port != server_args.port:
logger.info(
"SGLANG_RUST_SERVER: KV bootstrap registry is served on the api "
"port; disaggregation_bootstrap_port %d -> %d",
server_args.disaggregation_bootstrap_port,
server_args.port,
)
server_args.disaggregation_bootstrap_port = server_args.port
python/sglang/srt/entrypoints/http_server.py
+8 -6
View File
@@ -2251,10 +2251,6 @@ def _execute_server_warmup(server_args: ServerArgs):
_global_state.tokenizer_manager.server_status = ServerStatus.Up
else:
# TODO: @rainj-me fix this when Rust server supports disaggregation
assert (
not envs.SGLANG_RUST_SERVER.get()
), "Rust server is not supported for disaggregation warmup for now"
logger.info(f"Start of pd disaggregation warmup ...")
status_codes = asyncio.run(
_send_disaggregation_warmup_requests(
@@ -2272,14 +2268,20 @@ def _execute_server_warmup(server_args: ServerArgs):
server_args.dp_size,
)
logger.info("End of disaggregation warmup")
_global_state.tokenizer_manager.server_status = ServerStatus.Up
else:
logger.info(
"Disaggregation warmup failed (mode=%s), status codes: %s",
server_args.disaggregation_mode,
failed_status_codes,
)
_global_state.tokenizer_manager.server_status = ServerStatus.UnHealthy
# In rust-server mode there is no TokenizerManager (readiness is
# the Rust server's own /health), so skip the status update.
if not envs.SGLANG_RUST_SERVER.get():
_global_state.tokenizer_manager.server_status = (
ServerStatus.Up
if not failed_status_codes
else ServerStatus.UnHealthy
)
except Exception:
last_traceback = get_exception_traceback()
python/sglang/srt/managers/disagg_service.py
+24 -12
View File
@@ -27,18 +27,30 @@ def start_disagg_service(
host=server_args.host,
port=server_args.disaggregation_bootstrap_port,
)
is_create_store = (
server_args.node_rank == 0 and transfer_backend == TransferBackend.ASCEND
maybe_create_ascend_config_store(
server_args=server_args, transfer_backend=transfer_backend
)
if is_create_store:
try:
from memfabric_hybrid import create_config_store
ascend_url = os.getenv("ASCEND_MF_STORE_URL")
create_config_store(ascend_url)
except Exception as e:
error_message = f"Failed create mf store, invalid ascend_url."
error_message += f" With exception {e}"
raise error_message
return bootstrap_server
def maybe_create_ascend_config_store(
server_args: ServerArgs, transfer_backend: TransferBackend
) -> None:
"""Also called directly by the rust-server scheduler: there the KV
bootstrap registry is served by the embedded rust server's api listener
(one rust implementation covers every transfer backend — their
bootstrap-server subclasses are all plain ``CommonKVBootstrapServer``,
which the rust registry ports verbatim), leaving this store as the only
``start_disagg_service`` duty left to perform."""
if not (server_args.node_rank == 0 and transfer_backend == TransferBackend.ASCEND):
return
try:
from memfabric_hybrid import create_config_store
ascend_url = os.getenv("ASCEND_MF_STORE_URL")
create_config_store(ascend_url)
except Exception as e:
raise RuntimeError(
f"Failed create mf store, invalid ascend_url. With exception {e}"
)
python/sglang/srt/managers/scheduler.py
+40 -13
View File
@@ -99,6 +99,7 @@ from sglang.srt.layers.quantization.fp8_utils import initialize_fp8_gemm_config
from sglang.srt.layers.quantization.unquant import initialize_bf16_gemm_config
from sglang.srt.lora.lora_drainer import LoRADrainer
from sglang.srt.lora.lora_overlap_loader import LoRAOverlapLoader
from sglang.srt.managers.disagg_service import maybe_create_ascend_config_store
from sglang.srt.managers.hisparse_coordinator import HiSparseCoordinator
from sglang.srt.managers.io_struct import (
AbortReq,
@@ -590,6 +591,15 @@ class Scheduler(
# Init profiler
self.init_profiler()
# Start the embedded Rust frontend (rank 0). Must precede
# init_disaggregation: on PD prefill the rust api listener also serves
# the KV bootstrap registry, and the KVManagers built there register to
# it synchronously. (The listener is bound synchronously inside launch,
# so the registry is accepting once this returns.) Must also precede
# the request receiver, which reads self.recv_from_tokenizer to pick
# its ingress transport.
self.maybe_init_rust_server()
# Init prefill-decodedisaggregation
self.init_disaggregation()
@@ -618,10 +628,6 @@ class Scheduler(
self.maybe_init_scripted_scheduler_hook()
# Start the embedded Rust frontend (rank 0) before the request receiver,
# which reads self.rust_ring_recv to pick its ingress transport.
self.maybe_init_rust_server()
self.init_request_receiver()
self.init_dp_attn_adapter()
@@ -1214,6 +1220,21 @@ class Scheduler(
get_disagg().disaggregation_transfer_backend
)
# In rust-server mode the KV bootstrap registry is already serving on
# the rust api listener (maybe_init_rust_server runs before this
# method — the PrefillBootstrapQueue's KVManager below registers to it
# synchronously, and a failed registration only retries ~60s then logs,
# leaving every PD request unroutable). Only the ascend config store,
# which start_disagg_service would otherwise create, is left to do.
if (
self.disaggregation_mode == DisaggregationMode.PREFILL
and self._hosts_rust_server()
):
maybe_create_ascend_config_store(
server_args=self.server_args,
transfer_backend=self.transfer_backend,
)
# todo: should we fix this when enabling mtp or it doesn't matter since we only enable mtp in decode node thus we don't transfer draft kvs between P and D?
draft_token_to_kv_pool = kv_cache_builder.get_draft_kv_pool(
draft_worker=self.draft_worker,
@@ -1844,17 +1865,22 @@ class Scheduler(
else:
self.scripted_scheduler_hook = None
def _hosts_rust_server(self) -> bool:
"""Whether this scheduler rank embeds the Rust server (rank 0 only) —
and with it the server-process duties a Python ``TokenizerManager``
would otherwise own (e.g. serving the PD KV bootstrap registry)."""
return envs.SGLANG_RUST_SERVER.get() and (
self.ps.pp_rank == 0
and self.ps.attn_tp_rank == 0
and self.ps.attn_cp_rank == 0
)
def maybe_init_rust_server(self) -> None:
"""Start the embedded Rust server (rank 0) if ``SGLANG_RUST_SERVER`` is
set, and point the ingress receiver at it. All the plumbing lives in
``RustServer`` (scheduler_components/rust_scheduler.py)."""
is_rank_zero = (
self.ps.pp_rank == 0
and self.ps.attn_tp_rank == 0
and self.ps.attn_cp_rank == 0
)
if not (envs.SGLANG_RUST_SERVER.get() and is_rank_zero):
if not self._hosts_rust_server():
# Always define the attribute: init_output_streamer and the
# process_input_requests hook read self.rust_server unconditionally.
self.rust_server = None
@@ -2303,9 +2329,10 @@ class Scheduler(
f"bootstrap room id. {req.rid=}"
)
logger.error(error_msg)
recv_req.time_stats.trace_ctx.abort(
abort_info={"reason": error_msg}
)
if not envs.SGLANG_RUST_SERVER.get():
recv_req.time_stats.trace_ctx.abort(
abort_info={"reason": error_msg}
)
prepare_abort(req, error_msg, status_code=HTTPStatus.BAD_REQUEST)
self.output_streamer.stream_output([req], req.return_logprob)
return