Preserve transferred EAGLE state past metadata-slot reuse

Decode committed EAGLE top-k and hidden-state tensors as views into reusable metadata-buffer rows. The metadata index is freed immediately after transfer commit, while the request may wait before process_prebuilt consumes the draft state. Under concurrent cache-hit traffic a later transfer can overwrite the same row, leaving output_id copied correctly but EAGLE draft state corrupted, which matches low accept length despite successful KV/state registration. Constraint: Metadata slots are intentionally recycled right after transfer commit for throughput. Rejected: Hold metadata slots until process_prebuilt | larger lifetime change and reduces transfer capacity; cloning the small prebuilt EAGLE state is narrower. Confidence: high Scope-risk: narrow Directive: Do not store reusable metadata-buffer views on Req unless the slot lifetime is extended through all consumers. Tested: Local py_compile for decode.py and test_decode_queue_compaction.py. Tested: Remote g0034 container py_compile for decode.py and test_decode_queue_compaction.py. Tested: Remote g0034 focused clone-lifetime test: 1 passed. Tested: Remote g0034 test_decode_queue_compaction.py: 10 passed, 5 warnings. Not-tested: ETE cache-hit accept-length validation after restarting prefill/decode with this synced code.
2026-05-30 03:36:42 +08:00
parent b7364d23f9
commit 07c9544737
3 changed files with 121 additions and 5 deletions
--- a/docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md
+++ b/docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md
@@ -2188,3 +2188,38 @@ Checked point:
 Rejected:

 - Do not change prebuilt `num_states` to `topk * speculative_num_steps` for the non-multi-layer EAGLE path unless the launch also enables `enable_multi_layer_eagle`.  That would feed extra uninitialized/irrelevant metadata into `EAGLEWorkerV2.draft()`.
+
+### C45. 2026-05-30 new runtime did not recover; root-cause moved to decode metadata-buffer lifetime
+
+Runtime evidence:
+
+- New prefill process PID `2655723` started at `2026-05-30 03:18:38 CST`, after the C42 code sync.
+- Latest logs:
+  - prefill: `/mnt/beegfs/cjy/log/sglang_cp_hicache_20260529_191838.log`
+  - decode0: `/mnt/beegfs/cjy/log/decode0_20260529_191842.log`
+  - decode1: `/mnt/beegfs/cjy/log/decode1_20260529_191845.log`
+- The C42 change is active enough to remove the old fallback: `draft_partial_current_reuse_disabled=0` in the new prefill log.
+- Accept length still did not recover: parsed decode batches show average accept length about `1.35`, median about `1.25`, and many windows at `<=1.05`.  Transfer-free windows remain low as well.
+
+New root cause:
+
+- `DecodeTransferQueue._commit_transfer_to_req()` assigned `req.output_topk_p`, `req.output_topk_index`, and `req.hidden_states_tensor` directly from `metadata_buffers.get_buf(idx)`.
+- `metadata_buffers.get_buf(idx)` returns views into reusable metadata-buffer rows.
+- `DecodeTransferQueue.pop_transferred()` frees `metadata_buffer_index` immediately after `_commit_transfer_to_req()` returns, but the transferred request can sit in the scheduler waiting queue before `process_prebuilt()` consumes the EAGLE top-k/hidden state.
+- Therefore later transfers can overwrite the same metadata slot before decode prebuilt consumes it.  The scalar output id and cached-token count are copied with `.item()`, but the EAGLE draft state was not owned.  This explains why transfer and registration look correct while the first decode draft state is often wrong under concurrent/high-cache-hit traffic.
+
+Fix:
+
+- Clone the EAGLE metadata tensors at commit time before freeing the metadata slot:
+  - `output_topk_p.clone()`
+  - `output_topk_index.clone()`
+  - `output_hidden_states.clone()`
+- Added a unit test that mutates the original metadata-buffer tensors after commit and verifies the request still owns the original values and different storage.
+- Remote container verification on `g0034:/sgl-workspace/sglang-tai` passed:
+  - py_compile for `decode.py` and `test_decode_queue_compaction.py`
+  - focused clone-lifetime test: `1 passed`
+  - full `test_decode_queue_compaction.py`: `10 passed, 5 warnings`
+
+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.
--- a/python/sglang/srt/disaggregation/decode.py
+++ b/python/sglang/srt/disaggregation/decode.py
@@ -1175,9 +1175,15 @@ 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():
-            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
+            # ``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()

        _cp_draft_shared_kv_debug(
            "decode_transfer_commit rid=%s room=%s metadata_idx=%s cached_tokens=%s "
--- a/test/registered/unit/disaggregation/test_decode_queue_compaction.py
+++ b/test/registered/unit/disaggregation/test_decode_queue_compaction.py
@@ -5,6 +5,8 @@ from types import SimpleNamespace
 from typing import Any, cast
 from unittest.mock import patch

+import torch
+
 from sglang.srt.disaggregation.base import KVPoll
 from sglang.srt.disaggregation.decode import (
    DecodePreallocQueue,
@@ -338,6 +340,69 @@ 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):
+        queue = DecodeTransferQueue.__new__(DecodeTransferQueue)
+        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.metadata_buffers = cast(
+            Any,
+            SimpleNamespace(
+                get_buf=lambda idx: (
+                    torch.tensor([7], dtype=torch.int32),
+                    torch.tensor([320], dtype=torch.int32),
+                    None,
+                    None,
+                    None,
+                    None,
+                    output_topk_p,
+                    output_topk_index,
+                    output_hidden_states,
+                    torch.tensor([3], dtype=torch.int64),
+                )
+            ),
+        )
+        queue.scheduler = cast(
+            Any,
+            SimpleNamespace(
+                server_args=SimpleNamespace(disaggregation_transfer_backend="mooncake")
+            ),
+        )
+        queue.spec_algorithm = cast(Any, SimpleNamespace(is_none=lambda: False))
+
+        receiver = FakeReceiver()
+        decode_req = DecodeRequest(
+            req=cast(Any, FakeReq("eagle", 3)),
+            kv_receiver=cast(Any, receiver),
+            metadata_buffer_index=9,
+        )
+
+        should_remove = queue._commit_transfer_to_req(decode_req)
+        self.assertIs(should_remove, True)
+
+        output_topk_p.fill_(-1)
+        output_topk_index.fill_(-2)
+        output_hidden_states.fill_(-3)
+
+        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.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(
+            decode_req.req.output_topk_index.data_ptr(), output_topk_index.data_ptr()
+        )
+        self.assertNotEqual(
+            decode_req.req.hidden_states_tensor.data_ptr(),
+            output_hidden_states.data_ptr(),
+        )
+
    def test_resume_retracted_reqs_compacts_queue_in_one_pass(self):
        prealloc_queue = DecodePreallocQueue.__new__(DecodePreallocQueue)
        prealloc_queue.req_to_token_pool = cast(
@@ -380,6 +445,7 @@ class TestDecodeQueueCompaction(CustomTestCase):
            SimpleNamespace(
                available_size=lambda: 1,
                req_to_token=FakeTensor([[1, 2, 3, 4, 5, 6, 7, 8]]),
+                write=lambda *args, **kwargs: None,
            ),
        )
        queue.req_to_metadata_buffer_idx_allocator = cast(
@@ -387,6 +453,7 @@ class TestDecodeQueueCompaction(CustomTestCase):
        )
        queue.token_to_kv_pool_allocator = cast(Any, SimpleNamespace(page_size=1))
        queue.token_to_kv_pool = cast(Any, object())
+        queue.draft_token_to_kv_pool = None
        queue.num_reserved_decode_tokens = 1
        queue.scheduler = cast(
            Any,
@@ -402,7 +469,10 @@ class TestDecodeQueueCompaction(CustomTestCase):
        queue._allocatable_tokens = (
            lambda retractable_tokens=None, count_retracted=True: 2
        )
-        queue._pre_alloc = lambda req: setattr(req, "req_pool_idx", 0)
+        queue._pre_alloc = lambda req, **kwargs: (
+            setattr(req, "req_pool_idx", 0),
+            torch.arange(len(req.origin_input_ids), dtype=torch.int64),
+        )[1]

        blocked_req = FakeReq("blocked", 1)
        blocked_req.origin_input_ids = [1, 2, 3]
@@ -445,6 +515,7 @@ class TestDecodeQueueCompaction(CustomTestCase):
            SimpleNamespace(
                available_size=lambda: 1,
                req_to_token=FakeTensor([[1, 2, 3, 4, 5, 6, 7, 8]]),
+                write=lambda *args, **kwargs: None,
            ),
        )
        queue.req_to_metadata_buffer_idx_allocator = cast(
@@ -452,6 +523,7 @@ class TestDecodeQueueCompaction(CustomTestCase):
        )
        queue.token_to_kv_pool_allocator = cast(Any, SimpleNamespace(page_size=1))
        queue.token_to_kv_pool = cast(Any, object())
+        queue.draft_token_to_kv_pool = None
        queue.num_reserved_decode_tokens = 1
        queue.scheduler = cast(
            Any,
@@ -467,7 +539,10 @@ class TestDecodeQueueCompaction(CustomTestCase):
        queue._allocatable_tokens = (
            lambda retractable_tokens=None, count_retracted=True: 6
        )
-        queue._pre_alloc = lambda req: setattr(req, "req_pool_idx", 0)
+        queue._pre_alloc = lambda req, **kwargs: (
+            setattr(req, "req_pool_idx", 0),
+            torch.arange(len(req.origin_input_ids), dtype=torch.int64),
+        )[1]

        skipped = DecodeRequest(
            req=cast(Any, FakeReq("skip", 1)),