chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source 
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>
This commit is contained in:
leavelet
2026-06-07 09:51:04 +00:00
co-authored by Claude Opus 4.8
parent 1c6279eeac
commit a09aec30a3
3 changed files with 128 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/environ.py
+4
View File
@@ -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)