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

[PD-Disagg] Improve KVManager init across all backends (#19240)

Browse Source 
Signed-off-by: Shangming Cai <csmthu@gmail.com>
This commit is contained in:
Shangming Cai
2026-02-25 10:37:09 +08:00
committed by GitHub
parent c0fdfd4b92
commit 0fac2796b6
4 changed files with 25 additions and 51 deletions
Download Patch File Download Diff File Expand all files Collapse all files

24
python/sglang/srt/disaggregation/common/conn.py
View File

@@ -6,8 +6,9 @@ import logging
import socket
import threading
import time
from collections import defaultdict
from functools import cache
from typing import Dict, List, Optional, Tuple, Union
from typing import Dict, List, Optional, Set, Tuple, Union
import numpy as np
import numpy.typing as npt
@@ -107,11 +108,32 @@ class CommonKVManager(BaseKVManager):
self.transfer_infos = {}
self.decode_kv_args_table = {}
self.pp_group = get_pp_group()
# If a timeout happens on the prefill side, it means prefill instances
# fail to receive the KV indices from the decode instance of this request.
# These timeout requests should be aborted to release the tree cache.
self.bootstrap_timeout = envs.SGLANG_DISAGGREGATION_BOOTSTRAP_TIMEOUT.get()
elif self.disaggregation_mode == DisaggregationMode.DECODE:
self.connection_pool: Dict[str, Dict[str, Union[str, int]]] = {}
self.connection_lock = threading.Lock()
self.required_prefill_response_num_table: Dict[int, int] = {}
self.prefill_info_table: Dict[str, PrefillServerInfo] = {}
self.heartbeat_failures: Dict[str, int] = {}
self.session_pool: Dict = defaultdict(requests.Session)
self.session_pool_lock = threading.Lock()
self.addr_to_rooms_tracker: Dict[str, Set[int]] = defaultdict(set)
self.prefill_response_tracker: Dict[int, Set[int]] = defaultdict(set)
# Heartbeat interval should be at least 2 seconds
self.heartbeat_interval = max(
envs.SGLANG_DISAGGREGATION_HEARTBEAT_INTERVAL.get(), 2.0
)
# Heartbeat failure should be at least 1
self.max_failures = max(
envs.SGLANG_DISAGGREGATION_HEARTBEAT_MAX_FAILURE.get(), 1
)
# If a timeout happens on the decode side, it means decode instances
# fail to receive the KV Cache transfer done signal after bootstrapping.
# These timeout requests should be aborted to release the tree cache.
self.waiting_timeout = envs.SGLANG_DISAGGREGATION_WAITING_TIMEOUT.get()
else:
raise ValueError(
f"Unsupported DisaggregationMode: {self.disaggregation_mode}"

25
python/sglang/srt/disaggregation/mooncake/conn.py
View File

@@ -9,11 +9,10 @@ import struct
import threading
import time
from collections import defaultdict
from typing import Dict, List, Optional, Set, Tuple
from typing import List, Optional, Tuple
import numpy as np
import numpy.typing as npt
import requests
from sglang.srt.disaggregation.base.conn import KVArgs, KVPoll
from sglang.srt.disaggregation.common.conn import (
@@ -206,33 +205,11 @@ class MooncakeKVManager(CommonKVManager):
threading.Thread(
target=self.transfer_worker, args=(queue, executor), daemon=True
).start()
# If a timeout happens on the prefill side, it means prefill instances
# fail to receive the KV indices from the decode instance of this request.
# These timeout requests should be aborted to release the tree cache.
self.bootstrap_timeout = envs.SGLANG_DISAGGREGATION_BOOTSTRAP_TIMEOUT.get()
self.enable_custom_mem_pool, self.custom_mem_pool_type = (
check_mooncake_custom_mem_pool_enabled()
)
elif self.disaggregation_mode == DisaggregationMode.DECODE:
self.heartbeat_failures = {}
self.session_pool = defaultdict(requests.Session)
self.session_pool_lock = threading.Lock()
self.addr_to_rooms_tracker = defaultdict(set)
self.prefill_response_tracker: Dict[int, Set[int]] = defaultdict(set)
# Heartbeat interval should be at least 2 seconds
self.heartbeat_interval = max(
envs.SGLANG_DISAGGREGATION_HEARTBEAT_INTERVAL.get(), 2.0
)
# Heartbeat failure should be at least 1
self.max_failures = max(
envs.SGLANG_DISAGGREGATION_HEARTBEAT_MAX_FAILURE.get(), 1
)
self.start_decode_thread()
# If a timeout happens on the decode side, it means decode instances
# fail to receive the KV Cache transfer done signal after bootstrapping.
# These timeout requests should be aborted to release the tree cache.
self.waiting_timeout = envs.SGLANG_DISAGGREGATION_WAITING_TIMEOUT.get()
def init_engine(self):
self.engine = get_mooncake_transfer_engine()

11
python/sglang/srt/disaggregation/mori/conn.py
View File

@@ -7,8 +7,7 @@ import os
import struct
import threading
import time
from collections import defaultdict
from typing import Dict, List, Optional, Set, Tuple
from typing import Dict, List, Optional, Tuple
import msgspec
import numpy as np
@@ -194,16 +193,8 @@ class MoriKVManager(CommonKVManager):
self.transfer_lock = threading.Lock()
self._register_local_buffers()
if self.disaggregation_mode == DisaggregationMode.PREFILL:
self.bootstrap_timeout = get_int_env_var(
"SGLANG_DISAGGREGATION_BOOTSTRAP_TIMEOUT", 300
)
self._start_bootstrap_thread()
elif self.disaggregation_mode == DisaggregationMode.DECODE:
self.waiting_timeout = get_int_env_var(
"SGLANG_DISAGGREGATION_WAITING_TIMEOUT", 300
)
self.prefill_response_tracker: Dict[int, Set[int]] = defaultdict(set)
self.addr_to_rooms_tracker = defaultdict(set)
self.room_to_bootstrap_addr: Dict[int, str] = {}
self._start_decode_thread()

16
python/sglang/srt/disaggregation/nixl/conn.py
View File

@@ -11,7 +11,6 @@ from typing import Dict, List, Optional, Set
import numpy as np
import numpy.typing as npt
import requests
from sglang.srt.disaggregation.base.conn import KVArgs, KVPoll
from sglang.srt.disaggregation.common.conn import (
@@ -193,21 +192,6 @@ class NixlKVManager(CommonKVManager):
self.transfer_statuses: Dict[int, TransferStatus] = defaultdict(
TransferStatus
)
self.heartbeat_failures = {}
self.session_pool = defaultdict(requests.Session)
self.session_pool_lock = threading.Lock()
self.addr_to_rooms_tracker = defaultdict(set)
self.connection_lock = threading.Lock()
# Heartbeat interval should be at least 2 seconds
self.heartbeat_interval = max(
envs.SGLANG_DISAGGREGATION_HEARTBEAT_INTERVAL.get(), 2.0
)
# Heartbeat failure should be at least 1
self.max_failures = max(
envs.SGLANG_DISAGGREGATION_HEARTBEAT_MAX_FAILURE.get(), 1
)
self.waiting_timeout = envs.SGLANG_DISAGGREGATION_WAITING_TIMEOUT.get()
self._start_heartbeat_checker_thread()
else:
raise ValueError(