Stabilize disaggregated decode burst handoff
Burst decode can sit in disaggregation prealloc/transfer queues while health probes still need an immediate scheduler-alive response, and EAGLE prebuilt metadata must outlive the synthetic prebuilt step until the first real decode consume. The transfer path also needs to send final aux metadata even when there are no new KV pages in the final chunk. This keeps decode metadata slot ownership narrow instead of cloning EAGLE tensors, adds fail-safe cleanup for prebuilt exceptions/finished prebuilt requests, fixes final empty chunk metadata transfer, and records the investigation ledger for future debugging. Constraint: bs=1 historically worked, so decode compute/sampling behavior must not be changed without direct evidence. Rejected: Clone EAGLE metadata tensors on transfer commit | avoids lifetime issues but adds hot-path CPU/GPU memory traffic. Rejected: Treat prefill AbortReq logs as root cause | current evidence shows they can be downstream of decode/router aborts. Confidence: medium Scope-risk: moderate Directive: Do not move EAGLE metadata slot release earlier than first real decode result processing without proving the H2D/spec_info consume point changed. Tested: g0034 docker py_compile for touched scheduler/disagg/test files; PYTHONPATH=python python -m pytest -q test/registered/unit/disaggregation/test_decode_queue_compaction.py test/registered/unit/mem_cache/test_req_to_token_pool.py test/registered/unit/managers/test_scheduler_health_check.py -> 24 passed Not-tested: Fresh end-to-end burst traffic after restart; decode node_rank=1 persistent log capture.
This commit is contained in:
@@ -13,6 +13,7 @@ from sglang.srt.disaggregation.decode import (
|
||||
_kv_locs_to_page_indices_cpu,
|
||||
)
|
||||
from sglang.srt.disaggregation.prefill import (
|
||||
SchedulerDisaggregationPrefillMixin,
|
||||
_kv_locs_to_page_indices_cpu as _prefill_kv_locs_to_page_indices_cpu,
|
||||
)
|
||||
from sglang.srt.managers.schedule_batch import Req
|
||||
@@ -200,6 +201,61 @@ class TestDecodePreallocQueue(CustomTestCase):
|
||||
self.assertEqual(page_indices.dtype.name, "int32")
|
||||
self.assertEqual(page_indices.tolist(), [32, 33, 34])
|
||||
|
||||
|
||||
def test_prefill_final_empty_kv_chunk_still_sends_aux_metadata(self):
|
||||
class FakeAllocator:
|
||||
page_size = 4
|
||||
|
||||
def get_kvcache(self):
|
||||
return object()
|
||||
|
||||
class FakeMetadataBuffers:
|
||||
def __init__(self):
|
||||
self.requests = []
|
||||
|
||||
def set_buf(self, req):
|
||||
self.requests.append(req)
|
||||
|
||||
class FakeSender:
|
||||
def __init__(self):
|
||||
self.calls = []
|
||||
|
||||
def send(self, page_indices, state_indices):
|
||||
self.calls.append((page_indices, state_indices))
|
||||
|
||||
fake_scheduler = SimpleNamespace(
|
||||
token_to_kv_pool_allocator=FakeAllocator(),
|
||||
req_to_token_pool=SimpleNamespace(
|
||||
req_to_token=torch.arange(100, 112, dtype=torch.int64).reshape(1, -1)
|
||||
),
|
||||
disagg_metadata_buffers=FakeMetadataBuffers(),
|
||||
disagg_prefill_bootstrap_queue=SimpleNamespace(draft_token_to_kv_pool=None),
|
||||
)
|
||||
sender = FakeSender()
|
||||
req = SimpleNamespace(
|
||||
rid="r-final-empty",
|
||||
bootstrap_room=123,
|
||||
start_send_idx=8,
|
||||
fill_ids=list(range(8)),
|
||||
origin_input_ids=list(range(8)),
|
||||
req_pool_idx=0,
|
||||
disagg_kv_sender=sender,
|
||||
prefix_indices=[],
|
||||
host_hit_length=0,
|
||||
)
|
||||
|
||||
SchedulerDisaggregationPrefillMixin.send_kv_chunk(
|
||||
fake_scheduler, req, last_chunk=True
|
||||
)
|
||||
|
||||
self.assertEqual(req.start_send_idx, 8)
|
||||
self.assertEqual(fake_scheduler.disagg_metadata_buffers.requests, [req])
|
||||
self.assertEqual(len(sender.calls), 1)
|
||||
page_indices, state_indices = sender.calls[0]
|
||||
self.assertEqual(page_indices.dtype.name, "int32")
|
||||
self.assertEqual(page_indices.tolist(), [])
|
||||
self.assertIsNone(state_indices)
|
||||
|
||||
def test_prefill_kv_locs_to_page_indices_cpu_copies_only_page_starts(self):
|
||||
kv_locs = torch.tensor(
|
||||
[
|
||||
|
||||
Reference in New Issue
Block a user