diff --git a/python/sglang/srt/disaggregation/common/conn.py b/python/sglang/srt/disaggregation/common/conn.py index 6e0cd40a4..595bb3ef1 100644 --- a/python/sglang/srt/disaggregation/common/conn.py +++ b/python/sglang/srt/disaggregation/common/conn.py @@ -92,6 +92,7 @@ class CommonKVManager(BaseKVManager): self.prefill_attn_tp_size_table: Dict[str, int] = {} self.prefill_dp_size_table: Dict[str, int] = {} self.prefill_pp_size_table: Dict[str, int] = {} + self.prefill_page_size_table: Dict[str, Optional[int]] = {} else: raise ValueError( f"Unsupported DisaggregationMode: {self.disaggregation_mode}" @@ -127,6 +128,7 @@ class CommonKVManager(BaseKVManager): "system_dp_rank": self.system_dp_rank, "rank_ip": self.local_ip, "rank_port": self.rank_port, + "page_size": self.kv_args.page_size, } try: @@ -243,6 +245,7 @@ class CommonKVReceiver(BaseKVReceiver): self.prefill_attn_tp_size, self.prefill_dp_size, self.prefill_pp_size, + self.prefill_page_size, ) = self._get_prefill_parallel_info_from_server() if ( self.prefill_attn_tp_size is None @@ -256,19 +259,33 @@ class CommonKVReceiver(BaseKVReceiver): self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Failed) self.bootstrap_infos = None return - else: - logger.debug( - f"Fetch prefill parallel info from [{self.bootstrap_addr}]: DP size:{self.prefill_dp_size}, TP size:{self.prefill_attn_tp_size} PP size:{self.prefill_pp_size}" - ) - self.kv_mgr.prefill_attn_tp_size_table[self.bootstrap_addr] = ( - self.prefill_attn_tp_size - ) - self.kv_mgr.prefill_dp_size_table[self.bootstrap_addr] = ( - self.prefill_dp_size - ) - self.kv_mgr.prefill_pp_size_table[self.bootstrap_addr] = ( - self.prefill_pp_size - ) + + if self.prefill_page_size is not None: + decode_page_size = self.kv_mgr.kv_args.page_size + if self.prefill_page_size != decode_page_size: + error_msg = ( + f"Page size mismatch: prefill server has page_size={self.prefill_page_size}, " + f"but decode server has page_size={decode_page_size}. " + f"Both servers must use the same --page-size value." + ) + logger.error(error_msg) + raise RuntimeError(error_msg) + + logger.debug( + f"Fetch prefill parallel info from [{self.bootstrap_addr}]: DP size:{self.prefill_dp_size}, TP size:{self.prefill_attn_tp_size} PP size:{self.prefill_pp_size} Page size:{self.prefill_page_size}" + ) + self.kv_mgr.prefill_attn_tp_size_table[self.bootstrap_addr] = ( + self.prefill_attn_tp_size + ) + self.kv_mgr.prefill_dp_size_table[self.bootstrap_addr] = ( + self.prefill_dp_size + ) + self.kv_mgr.prefill_pp_size_table[self.bootstrap_addr] = ( + self.prefill_pp_size + ) + self.kv_mgr.prefill_page_size_table[self.bootstrap_addr] = ( + self.prefill_page_size + ) else: self.prefill_attn_tp_size = self.kv_mgr.prefill_attn_tp_size_table[ self.bootstrap_addr @@ -279,6 +296,9 @@ class CommonKVReceiver(BaseKVReceiver): self.prefill_pp_size = self.kv_mgr.prefill_pp_size_table[ self.bootstrap_addr ] + self.prefill_page_size = self.kv_mgr.prefill_page_size_table.get( + self.bootstrap_addr + ) # Handling for PD with different TP sizes per DP rank if self.kv_mgr.attn_tp_size == self.prefill_attn_tp_size: @@ -423,7 +443,7 @@ class CommonKVReceiver(BaseKVReceiver): def _get_prefill_parallel_info_from_server( self, - ) -> Tuple[Optional[int], Optional[int], Optional[int]]: + ) -> Tuple[Optional[int], Optional[int], Optional[int], Optional[int]]: """Fetch the prefill parallel info from the bootstrap server.""" try: url = f"http://{self.bootstrap_addr}/route?engine_rank={-1}&target_dp_group={-1}&target_pp_rank={-1}" @@ -434,15 +454,16 @@ class CommonKVReceiver(BaseKVReceiver): int(prefill_parallel_info["prefill_attn_tp_size"]), int(prefill_parallel_info["prefill_dp_size"]), int(prefill_parallel_info["prefill_pp_size"]), + int(prefill_parallel_info["prefill_page_size"]), ) else: logger.error( f"Failed to get prefill parallel info: {response.status_code}, {response.text}" ) - return None, None, None + return None, None, None, None except Exception as e: logger.error(f"Error fetching prefill parallel info from bootstrap: {e}") - return None, None, None + return None, None, None, None @classmethod def _connect(cls, endpoint: str, is_ipv6: bool = False): @@ -484,6 +505,7 @@ class CommonKVBootstrapServer(BaseKVBootstrapServer): self.pp_size = None self.attn_tp_size = None self.dp_size = None + self.page_size = None self.prefill_port_table: Dict[ int, Dict[int, Dict[int, Dict[str, Union[str, int]]]] ] = {} @@ -526,6 +548,7 @@ class CommonKVBootstrapServer(BaseKVBootstrapServer): system_dp_rank = data["system_dp_rank"] rank_ip = data["rank_ip"] rank_port = int(data["rank_port"]) + page_size = int(data["page_size"]) if self.attn_tp_size is None: self.attn_tp_size = attn_tp_size @@ -536,6 +559,9 @@ class CommonKVBootstrapServer(BaseKVBootstrapServer): if self.pp_size is None: self.pp_size = pp_size + if self.page_size is None and page_size is not None: + self.page_size = page_size + if role == "Prefill": if system_dp_size == 1: dp_group = attn_dp_rank @@ -576,6 +602,7 @@ class CommonKVBootstrapServer(BaseKVBootstrapServer): "prefill_attn_tp_size": self.attn_tp_size, "prefill_dp_size": self.dp_size, "prefill_pp_size": self.pp_size, + "prefill_page_size": self.page_size, } return web.json_response(prefill_parallel_info, status=200) diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index ad87b2354..51af67636 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -274,6 +274,7 @@ class DecodePreallocQueue: kv_args.kv_data_ptrs = kv_data_ptrs kv_args.kv_data_lens = kv_data_lens kv_args.kv_item_lens = kv_item_lens + kv_args.page_size = self.token_to_kv_pool.page_size kv_args.aux_data_ptrs, kv_args.aux_data_lens, kv_args.aux_item_lens = ( self.metadata_buffers.get_buf_infos()