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

feat(disagg): wire per-layer overlapped transfer end-to-end (lever A, A3-step3)

Browse Source 
Complete the lever-A hot-path integration behind SGLANG_CP_SHARED_KV_PER_LAYER_TRANSFER:
- PerLayerTransferContext: add num_layers + a completion event so finish() waits
  until ALL layers are processed before wait_batch_transfers (never races ahead of
  the worker threads and silently drops in-flight layers); times out to FAILURE.
- PerLayerTransferManager: has_room (for the swap) + drop (abort/failure drain so
  outstanding RDMA finishes before pages are reclaimed).
- MooncakeKVManager.register_per_layer_transfer: build + register a context before
  the forward, reusing send()'s exact CP filter (no re-derivation).
- transfer_worker: when a room is per-layer-active, wait those transfers (finish)
  instead of the monolithic send_kvcache -- no double-send; aux/state/completion
  unchanged. The skip path drops the context on abort/failure.
- prefill scheduler: _register_per_layer_transfers(batch) before run_batch, scoped
  (first impl) to single-forward, no-cached-prefix requests (the notifier transfers
  forward-written pages; chunked/cached-prefix are HiCache-loaded -> lever B).

Unit-tested (25 cases). e2e output-equality + TTFT verification next.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
孙远航
2026-06-07 01:01:40 +00:00
parent aa6acc9485
commit ae18e3adc8
4 changed files with 179 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files

80
python/sglang/srt/disaggregation/cp_per_layer_transfer.py
View File

@@ -50,56 +50,84 @@ def build_layer_blocks(
return src_addrs, dst_addrs, lengths
class PerLayerTransferContext:
"""Tracks the per-layer async transfers for one prefill request."""
# finish() default cap waiting for the forward's per-layer events to all be processed.
# Post-forward all events are already enqueued, so this is normally hit in <1s; the cap
# is only a safety net against a lost notifier (treated as failure, never silent loss).
_FINISH_TIMEOUT_S = 30.0
def __init__(self, engine, session_id: str, get_blocks: BlocksFn):
class PerLayerTransferContext:
"""Tracks the per-layer async transfers for one prefill request.
num_layers = the number of per-layer notifier fires expected (this pp stage's local
layer count). finish() blocks until all num_layers layers have been processed
(submitted or skipped) before waiting the batch_ids — otherwise it could race ahead
of the worker threads and silently miss layers still in flight (corrupt KV)."""
def __init__(self, engine, session_id: str, get_blocks: BlocksFn, num_layers: int):
self.engine = engine
self.session_id = session_id
self.get_blocks = get_blocks
self.num_layers = num_layers
self._batch_ids: List[int] = []
self._submitted_layers: set[int] = set()
self._processed = 0
self._failed = False
self._lock = threading.Lock()
self._all_processed = threading.Event()
if num_layers <= 0:
self._all_processed.set()
@property
def failed(self) -> bool:
with self._lock:
return self._failed
def _bump_processed_locked(self) -> None:
self._processed += 1
if self._processed >= self.num_layers:
self._all_processed.set()
def submit_layer(self, layer_id: int, event) -> None:
"""Wait the layer-L write event, then async-submit layer L's transfer.
Called on a BACKGROUND worker thread (event.synchronize() must not run on
the compute/forward thread). Idempotent per layer; no-op after failure."""
the compute/forward thread). Idempotent per layer; counts every unique layer
toward completion even when skipped/failed so finish() can't hang."""
with self._lock:
if self._failed or layer_id in self._submitted_layers:
return
if layer_id in self._submitted_layers:
return # duplicate fire — don't double-count
self._submitted_layers.add(layer_id)
if self._failed:
self._bump_processed_locked() # already failing: count + skip submit
return
if event is not None:
event.synchronize() # wait the GPU write kernel for layer_id to finish
blocks = self.get_blocks(layer_id)
if not blocks:
return
src_addrs, dst_addrs, lengths = blocks
if not src_addrs:
return
batch_id = self.engine.batch_transfer_async_submit(
self.session_id, list(src_addrs), list(dst_addrs), list(lengths)
)
batch_id = None
if blocks and blocks[0]:
src_addrs, dst_addrs, lengths = blocks
batch_id = self.engine.batch_transfer_async_submit(
self.session_id, list(src_addrs), list(dst_addrs), list(lengths)
)
with self._lock:
if batch_id < 0:
if batch_id is not None and batch_id < 0:
self._failed = True
else:
elif batch_id is not None:
self._batch_ids.append(batch_id)
self._bump_processed_locked()
def finish(self) -> int:
"""Wait for all submitted layer transfers. Returns 0 on success, non-zero on
failure (a submit failed, or wait reported failure). Drains regardless."""
def finish(self, timeout: float = _FINISH_TIMEOUT_S) -> int:
"""Wait for ALL layers to be processed, then wait the submitted transfers.
Returns 0 on success; non-zero on failure (a submit failed, the wait reported
failure, or layers never completed within the timeout)."""
completed = self._all_processed.wait(timeout)
with self._lock:
failed = self._failed
batch_ids = list(self._batch_ids)
wait_status = self.engine.wait_batch_transfers(batch_ids)
return -1 if failed else wait_status
if failed or not completed:
return -1
return wait_status
def mark_failed(self) -> None:
with self._lock:
@@ -159,11 +187,23 @@ class PerLayerTransferManager:
for ctx in ctxs:
self._q.put((ctx, layer_id, event))
def has_room(self, room) -> bool:
with self._active_lock:
return room in self._active
def finish(self, room) -> int:
with self._active_lock:
ctx = self._active.pop(room, None)
return ctx.finish() if ctx is not None else 0
def drop(self, room) -> None:
"""Abort path: drain + discard a room's context so any outstanding RDMA
completes before its KV pages are freed (never write into reclaimed pages)."""
with self._active_lock:
ctx = self._active.pop(room, None)
if ctx is not None:
ctx.finish()
def _worker_step(self, item) -> None:
ctx, layer_id, event = item
try:

69
python/sglang/srt/disaggregation/mooncake/conn.py
View File

@@ -391,7 +391,52 @@ class MooncakeKVManager(CommonKVManager):
dst_kv_blocks,
)
return PerLayerTransferContext(self.engine, mooncake_session_id, get_blocks)
return PerLayerTransferContext(
self.engine, mooncake_session_id, get_blocks, num_layers=n_layers
)
def register_per_layer_transfer(self, room, page_indices) -> bool:
"""Lever A: before the forward, build + register a per-layer transfer context
for `room` so the per-layer notifier overlaps its main-KV transfer with the
forward. Reuses send()'s CP filter exactly (no re-derivation). Scoped to
CP-shared-KV (the target scenario). Returns True iff a context was registered;
a False just falls back to the monolithic post-forward transfer (still correct).
page_indices = the request's full new-token logical page ids (req_to_token)."""
mgr = getattr(self, "per_layer_transfer_manager", None)
if mgr is None or not self.server_args.enable_nsa_prefill_cp_shared_kv:
return False
infos = self.transfer_infos.get(room)
if not infos:
return False
from sglang.srt.disaggregation.utils import filter_kv_pages_for_cp_shared_kv
from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
layout = CpSharedKVLayout(
page_size=self.kv_args.page_size,
cp_size=self.attn_cp_size,
cp_rank=self.attn_cp_rank,
)
pages = np.asarray(page_indices, dtype=np.int32)
for info in infos.values():
if getattr(info, "is_dummy", False):
continue
owned_pages, positions = filter_kv_pages_for_cp_shared_kv(
layout=layout, logical_pages=pages, chunk_page_start=0
)
dst_indices = np.asarray(info.dst_kv_indices, dtype=np.int32)[positions]
ctx = self.build_per_layer_context(
info.mooncake_session_id, owned_pages, dst_indices
)
if ctx is not None:
mgr.register(room, ctx)
logger.info(
"[CP_PER_LAYER_TRANSFER] registered room=%s owned_pages=%d",
room,
len(owned_pages),
)
return True
return False
return False
def _send_kvcache_generic(
self,
@@ -1090,6 +1135,11 @@ class MooncakeKVManager(CommonKVManager):
"Skipping chunk for room %s because it has already failed or been aborted",
kv_chunk.room,
)
# Lever A: drain + drop any per-layer context for this room so an
# outstanding RDMA completes before its KV pages can be reclaimed.
per_layer_mgr = getattr(self, "per_layer_transfer_manager", None)
if per_layer_mgr is not None:
per_layer_mgr.drop(kv_chunk.room)
continue
reqs_to_be_processed = (
self.transfer_infos[kv_chunk.room].values()
@@ -1171,7 +1221,22 @@ class MooncakeKVManager(CommonKVManager):
target_rank_registration_info: KVArgsRegisterInfo = (
self.decode_kv_args_table[req.mooncake_session_id]
)
if self.is_mla_backend or (
per_layer_mgr = getattr(
self, "per_layer_transfer_manager", None
)
if per_layer_mgr is not None and per_layer_mgr.has_room(
kv_chunk.room
):
# Lever A: the main KV was transferred per-layer, overlapped
# with the forward; wait those transfers here instead of the
# monolithic send (no double-send). aux/state below unchanged.
ret = per_layer_mgr.finish(kv_chunk.room)
logger.info(
"[CP_PER_LAYER_TRANSFER] finished room=%s ret=%s",
kv_chunk.room,
ret,
)
elif self.is_mla_backend or (
self.attn_tp_size
== target_rank_registration_info.dst_attn_tp_size
):

27
python/sglang/srt/disaggregation/prefill.py
View File

@@ -647,6 +647,31 @@ class SchedulerDisaggregationPrefillMixin:
return batch
def _register_per_layer_transfers(self: Scheduler, batch) -> None:
"""Lever A: before run_batch, register a per-layer transfer context for each
eligible request so the per-layer notifier overlaps its main-KV transfer with
the forward. Scoped (first impl) to single-forward, no-cached-prefix requests:
the per-layer notifier transfers FORWARD-written pages, so chunked or cached-
prefix requests (whose prefix is HiCache-loaded, not forward-written) are left
on the monolithic post-forward path. No-op unless the flag is on."""
kv_manager = self.disagg_prefill_bootstrap_queue.kv_manager
if getattr(kv_manager, "per_layer_transfer_manager", None) is None:
return
page_size = self.token_to_kv_pool_allocator.page_size
for req in batch.reqs:
if getattr(req, "disagg_kv_sender", None) is None:
continue
if getattr(req, "is_chunked", 0) != 0 or req.start_send_idx != 0:
continue # chunked / partially-sent: this forward isn't the full range
if len(getattr(req, "prefix_indices", []) or []) != 0:
continue # cached prefix is HiCache-loaded, not forward-written (lever B)
end_idx = min(len(req.fill_ids), len(req.origin_input_ids))
page_indices = _kv_locs_to_page_indices_cpu(
self.req_to_token_pool.req_to_token[req.req_pool_idx, 0:end_idx],
page_size,
)
kv_manager.register_per_layer_transfer(req.bootstrap_room, page_indices)
@torch.no_grad()
def event_loop_normal_disagg_prefill(self: Scheduler) -> None:
"""A normal scheduler loop for prefill worker in disaggregation mode."""
@@ -718,6 +743,7 @@ class SchedulerDisaggregationPrefillMixin:
_seq_summary(getattr(batch, "prefix_lens", None)),
len(self.disagg_prefill_inflight_queue),
)
self._register_per_layer_transfers(batch)
result = self.run_batch(batch)
_cp_shared_kv_bs_gt1_prefill_timing(
"event_loop_run_batch_done",
@@ -783,6 +809,7 @@ class SchedulerDisaggregationPrefillMixin:
# Launch the current batch
if batch:
self._register_per_layer_transfers(batch)
batch_result = self.run_batch(batch)
self.result_queue.append((batch.copy(), batch_result))
else:

31
test/registered/unit/disaggregation/test_cp_per_layer_transfer.py
View File

@@ -49,8 +49,8 @@ def _blocks(layer_id):
return ([base], [base + 5000], [64])
def _ctx(engine, blocks=_blocks):
return PerLayerTransferContext(engine, "sess", blocks)
def _ctx(engine, blocks=_blocks, num_layers=8):
return PerLayerTransferContext(engine, "sess", blocks, num_layers=num_layers)
class TestPerLayerTransferContext(unittest.TestCase):
@@ -81,7 +81,7 @@ class TestPerLayerTransferContext(unittest.TestCase):
def test_finish_waits_all_and_returns_zero(self):
eng = _FakeEngine()
ctx = _ctx(eng)
ctx = _ctx(eng, num_layers=3)
for L in range(3):
ctx.submit_layer(L, _FakeEvent())
self.assertEqual(ctx.finish(), 0)
@@ -89,19 +89,38 @@ class TestPerLayerTransferContext(unittest.TestCase):
def test_submit_failure_marks_failed_and_stops(self):
eng = _FakeEngine(submit_rets=[101, -1, 103])
ctx = _ctx(eng)
ctx = _ctx(eng, num_layers=3)
for L in range(3):
ctx.submit_layer(L, None)
self.assertTrue(ctx.failed)
self.assertEqual(len(eng.submits), 2) # stopped after the -1
self.assertEqual(len(eng.submits), 2) # stopped submitting after the -1
self.assertEqual(ctx.finish(), -1)
def test_finish_propagates_wait_failure(self):
eng = _FakeEngine(wait_ret=-9)
ctx = _ctx(eng)
ctx = _ctx(eng, num_layers=1)
ctx.submit_layer(0, None)
self.assertEqual(ctx.finish(), -9)
def test_finish_times_out_to_failure_when_layers_missing(self):
# only 1 of 3 layers ever processed -> finish must NOT silently succeed
eng = _FakeEngine()
ctx = _ctx(eng, num_layers=3)
ctx.submit_layer(0, None)
self.assertEqual(ctx.finish(timeout=0.05), -1) # incomplete -> failure
def test_skipped_and_failed_layers_still_count_toward_completion(self):
# layer 1 has no blocks (skip), layer 2 submit fails; all 3 still complete
eng = _FakeEngine(submit_rets=[101, -1])
def blocks(L):
return None if L == 1 else ([L], [L + 10], [8])
ctx = _ctx(eng, blocks, num_layers=3)
for L in range(3):
ctx.submit_layer(L, None)
self.assertEqual(ctx.finish(timeout=0.5), -1) # completes (no hang), reports failure
class TestBuildLayerBlocks(unittest.TestCase):
def test_addresses_and_lengths(self):