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

perf(disagg): coarse-grained per-layer transfer + skip re-registration (lever A)

Browse Source 
For large bs (production target: bs~10 x ~100k tokens), per-layer granularity does
10x79 = 790 submitTransfer calls + CUDA events + enqueues per forward on the forward
thread. Two overhead cuts:
- Group SGLANG_CP_SHARED_KV_PER_LAYER_GROUP (default 8) consecutive layers into ONE
  RDMA submit: ~num_layers/K submits + events + enqueues instead of per-layer; same
  bytes (page index lists are identical across layers). on_layer_end is O(1) at
  non-boundary layers. The last partial group enqueues via the num_layers boundary;
  any misses fall back to one batched sync submit.
- Scheduler hook skips reqs already registered (bs>1 batch-forming re-iterates the
  same reqs ~9x -> was rebuilding the CP filter + context every time).

27 unit tests pass incl. grouping-boundary + batched-fallback.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
孙远航
2026-06-07 05:31:46 +00:00
parent 5fdf439c7e
commit e12afe8ced
4 changed files with 119 additions and 59 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -106,43 +106,52 @@ class PerLayerTransferContext:
self._enqueued += 1
return True
def _submit_one(self, layer_id: int) -> None:
"""Submit layer_id's transfer (KV assumed written). Updates state; logs +
marks failed on any error."""
batch_id = None
def _submit_layers(self, layer_ids) -> None:
"""Submit the COMBINED blocks of several (consecutive) layers as ONE RDMA
batch — coarser granularity means far fewer submitTransfer calls + batch_ids,
the key CPU/overhead reduction at large batch sizes. The page index lists are
identical across layers (only the per-layer base ptr differs), so this is the
same bytes as per-layer, just batched. Updates state; logs + fails closed."""
src_all: List[int] = []
dst_all: List[int] = []
len_all: List[int] = []
try:
blocks = self.get_blocks(layer_id)
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)
)
for layer_id in layer_ids:
blocks = self.get_blocks(layer_id)
if blocks and blocks[0]:
src_all.extend(blocks[0])
dst_all.extend(blocks[1])
len_all.extend(blocks[2])
except Exception as e:
logger.warning(
"[CP_PER_LAYER_TRANSFER] submit_layer %d failed: %r", layer_id, e
)
logger.warning("[CP_PER_LAYER_TRANSFER] get_blocks failed: %r", e)
with self._cond:
self._failed = True
return
if not src_all:
return
batch_id = self.engine.batch_transfer_async_submit(
self.session_id, src_all, dst_all, len_all
)
with self._cond:
if batch_id is not None and batch_id < 0:
if batch_id < 0:
self._failed = True
elif batch_id is not None:
else:
self._batch_ids.append(batch_id)
def submit_layer(self, layer_id: int, event) -> None:
"""Worker thread: wait the layer-L write event, then submit. Idempotent per
layer; always advances _processed (the finally) so finish() can't hang."""
def submit_group(self, start: int, end: int, event) -> None:
"""Worker thread: wait the group's last-layer write event, then submit layers
[start..end] as one RDMA batch. Idempotent; always advances _processed (the
finally) so finish() can't hang. (Per-layer is the special case start==end.)"""
with self._cond:
if layer_id in self._submitted_layers:
return # duplicate fire
self._submitted_layers.add(layer_id)
new = [L for L in range(start, end + 1) if L not in self._submitted_layers]
for L in new:
self._submitted_layers.add(L)
skip = self._failed
try:
if not skip:
if not skip and new:
if event is not None:
event.synchronize() # wait the GPU write kernel for layer_id
self._submit_one(layer_id)
event.synchronize() # last layer of the group -> all written
self._submit_layers(new)
finally:
with self._cond:
self._processed += 1
@@ -164,12 +173,13 @@ class PerLayerTransferContext:
submitted_at_entry = len(self._batch_ids)
t1 = time.perf_counter() # workers drained
if missing:
for layer_id in missing:
with self._cond:
if layer_id in self._submitted_layers:
continue
self._submitted_layers.add(layer_id)
self._submit_one(layer_id) # post-forward: written, no event needed
with self._cond:
still = [L for L in missing if L not in self._submitted_layers]
for L in still:
self._submitted_layers.add(L)
if still:
# post-forward: all written, no event needed; one batch for the tail
self._submit_layers(still)
t2 = time.perf_counter() # fallback submitted
with self._cond:
failed = failed or self._failed
@@ -192,6 +202,10 @@ class PerLayerTransferContext:
return -1
return wait_status
def submit_layer(self, layer_id: int, event) -> None:
"""Per-layer convenience (group of one). The manager uses submit_group."""
self.submit_group(layer_id, layer_id, event)
def mark_failed(self) -> None:
with self._cond:
self._failed = True
@@ -216,13 +230,20 @@ class PerLayerTransferManager:
"""
def __init__(
self, num_workers=4, event_factory=None, current_stream=None, worker_init=None
self,
num_workers=4,
event_factory=None,
current_stream=None,
worker_init=None,
group_size=8,
):
import queue as _queue
self._event_factory = event_factory
self._current_stream = current_stream
self._worker_init = worker_init # called once per worker thread (e.g. set CUDA device)
self._group_size = max(1, int(group_size)) # layers per RDMA submit (overhead)
self._num_layers = 0 # learned from the first registered ctx (model layer count)
self._q = _queue.SimpleQueue()
self._active = {}
self._active_lock = threading.Lock()
@@ -234,12 +255,22 @@ class PerLayerTransferManager:
if room in self._active:
return # already registered for its forward; don't overwrite (leak)
self._active[room] = ctx
if self._num_layers == 0:
self._num_layers = ctx.num_layers
def has_active(self) -> bool:
with self._active_lock:
return bool(self._active)
def on_layer_end(self, layer_id: int) -> None:
# Group K consecutive layers into ONE submit. Only do work at a group boundary
# (every K layers, or the last layer); other layers are O(1). This is the CPU/
# overhead reduction at large bs: ~num_layers/K submits + events + enqueues
# instead of one per layer.
K = self._group_size
is_last = self._num_layers > 0 and layer_id == self._num_layers - 1
if (layer_id + 1) % K != 0 and not is_last:
return # not a group boundary
with self._active_lock:
if not self._active:
return
@@ -252,9 +283,10 @@ class PerLayerTransferManager:
event.record(stream)
else:
event.record()
start = (layer_id // K) * K # first layer of this group
for ctx in ctxs:
if ctx.note_enqueued(layer_id): # enqueue each layer ONCE per ctx
self._q.put((ctx, layer_id, event))
if ctx.note_enqueued(layer_id): # one group per (ctx, boundary) — dedup
self._q.put((ctx, start, layer_id, event))
def has_room(self, room) -> bool:
with self._active_lock:
@@ -274,9 +306,9 @@ class PerLayerTransferManager:
ctx.finish()
def _worker_step(self, item) -> None:
ctx, layer_id, event = item
ctx, start, end, event = item
try:
ctx.submit_layer(layer_id, event)
ctx.submit_group(start, end, event)
except Exception:
ctx.mark_failed()

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

@@ -294,6 +294,7 @@ class PrefillBootstrapQueue:
event_factory=torch.cuda.Event,
current_stream=torch.cuda.current_stream,
worker_init=lambda: torch.cuda.set_device(device),
group_size=envs.SGLANG_CP_SHARED_KV_PER_LAYER_GROUP.get(),
)
pool = self.token_to_kv_pool
if hasattr(pool, "register_layer_backup_notifier"):
@@ -657,13 +658,18 @@ class SchedulerDisaggregationPrefillMixin:
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:
mgr = getattr(kv_manager, "per_layer_transfer_manager", None)
if mgr is None:
return
page_size = self.token_to_kv_pool_allocator.page_size
for req in batch.reqs:
try:
if getattr(req, "disagg_kv_sender", None) is None:
continue
room = getattr(req, "bootstrap_room", None)
if room is None or mgr.has_room(room):
continue # already registered (bs>1 batch-forming re-iterates the
# same reqs many times) — skip the CP filter + build entirely
if getattr(req, "is_chunked", 0) != 0 or getattr(req, "start_send_idx", 0) != 0:
continue # chunked / partially-sent: this forward isn't the full range
# Cover the FULL sequence (cached prefix + new tokens). The per-layer
@@ -677,7 +683,7 @@ class SchedulerDisaggregationPrefillMixin:
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)
kv_manager.register_per_layer_transfer(room, page_indices)
except Exception as e:
# Never let lever-A setup crash the scheduler; fall back to the
# monolithic post-forward transfer for this request.

4
python/sglang/srt/environ.py
View File

@@ -215,6 +215,10 @@ class Envs:
# 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)
# Lever A granularity: group this many consecutive layers into one RDMA submit.
# Larger = fewer submitTransfer calls / events / enqueues (less CPU overhead at
# large bs) but coarser overlap. 1 = per-layer.
SGLANG_CP_SHARED_KV_PER_LAYER_GROUP = EnvInt(8)
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)

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

@@ -118,13 +118,15 @@ class TestPerLayerTransferContext(unittest.TestCase):
def test_finish_falls_back_for_notifier_missed_layers(self):
# only layer 0 fired by the notifier; finish must SYNCHRONOUSLY transfer the
# missing layers 1,2 (e.g. an MTP buffer not hooked) -> all 3 moved, success.
# missing layers 1,2 -> all 3 moved, success. The fallback BATCHES the missing
# layers into one submit.
eng = _FakeEngine()
ctx = _ctx(eng, num_layers=3)
ctx.submit_layer(0, None) # notifier fired layer 0 only
ctx.submit_layer(0, None) # notifier fired layer 0 only -> 1 submit
self.assertEqual(ctx.finish(timeout=0.5), 0)
self.assertEqual(len(eng.submits), 3) # layer 0 (overlap) + 1,2 (fallback)
self.assertEqual(eng.waits[-1], [101, 102, 103])
self.assertEqual(len(eng.submits), 2) # layer 0 + fallback batch [1,2]
self.assertEqual(eng.waits[-1], [101, 102])
self.assertEqual(eng.submits[1], ("sess", [1000, 2000], [6000, 7000], [64, 64]))
def test_failed_submit_still_reported_after_fallback(self):
# layer 2 submit fails during the overlap; finish reports -1 even though it
@@ -163,21 +165,24 @@ class TestBuildLayerBlocks(unittest.TestCase):
class _MockCtx:
num_layers = 8
def __init__(self):
self.submitted = []
self.submitted = [] # (start, end, event)
self.enqueued = 0
self.finished = False
self.failed = False
self._enq = set()
def note_enqueued(self, layer_id):
if layer_id in getattr(self, "_enq_layers", set()):
def note_enqueued(self, key):
if key in self._enq:
return False
self._enq_layers = getattr(self, "_enq_layers", set()) | {layer_id}
self._enq.add(key)
self.enqueued += 1
return True
def submit_layer(self, layer_id, event):
self.submitted.append((layer_id, event))
def submit_group(self, start, end, event):
self.submitted.append((start, end, event))
def finish(self):
self.finished = True
@@ -199,11 +204,14 @@ class _RecEvent:
self.synced = True
def _manager(num_workers=0):
def _manager(num_workers=0, group_size=1):
from sglang.srt.disaggregation.cp_per_layer_transfer import PerLayerTransferManager
return PerLayerTransferManager(
num_workers=num_workers, event_factory=_RecEvent, current_stream=lambda: "STREAM"
num_workers=num_workers,
event_factory=_RecEvent,
current_stream=lambda: "STREAM",
group_size=group_size,
)
@@ -221,23 +229,33 @@ def _drain(q):
class TestPerLayerTransferManager(unittest.TestCase):
def test_on_layer_end_enqueues_per_active_ctx_with_recorded_event(self):
m = _manager()
m = _manager(group_size=1) # per-layer: every layer is a group boundary
c1, c2 = _MockCtx(), _MockCtx()
m.register("r1", c1)
m.register("r2", c2)
m.on_layer_end(5)
items = _drain(m._q)
self.assertEqual(len(items), 2) # one per active context
for ctx, layer_id, ev in items:
self.assertEqual(layer_id, 5)
self.assertEqual(ev.recorded_stream, "STREAM") # event recorded on compute stream
for ctx, start, end, ev in items:
self.assertEqual((start, end), (5, 5))
self.assertEqual(ev.recorded_stream, "STREAM")
def test_worker_step_calls_submit_layer(self):
def test_grouping_only_enqueues_at_boundaries(self):
m = _manager(group_size=4)
c = _MockCtx()
m.register("r1", c)
for L in range(8):
m.on_layer_end(L)
items = _drain(m._q)
# boundaries at L=3 (group [0,3]) and L=7 (group [4,7]) -> 2 groups, not 8
self.assertEqual([(s, e) for _, s, e, _ in items], [(0, 3), (4, 7)])
def test_worker_step_calls_submit_group(self):
m = _manager()
c = _MockCtx()
ev = _RecEvent()
m._worker_step((c, 3, ev))
self.assertEqual(c.submitted, [(3, ev)])
m._worker_step((c, 0, 3, ev))
self.assertEqual(c.submitted, [(0, 3, ev)])
def test_worker_step_marks_failed_on_exception(self):
m = _manager()
@@ -246,14 +264,14 @@ class TestPerLayerTransferManager(unittest.TestCase):
def __init__(self):
self.failed = False
def submit_layer(self, layer_id, event):
def submit_group(self, start, end, event):
raise RuntimeError("boom")
def mark_failed(self):
self.failed = True
b = _Boom()
m._worker_step((b, 0, None))
m._worker_step((b, 0, 0, None))
self.assertTrue(b.failed)
def test_finish_pops_and_calls_ctx_finish(self):
@@ -267,7 +285,7 @@ class TestPerLayerTransferManager(unittest.TestCase):
self.assertEqual(m.finish("r1"), 0) # idempotent after pop
def test_on_layer_end_no_active_is_noop(self):
m = _manager()
m = _manager(group_size=1)
m.on_layer_end(0)
self.assertEqual(_drain(m._q), [])