From 501faa8f5b3cf5aafd89f2c3a5e92397dc2e1c14 Mon Sep 17 00:00:00 2001 From: leavelet Date: Sun, 7 Jun 2026 01:15:23 +0000 Subject: [PATCH] fix(disagg): tensor-safe prefix check + graceful per-layer register (lever A) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit E2E caught it: registered=8 (per-layer activated correctly) then the prefill crashed with "Boolean value of Tensor with no values is ambiguous" — req.prefix_indices is a tensor and `prefix or []` evaluated its truthiness. Use a None+len check (safe for None/list/tensor), and wrap each request's registration in try/except so a setup error degrades to the monolithic transfer instead of crashing the scheduler. Also guard start_send_idx with getattr. Co-Authored-By: Claude Opus 4.8 (1M context) --- python/sglang/srt/disaggregation/prefill.py | 35 ++++++++++++++------- 1 file changed, 23 insertions(+), 12 deletions(-) diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index 585034780..b27c4ec71 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -659,18 +659,29 @@ class SchedulerDisaggregationPrefillMixin: return page_size = self.token_to_kv_pool_allocator.page_size for req in batch.reqs: - if getattr(req, "disagg_kv_sender", None) is None: - continue - if getattr(req, "is_chunked", 0) != 0 or req.start_send_idx != 0: - continue # chunked / partially-sent: this forward isn't the full range - if len(getattr(req, "prefix_indices", []) or []) != 0: - continue # cached prefix is HiCache-loaded, not forward-written (lever B) - 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], - page_size, - ) - kv_manager.register_per_layer_transfer(req.bootstrap_room, page_indices) + try: + if getattr(req, "disagg_kv_sender", None) is None: + 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 + 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], + page_size, + ) + kv_manager.register_per_layer_transfer(req.bootstrap_room, page_indices) + except Exception as e: + # Never let lever-A setup crash the scheduler; fall back to the + # monolithic post-forward transfer for this request. + logger.warning( + "[CP_PER_LAYER_TRANSFER] register skipped for room %s: %r", + getattr(req, "bootstrap_room", None), + e, + ) @torch.no_grad() def event_loop_normal_disagg_prefill(self: Scheduler) -> None: