diff --git a/docs/advanced_features/nsa_prefill_cp_accept_len1_garbage_output_debug_ledger.md b/docs/advanced_features/nsa_prefill_cp_accept_len1_garbage_output_debug_ledger.md index e1667c12f..c9098253d 100644 --- a/docs/advanced_features/nsa_prefill_cp_accept_len1_garbage_output_debug_ledger.md +++ b/docs/advanced_features/nsa_prefill_cp_accept_len1_garbage_output_debug_ledger.md @@ -876,3 +876,138 @@ Post-cleanup verification note: - `_current_index_k_for_tai()` was made robust to PyTorch builds that may not expose every fp8 dtype symbol. - Re-ran remote full `test/registered/unit/layers/test_nsa_cp_utils.py`: `87 passed`. + +## 2026-06-04 full prefill-flow audit pass + +Scope: + +- Re-read the current prefill inference path after `f50e2b1e0` from scheduler admission through target forward, CP metadata, NSA index/MLA current reuse, EAGLE draft handoff, HiCache backup/insert, and Mooncake disaggregation transfer. +- This pass is an audit/ledger update only. No code was changed in this pass. + +Checked path and direct evidence: + +1. Scheduler/disaggregated prefill main loop + - `python/sglang/srt/disaggregation/prefill.py:529-551` receives bootstrapped requests, selects a prefill batch, runs `run_batch()`, then calls `process_batch_result()`. + - `python/sglang/srt/managers/schedule_batch.py:1481-1537` constructs logical extend inputs and calls `alloc_for_extend()` before model forward. + - `python/sglang/srt/managers/scheduler.py:2635-2703` registers CP HiCache write backups after KV slot allocation and immediately before model forward. + - `python/sglang/srt/disaggregation/prefill.py:635-668` updates radix/HiCache, stores EAGLE top-k/hidden side channel, then sends KV chunks. + +2. Target model CP metadata and hidden contract + - `python/sglang/srt/models/deepseek_v2.py:2237-2250` creates `forward_batch.nsa_cp_metadata` only when `can_cp_split(...)` succeeds. + - `python/sglang/srt/models/deepseek_v2.py:2076-2088` captures draft hidden before final target CP output collect, then target output is either narrow last-token collect or full gather/rerange. + - Inference: CP-local target hidden is a temporary side channel for draft model consumption; prefill-to-decode hidden should be batch-major after draft forward, not CP-local. + +3. EAGLE/draft handoff + - `python/sglang/srt/speculative/eagle_worker.py:284-308` enables CP draft shared KV from env/model-runner capability and extend mode; it does not explicitly check that current `forward_batch.nsa_cp_metadata` exists. + - `python/sglang/srt/speculative/eagle_worker.py:332-354` passes `logits_output.draft_hidden_states` as CP-local draft side channel when present. + - `python/sglang/srt/speculative/eagle_worker.py:1069-1075` now clears `cp_local_hidden_states` after `capture_for_decode()` replaces hidden with draft output. + - `python/sglang/srt/models/deepseek_nextn.py:195-307` still has a bs=1 compatibility branch that pads CP-local hidden by up to `cp_size` rows; bs>1 fail-fasts. + +4. CP split / compute padding / direct-write cache loc contract + - `python/sglang/srt/layers/attention/nsa/utils.py:500-755` builds request-first batch CP split metadata. Valid rows and compute-padded rows are tracked separately; compute pages are `max(valid_pages, cp_size)`, not `2*cp_size`. + - `python/sglang/srt/layers/attention/nsa/utils.py:783-1005` supports `split_kind="compute"` vs `split_kind="valid"`; cache writes must use valid rows. + - `python/sglang/srt/layers/attention/nsa/utils.py:1634-1805` direct-write local `out_cache_loc` uses valid split rows and slices only allowed static padding. + - `python/sglang/srt/mem_cache/cp_shared_kv_layout.py:56-101` preserves dummy page 0 and maps usable logical pages from page 1 by CP owner. + - Conclusion: current direct-write/page0/valid-vs-compute contract is internally consistent and is not the strongest current correctness suspect. + +5. NSA index/MLA current and partial-current reuse + - `python/sglang/srt/layers/attention/nsa/nsa_indexer.py:136-153` converts fp8 current-index K to `uint8` for the TAI ABI. + - `python/sglang/srt/layers/attention/nsa/nsa_indexer.py:155-192` builds request bases from CP-local batch-plan offsets when present. + - `python/sglang/srt/layers/attention/nsa/nsa_indexer.py:1916-2135` handles current-only and partial-current top-k paths, with current-index KV expected to be CP-local under batch plan. + - `python/sglang/srt/layers/attention/nsa_backend.py:1802-2125` selects valid current rows for compute padding, supports current-only page-slot compose, and falls back to synchronous partial-current compose when prefetch is absent. + - `python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py:2215-2305` batch prefix slot span may cover a contiguous flattened slot range that includes non-prefix/current slots between requests; those are later overwritten/masked. + - `python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py:1855-2035` TAI IPC materialize fast path only supports `start_slot == 0`; nonzero spans fall back to local materialize/collective path. + +6. Prefetch + - `python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py:334-450` MLA prefetcher creation currently supports only `batch_size == 1` and skips short prefixes/extends. + - `python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py:790-963` materializes the next-layer prefix on the current stream first, then enqueues async all-reduce on the prefetch stream. + - `python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py:1023-1030` attention-window wait is deliberately deferred; consume waits later. + - `python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py:1086-1155` index prefetcher also rejects bs>1. + - Conclusion: bs>1 currently loses async L2->L1 prefetch fast path. That is a performance risk more than the strongest garbage-output suspect. + +7. L1 owner-lane allocator + - `python/sglang/srt/mem_cache/common.py:447-633` requires CP owner page planning for CP shared-KV in-seq split; allocation retries after owner-lane eviction and raises `KVCapacityWaitError` on lane exhaustion. + - `python/sglang/srt/mem_cache/allocator.py:930-1045` allocates pages by CP owner lane and returns full logical `out_cache_loc` in request order. + - `python/sglang/srt/mem_cache/allocator.py:1047-1101` HiCache reload can allocate pages with saved owner pattern. + - Risk is mainly fragmentation/performance: owner buckets are consumed by sorted prefix selection, not by a best contiguous extent search. + +8. HiCache/radix/chunked backup + - `python/sglang/srt/mem_cache/hiradix_cache.py:2365-2388` stale-tail prune raises `HiCachePendingBackupSplit` if the stale subtree is still locked/host-referenced/write-pending. + - `python/sglang/srt/mem_cache/hiradix_cache.py:2444-2510` chunked CP backup is floored to completed pages and skipped when no full page exists. + - `python/sglang/srt/mem_cache/hiradix_cache.py:4068-4209` `insert()` pre-checks pending split state before `_split_node()`, but still calls `_cp_prune_stale_tail_after_page_floor()` after splitting. + - `python/sglang/srt/mem_cache/radix_cache.py:624-675` `cache_unfinished_req()` rolls back prepared backup if `insert()` returns `pending_backup_deferred_node`. + - Remaining risk: post-split stale-tail prune can still raise if the stale tail becomes/continues blocked after the pre-check. Previous chunked crash class is mitigated but not eliminated by construction. + +9. Disaggregation/Mooncake prefill-to-decode transfer + - `python/sglang/srt/disaggregation/prefill.py:870-1005` floors non-last chunk to page boundary, gathers page indices, and sends KV chunks. + - `python/sglang/srt/disaggregation/utils.py:633-651` filters pages by CP owner and returns absolute request page positions. + - `python/sglang/srt/disaggregation/mooncake/conn.py:1414-1453` sender replaces raw page slice with CP-owned logical pages plus logical positions. + - `python/sglang/srt/disaggregation/mooncake/conn.py:1023-1031` receiver maps those logical positions into decode destination pages. + - `python/sglang/srt/disaggregation/mooncake/conn.py:1055-1065` still truncates prefill pages if destination pages are shorter. The comment says this should never happen. + - `python/sglang/srt/disaggregation/mooncake/conn.py:338-450` transfer bandwidth depends on contiguous source/destination groups and whether layer transfers are batched. + - `python/sglang/srt/disaggregation/common/conn.py:50-56` runtime fingerprint covers only a narrow module subset and does not include `cp_shared_kv_runtime.py`, `cp_shared_kv_prefetch.py`, `eagle_worker.py`, `deepseek_nextn.py`, or tai-kernel version. + +Ranked remaining risks after this pass: + +1. Highest correctness risk: Mooncake CP shared-KV transfer mismatch is truncated instead of fail-fast. + - Evidence: `mooncake/conn.py:1055-1065` silently shortens `kv_chunk.prefill_kv_indices` when destination pages are fewer. + - Inference: a prefill/decode page-position mismatch can drop KV pages and still continue, which fits bad decode output symptoms better than a local shape mismatch that would fail-fast. + +2. EAGLE CP-local side-channel gate is not tied to actual metadata creation. + - Evidence: `_can_use_cp_draft_shared_kv()` checks env/model capability, not `forward_batch.nsa_cp_metadata`; target metadata is created only under `can_cp_split(...)`. + - Inference: unusual extend/padding/can-split edge cases can still mark hidden as CP-local without a fully proven CP metadata contract. + - Lower than risk 1 because the stale post-capture marker was fixed and bs>1 hidden mismatches fail-fast. + +3. HiCache chunked post-split stale-tail prune can still escape the deferred-insert path. + - Evidence: `insert()` calls `_cp_prune_stale_tail_after_page_floor()` after `_split_node()` even after pre-checks. + - Inference: chunked prefill plus in-flight backup can still hit uncaught `HiCachePendingBackupSplit` in rare timing windows. + +4. bs>1 partial-current prefix span can fall off TAI IPC and over-communicate. + - Evidence: batch prefix span may include contiguous slots beyond pure prefix; TAI IPC materialize fast path requires `start_slot == 0`. + - Inference: this is primarily a performance/regression risk and a complexity hotspot, not a proven correctness bug because current rows are later overwritten/masked. + +5. bs>1 async prefetch fast path is not implemented. + - Evidence: both MLA and index prefetcher creation skip `batch_size != 1`. + - Inference: short-extend batching can improve compute occupancy, but cache-hit paths may still pay synchronous L2->L1 materialize/collective costs. + +6. Runtime compatibility fingerprint is too narrow. + - Evidence: fingerprint excludes several files that define CP shared-KV/runtime/EAGLE behavior and excludes tai-kernel version. + - Inference: prefill/decode or host/container mixed-version runs can pass bootstrap while using incompatible CP contracts. + +7. Allocation/transfer contiguity is still best-effort. + - Evidence: owner-lane allocator takes sorted bucket prefixes; Mooncake groups contiguous runs and bandwidth depends on group count. + - Inference: this explains low transfer bandwidth under fragmented long-running workloads, but it is not the strongest garbage-output root cause. + +Checked-and-downranked items: + +- Dummy page 0 semantics are explicit and preserved by `CpSharedKVLayout`; current evidence does not point to page0 as the main issue. +- Direct-write local `out_cache_loc` now uses valid split rows and fail-fast checks; current evidence does not point to silent direct-write row misalignment. +- Current-index fp8 dtype/local-base issue was repaired in `f50e2b1e0` and has unit coverage; it remains a runtime item to verify, not the leading code-level suspect in this pass. + +Next discriminating probes/fixes to prioritize: + +1. Replace Mooncake CP shared-KV transfer truncation with fail-fast under CP shared-KV and add a focused unit test for mismatched logical positions vs destination pages. +2. Tie EAGLE CP-local draft enablement to actual `forward_batch.nsa_cp_metadata` / batch plan availability, or fail-fast when the marker is set without metadata. +3. Add a regression for HiCache chunked insertion where a stale tail is blocked by pending backup; expected behavior is deferred insert/rollback, not scheduler crash. +4. Extend runtime fingerprint to include CP shared-KV runtime/prefetch/EAGLE/NextN and tai-kernel version/build identity. +5. For performance work, measure bs>1 partial-current sync compose and transfer contiguous-group counts before changing allocator policy again. + +## 2026-06-04 Mooncake transfer page-count fail-fast implementation + +Change made: + +- Added `validate_transfer_page_count_or_raise()` in `python/sglang/srt/disaggregation/utils.py`. +- Replaced the Mooncake KV transfer truncation workaround in `python/sglang/srt/disaggregation/mooncake/conn.py` with the new validator. +- Added unit coverage in `test/registered/unit/disaggregation/test_cp_shared_kv_transfer_mapping.py`. + +Contract: + +- Source prefill pages and selected decode destination pages must have exactly the same count. +- Under CP shared-KV, a mismatch now raises `[CP_SHARED_KV_FAIL_FAST][mooncake_transfer_page_count_mismatch]` with room, cp_rank, source/destination lengths, logical positions, index slice, and compact page summaries. +- The old behavior truncated `prefill_kv_indices` to the destination length. That is no longer allowed because it can hide prefill/decode page-map corruption and let decode continue with incomplete KV. + +Verification status: + +- Local `py_compile` passed for `disaggregation/utils.py`, `disaggregation/mooncake/conn.py`, and the updated unit test. +- Local pytest is blocked by missing runtime dependency `orjson` in this environment. +- Remote RED/GREEN execution could not be completed in this pass because `ssh g0034` timed out. diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py index 0a531106d..2bbfbd6de 100644 --- a/python/sglang/srt/disaggregation/mooncake/conn.py +++ b/python/sglang/srt/disaggregation/mooncake/conn.py @@ -32,6 +32,7 @@ from sglang.srt.disaggregation.mooncake.utils import ( from sglang.srt.disaggregation.utils import ( DisaggregationMode, filter_kv_indices_for_cp_rank, + validate_transfer_page_count_or_raise, ) from sglang.srt.distributed.parallel_state import get_mooncake_transfer_engine from sglang.srt.environ import envs @@ -1052,17 +1053,18 @@ class MooncakeKVManager(CommonKVManager): kv_chunk.is_last_chunk, ) - # NOTE: This is temporarily a workaround to deal with the case where the prefill_kv_indices - # is mismatched with the dst_kv_indices when page size > 1, this should never happen. - if len(chunked_dst_kv_indice) < len( - kv_chunk.prefill_kv_indices - ): - logger.warning( - f"len(chunked_dst_kv_indice) = {len(chunked_dst_kv_indice)}, len(kv_chunk.prefill_kv_indices) = {len(kv_chunk.prefill_kv_indices)}" - ) - kv_chunk.prefill_kv_indices = kv_chunk.prefill_kv_indices[ - : len(chunked_dst_kv_indice) - ] + validate_transfer_page_count_or_raise( + prefill_indices=kv_chunk.prefill_kv_indices, + dst_indices=chunked_dst_kv_indice, + room=kv_chunk.room, + cp_rank=self.attn_cp_rank, + logical_page_positions=kv_chunk.logical_page_positions, + index_slice=kv_chunk.index_slice, + is_cp_shared_kv=( + kv_chunk.logical_page_positions is not None + ), + path="mooncake_kv", + ) target_rank_registration_info: KVArgsRegisterInfo = ( self.decode_kv_args_table[req.mooncake_session_id] diff --git a/python/sglang/srt/disaggregation/utils.py b/python/sglang/srt/disaggregation/utils.py index c15fff9a3..7e6b83c3f 100644 --- a/python/sglang/srt/disaggregation/utils.py +++ b/python/sglang/srt/disaggregation/utils.py @@ -668,6 +668,72 @@ def select_pages_by_request_positions( return pages[request_positions] +def _transfer_indices_summary(indices: Any) -> str: + if indices is None: + return "None" + arr = np.asarray(indices) + if arr.size == 0: + return f"shape={arr.shape} dtype={arr.dtype} size=0" + head = arr.reshape(-1)[: min(8, arr.size)].tolist() + tail = arr.reshape(-1)[-min(8, arr.size) :].tolist() + return ( + f"shape={arr.shape} dtype={arr.dtype} size={arr.size} " + f"head={head} tail={tail}" + ) + + +def _slice_summary(index_slice: Optional[slice]) -> str: + if index_slice is None: + return "None" + return ( + f"slice(start={index_slice.start}, stop={index_slice.stop}, " + f"step={index_slice.step})" + ) + + +def validate_transfer_page_count_or_raise( + *, + prefill_indices: Any, + dst_indices: Any, + room: Optional[int], + cp_rank: Optional[int], + logical_page_positions: Optional[Any], + index_slice: Optional[slice], + is_cp_shared_kv: bool, + path: str, +) -> None: + """Fail fast if a transfer chunk would drop or invent pages. + + Mooncake used to truncate source pages when the destination page selection + was shorter. Under CP shared-KV this hides a broken logical-page mapping + and lets decode continue with incomplete KV. Keep the check as a small + shared helper so tests can cover the contract without constructing the + transfer worker thread. + """ + + prefill_len = len(prefill_indices) if prefill_indices is not None else 0 + dst_len = len(dst_indices) if dst_indices is not None else 0 + if prefill_len == dst_len: + return + + prefix = ( + "[CP_SHARED_KV_FAIL_FAST]" + if is_cp_shared_kv or logical_page_positions is not None + else "[KV_TRANSFER_FAIL_FAST]" + ) + raise RuntimeError( + f"{prefix}[mooncake_transfer_page_count_mismatch] " + "source and destination page counts must match exactly; " + "continuing would silently drop or mis-map KV pages. " + f"path={path} room={room} cp_rank={cp_rank} " + f"prefill_pages={prefill_len} dst_pages={dst_len} " + f"logical_positions={_transfer_indices_summary(logical_page_positions)} " + f"index_slice={_slice_summary(index_slice)} " + f"prefill_indices={_transfer_indices_summary(prefill_indices)} " + f"dst_indices={_transfer_indices_summary(dst_indices)}" + ) + + ######################### # Misc ######################### diff --git a/test/registered/unit/disaggregation/test_cp_shared_kv_transfer_mapping.py b/test/registered/unit/disaggregation/test_cp_shared_kv_transfer_mapping.py index eceb72a84..db5a3961f 100644 --- a/test/registered/unit/disaggregation/test_cp_shared_kv_transfer_mapping.py +++ b/test/registered/unit/disaggregation/test_cp_shared_kv_transfer_mapping.py @@ -8,6 +8,7 @@ from sglang.srt.disaggregation.utils import ( append_cp_draft_state_buffers, filter_kv_pages_for_cp_shared_kv, select_pages_by_request_positions, + validate_transfer_page_count_or_raise, ) from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout from sglang.test.ci.ci_register import register_cpu_ci @@ -79,6 +80,34 @@ class TestCPSharedKVTransferMapping(unittest.TestCase): self.assertEqual(logical_positions.tolist(), [2, 6]) self.assertEqual(selected_decode_pages.tolist(), [43, 47]) + def test_cp_shared_kv_transfer_page_count_mismatch_fail_fast(self): + with self.assertRaisesRegex( + RuntimeError, + r"\[CP_SHARED_KV_FAIL_FAST\]\[mooncake_transfer_page_count_mismatch\].*room=17.*cp_rank=3.*prefill_pages=3.*dst_pages=2", + ): + validate_transfer_page_count_or_raise( + prefill_indices=np.array([11, 12, 13], dtype=np.int32), + dst_indices=np.array([21, 22], dtype=np.int32), + room=17, + cp_rank=3, + logical_page_positions=np.array([0, 1, 2], dtype=np.int32), + index_slice=None, + is_cp_shared_kv=True, + path="mooncake_kv", + ) + + def test_cp_shared_kv_transfer_page_count_match_passes(self): + validate_transfer_page_count_or_raise( + prefill_indices=np.array([11, 12], dtype=np.int32), + dst_indices=np.array([21, 22], dtype=np.int32), + room=18, + cp_rank=1, + logical_page_positions=np.array([4, 8], dtype=np.int32), + index_slice=None, + is_cp_shared_kv=True, + path="mooncake_kv", + ) + def test_append_cp_draft_state_buffers_registers_nsa_state_as_target_payload(self): kv_args = SimpleNamespace( state_type="nsa",