Hold decode EAGLE metadata until prebuilt consumption

Transferred EAGLE metadata buffers are reusable slot views. The previous clone-based mitigation protected correctness but added copies on the transfer hot path and hid the actual lifetime contract. This change makes the transferred request own the metadata slot while it sits in the decode waiting queue, then releases it immediately after process_prebuilt has consumed top-k and hidden state into the prebuilt batch. Abort paths also release any held decode metadata slot.

Constraint: Decode disaggregation metadata buffers are reusable slot views consumed later by process_prebuilt.

Rejected: Clone transferred EAGLE tensors at commit time | correct but less efficient and masks the ownership contract.

Rejected: Release in process_batch_result_prebuilt | holds slots across forward longer than needed.

Confidence: medium

Scope-risk: moderate

Directive: Do not free successful EAGLE transfer metadata in pop_transferred unless process_prebuilt consumption is also moved earlier.

Tested: Remote py_compile for decode.py, scheduler.py, scheduler_output_processor_mixin.py, and test_decode_queue_compaction.py.

Tested: Remote focused lifecycle tests passed: 3 passed.

Tested: Remote full test_decode_queue_compaction.py passed: 11 passed, 5 warnings.

Not-tested: Fresh ETE runtime validation of EAGLE accept-length recovery after the C48 sync.
This commit is contained in:
laoyao0822
2026-05-30 04:51:50 +08:00
parent 07c9544737
commit e9c341afe8
5 changed files with 196 additions and 29 deletions

View File

@@ -2223,3 +2223,103 @@ Fix:
Guardrail:
- Do not replace the clone with a view unless the allocator lifetime is changed so the metadata slot is held until after `process_prebuilt()` consumes the request.
### C46. 2026-05-30 decode EAGLE metadata lifetime must be ownership-based, not clone-based
Correction:
- The C45 clone fix is only a safety workaround. It protects the EAGLE top-k/hidden tensors from metadata-slot reuse, but it adds an extra device/host copy on the transfer hot path and hides the real ownership bug.
- The better contract is: when decode disaggregation commits an EAGLE transfer, the request owns the metadata-buffer slot until decode prebuilt has consumed `req.output_topk_p`, `req.output_topk_index`, and `req.hidden_states_tensor`.
- Therefore `DecodeTransferQueue.pop_transferred()` must not free a successful EAGLE request's `metadata_buffer_index` immediately. It should hand the slot index to the `Req`; `process_batch_result_prebuilt()` should free it after `process_prebuilt()` has already built `EagleDraftInput`.
Rejected:
- Do not keep clone as the final fix. It is correct but less efficient and does not encode the real lifetime contract.
Risk:
- Holding slots can reduce metadata-buffer availability when many transferred requests wait in the scheduler queue. That is expected backpressure because those requests still reference live metadata-buffer rows.
- Non-spec decode does not need the EAGLE tensor views, so it can continue freeing metadata slots at transfer commit.
### C47. 2026-05-30 held-slot implementation and verification status
Implemented:
- `_commit_transfer_to_req()` now stores EAGLE metadata as views and attaches `metadata_buffer_index` to the transferred `Req`.
- `pop_transferred()` keeps successful EAGLE metadata slots allocated; non-spec success and failed/corrupt transfers still free immediately.
- `process_batch_result_prebuilt()` frees any held decode metadata slot after prebuilt has consumed the EAGLE metadata.
- `Scheduler.abort_request()` also frees held decode metadata for requests aborted while waiting before prebuilt.
Verification:
- Remote red tests before implementation failed as expected:
- successful EAGLE transfer incorrectly freed metadata slot immediately
- prebuilt output processing did not free held metadata slot
- Remote container verification after implementation passed:
- py_compile for `decode.py`, `scheduler_output_processor_mixin.py`, `scheduler.py`, and `test_decode_queue_compaction.py`
- focused held-slot tests: `2 passed`
- full `test_decode_queue_compaction.py`: `11 passed, 5 warnings`
Open:
- ETE accept-length recovery still requires a restarted runtime with this code. Do not use logs from a process started before this sync to validate this fix.
### C48. 2026-05-30 decode EAGLE metadata slot release narrowed to consume point
Correction:
- C47 released held decode EAGLE metadata slots in `process_batch_result_prebuilt()`, after the prebuilt batch had already run forward. That is correct but holds metadata slots longer than necessary.
- The real consume point is `get_new_prebuilt_batch()` immediately after `new_batch.process_prebuilt(...)` returns. At that point `process_prebuilt()` has already materialized `EagleDraftInput` from `req.output_topk_p`, `req.output_topk_index`, and `req.hidden_states_tensor`; the reusable disaggregation metadata slot is no longer needed.
Implemented:
- `get_new_prebuilt_batch()` now frees held decode metadata slots in a `finally` block immediately after `process_prebuilt()`.
- `process_batch_result_prebuilt()` no longer owns metadata-slot cleanup.
- The metadata-free helper remains shared for the consume-point release and waiting-queue abort cleanup.
Verification:
- Remote red test first failed because `get_new_prebuilt_batch()` did not free metadata slots after `process_prebuilt()`.
- Remote container verification after implementation passed:
- py_compile for `decode.py`, `scheduler_output_processor_mixin.py`, `scheduler.py`, and `test_decode_queue_compaction.py`
- focused lifecycle tests: `3 passed`
- full `test_decode_queue_compaction.py`: `11 passed, 5 warnings`
Open:
- ETE/runtime validation still needs a fresh process started from this synced code.
### C49. 2026-05-30 remote log observation after C48 sync
Runtime identity:
- Active prefill process on `g0034` was PID `2729613`, started at `2026-05-30 04:36:56 CST`.
- Latest C48 source sync happened after that process started: remote source mtimes around `2026-05-30 04:43:56-04:43:58 CST`.
- Therefore this process cannot validate the C48 consume-point metadata-slot release. It may only validate the earlier code state present at launch.
Latest logs inspected:
- Prefill: `/mnt/beegfs/cjy/log/sglang_cp_hicache_20260529_203656.log`
- Decode0: `/mnt/beegfs/cjy/log/decode0_20260529_203700.log`
- Decode1: `/mnt/beegfs/cjy/log/decode1_20260529_203703.log`
Observed:
- Process is still alive and logs are still moving; no fatal traceback, health-check failure, memory leak, or OOM was found in the latest scan.
- Errors are mostly transfer aborts (`KVTransferError(...): Aborted by AbortReq`) and startup/import warnings.
- Decode accept remains low in this runtime:
- decode0 batch accept len avg `1.511`, p50 `1.36`, p10 `1.02`, p90 `2.19`
- decode1 batch accept len avg `1.457`, p50 `1.30`, p10 `1.02`, p90 `2.16`
- EAGLE debug avg_accept is much lower because it counts accepted draft tokens only: decode0 avg `0.415`, decode1 avg `0.298`; many lines have zero accepted drafts.
- EAGLE metadata presence itself is not missing: sampled debug lines show `has_pd_hidden=True`, `pd_hidden_shape=(6144,)`, `pd_topk_shape=(16,)`.
- Prefill is actively doing owner-lane eviction and CP shared-KV current reuse:
- `owner-lane evict END`: 768 lines
- `forward_partial_current_reuse`: 2688 lines
- `used_prefetch=True`: 1072 lines
- `used_prefetch=False`: 1616 lines
- `create_skip reason=prefix_below_min`: 1472 lines
Interpretation:
- Do not use this process to judge C48.
- If the earlier held-slot fix was active in this process, low accept length is not fully explained by missing EAGLE metadata lifetime alone; next investigation should compare a fresh C48 runtime and then inspect EAGLE state correctness/content, not just tensor presence.

View File

@@ -1175,15 +1175,13 @@ class DecodeTransferQueue:
decode_req.req.output_ids.append(output_id[0].item())
decode_req.req.cached_tokens = cached_tokens[0].item()
if not self.spec_algorithm.is_none():
# ``metadata_buffers.get_buf`` returns views into the reusable
# metadata slots. ``pop_transferred`` frees the slot immediately
# after this commit, while the request can stay in the waiting queue
# before ``process_prebuilt`` consumes the EAGLE state. Keep an
# owned copy so subsequent transfers cannot overwrite the prebuilt
# draft top-k/hidden state.
decode_req.req.output_topk_p = output_topk_p.clone()
decode_req.req.output_topk_index = output_topk_index.clone()
decode_req.req.hidden_states_tensor = output_hidden_states.clone()
# ``metadata_buffers.get_buf`` returns views into reusable metadata
# slots. Keep the slot owned by the request until prebuilt consumes
# the EAGLE state instead of cloning these tensors on the hot path.
decode_req.req.metadata_buffer_index = idx
decode_req.req.output_topk_p = output_topk_p
decode_req.req.output_topk_index = output_topk_index
decode_req.req.hidden_states_tensor = output_hidden_states
_cp_draft_shared_kv_debug(
"decode_transfer_commit rid=%s room=%s metadata_idx=%s cached_tokens=%s "
@@ -1224,7 +1222,7 @@ class DecodeTransferQueue:
)
transferred_reqs = []
completed_decode_reqs = []
decode_reqs_to_free_metadata = []
remaining_queue = []
for decode_req, poll in zip(self.queue, polls):
if rids_to_check is not None and decode_req.req.rid not in rids_to_check:
@@ -1247,14 +1245,13 @@ class DecodeTransferQueue:
)
# release pre-allocated kv cache, but don't insert into the tree since it's failed
release_kv_cache(decode_req.req, self.tree_cache, is_insert=False)
completed_decode_reqs.append(decode_req)
decode_reqs_to_free_metadata.append(decode_req)
if self.scheduler.enable_metrics:
self.scheduler.metrics_collector.increment_transfer_failed_reqs()
continue
elif poll == KVPoll.Success:
should_remove = self._commit_transfer_to_req(decode_req)
if should_remove:
completed_decode_reqs.append(decode_req)
# Check if request was aborted due to corruption
if isinstance(decode_req.req.finished_reason, FINISH_ABORT):
self.scheduler.stream_output(
@@ -1263,10 +1260,13 @@ class DecodeTransferQueue:
release_kv_cache(
decode_req.req, self.tree_cache, is_insert=False
)
decode_reqs_to_free_metadata.append(decode_req)
if self.scheduler.enable_metrics:
self.scheduler.metrics_collector.increment_transfer_failed_reqs()
else:
transferred_reqs.append(decode_req.req)
if self.spec_algorithm.is_none():
decode_reqs_to_free_metadata.append(decode_req)
else:
remaining_queue.append(decode_req)
elif poll in [
@@ -1281,7 +1281,7 @@ class DecodeTransferQueue:
if len(polls) < len(self.queue):
remaining_queue.extend(self.queue[len(polls) :])
for decode_req in completed_decode_reqs:
for decode_req in decode_reqs_to_free_metadata:
idx = decode_req.metadata_buffer_index
assert idx != -1
self.req_to_metadata_buffer_idx_allocator.free(idx)
@@ -1459,7 +1459,11 @@ class SchedulerDisaggregationDecodeMixin:
# construct fake completed prefill
new_batch.prepare_for_prebuilt()
new_batch.process_prebuilt(self.server_args, self.future_map)
try:
new_batch.process_prebuilt(self.server_args, self.future_map)
finally:
for req in can_run_list:
self._free_decode_metadata_index_if_held(req)
return new_batch

View File

@@ -3189,6 +3189,9 @@ class Scheduler(
if self.enable_hisparse:
self.hisparse_coordinator.request_finished(req)
release_kv_cache(req, self.tree_cache)
release_req_to_metadata_buffer(
req, self.req_to_metadata_buffer_idx_allocator
)
# For disaggregation prefill mode, free the metadata buffer index
if self.disaggregation_mode == DisaggregationMode.PREFILL:
release_req_to_metadata_buffer(

View File

@@ -51,6 +51,19 @@ class SchedulerOutputProcessorMixin:
storage_backend_type = type(storage_backend).__name__
return storage_backend_type
def _free_decode_metadata_index_if_held(self: Scheduler, req: Req) -> None:
idx = getattr(req, "metadata_buffer_index", -1)
if idx is None or idx < 0:
return
allocator = getattr(self, "req_to_metadata_buffer_idx_allocator", None)
assert allocator is not None, (
"decode request holds metadata_buffer_index but scheduler has no "
"req_to_metadata_buffer_idx_allocator"
)
allocator.free(idx)
req.metadata_buffer_index = -1
def _maybe_log_eagle_accept_debug(
self: Scheduler,
batch: ScheduleBatch,

View File

@@ -15,6 +15,9 @@ from sglang.srt.disaggregation.decode import (
SchedulerDisaggregationDecodeMixin,
)
from sglang.srt.disaggregation.utils import ReqToMetadataIdxAllocator
from sglang.srt.managers.scheduler_output_processor_mixin import (
SchedulerOutputProcessorMixin,
)
from sglang.srt.managers.schedule_batch import FINISH_ABORT
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
@@ -55,6 +58,12 @@ class FakeReq:
def set_forward_entry_time(self, ts):
self.forward_entry_time = ts
def set_decode_prebuilt_finish_time(self):
return None
def set_quick_finish_time(self):
return None
self.time_stats = FakeTimeStats()
def add_latency(self, stage):
@@ -66,6 +75,12 @@ class FakeReq:
def init_next_round_input(self, tree_cache):
self.init_next_round_calls.append(tree_cache)
def check_finished(self):
return None
def finished(self):
return self.finished_reason is not None
class FakeReceiver:
def __init__(self, should_fail=False):
@@ -340,11 +355,15 @@ class TestDecodeQueueCompaction(CustomTestCase):
self.assertEqual(len(aborted), 1)
self.assertEqual(aborted[0][0], "corrupt")
def test_commit_transfer_clones_eagle_metadata_before_reusing_slot(self):
def test_pop_transferred_holds_eagle_metadata_slot_until_prebuilt_consumes(self):
queue = DecodeTransferQueue.__new__(DecodeTransferQueue)
allocator = FakeAllocator()
output_topk_p = torch.arange(16, dtype=torch.float32)
output_topk_index = torch.arange(16, dtype=torch.int64)
output_hidden_states = torch.arange(8, dtype=torch.float32)
queue.gloo_group = None
queue.req_to_metadata_buffer_idx_allocator = allocator
queue.tp_rank = 0
queue.metadata_buffers = cast(
Any,
SimpleNamespace(
@@ -362,10 +381,13 @@ class TestDecodeQueueCompaction(CustomTestCase):
)
),
)
queue.tree_cache = cast(Any, object())
queue.scheduler = cast(
Any,
SimpleNamespace(
server_args=SimpleNamespace(disaggregation_transfer_backend="mooncake")
server_args=SimpleNamespace(disaggregation_transfer_backend="mooncake"),
stream_output=lambda reqs, return_logprob: None,
enable_metrics=False,
),
)
queue.spec_algorithm = cast(Any, SimpleNamespace(is_none=lambda: False))
@@ -376,33 +398,46 @@ class TestDecodeQueueCompaction(CustomTestCase):
kv_receiver=cast(Any, receiver),
metadata_buffer_index=9,
)
queue.queue = [decode_req]
should_remove = queue._commit_transfer_to_req(decode_req)
self.assertIs(should_remove, True)
with patch(
"sglang.srt.disaggregation.decode.poll_and_all_reduce",
return_value=[KVPoll.Success],
):
transferred = queue.pop_transferred()
output_topk_p.fill_(-1)
output_topk_index.fill_(-2)
output_hidden_states.fill_(-3)
self.assertEqual(transferred, [decode_req.req])
self.assertEqual(queue.queue, [])
self.assertEqual(allocator.freed, [])
self.assertEqual(decode_req.req.output_ids, [7])
self.assertEqual(decode_req.req.cached_tokens, 320)
self.assertEqual(decode_req.req.output_topk_p.tolist(), list(range(16)))
self.assertEqual(decode_req.req.output_topk_index.tolist(), list(range(16)))
self.assertEqual(decode_req.req.metadata_buffer_index, 9)
self.assertEqual(
decode_req.req.hidden_states_tensor.tolist(),
[float(i) for i in range(8)],
)
self.assertNotEqual(
decode_req.req.output_topk_p.data_ptr(), output_topk_p.data_ptr()
)
self.assertNotEqual(
self.assertEqual(
decode_req.req.output_topk_index.data_ptr(), output_topk_index.data_ptr()
)
self.assertNotEqual(
self.assertEqual(
decode_req.req.hidden_states_tensor.data_ptr(),
output_hidden_states.data_ptr(),
)
def test_free_decode_metadata_index_if_held_releases_once(self):
allocator = FakeAllocator()
req = FakeReq("eagle", 3)
req.metadata_buffer_index = 9
scheduler = SchedulerOutputProcessorMixin.__new__(SchedulerOutputProcessorMixin)
scheduler.req_to_metadata_buffer_idx_allocator = allocator
scheduler._free_decode_metadata_index_if_held(req)
scheduler._free_decode_metadata_index_if_held(req)
self.assertEqual(allocator.freed, [9])
self.assertEqual(req.metadata_buffer_index, -1)
def test_resume_retracted_reqs_compacts_queue_in_one_pass(self):
prealloc_queue = DecodePreallocQueue.__new__(DecodePreallocQueue)
prealloc_queue.req_to_token_pool = cast(
@@ -596,6 +631,7 @@ class TestDecodeQueueCompaction(CustomTestCase):
self.assertEqual(queue.queue, [skipped, waiting, blocked, tail])
def test_get_new_prebuilt_batch_slices_waiting_queue_prefix(self):
allocator = FakeAllocator()
scheduler = cast(Any, SimpleNamespace())
scheduler.grammar_manager = SimpleNamespace(
has_waiting_grammars=lambda: False,
@@ -613,6 +649,14 @@ class TestDecodeQueueCompaction(CustomTestCase):
scheduler.spec_algorithm = object()
scheduler.server_args = object()
scheduler.future_map = object()
scheduler.req_to_metadata_buffer_idx_allocator = allocator
scheduler._free_decode_metadata_index_if_held = (
SchedulerOutputProcessorMixin._free_decode_metadata_index_if_held.__get__(
scheduler
)
)
for i, req in enumerate(scheduler.waiting_queue[:3]):
req.metadata_buffer_index = 20 + i
captured = {}
@@ -642,6 +686,9 @@ class TestDecodeQueueCompaction(CustomTestCase):
captured["batch"].processed,
[(scheduler.server_args, scheduler.future_map)],
)
self.assertEqual(allocator.freed, [20, 21, 22])
for req in captured["reqs"]:
self.assertEqual(req.metadata_buffer_index, -1)
def test_get_new_prebuilt_batch_keeps_waiting_queue_when_no_capacity(self):
scheduler = cast(Any, SimpleNamespace())