From 5791684864b8a064a76f5fbcfd5b065a677056b5 Mon Sep 17 00:00:00 2001 From: leavelet Date: Sun, 7 Jun 2026 04:30:33 +0000 Subject: [PATCH] feat(disagg): broaden lever A to full sequence incl cached prefix (high cache-hit) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The premise holds at realistic context: at ~8k tokens + high cache-hit the KV transfer is ~30% of TTFT (712ms), sequential after the forward. Lever A was scoped to prefix==0 so it never activated there. Broaden it to cover the FULL sequence (cached prefix + new tokens): the per-layer backup hook fires after each layer's full processing, so the HiCache-loaded prefix KV and forward-written new KV are both final — transferring all of layer L's pages then overlaps the dominant prefix transfer with the forward. Add a sonnet-bench verify mode for output-equality. Gated by the flag; correctness verified next. Co-Authored-By: Claude Opus 4.8 (1M context) --- python/sglang/srt/disaggregation/prefill.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index 3d541d522..4de1bfdb7 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -666,10 +666,12 @@ class SchedulerDisaggregationPrefillMixin: continue if getattr(req, "is_chunked", 0) != 0 or getattr(req, "start_send_idx", 0) != 0: continue # chunked / partially-sent: this forward isn't the full range - prefix = getattr(req, "prefix_indices", None) - # NB: prefix_indices is a tensor — never use `or`/truthiness on it. - if prefix is not None and len(prefix) != 0: - continue # cached prefix is HiCache-loaded, not forward-written -> lever B + # Cover the FULL sequence (cached prefix + new tokens). The per-layer + # backup hook fires after each layer's full processing, by which point + # layer L's HiCache-loaded prefix KV and forward-written new KV are both + # final — so transferring all of layer L's pages then is correct, and it + # overlaps the (dominant, at high cache-hit) prefix transfer with the + # forward. Verified by output-equality. end_idx = min(len(req.fill_ids), len(req.origin_input_ids)) page_indices = _kv_locs_to_page_indices_cpu( self.req_to_token_pool.req_to_token[req.req_pool_idx, 0:end_idx],