[Disagg] GPU staging buffer with dynamic ring allocator for heterogeneous TP KV transfer

python/sglang/srt/disaggregation/common/conn.py

@@ -148,6 +148,7 @@ def __init__(
             # 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.enable_staging: bool = False
             self.connection_pool: Dict[str, Dict[str, Union[str, int]]] = {}
             self.connection_lock = threading.Lock()
             self.required_prefill_response_num_table: Dict[int, int] = {}
@@ -501,6 +502,7 @@ def __init__(
         self.bootstrap_addr = bootstrap_addr
         self.kv_mgr = mgr
         self.conclude_state: Optional[KVPoll] = None
+        self.require_staging: bool = False
         self.kv_mgr.addr_to_rooms_tracker[self.bootstrap_addr].add(self.bootstrap_room)
         self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Bootstrapping)
 
@@ -529,6 +531,12 @@ def init(self, prefill_dp_rank: int):
             self.required_prefill_response_num
         )
 
+        if self.kv_mgr.enable_staging:
+            self.require_staging = (
+                self.prefill_info.attn_tp_size != 0
+                and self.prefill_info.attn_tp_size != self.kv_mgr.attn_tp_size
+            )
+
         self.prefill_dp_rank = prefill_dp_rank
         self._setup_bootstrap_infos()
         self.kv_mgr.update_status(self.bootstrap_room, KVPoll.WaitingForInput)