Preserve draft KV across CP HiCache hits
Cache-hit prefill can skip draft forward for the prefix while PD transfer still reads draft KV for that same prefix. CP HiCache therefore needs to persist draft/MTP KV alongside target KV instead of relying on whatever remains in the draft GPU pool. Constraint: CP HiCache is host-only here; storage backends remain unsupported for CP shared KV. Constraint: CP shared KV must keep owner-page semantics and avoid falling back to full KV on every rank. Rejected: Recompute cached-prefix draft KV during prefill | loses the HiCache benefit and reintroduces the large hidden/KV footprint. Rejected: Change PD transfer to skip draft prefix KV | decode still needs draft cache continuity for MTP acceptance. Confidence: medium Scope-risk: moderate Directive: Keep target and draft CP HiCache metadata/load/write/evict paths in lockstep; changing one without the other can silently reduce MTP accept length. Tested: Remote g0034 container /sgl-workspace/sglang-tai: python3 -m pytest -q test/registered/unit/managers/test_hicache_controller_cp.py test/registered/unit/mem_cache/test_cp_hicache_metadata.py => 58 passed, 3 warnings Not-tested: Full multi-node HiCache+MTP serving benchmark and accept-length recovery.
This commit is contained in:
@@ -853,9 +853,18 @@ class SchedulerDisaggregationPrefillMixin:
|
||||
)
|
||||
return
|
||||
prefill_queue = getattr(self, "disagg_prefill_bootstrap_queue", None)
|
||||
has_draft_pool = (
|
||||
getattr(prefill_queue, "draft_token_to_kv_pool", None) is not None
|
||||
)
|
||||
prefix_len = len(getattr(req, "prefix_indices", ()))
|
||||
host_hit_length = int(getattr(req, "host_hit_length", 0) or 0)
|
||||
draft_prefix_overlap = max(0, min(end_idx, prefix_len) - start_idx)
|
||||
_cp_draft_shared_kv_debug(
|
||||
"prefill_send_kv_chunk rid=%s room=%s start_idx=%s end_idx=%s "
|
||||
"last_chunk=%s page_size=%s pages=%s state_pages=%s has_draft_pool=%s",
|
||||
"last_chunk=%s page_size=%s pages=%s state_pages=%s "
|
||||
"has_draft_pool=%s prefix_len=%s host_hit_length=%s "
|
||||
"cache_protected_len=%s extend_input_len=%s fill_len=%s "
|
||||
"origin_input_len=%s already_computed=%s draft_prefix_overlap=%s",
|
||||
req.rid,
|
||||
req.bootstrap_room,
|
||||
start_idx,
|
||||
@@ -864,6 +873,27 @@ class SchedulerDisaggregationPrefillMixin:
|
||||
page_size,
|
||||
_seq_summary(page_indices),
|
||||
_seq_summary(state_indices),
|
||||
getattr(prefill_queue, "draft_token_to_kv_pool", None) is not None,
|
||||
has_draft_pool,
|
||||
prefix_len,
|
||||
host_hit_length,
|
||||
getattr(req, "cache_protected_len", None),
|
||||
getattr(req, "extend_input_len", None),
|
||||
len(req.fill_ids),
|
||||
len(req.origin_input_ids),
|
||||
getattr(req, "already_computed", None),
|
||||
draft_prefix_overlap,
|
||||
)
|
||||
if has_draft_pool and draft_prefix_overlap > 0:
|
||||
_cp_draft_shared_kv_debug(
|
||||
"prefill_send_cachehit_draft_prefix rid=%s room=%s "
|
||||
"draft_prefix_overlap=%s prefix_len=%s host_hit_length=%s "
|
||||
"start_idx=%s end_idx=%s note=transfer_reads_draft_pool_for_cached_prefix",
|
||||
req.rid,
|
||||
req.bootstrap_room,
|
||||
draft_prefix_overlap,
|
||||
prefix_len,
|
||||
host_hit_length,
|
||||
start_idx,
|
||||
end_idx,
|
||||
)
|
||||
req.disagg_kv_sender.send(page_indices, state_indices)
|
||||
|
||||
Reference in New Issue
Block a user