Add staging buffer CI test and documentation for heterogeneous TP (#21921)

Co-authored-by: Shangming Cai <csmthu@gmail.com>

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

@@ -140,7 +140,7 @@ class StagingBuffer:
             alloc_method = "custom_mem_pool (cuMemCreate)"
         else:
             self.buffer = torch.empty(size_bytes, dtype=torch.uint8, device=device)
-            alloc_method = "cudaMalloc (NVLink incompatible!)"
+            alloc_method = "cudaMalloc"
         self.data_ptr = self.buffer.data_ptr()
 
         logger.info(

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

@@ -517,9 +517,10 @@ def init_staging_buffers(engine, kv_args, count: int) -> list:
 
     _, custom_mem_pool, pool_type = init_mooncake_custom_mem_pool(device)
     if custom_mem_pool is None:
-        logger.warning(
-            "No mooncake custom mem pool available for staging buffer. "
-            "NVLink transport will NOT work. Set SGLANG_MOONCAKE_CUSTOM_MEM_POOL."
+        logger.info(
+            "Staging buffer using cudaMalloc (no custom mem pool). "
+            "This works for all GPU architectures. "
+            "For NVLink/MNNVL transport, set SGLANG_MOONCAKE_CUSTOM_MEM_POOL."
         )
 
     buffers = []

python/sglang/srt/disaggregation/decode.py

@@ -293,6 +293,11 @@ class DecodePreallocQueue:
         self._ensure_last_attempt_time: Dict[str, float] = {}
         self._ensure_retry_interval: float = 1.0  # seconds
         self.enable_staging = envs.SGLANG_DISAGG_STAGING_BUFFER.get()
+        if self.enable_staging and self.is_mla_backend:
+            raise RuntimeError(
+                "SGLANG_DISAGG_STAGING_BUFFER is designed for non-MLA models "
+                "(e.g. GQA, MHA). MLA models should not set this flag."
+            )
         self.kv_manager = self._init_kv_manager()
         if self.enable_staging:
             self.transfer_queue._init_staging_handler(self.kv_manager)
@@ -944,7 +949,10 @@ class DecodeTransferQueue:
         self.queue.extend(decode_reqs)
         if self.enable_staging:
             for dr in decode_reqs:
-                if dr.kv_receiver.require_staging:
+                if (
+                    hasattr(dr.kv_receiver, "require_staging")
+                    and dr.kv_receiver.require_staging
+                ):
                     self.staging_handler.register_decode_req(dr.req.bootstrap_room, dr)
 
     def _commit_transfer_to_req(self, decode_req: DecodeRequest) -> bool:

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

@@ -85,6 +85,7 @@ class FakeKVReceiver(BaseKVReceiver):
     ):
         self.bootstrap_done = False
         self.has_sent_metadata = False
+        self.require_staging: bool = False
 
     def poll(self) -> KVPoll:
         if not self.bootstrap_done:

python/sglang/srt/disaggregation/prefill.py

@@ -122,6 +122,11 @@ class PrefillBootstrapQueue:
         self.max_total_num_tokens = max_total_num_tokens
         self.scheduler = scheduler
         self.transfer_backend = transfer_backend
+        if envs.SGLANG_DISAGG_STAGING_BUFFER.get() and self.is_mla_backend:
+            raise RuntimeError(
+                "SGLANG_DISAGG_STAGING_BUFFER is designed for non-MLA models "
+                "(e.g. GQA, MHA). MLA models should not set this flag."
+            )
         self.kv_manager = self._init_kv_manager()
 
         if self.scheduler.tp_worker.is_hybrid_swa: