feat: WIP NIXL transfer through CPU buffer for better performance with mixed TP sizes

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

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING, List, Optional
+from typing import TYPE_CHECKING, Any, List, Optional
 
 import numpy as np
 import numpy.typing as npt
@@ -40,6 +40,10 @@ class KVArgs:
     prefill_start_layer: int
     # for system dp
     system_dp_rank: int
+    # Optional tensor buffer references for CPU buffer KV transfer
+    k_buffers: Optional[List[Any]] = None  # List[torch.Tensor], one per layer
+    v_buffers: Optional[List[Any]] = None  # List[torch.Tensor], one per layer
+    head_dim: Optional[int] = None
 
 
 class KVPoll:

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

@@ -410,9 +410,16 @@ def __init__(
             self.target_tp_ranks = [self.target_tp_rank]
         elif self.kv_mgr.attn_tp_size > self.prefill_info.attn_tp_size:
             if not self.kv_mgr.is_mla_backend:
-                logger.warning_once(
-                    "Performance is NOT guaranteed when using different TP sizes for non-MLA models. "
-                )
+                if getattr(self.kv_mgr, "nixl_use_cpu_buffer", False):
+                    logger.info_once(
+                        "Mixed TP sizes detected (decode_tp > prefill_tp). "
+                        "CPU buffer transfer (--nixl-use-cpu-buffer) is enabled for correct head redistribution."
+                    )
+                else:
+                    logger.warning_once(
+                        "Performance is NOT guaranteed when using different TP sizes for non-MLA models. "
+                        "Consider running with --nixl-use-cpu-buffer for correct mixed-TP KV transfer."
+                    )
             self.target_tp_rank = (
                 self.kv_mgr.kv_args.engine_rank % self.kv_mgr.attn_tp_size
             ) // (self.kv_mgr.attn_tp_size // self.prefill_info.attn_tp_size)
@@ -425,9 +432,16 @@ def __init__(
             self.target_tp_ranks = [self.target_tp_rank]
         else:
             if not self.kv_mgr.is_mla_backend:
-                logger.warning_once(
-                    "Performance is NOT guaranteed when using different TP sizes for non-MLA models. "
-                )
+                if getattr(self.kv_mgr, "nixl_use_cpu_buffer", False):
+                    logger.info_once(
+                        "Mixed TP sizes detected (prefill_tp > decode_tp). "
+                        "CPU buffer transfer (--nixl-use-cpu-buffer) is enabled for correct head redistribution."
+                    )
+                else:
+                    logger.warning_once(
+                        "Performance is NOT guaranteed when using different TP sizes for non-MLA models. "
+                        "Consider running with --nixl-use-cpu-buffer for correct mixed-TP KV transfer."
+                    )
             # For non-MLA models, one decode rank needs to retrieve KVCache from multiple prefill ranks for non MLA models;
             self.target_tp_ranks = [
                 rank

python/sglang/srt/disaggregation/decode.py

@@ -304,6 +304,15 @@ def _init_kv_manager(self) -> CommonKVManager:
         kv_args.kv_item_lens = kv_item_lens
         kv_args.page_size = self.token_to_kv_pool.page_size
 
+        # Add tensor buffer references for CPU buffer KV transfer
+        if getattr(self.scheduler.server_args, "nixl_use_cpu_buffer", False):
+            if hasattr(self.token_to_kv_pool, "k_buffer") and hasattr(
+                self.token_to_kv_pool, "v_buffer"
+            ):
+                kv_args.k_buffers = self.token_to_kv_pool.k_buffer
+                kv_args.v_buffers = self.token_to_kv_pool.v_buffer
+                kv_args.head_dim = self.token_to_kv_pool.head_dim
+
         kv_args.aux_data_ptrs, kv_args.aux_data_lens, kv_args.aux_item_lens = (
             self.metadata_buffers.get_buf_infos()
         )

python/sglang/srt/disaggregation/nixl/__init__.py

@@ -4,3 +4,4 @@
     NixlKVReceiver,
     NixlKVSender,
 )
+from sglang.srt.disaggregation.nixl.pinned_buffer_pool import PinnedBufferPool