diff --git a/docs/advanced_features/pd_disaggregation.md b/docs/advanced_features/pd_disaggregation.md index 17b81b863..e1edc56b8 100644 --- a/docs/advanced_features/pd_disaggregation.md +++ b/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 diff --git a/docs/references/environment_variables.md b/docs/references/environment_variables.md index 1bfb7ce8f..e2e93b177 100644 --- a/docs/references/environment_variables.md +++ b/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.* diff --git a/python/sglang/srt/disaggregation/common/staging_buffer.py b/python/sglang/srt/disaggregation/common/staging_buffer.py index 4380c3ce7..b6fb6ba1a 100644 --- a/python/sglang/srt/disaggregation/common/staging_buffer.py +++ b/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( diff --git a/python/sglang/srt/disaggregation/common/staging_handler.py b/python/sglang/srt/disaggregation/common/staging_handler.py index b353f4609..c05e81511 100644 --- a/python/sglang/srt/disaggregation/common/staging_handler.py +++ b/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 = [] diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index 5c152feaa..ccf842ac0 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/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: diff --git a/python/sglang/srt/disaggregation/fake/conn.py b/python/sglang/srt/disaggregation/fake/conn.py index 03b79af18..60bf5465b 100644 --- a/python/sglang/srt/disaggregation/fake/conn.py +++ b/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: diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index 8eadf8195..8e3da245b 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/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: diff --git a/test/registered/distributed/test_disaggregation_different_tp.py b/test/registered/distributed/test_disaggregation_different_tp.py index bbf9f75aa..1bd1b5102 100644 --- a/test/registered/distributed/test_disaggregation_different_tp.py +++ b/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()