feat(disagg/mooncake): async per-layer KV transfer primitive (lever A foundation)

Add _transfer_layers_async: submit each layer's transfer non-blocking via the
async API (batch_transfer_async_submit), pipelining layers in the RDMA engine,
then wait for all once (wait_batch_transfers). Gated by the new
SGLANG_CP_SHARED_KV_PER_LAYER_TRANSFER env (default off); replaces the monolithic
all-layers batch_transfer_sync on that path. This is the transfer mechanism for
per-layer overlap (lever A) and removes the per-layer blocking-sync tax measured
in B1a; the forward-overlap hook (G2) builds on it next. Uses the safe async API,
never the OnCuda busy-wait/_exit path.

Unit-tested (test_per_layer_transfer.py, 5 cases): one submit per non-empty layer,
single wait-for-all, empty-layer skip, submit-failure drain + return -1, and
wait-status propagation.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

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

@@ -268,6 +268,9 @@ class MooncakeKVManager(CommonKVManager):
             self.enable_custom_mem_pool, self.custom_mem_pool_type = (
                 check_mooncake_custom_mem_pool_enabled()
             )
+            self.enable_per_layer_async_transfer = (
+                envs.SGLANG_CP_SHARED_KV_PER_LAYER_TRANSFER.get()
+            )
         elif self.disaggregation_mode == DisaggregationMode.DECODE:
             self.start_decode_thread()
 
@@ -321,6 +324,31 @@ class MooncakeKVManager(CommonKVManager):
             mooncake_session_id, list(src_addrs), list(dst_addrs), list(lengths)
         )
 
+    def _transfer_layers_async(
+        self, mooncake_session_id, layers_params, set_transfer_blocks
+    ):
+        # Submit each layer's transfer non-blocking (pipelined in the RDMA engine),
+        # then wait for all to complete once. Removes the per-layer blocking-sync tax
+        # (B1a) and is the transfer mechanism for per-layer overlap (lever A). Uses the
+        # safe async API (G1: batch_transfer_async_submit + wait_batch_transfers),
+        # never the OnCuda busy-wait/_exit path.
+        batch_ids = []
+        for src_ptr, dst_ptr, item_len in layers_params:
+            transfer_blocks = set_transfer_blocks(src_ptr, dst_ptr, item_len)
+            if not transfer_blocks:
+                continue
+            src_addrs, dst_addrs, lengths = zip(*transfer_blocks)
+            batch_id = self.engine.batch_transfer_async_submit(
+                mooncake_session_id, list(src_addrs), list(dst_addrs), list(lengths)
+            )
+            if batch_id < 0:
+                # Submit failed: drain whatever was already submitted, then fail.
+                if batch_ids:
+                    self.engine.wait_batch_transfers(batch_ids)
+                return -1
+            batch_ids.append(batch_id)
+        return self.engine.wait_batch_transfers(batch_ids)
+
     def _send_kvcache_generic(
         self,
         mooncake_session_id: str,
@@ -418,7 +446,11 @@ class MooncakeKVManager(CommonKVManager):
             return self._transfer_data(mooncake_session_id, transfer_blocks)
 
         start_time = time.perf_counter() if transfer_stats_enabled else 0.0
-        if self.enable_custom_mem_pool:
+        if self.enable_per_layer_async_transfer:
+            status = self._transfer_layers_async(
+                mooncake_session_id, layers_params, set_transfer_blocks
+            )
+        elif self.enable_custom_mem_pool:
             futures = [
                 executor.submit(
                     process_layer,

python/sglang/srt/environ.py

@@ -211,6 +211,10 @@ class Envs:
     SGLANG_DEBUG_SORT_NVTX = EnvBool(False)
     SGLANG_DEBUG_MOE_SORT_NVTX = EnvBool(False)
     SGLANG_CP_SHARED_KV_CURRENT_REUSE = EnvBool(False)
+    # Per-layer overlapped KV transfer (lever A): submit each layer's KV transfer
+    # asynchronously (pipelined batch_transfer_async) instead of one monolithic
+    # all-layers batch_transfer_sync. Default off until validated.
+    SGLANG_CP_SHARED_KV_PER_LAYER_TRANSFER = EnvBool(False)
     SGLANG_CP_SHARED_KV_USE_TAI_MATERIALIZE = EnvBool(False)
     SGLANG_CP_SHARED_KV_FUSED_MLA_STORE = EnvBool(False)
     SGLANG_CP_SHARED_KV_FUSED_INDEX_MQA_PREPARE = EnvBool(False)