diff --git a/docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md b/docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md index d5162d75e..cafcdcee9 100644 --- a/docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md +++ b/docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md @@ -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. diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index 83ce171ec..ee0e283d7 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -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 diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 5358a8f77..643bf3da7 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -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( diff --git a/python/sglang/srt/managers/scheduler_output_processor_mixin.py b/python/sglang/srt/managers/scheduler_output_processor_mixin.py index db9e8077b..5ab6661f4 100644 --- a/python/sglang/srt/managers/scheduler_output_processor_mixin.py +++ b/python/sglang/srt/managers/scheduler_output_processor_mixin.py @@ -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, diff --git a/test/registered/unit/disaggregation/test_decode_queue_compaction.py b/test/registered/unit/disaggregation/test_decode_queue_compaction.py index e9e11b86d..fb507afb1 100644 --- a/test/registered/unit/disaggregation/test_decode_queue_compaction.py +++ b/test/registered/unit/disaggregation/test_decode_queue_compaction.py @@ -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())