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

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

docs/advanced_features/pd_disaggregation.md

@@ -157,6 +157,58 @@ Please be aware that this setting will cause prefill instances to take a longer
 If a greater mean TTFT is acceptable, you can `export SGLANG_DISAGGREGATION_WAITING_TIMEOUT=600` (10 minutes) to relax the timeout condition.
 
 
+## Heterogeneous TP with GPU Staging Buffer
+
+When prefill and decode use different tensor parallelism (TP) sizes (e.g., prefill TP=4, decode DP attention with TP=1), the KV cache memory layout differs between the two sides. The **GPU staging buffer** solves this by gathering KV head slices into a contiguous buffer on the prefill side, performing bulk RDMA transfer, then scattering into the correct KV cache pages on the decode side. This provides **2–5x throughput improvement** over the default per-token slice approach at high concurrency and matches homogeneous TP baselines within ~5%.
+
+Enable the staging buffer when prefill and decode use **different TP sizes** with the **Mooncake** transfer backend. When both sides use the same TP size, staging is automatically bypassed even if enabled.
+
+> **Note:** The staging buffer is designed for non-MLA models (e.g. GQA, MHA). MLA models (e.g. DeepSeek-V2/V3) should not enable this flag.
+
+### Environment Variables
+
+| Variable | Description | Default |
+|:---------|:------------|:-------:|
+| **`SGLANG_DISAGG_STAGING_BUFFER`** | Enable GPU staging buffer for heterogeneous TP KV transfer | `False` |
+| **`SGLANG_DISAGG_STAGING_BUFFER_SIZE_MB`** | Prefill-side per-worker staging buffer size in MB | `64` |
+| **`SGLANG_DISAGG_STAGING_POOL_SIZE_MB`** | Decode-side ring buffer pool total size in MB | `4096` |
+
+### Usage Example
+
+```bash
+# Set staging buffer environment variables on BOTH prefill and decode
+export SGLANG_DISAGG_STAGING_BUFFER=1
+export SGLANG_DISAGG_STAGING_BUFFER_SIZE_MB=64
+export SGLANG_DISAGG_STAGING_POOL_SIZE_MB=4096
+
+# Prefill with TP=4
+python -m sglang.launch_server \
+  --model-path $MODEL_PATH \
+  --disaggregation-mode prefill \
+  --port 30000 \
+  --tp 4 \
+  --trust-remote-code \
+  --disaggregation-ib-device mlx5_1,mlx5_2
+
+# Decode with TP=1 (or DP attention with effective attention TP=1)
+python -m sglang.launch_server \
+  --model-path $MODEL_PATH \
+  --disaggregation-mode decode \
+  --port 30001 \
+  --tp 4 \
+  --dp 4 \
+  --enable-dp-attention \
+  --trust-remote-code \
+  --disaggregation-ib-device mlx5_3,mlx5_4
+
+# Router
+python -m sglang_router.launch_router \
+  --pd-disaggregation \
+  --prefill http://127.0.0.1:30000 \
+  --decode http://127.0.0.1:30001 \
+  --host 0.0.0.0 --port 8000
+```
+
 ## NIXL
 ### Requirements
 

docs/references/environment_variables.md

@@ -137,6 +137,15 @@ SGLang supports various environment variables that can be used to configure its
 | `SGLANG_PP_LAYER_PARTITION` | Pipeline parallel layer partition specification | Not set |
 | `SGLANG_ONE_VISIBLE_DEVICE_PER_PROCESS` | Set one visible device per process for distributed computing | `false` |
 
+## PD Disaggregation — Staging Buffer (Heterogeneous TP)
+
+| Environment Variable | Description | Default Value |
+| --- | --- | --- |
+| `SGLANG_DISAGG_STAGING_BUFFER` | Enable GPU staging buffer for heterogeneous TP KV transfer. Required when prefill and decode use different TP/attention-TP sizes. Only for non-MLA models (e.g. GQA, MHA). | `false` |
+| `SGLANG_DISAGG_STAGING_BUFFER_SIZE_MB` | Prefill-side per-worker staging buffer size in MB. Used for gathering KV head slices before bulk RDMA transfer. | `64` |
+| `SGLANG_DISAGG_STAGING_POOL_SIZE_MB` | Decode-side ring buffer pool total size in MB. Shared buffer receiving RDMA data from all prefill ranks. Larger values support higher concurrency. | `4096` |
+| `SGLANG_STAGING_USE_TORCH` | Force using PyTorch gather/scatter fallback instead of Triton fused kernels for staging operations. Useful for debugging. | `false` |
+
 ## Testing & Debugging (Internal/CI)
 
 *These variables are primarily used for internal testing, continuous integration, or debugging.*

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(
+        logger.info(
-            "No mooncake custom mem pool available for staging buffer. "
+            "Staging buffer using cudaMalloc (no custom mem pool). "
-            "NVLink transport will NOT work. Set SGLANG_MOONCAKE_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:

test/registered/distributed/test_disaggregation_different_tp.py

@@ -1,3 +1,4 @@
+import os
 import unittest
 from types import SimpleNamespace
 
@@ -318,5 +319,166 @@ class TestDisaggregationMooncakeMHADecodeLargerTP(PDDisaggregationServerBase):
         self.assertGreater(metrics["score"], 0.60)
 
 
+STAGING_ENV = {
+    "SGLANG_DISAGG_STAGING_BUFFER": "1",
+    "SGLANG_DISAGG_STAGING_BUFFER_SIZE_MB": "64",
+    "SGLANG_DISAGG_STAGING_POOL_SIZE_MB": "1024",
+}
+
+
+class TestDisaggregationStagingPrefillLargerTP(PDDisaggregationServerBase):
+    """Prefill TP=4 -> Decode TP=2 with staging buffer enabled (MHA model)."""
+
+    @classmethod
+    def setUpClass(cls):
+        super().setUpClass()
+        envs.SGLANG_ENABLE_JIT_DEEPGEMM.set(False)
+
+        cls.model = try_cached_model(DEFAULT_MODEL_NAME_FOR_TEST)
+
+        cls.start_prefill()
+        cls.start_decode()
+
+        cls.wait_server_ready(cls.prefill_url + "/health", process=cls.process_prefill)
+        cls.wait_server_ready(cls.decode_url + "/health", process=cls.process_decode)
+
+        cls.launch_lb()
+
+    @classmethod
+    def start_prefill(cls):
+        prefill_args = [
+            "--trust-remote-code",
+            "--disaggregation-mode",
+            "prefill",
+            "--disaggregation-bootstrap-port",
+            cls.bootstrap_port,
+            "--tp",
+            "4",
+        ]
+        prefill_args += cls.transfer_backend + cls.rdma_devices
+        env = {**os.environ, **STAGING_ENV}
+        cls.process_prefill = popen_launch_pd_server(
+            cls.model,
+            cls.prefill_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            other_args=prefill_args,
+            env=env,
+        )
+
+    @classmethod
+    def start_decode(cls):
+        decode_args = [
+            "--trust-remote-code",
+            "--disaggregation-mode",
+            "decode",
+            "--disaggregation-bootstrap-port",
+            cls.bootstrap_port,
+            "--tp",
+            "2",
+            "--base-gpu-id",
+            "4",
+        ]
+        decode_args += cls.transfer_backend + cls.rdma_devices
+        env = {**os.environ, **STAGING_ENV}
+        cls.process_decode = popen_launch_pd_server(
+            cls.model,
+            cls.decode_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            other_args=decode_args,
+            env=env,
+        )
+
+    def test_gsm8k(self):
+        args = SimpleNamespace(
+            base_url=self.base_url,
+            model=self.model,
+            eval_name="gsm8k",
+            api="completion",
+            max_tokens=512,
+            num_examples=200,
+            num_threads=128,
+        )
+        metrics = run_eval(args)
+        print(f"[Staging PrefillLargerTP] Evaluation metrics: {metrics}")
+        self.assertGreater(metrics["score"], 0.60)
+
+
+class TestDisaggregationStagingDecodeLargerTP(PDDisaggregationServerBase):
+    """Prefill TP=2 -> Decode TP=4 with staging buffer enabled (MHA model)."""
+
+    @classmethod
+    def setUpClass(cls):
+        super().setUpClass()
+        envs.SGLANG_ENABLE_JIT_DEEPGEMM.set(False)
+
+        cls.model = try_cached_model(DEFAULT_MODEL_NAME_FOR_TEST)
+
+        cls.start_prefill()
+        cls.start_decode()
+
+        cls.wait_server_ready(cls.prefill_url + "/health", process=cls.process_prefill)
+        cls.wait_server_ready(cls.decode_url + "/health", process=cls.process_decode)
+
+        cls.launch_lb()
+
+    @classmethod
+    def start_prefill(cls):
+        prefill_args = [
+            "--trust-remote-code",
+            "--disaggregation-mode",
+            "prefill",
+            "--disaggregation-bootstrap-port",
+            cls.bootstrap_port,
+            "--tp",
+            "2",
+        ]
+        prefill_args += cls.transfer_backend + cls.rdma_devices
+        env = {**os.environ, **STAGING_ENV}
+        cls.process_prefill = popen_launch_pd_server(
+            cls.model,
+            cls.prefill_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            other_args=prefill_args,
+            env=env,
+        )
+
+    @classmethod
+    def start_decode(cls):
+        decode_args = [
+            "--trust-remote-code",
+            "--disaggregation-mode",
+            "decode",
+            "--disaggregation-bootstrap-port",
+            cls.bootstrap_port,
+            "--tp",
+            "4",
+            "--base-gpu-id",
+            "4",
+        ]
+        decode_args += cls.transfer_backend + cls.rdma_devices
+        env = {**os.environ, **STAGING_ENV}
+        cls.process_decode = popen_launch_pd_server(
+            cls.model,
+            cls.decode_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            other_args=decode_args,
+            env=env,
+        )
+
+    def test_gsm8k(self):
+        args = SimpleNamespace(
+            base_url=self.base_url,
+            model=self.model,
+            eval_name="gsm8k",
+            api="completion",
+            max_tokens=512,
+            num_examples=200,
+            num_threads=128,
+        )
+        metrics = run_eval(args)
+        print(f"[Staging DecodeLargerTP] Evaluation metrics: {metrics}")
+        self.assertGreater(metrics["score"], 0.60)
+
+
 if __name__ == "__main__":
     unittest.main()