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

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
@@ -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()