chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

Expose CP shared-KV transfer page-count mismatches before decode

Browse Source 
Mooncake previously truncated source prefill pages when the selected decode destination page list was shorter. That made an invalid prefill/decode page mapping continue as an incomplete KV transfer, which can surface later as decode garbage instead of the original mapping error.\n\nThis changes the transfer contract to require exact source/destination page-count equality and records compact page summaries in the fail-fast error. The helper is shared so the page-count contract can be unit-tested without constructing the transfer worker thread.\n\nConstraint: CP shared-KV page ownership requires a one-to-one prefill source page to decode destination page mapping.\nRejected: Keep warning-and-truncate | masks page-map corruption and can silently drop KV.\nConfidence: medium\nScope-risk: moderate\nDirective: Do not restore source-page truncation; fix the mapping producer if this fail-fast triggers.\nTested: git diff --check; python -m py_compile python/sglang/srt/disaggregation/utils.py python/sglang/srt/disaggregation/mooncake/conn.py test/registered/unit/disaggregation/test_cp_shared_kv_transfer_mapping.py\nNot-tested: pytest blocked locally by missing orjson; remote g0034 was unavailable during this pass.
This commit is contained in:
laoyao0822
2026-06-04 23:14:22 +08:00
parent f50e2b1e00
commit e200091638
4 changed files with 243 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

135
docs/advanced_features/nsa_prefill_cp_accept_len1_garbage_output_debug_ledger.md
View File

@@ -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.

24
python/sglang/srt/disaggregation/mooncake/conn.py
View File

@@ -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]

66
python/sglang/srt/disaggregation/utils.py
View File

@@ -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
#########################

29
test/registered/unit/disaggregation/test_cp_shared_kv_transfer_mapping.py
View File

@@ -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",