Add an env var to allow transferring small metadata via TCP for PD (#16951)

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

@@ -549,7 +549,9 @@ class MooncakeKVManager(CommonKVManager):
         dst_aux_ptrs: list[int],
     ):
         # TODO(shangming): Fix me when nvlink_transport of Mooncake is bug-free
-        if self.enable_custom_mem_pool and self.custom_mem_pool_type == "NVLINK":
+        if (
+            self.enable_custom_mem_pool and self.custom_mem_pool_type == "NVLINK"
+        ) or envs.SGLANG_MOONCAKE_SEND_AUX_TCP.get():
             return self.send_aux_tcp(req, prefill_aux_index, dst_aux_ptrs)
 
         transfer_blocks = []

python/sglang/srt/environ.py

@@ -262,6 +262,7 @@ class Envs:
     SGLANG_MOONCAKE_CUSTOM_MEM_POOL = EnvStr(None)
     ENABLE_ASCEND_TRANSFER_WITH_MOONCAKE = EnvBool(False)
     ASCEND_NPU_PHY_ID = EnvInt(-1)
+    SGLANG_MOONCAKE_SEND_AUX_TCP = EnvBool(False)
 
     # Mooncake Store
     SGLANG_HICACHE_MOONCAKE_CONFIG_PATH = EnvStr(None)