8cfdb0466e
Per-layer prefill-to-decode transfer now finishes per request/rank/chunk, so success-path INFO logs can dominate production logs and hide actual failures. Keep successful finish breakdown and completion messages at DEBUG while preserving nonzero finish status as WARNING. Constraint: Per-layer transfer is a hot path under CP shared-KV and may produce many batch completions per request. Rejected: Disable CP per-layer transfer logging entirely | failures still need visible warning-level evidence. Confidence: high Scope-risk: narrow Directive: Do not promote successful per-request transfer completion logs back to INFO without rate limiting. Tested: PYTHONPATH=python python -m pytest -q test/registered/unit/disaggregation/test_cp_per_layer_transfer.py::TestPerLayerTransferContext::test_successful_finish_does_not_emit_hot_path_info_log Tested: python -m py_compile python/sglang/srt/disaggregation/cp_per_layer_transfer.py python/sglang/srt/disaggregation/mooncake/conn.py Not-tested: Full local disaggregation suite blocked by missing local orjson dependency. Co-authored-by: OmX <omx@oh-my-codex.dev>
482 lines
18 KiB
Python
482 lines
18 KiB
Python
"""Unit tests for A2 core: PerLayerTransferContext (per-request overlapped transfer).
|
|
|
|
Key correctness properties:
|
|
- submit_layer waits the layer's CUDA write event BEFORE issuing the RDMA (so the
|
|
transfer never reads a layer before its write kernel finished);
|
|
- idempotent per layer; skips layers with no owned blocks;
|
|
- a submit failure marks the context failed and stops further submits;
|
|
- finish waits all accumulated batch_ids and returns 0 / non-zero correctly.
|
|
|
|
Engine and CUDA events are mocked — no GPU/RDMA.
|
|
"""
|
|
import unittest
|
|
from unittest import mock
|
|
|
|
from sglang.srt.disaggregation.cp_per_layer_transfer import PerLayerTransferContext
|
|
|
|
|
|
class _FakeEvent:
|
|
def __init__(self):
|
|
self.synced = False
|
|
|
|
def synchronize(self):
|
|
self.synced = True
|
|
|
|
|
|
class _FakeEngine:
|
|
def __init__(self, submit_rets=None, wait_ret=0):
|
|
self.submits = []
|
|
self.waits = []
|
|
self._rets = submit_rets
|
|
self._wait = wait_ret
|
|
self._n = 0
|
|
self.submit_order = [] # (layer-derived addr, whether event was synced first)
|
|
|
|
def batch_transfer_async_submit(self, sid, src, dst, lens):
|
|
self.submits.append((sid, list(src), list(dst), list(lens)))
|
|
if self._rets is not None:
|
|
r = self._rets[self._n]
|
|
self._n += 1
|
|
return r
|
|
return 100 + len(self.submits)
|
|
|
|
def wait_batch_transfers(self, ids):
|
|
self.waits.append(list(ids))
|
|
return self._wait
|
|
|
|
|
|
def _blocks(layer_id):
|
|
base = 1000 * layer_id
|
|
return ([base], [base + 5000], [64])
|
|
|
|
|
|
def _ctx(engine, blocks=_blocks, num_layers=8):
|
|
return PerLayerTransferContext(engine, "sess", blocks, num_layers=num_layers)
|
|
|
|
|
|
class TestPerLayerTransferContext(unittest.TestCase):
|
|
def test_submit_waits_event_then_submits(self):
|
|
eng = _FakeEngine()
|
|
ev = _FakeEvent()
|
|
_ctx(eng).submit_layer(3, ev)
|
|
self.assertTrue(ev.synced) # write event waited BEFORE the RDMA submit
|
|
self.assertEqual(eng.submits, [("sess", [3000], [8000], [64])])
|
|
|
|
def test_idempotent_per_layer(self):
|
|
eng = _FakeEngine()
|
|
ctx = _ctx(eng)
|
|
ctx.submit_layer(0, _FakeEvent())
|
|
ctx.submit_layer(0, _FakeEvent())
|
|
self.assertEqual(len(eng.submits), 1)
|
|
|
|
def test_skip_layer_with_no_blocks(self):
|
|
eng = _FakeEngine()
|
|
|
|
def blocks(L):
|
|
return None if L == 1 else ([L], [L + 1], [8])
|
|
|
|
ctx = _ctx(eng, blocks)
|
|
for L in range(3):
|
|
ctx.submit_layer(L, None)
|
|
self.assertEqual(len(eng.submits), 2) # layer 1 (no blocks) skipped
|
|
|
|
def test_finish_waits_all_and_returns_zero(self):
|
|
eng = _FakeEngine()
|
|
ctx = _ctx(eng, num_layers=3)
|
|
for L in range(3):
|
|
ctx.submit_layer(L, _FakeEvent())
|
|
self.assertEqual(ctx.finish(), 0)
|
|
self.assertEqual(eng.waits, [[101, 102, 103]])
|
|
|
|
def test_submit_failure_marks_failed_and_stops(self):
|
|
eng = _FakeEngine(submit_rets=[101, -1, 103])
|
|
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 submitting after the -1
|
|
self.assertEqual(ctx.finish(), -1)
|
|
|
|
def test_finish_propagates_wait_failure(self):
|
|
eng = _FakeEngine(wait_ret=-9)
|
|
ctx = _ctx(eng, num_layers=1)
|
|
ctx.submit_layer(0, None)
|
|
self.assertEqual(ctx.finish(), -9)
|
|
|
|
def test_repeated_enqueues_dedup_so_finish_cannot_hang(self):
|
|
# the notifier re-fires every layer on every subsequent forward while the ctx
|
|
# is still active; note_enqueued must dedup so target stays capped at the layer
|
|
# count and finish() completes (the high-cache-hit hang was target>>layers).
|
|
eng = _FakeEngine()
|
|
ctx = _ctx(eng, num_layers=3)
|
|
for _forward in range(5): # 5 forwards all re-fire layers 0,1,2
|
|
for L in range(3):
|
|
self.assertEqual(ctx.note_enqueued(L), _forward == 0) # True only 1st
|
|
if _forward == 0:
|
|
ctx.submit_layer(L, None)
|
|
self.assertEqual(len(eng.submits), 3) # 3 unique despite 15 note_enqueued
|
|
self.assertEqual(ctx.finish(timeout=0.5), 0) # completes, no 30s timeout
|
|
|
|
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 -> 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 -> 1 submit
|
|
self.assertEqual(ctx.finish(timeout=0.5), 0)
|
|
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_successful_finish_does_not_emit_hot_path_info_log(self):
|
|
# Successful per-request transfer completion is hot-path and can fire once per
|
|
# request/rank/chunk. It must not emit INFO logs by default.
|
|
eng = _FakeEngine()
|
|
ctx = _ctx(eng, num_layers=1)
|
|
ctx.submit_layer(0, None)
|
|
with mock.patch(
|
|
"sglang.srt.disaggregation.cp_per_layer_transfer.logger.isEnabledFor",
|
|
return_value=True,
|
|
), mock.patch(
|
|
"sglang.srt.disaggregation.cp_per_layer_transfer.logger.info"
|
|
) as info:
|
|
self.assertEqual(ctx.finish(timeout=0.5), 0)
|
|
info.assert_not_called()
|
|
|
|
def test_failed_submit_still_reported_after_fallback(self):
|
|
# layer 2 submit fails during the overlap; finish reports -1 even though it
|
|
# also runs the fallback for any unfired layers.
|
|
eng = _FakeEngine(submit_rets=[101, -1, 103])
|
|
ctx = _ctx(eng, num_layers=3)
|
|
for L in range(3):
|
|
ctx.submit_layer(L, None)
|
|
self.assertEqual(ctx.finish(timeout=0.5), -1)
|
|
|
|
|
|
class TestBuildLayerBlocks(unittest.TestCase):
|
|
def test_addresses_and_lengths(self):
|
|
from sglang.srt.disaggregation.cp_per_layer_transfer import build_layer_blocks
|
|
|
|
src, dst, lens = build_layer_blocks(
|
|
1000, 5000, 64, [[0, 1, 2], [5, 6]], [[10, 11, 12], [20, 21]]
|
|
)
|
|
self.assertEqual(src, [1000 + 0 * 64, 1000 + 5 * 64])
|
|
self.assertEqual(dst, [5000 + 10 * 64, 5000 + 20 * 64])
|
|
self.assertEqual(lens, [64 * 3, 64 * 2]) # item_len * run length
|
|
|
|
def test_empty(self):
|
|
from sglang.srt.disaggregation.cp_per_layer_transfer import build_layer_blocks
|
|
|
|
self.assertEqual(build_layer_blocks(1, 2, 8, [], []), ([], [], []))
|
|
|
|
def test_only_base_ptr_changes_across_layers(self):
|
|
from sglang.srt.disaggregation.cp_per_layer_transfer import build_layer_blocks
|
|
|
|
s0, d0, l0 = build_layer_blocks(1000, 2000, 64, [[3]], [[7]])
|
|
s1, d1, l1 = build_layer_blocks(9000, 8000, 64, [[3]], [[7]])
|
|
self.assertEqual(s0, [1000 + 3 * 64])
|
|
self.assertEqual(s1, [9000 + 3 * 64])
|
|
self.assertEqual(l0, l1) # lengths identical across layers (the invariant)
|
|
|
|
|
|
class _MockCtx:
|
|
num_layers = 8
|
|
|
|
def __init__(self):
|
|
self.submitted = [] # (start, end, event)
|
|
self.enqueued = 0
|
|
self.finished = False
|
|
self.failed = False
|
|
self._enq = set()
|
|
|
|
def note_enqueued(self, key):
|
|
if key in self._enq:
|
|
return False
|
|
self._enq.add(key)
|
|
self.enqueued += 1
|
|
return True
|
|
|
|
def submit_group(self, start, end, event):
|
|
self.submitted.append((start, end, event))
|
|
|
|
def finish(self):
|
|
self.finished = True
|
|
return 0
|
|
|
|
def mark_failed(self):
|
|
self.failed = True
|
|
|
|
|
|
class _RecEvent:
|
|
def __init__(self):
|
|
self.recorded_stream = "UNSET"
|
|
self.synced = False
|
|
|
|
def record(self, stream=None):
|
|
self.recorded_stream = stream
|
|
|
|
def synchronize(self):
|
|
self.synced = True
|
|
|
|
|
|
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",
|
|
group_size=group_size,
|
|
)
|
|
|
|
|
|
def _drain(q):
|
|
import queue as _q
|
|
|
|
items = []
|
|
while True:
|
|
try:
|
|
items.append(q.get(block=False))
|
|
except _q.Empty:
|
|
break
|
|
return items
|
|
|
|
|
|
class TestPerLayerTransferManager(unittest.TestCase):
|
|
def test_on_layer_end_enqueues_per_active_ctx_with_recorded_event(self):
|
|
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, start, end, ev in items:
|
|
self.assertEqual((start, end), (5, 5))
|
|
self.assertEqual(ev.recorded_stream, "STREAM")
|
|
|
|
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_chunked_fifo_separate_contexts_and_dedup(self):
|
|
m = _manager(group_size=1)
|
|
c1, c2 = _MockCtx(), _MockCtx()
|
|
m.register("r1", c1, chunk_key=0) # chunk 1
|
|
m.register("r1", c1, chunk_key=0) # SAME chunk re-register -> deduped
|
|
m.register("r1", c2, chunk_key=64) # chunk 2 (different start_send_idx)
|
|
self.assertTrue(m.has_chunk("r1", 0))
|
|
self.assertTrue(m.has_chunk("r1", 64))
|
|
self.assertFalse(m.has_chunk("r1", 128))
|
|
m.on_layer_end(5) # both chunk contexts active -> both enqueued
|
|
self.assertEqual(len(_drain(m._q)), 2)
|
|
# finish pops FIFO: chunk 1 (c1) first, then chunk 2 (c2)
|
|
self.assertEqual(m.finish("r1"), 0)
|
|
self.assertTrue(c1.finished)
|
|
self.assertFalse(c2.finished)
|
|
self.assertEqual(m.finish("r1"), 0)
|
|
self.assertTrue(c2.finished)
|
|
self.assertFalse(m.has_room("r1"))
|
|
|
|
def test_worker_step_calls_submit_group(self):
|
|
m = _manager()
|
|
c = _MockCtx()
|
|
ev = _RecEvent()
|
|
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()
|
|
|
|
class _Boom:
|
|
def __init__(self):
|
|
self.failed = False
|
|
|
|
def submit_group(self, start, end, event):
|
|
raise RuntimeError("boom")
|
|
|
|
def mark_failed(self):
|
|
self.failed = True
|
|
|
|
b = _Boom()
|
|
m._worker_step((b, 0, 0, None))
|
|
self.assertTrue(b.failed)
|
|
|
|
def test_finish_pops_and_calls_ctx_finish(self):
|
|
m = _manager()
|
|
c = _MockCtx()
|
|
m.register("r1", c)
|
|
self.assertTrue(m.has_active())
|
|
self.assertEqual(m.finish("r1"), 0)
|
|
self.assertTrue(c.finished)
|
|
self.assertFalse(m.has_active())
|
|
self.assertEqual(m.finish("r1"), 0) # idempotent after pop
|
|
|
|
def test_on_layer_end_no_active_is_noop(self):
|
|
m = _manager(group_size=1)
|
|
m.on_layer_end(0)
|
|
self.assertEqual(_drain(m._q), [])
|
|
|
|
|
|
class TestChunkedDstMapping(unittest.TestCase):
|
|
"""Regression for the chunked per-layer dst-offset bug.
|
|
|
|
register_per_layer_transfer maps a chunk's owned source pages to decode pages
|
|
via dst_kv_indices[positions], where positions come from the CP filter built on
|
|
chunk_page_start. The monolithic send() uses chunk_page_start=index_slice.start
|
|
(the absolute, cumulative page offset) and the worker indexes the FULL-request
|
|
dst_kv_indices by those absolute positions. So a per-chunk registration MUST pass
|
|
chunk_page_start = start_send_idx // page_size; passing 0 for every chunk makes
|
|
chunk N>0's positions chunk-local and silently writes its KV onto chunk 0's
|
|
decode pages (observed as corrupted decode output)."""
|
|
|
|
def _layout(self, cp_rank, cp_size=8, page_size=64):
|
|
from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
|
|
|
|
return CpSharedKVLayout(
|
|
page_size=page_size, cp_size=cp_size, cp_rank=cp_rank
|
|
)
|
|
|
|
def _map_chunk(self, layout, pages, dst_kv_indices, chunk_page_start):
|
|
"""Mirror register_per_layer_transfer's source->dst mapping for one chunk."""
|
|
import numpy as np
|
|
|
|
from sglang.srt.disaggregation.utils import (
|
|
filter_kv_pages_for_cp_shared_kv,
|
|
)
|
|
|
|
owned, positions = filter_kv_pages_for_cp_shared_kv(
|
|
layout=layout,
|
|
logical_pages=np.asarray(pages, dtype=np.int32),
|
|
chunk_page_start=chunk_page_start,
|
|
)
|
|
dst = np.asarray(dst_kv_indices, dtype=np.int32)[positions]
|
|
return list(owned), list(dst)
|
|
|
|
def test_correct_offset_matches_whole_request_per_rank(self):
|
|
import numpy as np
|
|
|
|
page_size = 64
|
|
cp_size = 8
|
|
total_pages = 24 # logical pages 0..23
|
|
all_pages = np.arange(total_pages, dtype=np.int32)
|
|
# distinct decode page per absolute logical position
|
|
dst_kv_indices = (90000 + np.arange(total_pages)).astype(np.int32)
|
|
# 2 chunks split at logical page 16 (token 1024 with page_size 64)
|
|
chunk0 = all_pages[:16]
|
|
chunk1 = all_pages[16:]
|
|
for cp_rank in range(cp_size):
|
|
layout = self._layout(cp_rank, cp_size, page_size)
|
|
# whole-request ground truth (single shot, offset 0)
|
|
whole_src, whole_dst = self._map_chunk(
|
|
layout, all_pages, dst_kv_indices, chunk_page_start=0
|
|
)
|
|
# chunked with CORRECT absolute offsets -> identical pairing
|
|
s0, d0 = self._map_chunk(layout, chunk0, dst_kv_indices, 0)
|
|
s1, d1 = self._map_chunk(layout, chunk1, dst_kv_indices, 16)
|
|
self.assertEqual(whole_src, s0 + s1, f"src mismatch rank={cp_rank}")
|
|
self.assertEqual(whole_dst, d0 + d1, f"dst mismatch rank={cp_rank}")
|
|
|
|
def test_zero_offset_corrupts_second_chunk(self):
|
|
"""The pre-fix behaviour: chunk_page_start=0 for chunk 1 maps its owned
|
|
pages onto chunk 0's decode region (dst < the chunk's true start)."""
|
|
import numpy as np
|
|
|
|
page_size = 64
|
|
cp_size = 8
|
|
total_pages = 24
|
|
all_pages = np.arange(total_pages, dtype=np.int32)
|
|
dst_kv_indices = (90000 + np.arange(total_pages)).astype(np.int32)
|
|
chunk1 = all_pages[16:]
|
|
# find a rank that actually owns at least one page in chunk 1
|
|
for cp_rank in range(cp_size):
|
|
layout = self._layout(cp_rank, cp_size, page_size)
|
|
_, d_bug = self._map_chunk(layout, chunk1, dst_kv_indices, 0)
|
|
_, d_fixed = self._map_chunk(layout, chunk1, dst_kv_indices, 16)
|
|
if not d_fixed:
|
|
continue
|
|
# fixed dst pages are all in chunk1's region (>= 90000+16)
|
|
self.assertTrue(all(x >= 90000 + 16 for x in d_fixed))
|
|
# buggy dst pages land in chunk0's region (< 90000+16) -> corruption
|
|
self.assertTrue(all(x < 90000 + 16 for x in d_bug))
|
|
self.assertNotEqual(d_bug, d_fixed)
|
|
return
|
|
self.fail("no rank owned a chunk-1 page; adjust fixture")
|
|
|
|
|
|
class TestRegisterGuardSingleInfo(unittest.TestCase):
|
|
"""register_per_layer_transfer must register a context only when there is
|
|
exactly one non-dummy decode info. The transfer worker calls finish() once
|
|
per non-dummy info for the room while the per-layer path registers a single
|
|
context, so >1 non-dummy info (decode attn_tp < prefill attn_tp) must fall
|
|
back to the monolithic transfer."""
|
|
|
|
def _fake_self(self, infos, recorder):
|
|
import types
|
|
|
|
mgr = types.SimpleNamespace(
|
|
register=lambda room, ctx, chunk_key=0: recorder.append(
|
|
(room, chunk_key)
|
|
),
|
|
)
|
|
return types.SimpleNamespace(
|
|
per_layer_transfer_manager=mgr,
|
|
server_args=types.SimpleNamespace(enable_nsa_prefill_cp_shared_kv=True),
|
|
transfer_infos={7: infos},
|
|
kv_args=types.SimpleNamespace(page_size=64),
|
|
attn_cp_size=8,
|
|
attn_cp_rank=0,
|
|
build_per_layer_context=lambda sid, owned, dst: ("ctx", sid),
|
|
)
|
|
|
|
def _info(self, is_dummy):
|
|
import types
|
|
|
|
import numpy as np
|
|
|
|
return types.SimpleNamespace(
|
|
is_dummy=is_dummy,
|
|
dst_kv_indices=np.arange(64, dtype=np.int32),
|
|
mooncake_session_id="s",
|
|
)
|
|
|
|
def _call(self, infos, recorder):
|
|
try:
|
|
from sglang.srt.disaggregation.mooncake.conn import MooncakeKVManager
|
|
except Exception as e: # pragma: no cover - env without mooncake
|
|
self.skipTest(f"mooncake conn not importable: {e}")
|
|
s = self._fake_self(infos, recorder)
|
|
# logical pages incl. page 1 which cp_rank 0 owns ((1-1)%8==0)
|
|
return MooncakeKVManager.register_per_layer_transfer(
|
|
s, 7, [0, 1, 2, 3], chunk_key=0
|
|
)
|
|
|
|
def test_two_non_dummy_infos_falls_back(self):
|
|
rec = []
|
|
infos = {"a": self._info(False), "b": self._info(False)}
|
|
self.assertFalse(self._call(infos, rec))
|
|
self.assertEqual(rec, []) # nothing registered
|
|
|
|
def test_single_non_dummy_info_registers(self):
|
|
rec = []
|
|
infos = {"a": self._info(True), "b": self._info(False)} # 1 dummy + 1 real
|
|
self.assertTrue(self._call(infos, rec))
|
|
self.assertEqual(len(rec), 1)
|
|
|
|
def test_all_dummy_falls_back(self):
|
|
rec = []
|
|
infos = {"a": self._info(True), "b": self._info(True)}
|
|
self.assertFalse(self._call(infos, rec))
|
|
self.assertEqual(rec, [])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|